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

add return_full_text support

Browse Source
This commit is contained in:
OlivierDehaene 2023-02-27 19:22:09 +01:00
parent ed22912676
commit f3f9faca2f
2 changed files with 114 additions and 113 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
router/src/lib.rs
View File

@ -47,8 +47,8 @@ pub(crate) struct GenerateParameters {
#[schema(exclusive_minimum = 0, exclusive_maximum = 512, default = "20")] #[schema(exclusive_minimum = 0, exclusive_maximum = 512, default = "20")]
pub max_new_tokens: u32, pub max_new_tokens: u32,
#[serde(default)] #[serde(default)]
#[schema(default = "false", example = false)] #[schema(default = "None", example = false)]
pub return_full_text: bool, pub return_full_text: Option<bool>,
#[serde(default)] #[serde(default)]
#[schema(inline, max_items = 4, example = json ! (["photographer"]))] #[schema(inline, max_items = 4, example = json ! (["photographer"]))]
pub stop: Vec<String>, pub stop: Vec<String>,
@ -71,7 +71,7 @@ fn default_parameters() -> GenerateParameters {
top_p: None, top_p: None,
do_sample: false, do_sample: false,
max_new_tokens: default_max_new_tokens(), max_new_tokens: default_max_new_tokens(),
return_full_text: false, return_full_text: None,
stop: vec![], stop: vec![],
details: false, details: false,
seed: None, seed: None,

43
router/src/server.rs
View File

@ -29,31 +29,27 @@ use utoipa_swagger_ui::SwaggerUi;
/// Compatibility route with api-inference and AzureML /// Compatibility route with api-inference and AzureML
#[instrument(skip(infer))] #[instrument(skip(infer))]
async fn compat_generate( async fn compat_generate(
return_full_text: Extension<bool>, default_return_full_text: Extension<bool>,
infer: Extension<Infer>, infer: Extension<Infer>,
req: Json<CompatGenerateRequest>, req: Json<CompatGenerateRequest>,
) -> Result<impl IntoResponse, (StatusCode, Json<ErrorResponse>)> { ) -> Result<impl IntoResponse, (StatusCode, Json<ErrorResponse>)> {
// switch on stream // switch on stream
let req = req.0; let mut req = req.0;
if req.stream { if req.parameters.return_full_text.is_none() {
Ok(generate_stream(infer, Json(req.into())) req.parameters.return_full_text = Some(default_return_full_text.0)
.await
.into_response())
} else {
let mut add_prompt = None;
if return_full_text.0 {
add_prompt = Some(req.inputs.clone());
} }
let (headers, generation) = generate(infer, Json(req.into())).await?; if req.stream {
Ok(
let mut generation = generation.0; generate_stream(infer, Json(req.into()))
if let Some(prompt) = add_prompt { .await
generation.generated_text = prompt + &generation.generated_text; .into_response(),
}; )
} else {
let (headers, generation) =
generate(infer, Json(req.into())).await?;
// wrap generation inside a Vec to match api-inference // wrap generation inside a Vec to match api-inference
Ok((headers, Json(vec![generation])).into_response()) Ok((headers, Json(vec![generation.0])).into_response())
} }
} }
@ -75,7 +71,7 @@ async fn health(infer: Extension<Infer>) -> Result<(), (StatusCode, Json<ErrorRe
top_p: None, top_p: None,
do_sample: false, do_sample: false,
max_new_tokens: 1, max_new_tokens: 1,
return_full_text: false, return_full_text: None,
stop: Vec::new(), stop: Vec::new(),
details: false, details: false,
seed: None, seed: None,
@ -122,7 +118,12 @@ async fn generate(
let start_time = Instant::now(); let start_time = Instant::now();
let mut add_prompt = None; let mut add_prompt = None;
if req.0.parameters.return_full_text { if req
.0
.parameters
.return_full_text
.unwrap_or(false)
{
add_prompt = Some(req.0.inputs.clone()); add_prompt = Some(req.0.inputs.clone());
} }
@ -254,7 +255,7 @@ async fn generate_stream(
let mut error = false; let mut error = false;
let mut add_prompt = None; let mut add_prompt = None;
if req.0.parameters.return_full_text { if req.0.parameters.return_full_text.unwrap_or(false) {
add_prompt = Some(req.0.inputs.clone()); add_prompt = Some(req.0.inputs.clone());
} }
let details = req.0.parameters.details; let details = req.0.parameters.details;