feat: Return logprobs (IBM#8)

Author: OlivierDehaene

.github/workflows/server-tests.yaml

@@ -0,0 +1,30 @@
+name: Server Tests
+
+on:
+  pull_request:
+    paths:
+      - "server/**"
+      - "proto/**"
+
+jobs:
+  run_tests:
+    runs-on: ubuntu-20.04
+    steps:
+      - uses: actions/checkout@v2
+      - name: Set up Python
+        uses: actions/setup-python@v1
+        with:
+          python-version: 3.9
+      - name: Loading cache.
+        uses: actions/cache@v2
+        id: model_cache
+        with:
+          path: ~/.cache/huggingface/
+          key: models
+      - name: Install server dependencies
+        run: |
+          make install-server
+      - name: Run tests
+        run: |
+          pip install pytest
+          pytest -sv server/tests

README.md

@@ -17,6 +17,7 @@ to power Bloom, BloomZ and MT0-XXL api-inference widgets.
 - 45ms per token generation for BLOOM with 8xA100 80GB
 - Logits warpers (temperature scaling, topk ...)
 - Stop sequences
+- Log probabilities
 
 ## Officially supported models
 

proto/generate.proto

@@ -27,7 +27,7 @@ message ClearCacheRequest {}
 /// Empty response
 message ClearCacheResponse {}
 
-message LogitsWarperParameters {
+message NextTokenChooserParameters {
     /// exponential scaling output probability distribution
     float temperature = 1;
     /// restricting to the k highest probability elements
@@ -52,8 +52,8 @@ message Request {
     string inputs = 2;
     /// The number of tokens inside inputs
     uint32 input_length = 3;
-    /// Logits Warper Parameters
-    LogitsWarperParameters parameters = 4;
+    /// Next Token Chooser Parameters
+    NextTokenChooserParameters parameters = 4;
     /// Stopping Criteria Parameters
     StoppingCriteriaParameters stopping_parameters = 5;
 }
@@ -71,11 +71,17 @@ message GeneratedText {
     /// Request
     Request request = 1;
     /// Output
-    string output = 2;
+    string output_text = 2;
     /// Number of generated tokens
-    uint32 tokens = 3;
+    uint32 generated_tokens = 3;
+    /// Tokens
+    repeated string tokens = 4;
+    /// Token IDs
+    repeated uint32 token_ids = 5;
+    /// Logprobs
+    repeated float logprobs = 6;
     /// Finish reason
-    string finish_reason = 4;
+    string finish_reason = 7;
 }
 
 message GenerateRequest {

router/client/src/lib.rs

@@ -7,7 +7,7 @@ mod sharded_client;
 
 pub use client::Client;
 pub use pb::generate::v1::{
-    Batch, GeneratedText, LogitsWarperParameters, Request, StoppingCriteriaParameters,
+    Batch, GeneratedText, NextTokenChooserParameters, Request, StoppingCriteriaParameters,
 };
 pub use sharded_client::ShardedClient;
 use thiserror::Error;

router/src/batcher.rs

@@ -187,9 +187,13 @@ fn send_generated(finished: Vec<GeneratedText>, db: &Db) {
         let entry = db
             .remove(&output.request.unwrap().id)
             .expect("ID not found in db. This is a bug.");
+
         let response = InferResponse {
-            output: output.output,
+            output_text: output.output_text,
+            generated_tokens: output.generated_tokens,
+            token_ids: output.token_ids,
             tokens: output.tokens,
+            logprobs: output.logprobs,
             finish_reason: output.finish_reason,
             queued: entry.time,
             start: entry.batch_time.unwrap(), // unwrap is always valid
@@ -202,8 +206,11 @@ fn send_generated(finished: Vec<GeneratedText>, db: &Db) {
 
 #[derive(Debug)]
 pub(crate) struct InferResponse {
-    pub(crate) output: String,
-    pub(crate) tokens: u32,
+    pub(crate) output_text: String,
+    pub(crate) generated_tokens: u32,
+    pub(crate) token_ids: Vec<u32>,
+    pub(crate) tokens: Vec<String>,
+    pub(crate) logprobs: Vec<f32>,
     pub(crate) finish_reason: String,
     pub(crate) queued: Instant,
     pub(crate) start: Instant,

router/src/db.rs

@@ -5,7 +5,7 @@ use parking_lot::Mutex;
 use std::collections::BTreeMap;
 use std::sync::Arc;
 use text_generation_client::{
-    Batch, ClientError, LogitsWarperParameters, Request, StoppingCriteriaParameters,
+    Batch, ClientError, NextTokenChooserParameters, Request, StoppingCriteriaParameters,
 };
 use tokio::sync::oneshot::Sender;
 use tokio::time::Instant;
@@ -71,7 +71,7 @@ impl State {
                 id: *id,
                 inputs: entry.request.inputs.clone(),
                 input_length: entry.input_length as u32,
-                parameters: Some(LogitsWarperParameters::from(
+                parameters: Some(NextTokenChooserParameters::from(
                     entry.request.parameters.clone(),
                 )),
                 stopping_parameters: Some(StoppingCriteriaParameters::from(
@@ -162,7 +162,7 @@ impl Db {
     }
 }
 
-impl From<GenerateParameters> for LogitsWarperParameters {
+impl From<GenerateParameters> for NextTokenChooserParameters {
     fn from(parameters: GenerateParameters) -> Self {
         Self {
             temperature: parameters.temperature,

router/src/lib.rs

@@ -21,7 +21,10 @@ pub(crate) struct GenerateParameters {
     pub do_sample: bool,
     #[serde(default = "default_max_new_tokens")]
     pub max_new_tokens: u32,
+    #[serde(default)]
     pub stop: Vec<String>,
+    #[serde(default)]
+    pub details: bool,
 }
 
 fn default_temperature() -> f32 {
@@ -52,6 +55,7 @@ fn default_parameters() -> GenerateParameters {
         do_sample: default_do_sample(),
         max_new_tokens: default_max_new_tokens(),
         stop: vec![],
+        details: false,
     }
 }
 
@@ -62,10 +66,18 @@ pub(crate) struct GenerateRequest {
     pub parameters: GenerateParameters,
 }
 
+#[derive(Serialize)]
+pub(crate) struct Details {
+    pub finish_reason: String,
+    pub generated_tokens: u32,
+    pub tokens: Vec<(u32, String, f32)>,
+}
+
 #[derive(Serialize)]
 pub(crate) struct GeneratedText {
     pub generated_text: String,
-    pub finish_reason: String,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub details: Option<Details>,
 }
 
 #[derive(Serialize)]

router/src/server.rs

@@ -1,5 +1,5 @@
 use crate::{
-    Batcher, ErrorResponse, GenerateParameters, GenerateRequest, GeneratedText, Validation,
+    Batcher, Details, ErrorResponse, GenerateParameters, GenerateRequest, GeneratedText, Validation,
 };
 use axum::extract::Extension;
 use axum::http::{HeaderMap, StatusCode};
@@ -54,6 +54,7 @@ async fn health(state: Extension<ServerState>) -> Result<(), (StatusCode, Json<E
                     do_sample: false,
                     max_new_tokens: 1,
                     stop: vec![],
+                    details: false,
                 },
             },
         )
@@ -89,6 +90,7 @@ async fn generate(
     })?;
 
     // Validate request
+    let details = req.0.parameters.details;
     let (input_length, validated_request) =
         state.validation.validate(req.0).await.map_err(|err| {
             tracing::error!("{}", err.to_string());
@@ -105,12 +107,31 @@ async fn generate(
             err
         })?;
 
+    // Token details
+    let details = match details {
+        true => {
+            let tokens = response
+                .token_ids
+                .into_iter()
+                .zip(response.tokens.into_iter())
+                .zip(response.logprobs.into_iter())
+                .map(|((id, text), logprob)| (id, text, logprob))
+                .collect();
+            Some(Details {
+                finish_reason: response.finish_reason,
+                generated_tokens: response.generated_tokens,
+                tokens,
+            })
+        }
+        false => None,
+    };
+
     // Timings
     let total_time = start_time.elapsed();
     let validation_time = response.queued - start_time;
     let queue_time = response.start - response.queued;
     let inference_time = response.end - response.start;
-    let time_per_token = inference_time / response.tokens;
+    let time_per_token = inference_time / response.generated_tokens;
 
     // Headers
     let mut headers = HeaderMap::new();
@@ -141,12 +162,12 @@ async fn generate(
     tracing::Span::current().record("queue_time", format!("{:?}", queue_time));
     tracing::Span::current().record("inference_time", format!("{:?}", inference_time));
     tracing::Span::current().record("time_per_token", format!("{:?}", time_per_token));
-    tracing::info!("Output: {}", response.output);
+    tracing::info!("Output: {}", response.output_text);
 
     // Send response
     let response = vec![GeneratedText {
-        generated_text: response.output,
-        finish_reason: response.finish_reason,
+        generated_text: response.output_text,
+        details,
     }];
     Ok((headers, Json(response)))
 }

server/tests/conftest.py

@@ -7,7 +7,7 @@
 
 @pytest.fixture
 def default_pb_parameters():
-    return generate_pb2.LogitsWarperParameters(
+    return generate_pb2.NextTokenChooserParameters(
         temperature=1.0,
         top_k=0,
         top_p=1.0,

server/tests/models/test_bloom.py

@@ -128,10 +128,12 @@ def test_causal_lm_generate_token_completion(default_bloom, default_bloom_batch)
     assert next_batch is None
 
     assert len(generated_texts) == 1
-    assert generated_texts[0].output == "TestTestTestTestTestTestTestTestTestTestTest"
+    assert (
+        generated_texts[0].output_text == "TestTestTestTestTestTestTestTestTestTestTest"
+    )
     assert generated_texts[0].request == default_bloom_batch.requests[0]
     assert (
-        generated_texts[0].tokens
+        generated_texts[0].generated_tokens
         == default_bloom_batch.stopping_criterias[0].max_new_tokens
     )
 
@@ -151,10 +153,10 @@ def test_causal_lm_generate_token_completion_multi(
     assert next_batch is not None
 
     assert len(generated_texts) == 1
-    assert generated_texts[0].output == "TestTestTestTestTestTest"
+    assert generated_texts[0].output_text == "TestTestTestTestTestTest"
     assert generated_texts[0].request == default_multi_requests_bloom_batch.requests[1]
     assert (
-        generated_texts[0].tokens
+        generated_texts[0].generated_tokens
         == default_multi_requests_bloom_batch.stopping_criterias[1].max_new_tokens
     )
 
@@ -170,10 +172,12 @@ def test_causal_lm_generate_token_completion_multi(
     assert next_batch is None
 
     assert len(generated_texts) == 1
-    assert generated_texts[0].output == "TestTestTestTestTestTestTestTestTestTestTest"
+    assert (
+        generated_texts[0].output_text == "TestTestTestTestTestTestTestTestTestTestTest"
+    )
     assert generated_texts[0].request == default_multi_requests_bloom_batch.requests[0]
     assert (
-        generated_texts[0].tokens
+        generated_texts[0].generated_tokens
         == default_multi_requests_bloom_batch.stopping_criterias[0].max_new_tokens
     )
 
@@ -240,10 +244,10 @@ def test_batch_concatenate(
     assert next_batch is not None
 
     assert len(generated_texts) == 1
-    assert generated_texts[0].output == "TestTestTestTestTestTest"
+    assert generated_texts[0].output_text == "TestTestTestTestTestTest"
     assert generated_texts[0].request == default_multi_requests_bloom_batch.requests[1]
     assert (
-        generated_texts[0].tokens
+        generated_texts[0].generated_tokens
         == default_multi_requests_bloom_batch.stopping_criterias[1].max_new_tokens
     )
 
@@ -259,10 +263,12 @@ def test_batch_concatenate(
     assert next_batch is not None
 
     assert len(generated_texts) == 1
-    assert generated_texts[0].output == "TestTestTestTestTestTestTestTestTestTestTest"
+    assert (
+        generated_texts[0].output_text == "TestTestTestTestTestTestTestTestTestTestTest"
+    )
     assert generated_texts[0].request == default_bloom_batch.requests[0]
     assert (
-        generated_texts[0].tokens
+        generated_texts[0].generated_tokens
         == default_bloom_batch.stopping_criterias[0].max_new_tokens
     )
 
@@ -279,9 +285,11 @@ def test_batch_concatenate(
     assert next_batch is None
 
     assert len(generated_texts) == 1
-    assert generated_texts[0].output == "TestTestTestTestTestTestTestTestTestTestTest"
+    assert (
+        generated_texts[0].output_text == "TestTestTestTestTestTestTestTestTestTestTest"
+    )
     assert generated_texts[0].request == default_multi_requests_bloom_batch.requests[0]
     assert (
-        generated_texts[0].tokens
+        generated_texts[0].generated_tokens
         == default_multi_requests_bloom_batch.stopping_criterias[0].max_new_tokens
     )