feat(python): add cls and mean pooling (#402)

huggingface · Sep 17, 2024 · 80259c9 · 80259c9
1 parent c6c5e45
commit 80259c9
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/backends/python/server/text_embeddings_server/models/default_model.py b/backends/python/server/text_embeddings_server/models/default_model.py
@@ -43,7 +43,14 @@ def embed(self, batch: PaddedBatch) -> List[Embedding]:
             kwargs["position_ids"] = batch.position_ids
 
         output = self.model(**kwargs)
-        embedding = output[0][:, 0]
+
+        if self.pooling_mode == "cls":
+            embedding = output[0][:, 0]
+        elif self.pooling_mode == "mean":
+            embedding = output[0].mean(dim=1)
+        else:
+            raise NotImplementedError(f"Pooling {self.pooling_mode} is not implemented in the python backend")
+
         cpu_results = embedding.view(-1).tolist()
 
         return [