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

feat(router): add support for return_full_text

Browse Source
This commit is contained in:
OlivierDehaene 2023-02-27 17:08:16 +01:00
parent 1cddcbbc26
commit ed22912676
2 changed files with 38 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
router/src/lib.rs
View File

@ -40,13 +40,16 @@ pub(crate) struct GenerateParameters {
example = 0.95
)]
pub top_p: Option<f32>,
#[serde(default = "default_do_sample")]
#[serde(default)]
#[schema(default = "false", example = true)]
pub do_sample: bool,
#[serde(default = "default_max_new_tokens")]
#[schema(exclusive_minimum = 0, exclusive_maximum = 512, default = "20")]
pub max_new_tokens: u32,
#[serde(default)]
#[schema(default = "false", example = false)]
pub return_full_text: bool,
#[serde(default)]
#[schema(inline, max_items = 4, example = json ! (["photographer"]))]
pub stop: Vec<String>,
#[serde(default)]
@ -56,10 +59,6 @@ pub(crate) struct GenerateParameters {
pub seed: Option<u64>,
}
fn default_do_sample() -> bool {
false
}
fn default_max_new_tokens() -> u32 {
20
}
@ -70,8 +69,9 @@ fn default_parameters() -> GenerateParameters {
repetition_penalty: None,
top_k: None,
top_p: None,
do_sample: default_do_sample(),
do_sample: false,
max_new_tokens: default_max_new_tokens(),
return_full_text: false,
stop: vec![],
details: false,
seed: None,

39
router/src/server.rs
View File

@ -75,6 +75,7 @@ async fn health(infer: Extension<Infer>) -> Result<(), (StatusCode, Json<ErrorRe
top_p: None,
do_sample: false,
max_new_tokens: 1,
return_full_text: false,
stop: Vec::new(),
details: false,
seed: None,
@ -120,8 +121,14 @@ async fn generate(
let span = tracing::Span::current();
let start_time = Instant::now();
// Inference
let mut add_prompt = None;
if req.0.parameters.return_full_text {
add_prompt = Some(req.0.inputs.clone());
}
let details = req.0.parameters.details;
// Inference
let response = infer.generate(req.0).await?;
// Token details
@ -188,8 +195,13 @@ async fn generate(
);
// Send response
let mut output_text = response.generated_text.text;
if let Some(prompt) = add_prompt {
output_text = prompt + &output_text;
}
let response = GenerateResponse {
generated_text: response.generated_text.text,
generated_text: output_text,
details,
};
Ok((headers, Json(response)))
@ -240,6 +252,11 @@ async fn generate_stream(
// Inference
let mut end_reached = false;
let mut error = false;
let mut add_prompt = None;
if req.0.parameters.return_full_text {
add_prompt = Some(req.0.inputs.clone());
}
let details = req.0.parameters.details;
match infer.generate_stream(req.0).instrument(info_span!(parent: &span, "async_stream")).await {
@ -306,20 +323,28 @@ async fn generate_stream(
// StreamResponse
end_reached = true;
let mut output_text = generated_text.text;
if let Some(prompt) = add_prompt {
output_text = prompt + &output_text;
}
let stream_token = StreamResponse {
token,
generated_text: Some(generated_text.text),
generated_text: Some(output_text),
details
};
yield Ok(Event::default().json_data(stream_token).unwrap())
yield Ok(Event::default().json_data(stream_token).unwrap());
break;
}
}
}
// yield error
Err(err) => {
error = true;
yield Ok(Event::from(err))
yield Ok(Event::from(err));
break;
}
}
}
@ -327,7 +352,7 @@ async fn generate_stream(
// yield error
Err(err) => {
error = true;
yield Ok(Event::from(err))
yield Ok(Event::from(err));
}
}
// Check if generation reached the end
@ -336,7 +361,7 @@ async fn generate_stream(
let err = InferError::IncompleteGeneration;
metrics::increment_counter!("tgi_request_failure", "err" => "incomplete");
tracing::error!("{err}");
yield Ok(Event::from(err))
yield Ok(Event::from(err));
}
};