feat(router): refactor API and add openAPI schemas (#53)

router/src/lib.rs

 /// Text Generation Inference Webserver
-
 mod infer;
 mod queue;
 pub mod server;
@@ -8,45 +7,55 @@ mod validation;
 use infer::Infer;
 use queue::{Entry, Queue};
 use serde::{Deserialize, Serialize};
+use utoipa::ToSchema;
 use validation::Validation;
 
-#[derive(Clone, Debug, Deserialize)]
+#[derive(Clone, Debug, Deserialize, ToSchema)]
 pub(crate) struct GenerateParameters {
-    #[serde(default = "default_temperature")]
-    pub temperature: f32,
-    #[serde(default = "default_repetition_penalty")]
-    pub repetition_penalty: f32,
-    #[serde(default = "default_top_k")]
-    pub top_k: i32,
-    #[serde(default = "default_top_p")]
-    pub top_p: f32,
+    #[serde(default)]
+    #[schema(
+        exclusive_minimum = 0.0,
+        nullable = true,
+        default = "null",
+        example = 0.5
+    )]
+    pub temperature: Option<f32>,
+    #[serde(default)]
+    #[schema(
+        exclusive_minimum = 0.0,
+        nullable = true,
+        default = "null",
+        example = 1.03
+    )]
+    pub repetition_penalty: Option<f32>,
+    #[serde(default)]
+    #[schema(exclusive_minimum = 0, nullable = true, default = "null", example = 10)]
+    pub top_k: Option<i32>,
+    #[serde(default)]
+    #[schema(
+        exclusive_minimum = 0.0,
+        maximum = 1.0,
+        nullable = true,
+        default = "null",
+        example = 0.95
+    )]
+    pub top_p: Option<f32>,
     #[serde(default = "default_do_sample")]
+    #[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(inline, max_items = 4, example = json!(["photographer"]))]
     pub stop: Vec<String>,
     #[serde(default)]
+    #[schema(default = "true")]
     pub details: bool,
     #[serde(default)]
     pub seed: Option<u64>,
 }
 
-fn default_temperature() -> f32 {
-    1.0
-}
-fn default_repetition_penalty() -> f32 {
-    1.0
-}
-
-fn default_top_k() -> i32 {
-    0
-}
-
-fn default_top_p() -> f32 {
-    1.0
-}
-
 fn default_do_sample() -> bool {
     false
 }
