support mean pooling in python backend

Author: fxmarty

backends/python/server/text_embeddings_server/layers/pooling.py

@@ -0,0 +1,22 @@
+import torch
+from flash_attn.bert_padding import pad_input
+
+from loguru import logger
+
+def mean_pooling(embedding, cu_seqlens, max_s):
+    # Ideally, rust would pass `indices` to the FlashBatch.
+    seqlens = cu_seqlens[1:].clone()
+    seqlens[0] = cu_seqlens[1]
+    seqlens[1:] -= cu_seqlens[1:-1]
+    batch_size = len(seqlens)
+
+    # Example: indices = [0, 1, 2, 3, 7, 8, 9, 10, 11, 12, 13]
+    mask = torch.zeros(batch_size, max_s, dtype=torch.int32, device=cu_seqlens.device)
+    mask[torch.arange(max_s) < seqlens[:, None].cpu()] = 1
+    indices = torch.nonzero(mask.flatten(), as_tuple=False).flatten()
+
+    embedding_padded = pad_input(embedding, indices, batch_size, max_s)
+    
+    sum_embeddings = torch.sum(embedding_padded, 1)
+
+    return sum_embeddings / seqlens[:, None]

backends/python/server/text_embeddings_server/models/__init__.py

@@ -25,8 +25,6 @@
     __all__.append(FlashBert)
 
 
-class 
-
 def get_model(model_path: Path, dtype: Optional[str], pooling_mode: str):
     if dtype == "float32":
         dtype = torch.float32

backends/python/server/text_embeddings_server/models/default_model.py

@@ -8,14 +8,16 @@
 
 from text_embeddings_server.models import Model
 from text_embeddings_server.models.types import PaddedBatch, Embedding
+from typing import Optional
 
 tracer = trace.get_tracer(__name__)
 
 
 class DefaultModel(Model):
-    def __init__(self, model_path: Path, device: torch.device, dtype: torch.dtype):
+    def __init__(self, model_path: Path, device: torch.device, dtype: torch.dtype, pooling_mode: Optional[str]):
         model = AutoModel.from_pretrained(model_path).to(dtype).to(device)
         self.hidden_size = model.config.hidden_size
+        self.pooling_mode = pooling_mode
 
         self.has_position_ids = (
             inspect.signature(model.forward).parameters.get("position_ids", None)

backends/python/server/text_embeddings_server/models/flash_bert.py

@@ -12,7 +12,8 @@
 from text_embeddings_server.models.types import FlashBatch, Embedding
 from text_embeddings_server.layers.attention import attention
 from text_embeddings_server.layers.layernorm import FastLayerNorm
-from loguru import logger
+from text_embeddings_server.layers.pooling import mean_pooling
+from typing import Optional
 
 tracer = trace.get_tracer(__name__)
 
@@ -190,12 +191,13 @@ def forward(self, input_ids, token_type_ids, position_ids, cu_seqlens, max_s):
 
 
 class FlashBert(Model):
-    def __init__(self, model_path: Path, device: torch.device, dtype: torch.dtype):
+    def __init__(self, model_path: Path, device: torch.device, dtype: torch.dtype, pooling_mode: Optional[str]):
         config = BertConfig.from_pretrained(model_path)
         with safe_open(model_path / "model.safetensors", framework="pt") as f:
             model = FlashBertModel(f, device, dtype, config)
 
         self.hidden_size = config.hidden_size
+        self.pooling_mode = pooling_mode
 
         super(FlashBert, self).__init__(model=model, dtype=dtype, device=device)
 
@@ -205,7 +207,6 @@ def batch_type(self) -> Type[FlashBatch]:
 
     @tracer.start_as_current_span("embed")
     def embed(self, batch: FlashBatch) -> List[Embedding]:
-        logger.info(f"batch.input_ids {batch.input_ids}")
         embedding = self.model.forward(
             input_ids=batch.input_ids,
             token_type_ids=batch.token_type_ids,
@@ -214,9 +215,8 @@ def embed(self, batch: FlashBatch) -> List[Embedding]:
             max_s=batch.max_s,
         )
 
-        if True:
+        if self.pooling_mode == "cls":
             embedding = embedding[batch.cu_seqlens[:-1]]
-            logger.info(f"embedding {embedding.shape}")
             cpu_results = embedding.view(-1).tolist()
 
             return [
@@ -225,4 +225,14 @@ def embed(self, batch: FlashBatch) -> List[Embedding]:
                 )
                 for i in range(len(batch))
             ]
-        elif 
+        elif self.pooling_mode == "mean":
+            res = mean_pooling(embedding, batch.cu_seqlens, batch.max_s)
+            return [
+                Embedding(
+                    values=res[i]
+                )
+                for i in range(len(batch))
+            ]
+
+        else:
+            raise NotImplementedError(f"Pooling {self.pooling_mode} is not implemented in the python backend")

backends/python/src/lib.rs

