From 331207976e771e43fd92531c2236deb0e1a51ad0 Mon Sep 17 00:00:00 2001 From: Dmitrii Ogn Date: Fri, 14 Jun 2024 16:42:31 +0300 Subject: [PATCH] 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 --- experiments/attention_export.py | 4 +- experiments/try_attention_export.py | 8 +- fastembed/__init__.py | 4 +- fastembed/common/__init__.py | 2 +- fastembed/common/model_management.py | 29 +++-- fastembed/common/onnx_model.py | 19 +++- fastembed/common/preprocessor_utils.py | 6 +- fastembed/common/types.py | 2 +- fastembed/common/utils.py | 2 +- fastembed/image/__init__.py | 1 - fastembed/image/image_embedding.py | 13 ++- fastembed/image/onnx_embedding.py | 20 ++-- fastembed/image/onnx_image_model.py | 19 ++-- fastembed/image/transform/functional.py | 5 +- fastembed/image/transform/operators.py | 106 +++++++++++------- fastembed/late_interaction/__init__.py | 5 +- fastembed/late_interaction/colbert.py | 11 +- .../late_interaction_embedding_base.py | 4 +- .../late_interaction_text_embedding.py | 13 ++- fastembed/parallel_processor.py | 8 +- fastembed/sparse/bm25.py | 19 +++- fastembed/sparse/bm42.py | 33 ++++-- fastembed/sparse/sparse_embedding_base.py | 28 +++-- fastembed/sparse/sparse_text_embedding.py | 22 +++- fastembed/sparse/splade_pp.py | 12 +- fastembed/sparse/utils/tokenizer.py | 3 +- fastembed/text/clip_embedding.py | 14 ++- fastembed/text/e5_onnx_embedding.py | 10 +- fastembed/text/jina_onnx_embedding.py | 14 ++- fastembed/text/mini_lm_embedding.py | 69 ++++++++++++ fastembed/text/onnx_embedding.py | 26 ++--- fastembed/text/onnx_text_model.py | 13 ++- fastembed/text/text_embedding.py | 15 ++- fastembed/text/text_embedding_base.py | 4 +- tests/profiling.py | 16 ++- tests/test_attention_embeddings.py | 11 +- tests/test_late_interaction_embeddings.py | 13 ++- tests/test_sparse_embeddings.py | 8 +- tests/test_text_onnx_embeddings.py | 40 +++++-- 39 files changed, 460 insertions(+), 191 deletions(-) create mode 100644 fastembed/text/mini_lm_embedding.py diff --git a/experiments/attention_export.py b/experiments/attention_export.py index bd6e7d8..4507904 100644 --- a/experiments/attention_export.py +++ b/experiments/attention_export.py @@ -12,4 +12,6 @@ tokenizer = AutoTokenizer.from_pretrained(model_id) # 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 +) diff --git a/experiments/try_attention_export.py b/experiments/try_attention_export.py index 6f9ae55..e0301aa 100644 --- a/experiments/try_attention_export.py +++ b/experiments/try_attention_export.py @@ -17,10 +17,14 @@ input_ids = tokenizer_output["input_ids"] 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] diff --git a/fastembed/__init__.py b/fastembed/__init__.py index 03427ee..906cc1f 100644 --- a/fastembed/__init__.py +++ b/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") diff --git a/fastembed/common/__init__.py b/fastembed/common/__init__.py index 710fd37..5caf45f 100644 --- a/fastembed/common/__init__.py +++ b/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"] diff --git a/fastembed/common/model_management.py b/fastembed/common/model_management.py index d0c35ad..6a45061 100644 --- a/fastembed/common/model_management.py +++ b/fastembed/common/model_management.py @@ -2,13 +2,13 @@ import os 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 @@ class ModelManagement: 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 @@ class ModelManagement: # 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 @@ class ModelManagement: 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 @@ class ModelManagement: 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 diff --git a/fastembed/common/onnx_model.py b/fastembed/common/onnx_model.py index c6a0df5..7820016 100644 --- a/fastembed/common/onnx_model.py +++ b/fastembed/common/onnx_model.py @@ -1,7 +1,17 @@ +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 @@ -9,7 +19,6 @@ 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 @@ class OnnxModel(Generic[T]): 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: diff --git a/fastembed/common/preprocessor_utils.py b/fastembed/common/preprocessor_utils.py index 0f81a84..1b73404 100644 --- a/fastembed/common/preprocessor_utils.py +++ b/fastembed/common/preprocessor_utils.py @@ -2,7 +2,7 @@ import json 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"] ) diff --git a/fastembed/common/types.py b/fastembed/common/types.py index f658654..c428a3f 100644 --- a/fastembed/common/types.py +++ b/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 diff --git a/fastembed/common/utils.py b/fastembed/common/utils.py index 12a0f3c..ab6a858 100644 --- a/fastembed/common/utils.py +++ b/fastembed/common/utils.py @@ -2,7 +2,7 @@ import os 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 diff --git a/fastembed/image/__init__.py b/fastembed/image/__init__.py index ea9ca44..cc4f140 100644 --- a/fastembed/image/__init__.py +++ b/fastembed/image/__init__.py @@ -1,4 +1,3 @@ from fastembed.image.image_embedding import ImageEmbedding - __all__ = ["ImageEmbedding"] diff --git a/fastembed/image/image_embedding.py b/fastembed/image/image_embedding.py index b74f597..29df2e1 100644 --- a/fastembed/image/image_embedding.py +++ b/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 @@ class ImageEmbedding(ImageEmbeddingBase): 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 diff --git a/fastembed/image/onnx_embedding.py b/fastembed/image/onnx_embedding.py index 2cfe8ba..3437021 100644 --- a/fastembed/image/onnx_embedding.py +++ b/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 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]): 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 + ) diff --git a/fastembed/image/onnx_image_model.py b/fastembed/image/onnx_image_model.py index 7c8caf4..cfaa152 100644 --- a/fastembed/image/onnx_image_model.py +++ b/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 @@ class OnnxImageModel(OnnxModel[T]): 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 @@ class OnnxImageModel(OnnxModel[T]): 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 diff --git a/fastembed/image/transform/functional.py b/fastembed/image/transform/functional.py index 4e02473..0c8b054 100644 --- a/fastembed/image/transform/functional.py +++ b/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: diff --git a/fastembed/image/transform/operators.py b/fastembed/image/transform/operators.py index 2b3c8b2..854b091 100644 --- a/fastembed/image/transform/operators.py +++ b/fastembed/image/transform/operators.py @@ -1,14 +1,15 @@ -from typing import List, Tuple, Union, Any, Dict +from typing import Any, Dict, List, Tuple, Union import numpy as np from PIL import Image + from fastembed.image.transform.functional import ( center_crop, - normalize, - resize, convert_to_rgb, + normalize, + pil2ndarray, rescale, - pil2ndarray + resize, ) @@ -49,7 +50,9 @@ class Resize(Transform): self.resample = resample def __call__(self, images: List[Image.Image]) -> List[Image.Image]: - return [resize(image, size=self.size, resample=self.resample) for image in images] + return [ + resize(image, size=self.size, resample=self.resample) for image in images + ] class Rescale(Transform): @@ -59,15 +62,21 @@ class Rescale(Transform): def __call__(self, images: List[np.ndarray]) -> List[np.ndarray]: return [rescale(image, scale=self.scale) for image in images] + class PILtoNDarray(Transform): - def __call__(self, images: List[Union[Image.Image, np.ndarray]]) -> List[np.ndarray]: + def __call__( + self, images: List[Union[Image.Image, np.ndarray]] + ) -> List[np.ndarray]: return [pil2ndarray(image) for image in images] + class Compose: def __init__(self, transforms: List[Transform]): self.transforms = transforms - def __call__(self, images: Union[List[Image.Image], List[np.ndarray]]) -> Union[List[np.ndarray], List[Image.Image]]: + def __call__( + self, images: Union[List[Image.Image], List[np.ndarray]] + ) -> Union[List[np.ndarray], List[Image.Image]]: for transform in self.transforms: images = transform(images) return images @@ -75,25 +84,25 @@ class Compose: @classmethod def from_config(cls, config: Dict[str, Any]) -> "Compose": """Creates processor from a config dict. - Args: - config (Dict[str, Any]): Configuration dictionary. + Args: + config (Dict[str, Any]): Configuration dictionary. - Valid keys: - - do_resize - - size - - do_center_crop - - crop_size - - do_rescale - - rescale_factor - - do_normalize - - image_mean - - image_std - Valid size keys (nested): - - {"height", "width"} - - {"shortest_edge"} + Valid keys: + - do_resize + - size + - do_center_crop + - crop_size + - do_rescale + - rescale_factor + - do_normalize + - image_mean + - image_std + Valid size keys (nested): + - {"height", "width"} + - {"shortest_edge"} - Returns: - Compose: Image processor. + Returns: + Compose: Image processor. """ transforms = [] cls._get_convert_to_rgb(transforms, config) @@ -110,8 +119,8 @@ class Compose: @staticmethod def _get_resize(transforms: List[Transform], config: Dict[str, Any]): - mode = config.get('image_processor_type', 'CLIPImageProcessor') - if mode == 'CLIPImageProcessor': + mode = config.get("image_processor_type", "CLIPImageProcessor") + if mode == "CLIPImageProcessor": if config.get("do_resize", False): size = config["size"] if "shortest_edge" in size: @@ -119,25 +128,44 @@ class Compose: elif "height" in size and "width" in size: size = (size["height"], size["width"]) else: - raise ValueError("Size must contain either 'shortest_edge' or 'height' and 'width'.") - transforms.append(Resize(size=size, resample=config.get("resample", Image.Resampling.BICUBIC))) - elif mode == 'ConvNextFeatureExtractor': - if 'size' in config and "shortest_edge" not in config['size']: - raise ValueError(f"Size dictionary must contain 'shortest_edge' key. Got {config['size'].keys()}") - shortest_edge = config['size']["shortest_edge"] + raise ValueError( + "Size must contain either 'shortest_edge' or 'height' and 'width'." + ) + transforms.append( + Resize( + size=size, + resample=config.get("resample", Image.Resampling.BICUBIC), + ) + ) + elif mode == "ConvNextFeatureExtractor": + if "size" in config and "shortest_edge" not in config["size"]: + raise ValueError( + f"Size dictionary must contain 'shortest_edge' key. Got {config['size'].keys()}" + ) + shortest_edge = config["size"]["shortest_edge"] crop_pct = config.get("crop_pct", 0.875) if shortest_edge < 384: # maintain same ratio, resizing shortest edge to shortest_edge/crop_pct resize_shortest_edge = int(shortest_edge / crop_pct) - transforms.append(Resize(size=resize_shortest_edge, resample=config.get("resample", Image.Resampling.BICUBIC))) + transforms.append( + Resize( + size=resize_shortest_edge, + resample=config.get("resample", Image.Resampling.BICUBIC), + ) + ) transforms.append(CenterCrop(size=(shortest_edge, shortest_edge))) else: - transforms.append(Resize(size=(shortest_edge, shortest_edge), resample=config.get("resample", Image.Resampling.BICUBIC))) + transforms.append( + Resize( + size=(shortest_edge, shortest_edge), + resample=config.get("resample", Image.Resampling.BICUBIC), + ) + ) @staticmethod def _get_center_crop(transforms: List[Transform], config: Dict[str, Any]): - mode = config.get('image_processor_type', 'CLIPImageProcessor') - if mode == 'CLIPImageProcessor': + mode = config.get("image_processor_type", "CLIPImageProcessor") + if mode == "CLIPImageProcessor": if config.get("do_center_crop", False): crop_size = config["crop_size"] if isinstance(crop_size, int): @@ -147,7 +175,7 @@ class Compose: else: raise ValueError(f"Invalid crop size: {crop_size}") transforms.append(CenterCrop(size=crop_size)) - elif mode == 'ConvNextFeatureExtractor': + elif mode == "ConvNextFeatureExtractor": pass else: raise ValueError(f"Preprocessor {mode} is not supported") @@ -165,4 +193,6 @@ class Compose: @staticmethod def _get_normalize(transforms: List[Transform], config: Dict[str, Any]): if config.get("do_normalize", False): - transforms.append(Normalize(mean=config["image_mean"], std=config["image_std"])) + transforms.append( + Normalize(mean=config["image_mean"], std=config["image_std"]) + ) diff --git a/fastembed/late_interaction/__init__.py b/fastembed/late_interaction/__init__.py index 869dcd8..eae5d03 100644 --- a/fastembed/late_interaction/__init__.py +++ b/fastembed/late_interaction/__init__.py @@ -1,4 +1,5 @@ -from fastembed.late_interaction.late_interaction_text_embedding import LateInteractionTextEmbedding - +from fastembed.late_interaction.late_interaction_text_embedding import ( + LateInteractionTextEmbedding, +) __all__ = ["LateInteractionTextEmbedding"] diff --git a/fastembed/late_interaction/colbert.py b/fastembed/late_interaction/colbert.py index 11a7328..7cd5115 100644 --- a/fastembed/late_interaction/colbert.py +++ b/fastembed/late_interaction/colbert.py @@ -1,5 +1,5 @@ -from typing import Any, Dict, Iterable, List, Optional, Union, Type, Sequence import string +from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union import numpy as np from tokenizers import Encoding @@ -12,7 +12,6 @@ from fastembed.late_interaction.late_interaction_embedding_base import ( ) from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker - supported_colbert_models = [ { "model": "colbert-ir/colbertv2.0", @@ -44,7 +43,9 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]): if token_id in self.skip_list or token_id == self.pad_token_id: output.attention_mask[i, j] = 0 - output.model_output *= np.expand_dims(output.attention_mask, 2).astype(np.float32) + output.model_output *= np.expand_dims(output.attention_mask, 2).astype( + np.float32 + ) norm = np.linalg.norm(output.model_output, ord=2, axis=2, keepdims=True) norm_clamped = np.maximum(norm, 1e-12) output.model_output /= norm_clamped @@ -76,7 +77,9 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]): if self.tokenizer.padding: prev_padding = self.tokenizer.padding self.tokenizer.enable_padding( - pad_token=self.MASK_TOKEN, pad_id=self.mask_token_id, length=self.MIN_QUERY_LENGTH + pad_token=self.MASK_TOKEN, + pad_id=self.mask_token_id, + length=self.MIN_QUERY_LENGTH, ) encoded = self.tokenizer.encode_batch(query) if prev_padding is None: diff --git a/fastembed/late_interaction/late_interaction_embedding_base.py b/fastembed/late_interaction/late_interaction_embedding_base.py index 2a587e0..64fba49 100644 --- a/fastembed/late_interaction/late_interaction_embedding_base.py +++ b/fastembed/late_interaction/late_interaction_embedding_base.py @@ -42,7 +42,9 @@ class LateInteractionTextEmbeddingBase(ModelManagement): # This is model-specific, so that different models can have specialized implementations yield from self.embed(texts, **kwargs) - def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]: + def query_embed( + self, query: Union[str, Iterable[str]], **kwargs + ) -> Iterable[np.ndarray]: """ Embeds queries diff --git a/fastembed/late_interaction/late_interaction_text_embedding.py b/fastembed/late_interaction/late_interaction_text_embedding.py index 82ed975..6e84fec 100644 --- a/fastembed/late_interaction/late_interaction_text_embedding.py +++ b/fastembed/late_interaction/late_interaction_text_embedding.py @@ -1,12 +1,12 @@ -from typing import List, Type, Dict, Any, Union, Iterable, Optional, Sequence +from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union import numpy as np from fastembed.common import OnnxProvider +from fastembed.late_interaction.colbert import Colbert from fastembed.late_interaction.late_interaction_embedding_base import ( LateInteractionTextEmbeddingBase, ) -from fastembed.late_interaction.colbert import Colbert class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase): @@ -54,7 +54,10 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase): 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, providers=providers, **kwargs ) @@ -89,7 +92,9 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase): """ yield from self.model.embed(documents, batch_size, parallel, **kwargs) - def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]: + def query_embed( + self, query: Union[str, Iterable[str]], **kwargs + ) -> Iterable[np.ndarray]: """ Embeds queries diff --git a/fastembed/parallel_processor.py b/fastembed/parallel_processor.py index 03d20d5..1e3537b 100644 --- a/fastembed/parallel_processor.py +++ b/fastembed/parallel_processor.py @@ -83,7 +83,9 @@ def _worker( class ParallelWorkerPool: - def __init__(self, num_workers: int, worker: Type[Worker], start_method: Optional[str] = None): + def __init__( + self, num_workers: int, worker: Type[Worker], start_method: Optional[str] = None + ): self.worker_class = worker self.num_workers = num_workers self.input_queue: Optional[Queue] = None @@ -118,7 +120,9 @@ class ParallelWorkerPool: process.start() self.processes.append(process) - def ordered_map(self, stream: Iterable[Any], *args: Any, **kwargs: Any) -> Iterable[Any]: + def ordered_map( + self, stream: Iterable[Any], *args: Any, **kwargs: Any + ) -> Iterable[Any]: buffer = defaultdict(Any) next_expected = 0 diff --git a/fastembed/sparse/bm25.py b/fastembed/sparse/bm25.py index 0c70305..1b79712 100644 --- a/fastembed/sparse/bm25.py +++ b/fastembed/sparse/bm25.py @@ -3,7 +3,7 @@ import string from collections import defaultdict from multiprocessing import get_all_start_methods from pathlib import Path -from typing import Any, Dict, Iterable, List, Optional, Union, Type, Tuple +from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, Union import mmh3 import numpy as np @@ -11,10 +11,12 @@ from snowballstemmer import stemmer as get_stemmer from fastembed.common.utils import define_cache_dir, iter_batch from fastembed.parallel_processor import ParallelWorkerPool, Worker -from fastembed.sparse.sparse_embedding_base import SparseEmbedding, SparseTextEmbeddingBase +from fastembed.sparse.sparse_embedding_base import ( + SparseEmbedding, + SparseTextEmbeddingBase, +) from fastembed.sparse.utils.tokenizer import WordTokenizer - supported_bm25_models = [ { "model": "Qdrant/bm25", @@ -131,7 +133,9 @@ class Bm25(SparseTextEmbeddingBase): for batch in iter_batch(documents, batch_size): yield from self.raw_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, @@ -237,7 +241,9 @@ class Bm25(SparseTextEmbeddingBase): def compute_token_id(cls, token: str) -> int: return abs(mmh3.hash(token)) - def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[SparseEmbedding]: + def query_embed( + self, query: Union[str, Iterable[str]], **kwargs + ) -> Iterable[SparseEmbedding]: """To emulate BM25 behaviour, we don't need to use weights in the query, and it's enough to just hash the tokens and assign a weight of 1.0 to them. """ @@ -248,7 +254,8 @@ class Bm25(SparseTextEmbeddingBase): tokens = self.tokenizer.tokenize(text) stemmed_tokens = self._stem(tokens) token_ids = np.array( - [self.compute_token_id(token) for token in stemmed_tokens], dtype=np.float32 + [self.compute_token_id(token) for token in stemmed_tokens], + dtype=np.float32, ) values = np.ones_like(token_ids) yield SparseEmbedding(indices=token_ids, values=values) diff --git a/fastembed/sparse/bm42.py b/fastembed/sparse/bm42.py index 10bb39d..c043e7b 100644 --- a/fastembed/sparse/bm42.py +++ b/fastembed/sparse/bm42.py @@ -1,16 +1,19 @@ import math import string from pathlib import Path -from typing import Any, Dict, Iterable, List, Optional, Tuple, Union, Type, Sequence +from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Type, Union -import numpy as np import mmh3 +import numpy as np from snowballstemmer import stemmer as get_stemmer from fastembed.common import OnnxProvider from fastembed.common.onnx_model import OnnxOutputContext from fastembed.common.utils import define_cache_dir -from fastembed.sparse.sparse_embedding_base import SparseEmbedding, SparseTextEmbeddingBase +from fastembed.sparse.sparse_embedding_base import ( + SparseEmbedding, + SparseTextEmbeddingBase, +) from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker supported_bm42_models = [ @@ -103,7 +106,9 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]): self.stemmer = get_stemmer(MODEL_TO_LANGUAGE[model_name]) self.alpha = alpha - def _filter_pair_tokens(self, tokens: List[Tuple[str, Any]]) -> List[Tuple[str, Any]]: + def _filter_pair_tokens( + self, tokens: List[Tuple[str, Any]] + ) -> List[Tuple[str, Any]]: result = [] for token, value in tokens: if token in self.stopwords or token in self.punctuation: @@ -175,13 +180,19 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]): return new_vector - def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[SparseEmbedding]: + def _post_process_onnx_output( + self, output: OnnxOutputContext + ) -> Iterable[SparseEmbedding]: token_ids_batch = output.input_ids # attention_value shape: (batch_size, num_heads, num_tokens, num_tokens) - pooled_attention = np.mean(output.model_output[:, :, 0], axis=1) * output.attention_mask + pooled_attention = ( + np.mean(output.model_output[:, :, 0], axis=1) * output.attention_mask + ) - for document_token_ids, attention_value in zip(token_ids_batch, pooled_attention): + for document_token_ids, attention_value in zip( + token_ids_batch, pooled_attention + ): document_tokens_with_ids = ( (idx, self.invert_vocab[token_id]) for idx, token_id in enumerate(document_token_ids) @@ -261,7 +272,9 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]): result[token_id] = 1.0 return result - def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[SparseEmbedding]: + def query_embed( + self, query: Union[str, Iterable[str]], **kwargs + ) -> Iterable[SparseEmbedding]: """ To emulate BM25 behaviour, we don't need to use smart weights in the query, and it's enough to just hash the tokens and assign a weight of 1.0 to them. @@ -277,7 +290,9 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]): filtered = self._filter_pair_tokens(reconstructed) stemmed = self._stem_pair_tokens(filtered) - yield SparseEmbedding.from_dict(self._query_rehash(token for token, _ in stemmed)) + yield SparseEmbedding.from_dict( + self._query_rehash(token for token, _ in stemmed) + ) @classmethod def _get_worker_class(cls) -> Type[TextEmbeddingWorker]: diff --git a/fastembed/sparse/sparse_embedding_base.py b/fastembed/sparse/sparse_embedding_base.py index 775c286..a1cc053 100644 --- a/fastembed/sparse/sparse_embedding_base.py +++ b/fastembed/sparse/sparse_embedding_base.py @@ -28,11 +28,11 @@ class SparseEmbedding: class SparseTextEmbeddingBase(ModelManagement): def __init__( - self, - model_name: str, - cache_dir: Optional[str] = None, - threads: Optional[int] = None, - **kwargs, + self, + model_name: str, + cache_dir: Optional[str] = None, + threads: Optional[int] = None, + **kwargs, ): self.model_name = model_name self.cache_dir = cache_dir @@ -40,15 +40,17 @@ class SparseTextEmbeddingBase(ModelManagement): self._local_files_only = kwargs.pop("local_files_only", False) def embed( - self, - documents: Union[str, Iterable[str]], - batch_size: int = 256, - parallel: Optional[int] = None, - **kwargs, + self, + documents: Union[str, Iterable[str]], + batch_size: int = 256, + parallel: Optional[int] = None, + **kwargs, ) -> Iterable[SparseEmbedding]: raise NotImplementedError() - def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[SparseEmbedding]: + def passage_embed( + self, texts: Iterable[str], **kwargs + ) -> Iterable[SparseEmbedding]: """ Embeds a list of text passages into a list of embeddings. @@ -63,7 +65,9 @@ class SparseTextEmbeddingBase(ModelManagement): # This is model-specific, so that different models can have specialized implementations yield from self.embed(texts, **kwargs) - def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[SparseEmbedding]: + def query_embed( + self, query: Union[str, Iterable[str]], **kwargs + ) -> Iterable[SparseEmbedding]: """ Embeds queries diff --git a/fastembed/sparse/sparse_text_embedding.py b/fastembed/sparse/sparse_text_embedding.py index 684c311..37ea8e5 100644 --- a/fastembed/sparse/sparse_text_embedding.py +++ b/fastembed/sparse/sparse_text_embedding.py @@ -1,9 +1,12 @@ -from typing import List, Type, Dict, Any, Union, Iterable, Optional, Sequence +from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union from fastembed.common import OnnxProvider from fastembed.sparse.bm25 import Bm25 from fastembed.sparse.bm42 import Bm42 -from fastembed.sparse.sparse_embedding_base import SparseTextEmbeddingBase, SparseEmbedding +from fastembed.sparse.sparse_embedding_base import ( + SparseEmbedding, + SparseTextEmbeddingBase, +) from fastembed.sparse.splade_pp import SpladePP @@ -50,9 +53,16 @@ class SparseTextEmbedding(SparseTextEmbeddingBase): 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 @@ -85,7 +95,9 @@ class SparseTextEmbedding(SparseTextEmbeddingBase): """ yield from self.model.embed(documents, batch_size, parallel, **kwargs) - def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[SparseEmbedding]: + def query_embed( + self, query: Union[str, Iterable[str]], **kwargs + ) -> Iterable[SparseEmbedding]: """ Embeds queries diff --git a/fastembed/sparse/splade_pp.py b/fastembed/sparse/splade_pp.py index 4948ac1..859e32a 100644 --- a/fastembed/sparse/splade_pp.py +++ b/fastembed/sparse/splade_pp.py @@ -1,14 +1,16 @@ -from typing import Any, Dict, Iterable, List, Optional, Union, Type, Sequence +from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union import numpy as np from fastembed.common import OnnxProvider from fastembed.common.onnx_model import OnnxOutputContext from fastembed.common.utils import define_cache_dir -from fastembed.sparse.sparse_embedding_base import SparseEmbedding, SparseTextEmbeddingBase +from fastembed.sparse.sparse_embedding_base import ( + SparseEmbedding, + SparseTextEmbeddingBase, +) from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker - supported_splade_models = [ { "model": "prithvida/Splade_PP_en_v1", @@ -34,7 +36,9 @@ supported_splade_models = [ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]): - def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[SparseEmbedding]: + def _post_process_onnx_output( + self, output: OnnxOutputContext + ) -> Iterable[SparseEmbedding]: relu_log = np.log(1 + np.maximum(output.model_output, 0)) weighted_log = relu_log * np.expand_dims(output.attention_mask, axis=-1) diff --git a/fastembed/sparse/utils/tokenizer.py b/fastembed/sparse/utils/tokenizer.py index 354e282..c2f378b 100644 --- a/fastembed/sparse/utils/tokenizer.py +++ b/fastembed/sparse/utils/tokenizer.py @@ -68,7 +68,8 @@ class WordTokenizer: ) ] CONTRACTIONS3 = [ - re.compile(pattern) for pattern in (r"(?i) ('t)(?#X)(is)\b", r"(?i) ('t)(?#X)(was)\b") + re.compile(pattern) + for pattern in (r"(?i) ('t)(?#X)(is)\b", r"(?i) ('t)(?#X)(was)\b") ] @classmethod diff --git a/fastembed/text/clip_embedding.py b/fastembed/text/clip_embedding.py index 7da7068..843a5ec 100644 --- a/fastembed/text/clip_embedding.py +++ b/fastembed/text/clip_embedding.py @@ -1,4 +1,4 @@ -from typing import Type, List, Dict, Any, Iterable +from typing import Any, Dict, Iterable, List, Type import numpy as np @@ -34,10 +34,16 @@ class CLIPOnnxEmbedding(OnnxTextEmbedding): """ return supported_clip_models - def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]: + def _post_process_onnx_output( + self, output: OnnxOutputContext + ) -> Iterable[np.ndarray]: return output.model_output class CLIPEmbeddingWorker(OnnxTextEmbeddingWorker): - def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> OnnxTextEmbedding: - return CLIPOnnxEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs) + def init_embedding( + self, model_name: str, cache_dir: str, **kwargs + ) -> OnnxTextEmbedding: + return CLIPOnnxEmbedding( + model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs + ) diff --git a/fastembed/text/e5_onnx_embedding.py b/fastembed/text/e5_onnx_embedding.py index 8d6efa3..dc0048b 100644 --- a/fastembed/text/e5_onnx_embedding.py +++ b/fastembed/text/e5_onnx_embedding.py @@ -1,4 +1,4 @@ -from typing import Type, List, Dict, Any +from typing import Any, Dict, List, Type import numpy as np @@ -56,5 +56,9 @@ class E5OnnxEmbedding(OnnxTextEmbedding): class E5OnnxEmbeddingWorker(OnnxTextEmbeddingWorker): - def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> E5OnnxEmbedding: - return E5OnnxEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs) + def init_embedding( + self, model_name: str, cache_dir: str, **kwargs + ) -> E5OnnxEmbedding: + return E5OnnxEmbedding( + model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs + ) diff --git a/fastembed/text/jina_onnx_embedding.py b/fastembed/text/jina_onnx_embedding.py index 4237847..9f70fef 100644 --- a/fastembed/text/jina_onnx_embedding.py +++ b/fastembed/text/jina_onnx_embedding.py @@ -1,4 +1,4 @@ -from typing import Type, List, Dict, Any, Iterable +from typing import Any, Dict, Iterable, List, Type import numpy as np @@ -59,12 +59,18 @@ class JinaOnnxEmbedding(OnnxTextEmbedding): """ return supported_jina_models - def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]: + def _post_process_onnx_output( + self, output: OnnxOutputContext + ) -> Iterable[np.ndarray]: embeddings = output.model_output attn_mask = output.attention_mask return normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32) class JinaEmbeddingWorker(OnnxTextEmbeddingWorker): - def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> OnnxTextEmbedding: - return JinaOnnxEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs) + def init_embedding( + self, model_name: str, cache_dir: str, **kwargs + ) -> OnnxTextEmbedding: + return JinaOnnxEmbedding( + model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs + ) diff --git a/fastembed/text/mini_lm_embedding.py b/fastembed/text/mini_lm_embedding.py new file mode 100644 index 0000000..43f4d73 --- /dev/null +++ b/fastembed/text/mini_lm_embedding.py @@ -0,0 +1,69 @@ +from typing import Any, Dict, Iterable, List, Type + +import numpy as np + +from fastembed.common.onnx_model import OnnxOutputContext +from fastembed.common.utils import normalize +from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker +from fastembed.text.onnx_text_model import TextEmbeddingWorker + +supported_mini_lm_models = [ + { + "model": "sentence-transformers/all-MiniLM-L6-v2", + "dim": 384, + "description": "Sentence Transformer model, MiniLM-L6-v2", + "size_in_GB": 0.09, + "sources": { + "url": "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz", + "hf": "qdrant/all-MiniLM-L6-v2-onnx", + }, + "model_file": "model.onnx", + } +] + + +class MiniLMOnnxEmbedding(OnnxTextEmbedding): + @classmethod + def _get_worker_class(cls) -> Type[TextEmbeddingWorker]: + return MiniLMEmbeddingWorker + + @classmethod + def mean_pooling(self, model_output: np.ndarray, attention_mask: np.ndarray) -> np.ndarray: + token_embeddings = model_output + input_mask_expanded = np.expand_dims(attention_mask, axis=-1) + input_mask_expanded = np.tile( + input_mask_expanded, (1, 1, token_embeddings.shape[-1]) + ) + input_mask_expanded = input_mask_expanded.astype(float) + sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1) + sum_mask = np.sum(input_mask_expanded, axis=1) + pooled_embeddings = sum_embeddings / np.maximum(sum_mask, 1e-9) + return pooled_embeddings + + @classmethod + def list_supported_models(cls) -> List[Dict[str, Any]]: + """Lists the supported models. + + Returns: + List[Dict[str, Any]]: A list of dictionaries containing the model information. + """ + return supported_mini_lm_models + + def _post_process_onnx_output( + self, output: OnnxOutputContext + ) -> Iterable[np.ndarray]: + embeddings = output.model_output + attn_mask = output.attention_mask + return normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32) + + +class MiniLMEmbeddingWorker(OnnxTextEmbeddingWorker): + def init_embedding( + self, + model_name: str, + cache_dir: str, + **kwargs + ) -> OnnxTextEmbedding: + return MiniLMOnnxEmbedding( + model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs + ) diff --git a/fastembed/text/onnx_embedding.py b/fastembed/text/onnx_embedding.py index 39645b4..77b3754 100644 --- a/fastembed/text/onnx_embedding.py +++ b/fastembed/text/onnx_embedding.py @@ -1,14 +1,13 @@ -from typing import Dict, Optional, Union, Iterable, Type, List, Any, Sequence +from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union import numpy as np from fastembed.common import OnnxProvider from fastembed.common.onnx_model import OnnxOutputContext -from fastembed.common.utils import normalize, define_cache_dir -from fastembed.text.onnx_text_model import TextEmbeddingWorker, OnnxTextModel +from fastembed.common.utils import define_cache_dir, normalize +from fastembed.text.onnx_text_model import OnnxTextModel, TextEmbeddingWorker from fastembed.text.text_embedding_base import TextEmbeddingBase - supported_onnx_models = [ { "model": "BAAI/bge-base-en", @@ -71,17 +70,6 @@ supported_onnx_models = [ }, "model_file": "model_optimized.onnx", }, - { - "model": "sentence-transformers/all-MiniLM-L6-v2", - "dim": 384, - "description": "Sentence Transformer model, MiniLM-L6-v2", - "size_in_GB": 0.09, - "sources": { - "url": "https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz", - "hf": "qdrant/all-MiniLM-L6-v2-onnx", - }, - "model_file": "model.onnx", - }, { "model": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2", "dim": 384, @@ -286,7 +274,9 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]): """ 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]: embeddings = output.model_output return normalize(embeddings[:, 0]).astype(np.float32) @@ -298,4 +288,6 @@ class OnnxTextEmbeddingWorker(TextEmbeddingWorker): cache_dir: str, **kwargs, ) -> OnnxTextEmbedding: - return OnnxTextEmbedding(model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs) + return OnnxTextEmbedding( + model_name=model_name, cache_dir=cache_dir, threads=1, **kwargs + ) diff --git a/fastembed/text/onnx_text_model.py b/fastembed/text/onnx_text_model.py index b66a808..320a044 100644 --- a/fastembed/text/onnx_text_model.py +++ b/fastembed/text/onnx_text_model.py @@ -1,14 +1,14 @@ import os from multiprocessing import get_all_start_methods from pathlib import Path -from typing import Any, Dict, Iterable, List, Optional, Tuple, Type, Union, Sequence +from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple, Type, Union import numpy as np from tokenizers import Encoding from fastembed.common import OnnxProvider +from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T from fastembed.common.preprocessor_utils import load_tokenizer -from fastembed.common.onnx_model import OnnxModel, EmbeddingWorker, T, OnnxOutputContext from fastembed.common.utils import iter_batch from fastembed.parallel_processor import ParallelWorkerPool @@ -44,7 +44,10 @@ class OnnxTextModel(OnnxModel[T]): 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.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir) @@ -105,7 +108,9 @@ class OnnxTextModel(OnnxModel[T]): for batch in iter_batch(documents, 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 diff --git a/fastembed/text/text_embedding.py b/fastembed/text/text_embedding.py index c60585d..02d1b52 100644 --- a/fastembed/text/text_embedding.py +++ b/fastembed/text/text_embedding.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, Iterable, List, Optional, Type, Union, Sequence +from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union import numpy as np @@ -6,6 +6,7 @@ from fastembed.common import OnnxProvider from fastembed.text.clip_embedding import CLIPOnnxEmbedding from fastembed.text.e5_onnx_embedding import E5OnnxEmbedding from fastembed.text.jina_onnx_embedding import JinaOnnxEmbedding +from fastembed.text.mini_lm_embedding import MiniLMOnnxEmbedding from fastembed.text.onnx_embedding import OnnxTextEmbedding from fastembed.text.text_embedding_base import TextEmbeddingBase @@ -16,6 +17,7 @@ class TextEmbedding(TextEmbeddingBase): E5OnnxEmbedding, JinaOnnxEmbedding, CLIPOnnxEmbedding, + MiniLMOnnxEmbedding, ] @classmethod @@ -59,9 +61,16 @@ class TextEmbedding(TextEmbeddingBase): 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 diff --git a/fastembed/text/text_embedding_base.py b/fastembed/text/text_embedding_base.py index e0a4c4d..61963fc 100644 --- a/fastembed/text/text_embedding_base.py +++ b/fastembed/text/text_embedding_base.py @@ -42,7 +42,9 @@ class TextEmbeddingBase(ModelManagement): # This is model-specific, so that different models can have specialized implementations yield from self.embed(texts, **kwargs) - def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]: + def query_embed( + self, query: Union[str, Iterable[str]], **kwargs + ) -> Iterable[np.ndarray]: """ Embeds queries diff --git a/tests/profiling.py b/tests/profiling.py index 9a07b55..452be0a 100644 --- a/tests/profiling.py +++ b/tests/profiling.py @@ -57,7 +57,9 @@ class HF: self.tokenizer = AutoTokenizer.from_pretrained(model_id) def embed(self, texts: List[str]): - encoded_input = self.tokenizer(texts, max_length=512, padding=True, truncation=True, return_tensors="pt") + encoded_input = self.tokenizer( + texts, max_length=512, padding=True, truncation=True, return_tensors="pt" + ) model_output = self.model(**encoded_input) sentence_embeddings = model_output[0][:, 0] sentence_embeddings = F.normalize(sentence_embeddings) @@ -84,7 +86,9 @@ embedding_model = DefaultEmbedding() # %% -def calculate_time_stats(embed_func: Callable, documents: list, k: int) -> Tuple[float, float, float]: +def calculate_time_stats( + embed_func: Callable, documents: list, k: int +) -> Tuple[float, float, float]: times = [] for _ in range(k): # Timing the embed_func call @@ -101,13 +105,17 @@ def calculate_time_stats(embed_func: Callable, documents: list, k: int) -> Tuple # %% hf_stats = calculate_time_stats(hf.embed, documents, k=2) print(f"Huggingface Transformers (Average, Max, Min): {hf_stats}") -fst_stats = calculate_time_stats(lambda x: list(embedding_model.embed(x)), documents, k=2) +fst_stats = calculate_time_stats( + lambda x: list(embedding_model.embed(x)), documents, k=2 +) print(f"FastEmbed (Average, Max, Min): {fst_stats}") # %% def plot_character_per_second_comparison( - hf_stats: Tuple[float, float, float], fst_stats: Tuple[float, float, float], documents: list + hf_stats: Tuple[float, float, float], + fst_stats: Tuple[float, float, float], + documents: list, ): # Calculating total characters in documents total_characters = sum(len(doc) for doc in documents) diff --git a/tests/test_attention_embeddings.py b/tests/test_attention_embeddings.py index 1e6c6ce..8a0ff78 100644 --- a/tests/test_attention_embeddings.py +++ b/tests/test_attention_embeddings.py @@ -1,11 +1,12 @@ -import pytest - import numpy as np +import pytest from fastembed import SparseTextEmbedding -@pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"]) +@pytest.mark.parametrize( + "model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"] +) def test_attention_embeddings(model_name): model = SparseTextEmbedding(model_name=model_name) @@ -61,7 +62,9 @@ def test_attention_embeddings(model_name): assert len(result.indices) == 2 -@pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"]) +@pytest.mark.parametrize( + "model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"] +) def test_parallel_processing(model_name): model = SparseTextEmbedding(model_name=model_name) diff --git a/tests/test_late_interaction_embeddings.py b/tests/test_late_interaction_embeddings.py index 823ade0..0a871c7 100644 --- a/tests/test_late_interaction_embeddings.py +++ b/tests/test_late_interaction_embeddings.py @@ -1,7 +1,8 @@ import numpy as np -from fastembed.late_interaction.late_interaction_text_embedding import LateInteractionTextEmbedding - +from fastembed.late_interaction.late_interaction_text_embedding import ( + LateInteractionTextEmbedding, +) # vectors are abridged and rounded for brevity CANONICAL_COLUMN_VALUES = { @@ -76,7 +77,9 @@ def test_single_embedding(): for model_name, expected_result in CANONICAL_COLUMN_VALUES.items(): print("evaluating", model_name) - model = LateInteractionTextEmbedding(model_name=model_name, cache_dir="colbert-cache") + model = LateInteractionTextEmbedding( + model_name=model_name, cache_dir="colbert-cache" + ) result = next(iter(model.embed(docs_to_embed, batch_size=6))) token_num, abridged_dim = expected_result.shape assert np.allclose(result[:, :abridged_dim], expected_result, atol=10e-4) @@ -87,7 +90,9 @@ def test_single_embedding_query(): for model_name, expected_result in CANONICAL_QUERY_VALUES.items(): print("evaluating", model_name) - model = LateInteractionTextEmbedding(model_name=model_name, cache_dir="colbert-cache") + model = LateInteractionTextEmbedding( + model_name=model_name, cache_dir="colbert-cache" + ) result = next(iter(model.query_embed(queries_to_embed))) token_num, abridged_dim = expected_result.shape assert np.allclose(result[:, :abridged_dim], expected_result, atol=10e-4) diff --git a/tests/test_sparse_embeddings.py b/tests/test_sparse_embeddings.py index 370a5ae..06a5366 100644 --- a/tests/test_sparse_embeddings.py +++ b/tests/test_sparse_embeddings.py @@ -93,5 +93,9 @@ def test_parallel_processing(): == sparse_embedding_duo.indices.tolist() == sparse_embedding_all.indices.tolist() ) - assert np.allclose(sparse_embedding.values, sparse_embedding_duo.values, atol=1e-3) - assert np.allclose(sparse_embedding.values, sparse_embedding_all.values, atol=1e-3) + assert np.allclose( + sparse_embedding.values, sparse_embedding_duo.values, atol=1e-3 + ) + assert np.allclose( + sparse_embedding.values, sparse_embedding_all.values, atol=1e-3 + ) diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py index 6a420b6..7ca4005 100644 --- a/tests/test_text_onnx_embeddings.py +++ b/tests/test_text_onnx_embeddings.py @@ -26,18 +26,30 @@ CANONICAL_VECTOR_VALUES = { "BAAI/bge-large-en-v1.5-quantized": np.array( [0.03434538, 0.03316108, 0.02191251, -0.03713358, -0.01577825] ), - "sentence-transformers/all-MiniLM-L6-v2": np.array([0.0259, 0.0058, 0.0114, 0.0380, -0.0233]), + "sentence-transformers/all-MiniLM-L6-v2": np.array( + [-0.034478, 0.03102, 0.00673, 0.02611, -0.039362] + ), "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2": np.array( [0.0094, 0.0184, 0.0328, 0.0072, -0.0351] ), - "intfloat/multilingual-e5-large": np.array([0.0098, 0.0045, 0.0066, -0.0354, 0.0070]), + "intfloat/multilingual-e5-large": np.array( + [0.0098, 0.0045, 0.0066, -0.0354, 0.0070] + ), "sentence-transformers/paraphrase-multilingual-mpnet-base-v2": np.array( [-0.01341097, 0.0416553, -0.00480805, 0.02844842, 0.0505299] ), - "jinaai/jina-embeddings-v2-small-en": np.array([-0.0455, -0.0428, -0.0122, 0.0613, 0.0015]), - "jinaai/jina-embeddings-v2-base-en": np.array([-0.0332, -0.0509, 0.0287, -0.0043, -0.0077]), - "jinaai/jina-embeddings-v2-base-de": np.array([-0.0085, 0.0417, 0.0342, 0.0309, -0.0149]), - "nomic-ai/nomic-embed-text-v1": np.array([0.0061, 0.0103, -0.0296, -0.0242, -0.0170]), + "jinaai/jina-embeddings-v2-small-en": np.array( + [-0.0455, -0.0428, -0.0122, 0.0613, 0.0015] + ), + "jinaai/jina-embeddings-v2-base-en": np.array( + [-0.0332, -0.0509, 0.0287, -0.0043, -0.0077] + ), + "jinaai/jina-embeddings-v2-base-de": np.array( + [-0.0085, 0.0417, 0.0342, 0.0309, -0.0149] + ), + "nomic-ai/nomic-embed-text-v1": np.array( + [0.0061, 0.0103, -0.0296, -0.0242, -0.0170] + ), "nomic-ai/nomic-embed-text-v1.5": np.array( [-1.6531514e-02, 8.5380634e-05, -1.8171231e-01, -3.9333291e-03, 1.2763254e-02] ), @@ -50,13 +62,21 @@ CANONICAL_VECTOR_VALUES = { "mixedbread-ai/mxbai-embed-large-v1": np.array( [0.02295546, 0.03196154, 0.016512, -0.04031524, -0.0219634] ), - "snowflake/snowflake-arctic-embed-xs": np.array([0.0092, 0.0619, 0.0196, 0.009, -0.0114]), - "snowflake/snowflake-arctic-embed-s": np.array([-0.0416, -0.0867, 0.0209, 0.0554, -0.0272]), - "snowflake/snowflake-arctic-embed-m": np.array([-0.0329, 0.0364, 0.0481, 0.0016, 0.0328]), + "snowflake/snowflake-arctic-embed-xs": np.array( + [0.0092, 0.0619, 0.0196, 0.009, -0.0114] + ), + "snowflake/snowflake-arctic-embed-s": np.array( + [-0.0416, -0.0867, 0.0209, 0.0554, -0.0272] + ), + "snowflake/snowflake-arctic-embed-m": np.array( + [-0.0329, 0.0364, 0.0481, 0.0016, 0.0328] + ), "snowflake/snowflake-arctic-embed-m-long": np.array( [0.0080, -0.0266, -0.0335, 0.0282, 0.0143] ), - "snowflake/snowflake-arctic-embed-l": np.array([0.0189, -0.0673, 0.0183, 0.0124, 0.0146]), + "snowflake/snowflake-arctic-embed-l": np.array( + [0.0189, -0.0673, 0.0183, 0.0124, 0.0146] + ), "Qdrant/clip-ViT-B-32-text": np.array([0.0083, 0.0103, -0.0138, 0.0199, -0.0069]), }