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"]