@@ -23,6 +23,7 @@ impl PythonBackend {
         uds_path: String,
         otlp_endpoint: Option<String>,
         otlp_service_name: String,
+        pooling_mode: String,
     ) -> Result<Self, BackendError> {
         match model_type {
             ModelType::Classifier => {
@@ -31,8 +32,8 @@ impl PythonBackend {
                 ))
             }
             ModelType::Embedding(pool) => {
-                if pool != Pool::Cls {
-                    return Err(BackendError::Start(format!("{pool:?} is not supported")));
+                if pool != Pool::Cls && pool != Pool::Mean {
+                    return Err(BackendError::Start(format!("{pool:?} is not supported in the TEI Python backend. Please open an issue.")));
                 }
                 pool
             }
@@ -44,6 +45,7 @@ impl PythonBackend {
             &uds_path,
             otlp_endpoint,
             otlp_service_name,
+            pooling_mode,
         )?;
         let tokio_runtime = tokio::runtime::Builder::new_current_thread()
             .enable_all()

backends/python/src/management.rs

@@ -22,6 +22,7 @@ impl BackendProcess {
         uds_path: &str,
         otlp_endpoint: Option<String>,
         otlp_service_name: String,
+        pooling_mode: String,
     ) -> Result<Self, BackendError> {
         // Get UDS path
         let uds = Path::new(uds_path);
@@ -52,6 +53,9 @@ impl BackendProcess {
         python_server_args.push("--otlp-service-name".to_owned());
         python_server_args.push(otlp_service_name);
 
+        python_server_args.push("--pooling-mode".to_owned());
+        python_server_args.push(pooling_mode);
+
         // Copy current process env
         let envs: Vec<(OsString, OsString)> = env::vars_os().collect();
 

backends/src/lib.rs

@@ -39,6 +39,7 @@ impl Backend {
         uds_path: String,
         otlp_endpoint: Option<String>,
         otlp_service_name: String,
+        pooling_mode: String,
     ) -> Result<Self, BackendError> {
         let (backend_sender, backend_receiver) = mpsc::unbounded_channel();
 
@@ -49,6 +50,7 @@ impl Backend {
             uds_path,
             otlp_endpoint,
             otlp_service_name,
+            pooling_mode,
         )?;
         let padded_model = backend.is_padded();
         let max_batch_size = backend.max_batch_size();
@@ -138,6 +140,7 @@ fn init_backend(
     uds_path: String,
     otlp_endpoint: Option<String>,
     otlp_service_name: String,
+    pooling_mode: String,
 ) -> Result<Box<dyn CoreBackend + Send>, BackendError> {
     if cfg!(feature = "candle") {
         #[cfg(feature = "candle")]
@@ -158,6 +161,7 @@ fn init_backend(
                         uds_path,
                         otlp_endpoint,
                         otlp_service_name,
+                        pooling_mode,
                     )
                 })
                 .join()

router/src/lib.rs

@@ -105,7 +105,7 @@ pub async fn run(
         serde_json::from_str(&config).context("Failed to parse `config.json`")?;
 
     // Set model type from config
-    let backend_model_type = get_backend_model_type(&config, &model_root, pooling)?;
+    let backend_model_type = get_backend_model_type(&config, &model_root, &pooling)?;
 
     // Info model type
     let model_type = match &backend_model_type {
@@ -191,6 +191,11 @@ pub async fn run(
         }
     });
 
+    let pooling_str = match pooling {
+        Some(pool) => pool.to_string(),
+        None => "none".to_string(),
+    };
+
     // Create backend
     tracing::info!("Starting model backend");
     let backend = text_embeddings_backend::Backend::new(
@@ -200,7 +205,7 @@ pub async fn run(
         uds_path.unwrap_or("/tmp/text-embeddings-inference-server".to_string()),
         otlp_endpoint.clone(),
         otlp_service_name.clone(),
-        pooling.to_string(),
+        pooling_str,
     )
     .context("Could not create backend")?;
     backend
@@ -307,10 +312,10 @@ pub async fn run(
 fn get_backend_model_type(
     config: &ModelConfig,
     model_root: &Path,
-    pooling: Option<text_embeddings_backend::Pool>,
+    pooling: &Option<text_embeddings_backend::Pool>,
 ) -> Result<text_embeddings_backend::ModelType> {
     for arch in &config.architectures {
-        if Some(text_embeddings_backend::Pool::Splade) == pooling && arch.ends_with("MaskedLM") {
+        if Some(text_embeddings_backend::Pool::Splade) == *pooling && arch.ends_with("MaskedLM") {
             return Ok(text_embeddings_backend::ModelType::Embedding(
                 text_embeddings_backend::Pool::Splade,
             ));
@@ -324,15 +329,15 @@ fn get_backend_model_type(
         }
     }
 
-    if Some(text_embeddings_backend::Pool::Splade) == pooling {
+    if Some(text_embeddings_backend::Pool::Splade) == *pooling {
         return Err(anyhow!(
             "Splade pooling is not supported: model is not a ForMaskedLM model"
         ));
     }
 
     // Set pooling
     let pool = match pooling {
-        Some(pool) => pool,
+        Some(pool) => pool.clone(),
         None => {
             // Load pooling config
             let config_path = model_root.join("1_Pooling/config.json");