✨ Chunker client

Signed-off-by: Evaline Ju <69598118+evaline-ju@users.noreply.github.com>

Author: evaline-ju

build.rs

@@ -9,6 +9,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
         .include_file("mod.rs")
         .compile(
             &[
+                "protos/caikit_runtime_Chunkers.proto",
                 "protos/caikit_runtime_Nlp.proto",
                 "protos/generation.proto",
                 "protos/caikit_data_model_caikit_nlp.proto",

protos/caikit_runtime_Chunkers.proto

@@ -0,0 +1,31 @@
+
+/*------------------------------------------------------------------------------
+ * AUTO GENERATED
+ *----------------------------------------------------------------------------*/
+
+syntax = "proto3";
+package caikit.runtime.Chunkers;
+import "caikit_data_model_nlp.proto";
+
+
+/*-- MESSAGES ----------------------------------------------------------------*/
+
+message BidiStreamingTokenizationTaskRequest {
+
+  /*-- fields --*/
+  string text_stream = 1;
+}
+
+message TokenizationTaskRequest {
+
+  /*-- fields --*/
+  string text = 1;
+}
+
+
+/*-- SERVICES ----------------------------------------------------------------*/
+
+service ChunkersService {
+  rpc BidiStreamingTokenizationTaskPredict(stream caikit.runtime.Chunkers.BidiStreamingTokenizationTaskRequest) returns (stream caikit_data_model.nlp.TokenizationStreamResult);
+  rpc TokenizationTaskPredict(caikit.runtime.Chunkers.TokenizationTaskRequest) returns (caikit_data_model.nlp.TokenizationResults);
+}

protos/caikit_runtime_Nlp.proto

@@ -83,7 +83,6 @@
  /*-- SERVICES ----------------------------------------------------------------*/
 
  service NlpService {
-   rpc BidiStreamingTokenizationTaskPredict(stream caikit.runtime.Nlp.BidiStreamingTokenizationTaskRequest) returns (stream caikit_data_model.nlp.TokenizationStreamResult);
    rpc ServerStreamingTextGenerationTaskPredict(caikit.runtime.Nlp.ServerStreamingTextGenerationTaskRequest) returns (stream caikit_data_model.nlp.GeneratedTextStreamResult);
    rpc TextGenerationTaskPredict(caikit.runtime.Nlp.TextGenerationTaskRequest) returns (caikit_data_model.nlp.GeneratedTextResult);
    rpc TokenizationTaskPredict(caikit.runtime.Nlp.TokenizationTaskRequest) returns (caikit_data_model.nlp.TokenizationResults);

src/clients.rs

@@ -7,6 +7,9 @@ use url::Url;
 
 use crate::config::{ServiceConfig, Tls};
 
+pub mod chunker;
+pub use chunker::ChunkerClient;
+
 pub mod detector;
 pub use detector::DetectorClient;
 
@@ -18,6 +21,7 @@ pub use nlp::NlpClient;
 
 pub const DEFAULT_TGIS_PORT: u16 = 8033;
 pub const DEFAULT_CAIKIT_NLP_PORT: u16 = 8085;
+pub const DEFAULT_CHUNKER_PORT: u16 = 8085;
 pub const DEFAULT_DETECTOR_PORT: u16 = 8080;
 const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
 const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(10);

src/clients/chunker.rs

@@ -0,0 +1,85 @@
+use std::{collections::HashMap, pin::Pin};
+
+use futures::{Stream, StreamExt};
+use ginepro::LoadBalancedChannel;
+use tokio::sync::mpsc;
+use tokio_stream::wrappers::ReceiverStream;
+use tonic::Request;
+
+use super::{create_grpc_clients, Error};
+use crate::{
+    config::ServiceConfig,
+    pb::{
+        caikit::runtime::chunkers::{
+            chunkers_service_client::ChunkersServiceClient, BidiStreamingTokenizationTaskRequest,
+            TokenizationTaskRequest,
+        },
+        caikit_data_model::nlp::{
+            TokenizationResults, TokenizationStreamResult,
+        },
+    },
+};
+
+const MODEL_ID_HEADER_NAME: &str = "mm-model-id";
+
+#[derive(Clone)]
+pub struct ChunkerClient {
+    clients: HashMap<String, ChunkersServiceClient<LoadBalancedChannel>>,
+}
+
+impl ChunkerClient {
+    pub async fn new(default_port: u16, config: &[(String, ServiceConfig)]) -> Result<Self, Error> {
+        let clients = create_grpc_clients(default_port, config, ChunkersServiceClient::new).await?;
+        Ok(Self { clients })
+    }
+
+    fn client(&self, model_id: &str) -> Result<ChunkersServiceClient<LoadBalancedChannel>, Error> {
+        Ok(self
+            .clients
+            .get(model_id)
+            .ok_or_else(|| Error::ModelNotFound(model_id.into()))?
+            .clone())
+    }
+
+    pub async fn tokenization_task_predict(
+        &self,
+        model_id: &str,
+        request: TokenizationTaskRequest,
+    ) -> Result<TokenizationResults, Error> {
+        let request = request_with_model_id(request, model_id);
+        Ok(self
+            .client(model_id)?
+            .tokenization_task_predict(request)
+            .await?
+            .into_inner())
+    }
+
+    pub async fn bidi_streaming_tokenization_task_predict(
+        &self,
+        model_id: &str,
+        request: Pin<Box<dyn Stream<Item = BidiStreamingTokenizationTaskRequest> + Send + 'static>>,
+    ) -> Result<ReceiverStream<TokenizationStreamResult>, Error> {
+        let request = request_with_model_id(request, model_id);
+        let mut response_stream = self
+            .client(model_id)?
+            .bidi_streaming_tokenization_task_predict(request)
+            .await?
+            .into_inner();
+        let (tx, rx) = mpsc::channel(128);
+        tokio::spawn(async move {
+            while let Some(Ok(message)) = response_stream.next().await {
+                let _ = tx.send(message).await;
+            }
+        });
+        Ok(ReceiverStream::new(rx))
+    }
+
+}
+
+fn request_with_model_id<T>(request: T, model_id: &str) -> Request<T> {
+    let mut request = Request::new(request);
+    request
+        .metadata_mut()
+        .insert(MODEL_ID_HEADER_NAME, model_id.parse().unwrap());
+    request
+}

src/clients/nlp.rs

@@ -11,13 +11,13 @@ use crate::{
     config::ServiceConfig,
     pb::{
         caikit::runtime::nlp::{
-            nlp_service_client::NlpServiceClient, BidiStreamingTokenizationTaskRequest,
+            nlp_service_client::NlpServiceClient,
             ServerStreamingTextGenerationTaskRequest, TextGenerationTaskRequest,
             TokenClassificationTaskRequest, TokenizationTaskRequest,
         },
         caikit_data_model::nlp::{
             GeneratedTextResult, GeneratedTextStreamResult, TokenClassificationResults,
-            TokenizationResults, TokenizationStreamResult,
+            TokenizationResults,
         },
     },
 };
@@ -56,26 +56,6 @@ impl NlpClient {
             .into_inner())
     }
 
-    pub async fn bidi_streaming_tokenization_task_predict(
-        &self,
-        model_id: &str,
-        request: Pin<Box<dyn Stream<Item = BidiStreamingTokenizationTaskRequest> + Send + 'static>>,
-    ) -> Result<ReceiverStream<TokenizationStreamResult>, Error> {
-        let request = request_with_model_id(request, model_id);
-        let mut response_stream = self
-            .client(model_id)?
-            .bidi_streaming_tokenization_task_predict(request)
-            .await?
-            .into_inner();
-        let (tx, rx) = mpsc::channel(128);
-        tokio::spawn(async move {
-            while let Some(Ok(message)) = response_stream.next().await {
-                let _ = tx.send(message).await;
-            }
-        });
-        Ok(ReceiverStream::new(rx))
-    }
-
     pub async fn token_classification_task_predict(
         &self,
         model_id: &str,

src/orchestrator.rs

@@ -10,7 +10,7 @@ use uuid::Uuid;
 
 use crate::{
     clients::{
-        self, detector::DetectorRequest, DetectorClient, GenerationClient, NlpClient, TgisClient,
+        self, detector::DetectorRequest, ChunkerClient, DetectorClient, GenerationClient, NlpClient, TgisClient,
     },
     config::{GenerationProvider, OrchestratorConfig},
     models::{
@@ -19,6 +19,7 @@ use crate::{
         InputWarningReason, TextGenTokenClassificationResults, TokenClassificationResult,
     },
     pb::{
+        caikit::runtime::chunkers::TokenizationTaskRequest as ChunkersTokenizationTaskRequest,
         caikit::runtime::nlp::{TextGenerationTaskRequest, TokenizationTaskRequest},
         fmaas::{
             BatchedGenerationRequest, BatchedTokenizeRequest, GenerationRequest, TokenizeRequest,
@@ -34,7 +35,7 @@ const UNSUITABLE_INPUT_MESSAGE: &str = "Unsuitable input detected. \
 struct Context {
     config: OrchestratorConfig,
     generation_client: GenerationClient,
-    chunker_client: NlpClient,
+    chunker_client: ChunkerClient,
     detector_client: DetectorClient,
 }
 
@@ -248,7 +249,7 @@ async fn handle_chunk_task(
             let ctx = ctx.clone();
             let chunker_id = chunker_id.clone();
             async move {
-                let request = TokenizationTaskRequest { text };
+                let request = ChunkersTokenizationTaskRequest { text };
                 debug!(
                     %chunker_id,
                     ?request,
@@ -495,7 +496,7 @@ fn apply_masks(text: &str, masks: &[(usize, usize)]) -> Vec<(usize, String)> {
 
 async fn create_clients(
     config: &OrchestratorConfig,
-) -> Result<(GenerationClient, NlpClient, DetectorClient), Error> {
+) -> Result<(GenerationClient, ChunkerClient, DetectorClient), Error> {
     // TODO: create better solution for routers
     let generation_client = match config.generation.provider {
         GenerationProvider::Tgis => {
@@ -521,7 +522,7 @@ async fn create_clients(
         .iter()
         .map(|(chunker_id, config)| (chunker_id.clone(), config.service.clone()))
         .collect::<Vec<_>>();
-    let chunker_client = NlpClient::new(clients::DEFAULT_CAIKIT_NLP_PORT, &chunker_config).await?;
+    let chunker_client = ChunkerClient::new(clients::DEFAULT_CHUNKER_PORT, &chunker_config).await?;
 
     let detector_config = config
         .detectors