Add support for batching to embedder models (#503)

Co-authored-by: Magdy Saleh <[email protected]>

proto/generate.proto

@@ -232,16 +232,24 @@ message DecodeResponse {
     optional CachedBatch batch = 2;
 }
 
-message EmbedRequest {
-    string inputs = 1;
+message Embedding {
+    /// Request ID
+    uint64 request_id = 1;
+
+    /// Embedding values
+    repeated float values = 2;
 }
 
-message Embedding {
-    repeated float values = 1;
+message EmbedRequest {
+    /// Batch
+    Batch batch = 1;
 }
 
 message EmbedResponse {
-    Embedding embeddings = 1;
+    /// Embeddings
+    repeated Embedding embeddings = 1;
+
+    /// Error message on failure
     string errorMsg = 2;
 }
 

router/Cargo.toml

@@ -46,10 +46,11 @@ utoipa-swagger-ui = { version = "3.0.2", features = ["axum"] }
 ngrok = { version = "0.12.3", features = ["axum"], optional = true }
 once_cell = "1.19.0"
 itertools = "0.12.1"
+async-trait = "0.1.80"
 
 [build-dependencies]
 vergen = { version = "8.0.0", features = ["build", "git", "gitcl"] }
 
 [features]
 default = ["ngrok"]
-ngrok = ["dep:ngrok"]
+ngrok = ["dep:ngrok"]

router/client/src/client.rs

@@ -64,14 +64,6 @@ impl Client {
         Ok(response)
     }
 
-    /// Embed
-    #[instrument(skip(self))]
-    pub async fn embed(&mut self, inputs: String) -> Result<EmbedResponse> {
-        let request = tonic::Request::new(EmbedRequest { inputs }).inject_context();
-        let response = self.stub.embed(request).await?.into_inner();
-        Ok(response)
-    }
-
     /// Get model health
     #[instrument(skip(self))]
     pub async fn health(&mut self) -> Result<HealthResponse> {
@@ -196,6 +188,14 @@ impl Client {
         Ok((response.generations, response.batch))
     }
 
+    /// Embed
+    #[instrument(skip(self))]
+    pub async fn embed(&mut self, batch: Batch) -> Result<Vec<Embedding>> {
+        let request = tonic::Request::new(EmbedRequest { batch: Some(batch) }).inject_context();
+        let response = self.stub.embed(request).await?.into_inner();
+        Ok(response.embeddings)
+    }
+
     /// Downloads the weights for an adapter.
     pub async fn download_adapter(
         &mut self,

router/client/src/lib.rs

@@ -9,7 +9,7 @@ pub use client::Client;
 pub use pb::generate::v1::HealthResponse;
 pub use pb::generate::v1::InfoResponse as ShardInfo;
 pub use pb::generate::v1::{
-    AdapterParameters, AlternativeTokens, Batch, CachedBatch, DownloadAdapterResponse,
+    AdapterParameters, AlternativeTokens, Batch, CachedBatch, DownloadAdapterResponse, Embedding,
     FinishReason, GeneratedText, Generation, MajoritySignMethod, MergeStrategy,
     NextTokenChooserParameters, NextTokens, PrefillTokens, Request, StoppingCriteriaParameters,
 };

router/client/src/sharded_client.rs

@@ -1,4 +1,4 @@
-use crate::pb::generate::v1::EmbedResponse;
+use crate::pb::generate::v1::{EmbedResponse, Embedding};
 /// Multi shard Client
 use crate::{
     AdapterParameters, Batch, CachedBatch, Client, DownloadAdapterResponse, Generation,
@@ -154,15 +154,16 @@ impl ShardedClient {
         merge_generations(results?)
     }
 
-    /// Get the model info
+    /// Embed the given batch
     #[instrument(skip(self))]
-    pub async fn embed(&mut self, inputs: String) -> Result<Vec<EmbedResponse>> {
+    pub async fn embed(&mut self, batch: Batch) -> Result<Vec<Embedding>> {
         let futures: Vec<_> = self
             .clients
             .iter_mut()
-            .map(|client| Box::pin(client.embed(inputs.clone())))
+            .map(|client| Box::pin(client.embed(batch.clone())))
             .collect();
-        join_all(futures).await.into_iter().collect()
+        let results: Result<Vec<Vec<Embedding>>> = join_all(futures).await.into_iter().collect();
+        Ok(results?.into_iter().flatten().collect())
     }
 
     pub async fn download_adapter(