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,