fix: Update voyage, mistral, nomic embedding functions. (#41)

Signed-off-by: wxywb <[email protected]>
milvus-io · Sep 20, 2024 · 1f1d4e1 · 1f1d4e1
1 parent 27259a0
commit 1f1d4e1
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 19 deletions.
diff --git a/milvus_model/dense/mistralai.py b/milvus_model/dense/mistralai.py
@@ -56,8 +56,8 @@ def _encode_document(self, document: str) -> np.array:
 
     def _call_mistral_api(self, texts: List[str]):
         embeddings_batch_response = self.client.embeddings.create(
-            model=self.model_name,
             inputs=texts,
+            **self._encode_config
         )
         return [np.array(data.embedding) for data in embeddings_batch_response.data]
 

diff --git a/milvus_model/dense/nomic.py b/milvus_model/dense/nomic.py
@@ -12,7 +12,6 @@
 class NomicEmbeddingFunction(BaseEmbeddingFunction):
     def __init__(
         self,
-        api_key: str,
         model_name: str = "nomic-embed-text-v1.5",
         task_type: str = "search_document",
         dimensions: int = 768,
@@ -21,18 +20,6 @@ def __init__(
         self._nomic_model_meta_info = defaultdict(dict)
         self._nomic_model_meta_info[model_name]["dim"] = dimensions  # set the dimension
 
-        if api_key is None:
-            if "NOMIC_API_KEY" in os.environ and os.environ["NOMIC_API_KEY"]:
-                self.api_key = os.environ["NOMIC_API_KEY"]
-            else:
-                error_message = (
-                    "Did not find api_key, please add an environment variable"
-                    " `NOMIC_API_KEY` which contains it, or pass"
-                    "  `api_key` as a named parameter."
-                )
-                raise ValueError(error_message)
-        else:
-            self.api_key = api_key
         self.model_name = model_name
         self.task_type = task_type
         self.dimensionality = dimensions
@@ -67,12 +54,9 @@ def _encode_document(self, document: str) -> np.array:
         return self._encode([document], task_type="search_document")[0]
 
     def _call_nomic_api(self, texts: List[str], task_type: str):
-        headers = {"Authorization": f"Bearer {self.api_key}"}
         embeddings_batch_response = embed.text(
             texts=texts,
-            model=self.model_name,
-            task_type=task_type,
-            dimensionality=self.dimensionality,
+            **self._encode_config
         )
         return [np.array(embedding) for embedding in embeddings_batch_response["embeddings"]]
 

diff --git a/milvus_model/dense/voyageai.py b/milvus_model/dense/voyageai.py
@@ -11,9 +11,16 @@
 
 
 class VoyageEmbeddingFunction(BaseEmbeddingFunction):
-    def __init__(self, model_name: str = "voyage-2", api_key: Optional[str] = None, **kwargs):
+    def __init__(self, model_name: str = "voyage-3", api_key: Optional[str] = None, **kwargs):
         self.model_name = model_name
         self._voyageai_model_meta_info = defaultdict(dict)
+        self._voyageai_model_meta_info["voyage-3"]["dim"] = 1024
+        self._voyageai_model_meta_info["voyage-3-lite"]["dim"] = 512
+        self._voyageai_model_meta_info["voyage-finance-2"]["dim"] = 1024
+        self._voyageai_model_meta_info["voyage-multilingual-2"]["dim"] = 1024
+        self._voyageai_model_meta_info["voyage-law-2"]["dim"] = 1024
+        self._voyageai_model_meta_info["voyage-code-2"]["dim"] = 1536
+        #old model
         self._voyageai_model_meta_info["voyage-large-2"]["dim"] = 1536
         self._voyageai_model_meta_info["voyage-code-2"]["dim"] = 1536
         self._voyageai_model_meta_info["voyage-2"]["dim"] = 1024