pymilvus.model 0.3.0__py3-none-any.whl

pymilvus/model/__init__.py

@@ -0,0 +1,5 @@
+__all__ = ["DefaultEmbeddingFunction", "dense", "sparse", "hybrid", "reranker", "utils"]
+
+from . import dense, hybrid, sparse, reranker, utils
+
+DefaultEmbeddingFunction = dense.onnx.OnnxEmbeddingFunction

pymilvus/model/base.py

@@ -0,0 +1,34 @@
+from abc import abstractmethod
+from typing import List
+
+
+class BaseEmbeddingFunction:
+    @abstractmethod
+    def __call__(self, texts: List[str]):
+        """ """
+
+    @abstractmethod
+    def encode_queries(self, queries: List[str]):
+        """ """
+
+
+class BaseRerankFunction:
+    @abstractmethod
+    def __call__(self, query: str, documents: List[str], top_k: int):
+        """ """
+
+
+class RerankResult:
+    def __init__(self, text: str, score: float, index: int):
+        self.text = text
+        self.score = score
+        self.index = index
+
+    def to_dict(self):
+        return {"text": self.text, "score": self.score, "index": self.index}
+
+    def __str__(self):
+        return f"RerankResult(text={self.text!r}, score={self.score}, index={self.index})"
+
+    def __repr__(self):
+        return f"RerankResult(text={self.text!r}, score={self.score}, index={self.index})"

pymilvus/model/dense/__init__.py

@@ -0,0 +1,21 @@
+from pymilvus.model.dense.openai import OpenAIEmbeddingFunction
+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.onnx import OnnxEmbeddingFunction
+from pymilvus.model.dense.cohere import CohereEmbeddingFunction
+from pymilvus.model.dense.mistralai import MistralAIEmbeddingFunction
+from pymilvus.model.dense.nomic import NomicEmbeddingFunction
+from pymilvus.model.dense.instructor import InstructorEmbeddingFunction
+
+__all__ = [
+    "OpenAIEmbeddingFunction",
+    "SentenceTransformerEmbeddingFunction",
+    "VoyageEmbeddingFunction",
+    "JinaEmbeddingFunction",
+    "OnnxEmbeddingFunction",
+    "CohereEmbeddingFunction",
+    "MistralAIEmbeddingFunction",
+    "NomicEmbeddingFunction",
+    "InstructorEmbeddingFunction"
+]

pymilvus/model/dense/cohere.py

@@ -0,0 +1,79 @@
+from typing import List, Optional
+import struct
+from collections import defaultdict
+import numpy as np
+
+from pymilvus.model.base import BaseEmbeddingFunction
+from pymilvus.model.utils import import_cohere
+
+
+
+class CohereEmbeddingFunction(BaseEmbeddingFunction):
+    def __init__(self,
+                 model_name: str = "embed-english-light-v3.0",
+                 api_key: Optional[str] = None,
+                 input_type: str = "search_document",
+                 embedding_types: Optional[List[str]] = None,
+                 truncate: Optional[str] = None,
+                 **kwargs):
+        self.model_name = model_name
+        self.input_type = input_type
+        self.embedding_types = embedding_types
+        self.truncate = truncate
+
+        import_cohere()
+        import cohere
+
+        if isinstance(embedding_types, list):
+            if len(embedding_types) > 1:
+                raise ValueError("Only one embedding type can be specified using current PyMilvus model library.")
+            elif embedding_types[0] == "int8" or embedding_types[0] == "uint8":
+                raise ValueError("Currently int8 or uint8 is not supported with PyMilvus model library.")
+            else:
+                pass
+    
+        self.client = cohere.Client(api_key, **kwargs)
+        self._cohereai_model_meta_info = defaultdict(dict)
+        self._cohereai_model_meta_info["embed-english-v3.0"]["dim"] = 1024
+        self._cohereai_model_meta_info["embed-english-light-v3.0"]["dim"] = 384
+        self._cohereai_model_meta_info["embed-english-v2.0"]["dim"] = 4096
+        self._cohereai_model_meta_info["embed-english-light-v2.0"]["dim"] = 1024
+        self._cohereai_model_meta_info["embed-multilingual-v3.0"]["dim"] = 1024
+        self._cohereai_model_meta_info["embed-multilingual-light-v3.0"]["dim"] = 384
+        self._cohereai_model_meta_info["embed-multilingual-v2.0"]["dim"] = 768
+
+    def _call_cohere_api(self, texts: List[str], input_type: str) -> List[np.array]:
+        embeddings = self.client.embed(
+            texts=texts,
+            model=self.model_name,
+            input_type=input_type,
+            embedding_types=self.embedding_types,
+            truncate=self.truncate
+        ).embeddings
+        if self.embedding_types is None:
+            results = [np.array(data, dtype=np.float32) for data in embeddings]
+        else:
+            results = getattr(embeddings, self.embedding_types[0], None)
+            if self.embedding_types[0] == "binary":
+                results = [struct.pack('b' * len(int8_vector), *int8_vector) for int8_vector in results] 
+            elif self.embedding_types[0] == "ubinary":
+                results = [struct.pack('B' * len(uint8_vector), *uint8_vector) for uint8_vector in results] 
+            elif self.embedding_types[0] == "float":
+                results = [np.array(result, dtype=np.float32) for result in results]
+            else:
+                pass
+        return results
+
+    def encode_documents(self, documents: List[str]) -> List[np.array]:
+        return self._call_cohere_api(documents, input_type="search_document")
+
+    def encode_queries(self, queries: List[str]) -> List[np.array]:
+        return self._call_cohere_api(queries, input_type="search_query")
+
+    def __call__(self, texts: List[str]) -> List[np.array]:
+        return self._call_cohere_api(texts, self.input_type)
+
+    @property
+    def dim(self):
+        return self._cohereai_model_meta_info[self.model_name]["dim"]
+

