huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-09-11 12:24:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix: adjust logprob response logic

Browse Source
This commit is contained in:
drbh 2024-03-28 00:00:55 +00:00
parent 6c4496a1a3
commit 01ebb77d12
2 changed files with 22 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

32
router/src/lib.rs
View File

@ -385,23 +385,33 @@ impl From<(Token, Vec<Token>)> for ChatCompletionLogprobs {
impl From<(Vec<Token>, Vec<Vec<Token>>)> for ChatCompletionLogprobs {
fn from(value: (Vec<Token>, Vec<Vec<Token>>)) -> 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 }
}
}

2
router/src/server.rs
View File

@ -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(),
},
};