diff --git a/router/src/lib.rs b/router/src/lib.rs index b231a31f..5415a956 100644 --- a/router/src/lib.rs +++ b/router/src/lib.rs @@ -385,23 +385,33 @@ impl From<(Token, Vec)> for ChatCompletionLogprobs { impl From<(Vec, Vec>)> for ChatCompletionLogprobs { fn from(value: (Vec, Vec>)) -> Self { let (tokens, top_tokens) = value; - Self { - content: tokens - .into_iter() - .zip(top_tokens) - .map(|(t, top_t)| ChatCompletionLogprob { - token: t.text, - logprob: t.logprob, - top_logprobs: top_t + + // Create an iterator that produces None for top_tokens once it's exhausted + let top_tokens_iter = top_tokens + .into_iter() + .map(Some) + .chain(std::iter::repeat(None)); + + let content = tokens + .into_iter() + .zip(top_tokens_iter) + .map(|(t, top_t_option)| ChatCompletionLogprob { + token: t.text, + logprob: t.logprob, + top_logprobs: match top_t_option { + Some(top_t) => top_t .into_iter() .map(|t| ChatCompletionTopLogprob { token: t.text, logprob: t.logprob, }) .collect(), - }) - .collect(), - } + None => vec![], // Handle the case where there are no top tokens + }, + }) + .collect(); + + Self { content } } } diff --git a/router/src/server.rs b/router/src/server.rs index 652bc699..b328a8e8 100644 --- a/router/src/server.rs +++ b/router/src/server.rs @@ -857,7 +857,7 @@ async fn chat_completions( details: true, decoder_input_details: !stream, seed, - top_n_tokens: None, + top_n_tokens: req.top_logprobs, grammar: tool_grammar.clone(), }, };