ray-embedding 0.9.9__py3-none-any.whl → 0.10.0__py3-none-any.whl

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.dist-info → ray_embedding-0.10.0.dist-info}/METADATA

@@ -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.10.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+ray_embedding/__init__.py,sha256=OYJT0rVaaGzY613JqgfktsCgroDnBkGOHxR2FE9UtRU,49
+ray_embedding/deploy.py,sha256=oqOhMxS5UyZ4oGhfpL7kqvrxLO8QW41sF_FHbNSJe-w,1858
+ray_embedding/dto.py,sha256=e91ejZbM_NB9WTjF1YnfuV71cajYIh0vOX8oV_g2OwM,595
+ray_embedding/embedding_model.py,sha256=cyGH7CZRCAV50T51aTAxTSRCU0haCV_mX42tyhmFa6U,7007
+ray_embedding-0.10.0.dist-info/METADATA,sha256=8uU8yj2bnbxZPlL1WEaU0V6ioi0We6rcUUzTGd62yTk,605
+ray_embedding-0.10.0.dist-info/WHEEL,sha256=GHB6lJx2juba1wDgXDNlMTyM13ckjBMKf-OnwgKOCtA,91
+ray_embedding-0.10.0.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
+ray_embedding-0.10.0.dist-info/RECORD,,

{ray_embedding-0.9.9.dist-info → ray_embedding-0.10.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (78.1.0)
+Generator: setuptools (80.3.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

ray_embedding-0.9.9.dist-info/RECORD

@@ -1,8 +0,0 @@
-ray_embedding/__init__.py,sha256=OYJT0rVaaGzY613JqgfktsCgroDnBkGOHxR2FE9UtRU,49
-ray_embedding/deploy.py,sha256=oqOhMxS5UyZ4oGhfpL7kqvrxLO8QW41sF_FHbNSJe-w,1858
-ray_embedding/dto.py,sha256=e91ejZbM_NB9WTjF1YnfuV71cajYIh0vOX8oV_g2OwM,595
-ray_embedding/embedding_model.py,sha256=PBBuntLrJiwstD9AbD2n_ElU7rYZAYK0fMSESkMsFTU,5797
-ray_embedding-0.9.9.dist-info/METADATA,sha256=XRtQFNU7mhhSXFZF1YMyIII3zW9PfTXkiUN6AEHVAkA,604
-ray_embedding-0.9.9.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-ray_embedding-0.9.9.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
-ray_embedding-0.9.9.dist-info/RECORD,,