ray-embedding 0.9.9__tar.gz → 0.10.0__tar.gz

{ray_embedding-0.9.9 → ray_embedding-0.10.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ray-embedding
-Version: 0.9.9
+Version: 0.10.0
 Summary: Deploy SentenceTransformers embedding models to a ray cluster
 Author: Crispin Almodovar
 Author-email: 

{ray_embedding-0.9.9 → ray_embedding-0.10.0}/ray_embedding/embedding_model.py

@@ -69,25 +69,51 @@ class EmbeddingModel:
             assert request.model == self.served_model_name, (
                 f"Model '{request.model}' is not supported. Use '{self.served_model_name}' instead."
             )
-            return await self.__create_embeddings_batch(request)
+
+            if isinstance(request.input, str):
+                request.input = [request.input]
+
+            truncate_dim = request.dimensions or self.matryoshka_dim
+
+            # Compute embeddings and convert to a PyTorch tensor on the GPU
+            embeddings = self.embedding_model.encode(
+                request.input, convert_to_tensor=True, normalize_embeddings=True
+            ).to(self.torch_device)
+
+            if truncate_dim is not None:
+                # Truncate and re-normalize the embeddings
+                embeddings = embeddings[:, :truncate_dim]
+                embeddings = embeddings / torch.norm(embeddings, dim=1, keepdim=True)
+
+            # Move all embeddings to CPU at once before conversion
+            embeddings = embeddings.cpu().tolist()
+
+            # Convert embeddings to list format for response
+            response_data = [
+                {"index": idx, "embedding": emb}
+                for idx, emb in enumerate(embeddings)
+            ]
+            return EmbeddingResponse(object="list", data=response_data, model=request.model)
+
         except Exception as e:
             self.logger.error(e)
             raise HTTPException(status_code=500, detail=str(e))
 
     @serve.batch(max_batch_size=8, batch_wait_timeout_s=0.25)
     async def __create_embeddings_batch(self, requests_batch: List[EmbeddingRequest]) -> List[EmbeddingResponse]:
-        self_0 = self[0] if isinstance(self, list) else self
+        self_0 = self[0] if isinstance(self, list) else self  # Ray also passes an array of self refs; just take the first one
+        embedding_model, matryoshka_dim, torch_device = self_0.embedding_model, self_0.matryoshka_dim, self_0.torch_device
 
         inputs, truncate_dims, num_inputs_list = [], [], []
         for request in requests_batch:
             input_text = request.input if isinstance(request.input, list) else [request.input] # Can be a list of texts
             inputs.extend(input_text)
             num_inputs_list.append(len(input_text))
-            truncate_dims.append(request.dimensions or self_0.matryoshka_dim)
+            truncate_dims.append(request.dimensions or matryoshka_dim)
 
-        embeddings = self_0.embedding_model.encode(
+        embeddings = embedding_model.encode(
             inputs, convert_to_tensor=True, normalize_embeddings=True, show_progress_bar=False,
-        ).to(self_0.torch_device)
+        ).to(torch_device)
 
         model_name = requests_batch[0].model
         truncate_needed = any(dim is not None for dim in truncate_dims)

{ray_embedding-0.9.9 → ray_embedding-0.10.0}/ray_embedding.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ray-embedding
-Version: 0.9.9
+Version: 0.10.0
 Summary: Deploy SentenceTransformers embedding models to a ray cluster
 Author: Crispin Almodovar
 Author-email: 

{ray_embedding-0.9.9 → ray_embedding-0.10.0}/setup.cfg

@@ -1,6 +1,6 @@
 [metadata]
 name = ray-embedding
-version = 0.9.9
+version = 0.10.0
 author = Crispin Almodovar
 author_email = 
 description = Deploy SentenceTransformers embedding models to a ray cluster