From 50289c62ed33fb859792ce4618ab54597370b06b Mon Sep 17 00:00:00 2001 From: George Date: Wed, 15 Jan 2025 17:39:06 +0100 Subject: [PATCH 01/11] new: move onnx dependency to dev (#439) * new: move onnx dependency to dev * update poetry install, add more groups in pyproject --- .github/workflows/python-tests.yml | 2 +- pyproject.toml | 6 ++++-- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/.github/workflows/python-tests.yml b/.github/workflows/python-tests.yml index 0c1c52c..5c33d23 100644 --- a/.github/workflows/python-tests.yml +++ b/.github/workflows/python-tests.yml @@ -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: | diff --git a/pyproject.toml b/pyproject.toml index 7d2b8ba..7548f85 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" From fa11d0f0c77b788f51e27026cc6dcc3a94ea5884 Mon Sep 17 00:00:00 2001 From: George Panchuk Date: Thu, 16 Jan 2025 11:11:43 +0100 Subject: [PATCH 02/11] bump version to 0.5.1 --- .github/ISSUE_TEMPLATE/bug-report.yml | 2 +- pyproject.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/ISSUE_TEMPLATE/bug-report.yml b/.github/ISSUE_TEMPLATE/bug-report.yml index 2396391..ca0bffc 100644 --- a/.github/ISSUE_TEMPLATE/bug-report.yml +++ b/.github/ISSUE_TEMPLATE/bug-report.yml @@ -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 diff --git a/pyproject.toml b/pyproject.toml index 7548f85..f7b40f9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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 ", "NirantK "] license = "Apache License" From ae37da3bd4790c3628041614584bc0ea89cb26a9 Mon Sep 17 00:00:00 2001 From: Hossam Hagag <90828745+hh-space-invader@users.noreply.github.com> Date: Mon, 27 Jan 2025 12:25:04 +0200 Subject: [PATCH 03/11] fix: Update nomic ai model (#441) * fix: Updated nomic ai with mean pooling * chore: Updated warning message * nit * fix: Fix ci --- fastembed/text/text_embedding.py | 8 ++++++++ tests/test_text_onnx_embeddings.py | 8 ++++++-- 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/fastembed/text/text_embedding.py b/fastembed/text/text_embedding.py index 960d68f..962f2b9 100644 --- a/fastembed/text/text_embedding.py +++ b/fastembed/text/text_embedding.py @@ -1,3 +1,4 @@ +import warnings from typing import Any, Iterable, Optional, Sequence, Type, Union import numpy as np @@ -62,6 +63,13 @@ class TextEmbedding(TextEmbeddingBase): **kwargs, ): 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, + ) 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): diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py index ac2c41a..b7ad843 100644 --- a/tests/test_text_onnx_embeddings.py +++ b/tests/test_text_onnx_embeddings.py @@ -1,4 +1,5 @@ import os +import platform import numpy as np import pytest @@ -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] @@ -71,9 +72,12 @@ CANONICAL_VECTOR_VALUES = { def test_embedding(): 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 ( + is_mac and model_desc["model"] == "nomic-ai/nomic-embed-text-v1.5-Q" + ): continue dim = model_desc["dim"] From 54f6cd9cbc7609bb3413d16ec2afc07097c55e05 Mon Sep 17 00:00:00 2001 From: Hossam Hagag <90828745+hh-space-invader@users.noreply.github.com> Date: Mon, 27 Jan 2025 23:34:39 +0200 Subject: [PATCH 04/11] Improve progress bar new (#440) * improve: Improve progress bar * fix: Fix error downloading when internet connection down * new: Added file hash computation to track new versions * refactor: Removed redundant hash check fix: Fix ci * new: Verify using hf_api * new: Improve progress bar * refactor new progress bar (#446) * refactor * chore: Remove redundant enable progress bar --------- Co-authored-by: hh-space-invader * refactor comments --------- Co-authored-by: George --- fastembed/common/model_management.py | 134 +++++++++++++++++++++++++-- 1 file changed, 125 insertions(+), 9 deletions(-) diff --git a/fastembed/common/model_management.py b/fastembed/common/model_management.py index 33f513c..d85406f 100644 --- a/fastembed/common/model_management.py +++ b/fastembed/common/model_management.py @@ -1,12 +1,14 @@ import os import time +import json import shutil import tarfile from pathlib import Path -from typing import Any, Optional +from typing import Any 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,7 +102,7 @@ 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, ) -> str: @@ -107,12 +111,63 @@ class ModelManagement: 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. 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 +175,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,6 +235,24 @@ 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): """ From b05877de93b4eeff6bec315c8c45b6d9ba4e890c Mon Sep 17 00:00:00 2001 From: Hossam Hagag <90828745+hh-space-invader@users.noreply.github.com> Date: Tue, 28 Jan 2025 00:27:23 +0200 Subject: [PATCH 05/11] new: Added jina embedding v3 (#428) * new: Added jina embedding v3 * refactor: Changed dim to int value * new: Updated notice * new: Extended text embedding with query embed and passage embed * fix: Fix lazy load in query and passage embed * tests: Added test for multitask embeddings * nit: Remove cache dir from tests * tests: Updated tests * improve: Improve task selection * fix: Fix ci * fix: Update fastembed/text/multitask_embedding.py Co-authored-by: George * Update fastembed/text/multitask_embedding.py Co-authored-by: George * fix: Pass task id using kwargs to parallel processor * tests: Added test for task assignment * prefer enums over ints * tests: Added test for parallel * improve: Updated model description * fix: Fix ci * fix: Fix ci * refactor: Refactor query_embed and passage_embed * tests: Added task propagation to parallel * refactor: Set default task as retrieval passage * chore: Update default task in tests --------- Co-authored-by: George --- NOTICE | 2 + fastembed/text/multitask_embedding.py | 98 +++++++++ fastembed/text/text_embedding.py | 29 +++ tests/test_text_multitask_embeddings.py | 253 ++++++++++++++++++++++++ tests/test_text_onnx_embeddings.py | 8 +- 5 files changed, 388 insertions(+), 2 deletions(-) create mode 100644 fastembed/text/multitask_embedding.py create mode 100644 tests/test_text_multitask_embeddings.py diff --git a/NOTICE b/NOTICE index bfa9618..caa664b 100644 --- a/NOTICE +++ b/NOTICE @@ -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. diff --git a/fastembed/text/multitask_embedding.py b/fastembed/text/multitask_embedding.py new file mode 100644 index 0000000..f34efd1 --- /dev/null +++ b/fastembed/text/multitask_embedding.py @@ -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, **kwargs): + 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 diff --git a/fastembed/text/text_embedding.py b/fastembed/text/text_embedding.py index 962f2b9..fe5307d 100644 --- a/fastembed/text/text_embedding.py +++ b/fastembed/text/text_embedding.py @@ -8,6 +8,7 @@ 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 @@ -19,6 +20,7 @@ class TextEmbedding(TextEmbeddingBase): CLIPOnnxEmbedding, PooledNormalizedEmbedding, PooledEmbedding, + JinaEmbeddingV3, ] @classmethod @@ -113,3 +115,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) -> 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) -> 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) diff --git a/tests/test_text_multitask_embeddings.py b/tests/test_text_multitask_embeddings.py new file mode 100644 index 0000000..193a5f5 --- /dev/null +++ b/tests/test_text_multitask_embeddings.py @@ -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): + 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) diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py index b7ad843..96baf48 100644 --- a/tests/test_text_onnx_embeddings.py +++ b/tests/test_text_onnx_embeddings.py @@ -69,14 +69,18 @@ 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(): 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) or ( - is_mac and model_desc["model"] == "nomic-ai/nomic-embed-text-v1.5-Q" + 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 From c2f6fd1c90c4b0770d8676478d7444e98d421f8c Mon Sep 17 00:00:00 2001 From: Hossam Hagag <90828745+hh-space-invader@users.noreply.github.com> Date: Tue, 28 Jan 2025 00:27:47 +0200 Subject: [PATCH 06/11] Fix paraphrase minilm (#436) * fix: Fix minilm paraphrase by adding it to pool models * tests: Updated minilm paraphrase canonical vector * chore: Added a warning message for updating the model * chore: Added version where model will be removed --- fastembed/text/onnx_embedding.py | 11 ----------- fastembed/text/pooled_embedding.py | 11 +++++++++++ fastembed/text/text_embedding.py | 8 ++++++++ tests/test_text_onnx_embeddings.py | 2 +- 4 files changed, 20 insertions(+), 12 deletions(-) diff --git a/fastembed/text/onnx_embedding.py b/fastembed/text/onnx_embedding.py index 48b93da..22004f6 100644 --- a/fastembed/text/onnx_embedding.py +++ b/fastembed/text/onnx_embedding.py @@ -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, diff --git a/fastembed/text/pooled_embedding.py b/fastembed/text/pooled_embedding.py index 526f2df..3923957 100644 --- a/fastembed/text/pooled_embedding.py +++ b/fastembed/text/pooled_embedding.py @@ -40,6 +40,17 @@ 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", + }, ] diff --git a/fastembed/text/text_embedding.py b/fastembed/text/text_embedding.py index fe5307d..413943c 100644 --- a/fastembed/text/text_embedding.py +++ b/fastembed/text/text_embedding.py @@ -72,6 +72,14 @@ class TextEmbedding(TextEmbeddingBase): 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, + ) 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): diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py index 96baf48..5addf15 100644 --- a/tests/test_text_onnx_embeddings.py +++ b/tests/test_text_onnx_embeddings.py @@ -32,7 +32,7 @@ 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]), "sentence-transformers/paraphrase-multilingual-mpnet-base-v2": np.array( From 314842121d458483c428a813ebed7b45681af4a1 Mon Sep 17 00:00:00 2001 From: Dmitrii Ogn Date: Tue, 28 Jan 2025 14:46:16 +0300 Subject: [PATCH 07/11] Load from local dir (#443) * HF sources for all models * Specific_model_path model path support * Fix hf download * fix: rollback incorrect model replacement * refactor: remove redundant type imports * refactor: replace List with list * fix: remove redundant param in late interaction text embedding * Update fastembed/common/model_management.py * fix: rollback post process onnx output --------- Co-authored-by: George Panchuk --- fastembed/common/model_management.py | 4 ++++ fastembed/image/onnx_embedding.py | 10 +++++++--- fastembed/late_interaction/colbert.py | 7 ++++++- .../rerank/cross_encoder/onnx_text_cross_encoder.py | 3 +++ fastembed/sparse/bm25.py | 6 +++++- fastembed/sparse/bm42.py | 7 ++++++- fastembed/sparse/splade_pp.py | 7 ++++++- fastembed/text/onnx_embedding.py | 7 ++++++- fastembed/text/onnx_text_model.py | 1 + 9 files changed, 44 insertions(+), 8 deletions(-) diff --git a/fastembed/common/model_management.py b/fastembed/common/model_management.py index d85406f..f80d72a 100644 --- a/fastembed/common/model_management.py +++ b/fastembed/common/model_management.py @@ -114,6 +114,7 @@ class ModelManagement: 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. """ @@ -364,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") diff --git a/fastembed/image/onnx_embedding.py b/fastembed/image/onnx_embedding.py index 47beae2..82b21d3 100644 --- a/fastembed/image/onnx_embedding.py +++ b/fastembed/image/onnx_embedding.py @@ -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,6 +77,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]): device_ids: Optional[list[int]] = None, lazy_load: bool = False, device_id: Optional[int] = None, + specific_model_path: Optional[str] = None, **kwargs, ): """ @@ -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 / 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 diff --git a/fastembed/late_interaction/colbert.py b/fastembed/late_interaction/colbert.py index 97c3857..e9da217 100644 --- a/fastembed/late_interaction/colbert.py +++ b/fastembed/late_interaction/colbert.py @@ -124,6 +124,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]): device_ids: Optional[list[int]] = None, lazy_load: bool = False, device_id: Optional[int] = None, + specific_model_path: Optional[str] = None, **kwargs, ): """ @@ -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 / 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 diff --git a/fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py b/fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py index 3b71671..4b31063 100644 --- a/fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py +++ b/fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py @@ -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 / 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: diff --git a/fastembed/sparse/bm25.py b/fastembed/sparse/bm25.py index 9db6c8d..789a44d 100644 --- a/fastembed/sparse/bm25.py +++ b/fastembed/sparse/bm25.py @@ -108,6 +108,7 @@ class Bm25(SparseTextEmbeddingBase): language: str = "english", token_max_length: int = 40, disable_stemmer: bool = False, + specific_model_path: Optional[str] = None, **kwargs, ): 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 diff --git a/fastembed/sparse/bm42.py b/fastembed/sparse/bm42.py index 3de1900..171a16f 100644 --- a/fastembed/sparse/bm42.py +++ b/fastembed/sparse/bm42.py @@ -66,6 +66,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]): device_ids: Optional[list[int]] = None, lazy_load: bool = False, device_id: Optional[int] = None, + specific_model_path: Optional[str] = None, **kwargs, ): """ @@ -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 / 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 = {} diff --git a/fastembed/sparse/splade_pp.py b/fastembed/sparse/splade_pp.py index cbd88be..1a8bbd2 100644 --- a/fastembed/sparse/splade_pp.py +++ b/fastembed/sparse/splade_pp.py @@ -73,6 +73,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]): device_ids: Optional[list[int]] = None, lazy_load: bool = False, device_id: Optional[int] = None, + specific_model_path: Optional[str] = None, **kwargs, ): """ @@ -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 / 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: diff --git a/fastembed/text/onnx_embedding.py b/fastembed/text/onnx_embedding.py index 22004f6..967c9c6 100644 --- a/fastembed/text/onnx_embedding.py +++ b/fastembed/text/onnx_embedding.py @@ -193,6 +193,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]): device_ids: Optional[list[int]] = None, lazy_load: bool = False, device_id: Optional[int] = None, + specific_model_path: Optional[str] = None, **kwargs, ): """ @@ -211,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 / e.g. BAAI/bge-base-en. @@ -234,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: diff --git a/fastembed/text/onnx_text_model.py b/fastembed/text/onnx_text_model.py index ba3e151..b598c40 100644 --- a/fastembed/text/onnx_text_model.py +++ b/fastembed/text/onnx_text_model.py @@ -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, From bb815405aa9cfd1863f837c9e0413fc03d8b7d91 Mon Sep 17 00:00:00 2001 From: Hossam Hagag <90828745+hh-space-invader@users.noreply.github.com> Date: Tue, 28 Jan 2025 23:49:37 +0200 Subject: [PATCH 08/11] chore: Add any to kwargs (#450) * new: Add mypy and pyright deps * chore: Add any to kwargs * chore: Add any to args * add missing kwargs --------- Co-authored-by: George Panchuk --- fastembed/common/model_management.py | 4 ++-- fastembed/common/onnx_model.py | 8 ++++---- fastembed/embedding.py | 4 ++-- fastembed/image/image_embedding.py | 4 ++-- fastembed/image/image_embedding_base.py | 6 +++--- fastembed/image/onnx_embedding.py | 8 ++++---- fastembed/image/onnx_image_model.py | 6 +++--- fastembed/late_interaction/colbert.py | 8 ++++---- fastembed/late_interaction/jina_colbert.py | 2 +- .../late_interaction_embedding_base.py | 12 +++++------- .../late_interaction_text_embedding.py | 6 +++--- .../rerank/cross_encoder/onnx_text_cross_encoder.py | 2 +- .../rerank/cross_encoder/text_cross_encoder_base.py | 4 ++-- fastembed/sparse/bm25.py | 12 +++++++----- fastembed/sparse/bm42.py | 10 ++++++---- fastembed/sparse/sparse_embedding_base.py | 12 +++++++----- fastembed/sparse/sparse_text_embedding.py | 8 +++++--- fastembed/sparse/splade_pp.py | 6 +++--- fastembed/text/clip_embedding.py | 2 +- fastembed/text/e5_onnx_embedding.py | 4 ++-- fastembed/text/multitask_embedding.py | 2 +- fastembed/text/onnx_embedding.py | 8 ++++---- fastembed/text/onnx_text_model.py | 8 ++++---- fastembed/text/pooled_embedding.py | 2 +- fastembed/text/pooled_normalized_embedding.py | 2 +- fastembed/text/text_embedding.py | 8 ++++---- fastembed/text/text_embedding_base.py | 12 +++++------- pyproject.toml | 6 ++++++ 28 files changed, 93 insertions(+), 83 deletions(-) diff --git a/fastembed/common/model_management.py b/fastembed/common/model_management.py index f80d72a..943f85d 100644 --- a/fastembed/common/model_management.py +++ b/fastembed/common/model_management.py @@ -104,7 +104,7 @@ class ModelManagement: cache_dir: str, extra_patterns: list[str], local_files_only: bool = False, - **kwargs, + **kwargs: Any, ) -> str: """ Downloads a model from HuggingFace Hub. @@ -338,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. diff --git a/fastembed/common/onnx_model.py b/fastembed/common/onnx_model.py index 9e9376f..8a03e3d 100644 --- a/fastembed/common/onnx_model.py +++ b/fastembed/common/onnx_model.py @@ -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) diff --git a/fastembed/embedding.py b/fastembed/embedding.py index 8140265..4266de4 100644 --- a/fastembed/embedding.py +++ b/fastembed/embedding.py @@ -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) diff --git a/fastembed/image/image_embedding.py b/fastembed/image/image_embedding.py index aa4c91b..490f88a 100644 --- a/fastembed/image/image_embedding.py +++ b/fastembed/image/image_embedding.py @@ -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. diff --git a/fastembed/image/image_embedding_base.py b/fastembed/image/image_embedding_base.py index 6a23df1..c4c13a1 100644 --- a/fastembed/image/image_embedding_base.py +++ b/fastembed/image/image_embedding_base.py @@ -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. diff --git a/fastembed/image/onnx_embedding.py b/fastembed/image/onnx_embedding.py index 82b21d3..ae7bfce 100644 --- a/fastembed/image/onnx_embedding.py +++ b/fastembed/image/onnx_embedding.py @@ -78,7 +78,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[np.ndarray]): lazy_load: bool = False, device_id: Optional[int] = None, specific_model_path: Optional[str] = None, - **kwargs, + **kwargs: Any, ): """ Args: @@ -158,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. @@ -193,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. @@ -206,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, diff --git a/fastembed/image/onnx_image_model.py b/fastembed/image/onnx_image_model.py index 895f006..067847c 100644 --- a/fastembed/image/onnx_image_model.py +++ b/fastembed/image/onnx_image_model.py @@ -29,7 +29,7 @@ class OnnxImageModel(OnnxModel[T]): 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 diff --git a/fastembed/late_interaction/colbert.py b/fastembed/late_interaction/colbert.py index e9da217..a9840d0 100644 --- a/fastembed/late_interaction/colbert.py +++ b/fastembed/late_interaction/colbert.py @@ -125,7 +125,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[np.ndarray]): lazy_load: bool = False, device_id: Optional[int] = None, specific_model_path: Optional[str] = None, - **kwargs, + **kwargs: Any, ): """ Args: @@ -205,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. @@ -234,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] @@ -252,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, diff --git a/fastembed/late_interaction/jina_colbert.py b/fastembed/late_interaction/jina_colbert.py index 9e19e0b..566e4ae 100644 --- a/fastembed/late_interaction/jina_colbert.py +++ b/fastembed/late_interaction/jina_colbert.py @@ -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, diff --git a/fastembed/late_interaction/late_interaction_embedding_base.py b/fastembed/late_interaction/late_interaction_embedding_base.py index 64fba49..bc8c6da 100644 --- a/fastembed/late_interaction/late_interaction_embedding_base.py +++ b/fastembed/late_interaction/late_interaction_embedding_base.py @@ -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 diff --git a/fastembed/late_interaction/late_interaction_text_embedding.py b/fastembed/late_interaction/late_interaction_text_embedding.py index 58c8841..993924f 100644 --- a/fastembed/late_interaction/late_interaction_text_embedding.py +++ b/fastembed/late_interaction/late_interaction_text_embedding.py @@ -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 diff --git a/fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py b/fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py index 4b31063..0d2c9e0 100644 --- a/fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py +++ b/fastembed/rerank/cross_encoder/onnx_text_cross_encoder.py @@ -217,7 +217,7 @@ class TextCrossEncoderWorker(TextRerankerWorker): self, model_name: str, cache_dir: str, - **kwargs, + **kwargs: Any, ) -> OnnxTextCrossEncoder: return OnnxTextCrossEncoder( model_name=model_name, diff --git a/fastembed/rerank/cross_encoder/text_cross_encoder_base.py b/fastembed/rerank/cross_encoder/text_cross_encoder_base.py index 79d0210..cadf6d0 100644 --- a/fastembed/rerank/cross_encoder/text_cross_encoder_base.py +++ b/fastembed/rerank/cross_encoder/text_cross_encoder_base.py @@ -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. diff --git a/fastembed/sparse/bm25.py b/fastembed/sparse/bm25.py index 789a44d..0033089 100644 --- a/fastembed/sparse/bm25.py +++ b/fastembed/sparse/bm25.py @@ -109,7 +109,7 @@ class Bm25(SparseTextEmbeddingBase): token_max_length: int = 40, disable_stemmer: bool = False, specific_model_path: Optional[str] = None, - **kwargs, + **kwargs: Any, ): super().__init__(model_name, cache_dir, **kwargs) @@ -213,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. @@ -305,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. """ @@ -333,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) @@ -347,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) diff --git a/fastembed/sparse/bm42.py b/fastembed/sparse/bm42.py index 171a16f..d593e84 100644 --- a/fastembed/sparse/bm42.py +++ b/fastembed/sparse/bm42.py @@ -67,7 +67,7 @@ class Bm42(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]): lazy_load: bool = False, device_id: Optional[int] = None, specific_model_path: Optional[str] = None, - **kwargs, + **kwargs: Any, ): """ Args: @@ -273,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. @@ -310,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. @@ -337,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, diff --git a/fastembed/sparse/sparse_embedding_base.py b/fastembed/sparse/sparse_embedding_base.py index 740cd9d..ee6761d 100644 --- a/fastembed/sparse/sparse_embedding_base.py +++ b/fastembed/sparse/sparse_embedding_base.py @@ -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 diff --git a/fastembed/sparse/sparse_text_embedding.py b/fastembed/sparse/sparse_text_embedding.py index d8c013e..66b60c3 100644 --- a/fastembed/sparse/sparse_text_embedding.py +++ b/fastembed/sparse/sparse_text_embedding.py @@ -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 diff --git a/fastembed/sparse/splade_pp.py b/fastembed/sparse/splade_pp.py index 1a8bbd2..807d3a9 100644 --- a/fastembed/sparse/splade_pp.py +++ b/fastembed/sparse/splade_pp.py @@ -74,7 +74,7 @@ class SpladePP(SparseTextEmbeddingBase, OnnxTextModel[SparseEmbedding]): lazy_load: bool = False, device_id: Optional[int] = None, specific_model_path: Optional[str] = None, - **kwargs, + **kwargs: Any, ): """ Args: @@ -141,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. @@ -176,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, diff --git a/fastembed/text/clip_embedding.py b/fastembed/text/clip_embedding.py index a757d87..b0721bf 100644 --- a/fastembed/text/clip_embedding.py +++ b/fastembed/text/clip_embedding.py @@ -44,7 +44,7 @@ class CLIPEmbeddingWorker(OnnxTextEmbeddingWorker): self, model_name: str, cache_dir: str, - **kwargs, + **kwargs: Any, ) -> OnnxTextEmbedding: return CLIPOnnxEmbedding( model_name=model_name, diff --git a/fastembed/text/e5_onnx_embedding.py b/fastembed/text/e5_onnx_embedding.py index 96a379d..7dea36e 100644 --- a/fastembed/text/e5_onnx_embedding.py +++ b/fastembed/text/e5_onnx_embedding.py @@ -48,7 +48,7 @@ class E5OnnxEmbedding(OnnxTextEmbedding): return supported_multilingual_e5_models 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. @@ -62,7 +62,7 @@ class E5OnnxEmbeddingWorker(OnnxTextEmbeddingWorker): self, model_name: str, cache_dir: str, - **kwargs, + **kwargs: Any, ) -> E5OnnxEmbedding: return E5OnnxEmbedding( model_name=model_name, diff --git a/fastembed/text/multitask_embedding.py b/fastembed/text/multitask_embedding.py index f34efd1..cc479c2 100644 --- a/fastembed/text/multitask_embedding.py +++ b/fastembed/text/multitask_embedding.py @@ -42,7 +42,7 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding): PASSAGE_TASK = Task.RETRIEVAL_PASSAGE QUERY_TASK = Task.RETRIEVAL_QUERY - def __init__(self, *args, **kwargs): + def __init__(self, *args: Any, **kwargs: Any): super().__init__(*args, **kwargs) self._current_task_id = self.PASSAGE_TASK diff --git a/fastembed/text/onnx_embedding.py b/fastembed/text/onnx_embedding.py index 967c9c6..a93920c 100644 --- a/fastembed/text/onnx_embedding.py +++ b/fastembed/text/onnx_embedding.py @@ -194,7 +194,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[np.ndarray]): lazy_load: bool = False, device_id: Optional[int] = None, specific_model_path: Optional[str] = None, - **kwargs, + **kwargs: Any, ): """ Args: @@ -250,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. @@ -284,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. @@ -317,7 +317,7 @@ class OnnxTextEmbeddingWorker(TextEmbeddingWorker): self, model_name: str, cache_dir: str, - **kwargs, + **kwargs: Any, ) -> OnnxTextEmbedding: return OnnxTextEmbedding( model_name=model_name, diff --git a/fastembed/text/onnx_text_model.py b/fastembed/text/onnx_text_model.py index b598c40..7c8a2c8 100644 --- a/fastembed/text/onnx_text_model.py +++ b/fastembed/text/onnx_text_model.py @@ -29,7 +29,7 @@ class OnnxTextModel(OnnxModel[T]): 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. @@ -59,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]) @@ -99,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 diff --git a/fastembed/text/pooled_embedding.py b/fastembed/text/pooled_embedding.py index 3923957..71f1cbe 100644 --- a/fastembed/text/pooled_embedding.py +++ b/fastembed/text/pooled_embedding.py @@ -93,7 +93,7 @@ class PooledEmbeddingWorker(OnnxTextEmbeddingWorker): self, model_name: str, cache_dir: str, - **kwargs, + **kwargs: Any, ) -> OnnxTextEmbedding: return PooledEmbedding( model_name=model_name, diff --git a/fastembed/text/pooled_normalized_embedding.py b/fastembed/text/pooled_normalized_embedding.py index 47ec91a..5d5cef7 100644 --- a/fastembed/text/pooled_normalized_embedding.py +++ b/fastembed/text/pooled_normalized_embedding.py @@ -115,7 +115,7 @@ class PooledNormalizedEmbeddingWorker(OnnxTextEmbeddingWorker): self, model_name: str, cache_dir: str, - **kwargs, + **kwargs: Any, ) -> OnnxTextEmbedding: return PooledNormalizedEmbedding( model_name=model_name, diff --git a/fastembed/text/text_embedding.py b/fastembed/text/text_embedding.py index 413943c..b1cf414 100644 --- a/fastembed/text/text_embedding.py +++ b/fastembed/text/text_embedding.py @@ -62,7 +62,7 @@ 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": @@ -105,7 +105,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. @@ -124,7 +124,7 @@ class TextEmbedding(TextEmbeddingBase): """ 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 @@ -137,7 +137,7 @@ class TextEmbedding(TextEmbeddingBase): # 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) -> 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. diff --git a/fastembed/text/text_embedding_base.py b/fastembed/text/text_embedding_base.py index 61963fc..9d4cc4d 100644 --- a/fastembed/text/text_embedding_base.py +++ b/fastembed/text/text_embedding_base.py @@ -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,7 @@ 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 diff --git a/pyproject.toml b/pyproject.toml index f7b40f9..7111c22 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,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 From 105d6cfb97497b25ab76bb087195fafb21594019 Mon Sep 17 00:00:00 2001 From: Dmitrii Ogn Date: Wed, 29 Jan 2025 01:14:43 +0300 Subject: [PATCH 09/11] E5 pooling fix (#445) * HF sources for all models * Proper normalization for e5 models * Rollback to origin/master * Warning * Tests fix * Logging + model refactoring * fix: refactor warnings, make e5-large non-normalized * remove redundant code, update canonical values for e5 * align warning style --------- Co-authored-by: George Panchuk --- fastembed/text/e5_onnx_embedding.py | 72 ----------------------------- fastembed/text/pooled_embedding.py | 24 ++++++++++ fastembed/text/text_embedding.py | 13 ++++-- tests/test_text_onnx_embeddings.py | 4 +- 4 files changed, 36 insertions(+), 77 deletions(-) delete mode 100644 fastembed/text/e5_onnx_embedding.py diff --git a/fastembed/text/e5_onnx_embedding.py b/fastembed/text/e5_onnx_embedding.py deleted file mode 100644 index 7dea36e..0000000 --- a/fastembed/text/e5_onnx_embedding.py +++ /dev/null @@ -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: Any - ) -> 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: Any, - ) -> E5OnnxEmbedding: - return E5OnnxEmbedding( - model_name=model_name, - cache_dir=cache_dir, - threads=1, - **kwargs, - ) diff --git a/fastembed/text/pooled_embedding.py b/fastembed/text/pooled_embedding.py index 71f1cbe..063c47b 100644 --- a/fastembed/text/pooled_embedding.py +++ b/fastembed/text/pooled_embedding.py @@ -51,6 +51,30 @@ supported_pooled_models = [ }, "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"], + }, ] diff --git a/fastembed/text/text_embedding.py b/fastembed/text/text_embedding.py index b1cf414..30e4bba 100644 --- a/fastembed/text/text_embedding.py +++ b/fastembed/text/text_embedding.py @@ -2,10 +2,8 @@ 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 @@ -16,7 +14,6 @@ from fastembed.text.text_embedding_base import TextEmbeddingBase class TextEmbedding(TextEmbeddingBase): EMBEDDINGS_REGISTRY: list[Type[TextEmbeddingBase]] = [ OnnxTextEmbedding, - E5OnnxEmbedding, CLIPOnnxEmbedding, PooledNormalizedEmbedding, PooledEmbedding, @@ -80,6 +77,16 @@ class TextEmbedding(TextEmbeddingBase): 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): diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py index 5addf15..5df1d7c 100644 --- a/tests/test_text_onnx_embeddings.py +++ b/tests/test_text_onnx_embeddings.py @@ -34,9 +34,9 @@ CANONICAL_VECTOR_VALUES = { "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2": np.array( [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]), From 73e1e5ecb98119190960f7986677358b549bcec4 Mon Sep 17 00:00:00 2001 From: Hossam Hagag <90828745+hh-space-invader@users.noreply.github.com> Date: Wed, 29 Jan 2025 09:33:49 +0200 Subject: [PATCH 10/11] chore: Add missing returns in defs (#451) * chore: Add missing returns in defs * remove return type from init * remove incorrect ndarray specifier --------- Co-authored-by: George Panchuk --- docs/examples/ColBERT_with_FastEmbed.ipynb | 17 ++++++++++++++--- docs/examples/FastEmbed_vs_HF_Comparison.ipynb | 2 +- docs/examples/Hybrid_Search.ipynb | 6 +++--- fastembed/common/model_management.py | 2 +- fastembed/image/onnx_image_model.py | 2 +- fastembed/image/transform/functional.py | 2 +- fastembed/image/transform/operators.py | 14 +++++++------- .../rerank/cross_encoder/onnx_text_model.py | 7 +++++-- fastembed/text/onnx_text_model.py | 2 +- fastembed/text/text_embedding_base.py | 1 + tests/test_attention_embeddings.py | 10 +++++----- tests/test_image_onnx_embeddings.py | 8 ++++---- tests/test_multi_gpu.py | 6 +++--- tests/test_sparse_embeddings.py | 16 ++++++++-------- tests/test_text_cross_encoder.py | 8 ++++---- tests/test_text_onnx_embeddings.py | 8 ++++---- tests/utils.py | 2 +- 17 files changed, 64 insertions(+), 49 deletions(-) diff --git a/docs/examples/ColBERT_with_FastEmbed.ipynb b/docs/examples/ColBERT_with_FastEmbed.ipynb index b3f51e2..4cffec4 100644 --- a/docs/examples/ColBERT_with_FastEmbed.ipynb +++ b/docs/examples/ColBERT_with_FastEmbed.ipynb @@ -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", diff --git a/docs/examples/FastEmbed_vs_HF_Comparison.ipynb b/docs/examples/FastEmbed_vs_HF_Comparison.ipynb index d87af16..32c5318 100644 --- a/docs/examples/FastEmbed_vs_HF_Comparison.ipynb +++ b/docs/examples/FastEmbed_vs_HF_Comparison.ipynb @@ -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", diff --git a/docs/examples/Hybrid_Search.ipynb b/docs/examples/Hybrid_Search.ipynb index 32b822b..95182d8 100644 --- a/docs/examples/Hybrid_Search.ipynb +++ b/docs/examples/Hybrid_Search.ipynb @@ -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", diff --git a/fastembed/common/model_management.py b/fastembed/common/model_management.py index 943f85d..2f5afec 100644 --- a/fastembed/common/model_management.py +++ b/fastembed/common/model_management.py @@ -255,7 +255,7 @@ class ModelManagement: 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. diff --git a/fastembed/image/onnx_image_model.py b/fastembed/image/onnx_image_model.py index 067847c..188bffb 100644 --- a/fastembed/image/onnx_image_model.py +++ b/fastembed/image/onnx_image_model.py @@ -24,7 +24,7 @@ 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 diff --git a/fastembed/image/transform/functional.py b/fastembed/image/transform/functional.py index afefe4b..9580612 100644 --- a/fastembed/image/transform/functional.py +++ b/fastembed/image/transform/functional.py @@ -118,7 +118,7 @@ def rescale(image: np.ndarray, scale: float, dtype=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 diff --git a/fastembed/image/transform/operators.py b/fastembed/image/transform/operators.py index bac65e0..7c810f9 100644 --- a/fastembed/image/transform/operators.py +++ b/fastembed/image/transform/operators.py @@ -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 diff --git a/fastembed/rerank/cross_encoder/onnx_text_model.py b/fastembed/rerank/cross_encoder/onnx_text_model.py index f3a02b6..e3e8f9c 100644 --- a/fastembed/rerank/cross_encoder/onnx_text_model.py +++ b/fastembed/rerank/cross_encoder/onnx_text_model.py @@ -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), diff --git a/fastembed/text/onnx_text_model.py b/fastembed/text/onnx_text_model.py index 7c8a2c8..269bfef 100644 --- a/fastembed/text/onnx_text_model.py +++ b/fastembed/text/onnx_text_model.py @@ -23,7 +23,7 @@ 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 = {} diff --git a/fastembed/text/text_embedding_base.py b/fastembed/text/text_embedding_base.py index 9d4cc4d..ed6ba18 100644 --- a/fastembed/text/text_embedding_base.py +++ b/fastembed/text/text_embedding_base.py @@ -43,6 +43,7 @@ class TextEmbeddingBase(ModelManagement): yield from self.embed(texts, **kwargs) def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[np.ndarray]: + """ Embeds queries diff --git a/tests/test_attention_embeddings.py b/tests/test_attention_embeddings.py index ac8aa72..f5c3692 100644 --- a/tests/test_attention_embeddings.py +++ b/tests/test_attention_embeddings.py @@ -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) -> 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) -> 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) -> 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) -> 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) -> None: model = SparseTextEmbedding(model_name=model_name, lazy_load=True) assert not hasattr(model.model, "model") docs = ["hello world", "flag embedding"] diff --git a/tests/test_image_onnx_embeddings.py b/tests/test_image_onnx_embeddings.py index a5fb8e3..23b1298 100644 --- a/tests/test_image_onnx_embeddings.py +++ b/tests/test_image_onnx_embeddings.py @@ -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, model_name) -> 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, model_name) -> 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) -> None: is_ci = os.getenv("CI") model = ImageEmbedding(model_name=model_name, lazy_load=True) assert not hasattr(model.model, "model") diff --git a/tests/test_multi_gpu.py b/tests/test_multi_gpu.py index a5ed830..a6e9dc3 100644 --- a/tests/test_multi_gpu.py +++ b/tests/test_multi_gpu.py @@ -13,7 +13,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) -> None: docs = ["hello world", "flag embedding"] device_id = device_id if device_id is not None else 0 @@ -85,7 +85,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) -> None: docs = ["hello world", "flag embedding"] device_id = device_ids[0] if device_ids else 0 embedding_model = TextEmbedding( @@ -170,7 +170,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, parallel) -> None: docs = ["hello world", "flag embedding"] * 100 batch_size = 5 diff --git a/tests/test_sparse_embeddings.py b/tests/test_sparse_embeddings.py index 236b1de..943546d 100644 --- a/tests/test_sparse_embeddings.py +++ b/tests/test_sparse_embeddings.py @@ -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) -> 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) -> 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) -> 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) -> None: is_ci = os.getenv("CI") model = SparseTextEmbedding(model_name=model_name, lazy_load=True) assert not hasattr(model.model, "model") diff --git a/tests/test_text_cross_encoder.py b/tests/test_text_cross_encoder.py index 15ab373..58cc6df 100644 --- a/tests/test_text_cross_encoder.py +++ b/tests/test_text_cross_encoder.py @@ -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) -> 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) -> 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) -> 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) -> None: is_ci = os.getenv("CI") model = TextCrossEncoder(model_name=model_name) diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py index 5df1d7c..875c707 100644 --- a/tests/test_text_onnx_embeddings.py +++ b/tests/test_text_onnx_embeddings.py @@ -72,7 +72,7 @@ CANONICAL_VECTOR_VALUES = { MULTI_TASK_MODELS = ["jinaai/jina-embeddings-v3"] -def test_embedding(): +def test_embedding() -> None: is_ci = os.getenv("CI") is_mac = platform.system() == "Darwin" @@ -104,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, model_name) -> None: is_ci = os.getenv("CI") model = TextEmbedding(model_name=model_name) @@ -121,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, model_name) -> None: is_ci = os.getenv("CI") model = TextEmbedding(model_name=model_name) @@ -147,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) -> None: is_ci = os.getenv("CI") model = TextEmbedding(model_name=model_name, lazy_load=True) assert not hasattr(model.model, "model") diff --git a/tests/utils.py b/tests/utils.py index 8698c8f..0aa93a7 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -16,7 +16,7 @@ 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, path, exc_info) -> None: print("Failed to remove: ", path) print("Exception: ", exc_info) traceback.print_exception(*exc_info) From 993dcd5f6895980cf0b7832b1059db2b68969a2b Mon Sep 17 00:00:00 2001 From: Hossam Hagag <90828745+hh-space-invader@users.noreply.github.com> Date: Wed, 29 Jan 2025 13:08:20 +0200 Subject: [PATCH 11/11] chore: Add missing type hints in functions (#453) * chore: Add missing type hints in functions * add missing import, small type refactor --------- Co-authored-by: George Panchuk --- fastembed/common/model_management.py | 2 +- fastembed/common/types.py | 4 ++-- fastembed/common/utils.py | 8 +++++--- fastembed/image/transform/functional.py | 2 +- tests/test_attention_embeddings.py | 10 +++++----- tests/test_image_onnx_embeddings.py | 6 +++--- tests/test_late_interaction_embeddings.py | 2 +- tests/test_multi_gpu.py | 7 ++++--- tests/test_sparse_embeddings.py | 8 ++++---- tests/test_text_cross_encoder.py | 8 ++++---- tests/test_text_multitask_embeddings.py | 2 +- tests/test_text_onnx_embeddings.py | 6 +++--- tests/utils.py | 9 +++++++-- 13 files changed, 41 insertions(+), 33 deletions(-) diff --git a/fastembed/common/model_management.py b/fastembed/common/model_management.py index 2f5afec..2db5262 100644 --- a/fastembed/common/model_management.py +++ b/fastembed/common/model_management.py @@ -4,7 +4,7 @@ import json import shutil import tarfile from pathlib import Path -from typing import Any +from typing import Any, Optional import requests from huggingface_hub import snapshot_download, model_info, list_repo_tree diff --git a/fastembed/common/types.py b/fastembed/common/types.py index 98536e8..0750cfc 100644 --- a/fastembed/common/types.py +++ b/fastembed/common/types.py @@ -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] diff --git a/fastembed/common/utils.py b/fastembed/common/utils.py index 3a8bb7f..45a4e59 100644 --- a/fastembed/common/utils.py +++ b/fastembed/common/utils.py @@ -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]] diff --git a/fastembed/image/transform/functional.py b/fastembed/image/transform/functional.py index 9580612..b15234b 100644 --- a/fastembed/image/transform/functional.py +++ b/fastembed/image/transform/functional.py @@ -114,7 +114,7 @@ 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) diff --git a/tests/test_attention_embeddings.py b/tests/test_attention_embeddings.py index f5c3692..a3598e0 100644 --- a/tests/test_attention_embeddings.py +++ b/tests/test_attention_embeddings.py @@ -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) -> None: +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) -> None: @pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions", "Qdrant/bm25"]) -def test_parallel_processing(model_name) -> None: +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) -> None: @pytest.mark.parametrize("model_name", ["Qdrant/bm25"]) -def test_multilanguage(model_name) -> None: +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) -> None: @pytest.mark.parametrize("model_name", ["Qdrant/bm25"]) -def test_special_characters(model_name) -> None: +def test_special_characters(model_name: str) -> None: is_ci = os.getenv("CI") docs = [ @@ -145,7 +145,7 @@ def test_special_characters(model_name) -> None: @pytest.mark.parametrize("model_name", ["Qdrant/bm42-all-minilm-l6-v2-attentions"]) -def test_lazy_load(model_name) -> None: +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"] diff --git a/tests/test_image_onnx_embeddings.py b/tests/test_image_onnx_embeddings.py index 23b1298..1488b27 100644 --- a/tests/test_image_onnx_embeddings.py +++ b/tests/test_image_onnx_embeddings.py @@ -61,7 +61,7 @@ def test_embedding() -> None: @pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")]) -def test_batch_embedding(n_dims, model_name) -> None: +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) -> None: @pytest.mark.parametrize("n_dims,model_name", [(512, "Qdrant/clip-ViT-B-32-vision")]) -def test_parallel_processing(n_dims, model_name) -> None: +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) -> None: @pytest.mark.parametrize("model_name", ["Qdrant/clip-ViT-B-32-vision"]) -def test_lazy_load(model_name) -> None: +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") diff --git a/tests/test_late_interaction_embeddings.py b/tests/test_late_interaction_embeddings.py index 470ea2d..613895d 100644 --- a/tests/test_late_interaction_embeddings.py +++ b/tests/test_late_interaction_embeddings.py @@ -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) diff --git a/tests/test_multi_gpu.py b/tests/test_multi_gpu.py index a6e9dc3..32d6c94 100644 --- a/tests/test_multi_gpu.py +++ b/tests/test_multi_gpu.py @@ -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) -> None: +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) -> None: @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) -> None: +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) -> None: @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) -> None: +def test_multi_gpu_parallel_inference(device_ids: Optional[list[int]], parallel: int) -> None: docs = ["hello world", "flag embedding"] * 100 batch_size = 5 diff --git a/tests/test_sparse_embeddings.py b/tests/test_sparse_embeddings.py index 943546d..48cc14f 100644 --- a/tests/test_sparse_embeddings.py +++ b/tests/test_sparse_embeddings.py @@ -119,7 +119,7 @@ def bm25_instance() -> None: delete_model_cache(model._model_dir) -def test_stem_with_stopwords_and_punctuation(bm25_instance) -> None: +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) -> None: assert result == expected, f"Expected {expected}, but got {result}" -def test_stem_case_insensitive_stopwords(bm25_instance) -> None: +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) -> None: @pytest.mark.parametrize("disable_stemmer", [True, False]) -def test_disable_stemmer_behavior(disable_stemmer) -> None: +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) -> None: "model_name", ["prithivida/Splade_PP_en_v1"], ) -def test_lazy_load(model_name) -> None: +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") diff --git a/tests/test_text_cross_encoder.py b/tests/test_text_cross_encoder.py index 58cc6df..680c1a0 100644 --- a/tests/test_text_cross_encoder.py +++ b/tests/test_text_cross_encoder.py @@ -26,7 +26,7 @@ SELECTED_MODELS = { "model_name", [model_name for model_name in CANONICAL_SCORE_VALUES], ) -def test_rerank(model_name) -> None: +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) -> None: "model_name", [model_name for model_name in SELECTED_MODELS.values()], ) -def test_batch_rerank(model_name) -> None: +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) -> None: "model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"], ) -def test_lazy_load(model_name) -> None: +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) -> None: "model_name", [model_name for model_name in SELECTED_MODELS.values()], ) -def test_rerank_pairs_parallel(model_name) -> None: +def test_rerank_pairs_parallel(model_name: str) -> None: is_ci = os.getenv("CI") model = TextCrossEncoder(model_name=model_name) diff --git a/tests/test_text_multitask_embeddings.py b/tests/test_text_multitask_embeddings.py index 193a5f5..2292dd3 100644 --- a/tests/test_text_multitask_embeddings.py +++ b/tests/test_text_multitask_embeddings.py @@ -241,7 +241,7 @@ def test_task_assignment(): "model_name", ["jinaai/jina-embeddings-v3"], ) -def test_lazy_load(model_name): +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") diff --git a/tests/test_text_onnx_embeddings.py b/tests/test_text_onnx_embeddings.py index 875c707..757e097 100644 --- a/tests/test_text_onnx_embeddings.py +++ b/tests/test_text_onnx_embeddings.py @@ -104,7 +104,7 @@ def test_embedding() -> None: "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) -> None: +def test_batch_embedding(n_dims: int, model_name: str) -> None: is_ci = os.getenv("CI") model = TextEmbedding(model_name=model_name) @@ -121,7 +121,7 @@ def test_batch_embedding(n_dims, model_name) -> None: "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) -> None: +def test_parallel_processing(n_dims: int, model_name: str) -> None: is_ci = os.getenv("CI") model = TextEmbedding(model_name=model_name) @@ -147,7 +147,7 @@ def test_parallel_processing(n_dims, model_name) -> None: "model_name", ["BAAI/bge-small-en-v1.5"], ) -def test_lazy_load(model_name) -> None: +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") diff --git a/tests/utils.py b/tests/utils.py index 0aa93a7..cfd6ae8 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -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) -> None: + 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)