pymilvus/model/dense/instructor.py

@@ -0,0 +1,76 @@
+from typing import List, Optional
+import numpy as np
+
+from pymilvus.model.base import BaseEmbeddingFunction
+
+
+class InstructorEmbeddingFunction(BaseEmbeddingFunction):
+    def __init__(
+        self,
+        model_name: str = "hkunlp/instructor-xl",
+        batch_size: int = 32,
+        query_instruction: str = "Represent the question for retrieval:",
+        doc_instruction: str = "Represent the document for retrieval:",
+        device: str = "cpu",
+        normalize_embeddings: bool = True,
+        **kwargs,
+    ):
+        from .instructor_embedding.instructor_impl import Instructor
+
+        self.model_name = model_name
+        self.query_instruction = query_instruction
+        self.doc_instruction = doc_instruction
+        self.batch_size = batch_size
+        self.normalize_embeddings = normalize_embeddings
+
+        _model_config = dict({"model_name_or_path": model_name, "device": device}, **kwargs)
+        self.model = Instructor(**_model_config)
+
+    def __call__(self, texts: List[str]) -> List[np.array]:
+        return self._encode([[self.doc_instruction, text] for text in texts])
+
+    def _encode(self, texts: List[str]) -> List[np.array]:
+        embs = self.model.encode(
+            texts,
+            batch_size=self.batch_size,
+            show_progress_bar=False,
+            convert_to_numpy=True,
+            normalize_embeddings=self.normalize_embeddings
+        )
+        return list(embs)
+
+    @property
+    def dim(self):
+        return self.model.get_sentence_embedding_dimension()
+
+    def encode_queries(self, queries: List[str]) -> List[np.array]:
+        instructed_queries = [[self.query_instruction, query] for query in queries]
+        return self._encode(instructed_queries)
+
+    def encode_documents(self, documents: List[str]) -> List[np.array]:
+        instructed_documents = [[self.doc_instruction, document] for document in documents]
+        return self._encode(instructed_documents)
+
+    def _encode_query(self, query: str) -> np.array:
+        instructed_query = self.query_instruction + query
+        embs = self.model.encode(
+            sentences=[instructed_query],
+            batch_size=1,
+            show_progress_bar=False,
+            convert_to_numpy=True,
+            normalize_embeddings=self.normalize_embeddings,
+        )
+        return embs[0]
+
+    def _encode_document(self, document: str) -> np.array:
+        instructed_document = self.doc_instruction + document
+        embs = self.model.encode(
+            sentences=[instructed_document],
+            batch_size=1,
+            show_progress_bar=False,
+            convert_to_numpy=True,
+            normalize_embeddings=self.normalize_embeddings,
+        )
+        return embs[0]
+ 
+