huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-04-29 05:52:11 +00:00
Code Issues Packages Projects Releases Wiki Activity

Only add token when it is defined.

Browse Source
This commit is contained in:
Nicolas Patry 2025-03-05 11:16:55 +01:00
parent 491ed9e11d
commit 6777b1075d
No known key found for this signature in database
GPG Key ID: 4242CEF24CB6DBF9
1 changed files with 18 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

21
router/src/server.rs
View File

@ -1287,6 +1287,20 @@ pub(crate) async fn chat_completions(
match result{ match result{
Ok(stream_token) => { Ok(stream_token) => {
let token_text = &stream_token.token.text.clone(); let token_text = &stream_token.token.text.clone();
<<<<<<< Updated upstream
=======
let usage = stream_token.details.as_ref().map(|details| {
let completion_tokens = details.generated_tokens;
let prompt_tokens = details.input_length;
let total_tokens = prompt_tokens + completion_tokens;
Usage {
completion_tokens,
prompt_tokens,
total_tokens,
}
});
>>>>>>> Stashed changes
match state { match state {
StreamState::Buffering => { StreamState::Buffering => {
json_buffer.push_str(&token_text.replace(" ", "")); json_buffer.push_str(&token_text.replace(" ", ""));
@ -1711,9 +1725,10 @@ pub async fn run(
// Shared API builder initialization // Shared API builder initialization
let api_builder = || { let api_builder = || {
let mut builder = ApiBuilder::new() let mut builder = ApiBuilder::new().with_progress(false);
.with_progress(false) if let Some(token) = authorization_token {
.with_token(authorization_token); builder = builder.with_token(Some(token));
}
if let Ok(cache_dir) = std::env::var("HUGGINGFACE_HUB_CACHE") { if let Ok(cache_dir) = std::env::var("HUGGINGFACE_HUB_CACHE") {
builder = builder.with_cache_dir(cache_dir.into()); builder = builder.with_cache_dir(cache_dir.into());