ray-embedding 1.0.0__py3-none-any.whl

ray_embedding/__init__.py

@@ -0,0 +1,2 @@
+from ray_embedding.deploy import deploy_model
+

ray_embedding/deploy.py

@@ -0,0 +1,34 @@
+from typing import Dict, Any, Optional
+from ray.serve import Application
+from ray_embedding.embedding_model import EmbeddingModel
+import torch
+
+
+def deploy_model(args: Dict[str, Any]) -> Application:
+    assert args
+    deployment_name: str = args.pop("deployment", "")
+    assert deployment_name
+
+    model: str = args.pop("model", "")
+    assert model
+
+    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)
+    model_kwargs: Dict[str, Any] = args.pop("model_kwargs", {})
+    if "torch_dtype" in model_kwargs:
+        model_kwargs["torch_dtype"] = model_kwargs["torch_dtype"].strip()
+        if model_kwargs["torch_dtype"] == "float16":
+            model_kwargs["torch_dtype"] = torch.float16
+        elif model_kwargs["torch_dtype"] == "bfloat16":
+            model_kwargs["torch_dtype"] = torch.bfloat16
+        else:
+            del model_kwargs["torch_dtype"]
+
+    deployment = EmbeddingModel.options(name=deployment_name).bind(model=model,
+                                                                   backend=backend,
+                                                                   matryoshka_dim=matryoshka_dim,
+                                                                   trust_remote_code=trust_remote_code,
+                                                                   model_kwargs=model_kwargs
+                                                                   )
+    return deployment

ray_embedding/dto.py

@@ -0,0 +1,16 @@
+from typing import Union, List, Optional
+from pydantic import BaseModel
+
+
+class EmbeddingRequest(BaseModel):
+    """Schema of embedding requests (compatible with OpenAI)"""
+    model: str  # Model name (for compatibility; only one model is used here)
+    input: Union[str, List[str]]  # List of strings to embed
+    dimensions: Optional[int] = None
+
+
+class EmbeddingResponse(BaseModel):
+    """Schema of embedding response (compatible with OpenAI)"""
+    object: str
+    data: List[dict]  # Embedding data including index and vector
+    model: str  # Model name used for embedding

ray_embedding/embedding_model.py

@@ -0,0 +1,93 @@
+import logging
+import os.path
+import time
+from typing import Optional, Dict, Any
+
+import torch
+from fastapi import FastAPI, HTTPException
+from ray import serve
+from sentence_transformers import SentenceTransformer
+
+from ray_embedding.dto import EmbeddingResponse, EmbeddingRequest
+
+web_api = FastAPI(title=f"Ray Embeddings - OpenAI-compatible API")
+
+
+@serve.deployment(
+    ray_actor_options={"num_gpus": 1},
+    autoscaling_config={
+        "target_ongoing_requests": 2,
+        "min_replicas": 0,
+        "initial_replicas": 1,
+        "max_replicas": 1,
+    }
+)
+@serve.ingress(web_api)
+class EmbeddingModel:
+    def __init__(self, model: str, 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.backend = backend or "torch"
+        self.matryoshka_dim = matryoshka_dim
+        self.trust_remote_code = trust_remote_code or False
+        self.model_kwargs = model_kwargs or {}
+        self.torch_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        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,
+                                                   model_kwargs=self.model_kwargs)
+
+        self.served_model_name = os.path.basename(self.model)
+        self.available_models = [
+            {"id": self.served_model_name,
+             "object": "model",
+             "created": int(time.time()),
+             "owned_by": "openai",
+             "permission": []}
+        ]
+        self.logger.info(f"Successfully initialized embedding model {self.model}")
+
+    @web_api.post("/v1/embeddings", response_model=EmbeddingResponse)
+    async def create_embeddings(self, request: EmbeddingRequest):
+        """Generate embeddings for the input text using the specified model."""
+        try:
+            assert request.model == self.served_model_name, (
+                f"Model '{request.model}' is not supported. Use '{self.served_model_name}' instead."
+            )
+
+            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:
+            raise HTTPException(status_code=500, detail=str(e))
+
+    @web_api.get("/v1/models")
+    async def list_models(self):
+        """Returns the list of available models in OpenAI-compatible format."""
+        return {"object": "list", "data": self.available_models}
+
+
+

ray_embedding-1.0.0.dist-info/METADATA

@@ -0,0 +1,24 @@
+Metadata-Version: 2.4
+Name: ray-embedding
+Version: 1.0.0
+Summary: "Deploy SentenceTransformers models  Ray-based embedding service with SentenceTransformers"
+Author: Crispin Almodovar
+Author-email: crispin.almodovar@docorto.ai
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.12
+Description-Content-Type: text/markdown
+
+# ray-embedding
+
+A tool for deploying SentenceTransformers models to a ray cluster.
+
+### Supports the following backends
+
+- sbert-pytorch-gpu
+- sbert-pytorch-cpu
+- sbert-onnx-gpu
+- sbert-onnx-cpu
+- sbert-openvino-cpu
+- fastembed-onnx-cpu

ray_embedding-1.0.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+ray_embedding/__init__.py,sha256=OYJT0rVaaGzY613JqgfktsCgroDnBkGOHxR2FE9UtRU,49
+ray_embedding/deploy.py,sha256=E79J0bcVNXAFlFMVrjZTaSXLnrMZ6LvtRdD4d1DKu1w,1598
+ray_embedding/dto.py,sha256=e91ejZbM_NB9WTjF1YnfuV71cajYIh0vOX8oV_g2OwM,595
+ray_embedding/embedding_model.py,sha256=1sx5jXo61UgLLL8BtFBPflsbBay6J3yG1u2UMIPgAtk,3768
+ray_embedding-1.0.0.dist-info/METADATA,sha256=8DBIYxBOxqsXrtHsDFd4JvIThr_w4yO0iLqo5Ui8UTM,676
+ray_embedding-1.0.0.dist-info/WHEEL,sha256=1tXe9gY0PYatrMPMDd6jXqjfpz_B-Wqm32CPfRC58XU,91
+ray_embedding-1.0.0.dist-info/top_level.txt,sha256=ziCblpJq1YsrryshFqxTRuRMgNuO1_tgvAAkGShATNA,14
+ray_embedding-1.0.0.dist-info/RECORD,,

ray_embedding-1.0.0.dist-info/WHEEL

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

ray_embedding-1.0.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+ray_embedding