🏷️ Update based on open-sourced types (#185)

* 🏷️ Update based on open-sourced types Signed-off-by: Evaline Ju <[email protected]> * 💡 Add caikit reference as comments Signed-off-by: Evaline Ju <[email protected]> --------- Signed-off-by: Evaline Ju <[email protected]>
foundation-model-stack · Sep 6, 2024 · f970d56 · f970d56
1 parent 4d136ad
commit f970d56
Show file tree

Hide file tree

Showing 5 changed files with 19 additions and 23 deletions.
diff --git a/protos/caikit_data_model_nlp.proto b/protos/caikit_data_model_nlp.proto
@@ -1,4 +1,4 @@
-
+// Ref. https://github.com/caikit/caikit/blob/main/caikit/interfaces/nlp/data_model/text.py
 /*------------------------------------------------------------------------------
  * AUTO GENERATED
  *----------------------------------------------------------------------------*/
@@ -27,6 +27,16 @@
 
  /*-- MESSAGES ----------------------------------------------------------------*/
 
+  message ChunkerTokenizationStreamResult {
+
+    /*-- fields --*/
+    repeated caikit_data_model.nlp.Token results = 1;
+    int64 token_count = 4;
+    int64 processed_index = 2;
+    int64 start_index = 3;
+    int64 input_start_index = 20;
+    int64 input_end_index = 21;
+  }
 
  message GeneratedTextResult {
 

diff --git a/protos/caikit_runtime_Chunkers.proto b/protos/caikit_runtime_Chunkers.proto
@@ -1,4 +1,4 @@
-
+// Ref. https://github.com/caikit/caikit/blob/main/caikit/interfaces/nlp/tasks.py
 /*------------------------------------------------------------------------------
  * AUTO GENERATED
  *----------------------------------------------------------------------------*/
@@ -23,21 +23,10 @@ message ChunkerTokenizationTaskRequest {
   string text = 1;
 }
 
-message ChunkerTokenizationStreamResult {
-
-  /*-- fields --*/
-  repeated caikit_data_model.nlp.Token results = 1;
-  int64 token_count = 4;
-  int64 processed_index = 2;
-  int64 start_index = 3;
-  int64 input_start_index = 5;
-  int64 input_end_index = 6;
-}
-
 
 /*-- SERVICES ----------------------------------------------------------------*/
 
 service ChunkersService {
-  rpc BidiStreamingChunkerTokenizationTaskPredict(stream BidiStreamingChunkerTokenizationTaskRequest) returns (stream ChunkerTokenizationStreamResult);
-  rpc ChunkerTokenizationTaskPredict(ChunkerTokenizationTaskRequest) returns (caikit_data_model.nlp.TokenizationResults);
+  rpc BidiStreamingChunkerTokenizationTaskPredict(stream caikit.runtime.Chunkers.BidiStreamingChunkerTokenizationTaskRequest) returns (stream caikit_data_model.nlp.ChunkerTokenizationStreamResult);
+  rpc ChunkerTokenizationTaskPredict(caikit.runtime.Chunkers.ChunkerTokenizationTaskRequest) returns (caikit_data_model.nlp.TokenizationResults);
 }
diff --git a/src/clients/chunker.rs b/src/clients/chunker.rs
@@ -30,10 +30,9 @@ use crate::{
     pb::{
         caikit::runtime::chunkers::{
             chunkers_service_client::ChunkersServiceClient,
-            BidiStreamingChunkerTokenizationTaskRequest, ChunkerTokenizationStreamResult,
-            ChunkerTokenizationTaskRequest,
+            BidiStreamingChunkerTokenizationTaskRequest, ChunkerTokenizationTaskRequest,
         },
-        caikit_data_model::nlp::{Token, TokenizationResults},
+        caikit_data_model::nlp::{ChunkerTokenizationStreamResult, Token, TokenizationResults},
     },
 };
 

diff --git a/src/orchestrator/streaming.rs b/src/orchestrator/streaming.rs
@@ -37,7 +37,8 @@ use crate::{
         unary::{input_detection_task, tokenize},
         UNSUITABLE_INPUT_MESSAGE,
     },
-    pb::caikit::runtime::chunkers::{self, ChunkerTokenizationStreamResult},
+    pb::caikit::runtime::chunkers,
+    pb::caikit_data_model::nlp::ChunkerTokenizationStreamResult,
 };
 
 pub type Chunk = ChunkerTokenizationStreamResult;

diff --git a/src/orchestrator/streaming/aggregator.rs b/src/orchestrator/streaming/aggregator.rs
@@ -446,10 +446,7 @@ mod tests {
     use super::*;
     use crate::{
         models::TokenClassificationResult,
-        pb::{
-            caikit::runtime::chunkers::ChunkerTokenizationStreamResult,
-            caikit_data_model::nlp::Token,
-        },
+        pb::caikit_data_model::nlp::{ChunkerTokenizationStreamResult, Token},
     };
 
     fn get_detection_obj(