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

Merge branch 'main' into mean_pooling

Browse Source
This commit is contained in:
Dmitrii Ogn
2025-01-30 15:59:50 +03:00
committed by GitHub
parent e41197dc0b 993dcd5f68
commit a32500fa6d
48 changed files with 799 additions and 247 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
.github/ISSUE_TEMPLATE/bug-report.yml vendored
View File

@@ -39,7 +39,7 @@ body:
attributes:
label: FastEmbed version
description: What version of FastEmbed are you running? python -c "import fastembed; print(fastembed.__version__)". If you're not on the latest, please upgrade and see if the problem persists.
placeholder: v0.4.2
placeholder: v0.5.1
validations:
required: true
- type: dropdown

2
.github/workflows/python-tests.yml vendored
View File

@@ -38,7 +38,7 @@ jobs:
run: |
python -m pip install poetry
poetry config virtualenvs.create false
poetry install --no-interaction --no-ansi --without docs
poetry install --no-interaction --no-ansi --without dev,docs
- name: Run pytest
run: |

2
NOTICE
View File

@@ -7,6 +7,8 @@ This distribution includes the following Jina AI models, each with its respectiv
- License: cc-by-nc-4.0
- jinaai/jina-reranker-v2-base-multilingual
- License: cc-by-nc-4.0
- jinaai/jina-embeddings-v3
- License: cc-by-nc-4.0
These models are developed by Jina (https://jina.ai/) and are subject to Jina AI's licensing terms.

17
docs/examples/ColBERT_with_FastEmbed.ipynb
View File

@@ -54,7 +54,14 @@
},
{
"data": {
"text/plain": "[{'model': 'colbert-ir/colbertv2.0',\n 'dim': 128,\n 'description': 'Late interaction model',\n 'size_in_GB': 0.44,\n 'sources': {'hf': 'colbert-ir/colbertv2.0'},\n 'model_file': 'model.onnx'}]"
"text/plain": [
"[{'model': 'colbert-ir/colbertv2.0',\n",
" 'dim': 128,\n",
" 'description': 'Late interaction model',\n",
" 'size_in_GB': 0.44,\n",
" 'sources': {'hf': 'colbert-ir/colbertv2.0'},\n",
" 'model_file': 'model.onnx'}]"
]
},
"execution_count": 1,
"metadata": {},
@@ -212,7 +219,9 @@
"outputs": [
{
"data": {
"text/plain": "((26, 128), (32, 128))"
"text/plain": [
"((26, 128), (32, 128))"
]
},
"execution_count": 18,
"metadata": {},
@@ -271,7 +280,9 @@
"import numpy as np\n",
"\n",
"\n",
"def compute_relevance_scores(query_embedding: np.array, document_embeddings: np.array, k: int):\n",
"def compute_relevance_scores(\n",
" query_embedding: np.array, document_embeddings: np.array, k: int\n",
") -> list[int]:\n",
" \"\"\"\n",
" Compute relevance scores for top-k documents given a query.\n",
"\n",

2
docs/examples/FastEmbed_vs_HF_Comparison.ipynb
View File

@@ -152,7 +152,7 @@
" HuggingFace Transformer implementation of FlagEmbedding\n",
" \"\"\"\n",
"\n",
" def __init__(self, model_id: str):\n",
" def __init__(self, model_id: str) -> None:\n",
" self.model = AutoModel.from_pretrained(model_id)\n",
" self.tokenizer = AutoTokenizer.from_pretrained(model_id)\n",
"\n",

6
docs/examples/Hybrid_Search.ipynb
View File

@@ -488,7 +488,7 @@
}
],
"source": [
"def make_sparse_embedding(texts: list[str]):\n",
"def make_sparse_embedding(texts: list[str]) -> list[SparseEmbedding]:\n",
" return list(sparse_model.embed(texts, batch_size=32))\n",
"\n",
"\n",
@@ -615,7 +615,7 @@
}
],
"source": [
"def get_tokens_and_weights(sparse_embedding, model_name):\n",
"def get_tokens_and_weights(sparse_embedding, model_name) -> dict[str, float]:\n",
" # Find the tokenizer for the model\n",
" tokenizer_source = None\n",
" for model_info in SparseTextEmbedding.list_supported_models():\n",
@@ -626,7 +626,7 @@
" raise ValueError(f\"Model {model_name} not found in the supported models.\")\n",
"\n",
" tokenizer = AutoTokenizer.from_pretrained(tokenizer_source)\n",
" token_weight_dict = {}\n",
" token_weight_dict: dict[str, float] = {}\n",
" for i in range(len(sparse_embedding.indices)):\n",
" token = tokenizer.decode([sparse_embedding.indices[i]])\n",
" weight = sparse_embedding.values[i]\n",

142
fastembed/common/model_management.py
View File

@@ -1,12 +1,14 @@
import os
import time
import json
import shutil
import tarfile
from pathlib import Path
from typing import Any, Optional
import requests
from huggingface_hub import snapshot_download
from huggingface_hub import snapshot_download, model_info, list_repo_tree
from huggingface_hub.hf_api import RepoFile
from huggingface_hub.utils import (
RepositoryNotFoundError,
disable_progress_bars,
@@ -17,6 +19,8 @@ from tqdm import tqdm
class ModelManagement:
METADATA_FILE = "files_metadata.json"
@classmethod
def list_supported_models(cls) -> list[dict[str, Any]]:
"""Lists the supported models.
@@ -98,21 +102,73 @@ class ModelManagement:
cls,
hf_source_repo: str,
cache_dir: str,
extra_patterns: Optional[list[str]] = None,
extra_patterns: list[str],
local_files_only: bool = False,
**kwargs,
**kwargs: Any,
) -> str:
"""
Downloads a model from HuggingFace Hub.
Args:
hf_source_repo (str): Name of the model on HuggingFace Hub, e.g. "qdrant/all-MiniLM-L6-v2-onnx".
cache_dir (Optional[str]): The path to the cache directory.
extra_patterns (Optional[list[str]]): extra patterns to allow in the snapshot download, typically
extra_patterns (list[str]): extra patterns to allow in the snapshot download, typically
includes the required model files.
local_files_only (bool, optional): Whether to only use local files. Defaults to False.
specific_model_path (Optional[str], optional): The path to the model dir already pooled from external source
Returns:
Path: The path to the model directory.
"""
def _verify_files_from_metadata(
model_dir: Path, stored_metadata: dict[str, Any], repo_files: list[RepoFile]
) -> bool:
try:
for rel_path, meta in stored_metadata.items():
file_path = model_dir / rel_path
if not file_path.exists():
return False
if repo_files: # online verification
file_info = next((f for f in repo_files if f.path == file_path.name), None)
if (
not file_info
or file_info.size != meta["size"]
or file_info.blob_id != meta["blob_id"]
):
return False
else: # offline verification
if file_path.stat().st_size != meta["size"]:
return False
return True
except (OSError, KeyError) as e:
logger.error(f"Error verifying files: {str(e)}")
return False
def _collect_file_metadata(
model_dir: Path, repo_files: list[RepoFile]
) -> dict[str, dict[str, int]]:
meta = {}
file_info_map = {f.path: f for f in repo_files}
for file_path in model_dir.rglob("*"):
if file_path.is_file() and file_path.name != cls.METADATA_FILE:
repo_file = file_info_map.get(file_path.name)
if repo_file:
meta[str(file_path.relative_to(model_dir))] = {
"size": repo_file.size,
"blob_id": repo_file.blob_id,
}
return meta
def _save_file_metadata(model_dir: Path, meta: dict[str, dict[str, int]]) -> None:
try:
if not model_dir.exists():
model_dir.mkdir(parents=True, exist_ok=True)
(model_dir / cls.METADATA_FILE).write_text(json.dumps(meta))
except (OSError, ValueError) as e:
logger.warning(f"Error saving metadata: {str(e)}")
allow_patterns = [
"config.json",
"tokenizer.json",
@@ -120,16 +176,59 @@ class ModelManagement:
"special_tokens_map.json",
"preprocessor_config.json",
]
if extra_patterns is not None:
allow_patterns.extend(extra_patterns)
allow_patterns.extend(extra_patterns)
snapshot_dir = Path(cache_dir) / f"models--{hf_source_repo.replace('/', '--')}"
is_cached = snapshot_dir.exists()
metadata_file = snapshot_dir / cls.METADATA_FILE
if is_cached:
if local_files_only:
disable_progress_bars()
if metadata_file.exists():
metadata = json.loads(metadata_file.read_text())
verified = _verify_files_from_metadata(snapshot_dir, metadata, repo_files=[])
if not verified:
logger.warning(
"Local file sizes do not match the metadata."
) # do not raise, still make an attempt to load the model
else:
logger.warning(
"Metadata file not found. Proceeding without checking local files."
) # if users have downloaded models from hf manually, or they're updating from previous versions of
# fastembed
result = snapshot_download(
repo_id=hf_source_repo,
allow_patterns=allow_patterns,
cache_dir=cache_dir,
local_files_only=local_files_only,
**kwargs,
)
return result
repo_revision = model_info(hf_source_repo).sha
repo_tree = list(list_repo_tree(hf_source_repo, revision=repo_revision, repo_type="model"))
allowed_extensions = {".json", ".onnx", ".txt"}
repo_files = (
[
f
for f in repo_tree
if isinstance(f, RepoFile) and Path(f.path).suffix in allowed_extensions
]
if repo_tree
else []
)
verified_metadata = False
if snapshot_dir.exists() and metadata_file.exists():
metadata = json.loads(metadata_file.read_text())
verified_metadata = _verify_files_from_metadata(snapshot_dir, metadata, repo_files)
if verified_metadata:
disable_progress_bars()
return snapshot_download(
result = snapshot_download(
repo_id=hf_source_repo,
allow_patterns=allow_patterns,
cache_dir=cache_dir,
@@ -137,8 +236,26 @@ class ModelManagement:
**kwargs,
)
if (
not verified_metadata
): # metadata is not up-to-date, update it and check whether the files have been
# downloaded correctly
metadata = _collect_file_metadata(snapshot_dir, repo_files)
download_successful = _verify_files_from_metadata(
snapshot_dir, metadata, repo_files=[]
) # offline verification
if not download_successful:
raise ValueError(
"Files have been corrupted during downloading process. "
"Please check your internet connection and try again."
)
_save_file_metadata(snapshot_dir, metadata)
return result
@classmethod
def decompress_to_cache(cls, targz_path: str, cache_dir: str):
def decompress_to_cache(cls, targz_path: str, cache_dir: str) -> str:
"""
Decompresses a .tar.gz file to a cache directory.
@@ -221,7 +338,7 @@ class ModelManagement:
@classmethod
def download_model(
cls, model: dict[str, Any], cache_dir: Path, retries: int = 3, **kwargs
cls, model: dict[str, Any], cache_dir: Path, retries: int = 3, **kwargs: Any
) -> Path:
"""
Downloads a model from HuggingFace Hub or Google Cloud Storage.
@@ -248,6 +365,9 @@ class ModelManagement:
Path: The path to the downloaded model directory.
"""
local_files_only = kwargs.get("local_files_only", False)
specific_model_path: Optional[str] = kwargs.pop("specific_model_path", None)
if specific_model_path:
return Path(specific_model_path)
retries = 1 if local_files_only else retries
hf_source = model.get("sources", {}).get("hf")
url_source = model.get("sources", {}).get("url")

8
fastembed/common/onnx_model.py
View File

@@ -33,7 +33,7 @@ class OnnxModel(Generic[T]):
self.tokenizer = None
def _preprocess_onnx_input(
self, onnx_input: dict[str, np.ndarray], **kwargs
self, onnx_input: dict[str, np.ndarray], **kwargs: Any
) -> dict[str, np.ndarray]:
"""
Preprocess the onnx input.
@@ -103,7 +103,7 @@ class OnnxModel(Generic[T]):
def load_onnx_model(self) -> None:
raise NotImplementedError("Subclasses must implement this method")
def onnx_embed(self, *args, **kwargs) -> OnnxOutputContext:
def onnx_embed(self, *args: Any, **kwargs: Any) -> OnnxOutputContext:
raise NotImplementedError("Subclasses must implement this method")
@@ -112,7 +112,7 @@ class EmbeddingWorker(Worker):
self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
) -> OnnxModel:
raise NotImplementedError()
@@ -120,7 +120,7 @@ class EmbeddingWorker(Worker):
self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
):
self.model = self.init_embedding(model_name, cache_dir, **kwargs)

4
fastembed/common/types.py
View File

@@ -1,4 +1,4 @@
import os
from pathlib import Path
import sys
from PIL import Image
from typing import Any, Iterable, Union
@@ -9,7 +9,7 @@ else:
from typing_extensions import TypeAlias
PathInput: TypeAlias = Union[str, os.PathLike]
PathInput: TypeAlias = Union[str, Path]
PilInput: TypeAlias = Union[Image.Image, Iterable[Image.Image]]
ImageInput: TypeAlias = Union[PathInput, Iterable[PathInput], PilInput]

8
fastembed/common/utils.py
View File

@@ -5,12 +5,14 @@ import tempfile
import unicodedata
from pathlib import Path
from itertools import islice
from typing import Generator, Iterable, Optional, Union
from typing import Iterable, Optional, TypeVar
import numpy as np
T = TypeVar("T")
def normalize(input_array, p=2, dim=1, eps=1e-12) -> np.ndarray:
def normalize(input_array: np.ndarray, p: int = 2, dim: int = 1, eps: float = 1e-12) -> np.ndarray:
# Calculate the Lp norm along the specified dimension
norm = np.linalg.norm(input_array, ord=p, axis=dim, keepdims=True)
norm = np.maximum(norm, eps) # Avoid division by zero
@@ -18,7 +20,7 @@ def normalize(input_array, p=2, dim=1, eps=1e-12) -> np.ndarray:
return normalized_array
def iter_batch(iterable: Union[Iterable, Generator], size: int) -> Iterable:
def iter_batch(iterable: Iterable[T], size: int) -> Iterable[list[T]]:
"""
>>> list(iter_batch([1,2,3,4,5], 3))
[[1, 2, 3], [4, 5]]

4
fastembed/embedding.py
View File

@@ -1,4 +1,4 @@
from typing import Optional
from typing import Optional, Any
from loguru import logger
@@ -19,6 +19,6 @@ class JinaEmbedding(TextEmbedding):
model_name: str = "jinaai/jina-embeddings-v2-base-en",
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
**kwargs,
**kwargs: Any,
):
super().__init__(model_name, cache_dir, threads, **kwargs)

4
fastembed/image/image_embedding.py
View File

@@ -49,7 +49,7 @@ class ImageEmbedding(ImageEmbeddingBase):
cuda: bool = False,
device_ids: Optional[list[int]] = None,
lazy_load: bool = False,
**kwargs,
**kwargs: Any,
):
super().__init__(model_name, cache_dir, threads, **kwargs)
for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
@@ -77,7 +77,7 @@ class ImageEmbedding(ImageEmbeddingBase):
images: ImageInput,
batch_size: int = 16,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[np.ndarray]:
"""
Encode a list of documents into list of embeddings.

6
fastembed/image/image_embedding_base.py
View File

@@ -1,4 +1,4 @@
from typing import Iterable, Optional
from typing import Iterable, Optional, Any
import numpy as np
@@ -12,7 +12,7 @@ class ImageEmbeddingBase(ModelManagement):
model_name: str,
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
**kwargs,
**kwargs: Any,
):
self.model_name = model_name
self.cache_dir = cache_dir
@@ -24,7 +24,7 @@ class ImageEmbeddingBase(ModelManagement):
images: ImageInput,
batch_size: int = 16,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[np.ndarray]:
"""
Embeds a list of images into a list of embeddings.

