huggingface/text-generation-inference
5
0
Fork 0
mirror of https://github.com/huggingface/text-generation-inference.git synced 2025-04-24 08:22:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat(backend): handle all the tokenization failure and send back to the client

Browse Source
This commit is contained in:
Morgan Funtowicz 2024-11-06 17:46:46 +01:00
parent 20652824d9
commit 26d0266cec
1 changed files with 36 additions and 33 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

37
backends/llamacpp/src/backend.rs
View File

@ -124,29 +124,31 @@ fn llama_generate_callback(
let token = match ctx.tokenizer.decode(&[new_token_id], false) { let token = match ctx.tokenizer.decode(&[new_token_id], false) {
Ok(text) => { Ok(text) => {
let special = ctx.tokenizer.get_added_vocabulary().is_special_token(&text); let special = ctx.tokenizer.get_added_vocabulary().is_special_token(&text);
Token { Ok(Token {
id: new_token_id, id: new_token_id,
text, text,
logprob: new_token_logit, logprob: new_token_logit,
special, special,
})
} }
} Err(ref err) => Err(InferError::GenerationError(err.to_string())),
Err(_) => panic!("Failed to decode token"),
}; };
// Create the streamed response // Create the streamed response
let response = match is_final { let response = match token {
false => InferStreamResponse::Intermediate { Ok(token) => {
match is_final {
false => Ok(InferStreamResponse::Intermediate {
token, token,
top_tokens: vec![], top_tokens: vec![],
}, }),
true => { true => {
// Decode the whole text // Decode the whole text
match ctx match ctx
.tokenizer .tokenizer
.decode(&ctx.generation.generated_tokens, false) .decode(&ctx.generation.generated_tokens, false)
{ {
Ok(text) => InferStreamResponse::End { Ok(text) => Ok(InferStreamResponse::End {
token, token,
top_tokens: vec![], top_tokens: vec![],
generated_text: GeneratedText { generated_text: GeneratedText {
@ -157,23 +159,24 @@ fn llama_generate_callback(
}, },
start: ctx.start, start: ctx.start,
queued: ctx.start, queued: ctx.start,
}, }),
Err(_) => panic!("Failed to decode token"), Err(err) => Err(InferError::GenerationError(err.to_string())),
} }
// Stream end response
} }
}
}
Err(err) => Err(err),
}; };
// Send back to the client // Send back to the client
if let Err(ref _err) = ctx.stream.send(Ok(response)) { let should_stop = if let Err(ref _err) = ctx.stream.send(response) {
error!("Failed to send back the response to the client, cancelling request"); error!("Failed to send back the response to the client, cancelling request");
// TODO: cancel the request true
return true; // should_stop } else {
} true
};
// should_stop should_stop
false
} }
unsafe fn scheduler_loop( unsafe fn scheduler_loop(