@@ -57,10 +66,10 @@ fn default_max_new_tokens() -> u32 {
 
 fn default_parameters() -> GenerateParameters {
     GenerateParameters {
-        temperature: default_temperature(),
-        repetition_penalty: default_repetition_penalty(),
-        top_k: default_top_k(),
-        top_p: default_top_p(),
+        temperature: None,
+        repetition_penalty: None,
+        top_k: None,
+        top_p: None,
         do_sample: default_do_sample(),
         max_new_tokens: default_max_new_tokens(),
         stop: vec![],
@@ -69,42 +78,77 @@ fn default_parameters() -> GenerateParameters {
     }
 }
 
-#[derive(Clone, Debug, Deserialize)]
+#[derive(Clone, Debug, Deserialize, ToSchema)]
 pub(crate) struct GenerateRequest {
+    #[schema(example = "My name is Olivier and I")]
     pub inputs: String,
     #[serde(default = "default_parameters")]
     pub parameters: GenerateParameters,
 }
 
-#[derive(Debug, Serialize)]
-pub struct Token(u32, String, f32);
+#[derive(Debug, Serialize, ToSchema)]
+pub struct Token {
+    #[schema(example = 0)]
+    id: u32,
+    #[schema(example = "test")]
+    text: String,
+    #[schema(nullable = true, example = -0.34)]
+    logprob: f32,
+}
+
+#[derive(Serialize, ToSchema)]
+#[serde(rename_all(serialize = "snake_case"))]
+pub(crate) enum FinishReason {
+    #[schema(rename = "length")]
+    Length,
+    #[serde(rename = "eos_token")]
+    #[schema(rename = "eos_token")]
+    EndOfSequenceToken,
+    #[schema(rename = "stop_sequence")]
+    StopSequence,
+}
 
-#[derive(Serialize)]
+#[derive(Serialize, ToSchema)]
 pub(crate) struct Details {
-    pub finish_reason: String,
+    #[schema(example = "length")]
+    pub finish_reason: FinishReason,
+    #[schema(example = 1)]
     pub generated_tokens: u32,
+    #[schema(example = 42)]
     pub seed: Option<u64>,
-    #[serde(skip_serializing_if = "Option::is_none")]
     pub prefill: Option<Vec<Token>>,
-    #[serde(skip_serializing_if = "Option::is_none")]
     pub tokens: Option<Vec<Token>>,
 }
 
-#[derive(Serialize)]
+#[derive(Serialize, ToSchema)]
 pub(crate) struct GenerateResponse {
+    #[schema(example = "test")]
     pub generated_text: String,
     #[serde(skip_serializing_if = "Option::is_none")]
     pub details: Option<Details>,
 }
 
-#[derive(Serialize)]
+#[derive(Serialize, ToSchema)]
+pub(crate) struct StreamDetails {
+    #[schema(example = "length")]
+    pub finish_reason: FinishReason,
+    #[schema(example = 1)]
+    pub generated_tokens: u32,
+    #[schema(example = 42)]
+    pub seed: Option<u64>,
+}
+
+#[derive(Serialize, ToSchema)]
 pub(crate) struct StreamResponse {
     pub token: Token,
+    #[schema(nullable = true, default = "null", example = "test")]
     pub generated_text: Option<String>,
-    pub details: Option<Details>,
+    #[schema(nullable = true, default = "null")]
+    pub details: Option<StreamDetails>,
 }
 
-#[derive(Serialize)]
+#[derive(Serialize, ToSchema)]
 pub(crate) struct ErrorResponse {
+    #[schema(inline)]
     pub error: String,
 }

router/src/server.rs

 /// HTTP Server logic
 use crate::infer::{InferError, InferStreamResponse};
 use crate::{
-    Details, ErrorResponse, GenerateParameters, GenerateRequest, GenerateResponse, Infer,
-    StreamResponse, Validation,
+    Details, ErrorResponse, FinishReason, GenerateParameters, GenerateRequest, GenerateResponse,
+    Infer, StreamDetails, StreamResponse, Token, Validation,
 };
 use axum::extract::Extension;
 use axum::http::{HeaderMap, StatusCode};
@@ -19,6 +19,8 @@ use tokio::signal;
 use tokio::time::Instant;
 use tokio_stream::StreamExt;
 use tracing::instrument;
+use utoipa::OpenApi;
+use utoipa_swagger_ui::SwaggerUi;
 
 /// Health check method
 #[instrument(skip(infer))]
@@ -32,13 +34,13 @@ async fn health(infer: Extension<Infer>) -> Result<(), (StatusCode, Json<ErrorRe
         .generate(GenerateRequest {
             inputs: "liveness".to_string(),
             parameters: GenerateParameters {
-                temperature: 1.0,
-                repetition_penalty: 1.0,
-                top_k: 0,
-                top_p: 1.0,
+                temperature: None,
+                repetition_penalty: None,
+                top_k: None,
+                top_p: None,
                 do_sample: false,
                 max_new_tokens: 1,
-                stop: vec![],
+                stop: Vec::new(),
                 details: false,
                 seed: None,
             },
@@ -47,7 +49,24 @@ async fn health(infer: Extension<Infer>) -> Result<(), (StatusCode, Json<ErrorRe
     Ok(())
 }
 
-/// Generate method
+/// Generate tokens
+#[utoipa::path(
+    post,
+    tag = "Text Generation Inference",
+    path = "/generate",
+    request_body = GenerateRequest,
+    responses(
+        (status = 200, description = "Generated Text", body = [GenerateResponse]),
+        (status = 424, description = "Generation Error", body = [ErrorResponse],
+            example = json!({"error": "Request failed during generation"})),
+        (status = 429, description = "Model is overloaded", body = [ErrorResponse],
+            example = json!({"error": "Model is overloaded"})),
+        (status = 422, description = "Input validation error", body = [ErrorResponse],
+            example = json!({"error": "Input validation error"})),
+        (status = 500, description = "Incomplete generation", body = [ErrorResponse],
+            example = json!({"error": "Incomplete generation"})),
+    )
+)]
 #[instrument(
     skip(infer),
     fields(
@@ -76,7 +95,7 @@ async fn generate(
     // Token details
     let details = match details {
         true => Some(Details {
-            finish_reason: response.generated_text.finish_reason,
+            finish_reason: FinishReason::from(response.generated_text.finish_reason),
             generated_tokens: response.generated_text.generated_tokens,
             prefill: Some(response.prefill),
             tokens: Some(response.tokens),
@@ -132,7 +151,29 @@ async fn generate(
     Ok((headers, Json(response)))
 }
 
-/// Generate stream method
+/// Generate a stream of token using Server Side Events
+#[utoipa::path(
+    post,
+    tag = "Text Generation Inference",
+    path = "/generate_stream",
+    request_body = GenerateRequest,
+    responses(
+        (status = 200, description = "Generated Text", body = [StreamResponse],
+            content_type="text/event-stream "),
+        (status = 424, description = "Generation Error", body = [ErrorResponse],
+            example = json!({"error": "Request failed during generation"}),
+            content_type="text/event-stream "),
+        (status = 429, description = "Model is overloaded", body = [ErrorResponse],
+            example = json!({"error": "Model is overloaded"}),
+            content_type="text/event-stream "),
+        (status = 422, description = "Input validation error", body = [ErrorResponse],
+            example = json!({"error": "Input validation error"}),
+            content_type="text/event-stream "),
+        (status = 500, description = "Incomplete generation", body = [ErrorResponse],
+            example = json!({"error": "Incomplete generation"}),
+            content_type="text/event-stream "),
+    )
+)]
 #[instrument(
     skip(infer),
     fields(
@@ -185,11 +226,9 @@ async fn generate_stream(
                                 } => {
                                     // Token details
                                     let details = match details {
-                                        true => Some(Details {
-                                            finish_reason: generated_text.finish_reason,
+                                        true => Some(StreamDetails {
+                                            finish_reason: FinishReason::from(generated_text.finish_reason),
                                             generated_tokens: generated_text.generated_tokens,
-                                            prefill: None,
-                                            tokens: None,
                                             seed: generated_text.seed,
                                         }),
                                         false => None,
@@ -265,6 +304,39 @@ pub async fn run(
     validation_workers: usize,
     addr: SocketAddr,
 ) {
+    // OpenAPI documentation
+    #[derive(OpenApi)]
+    #[openapi(
+        paths(
+            generate,
+            generate_stream,
+        ),
+        components(
+            schemas(
+                GenerateRequest,
+                GenerateParameters,
+                Token,
+                GenerateResponse,
+                Details,
+                FinishReason,
+                StreamResponse,
+                StreamDetails,
+                ErrorResponse,
+            )
+        ),
+        tags(
+            (name = "Text Generation Inference", description = "Hugging Face Text Generation Inference API")
+        ),
+        info(
+            title = "Text Generation Inference",
+            license(
+                name = "Apache 2.0",
+                url = "https://www.apache.org/licenses/LICENSE-2.0"
+            )
+        )
+    )]
+    struct ApiDoc;
+
     // Create state
     let validation = Validation::new(validation_workers, tokenizer, max_input_length);
     let infer = Infer::new(
@@ -277,6 +349,7 @@ pub async fn run(
 
     // Create router
     let app = Router::new()
+        .merge(SwaggerUi::new("/docs").url("/api-doc/openapi.json", ApiDoc::openapi()))
         .route("/", post(generate))
         .route("/generate", post(generate))
         .route("/generate_stream", post(generate_stream))
@@ -320,6 +393,17 @@ async fn shutdown_signal() {
     tracing::info!("signal received, starting graceful shutdown");
 }
 
+impl From<i32> for FinishReason {
+    fn from(finish_reason: i32) -> Self {
+        let finish_reason = text_generation_client::FinishReason::from_i32(finish_reason).unwrap();
+        match finish_reason {
+            text_generation_client::FinishReason::Length => FinishReason::Length,
+            text_generation_client::FinishReason::EosToken => FinishReason::EndOfSequenceToken,
+            text_generation_client::FinishReason::StopSequence => FinishReason::StopSequence,
+        }
+    }
+}
+
 /// Convert to Axum supported formats
 impl From<InferError> for (StatusCode, Json<ErrorResponse>) {
     fn from(err: InferError) -> Self {

router/src/validation.rs

@@ -110,30 +110,58 @@ fn validate(
     max_input_length: usize,
     rng: &mut ThreadRng,
 ) -> Result<ValidGenerateRequest, ValidationError> {
-    if request.parameters.temperature <= 0.0 {
+    let GenerateParameters {
+        temperature,
+        repetition_penalty,
+        top_k,
+        top_p,
+        do_sample,
+        max_new_tokens,
+        stop: stop_sequences,
+        seed,
+        ..
+    } = request.parameters;
+
+    let temperature = temperature.unwrap_or(1.0);
+    if temperature <= 0.0 {
         return Err(ValidationError::Temperature);
     }
-    if request.parameters.repetition_penalty <= 0.0 {
+
+    let repetition_penalty = repetition_penalty.unwrap_or(1.0);
+    if repetition_penalty <= 0.0 {
         return Err(ValidationError::RepetitionPenalty);
     }
-    if request.parameters.top_p <= 0.0 || request.parameters.top_p > 1.0 {
+
+    let top_p = top_p.unwrap_or(1.0);
+    if top_p <= 0.0 || top_p > 1.0 {
         return Err(ValidationError::TopP);
     }
-    if request.parameters.top_k < 0 {
+
+    // Different because the proto default value is 0 while it is not a valid value
+    // for the user
+    let top_k: u32 = match top_k {
+        None => Ok(0),
+        Some(top_k) => {
+            if top_k <= 0 {
                 return Err(ValidationError::TopK);
             }
-    if request.parameters.max_new_tokens > MAX_MAX_NEW_TOKENS {
+            Ok(top_k as u32)
+        }
+    }?;
+
+    if max_new_tokens == 0 || max_new_tokens > MAX_MAX_NEW_TOKENS {
         return Err(ValidationError::MaxNewTokens(MAX_MAX_NEW_TOKENS));
     }
-    if request.parameters.stop.len() > MAX_STOP_SEQUENCES {
+
+    if stop_sequences.len() > MAX_STOP_SEQUENCES {
         return Err(ValidationError::StopSequence(
             MAX_STOP_SEQUENCES,
-            request.parameters.stop.len(),
+            stop_sequences.len(),
         ));
     }
 
     // If seed is None, assign a random one
-    let seed = match request.parameters.seed {
+    let seed = match seed {
         None => rng.gen(),
         Some(seed) => seed,
     };
@@ -147,21 +175,10 @@ fn validate(
                 Err(ValidationError::InputLength(input_length, max_input_length))
             } else {
                 // Return ValidGenerateRequest
-                let GenerateParameters {
-                    temperature,
-                    repetition_penalty,
-                    top_k,
-                    top_p,
-                    do_sample,
-                    max_new_tokens,
-                    stop: stop_sequences,
-                    ..
-                } = request.parameters;
-
                 let parameters = NextTokenChooserParameters {
                     temperature,
                     repetition_penalty,
-                    top_k: top_k as u32,
+                    top_k,
                     top_p,
                     do_sample,
                     seed,
@@ -206,7 +223,7 @@ pub enum ValidationError {
     TopP,
     #[error("top_k must be strictly positive")]
     TopK,
-    #[error("max_new_tokens must be <= {0}")]
+    #[error("max_new_tokens must be strictly positive and <= {0}")]
     MaxNewTokens(u32),
     #[error("inputs must have less than {1} tokens. Given: {0}")]
     InputLength(usize, usize),

server/README.md

-# BLOOM Inference Python gRPC Server
+# Text Generation Inference Python gRPC Server
 
-A Python gRPC server for BLOOM Inference
+A Python gRPC server for Text Generation Inference
 
 ## Install
 

server/pyproject.toml

 [tool.poetry]
 name = "text-generation"
-version = "0.1.0"
-description = "BLOOM Inference Python gRPC Server"
+version = "0.2.0"
+description = "Text Generation Inference Python gRPC Server"
 authors = ["Olivier Dehaene <olivier@huggingface.co>"]
 
 [tool.poetry.scripts]

server/tests/models/test_bloom.py

@@ -140,8 +140,7 @@ def test_causal_lm_generate_token_completion(default_bloom, default_bloom_batch)
 
     assert len(generations) == 1
     assert (
-        generations[0].generated_text.text
-        == "TestTestTestTestTestTestTestTestTestTest"
+        generations[0].generated_text.text == "TestTestTestTestTestTestTestTestTestTest"
     )
     assert generations[0].request_id == default_bloom_batch.requests[0].id
     assert (
@@ -187,8 +186,7 @@ def test_causal_lm_generate_token_completion_multi(
 
     assert len(generations) == 1
     assert (
-        generations[0].generated_text.text
-        == "TestTestTestTestTestTestTestTestTestTest"
+        generations[0].generated_text.text == "TestTestTestTestTestTestTestTestTestTest"
     )
     assert (
         generations[0].request_id == default_multi_requests_bloom_batch.requests[0].id
@@ -283,8 +281,7 @@ def test_batch_concatenate(
 
     assert len(generations) == 2
     assert (
-        generations[0].generated_text.text
-        == "TestTestTestTestTestTestTestTestTestTest"
+        generations[0].generated_text.text == "TestTestTestTestTestTestTestTestTestTest"
     )
     assert generations[0].request_id == default_bloom_batch.requests[0].id
     assert (
@@ -306,8 +303,7 @@ def test_batch_concatenate(
 
     assert len(generations) == 1
     assert (
-        generations[0].generated_text.text
-        == "TestTestTestTestTestTestTestTestTestTest"
+        generations[0].generated_text.text == "TestTestTestTestTestTestTestTestTestTest"
     )
     assert (
         generations[0].request_id == default_multi_requests_bloom_batch.requests[0].id

server/tests/test_utils.py

@@ -9,6 +9,7 @@ from text_generation.utils import (
     StopSequenceCriteria,
     StoppingCriteria,
     LocalEntryNotFoundError,
+    FinishReason,
 )
 
 
@@ -24,13 +25,13 @@ def test_stop_sequence_criteria():
 def test_stopping_criteria():
     criteria = StoppingCriteria(0, [StopSequenceCriteria("/test;")], max_new_tokens=5)
     assert criteria(65827, "/test") == (False, None)
-    assert criteria(30, ";") == (True, "stop_sequence")
+    assert criteria(30, ";") == (True, FinishReason.FINISH_REASON_STOP_SEQUENCE)
 
 
 def test_stopping_criteria_eos():
     criteria = StoppingCriteria(0, [StopSequenceCriteria("/test;")], max_new_tokens=5)
     assert criteria(1, "") == (False, None)
-    assert criteria(0, "") == (True, "eos_token")
+    assert criteria(0, "") == (True, FinishReason.FINISH_REASON_EOS_TOKEN)
 
 
 def test_stopping_criteria_max():
@@ -39,7 +40,7 @@ def test_stopping_criteria_max():
     assert criteria(1, "") == (False, None)
     assert criteria(1, "") == (False, None)
     assert criteria(1, "") == (False, None)
-    assert criteria(1, "") == (True, "length")
+    assert criteria(1, "") == (True, FinishReason.FINISH_REASON_LENGTH)
 
 
 def test_weight_hub_files():

server/text_generation/cli.py

@@ -13,7 +13,7 @@ app = typer.Typer()
 
 @app.command()
 def serve(
-    model_name: str,
+    model_id: str,
     revision: Optional[str] = None,
     sharded: bool = False,
     quantize: bool = False,
@@ -46,16 +46,16 @@ def serve(
             os.getenv("MASTER_PORT", None) is not None
         ), "MASTER_PORT must be set when sharded is True"
 
-    server.serve(model_name, revision, sharded, quantize, uds_path)
+    server.serve(model_id, revision, sharded, quantize, uds_path)
 
 
 @app.command()
 def download_weights(
-    model_name: str,
+    model_id: str,
     revision: Optional[str] = None,
     extension: str = ".safetensors",
 ):
-    utils.download_weights(model_name, revision, extension)
+    utils.download_weights(model_id, revision, extension)
 
 
 if __name__ == "__main__":

server/text_generation/models/__init__.py

@@ -30,31 +30,31 @@ torch.backends.cudnn.allow_tf32 = True
 
 
 def get_model(
-    model_name: str, revision: Optional[str], sharded: bool, quantize: bool
+    model_id: str, revision: Optional[str], sharded: bool, quantize: bool
 ) -> Model:
-    config = AutoConfig.from_pretrained(model_name, revision=revision)
+    config = AutoConfig.from_pretrained(model_id, revision=revision)
 
     if config.model_type == "bloom":
         if sharded:
-            return BLOOMSharded(model_name, revision, quantize=quantize)
+            return BLOOMSharded(model_id, revision, quantize=quantize)
         else:
-            return BLOOM(model_name, revision, quantize=quantize)
+            return BLOOM(model_id, revision, quantize=quantize)
     elif config.model_type == "gpt_neox":
         if sharded:
-            return GPTNeoxSharded(model_name, revision, quantize=quantize)
+            return GPTNeoxSharded(model_id, revision, quantize=quantize)
         else:
-            return GPTNeox(model_name, revision, quantize=quantize)
-    elif model_name.startswith("facebook/galactica"):
+            return GPTNeox(model_id, revision, quantize=quantize)
+    elif model_id.startswith("facebook/galactica"):
         if sharded:
-            return GalacticaSharded(model_name, revision, quantize=quantize)
+            return GalacticaSharded(model_id, revision, quantize=quantize)
         else:
-            return Galactica(model_name, revision, quantize=quantize)
-    elif "santacoder" in model_name:
-        return SantaCoder(model_name, revision, quantize)
+            return Galactica(model_id, revision, quantize=quantize)
+    elif "santacoder" in model_id:
+        return SantaCoder(model_id, revision, quantize)
     else:
         if sharded:
             raise ValueError("sharded is not supported for AutoModel")
         try:
-            return CausalLM(model_name, revision, quantize=quantize)
+            return CausalLM(model_id, revision, quantize=quantize)
         except Exception:
-            return Seq2SeqLM(model_name, revision, quantize=quantize)
+            return Seq2SeqLM(model_id, revision, quantize=quantize)

server/text_generation/models/bloom.py

@@ -57,10 +57,10 @@ class BLOOM(CausalLM):
 
 class BLOOMSharded(BLOOM):
     def __init__(
-        self, model_name: str, revision: Optional[str] = None, quantize: bool = False
+        self, model_id: str, revision: Optional[str] = None, quantize: bool = False
     ):
-        if not model_name.startswith("bigscience/bloom"):
-            raise ValueError(f"Model {model_name} is not supported")
+        if not model_id.startswith("bigscience/bloom"):
+            raise ValueError(f"Model {model_id} is not supported")
 
         self.process_group, self.rank, self.world_size = initialize_torch_distributed()
         self.master = self.rank == 0
@@ -72,22 +72,20 @@ class BLOOMSharded(BLOOM):
             dtype = torch.float32
 
         tokenizer = AutoTokenizer.from_pretrained(
-            model_name, revision=revision, padding_side="left"
+            model_id, revision=revision, padding_side="left"
         )
 
         config = AutoConfig.from_pretrained(
-            model_name, revision=revision, slow_but_exact=False, tp_parallel=True
+            model_id, revision=revision, slow_but_exact=False, tp_parallel=True
         )
         config.pad_token_id = 3
 
         # Only download weights for small models
-        if self.master and model_name == "bigscience/bloom-560m":
-            download_weights(model_name, revision=revision, extension=".safetensors")
+        if self.master and model_id == "bigscience/bloom-560m":
+            download_weights(model_id, revision=revision, extension=".safetensors")
 
         torch.distributed.barrier(group=self.process_group)
-        filenames = weight_files(
-            model_name, revision=revision, extension=".safetensors"
-        )
+        filenames = weight_files(model_id, revision=revision, extension=".safetensors")
         if not filenames:
             raise ValueError("No safetensors weights found")
 

server/text_generation/models/causal_lm.py

@@ -232,7 +232,7 @@ class CausalLMBatch(Batch):
 
 
 class CausalLM(Model):
-    def __init__(self, model_name: str, revision: Optional[str] = None, quantize=False):
+    def __init__(self, model_id: str, revision: Optional[str] = None, quantize=False):
         if torch.cuda.is_available():
             device = torch.device("cuda")
             dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32
@@ -244,10 +244,10 @@ class CausalLM(Model):
             dtype = torch.float32
 
         tokenizer = AutoTokenizer.from_pretrained(
-            model_name, revision=revision, padding_side="left"
+            model_id, revision=revision, padding_side="left"
         )
         self.model = AutoModelForCausalLM.from_pretrained(
-            model_name,
+            model_id,
             revision=revision,
             torch_dtype=dtype,
             device_map="auto" if torch.cuda.is_available() else None,

server/text_generation/models/galactica.py

@@ -149,10 +149,10 @@ class Galactica(CausalLM):
 
 class GalacticaSharded(Galactica):
     def __init__(
-        self, model_name: str, revision: Optional[str] = None, quantize: bool = False
+        self, model_id: str, revision: Optional[str] = None, quantize: bool = False
     ):
-        if not model_name.startswith("facebook/galactica"):
-            raise ValueError(f"Model {model_name} is not supported")
+        if not model_id.startswith("facebook/galactica"):
+            raise ValueError(f"Model {model_id} is not supported")
 
         self.process_group, self.rank, self.world_size = initialize_torch_distributed()
         self.master = self.rank == 0
@@ -164,22 +164,20 @@ class GalacticaSharded(Galactica):
             dtype = torch.float32
 
         tokenizer = AutoTokenizer.from_pretrained(
-            model_name, revision=revision, padding_side="left"
+            model_id, revision=revision, padding_side="left"
         )
 
         config = AutoConfig.from_pretrained(
-            model_name, revision=revision, tp_parallel=True
+            model_id, revision=revision, tp_parallel=True
         )
         tokenizer.pad_token_id = config.pad_token_id
 
         # Only download weights for small models
-        if self.master and model_name == "facebook/galactica-125m":
-            download_weights(model_name, revision=revision, extension=".safetensors")
+        if self.master and model_id == "facebook/galactica-125m":
+            download_weights(model_id, revision=revision, extension=".safetensors")
 
         torch.distributed.barrier(group=self.process_group)
-        filenames = weight_files(
-            model_name, revision=revision, extension=".safetensors"
-        )
+        filenames = weight_files(model_id, revision=revision, extension=".safetensors")
         if not filenames:
             raise ValueError("No safetensors weights found")
 

server/text_generation/models/gpt_neox.py

@@ -49,7 +49,7 @@ class GPTNeox(CausalLM):
 
 class GPTNeoxSharded(GPTNeox):
     def __init__(
-        self, model_name: str, revision: Optional[str] = None, quantize: bool = False
+        self, model_id: str, revision: Optional[str] = None, quantize: bool = False
     ):
         self.process_group, self.rank, self.world_size = initialize_torch_distributed()
         self.master = self.rank == 0
@@ -61,22 +61,20 @@ class GPTNeoxSharded(GPTNeox):
             dtype = torch.float32
 
         tokenizer = AutoTokenizer.from_pretrained(
-            model_name, revision=revision, padding_side="left"
+            model_id, revision=revision, padding_side="left"
         )
         tokenizer.pad_token = tokenizer.eos_token
 
         config = AutoConfig.from_pretrained(
-            model_name, revision=revision, tp_parallel=True
+            model_id, revision=revision, tp_parallel=True
         )
 
         # Only master download weights
         if self.master:
-            download_weights(model_name, revision=revision, extension=".safetensors")
+            download_weights(model_id, revision=revision, extension=".safetensors")
 
         torch.distributed.barrier(group=self.process_group)
-        filenames = weight_files(
-            model_name, revision=revision, extension=".safetensors"
-        )
+        filenames = weight_files(model_id, revision=revision, extension=".safetensors")
         if not filenames:
             raise ValueError("No safetensors weights found")
 

server/text_generation/models/santacoder.py

@@ -14,7 +14,7 @@ EOD = "<|endoftext|>"
 
 
 class SantaCoder(CausalLM):
-    def __init__(self, model_name: str, revision: Optional[str] = None, quantize=False):
+    def __init__(self, model_id: str, revision: Optional[str] = None, quantize=False):
         if torch.cuda.is_available():
             device = torch.device("cuda")
             dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32
@@ -26,7 +26,7 @@ class SantaCoder(CausalLM):
             dtype = torch.float32
 
         tokenizer = AutoTokenizer.from_pretrained(
-            model_name, revision=revision, padding_side="left"
+            model_id, revision=revision, padding_side="left"
         )
         tokenizer.add_special_tokens(
             {
@@ -43,7 +43,7 @@ class SantaCoder(CausalLM):
 
         self.model = (
             AutoModelForCausalLM.from_pretrained(
-                model_name,
+                model_id,
                 revision=revision,
                 torch_dtype=dtype,
                 load_in_8bit=quantize,

server/text_generation/models/seq2seq_lm.py

@@ -289,7 +289,7 @@ class Seq2SeqLMBatch(Batch):
 
 
 class Seq2SeqLM(Model):
-    def __init__(self, model_name: str, revision: Optional[str] = None, quantize=False):
+    def __init__(self, model_id: str, revision: Optional[str] = None, quantize=False):
         if torch.cuda.is_available():
             device = torch.device("cuda")
             dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32
@@ -301,14 +301,14 @@ class Seq2SeqLM(Model):
             dtype = torch.float32
 
         self.model = AutoModelForSeq2SeqLM.from_pretrained(
-            model_name,
+            model_id,
             revision=revision,
             torch_dtype=dtype,
             device_map="auto" if torch.cuda.is_available() else None,
             load_in_8bit=quantize,
         ).eval()
         tokenizer = AutoTokenizer.from_pretrained(
-            model_name, revision=revision, padding_side="left"
+            model_id, revision=revision, padding_side="left"
         )
         tokenizer.bos_token_id = self.model.config.decoder_start_token_id
 

server/text_generation/models/types.py

@@ -7,6 +7,7 @@ from typing import List, Optional
 from transformers import PreTrainedTokenizerBase
 
 from text_generation.pb import generate_pb2
+from text_generation.pb.generate_pb2 import FinishReason
 
 
 class Batch(ABC):
@@ -38,7 +39,7 @@ class Batch(ABC):
 class GeneratedText:
     text: str
     generated_tokens: int
-    finish_reason: str
+    finish_reason: FinishReason
     seed: Optional[int]
 
     def to_pb(self) -> generate_pb2.GeneratedText:

server/text_generation/server.py

@@ -66,14 +66,14 @@ class TextGenerationService(generate_pb2_grpc.TextGenerationServiceServicer):
 
 
 def serve(
-    model_name: str,
+    model_id: str,
     revision: Optional[str],
     sharded: bool,
     quantize: bool,
     uds_path: Path,
 ):
     async def serve_inner(
-        model_name: str,
+        model_id: str,
         revision: Optional[str],
         sharded: bool = False,
         quantize: bool = False,
@@ -89,7 +89,7 @@ def serve(
             local_url = unix_socket_template.format(uds_path, 0)
             server_urls = [local_url]
 
-        model = get_model(model_name, revision, sharded, quantize)
+        model = get_model(model_id, revision, sharded, quantize)
 
         server = aio.server(interceptors=[ExceptionInterceptor()])
         generate_pb2_grpc.add_TextGenerationServiceServicer_to_server(
@@ -109,4 +109,4 @@ def serve(
             logger.info("Signal received. Shutting down")
             await server.stop(0)
 
-    asyncio.run(serve_inner(model_name, revision, sharded, quantize))
+    asyncio.run(serve_inner(model_id, revision, sharded, quantize))

server/text_generation/utils.py

@@ -24,9 +24,11 @@ from transformers.generation.logits_process import (
 )
 
 from text_generation.pb import generate_pb2
+from text_generation.pb.generate_pb2 import FinishReason
 
 WEIGHTS_CACHE_OVERRIDE = os.getenv("WEIGHTS_CACHE_OVERRIDE", None)
 
+
 class Sampling:
     def __init__(self, seed: int, device: str = "cpu"):
         self.generator = torch.Generator(device)
@@ -129,15 +131,15 @@ class StoppingCriteria:
     def __call__(self, last_token: int, last_output: str) -> Tuple[bool, Optional[str]]:
         self.current_tokens += 1
         if self.current_tokens >= self.max_new_tokens:
-            return True, "length"
+            return True, FinishReason.FINISH_REASON_LENGTH
 
         if last_token == self.eos_token_id:
-            return True, "eos_token"
+            return True, FinishReason.FINISH_REASON_EOS_TOKEN
 
         self.current_output += last_output
         for stop_sequence_criteria in self.stop_sequence_criterias:
             if stop_sequence_criteria(self.current_output):
-                return True, "stop_sequence"
+                return True, FinishReason.FINISH_REASON_STOP_SEQUENCE
 
         return False, None
 
@@ -180,20 +182,20 @@ def initialize_torch_distributed():
     return torch.distributed.distributed_c10d._get_default_group(), rank, world_size
 
 
-def weight_hub_files(model_name, revision=None, extension=".safetensors"):
+def weight_hub_files(model_id, revision=None, extension=".safetensors"):
     """Get the safetensors filenames on the hub"""
     api = HfApi()
-    info = api.model_info(model_name, revision=revision)
+    info = api.model_info(model_id, revision=revision)
     filenames = [s.rfilename for s in info.siblings if s.rfilename.endswith(extension)]
     return filenames
 
 
-def try_to_load_from_cache(model_name, revision, filename):
+def try_to_load_from_cache(model_id, revision, filename):
     """Try to load a file from the Hugging Face cache"""
     if revision is None:
         revision = "main"
 
-    object_id = model_name.replace("/", "--")
+    object_id = model_id.replace("/", "--")
     repo_cache = Path(HUGGINGFACE_HUB_CACHE) / f"models--{object_id}"
 
     if not repo_cache.is_dir():
@@ -228,38 +230,38 @@ def try_to_load_from_cache(model_name, revision, filename):
     return str(cached_file) if cached_file.is_file() else None
 
 
-def weight_files(model_name, revision=None, extension=".safetensors"):
+def weight_files(model_id, revision=None, extension=".safetensors"):
     """Get the local safetensors filenames"""
     if WEIGHTS_CACHE_OVERRIDE is not None:
         return list(Path(WEIGHTS_CACHE_OVERRIDE).glob(f"*{extension}"))
 
-    filenames = weight_hub_files(model_name, revision, extension)
+    filenames = weight_hub_files(model_id, revision, extension)
     files = []
     for filename in filenames:
         cache_file = try_to_load_from_cache(
-            model_name, revision=revision, filename=filename
+            model_id, revision=revision, filename=filename
         )
         if cache_file is None:
             raise LocalEntryNotFoundError(
-                f"File {filename} of model {model_name} not found in "
+                f"File {filename} of model {model_id} not found in "
                 f"{os.getenv('HUGGINGFACE_HUB_CACHE', 'the local cache')}. "
-                f"Please run `text-generation-server download-weights {model_name}` first."
+                f"Please run `text-generation-server download-weights {model_id}` first."
             )
         files.append(cache_file)
 
     return files
 
 
-def download_weights(model_name, revision=None, extension=".safetensors"):
+def download_weights(model_id, revision=None, extension=".safetensors"):
     """Download the safetensors files from the hub"""
     if WEIGHTS_CACHE_OVERRIDE is not None:
         return list(Path(WEIGHTS_CACHE_OVERRIDE).glob(f"*{extension}"))
 
-    filenames = weight_hub_files(model_name, revision, extension)
+    filenames = weight_hub_files(model_id, revision, extension)
 
     download_function = partial(
         hf_hub_download,
-        repo_id=model_name,
+        repo_id=model_id,
         local_files_only=False,
     )