18
fastembed/image/onnx_embedding.py
View File

@@ -1,7 +1,6 @@
from typing import Any, Iterable, Optional, Sequence, Type
import numpy as np
from fastembed.common import ImageInput, OnnxProvider
from fastembed.common.onnx_model import OnnxOutputContext
from fastembed.common.utils import define_cache_dir, normalize
@@ -78,7 +77,8 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
device_ids: Optional[list[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
**kwargs,
specific_model_path: Optional[str] = None,
**kwargs: Any,
):
"""
Args:
@@ -96,6 +96,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
Raises:
ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
@@ -120,7 +121,10 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
self.model_description = self._get_model_description(model_name)
self.cache_dir = define_cache_dir(cache_dir)
self._model_dir = self.download_model(
self.model_description, self.cache_dir, local_files_only=self._local_files_only
self.model_description,
self.cache_dir,
local_files_only=self._local_files_only,
specific_model_path=specific_model_path,
)
if not self.lazy_load:
@@ -145,7 +149,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
Lists the supported models.
Returns:
list[Dict[str, Any]]: A list of dictionaries containing the model information.
list[dict[str, Any]]: A list of dictionaries containing the model information.
"""
return supported_onnx_models
@@ -154,7 +158,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
images: ImageInput,
batch_size: int = 16,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[np.ndarray]:
"""
Encode a list of images into list of embeddings.
@@ -189,7 +193,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
return OnnxImageEmbeddingWorker
def _preprocess_onnx_input(
self, onnx_input: dict[str, np.ndarray], **kwargs
self, onnx_input: dict[str, np.ndarray], **kwargs: Any
) -> dict[str, np.ndarray]:
"""
Preprocess the onnx input.
@@ -202,7 +206,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]):
class OnnxImageEmbeddingWorker(ImageEmbeddingWorker):
def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> OnnxImageEmbedding:
def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> OnnxImageEmbedding:
return OnnxImageEmbedding(
model_name=model_name,
cache_dir=cache_dir,

8
fastembed/image/onnx_image_model.py
View File

@@ -24,12 +24,12 @@ class OnnxImageModel(OnnxModel[T]):
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
raise NotImplementedError("Subclasses must implement this method")
def __init__(self) -> None:
def __init__(self):
super().__init__()
self.processor = None
def _preprocess_onnx_input(
self, onnx_input: dict[str, np.ndarray], **kwargs
self, onnx_input: dict[str, np.ndarray], **kwargs: Any
) -> dict[str, np.ndarray]:
"""
Preprocess the onnx input.
@@ -61,7 +61,7 @@ class OnnxImageModel(OnnxModel[T]):
def _build_onnx_input(self, encoded: np.ndarray) -> dict[str, np.ndarray]:
return {node.name: encoded for node in self.model.get_inputs()}
def onnx_embed(self, images: list[ImageInput], **kwargs) -> OnnxOutputContext:
def onnx_embed(self, images: list[ImageInput], **kwargs: Any) -> OnnxOutputContext:
with contextlib.ExitStack():
image_files = [
Image.open(image) if not isinstance(image, Image.Image) else image
@@ -84,7 +84,7 @@ class OnnxImageModel(OnnxModel[T]):
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[list[int]] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[T]:
is_small = False

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

@@ -114,11 +114,11 @@ def resize(
return image.resize(new_size, resample)
def rescale(image: np.ndarray, scale: float, dtype=np.float32) -> np.ndarray:
def rescale(image: np.ndarray, scale: float, dtype: type = np.float32) -> np.ndarray:
return (image * scale).astype(dtype)
def pil2ndarray(image: Union[Image.Image, np.ndarray]):
def pil2ndarray(image: Union[Image.Image, np.ndarray]) -> np.ndarray:
if isinstance(image, Image.Image):
return np.asarray(image).transpose((2, 0, 1))
return image

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

@@ -133,11 +133,11 @@ class Compose:
return cls(transforms=transforms)
@staticmethod
def _get_convert_to_rgb(transforms: list[Transform], config: dict[str, Any]):
def _get_convert_to_rgb(transforms: list[Transform], config: dict[str, Any]) -> None:
transforms.append(ConvertToRGB())
@classmethod
def _get_resize(cls, transforms: list[Transform], config: dict[str, Any]):
def _get_resize(cls, transforms: list[Transform], config: dict[str, Any]) -> None:
mode = config.get("image_processor_type", "CLIPImageProcessor")
if mode == "CLIPImageProcessor":
if config.get("do_resize", False):
@@ -200,7 +200,7 @@ class Compose:
raise ValueError(f"Preprocessor {mode} is not supported")
@staticmethod
def _get_center_crop(transforms: list[Transform], config: dict[str, Any]):
def _get_center_crop(transforms: list[Transform], config: dict[str, Any]) -> None:
mode = config.get("image_processor_type", "CLIPImageProcessor")
if mode == "CLIPImageProcessor":
if config.get("do_center_crop", False):
@@ -220,24 +220,24 @@ class Compose:
raise ValueError(f"Preprocessor {mode} is not supported")
@staticmethod
def _get_pil2ndarray(transforms: list[Transform], config: dict[str, Any]):
def _get_pil2ndarray(transforms: list[Transform], config: dict[str, Any]) -> None:
transforms.append(PILtoNDarray())
@staticmethod
def _get_rescale(transforms: list[Transform], config: dict[str, Any]):
def _get_rescale(transforms: list[Transform], config: dict[str, Any]) -> None:
if config.get("do_rescale", True):
rescale_factor = config.get("rescale_factor", 1 / 255)
transforms.append(Rescale(scale=rescale_factor))
@staticmethod
def _get_normalize(transforms: list[Transform], config: dict[str, Any]):
def _get_normalize(transforms: list[Transform], config: dict[str, Any]) -> None:
if config.get("do_normalize", False):
transforms.append(Normalize(mean=config["image_mean"], std=config["image_std"]))
elif "mean" in config and "std" in config:
transforms.append(Normalize(mean=config["mean"], std=config["std"]))
@staticmethod
def _get_pad2square(transforms: list[Transform], config: dict[str, Any]):
def _get_pad2square(transforms: list[Transform], config: dict[str, Any]) -> None:
mode = config.get("image_processor_type", "CLIPImageProcessor")
if mode == "CLIPImageProcessor":
pass

15
fastembed/late_interaction/colbert.py
View File

@@ -124,7 +124,8 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
device_ids: Optional[list[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
**kwargs,
specific_model_path: Optional[str] = None,
**kwargs: Any,
):
"""
Args:
@@ -142,6 +143,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
Raises:
ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
@@ -167,7 +169,10 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
self.cache_dir = define_cache_dir(cache_dir)
self._model_dir = self.download_model(
self.model_description, self.cache_dir, local_files_only=self._local_files_only
self.model_description,
self.cache_dir,
local_files_only=self._local_files_only,
specific_model_path=specific_model_path,
)
self.mask_token_id = None
self.pad_token_id = None
@@ -200,7 +205,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[np.ndarray]:
"""
Encode a list of documents into list of embeddings.
@@ -229,7 +234,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
**kwargs,
)
def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]:
def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]:
if isinstance(query, str):
query = [query]
@@ -247,7 +252,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]):
class ColbertEmbeddingWorker(TextEmbeddingWorker):
def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> Colbert:
def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> Colbert:
return Colbert(
model_name=model_name,
cache_dir=cache_dir,

2
fastembed/late_interaction/jina_colbert.py
View File

@@ -53,7 +53,7 @@ class JinaColbert(Colbert):
class JinaColbertEmbeddingWorker(TextEmbeddingWorker):
def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> JinaColbert:
def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> JinaColbert:
return JinaColbert(
model_name=model_name,
cache_dir=cache_dir,

12
fastembed/late_interaction/late_interaction_embedding_base.py
View File

@@ -1,4 +1,4 @@
from typing import Iterable, Optional, Union
from typing import Iterable, Optional, Union, Any
import numpy as np
@@ -11,7 +11,7 @@ class LateInteractionTextEmbeddingBase(ModelManagement):
model_name: str,
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
**kwargs,
**kwargs: Any,
):
self.model_name = model_name
self.cache_dir = cache_dir
@@ -23,11 +23,11 @@ class LateInteractionTextEmbeddingBase(ModelManagement):
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[np.ndarray]:
raise NotImplementedError()
def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[np.ndarray]:
def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[np.ndarray]:
"""
Embeds a list of text passages into a list of embeddings.
@@ -42,9 +42,7 @@ 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: Any) -> Iterable[np.ndarray]:
"""
Embeds queries

6
fastembed/late_interaction/late_interaction_text_embedding.py
View File

@@ -52,7 +52,7 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
cuda: bool = False,
device_ids: Optional[list[int]] = None,
lazy_load: bool = False,
**kwargs,
**kwargs: Any,
):
super().__init__(model_name, cache_dir, threads, **kwargs)
for EMBEDDING_MODEL_TYPE in self.EMBEDDINGS_REGISTRY:
@@ -80,7 +80,7 @@ class LateInteractionTextEmbedding(LateInteractionTextEmbeddingBase):
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[np.ndarray]:
"""
Encode a list of documents into list of embeddings.
@@ -99,7 +99,7 @@ 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: Any) -> Iterable[np.ndarray]:
"""
Embeds queries

5
fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py
View File

@@ -95,6 +95,7 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
device_ids: Optional[list[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
specific_model_path: Optional[str] = None,
**kwargs: Any,
):
"""
@@ -113,6 +114,7 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
Raises:
ValueError: If the model_name is not in the format <org>/<model> e.g. Xenova/ms-marco-MiniLM-L-6-v2.
@@ -145,6 +147,7 @@ class OnnxTextCrossEncoder(TextCrossEncoderBase, OnnxCrossEncoderModel):
self.model_description,
self.cache_dir,
local_files_only=self._local_files_only,
specific_model_path=specific_model_path,
)
if not self.lazy_load:
@@ -214,7 +217,7 @@ class TextCrossEncoderWorker(TextRerankerWorker):
self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
) -> OnnxTextCrossEncoder:
return OnnxTextCrossEncoder(
model_name=model_name,

7
fastembed/rerank/cross_encoder/onnx_text_model.py
View File

@@ -1,9 +1,10 @@
import os
from multiprocessing import get_all_start_methods
from pathlib import Path
from typing import Any, Iterable, Optional, Sequence, Type
from typing import Any, Iterable, Optional, Sequence, Type, Union
import numpy as np
from numpy.typing import NDArray
from tokenizers import Encoding
from fastembed.common.onnx_model import (
@@ -46,7 +47,9 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
def tokenize(self, pairs: list[tuple[str, str]], **_: Any) -> list[Encoding]:
return self.tokenizer.encode_batch(pairs)
def _build_onnx_input(self, tokenized_input):
def _build_onnx_input(
self, tokenized_input
) -> dict[str, NDArray[Union[np.float32, np.int64]]]:
input_names = {node.name for node in self.model.get_inputs()}
inputs = {
"input_ids": np.array([enc.ids for enc in tokenized_input], dtype=np.int64),

4
fastembed/rerank/cross_encoder/text_cross_encoder_base.py
View File

@@ -9,7 +9,7 @@ class TextCrossEncoderBase(ModelManagement):
model_name: str,
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
**kwargs,
**kwargs: Any,
):
self.model_name = model_name
self.cache_dir = cache_dir
@@ -21,7 +21,7 @@ class TextCrossEncoderBase(ModelManagement):
query: str,
documents: Iterable[str],
batch_size: int = 64,
**kwargs,
**kwargs: Any,
) -> Iterable[float]:
"""Rerank a list of documents given a query.

18
fastembed/sparse/bm25.py
View File

@@ -108,7 +108,8 @@ class Bm25(SparseTextEmbeddingBase):
language: str = "english",
token_max_length: int = 40,
disable_stemmer: bool = False,
**kwargs,
specific_model_path: Optional[str] = None,
**kwargs: Any,
):
super().__init__(model_name, cache_dir, **kwargs)
@@ -125,7 +126,10 @@ class Bm25(SparseTextEmbeddingBase):
self.cache_dir = define_cache_dir(cache_dir)
self._model_dir = self.download_model(
model_description, self.cache_dir, local_files_only=self._local_files_only
model_description,
self.cache_dir,
local_files_only=self._local_files_only,
specific_model_path=specific_model_path,
)
self.token_max_length = token_max_length
@@ -209,7 +213,7 @@ class Bm25(SparseTextEmbeddingBase):
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[SparseEmbedding]:
"""
Encode a list of documents into list of embeddings.
@@ -301,7 +305,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: Any
) -> 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.
"""
@@ -329,7 +335,7 @@ class Bm25Worker(Worker):
self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
):
self.model = self.init_embedding(model_name, cache_dir, **kwargs)
@@ -343,5 +349,5 @@ class Bm25Worker(Worker):
yield idx, onnx_output
@staticmethod
def init_embedding(model_name: str, cache_dir: str, **kwargs) -> Bm25:
def init_embedding(model_name: str, cache_dir: str, **kwargs: Any) -> Bm25:
return Bm25(model_name=model_name, cache_dir=cache_dir, **kwargs)

17
fastembed/sparse/bm42.py
View File

@@ -66,7 +66,8 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
device_ids: Optional[list[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
**kwargs,
specific_model_path: Optional[str] = None,
**kwargs: Any,
):
"""
Args:
@@ -86,6 +87,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
Raises:
ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
@@ -111,7 +113,10 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
self.cache_dir = define_cache_dir(cache_dir)
self._model_dir = self.download_model(
self.model_description, self.cache_dir, local_files_only=self._local_files_only
self.model_description,
self.cache_dir,
local_files_only=self._local_files_only,
specific_model_path=specific_model_path,
)
self.invert_vocab = {}
@@ -268,7 +273,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[SparseEmbedding]:
"""
Encode a list of documents into list of embeddings.
@@ -305,7 +310,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: Any
) -> 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.
@@ -332,7 +339,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
class Bm42TextEmbeddingWorker(TextEmbeddingWorker):
def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> Bm42:
def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> Bm42:
return Bm42(
model_name=model_name,
cache_dir=cache_dir,

12
fastembed/sparse/sparse_embedding_base.py
View File

@@ -1,5 +1,5 @@
from dataclasses import dataclass
from typing import Iterable, Optional, Union
from typing import Iterable, Optional, Union, Any
import numpy as np
@@ -34,7 +34,7 @@ class SparseTextEmbeddingBase(ModelManagement):
model_name: str,
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
**kwargs,
**kwargs: Any,
):
self.model_name = model_name
self.cache_dir = cache_dir
@@ -46,11 +46,11 @@ class SparseTextEmbeddingBase(ModelManagement):
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[SparseEmbedding]:
raise NotImplementedError()
def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[SparseEmbedding]:
def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[SparseEmbedding]:
"""
Embeds a list of text passages into a list of embeddings.
@@ -65,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: Any
) -> Iterable[SparseEmbedding]:
"""
Embeds queries

8
fastembed/sparse/sparse_text_embedding.py
View File

@@ -52,7 +52,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
cuda: bool = False,
device_ids: Optional[list[int]] = None,
lazy_load: bool = False,
**kwargs,
**kwargs: Any,
):
super().__init__(model_name, cache_dir, threads, **kwargs)
if model_name == "prithvida/Splade_PP_en_v1":
@@ -89,7 +89,7 @@ class SparseTextEmbedding(SparseTextEmbeddingBase):
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[SparseEmbedding]:
"""
Encode a list of documents into list of embeddings.
@@ -108,7 +108,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: Any
) -> Iterable[SparseEmbedding]:
"""
Embeds queries

13
fastembed/sparse/splade_pp.py
View File

@@ -73,7 +73,8 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
device_ids: Optional[list[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
**kwargs,
specific_model_path: Optional[str] = None,
**kwargs: Any,
):
"""
Args:
@@ -91,6 +92,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
Raises:
ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
@@ -115,7 +117,10 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
self.cache_dir = define_cache_dir(cache_dir)
self._model_dir = self.download_model(
self.model_description, self.cache_dir, local_files_only=self._local_files_only
self.model_description,
self.cache_dir,
local_files_only=self._local_files_only,
specific_model_path=specific_model_path,
)
if not self.lazy_load:
@@ -136,7 +141,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[SparseEmbedding]:
"""
Encode a list of documents into list of embeddings.
@@ -171,7 +176,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]):
class SpladePPEmbeddingWorker(TextEmbeddingWorker):
def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> SpladePP:
def init_embedding(self, model_name: str, cache_dir: str, **kwargs: Any) -> SpladePP:
return SpladePP(
model_name=model_name,
cache_dir=cache_dir,

2
fastembed/text/clip_embedding.py
View File

@@ -44,7 +44,7 @@ class CLIPEmbeddingWorker(OnnxTextEmbeddingWorker):
self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
) -> OnnxTextEmbedding:
return CLIPOnnxEmbedding(
model_name=model_name,

72
fastembed/text/e5_onnx_embedding.py
View File

@@ -1,72 +0,0 @@
from typing import Any, Type
import numpy as np
from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
from fastembed.text.onnx_text_model import TextEmbeddingWorker
supported_multilingual_e5_models = [
{
"model": "intfloat/multilingual-e5-large",
"dim": 1024,
"description": "Text embeddings, Unimodal (text), Multilingual (~100 languages), 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
"license": "mit",
"size_in_GB": 2.24,
"sources": {
"url": "https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz",
"hf": "qdrant/multilingual-e5-large-onnx",
},
"model_file": "model.onnx",
"additional_files": ["model.onnx_data"],
},
{
"model": "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
"dim": 768,
"description": "Text embeddings, Unimodal (text), Multilingual (~50 languages), 384 input tokens truncation, Prefixes for queries/documents: not necessary, 2021 year.",
"license": "apache-2.0",
"size_in_GB": 1.00,
"sources": {
"hf": "xenova/paraphrase-multilingual-mpnet-base-v2",
},
"model_file": "onnx/model.onnx",
},
]
class E5OnnxEmbedding(OnnxTextEmbedding):
@classmethod
def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
return E5OnnxEmbeddingWorker
@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_multilingual_e5_models
def _preprocess_onnx_input(
self, onnx_input: dict[str, np.ndarray], **kwargs
) -> dict[str, np.ndarray]:
"""
Preprocess the onnx input.
"""
onnx_input.pop("token_type_ids", None)
return onnx_input
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,
)

98
fastembed/text/multitask_embedding.py Normal file
View File

@@ -0,0 +1,98 @@
from enum import Enum
from typing import Any, Type, Iterable, Union, Optional
import numpy as np
from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
from fastembed.text.onnx_embedding import OnnxTextEmbeddingWorker
from fastembed.text.onnx_text_model import TextEmbeddingWorker
supported_multitask_models = [
{
"model": "jinaai/jina-embeddings-v3",
"dim": 1024,
"tasks": {
"retrieval.query": 0,
"retrieval.passage": 1,
"separation": 2,
"classification": 3,
"text-matching": 4,
},
"description": "Multi-task unimodal (text) embedding model, multi-lingual (~100), 1024 tokens truncation, and 8192 sequence length. Prefixes for queries/documents: not necessary, 2024 year.",
"license": "cc-by-nc-4.0",
"size_in_GB": 2.29,
"sources": {
"hf": "jinaai/jina-embeddings-v3",
},
"model_file": "onnx/model.onnx",
"additional_files": ["onnx/model.onnx_data"],
},
]
class Task(int, Enum):
RETRIEVAL_QUERY = 0
RETRIEVAL_PASSAGE = 1
SEPARATION = 2
CLASSIFICATION = 3
TEXT_MATCHING = 4
class JinaEmbeddingV3(PooledNormalizedEmbedding):
PASSAGE_TASK = Task.RETRIEVAL_PASSAGE
QUERY_TASK = Task.RETRIEVAL_QUERY
def __init__(self, *args: Any, **kwargs: Any):
super().__init__(*args, **kwargs)
self._current_task_id = self.PASSAGE_TASK
@classmethod
def _get_worker_class(cls) -> Type["TextEmbeddingWorker"]:
return JinaEmbeddingV3Worker
@classmethod
def list_supported_models(cls) -> list[dict[str, Any]]:
return supported_multitask_models
def _preprocess_onnx_input(
self, onnx_input: dict[str, np.ndarray], **kwargs
) -> dict[str, np.ndarray]:
onnx_input["task_id"] = np.array(self._current_task_id, dtype=np.int64)
return onnx_input
def embed(
self,
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
task_id: int = PASSAGE_TASK,
**kwargs,
) -> Iterable[np.ndarray]:
self._current_task_id = task_id
kwargs["task_id"] = task_id
yield from super().embed(documents, batch_size, parallel, **kwargs)
def query_embed(self, query: Union[str, Iterable[str]], **kwargs) -> Iterable[np.ndarray]:
self._current_task_id = self.QUERY_TASK
yield from super().embed(query, **kwargs)
def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[np.ndarray]:
self._current_task_id = self.PASSAGE_TASK
yield from super().embed(texts, **kwargs)
class JinaEmbeddingV3Worker(OnnxTextEmbeddingWorker):
def init_embedding(
self,
model_name: str,
cache_dir: str,
**kwargs,
) -> JinaEmbeddingV3:
model = JinaEmbeddingV3(
model_name=model_name,
cache_dir=cache_dir,
threads=1,
**kwargs,
)
model._current_task_id = kwargs["task_id"]
return model

26
fastembed/text/onnx_embedding.py
View File

@@ -79,17 +79,6 @@ supported_onnx_models = [
},
"model_file": "model_optimized.onnx",
},
{
"model": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
"dim": 384,
"description": "Text embeddings, Unimodal (text), Multilingual (~50 languages), 512 input tokens truncation, Prefixes for queries/documents: not necessary, 2019 year.",
"license": "apache-2.0",
"size_in_GB": 0.22,
"sources": {
"hf": "qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q",
},
"model_file": "model_optimized.onnx",
},
{
"model": "thenlper/gte-large",
"dim": 1024,
@@ -204,7 +193,8 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
device_ids: Optional[list[int]] = None,
lazy_load: bool = False,
device_id: Optional[int] = None,
**kwargs,
specific_model_path: Optional[str] = None,
**kwargs: Any,
):
"""
Args:
@@ -222,6 +212,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
device_id (Optional[int], optional): The device id to use for loading the model in the worker process.
specific_model_path (Optional[str], optional): The specific path to the onnx model dir if it should be imported from somewhere else
Raises:
ValueError: If the model_name is not in the format <org>/<model> e.g. BAAI/bge-base-en.
@@ -245,7 +236,10 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
self.model_description = self._get_model_description(model_name)
self.cache_dir = define_cache_dir(cache_dir)
self._model_dir = self.download_model(
self.model_description, self.cache_dir, local_files_only=self._local_files_only
self.model_description,
self.cache_dir,
local_files_only=self._local_files_only,
specific_model_path=specific_model_path,
)
if not self.lazy_load:
@@ -256,7 +250,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[np.ndarray]:
"""
Encode a list of documents into list of embeddings.
@@ -290,7 +284,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]):
return OnnxTextEmbeddingWorker
def _preprocess_onnx_input(
self, onnx_input: dict[str, np.ndarray], **kwargs
self, onnx_input: dict[str, np.ndarray], **kwargs: Any
) -> dict[str, np.ndarray]:
"""
Preprocess the onnx input.
@@ -323,7 +317,7 @@ class OnnxTextEmbeddingWorker(TextEmbeddingWorker):
self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
) -> OnnxTextEmbedding:
return OnnxTextEmbedding(
model_name=model_name,

11
fastembed/text/onnx_text_model.py
View File

@@ -23,13 +23,13 @@ class OnnxTextModel(OnnxModel[T]):
def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[T]:
raise NotImplementedError("Subclasses must implement this method")
def __init__(self) -> None:
def __init__(self):
super().__init__()
self.tokenizer = None
self.special_token_to_id = {}
def _preprocess_onnx_input(
self, onnx_input: dict[str, np.ndarray], **kwargs
self, onnx_input: dict[str, np.ndarray], **kwargs: Any
) -> dict[str, np.ndarray]:
"""
Preprocess the onnx input.
@@ -44,6 +44,7 @@ class OnnxTextModel(OnnxModel[T]):
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_id: Optional[int] = None,
specific_model_path: Optional[str] = None,
) -> None:
super()._load_onnx_model(
model_dir=model_dir,
@@ -58,13 +59,13 @@ class OnnxTextModel(OnnxModel[T]):
def load_onnx_model(self) -> None:
raise NotImplementedError("Subclasses must implement this method")
def tokenize(self, documents: list[str], **kwargs) -> list[Encoding]:
def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
return self.tokenizer.encode_batch(documents)
def onnx_embed(
self,
documents: list[str],
**kwargs,
**kwargs: Any,
) -> OnnxOutputContext:
encoded = self.tokenize(documents, **kwargs)
input_ids = np.array([e.ids for e in encoded])
@@ -98,7 +99,7 @@ class OnnxTextModel(OnnxModel[T]):
providers: Optional[Sequence[OnnxProvider]] = None,
cuda: bool = False,
device_ids: Optional[list[int]] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[T]:
is_small = False

37
fastembed/text/pooled_embedding.py
View File

@@ -40,6 +40,41 @@ supported_pooled_models = [
},
"model_file": "onnx/model.onnx",
},
{
"model": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
"dim": 384,
"description": "Text embeddings, Unimodal (text), Multilingual (~50 languages), 512 input tokens truncation, Prefixes for queries/documents: not necessary, 2019 year.",
"license": "apache-2.0",
"size_in_GB": 0.22,
"sources": {
"hf": "qdrant/paraphrase-multilingual-MiniLM-L12-v2-onnx-Q",
},
"model_file": "model_optimized.onnx",
},
{
"model": "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
"dim": 768,
"description": "Text embeddings, Unimodal (text), Multilingual (~50 languages), 384 input tokens truncation, Prefixes for queries/documents: not necessary, 2021 year.",
"license": "apache-2.0",
"size_in_GB": 1.00,
"sources": {
"hf": "xenova/paraphrase-multilingual-mpnet-base-v2",
},
"model_file": "onnx/model.onnx",
},
{
"model": "intfloat/multilingual-e5-large",
"dim": 1024,
"description": "Text embeddings, Unimodal (text), Multilingual (~100 languages), 512 input tokens truncation, Prefixes for queries/documents: necessary, 2024 year.",
"license": "mit",
"size_in_GB": 2.24,
"sources": {
"url": "https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz",
"hf": "qdrant/multilingual-e5-large-onnx",
},
"model_file": "model.onnx",
"additional_files": ["model.onnx_data"],
},
]
@@ -82,7 +117,7 @@ class PooledEmbeddingWorker(OnnxTextEmbeddingWorker):
self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
) -> OnnxTextEmbedding:
return PooledEmbedding(
model_name=model_name,

2
fastembed/text/pooled_normalized_embedding.py
View File

@@ -115,7 +115,7 @@ class PooledNormalizedEmbeddingWorker(OnnxTextEmbeddingWorker):
self,
model_name: str,
cache_dir: str,
**kwargs,
**kwargs: Any,
) -> OnnxTextEmbedding:
return PooledNormalizedEmbedding(
model_name=model_name,

62
fastembed/text/text_embedding.py
View File

@@ -1,12 +1,12 @@
import warnings
from typing import Any, Iterable, Optional, Sequence, Type, Union
import numpy as np
from fastembed.common import OnnxProvider
from fastembed.text.clip_embedding import CLIPOnnxEmbedding
from fastembed.text.e5_onnx_embedding import E5OnnxEmbedding
from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
from fastembed.text.pooled_embedding import PooledEmbedding
from fastembed.text.multitask_embedding import JinaEmbeddingV3
from fastembed.text.onnx_embedding import OnnxTextEmbedding
from fastembed.text.text_embedding_base import TextEmbeddingBase
@@ -14,10 +14,10 @@ from fastembed.text.text_embedding_base import TextEmbeddingBase
class TextEmbedding(TextEmbeddingBase):
EMBEDDINGS_REGISTRY: list[Type[TextEmbeddingBase]] = [
OnnxTextEmbedding,
E5OnnxEmbedding,
CLIPOnnxEmbedding,
PooledNormalizedEmbedding,
PooledEmbedding,
JinaEmbeddingV3,
]
@classmethod
@@ -59,9 +59,34 @@ class TextEmbedding(TextEmbeddingBase):
cuda: bool = False,
device_ids: Optional[list[int]] = None,
lazy_load: bool = False,
**kwargs,
**kwargs: Any,
):
super().__init__(model_name, cache_dir, threads, **kwargs)
if model_name == "nomic-ai/nomic-embed-text-v1.5-Q":
warnings.warn(
"The model 'nomic-ai/nomic-embed-text-v1.5-Q' has been updated on HuggingFace. "
"Please review the latest documentation and release notes to ensure compatibility with your workflow. ",
UserWarning,
stacklevel=2,
)
if model_name == "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2":
warnings.warn(
"The model 'sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2' has been updated to "
"include a mean pooling layer. Please ensure your usage aligns with the new functionality. "
"Support for the previous version without mean pooling will be removed as of version 0.5.2.",
UserWarning,
stacklevel=2,
)
if model_name in {
"sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
"intfloat/multilingual-e5-large",
}:
warnings.warn(
f"{model_name} has been updated as of fastembed 0.5.2, outputs are now average pooled.",
UserWarning,
stacklevel=2,
)
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):
@@ -87,7 +112,7 @@ class TextEmbedding(TextEmbeddingBase):
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[np.ndarray]:
"""
Encode a list of documents into list of embeddings.
@@ -105,3 +130,30 @@ class TextEmbedding(TextEmbeddingBase):
List of embeddings, one per document
"""
yield from self.model.embed(documents, batch_size, parallel, **kwargs)
def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]:
"""
Embeds queries
Args:
query (Union[str, Iterable[str]]): The query to embed, or an iterable e.g. list of queries.
Returns:
Iterable[np.ndarray]: The embeddings.
"""
# This is model-specific, so that different models can have specialized implementations
yield from self.model.query_embed(query, **kwargs)
def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[np.ndarray]:
"""
Embeds a list of text passages into a list of embeddings.
Args:
texts (Iterable[str]): The list of texts to embed.
**kwargs: Additional keyword argument to pass to the embed method.
Yields:
Iterable[SparseEmbedding]: The sparse embeddings.
"""
# This is model-specific, so that different models can have specialized implementations
yield from self.model.passage_embed(texts, **kwargs)

13
fastembed/text/text_embedding_base.py
View File

@@ -1,4 +1,4 @@
from typing import Iterable, Optional, Union
from typing import Iterable, Optional, Union, Any
import numpy as np
@@ -11,7 +11,7 @@ class TextEmbeddingBase(ModelManagement):
model_name: str,
cache_dir: Optional[str] = None,
threads: Optional[int] = None,
**kwargs,
**kwargs: Any,
):
self.model_name = model_name
self.cache_dir = cache_dir
@@ -23,11 +23,11 @@ class TextEmbeddingBase(ModelManagement):
documents: Union[str, Iterable[str]],
batch_size: int = 256,
parallel: Optional[int] = None,
**kwargs,
**kwargs: Any,
) -> Iterable[np.ndarray]:
raise NotImplementedError()
def passage_embed(self, texts: Iterable[str], **kwargs) -> Iterable[np.ndarray]:
def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[np.ndarray]:
"""
Embeds a list of text passages into a list of embeddings.
@@ -42,9 +42,8 @@ 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: Any) -> Iterable[np.ndarray]:
"""
Embeds queries

14
pyproject.toml
View File

@@ -1,6 +1,6 @@
[tool.poetry]
name = "fastembed"
version = "0.5.0"
version = "0.5.1"
description = "Fast, light, accurate library built for retrieval embedding generation"
authors = ["Qdrant Team <info@qdrant.tech>", "NirantK <nirant.bits@gmail.com>"]
license = "Apache License"
@@ -12,7 +12,6 @@ keywords = ["vector", "embedding", "neural", "search", "qdrant", "sentence-trans
[tool.poetry.dependencies]
python = ">=3.9.0"
onnx = ">=1.15.0"
numpy = [
{ version = ">=1.21,<2.1.0", python = "<3.10" },
{ version = ">=1.21", python = ">=3.10,<3.12" },
@@ -33,11 +32,14 @@ pillow = "^10.3.0"
mmh3 = "^4.1.0"
py-rust-stemmers = "^0.1.0"
[tool.poetry.group.dev.dependencies]
[tool.poetry.group.test.dependencies]
pytest = "^7.4.2"
ruff = ">=0.3.1,<1.0"
[tool.poetry.group.dev.dependencies]
notebook = ">=7.0.2"
pre-commit = "^3.6.2"
onnx = ">=1.15.0"
[tool.poetry.group.docs.dependencies]
mkdocs-material = "^9.5.10"
@@ -46,10 +48,16 @@ pillow = "^10.2.0"
cairosvg = "^2.7.1"
mknotebooks = "^0.8.0"
[tool.poetry.group.types.dependencies]
pyright = ">=1.1.293"
mypy = "^1.0.0"
[build-system]
requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"
[tool.pyright]
typeCheckingMode = "strict"
[tool.ruff]
line-length = 99

10
tests/test_attention_embeddings.py
View File

@@ -8,7 +8,7 @@ from tests.utils import delete_model_cache
@pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"])
def test_attention_embeddings(model_name):
def test_attention_embeddings(model_name: str) -> None:
is_ci = os.getenv("CI")
model = SparseTextEmbedding(model_name=model_name)
@@ -71,7 +71,7 @@ def test_attention_embeddings(model_name):
@pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"])
def test_parallel_processing(model_name):
def test_parallel_processing(model_name: str) -> None:
is_ci = os.getenv("CI")
model = SparseTextEmbedding(model_name=model_name)
@@ -96,7 +96,7 @@ def test_parallel_processing(model_name):
@pytest.mark.parametrize("model_name", ["Qdrant/bm25"])
def test_multilanguage(model_name):
def test_multilanguage(model_name: str) -> None:
is_ci = os.getenv("CI")
docs = ["Mangez-vous vraiment des grenouilles?", "Je suis au lit"]
@@ -122,7 +122,7 @@ def test_multilanguage(model_name):
@pytest.mark.parametrize("model_name", ["Qdrant/bm25"])
def test_special_characters(model_name):
def test_special_characters(model_name: str) -> None:
is_ci = os.getenv("CI")
docs = [
@@ -145,7 +145,7 @@ def test_special_characters(model_name):
@pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions"])
def test_lazy_load(model_name):
def test_lazy_load(model_name: str) -> None:
model = SparseTextEmbedding(model_name=model_name, lazy_load=True)
assert not hasattr(model.model, "model")
docs = ["hello world", "flag embedding"]

8
tests/test_image_onnx_embeddings.py
View File

@@ -27,7 +27,7 @@ CANONICAL_VECTOR_VALUES = {
}
def test_embedding():
def test_embedding() -> None:
is_ci = os.getenv("CI")
for model_desc in ImageEmbedding.list_supported_models():
@@ -61,7 +61,7 @@ def test_embedding():
@pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")])
def test_batch_embedding(n_dims, model_name):
def test_batch_embedding(n_dims: int, model_name: str) -> None:
is_ci = os.getenv("CI")
model = ImageEmbedding(model_name=model_name)
n_images = 32
@@ -81,7 +81,7 @@ def test_batch_embedding(n_dims, model_name):
@pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")])
def test_parallel_processing(n_dims, model_name):
def test_parallel_processing(n_dims: int, model_name: str) -> None:
is_ci = os.getenv("CI")
model = ImageEmbedding(model_name=model_name)
@@ -109,7 +109,7 @@ def test_parallel_processing(n_dims, model_name):
@pytest.mark.parametrize("model_name", ["Qdrant/clip-ViT-B-32-vision"])
def test_lazy_load(model_name):
def test_lazy_load(model_name: str) -> None:
is_ci = os.getenv("CI")
model = ImageEmbedding(model_name=model_name, lazy_load=True)
assert not hasattr(model.model, "model")

2
tests/test_late_interaction_embeddings.py
View File

@@ -226,7 +226,7 @@ def test_parallel_processing():
"model_name",
["colbert-ir/colbertv2.0"],
)
def test_lazy_load(model_name):
def test_lazy_load(model_name: str):
is_ci = os.getenv("CI")
model = LateInteractionTextEmbedding(model_name=model_name, lazy_load=True)

7
tests/test_multi_gpu.py
View File

@@ -1,4 +1,5 @@
import pytest
from typing import Optional
from fastembed import (
TextEmbedding,
SparseTextEmbedding,
@@ -13,7 +14,7 @@ CACHE_DIR = "../model_cache"
@pytest.mark.skip(reason="Requires a multi-gpu server")
@pytest.mark.parametrize("device_id", [None, 0, 1])
def test_gpu_via_providers(device_id):
def test_gpu_via_providers(device_id: Optional[int]) -> None:
docs = ["hello world", "flag embedding"]
device_id = device_id if device_id is not None else 0
@@ -85,7 +86,7 @@ def test_gpu_via_providers(device_id):
@pytest.mark.skip(reason="Requires a multi-gpu server")
@pytest.mark.parametrize("device_ids", [None, [0], [1], [0, 1]])
def test_gpu_cuda_device_ids(device_ids):
def test_gpu_cuda_device_ids(device_ids: Optional[list[int]]) -> None:
docs = ["hello world", "flag embedding"]
device_id = device_ids[0] if device_ids else 0
embedding_model = TextEmbedding(
@@ -170,7 +171,7 @@ def test_gpu_cuda_device_ids(device_ids):
@pytest.mark.parametrize(
"device_ids,parallel", [(None, None), (None, 2), ([1], None), ([1], 1), ([1], 2), ([0, 1], 2)]
)
def test_multi_gpu_parallel_inference(device_ids, parallel):
def test_multi_gpu_parallel_inference(device_ids: Optional[list[int]], parallel: int) -> None:
docs = ["hello world", "flag embedding"] * 100
batch_size = 5

16
tests/test_sparse_embeddings.py
View File

@@ -49,7 +49,7 @@ CANONICAL_COLUMN_VALUES = {
docs = ["Hello World"]
def test_batch_embedding():
def test_batch_embedding() -> None:
is_ci = os.getenv("CI")
docs_to_embed = docs * 10
@@ -64,7 +64,7 @@ def test_batch_embedding():
delete_model_cache(model.model._model_dir)
def test_single_embedding():
def test_single_embedding() -> None:
is_ci = os.getenv("CI")
for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
model = SparseTextEmbedding(model_name=model_name)
@@ -80,7 +80,7 @@ def test_single_embedding():
delete_model_cache(model.model._model_dir)
def test_parallel_processing():
def test_parallel_processing() -> None:
is_ci = os.getenv("CI")
model = SparseTextEmbedding(model_name="prithivida/Splade_PP_en_v1")
docs = ["hello world", "flag embedding"] * 30
@@ -111,7 +111,7 @@ def test_parallel_processing():
@pytest.fixture
def bm25_instance():
def bm25_instance() -> None:
ci = os.getenv("CI", True)
model = Bm25("Qdrant/bm25", language="english")
yield model
@@ -119,7 +119,7 @@ def bm25_instance():
delete_model_cache(model._model_dir)
def test_stem_with_stopwords_and_punctuation(bm25_instance):
def test_stem_with_stopwords_and_punctuation(bm25_instance: Bm25) -> None:
# Setup
bm25_instance.stopwords = {"the", "is", "a"}
bm25_instance.punctuation = {".", ",", "!"}
@@ -135,7 +135,7 @@ def test_stem_with_stopwords_and_punctuation(bm25_instance):
assert result == expected, f"Expected {expected}, but got {result}"
def test_stem_case_insensitive_stopwords(bm25_instance):
def test_stem_case_insensitive_stopwords(bm25_instance: Bm25) -> None:
# Setup
bm25_instance.stopwords = {"the", "is", "a"}
bm25_instance.punctuation = {".", ",", "!"}
@@ -152,7 +152,7 @@ def test_stem_case_insensitive_stopwords(bm25_instance):
@pytest.mark.parametrize("disable_stemmer", [True, False])
def test_disable_stemmer_behavior(disable_stemmer):
def test_disable_stemmer_behavior(disable_stemmer: bool) -> None:
# Setup
model = Bm25("Qdrant/bm25", language="english", disable_stemmer=disable_stemmer)
model.stopwords = {"the", "is", "a"}
@@ -176,7 +176,7 @@ def test_disable_stemmer_behavior(disable_stemmer):
"model_name",
["prithivida/Splade_PP_en_v1"],
)
def test_lazy_load(model_name):
def test_lazy_load(model_name: str) -> None:
is_ci = os.getenv("CI")
model = SparseTextEmbedding(model_name=model_name, lazy_load=True)
assert not hasattr(model.model, "model")

8
tests/test_text_cross_encoder.py
View File

@@ -26,7 +26,7 @@ SELECTED_MODELS = {
"model_name",
[model_name for model_name in CANONICAL_SCORE_VALUES],
)
def test_rerank(model_name):
def test_rerank(model_name: str) -> None:
is_ci = os.getenv("CI")
model = TextCrossEncoder(model_name=model_name)
@@ -53,7 +53,7 @@ def test_rerank(model_name):
"model_name",
[model_name for model_name in SELECTED_MODELS.values()],
)
def test_batch_rerank(model_name):
def test_batch_rerank(model_name: str) -> None:
is_ci = os.getenv("CI")
model = TextCrossEncoder(model_name=model_name)
@@ -82,7 +82,7 @@ def test_batch_rerank(model_name):
"model_name",
["Xenova/ms-marco-MiniLM-L-6-v2"],
)
def test_lazy_load(model_name):
def test_lazy_load(model_name: str) -> None:
is_ci = os.getenv("CI")
model = TextCrossEncoder(model_name=model_name, lazy_load=True)
assert not hasattr(model.model, "model")
@@ -99,7 +99,7 @@ def test_lazy_load(model_name):
"model_name",
[model_name for model_name in SELECTED_MODELS.values()],
)
def test_rerank_pairs_parallel(model_name):
def test_rerank_pairs_parallel(model_name: str) -> None:
is_ci = os.getenv("CI")
model = TextCrossEncoder(model_name=model_name)

253
tests/test_text_multitask_embeddings.py Normal file
View File

@@ -0,0 +1,253 @@
import os
import numpy as np
import pytest
from fastembed import TextEmbedding
from fastembed.text.multitask_embedding import Task
from tests.utils import delete_model_cache
CANONICAL_VECTOR_VALUES = {
"jinaai/jina-embeddings-v3": [
{
"task_id": Task.RETRIEVAL_QUERY,
"vectors": np.array(
[
[0.0623, -0.0402, 0.1706, -0.0143, 0.0617],
[-0.1064, -0.0733, 0.0353, 0.0096, 0.0667],
]
),
},
{
"task_id": Task.RETRIEVAL_PASSAGE,
"vectors": np.array(
[
[0.0513, -0.0247, 0.1751, -0.0075, 0.0679],
[-0.0987, -0.0786, 0.09, 0.0087, 0.0577],
]
),
},
{
"task_id": Task.SEPARATION,
"vectors": np.array(
[
[0.094, -0.1065, 0.1305, 0.0547, 0.0556],
[0.0315, -0.1468, 0.065, 0.0568, 0.0546],
]
),
},
{
"task_id": Task.CLASSIFICATION,
"vectors": np.array(
[
[0.0606, -0.0877, 0.1384, 0.0065, 0.0722],
[-0.0502, -0.119, 0.032, 0.0514, 0.0689],
]
),
},
{
"task_id": Task.TEXT_MATCHING,
"vectors": np.array(
[
[0.0911, -0.0341, 0.1305, -0.026, 0.0576],
[-0.1432, -0.05, 0.0133, 0.0464, 0.0789],
]
),
},
]
}
docs = ["Hello World", "Follow the white rabbit."]
def test_batch_embedding():
is_ci = os.getenv("CI")
docs_to_embed = docs * 10
default_task = Task.RETRIEVAL_PASSAGE
for model_desc in TextEmbedding.list_supported_models():
if not is_ci and model_desc["size_in_GB"] > 1:
continue
model_name = model_desc["model"]
dim = model_desc["dim"]
if model_name not in CANONICAL_VECTOR_VALUES.keys():
continue
model = TextEmbedding(model_name=model_name)
print(f"evaluating {model_name} default task")
embeddings = list(model.embed(documents=docs_to_embed, batch_size=6))
embeddings = np.stack(embeddings, axis=0)
assert embeddings.shape == (len(docs_to_embed), dim)
canonical_vector = CANONICAL_VECTOR_VALUES[model_name][default_task]["vectors"]
assert np.allclose(
embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
), model_desc["model"]
if is_ci:
delete_model_cache(model.model._model_dir)
def test_single_embedding():
is_ci = os.getenv("CI")
for model_desc in TextEmbedding.list_supported_models():
if not is_ci and model_desc["size_in_GB"] > 1:
continue
model_name = model_desc["model"]
dim = model_desc["dim"]
if model_name not in CANONICAL_VECTOR_VALUES.keys():
continue
model = TextEmbedding(model_name=model_name)
for task in CANONICAL_VECTOR_VALUES[model_name]:
print(f"evaluating {model_name} task_id: {task['task_id']}")
embeddings = list(model.embed(documents=docs, task_id=task["task_id"]))
embeddings = np.stack(embeddings, axis=0)
assert embeddings.shape == (len(docs), dim)
canonical_vector = task["vectors"]
assert np.allclose(
embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
), model_desc["model"]
if is_ci:
delete_model_cache(model.model._model_dir)
def test_single_embedding_query():
is_ci = os.getenv("CI")
task_id = Task.RETRIEVAL_QUERY
for model_desc in TextEmbedding.list_supported_models():
if not is_ci and model_desc["size_in_GB"] > 1:
continue
model_name = model_desc["model"]
dim = model_desc["dim"]
if model_name not in CANONICAL_VECTOR_VALUES.keys():
continue
model = TextEmbedding(model_name=model_name)
print(f"evaluating {model_name} query_embed task_id: {task_id}")
embeddings = list(model.query_embed(query=docs))
embeddings = np.stack(embeddings, axis=0)
assert embeddings.shape == (len(docs), dim)
canonical_vector = CANONICAL_VECTOR_VALUES[model_name][task_id]["vectors"]
assert np.allclose(
embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
), model_desc["model"]
if is_ci:
delete_model_cache(model.model._model_dir)
def test_single_embedding_passage():
is_ci = os.getenv("CI")
task_id = Task.RETRIEVAL_PASSAGE
for model_desc in TextEmbedding.list_supported_models():
if not is_ci and model_desc["size_in_GB"] > 1:
continue
model_name = model_desc["model"]
dim = model_desc["dim"]
if model_name not in CANONICAL_VECTOR_VALUES.keys():
continue
model = TextEmbedding(model_name=model_name)
print(f"evaluating {model_name} passage_embed task_id: {task_id}")
embeddings = list(model.passage_embed(texts=docs))
embeddings = np.stack(embeddings, axis=0)
assert embeddings.shape == (len(docs), dim)
canonical_vector = CANONICAL_VECTOR_VALUES[model_name][task_id]["vectors"]
assert np.allclose(
embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
), model_desc["model"]
if is_ci:
delete_model_cache(model.model._model_dir)
def test_parallel_processing():
is_ci = os.getenv("CI")
docs = ["Hello World", "Follow the white rabbit."] * 10
model_name = "jinaai/jina-embeddings-v3"
dim = 1024
model = TextEmbedding(model_name=model_name)
task_id = Task.SEPARATION
embeddings_1 = list(model.embed(docs, batch_size=10, parallel=None, task_id=task_id))
embeddings_1 = np.stack(embeddings_1, axis=0)
embeddings_2 = list(model.embed(docs, batch_size=10, parallel=1, task_id=task_id))
embeddings_2 = np.stack(embeddings_2, axis=0)
assert embeddings_1.shape[0] == len(docs) and embeddings_1.shape[-1] == dim
assert np.allclose(embeddings_1, embeddings_2, atol=1e-4)
canonical_vector = CANONICAL_VECTOR_VALUES[model_name][task_id]["vectors"]
assert np.allclose(embeddings_2[:2, : canonical_vector.shape[1]], canonical_vector, atol=1e-4)
if is_ci:
delete_model_cache(model.model._model_dir)
def test_task_assignment():
is_ci = os.getenv("CI")
for model_desc in TextEmbedding.list_supported_models():
if not is_ci and model_desc["size_in_GB"] > 1:
continue
model_name = model_desc["model"]
if model_name not in CANONICAL_VECTOR_VALUES.keys():
continue
model = TextEmbedding(model_name=model_name)
for i, task_id in enumerate(Task):
_ = list(model.embed(documents=docs, batch_size=1, task_id=i))
assert model.model._current_task_id == task_id
if is_ci:
delete_model_cache(model.model._model_dir)
@pytest.mark.parametrize(
"model_name",
["jinaai/jina-embeddings-v3"],
)
def test_lazy_load(model_name: str):
is_ci = os.getenv("CI")
model = TextEmbedding(model_name=model_name, lazy_load=True)
assert not hasattr(model.model, "model")
list(model.embed(docs))
assert hasattr(model.model, "model")
if is_ci:
delete_model_cache(model.model._model_dir)

26
tests/test_text_onnx_embeddings.py
View File

@@ -1,4 +1,5 @@
import os
import platform
import numpy as np
import pytest
@@ -31,11 +32,11 @@ CANONICAL_VECTOR_VALUES = {
[-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]
[0.0361, 0.1862, 0.2776, 0.2461, -0.1904]
),
"intfloat/multilingual-e5-large": np.array([0.0098, 0.0045, 0.0066, -0.0354, 0.0070]),
"intfloat/multilingual-e5-large": np.array([0.4544, -0.0968, 0.1054, -1.3753, 0.1500]),
"sentence-transformers/paraphrase-multilingual-mpnet-base-v2": np.array(
[-0.01341097, 0.0416553, -0.00480805, 0.02844842, 0.0505299]
[0.0047, 0.1334, -0.0102, 0.0714, 0.1930]
),
"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]),
@@ -48,7 +49,7 @@ CANONICAL_VECTOR_VALUES = {
[-0.15407836, -0.03053198, -3.9138033, 0.1910364, 0.13224715]
),
"nomic-ai/nomic-embed-text-v1.5-Q": np.array(
[-0.12525563, 0.38030425, -3.961622, 0.04176439, -0.0758301]
[0.0802303, 0.3700881, -4.3053818, 0.4431803, -0.271572]
),
"thenlper/gte-large": np.array(
[-0.01920587, 0.00113156, -0.00708992, -0.00632304, -0.04025577]
@@ -68,12 +69,19 @@ CANONICAL_VECTOR_VALUES = {
"jinaai/jina-clip-v1": np.array([-0.0862, -0.0101, -0.0056, 0.0375, -0.0472]),
}
MULTI_TASK_MODELS = ["jinaai/jina-embeddings-v3"]
def test_embedding():
def test_embedding() -> None:
is_ci = os.getenv("CI")
is_mac = platform.system() == "Darwin"
for model_desc in TextEmbedding.list_supported_models():
if not is_ci and model_desc["size_in_GB"] > 1:
if (
(not is_ci and model_desc["size_in_GB"] > 1)
or model_desc["model"] in MULTI_TASK_MODELS
or (is_mac and model_desc["model"] == "nomic-ai/nomic-embed-text-v1.5-Q")
):
continue
dim = model_desc["dim"]
@@ -96,7 +104,7 @@ def test_embedding():
"n_dims,model_name",
[(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")],
)
def test_batch_embedding(n_dims, model_name):
def test_batch_embedding(n_dims: int, model_name: str) -> None:
is_ci = os.getenv("CI")
model = TextEmbedding(model_name=model_name)
@@ -113,7 +121,7 @@ def test_batch_embedding(n_dims, model_name):
"n_dims,model_name",
[(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")],
)
def test_parallel_processing(n_dims, model_name):
def test_parallel_processing(n_dims: int, model_name: str) -> None:
is_ci = os.getenv("CI")
model = TextEmbedding(model_name=model_name)
@@ -139,7 +147,7 @@ def test_parallel_processing(n_dims, model_name):
"model_name",
["BAAI/bge-small-en-v1.5"],
)
def test_lazy_load(model_name):
def test_lazy_load(model_name: str) -> None:
is_ci = os.getenv("CI")
model = TextEmbedding(model_name=model_name, lazy_load=True)
assert not hasattr(model.model, "model")

9
tests/utils.py
View File

@@ -2,7 +2,8 @@ import shutil
import traceback
from pathlib import Path
from typing import Union
from types import TracebackType
from typing import Union, Callable, Any, Type
def delete_model_cache(model_dir: Union[str, Path]) -> None:
@@ -16,7 +17,11 @@ def delete_model_cache(model_dir: Union[str, Path]) -> None:
model_dir (Union[str, Path]): The path to the model cache directory.
"""
def on_error(func, path, exc_info):
def on_error(
func: Callable[..., Any],
path: str,
exc_info: tuple[Type[BaseException], BaseException, TracebackType],
) -> None:
print("Failed to remove: ", path)
print("Exception: ", exc_info)
traceback.print_exception(*exc_info)