Add support for open source models based on text-embeddings-inference

README.md

@@ -2,7 +2,7 @@
 
 The `milvus-model` library provides the integration with common embedding and reranker models for Milvus, a high performance open-source vector database built for AI applications.  `milvus-model` lib is included as a dependency in `pymilvus`, the Python SDK of Milvus.
 
-`milvus-model` supports embedding and reranker models from service providers like OpenAI, Voyage AI, Cohere, and open-source models through SentenceTransformers.
+`milvus-model` supports embedding and reranker models from service providers like OpenAI, Voyage AI, Cohere, and open-source models through SentenceTransformers or [text-embeddings-inference](https://github.com/huggingface/text-embeddings-inference) .
 
 `milvus-model` supports Python 3.8 and above.
 

src/pymilvus/model/dense/__init__.py

@@ -2,6 +2,7 @@
 from pymilvus.model.dense.sentence_transformer import SentenceTransformerEmbeddingFunction
 from pymilvus.model.dense.voyageai import VoyageEmbeddingFunction
 from pymilvus.model.dense.jinaai import JinaEmbeddingFunction
+from pymilvus.model.dense.opensource import OpenSourceEmbeddingFunction
 from pymilvus.model.dense.onnx import OnnxEmbeddingFunction
 from pymilvus.model.dense.cohere import CohereEmbeddingFunction
 from pymilvus.model.dense.mistralai import MistralAIEmbeddingFunction
@@ -13,6 +14,7 @@
     "SentenceTransformerEmbeddingFunction",
     "VoyageEmbeddingFunction",
     "JinaEmbeddingFunction",
+    "OpenSourceEmbeddingFunction",
     "OnnxEmbeddingFunction",
     "CohereEmbeddingFunction",
     "MistralAIEmbeddingFunction",

src/pymilvus/model/dense/opensource.py

@@ -0,0 +1,47 @@
+from typing import List, Optional
+
+import numpy as np
+import requests
+
+from pymilvus.model.base import BaseEmbeddingFunction
+
+
+class OpenSourceEmbeddingFunction(BaseEmbeddingFunction):
+    def __init__(
+        self,
+        api_url: str,
+        dimensions: Optional[int] = None,
+    ):
+        self.api_url = api_url + "/v1/embeddings"
+        self._session = requests.Session()
+        self._dim = dimensions
+
+    @property
+    def dim(self):
+        if self._dim is None:
+            self._dim = self._call_api(["get dim"])[0].shape[0]
+        return self._dim
+
+    def encode_queries(self, queries: List[str]) -> List[np.array]:
+        return self._call_api(queries)
+
+    def encode_documents(self, documents: List[str]) -> List[np.array]:
+        return self._call_api(documents)
+
+    def __call__(self, texts: List[str]) -> List[np.array]:
+        return self._call_api(texts)
+
+    def _call_api(self, texts: List[str]):
+        data = {"input": texts}
+        resp = self._session.post(  # type: ignore[assignment]
+            self.api_url,
+            json=data,
+        ).json()
+        if "data" not in resp:
+            raise RuntimeError(resp["message"])
+
+        embeddings = resp["data"]
+
+        # Sort resulting embeddings by index
+        sorted_embeddings = sorted(embeddings, key=lambda e: e["index"])  # type: ignore[no-any-return]
+        return [np.array(result["embedding"]) for result in sorted_embeddings]

src/pymilvus/model/reranker/__init__.py

@@ -3,11 +3,13 @@
 from pymilvus.model.reranker.voyageai import VoyageRerankFunction
 from pymilvus.model.reranker.cross_encoder import CrossEncoderRerankFunction
 from pymilvus.model.reranker.jinaai import JinaRerankFunction
+from pymilvus.model.reranker.opensource import OpenSourceRerankFunction
 
 __all__ = [
     "CohereRerankFunction",
     "BGERerankFunction",
     "VoyageRerankFunction",
     "CrossEncoderRerankFunction",
     "JinaRerankFunction",
+    "OpenSourceRerankFunction",
 ]

src/pymilvus/model/reranker/opensource.py

@@ -0,0 +1,30 @@
+from typing import List
+
+import requests
+
+from pymilvus.model.base import BaseRerankFunction, RerankResult
+
+
+class OpenSourceRerankFunction(BaseRerankFunction):
+    def __init__(self, api_url: str):
+        self.api_url = api_url + "/rerank"
+        self._session = requests.Session()
+
+    def __call__(self, query: str, documents: List[str], top_k: int = 5) -> List[RerankResult]:
+        resp = self._session.post(  # type: ignore[assignment]
+            self.api_url,
+            json={
+                "query": query,
+                "raw_scores": False,
+                "return_text": True,
+                "texts": documents,
+                "truncate": False,
+            },
+        ).json()
+        if "error" in resp:
+            raise RuntimeError(resp["error"])
+
+        results = []
+        for res in resp:
+            results.append(RerankResult(text=res["text"], score=res["score"], index=res["index"]))
+        return results