From c4e874e97db6e11737e076675dd1e112e10ea7f5 Mon Sep 17 00:00:00 2001 From: Evan Quiney Date: Wed, 25 Feb 2026 17:44:15 +0000 Subject: [PATCH] skip nan logprobs on tokens (#1622) sometimes we generate NaN logprobs for tokens, this causes pydantic validation errors on the receiving end. in this case we should just not send the logprob items --- src/exo/worker/engines/mlx/generator/generate.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/exo/worker/engines/mlx/generator/generate.py b/src/exo/worker/engines/mlx/generator/generate.py index 6e926205..e7307b8d 100644 --- a/src/exo/worker/engines/mlx/generator/generate.py +++ b/src/exo/worker/engines/mlx/generator/generate.py @@ -1,3 +1,4 @@ +import math import time from copy import deepcopy from typing import Callable, Generator, cast, get_args @@ -248,6 +249,9 @@ def extract_top_logprobs( for i in range(top_logprobs): token_id = int(top_indices[i].item()) token_logprob = float(top_values[i].item()) + if math.isnan(token_logprob): + continue + # Decode token ID to string token_str = tokenizer.decode([token_id]) # Get byte representation