fix: Fix normalize embedding, Mistral's client API, changes BM25 defa… (

#31)

…ult worker num.

Signed-off-by: wxywb <[email protected]>

milvus_model/dense/instructor.py

@@ -1,6 +1,4 @@
 from typing import List, Optional
-import struct
-from collections import defaultdict
 import numpy as np
 
 from milvus_model.base import BaseEmbeddingFunction
@@ -36,7 +34,11 @@ def __call__(self, texts: List[str]) -> List[np.array]:
 
     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,
+            texts,
+            batch_size=self.batch_size,
+            show_progress_bar=False,
+            convert_to_numpy=True,
+            normalize_embeddings=self.normalize_embeddings
         )
         return list(embs)
 

milvus_model/dense/mistralai.py

@@ -32,7 +32,7 @@ def __init__(
         else:
             self.api_key = api_key
         self.model_name = model_name
-        self.client = MistralClient(api_key=api_key)
+        self.client = Mistral(api_key=api_key)
         self._encode_config = {"model": model_name, **kwargs}
 
     def encode_queries(self, queries: List[str]) -> List[np.array]:

milvus_model/dense/sentence_transformer.py

@@ -33,7 +33,11 @@ def __call__(self, texts: List[str]) -> List[np.array]:
 
     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,
+            texts, 
+            batch_size=self.batch_size,
+            show_progress_bar=False,
+            convert_to_numpy=True,
+            normalize_embeddings=self.normalize_embeddings
         )
         return list(embs)
 

milvus_model/sparse/bm25/bm25.py

@@ -18,7 +18,7 @@
 import logging
 import math
 from collections import defaultdict
-from multiprocessing import Pool, cpu_count
+from multiprocessing import Pool
 from pathlib import Path
 from typing import Dict, List, Optional
 
@@ -44,7 +44,7 @@ def __init__(
         k1: float = 1.5,
         b: float = 0.75,
         epsilon: float = 0.25,
-        num_workers: Optional[int] = None,
+        num_workers: int = 1,
     ):
         if analyzer is None:
             analyzer = build_default_analyzer(language="en")
@@ -55,8 +55,6 @@ def __init__(
         self.k1 = k1
         self.b = b
         self.epsilon = epsilon
-        if num_workers is None:
-            self.num_workers = cpu_count()
         self.num_workers = num_workers
 
         if analyzer and corpus is not None:

milvus_model/utils/__init__.py

@@ -31,6 +31,7 @@ def import_sentence_transformers():
     _check_library("sentence_transformers", package="sentence-transformers")
 
 def import_FlagEmbedding():
+    _check_library("peft", package="peft")
     _check_library("FlagEmbedding", package="FlagEmbedding>=1.2.2")
 
 def import_nltk():