Add return_offsets and truncate_input_tokens tokenize API options (#57)

By popular request

Signed-off-by: Nick Hill <[email protected]>

Author: njhill

integration_tests/test_cases_mt0small.yaml

@@ -50,6 +50,100 @@
           - </s>
 
 
+# Tokenize with offsets
+- name: Tokenize with offsets
+  request_type: tokenize
+  request:
+    return_offsets: true
+    requests:
+      - {"text": "The very long story is written"}
+  response:
+    responses:
+      - offsets:
+          - end: 3
+          - end: 4
+            start: 3
+          - end: 8
+            start: 4
+          - end: 13
+            start: 8
+          - end: 19
+            start: 13
+          - end: 22
+            start: 19
+          - end: 23
+            start: 22
+          - end: 30
+            start: 23
+          - {}
+        tokenCount: 9
+
+
+# Tokenize with tokens and offsets
+- name: Tokenize with tokens and offsets
+  request_type: tokenize
+  request:
+    return_tokens: true
+    return_offsets: true
+    requests:
+      - { "text": "The very long story is written" }
+  response:
+    responses:
+      - offsets:
+          - end: 3
+          - end: 4
+            start: 3
+          - end: 8
+            start: 4
+          - end: 13
+            start: 8
+          - end: 19
+            start: 13
+          - end: 22
+            start: 19
+          - end: 23
+            start: 22
+          - end: 30
+            start: 23
+          - {}
+        tokenCount: 9
+        tokens:
+          - "\u2581The"
+          - "\u2581"
+          - very
+          - "\u2581long"
+          - "\u2581story"
+          - "\u2581is"
+          - "\u2581"
+          - written
+          - </s>
+
+
+  # Tokenize with truncate
+- name: Tokenize with tokens and truncation
+  request_type: tokenize
+  request:
+    return_tokens: true
+    truncate_input_tokens: 10
+    requests:
+      - {"text": "The very long story is written by a very long story"}
+  response:
+    responses:
+      - tokenCount: 10
+        # Truncation happens on the left
+        tokens:
+          - "\u2581"
+          - written
+          - "\u2581by"
+          - "\u2581"
+          - a
+          - "\u2581"
+          - very
+          - "\u2581long"
+          - "\u2581story"
+          - </s>
+
+
 # Basic Greedy (implicit)
 - name: Basic Greedy, max new tokens (implicit)
   request:

proto/generation.proto

@@ -197,7 +197,11 @@ message TokenInfo {
 message BatchedTokenizeRequest {
   string model_id = 1;
   repeated TokenizeRequest requests = 2;
-  bool return_tokens = 3; //TBD
+  bool return_tokens = 3;
+  bool return_offsets = 4;
+
+  // Zero means don't truncate.
+  uint32 truncate_input_tokens = 5;
 }
 
 message BatchedTokenizeResponse {
@@ -209,10 +213,17 @@ message TokenizeRequest {
 }
 
 message TokenizeResponse {
+  message Offset {
+    uint32 start = 1;
+    uint32 end = 2;
+  }
+
   uint32 token_count = 1;
-  repeated string tokens = 2; // if include_tokens = true
 
-  // We'll possibly add more later
+  // if return_tokens = true
+  repeated string tokens = 2;
+  // if return_tokens = true
+  repeated Offset offsets = 3;
 }
 
 

router/src/grpc_server.rs

@@ -31,6 +31,7 @@ use crate::{
     validation::{RequestSize, ValidationError},
     GenerateParameters, GenerateRequest,
 };
+use crate::pb::fmaas::tokenize_response::Offset;
 
 /// Whether to fail if sampling parameters are provided in greedy-mode requests
 /// or to silently ignore them.
@@ -339,17 +340,37 @@ impl GenerationService for GenerationServicer {
         let br = request.into_inner();
         metrics::increment_counter!("tgi_tokenize_request_count");
         let start_time = Instant::now();
-        self.tokenize_input_counter
-            .increment(br.requests.len() as u64);
-
-        let responses = try_join_all(br.requests.into_iter().map(|tr| {
-            self.tokenizer.tokenize(tr.text, br.return_tokens).map_ok(
-                |(_, token_count, encoding)| TokenizeResponse {
-                    token_count: token_count as u32,
-                    tokens: encoding.map_or_else(Vec::new, |e| e.get_tokens().to_vec()),
+        self.tokenize_input_counter.increment(br.requests.len() as u64);
+
+        let truncate_to = match br.truncate_input_tokens {
+            0 => u32::MAX,
+            length => length,
+        };
+        let include_encoding = br.return_tokens || br.return_offsets;
+        let responses = try_join_all(br.requests.into_iter().map(|tr|
+            self.tokenizer.tokenize(tr.text, include_encoding).map_ok(
+                |(_, token_count, encoding)| {
+                    let token_count = token_count as u32;
+                    let from = token_count.saturating_sub(truncate_to) as usize;
+                    TokenizeResponse {
+                        token_count: token_count.min(truncate_to),
+                        tokens: match br.return_tokens {
+                            true => encoding.as_ref().unwrap().get_tokens()[from..].to_vec(),
+                            false => vec![],
+                        },
+                        offsets: match br.return_offsets {
+                            true => encoding.unwrap().get_offsets()[from..].iter().map(
+                                |(start, end)| Offset{
+                                    start: *start as u32,
+                                    end: *end as u32,
+                                }
+                            ).collect(),
+                            false => vec![],
+                        },
+                    }
                 },
             )
-        }))
+        ))
         .map_err(Status::from_error)
         .await?;