ray-embedding 0.10.5__py3-none-any.whl → 0.10.11__py3-none-any.whl

ray_embedding/deploy.py

@@ -16,10 +16,10 @@ def deploy_model(args: Dict[str, Any]) -> Application:
     model: str = args.pop("model", "")
     assert model
 
+    device: Optional[str] = args.pop("device", None)
     backend: Optional[str] = args.pop("backend", "torch")
     matryoshka_dim: Optional[int] = args.pop("matryoshka_dim", None)
     trust_remote_code: Optional[bool] = args.pop("trust_remote_code", False)
-    device: Optional[str] = args.pop("device", None)
     model_kwargs: Dict[str, Any] = args.pop("model_kwargs", {})
     if "torch_dtype" in model_kwargs:
         torch_dtype = model_kwargs["torch_dtype"].strip()
@@ -33,8 +33,8 @@ def deploy_model(args: Dict[str, Any]) -> Application:
             raise ValueError(f"Invalid torch_dtype: '{torch_dtype}'")
 
     deployment = EmbeddingModel.options(name=deployment_name).bind(model=model,
-                                                                   backend=backend,
                                                                    device=device,
+                                                                   backend=backend,
                                                                    matryoshka_dim=matryoshka_dim,
                                                                    trust_remote_code=trust_remote_code,
                                                                    model_kwargs=model_kwargs

ray_embedding/embedding_model.py

@@ -28,23 +28,24 @@ web_api = FastAPI(title=f"Ray Embeddings - OpenAI-compatible API")
 )
 @serve.ingress(web_api)
 class EmbeddingModel:
-    def __init__(self, model: str, backend: Optional[str] = "torch", device: Optional[str] = None,
+    def __init__(self, model: str, device: Optional[str] = None, backend: Optional[str] = "torch",
                  matryoshka_dim: Optional[int] = None, trust_remote_code: Optional[bool] = False,
                  model_kwargs: Dict[str, Any] = None):
         logging.basicConfig(level=logging.INFO)
         self.logger = logging.getLogger(__name__)
         self.model = model
         self.init_device = device
+        if self.init_device is None or self.init_device == "auto":
+            self.init_device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.check_health()
+        self.torch_device = torch.device(self.init_device)
         self.backend = backend or "torch"
         self.matryoshka_dim = matryoshka_dim
         self.trust_remote_code = trust_remote_code or False
         self.model_kwargs = model_kwargs or {}
-        if self.init_device is None:
-            self.torch_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-        else:
-            self.torch_device = torch.device(self.init_device)
         self.logger.info(f"Initializing embedding model: {self.model}")
-        self.embedding_model = SentenceTransformer(self.model, backend=self.backend, trust_remote_code=self.trust_remote_code,
+        self.embedding_model = SentenceTransformer(self.model, device=self.init_device, backend=self.backend,
+                                                   trust_remote_code=self.trust_remote_code,
                                                    model_kwargs=self.model_kwargs)
 
         self.served_model_name = os.path.basename(self.model)

{ray_embedding-0.10.5.dist-info → ray_embedding-0.10.11.dist-info}/METADATA

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

ray_embedding-0.10.11.dist-info/RECORD

@@ -0,0 +1,8 @@
+ray_embedding/__init__.py,sha256=OYJT0rVaaGzY613JqgfktsCgroDnBkGOHxR2FE9UtRU,49
+ray_embedding/deploy.py,sha256=ZGxcG4589WcRtaM6H84YJarw0m1XqHNgfOf3PLAhM5M,1995
+ray_embedding/dto.py,sha256=e91ejZbM_NB9WTjF1YnfuV71cajYIh0vOX8oV_g2OwM,595
+ray_embedding/embedding_model.py,sha256=SXRc3PtfJkEIHdLvHFCTSTXo8c_NYaL3-5xkxHr3VD4,4603
+ray_embedding-0.10.11.dist-info/METADATA,sha256=413SNOlhdIf-k8kOpzWoQZ5Z_gCFAbova7NTvulnYb8,1095
+ray_embedding-0.10.11.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
+ray_embedding-0.10.11.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
+ray_embedding-0.10.11.dist-info/RECORD,,

ray_embedding-0.10.5.dist-info/RECORD

@@ -1,8 +0,0 @@
-ray_embedding/__init__.py,sha256=OYJT0rVaaGzY613JqgfktsCgroDnBkGOHxR2FE9UtRU,49
-ray_embedding/deploy.py,sha256=GEKJAaV25DOlXd7Pqj7bqASM-FcP0FtFsJlI7u_U6Iw,1995
-ray_embedding/dto.py,sha256=e91ejZbM_NB9WTjF1YnfuV71cajYIh0vOX8oV_g2OwM,595
-ray_embedding/embedding_model.py,sha256=i0sno946g2nbnhhhcD-bZ3WoWJi5qScbm716-i6tI1I,4501
-ray_embedding-0.10.5.dist-info/METADATA,sha256=L77LrGKHM2GYVgpXGVqkRYf5LoSCkbHpar6vkwlxIp4,1094
-ray_embedding-0.10.5.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
-ray_embedding-0.10.5.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
-ray_embedding-0.10.5.dist-info/RECORD,,