MiniLM fix (#275)

* MiniLM fix

* Added MiniLM to text embedding
Fixed MiniLM source destination
Black + isort for repo

* Fixed model all-MiniLM-L6-v2 description
Recomputed canonical vector for all-MiniLM-L6-v2 in test

---------

Co-authored-by: d.rudenko <[email protected]>

Author: I8dNLo

experiments/attention_export.py

@@ -12,4 +12,6 @@
 #     print("Model already exported")
 # except FileNotFoundError:
 print(f"Exporting model to {output_dir}")
-main_export(model_id, output=output_dir, no_post_process=True, model_kwargs=model_kwargs)
+main_export(
+    model_id, output=output_dir, no_post_process=True, model_kwargs=model_kwargs
+)

experiments/try_attention_export.py

@@ -17,10 +17,14 @@
 attention_mask = tokenizer_output["attention_mask"]
 print(attention_mask)
 # Prepare the input
-input_ids = np.array(input_ids).astype(np.int64)  # Replace your_input_ids with actual input data
+input_ids = np.array(input_ids).astype(
+    np.int64
+)  # Replace your_input_ids with actual input data
 
 # Run the ONNX model
-outputs = ort_session.run(None, {"input_ids": input_ids, "attention_mask": attention_mask})
+outputs = ort_session.run(
+    None, {"input_ids": input_ids, "attention_mask": attention_mask}
+)
 
 # Get the attention weights
 attentions = outputs[-1]

fastembed/__init__.py

@@ -1,9 +1,9 @@
 import importlib.metadata
 
 from fastembed.image import ImageEmbedding
-from fastembed.text import TextEmbedding
-from fastembed.sparse import SparseTextEmbedding, SparseEmbedding
 from fastembed.late_interaction import LateInteractionTextEmbedding
+from fastembed.sparse import SparseEmbedding, SparseTextEmbedding
+from fastembed.text import TextEmbedding
 
 try:
     version = importlib.metadata.version("fastembed")

fastembed/common/__init__.py

@@ -1,3 +1,3 @@
-from fastembed.common.types import OnnxProvider, ImageInput, PathInput
+from fastembed.common.types import ImageInput, OnnxProvider, PathInput
 
 __all__ = ["OnnxProvider", "ImageInput", "PathInput"]

fastembed/common/model_management.py

@@ -2,13 +2,13 @@
 import shutil
 import tarfile
 from pathlib import Path
-from typing import List, Optional, Dict, Any
+from typing import Any, Dict, List, Optional
 
 import requests
 from huggingface_hub import snapshot_download
 from huggingface_hub.utils import RepositoryNotFoundError
-from tqdm import tqdm
 from loguru import logger
+from tqdm import tqdm
 
 
 class ModelManagement:
@@ -42,7 +42,9 @@ def _get_model_description(cls, model_name: str) -> Dict[str, Any]:
         raise ValueError(f"Model {model_name} is not supported in {cls.__name__}.")
 
     @classmethod
-    def download_file_from_gcs(cls, url: str, output_path: str, show_progress: bool = True) -> str:
+    def download_file_from_gcs(
+        cls, url: str, output_path: str, show_progress: bool = True
+    ) -> str:
         """
         Downloads a file from Google Cloud Storage.
 
@@ -71,12 +73,17 @@ def download_file_from_gcs(cls, url: str, output_path: str, show_progress: bool
 
         # Warn if the total size is zero
         if total_size_in_bytes == 0:
-            print(f"Warning: Content-length header is missing or zero in the response from {url}.")
+            print(
+                f"Warning: Content-length header is missing or zero in the response from {url}."
+            )
 
         show_progress = total_size_in_bytes and show_progress
 
         with tqdm(
-            total=total_size_in_bytes, unit="iB", unit_scale=True, disable=not show_progress
+            total=total_size_in_bytes,
+            unit="iB",
+            unit_scale=True,
+            disable=not show_progress,
         ) as progress_bar:
             with open(output_path, "wb") as file:
                 for chunk in response.iter_content(chunk_size=1024):
@@ -156,7 +163,9 @@ def decompress_to_cache(cls, targz_path: str, cache_dir: str):
         return cache_dir
 
     @classmethod
-    def retrieve_model_gcs(cls, model_name: str, source_url: str, cache_dir: str) -> Path:
+    def retrieve_model_gcs(
+        cls, model_name: str, source_url: str, cache_dir: str
+    ) -> Path:
         fast_model_name = f"fast-{model_name.split('/')[-1]}"
 
         cache_tmp_dir = Path(cache_dir) / "tmp"
@@ -182,8 +191,12 @@ def retrieve_model_gcs(cls, model_name: str, source_url: str, cache_dir: str) ->
             output_path=str(model_tar_gz),
         )
 
-        cls.decompress_to_cache(targz_path=str(model_tar_gz), cache_dir=str(cache_tmp_dir))
-        assert model_tmp_dir.exists(), f"Could not find {model_tmp_dir} in {cache_tmp_dir}"
+        cls.decompress_to_cache(
+            targz_path=str(model_tar_gz), cache_dir=str(cache_tmp_dir)
+        )
+        assert (
+            model_tmp_dir.exists()
+        ), f"Could not find {model_tmp_dir} in {cache_tmp_dir}"
 
         model_tar_gz.unlink()
         # Rename from tmp to final name is atomic

fastembed/common/onnx_model.py

@@ -1,15 +1,24 @@
+import warnings
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Any, Dict, Generic, Iterable, Optional, Tuple, Type, TypeVar, Sequence
-import warnings
+from typing import (
+    Any,
+    Dict,
+    Generic,
+    Iterable,
+    Optional,
+    Sequence,
+    Tuple,
+    Type,
+    TypeVar,
+)
 
 import numpy as np
 import onnxruntime as ort
 
 from fastembed.common.types import OnnxProvider
 from fastembed.parallel_processor import Worker
 
-
 # Holds type of the embedding result
 T = TypeVar("T")
 
@@ -51,7 +60,9 @@ def load_onnx_model(
         model_path = model_dir / model_file
         # List of Execution Providers: https://onnxruntime.ai/docs/execution-providers
 
-        onnx_providers = ["CPUExecutionProvider"] if providers is None else list(providers)
+        onnx_providers = (
+            ["CPUExecutionProvider"] if providers is None else list(providers)
+        )
         available_providers = ort.get_available_providers()
         requested_provider_names = []
         for provider in onnx_providers:

fastembed/common/preprocessor_utils.py

@@ -2,7 +2,7 @@
 from pathlib import Path
 from typing import Tuple
 
-from tokenizers import Tokenizer, AddedToken
+from tokenizers import AddedToken, Tokenizer
 
 from fastembed.image.transform.operators import Compose
 
@@ -40,7 +40,9 @@ def load_tokenizer(model_dir: Path, max_length: int = 512) -> Tuple[Tokenizer, d
     tokens_map = load_special_tokens(model_dir)
 
     tokenizer = Tokenizer.from_file(str(tokenizer_path))
-    tokenizer.enable_truncation(max_length=min(tokenizer_config["model_max_length"], max_length))
+    tokenizer.enable_truncation(
+        max_length=min(tokenizer_config["model_max_length"], max_length)
+    )
     tokenizer.enable_padding(
         pad_id=config.get("pad_token_id", 0), pad_token=tokenizer_config["pad_token"]
     )

fastembed/common/types.py

@@ -1,6 +1,6 @@
 import os
 import sys
-from typing import Union, Iterable, Tuple, Dict, Any
+from typing import Any, Dict, Iterable, Tuple, Union
 
 if sys.version_info >= (3, 10):
     from typing import TypeAlias

fastembed/common/utils.py

@@ -2,7 +2,7 @@
 import tempfile
 from itertools import islice
 from pathlib import Path
-from typing import Union, Iterable, Generator, Optional
+from typing import Generator, Iterable, Optional, Union
 
 import numpy as np
 

fastembed/image/__init__.py

@@ -1,4 +1,3 @@
 from fastembed.image.image_embedding import ImageEmbedding
 
-
 __all__ = ["ImageEmbedding"]

fastembed/image/image_embedding.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Iterable, List, Optional, Type, Sequence
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Type
 
 import numpy as np
 
@@ -51,9 +51,16 @@ def __init__(
 
         for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
             supported_models = EMBEDDING_MODEL_TYPE.list_supported_models()
-            if any(model_name.lower() == model["model"].lower() for model in supported_models):
+            if any(
+                model_name.lower() == model["model"].lower()
+                for model in supported_models
+            ):
                 self.model = EMBEDDING_MODEL_TYPE(
-                    model_name, cache_dir, threads=threads, providers=providers, **kwargs
+                    model_name,
+                    cache_dir,
+                    threads=threads,
+                    providers=providers,
+                    **kwargs,
                 )
                 return
 

fastembed/image/onnx_embedding.py

@@ -1,12 +1,12 @@
-from typing import Dict, Optional, Iterable, Type, List, Any, Sequence
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Type
 
 import numpy as np
 
-from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.utils import normalize, define_cache_dir
 from fastembed.common import ImageInput, OnnxProvider
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.utils import define_cache_dir, normalize
 from fastembed.image.image_embedding_base import ImageEmbeddingBase
-from fastembed.image.onnx_image_model import OnnxImageModel, ImageEmbeddingWorker
+from fastembed.image.onnx_image_model import ImageEmbeddingWorker, OnnxImageModel
 
 supported_onnx_models = [
     {
@@ -122,10 +122,16 @@ def _preprocess_onnx_input(
 
         return onnx_input
 
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext
+    ) -> Iterable[np.ndarray]:
         return normalize(output.model_output).astype(np.float32)
 
 
 class OnnxImageEmbeddingWorker(ImageEmbeddingWorker):
-    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> OnnxImageEmbedding:
-        return OnnxImageEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs)
+    def init_embedding(
+        self, model_name: str, cache_dir: str, **kwargs
+    ) -> OnnxImageEmbedding:
+        return OnnxImageEmbedding(
+            model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs
+        )

fastembed/image/onnx_image_model.py

@@ -1,15 +1,15 @@
-import os
 import contextlib
+import os
 from multiprocessing import get_all_start_methods
 from pathlib import Path
-from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, Sequence
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Type
 
-from PIL import Image
 import numpy as np
+from PIL import Image
 
+from fastembed.common import ImageInput, OnnxProvider, PathInput
+from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_preprocessor
-from fastembed.common.onnx_model import OnnxModel, EmbeddingWorker, T, OnnxOutputContext
-from fastembed.common import PathInput, ImageInput, OnnxProvider
 from fastembed.common.utils import iter_batch
 from fastembed.parallel_processor import ParallelWorkerPool
 
@@ -44,7 +44,10 @@ def load_onnx_model(
         providers: Optional[Sequence[OnnxProvider]] = None,
     ) -> None:
         super().load_onnx_model(
-            model_dir=model_dir, model_file=model_file, threads=threads, providers=providers
+            model_dir=model_dir,
+            model_file=model_file,
+            threads=threads,
+            providers=providers,
         )
         self.processor = load_preprocessor(model_dir=model_dir)
 
@@ -87,7 +90,9 @@ def _embed_images(
             for batch in iter_batch(images, batch_size):
                 yield from self._post_process_onnx_output(self.onnx_embed(batch))
         else:
-            start_method = "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
+            start_method = (
+                "forkserver" if "forkserver" in get_all_start_methods() else "spawn"
+            )
             params = {"model_name": model_name, "cache_dir": cache_dir, **kwargs}
             pool = ParallelWorkerPool(
                 parallel, self._get_worker_class(), start_method=start_method

fastembed/image/transform/functional.py

@@ -1,8 +1,7 @@
-from typing import Union, Tuple, Sized
-
-from PIL import Image
+from typing import Sized, Tuple, Union
 
 import numpy as np
+from PIL import Image
 
 
 def convert_to_rgb(image: Image.Image) -> Image.Image: