upstream-archive/fastembed
1
0
Fork 0
mirror of https://github.com/qdrant/fastembed.git synced 2026-07-23 11:20:51 -05:00
Code Issues Packages Projects Releases Wiki Activity

MiniLM fix (#275)

Browse Source 
* 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 <dimitriyrudenk@gmail.com>
This commit is contained in:
Dmitrii Ogn
2024-06-14 16:42:31 +03:00
committed by GitHub
parent fd0b26f009
commit 331207976e
39 changed files with 460 additions and 191 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
experiments/attention_export.py
View File

@@ -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
)

8
experiments/try_attention_export.py
View File

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

4
fastembed/__init__.py
View File

@@ -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")

2
fastembed/common/__init__.py
View File

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

29
fastembed/common/model_management.py
View File

@@ -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

19
fastembed/common/onnx_model.py
View File

@@ -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:

6
fastembed/common/preprocessor_utils.py
View File

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

2
fastembed/common/types.py
View File

@@ -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

2
fastembed/common/utils.py
View File

@@ -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

1
fastembed/image/__init__.py
View File

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

13
fastembed/image/image_embedding.py
View File

@@ -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

20
fastembed/image/onnx_embedding.py
View File

@@ -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
)

19
fastembed/image/onnx_image_model.py
View File

@@ -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

5
fastembed/image/transform/functional.py
View File

@@ -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:

106
fastembed/image/transform/operators.py
View File

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

5
fastembed/late_interaction/__init__.py
View File

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

11
fastembed/late_interaction/colbert.py
View File

@@ -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:

4
fastembed/late_interaction/late_interaction_embedding_base.py
View File

@@ -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

13
fastembed/late_interaction/late_interaction_text_embedding.py
View File

@@ -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

8
fastembed/parallel_processor.py
View File

@@ -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

19
fastembed/sparse/bm25.py
View File

@@ -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)

33
fastembed/sparse/bm42.py
View File

@@ -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]:

28
fastembed/sparse/sparse_embedding_base.py
View File

@@ -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

22
fastembed/sparse/sparse_text_embedding.py
View File

@@ -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

12
fastembed/sparse/splade_pp.py
View File

@@ -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)

3
fastembed/sparse/utils/tokenizer.py
View File

@@ -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

14
fastembed/text/clip_embedding.py
View File

@@ -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
)

10
fastembed/text/e5_onnx_embedding.py
View File

@@ -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
)

14
fastembed/text/jina_onnx_embedding.py
View File

@@ -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
)

69
fastembed/text/mini_lm_embedding.py Normal file
View File

@@ -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
)

26
fastembed/text/onnx_embedding.py
View File

@@ -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
)

13
fastembed/text/onnx_text_model.py
View File

@@ -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

15
fastembed/text/text_embedding.py
View File

@@ -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

4
fastembed/text/text_embedding_base.py
View File

@@ -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

16
tests/profiling.py
View File

@@ -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)

11
tests/test_attention_embeddings.py
View File

@@ -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)

13
tests/test_late_interaction_embeddings.py
View File

@@ -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)

8
tests/test_sparse_embeddings.py
View File

@@ -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
)

40
tests/test_text_onnx_embeddings.py
View File

@@ -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]),
}