ray-embedding 0.9.8__py3-none-any.whl → 0.9.10__py3-none-any.whl

ray_embedding/deploy.py

@@ -29,7 +29,7 @@ def deploy_model(args: Dict[str, Any]) -> Application:
         elif torch_dtype == "float32":
             model_kwargs["torch_dtype"] = torch.float32
         else:
-            del model_kwargs["torch_dtype"] # Remove
+            raise ValueError(f"Invalid torch_dtype: '{torch_dtype}'")
 
     deployment = EmbeddingModel.options(name=deployment_name).bind(model=model,
                                                                    backend=backend,

ray_embedding/embedding_model.py

@@ -76,42 +76,51 @@ class EmbeddingModel:
 
     @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
 
-        # Batch the text inputs
-        inputs = []
-        truncate_dims = []
+        inputs, truncate_dims, num_inputs_list = [], [], []
         for request in requests_batch:
-            if isinstance(request.input, str):
-                request.input = [request.input]
-            inputs.extend(request.input)
-            truncate_dims.append(request.dimensions or self_0.matryoshka_dim)
+            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 matryoshka_dim)
 
-        # Compute embeddings for the batch of text inputs
-        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)
-
-        # Truncate the embeddings; note that the truncate_dim can be different for each request
-        # so we need to this step one by one
-        results = []
-        ix = 0
-        for truncate_dim, request in zip(truncate_dims, requests_batch):
-            num_inputs = len(request.input)
-            batch_embeddings = embeddings[ix: ix + num_inputs]
-            ix += num_inputs
-            if truncate_dim is not None:
-                batch_embeddings = batch_embeddings[:, :truncate_dim]
-                batch_embeddings = batch_embeddings / torch.norm(batch_embeddings, dim=1, keepdim=True)
-
-            batch_embeddings = batch_embeddings.cpu().tolist()
-            response_data = [
-                {"index": emb_ix, "embedding": emb}
-                for emb_ix, emb in enumerate(batch_embeddings)
-            ]
-            results.append(EmbeddingResponse(object="list", data=response_data, model=request.model))
-
-        return results
+        ).to(torch_device)
+
+        model_name = requests_batch[0].model
+        truncate_needed = any(dim is not None for dim in truncate_dims)
+        results_batch, ix = [], 0
+
+        if truncate_needed:
+            for truncate_dim, num_inputs in zip(truncate_dims, num_inputs_list):
+                batch_embeddings = embeddings[ix: ix + num_inputs]
+                ix += num_inputs
+
+                if truncate_dim is not None:
+                    # Truncate and normalize using pytorch (faster)
+                    batch_embeddings = batch_embeddings[:, :truncate_dim]
+                    batch_embeddings = batch_embeddings / torch.norm(batch_embeddings, dim=1, keepdim=True)
+
+                batch_embeddings = batch_embeddings.cpu().tolist()
+                response_data = [
+                    {"index": emb_ix, "embedding": emb} for emb_ix, emb in enumerate(batch_embeddings)
+                ]
+                results_batch.append(EmbeddingResponse(object="list", data=response_data, model=model_name))
+        else:
+            embeddings_list = embeddings.cpu().tolist() # Move everything to CPU
+            for num_inputs in num_inputs_list:
+                batch_embeddings = embeddings_list[ix: ix + num_inputs]
+                ix += num_inputs
+
+                response_data = [
+                    {"index": emb_ix, "embedding": emb} for emb_ix, emb in enumerate(batch_embeddings)
+                ]
+                results_batch.append(EmbeddingResponse(object="list", data=response_data, model=model_name))
+
+        return results_batch
 
     @web_api.get("/v1/models")
     async def list_models(self):

{ray_embedding-0.9.8.dist-info → ray_embedding-0.9.10.dist-info}/METADATA

@@ -1,9 +1,9 @@
 Metadata-Version: 2.4
 Name: ray-embedding
-Version: 0.9.8
+Version: 0.9.10
 Summary: Deploy SentenceTransformers embedding models to a ray cluster
 Author: Crispin Almodovar
-Author-email: crispin.almodovar@docorto.ai
+Author-email: 
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Operating System :: OS Independent
@@ -24,11 +24,5 @@ A tool for deploying SentenceTransformers models to a ray cluster.
 - onnx-cpu
 - openvino-cpu
 - fastembed-onnx-cpu
- 
-- spot instances
-- grpc
 
-### To build:
-- python -m build
-- twine upload dist/*
 

ray_embedding-0.9.10.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=DyhFO0kQnyoFq_VGJQYDEZKuvtwqhuXr83Vu23yb4ds,5966
+ray_embedding-0.9.10.dist-info/METADATA,sha256=aDU4KlY1_-mtJ80tPb9FGGAnB9o2VSjQtsUR3Xt60_k,605
+ray_embedding-0.9.10.dist-info/WHEEL,sha256=SmOxYU7pzNKBqASvQJ7DjX3XGUF92lrGhMb3R6_iiqI,91
+ray_embedding-0.9.10.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
+ray_embedding-0.9.10.dist-info/RECORD,,

{ray_embedding-0.9.8.dist-info → ray_embedding-0.9.10.dist-info}/WHEEL

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

ray_embedding-0.9.8.dist-info/RECORD

@@ -1,8 +0,0 @@
-ray_embedding/__init__.py,sha256=OYJT0rVaaGzY613JqgfktsCgroDnBkGOHxR2FE9UtRU,49
-ray_embedding/deploy.py,sha256=YD_udSm13QbFPgSAkCrTQso15DmtIn0QEhErOFNg7jM,1841
-ray_embedding/dto.py,sha256=e91ejZbM_NB9WTjF1YnfuV71cajYIh0vOX8oV_g2OwM,595
-ray_embedding/embedding_model.py,sha256=aeNgO01Qu8tnhi8ScrdBJDmIoJnNwutRqKUVjqvZM08,5211
-ray_embedding-0.9.8.dist-info/METADATA,sha256=XjQh_exJnrlmNryxnOFgVmzlU9XugcNK_0PnUquD8h0,712
-ray_embedding-0.9.8.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-ray_embedding-0.9.8.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
-ray_embedding-0.9.8.dist-info/RECORD,,