Added prompt_tokens to the response (#165)

clients/python/README.md

@@ -181,6 +181,8 @@ class BestOfSequence:
 class Details:
     # Generation finish reason
     finish_reason: FinishReason
+    # Number of prompt tokens
+    prompt_tokens: int
     # Number of generated tokens
     generated_tokens: int
     # Sampling seed if sampling was activated
@@ -205,6 +207,8 @@ class Response:
 class StreamDetails:
     # Generation finish reason
     finish_reason: FinishReason
+    # Number of prompt tokens
+    prompt_tokens: int
     # Number of generated tokens
     generated_tokens: int
     # Sampling seed if sampling was activated

clients/python/lorax/types.py

@@ -198,6 +198,8 @@ class BestOfSequence(BaseModel):
 class Details(BaseModel):
     # Generation finish reason
     finish_reason: FinishReason
+    # Number of prompt tokens
+    prompt_tokens: int
     # Number of generated tokens
     generated_tokens: int
     # Sampling seed if sampling was activated
@@ -222,6 +224,8 @@ class Response(BaseModel):
 class StreamDetails(BaseModel):
     # Generation finish reason
     finish_reason: FinishReason
+    # Number of prompt tokens
+    prompt_tokens: int
     # Number of generated tokens
     generated_tokens: int
     # Sampling seed if sampling was activated

clients/python/tests/conftest.py

@@ -1,51 +0,0 @@
-import pytest
-
-from lorax import __version__
-from huggingface_hub.utils import build_hf_headers
-
-
-@pytest.fixture
-def flan_t5_xxl():
-    return "google/flan-t5-xxl"
-
-
-@pytest.fixture
-def fake_model():
-    return "fake/model"
-
-
-@pytest.fixture
-def unsupported_model():
-    return "gpt2"
-
-
-@pytest.fixture
-def base_url():
-    return "https://api-inference.huggingface.co/models"
-
-
-@pytest.fixture
-def bloom_url(base_url, bloom_model):
-    return f"{base_url}/{bloom_model}"
-
-
-@pytest.fixture
-def flan_t5_xxl_url(base_url, flan_t5_xxl):
-    return f"{base_url}/{flan_t5_xxl}"
-
-
-@pytest.fixture
-def fake_url(base_url, fake_model):
-    return f"{base_url}/{fake_model}"
-
-
-@pytest.fixture
-def unsupported_url(base_url, unsupported_model):
-    return f"{base_url}/{unsupported_model}"
-
-
-@pytest.fixture(scope="session")
-def hf_headers():
-    return build_hf_headers(
-        library_name="lorax-tests", library_version=__version__
-    )

docs/reference/openapi.json

@@ -413,6 +413,7 @@
         "type": "object",
         "required": [
           "finish_reason",
+          "prompt_tokens",
           "generated_tokens",
           "prefill",
           "tokens"
@@ -428,6 +429,12 @@
           "finish_reason": {
             "$ref": "#/components/schemas/FinishReason"
           },
+          "prompt_tokens": {
+            "type": "integer",
+            "format": "int32",
+            "example": 1,
+            "minimum": 0.0
+          },
           "generated_tokens": {
             "type": "integer",
             "format": "int32",
@@ -773,12 +780,19 @@
         "type": "object",
         "required": [
           "finish_reason",
+          "prompt_tokens",
           "generated_tokens"
         ],
         "properties": {
           "finish_reason": {
             "$ref": "#/components/schemas/FinishReason"
           },
+          "prompt_tokens": {
+            "type": "integer",
+            "format": "int32",
+            "example": 1,
+            "minimum": 0.0
+          },
           "generated_tokens": {
             "type": "integer",
             "format": "int32",

docs/reference/python_client.md

@@ -181,6 +181,8 @@ class BestOfSequence:
 class Details:
     # Generation finish reason
     finish_reason: FinishReason
+    # Number of prompt tokens
+    prompt_tokens: int
     # Number of generated tokens
     generated_tokens: int
     # Sampling seed if sampling was activated
@@ -205,6 +207,8 @@ class Response:
 class StreamDetails:
     # Generation finish reason
     finish_reason: FinishReason
+    # Number of prompt tokens
+    prompt_tokens: int
     # Number of generated tokens
     generated_tokens: int
     # Sampling seed if sampling was activated

proto/generate.proto

@@ -167,6 +167,8 @@ message Generation {
     bool token_is_special = 6;
     /// Complete generated text
     optional GeneratedText generated_text = 7;
+    /// Prefill tokens length
+    uint32 prefill_tokens_length = 8;
 }
 
 message FilterBatchRequest {

router/src/infer.rs

@@ -189,6 +189,7 @@ impl Infer {
         // Return values
         let mut result_prefill = Vec::new();
         let mut result_tokens = Vec::new();
+        let mut result_prefill_length = 0;
         let mut result_generated_text = None;
         let mut result_start = None;
         let mut result_queued = None;
@@ -197,16 +198,22 @@ impl Infer {
         while let Some(response) = stream.next().await {
             match response? {
                 // Add prefill tokens
-                InferStreamResponse::Prefill(tokens) => {
+                InferStreamResponse::Prefill {
+                    tokens,
+                    tokens_length,
+                } => {
                     // Create Token objects
                     // We do that here instead of in the Python code as Rust for loops are faster
-                    result_prefill = tokens
-                        .ids
-                        .into_iter()
-                        .zip(tokens.logprobs.into_iter())
-                        .zip(tokens.texts.into_iter())
-                        .map(|((id, logprob), text)| PrefillToken { id, text, logprob })
-                        .collect();
+                    if let Some(tokens_val) = tokens {
+                        result_prefill = tokens_val
+                            .ids
+                            .into_iter()
+                            .zip(tokens_val.logprobs.into_iter())
+                            .zip(tokens_val.texts.into_iter())
+                            .map(|((id, logprob), text)| PrefillToken { id, text, logprob })
+                            .collect();
+                    }
+                    result_prefill_length = tokens_length;
                 }
                 // Push last token
                 InferStreamResponse::Token(token) => result_tokens.push(token),
@@ -233,6 +240,7 @@ impl Infer {
             Ok(InferResponse {
                 prefill: result_prefill,
                 tokens: result_tokens,
+                prompt_tokens: result_prefill_length,
                 generated_text,
                 queued,
                 start,
@@ -569,10 +577,13 @@ fn send_responses(
 
     let mut stopped = false;
 
-    if let Some(prefill_tokens) = generation.prefill_tokens {
+    if generation.prefill_tokens_length > 0 {
         // Send message
         entry.response_tx.send_timeout(
-            Ok(InferStreamResponse::Prefill(prefill_tokens)),
+            Ok(InferStreamResponse::Prefill {
+                tokens: generation.prefill_tokens,
+                tokens_length: generation.prefill_tokens_length,
+            }),
             Duration::from_millis(10),
         )?;
     }
@@ -629,7 +640,10 @@ fn send_errors(error: ClientError, entries: &mut IntMap<u64, Entry>) {
 #[derive(Debug)]
 pub(crate) enum InferStreamResponse {
     // Optional first message
-    Prefill(PrefillTokens),
+    Prefill {
+        tokens: Option<PrefillTokens>,
+        tokens_length: u32,
+    },
     // Intermediate messages
     Token(Token),
     // Last message
@@ -645,6 +659,7 @@ pub(crate) enum InferStreamResponse {
 pub(crate) struct InferResponse {
     pub(crate) prefill: Vec<PrefillToken>,
     pub(crate) tokens: Vec<Token>,
+    pub(crate) prompt_tokens: u32,
     pub(crate) generated_text: GeneratedText,
     pub(crate) queued: Instant,
     pub(crate) start: Instant,