feat: Support stop sequences (IBM#7)

Author: OlivierDehaene

README.md

@@ -15,6 +15,8 @@ to power Bloom, BloomZ and MT0-XXL api-inference widgets.
 - Quantization with [bitsandbytes](https://github.com/TimDettmers/bitsandbytes)
 - [Safetensors](https://github.com/huggingface/safetensors) weight loading
 - 45ms per token generation for BLOOM with 8xA100 80GB
+- Logits warpers (temperature scaling, topk ...)
+- Stop sequences
 
 ## Officially supported models
 

proto/generate.proto

@@ -28,12 +28,23 @@ message ClearCacheRequest {}
 message ClearCacheResponse {}
 
 message LogitsWarperParameters {
+    /// exponential scaling output probability distribution
     float temperature = 1;
+    /// restricting to the k highest probability elements
     uint32 top_k = 2;
+    /// restricting to top tokens summing to prob_cut_off <= prob_cut_off
     float top_p = 3;
+    /// apply sampling on the logits
     bool do_sample = 4;
 }
 
+message StoppingCriteriaParameters {
+    /// Maximum number of generated tokens
+    uint32 max_new_tokens = 1;
+    /// Optional stopping sequences
+    repeated string stop_sequences = 2;
+}
+
 message Request {
     /// Request ID
     uint64 id = 1;
@@ -43,8 +54,8 @@ message Request {
     uint32 input_length = 3;
     /// Logits Warper Parameters
     LogitsWarperParameters parameters = 4;
-    /// Stopping criteria
-    uint32 max_new_tokens = 5;
+    /// Stopping Criteria Parameters
+    StoppingCriteriaParameters stopping_parameters = 5;
 }
 
 message Batch {
@@ -63,6 +74,8 @@ message GeneratedText {
     string output = 2;
     /// Number of generated tokens
     uint32 tokens = 3;
+    /// Finish reason
+    string finish_reason = 4;
 }
 
 message GenerateRequest {

router/client/src/lib.rs

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

router/src/batcher.rs

@@ -190,6 +190,7 @@ fn send_generated(finished: Vec<GeneratedText>, db: &Db) {
         let response = InferResponse {
             output: output.output,
             tokens: output.tokens,
+            finish_reason: output.finish_reason,
             queued: entry.time,
             start: entry.batch_time.unwrap(), // unwrap is always valid
             end: Instant::now(),
@@ -203,6 +204,7 @@ fn send_generated(finished: Vec<GeneratedText>, db: &Db) {
 pub(crate) struct InferResponse {
     pub(crate) output: String,
     pub(crate) tokens: u32,
+    pub(crate) finish_reason: String,
     pub(crate) queued: Instant,
     pub(crate) start: Instant,
     pub(crate) end: Instant,

router/src/db.rs

@@ -4,7 +4,9 @@ use crate::{GenerateParameters, GenerateRequest};
 use parking_lot::Mutex;
 use std::collections::BTreeMap;
 use std::sync::Arc;
-use text_generation_client::{Batch, ClientError, LogitsWarperParameters, Request};
+use text_generation_client::{
+    Batch, ClientError, LogitsWarperParameters, Request, StoppingCriteriaParameters,
+};
 use tokio::sync::oneshot::Sender;
 use tokio::time::Instant;
 
@@ -72,7 +74,9 @@ impl State {
                 parameters: Some(LogitsWarperParameters::from(
                     entry.request.parameters.clone(),
                 )),
-                max_new_tokens: entry.request.parameters.max_new_tokens,
+                stopping_parameters: Some(StoppingCriteriaParameters::from(
+                    entry.request.parameters.clone(),
+                )),
             });
 
             ids.push(*id);
@@ -168,3 +172,12 @@ impl From<GenerateParameters> for LogitsWarperParameters {
         }
     }
 }
+
+impl From<GenerateParameters> for StoppingCriteriaParameters {
+    fn from(parameters: GenerateParameters) -> Self {
+        Self {
+            stop_sequences: parameters.stop,
+            max_new_tokens: parameters.max_new_tokens,
+        }
+    }
+}

router/src/lib.rs

@@ -21,6 +21,7 @@ pub(crate) struct GenerateParameters {
     pub do_sample: bool,
     #[serde(default = "default_max_new_tokens")]
     pub max_new_tokens: u32,
+    pub stop: Vec<String>,
 }
 
 fn default_temperature() -> f32 {
@@ -50,6 +51,7 @@ fn default_parameters() -> GenerateParameters {
         top_p: default_top_p(),
         do_sample: default_do_sample(),
         max_new_tokens: default_max_new_tokens(),
+        stop: vec![],
     }
 }
 
@@ -63,6 +65,7 @@ pub(crate) struct GenerateRequest {
 #[derive(Serialize)]
 pub(crate) struct GeneratedText {
     pub generated_text: String,
+    pub finish_reason: String,
 }
 
 #[derive(Serialize)]

router/src/server.rs

@@ -53,6 +53,7 @@ async fn health(state: Extension<ServerState>) -> Result<(), (StatusCode, Json<E
                     top_p: 1.0,
                     do_sample: false,
                     max_new_tokens: 1,
+                    stop: vec![],
                 },
             },
         )
@@ -88,11 +89,8 @@ async fn generate(
     })?;
 
     // Validate request
-    let (input_length, validated_request) = state
-        .validation
-        .validate(req.0)
-        .await
-        .map_err(|err| {
+    let (input_length, validated_request) =
+        state.validation.validate(req.0).await.map_err(|err| {
             tracing::error!("{}", err.to_string());
             err
         })?;
@@ -148,6 +146,7 @@ async fn generate(
     // Send response
     let response = vec![GeneratedText {
         generated_text: response.output,
+        finish_reason: response.finish_reason,
     }];
     Ok((headers, Json(response)))
 }

router/src/validation.rs

@@ -121,6 +121,14 @@ fn validation_worker(
                 .unwrap_or(());
             continue;
         }
+        if request.parameters.stop.len() > 4 {
+            response_tx
+                .send(Err(ValidationError::StopSequence(
+                    request.parameters.stop.len(),
+                )))
+                .unwrap_or(());
+            continue;
+        }
 
         // Get the number of tokens in the input
         match tokenizer.encode(request.inputs.clone(), false) {
@@ -163,6 +171,8 @@ pub enum ValidationError {
     MaxNewTokens,
     #[error("inputs must have less than {1} tokens. Given: {0}")]
     InputLength(usize, usize),
+    #[error("stop supports up to 4 stop sequences. Given: {0}")]
+    StopSequence(usize),
     #[error("tokenizer error {0}")]
     Tokenizer(String),
 }

server/tests/conftest.py

@@ -15,6 +15,11 @@ def default_pb_parameters():
     )
 
 
+@pytest.fixture
+def default_pb_stop_parameters():
+    return generate_pb2.StoppingCriteriaParameters(stop_sequences=[], max_new_tokens=10)
+
+
 @pytest.fixture(scope="session")
 def bloom_560m_tokenizer():
     return AutoTokenizer.from_pretrained("bigscience/bloom-560m", padding_side="left")

server/tests/models/test_bloom.py

@@ -9,13 +9,13 @@
 
 
 @pytest.fixture
-def default_pb_request(default_pb_parameters):
+def default_pb_request(default_pb_parameters, default_pb_stop_parameters):
     return generate_pb2.Request(
         id=0,
         inputs="Test",
         input_length=1,
         parameters=default_pb_parameters,
-        max_new_tokens=10,
+        stopping_parameters=default_pb_stop_parameters,
     )
 
 
@@ -36,7 +36,7 @@ def default_multi_requests_bloom_batch(default_pb_request, bloom_560m_tokenizer)
     req_0 = copy(default_pb_request)
     req_1 = default_pb_request
     req_1.id = 1
-    req_1.max_new_tokens = 5
+    req_1.stopping_parameters.max_new_tokens = 5
 
     batch_pb = generate_pb2.Batch(id=0, requests=[req_0, req_1], size=2)
     return BloomCausalLMBatch.from_pb(
@@ -56,7 +56,6 @@ def test_batch_from_pb(default_pb_batch, default_bloom_batch):
     assert batch.requests == default_pb_batch.requests
 
     assert len(batch.input_ids) == default_pb_batch.size
-    assert len(batch.input_ids[0]) == 8
     assert batch.input_ids[0][-1] == 10264
     assert torch.all(batch.input_ids[0][:-1] == 3)
 
@@ -85,14 +84,19 @@ def test_causal_lm_batch_type(default_bloom):
 
 
 def test_causal_lm_generate_token(default_bloom, default_bloom_batch):
+    sequence_length = len(default_bloom_batch.all_input_ids[0])
     generated_texts, next_batch = default_bloom.generate_token(default_bloom_batch)
 
     assert generated_texts == []
     assert isinstance(next_batch, CausalLMBatch)
     assert not next_batch.keys_head_dim_last
 
     assert len(next_batch.all_input_ids) == next_batch.size
-    assert len(next_batch.all_input_ids[0]) == len(next_batch.attention_mask[0]) == 9
+    assert (
+        len(next_batch.all_input_ids[0])
+        == len(next_batch.attention_mask[0])
+        == sequence_length + 1
+    )
     assert torch.all(next_batch.all_input_ids[0][-2:] == 10264)
     assert torch.all(next_batch.all_input_ids[0][:-2] == 3)
 
@@ -106,8 +110,12 @@ def test_causal_lm_generate_token(default_bloom, default_bloom_batch):
     assert next_batch.max_sequence_length == next_batch.input_lengths[0]
 
     assert next_batch.past_key_values is not None
-    assert all([p[0].shape == (16, 64, 8) for p in next_batch.past_key_values])
-    assert all([p[1].shape == (16, 8, 64) for p in next_batch.past_key_values])
+    assert all(
+        [p[0].shape == (16, 64, sequence_length) for p in next_batch.past_key_values]
+    )
+    assert all(
+        [p[1].shape == (16, sequence_length, 64) for p in next_batch.past_key_values]
+    )
 
 
 def test_causal_lm_generate_token_completion(default_bloom, default_bloom_batch):