fix parallel worker init

* new: improve task setter in jina v3

* refactor

* new: add hf_token secret

* fix: cross platform env propagation

.github/workflows/python-tests.yml

@@ -1,10 +1,10 @@
 name: Tests
-run-name: Tests (gpu)
 
 on:
-  push:
-    branches: [ master, main, gpu ]
   pull_request:
+    branches: [ master, main, gpu ]
+  workflow_dispatch:
+  
 
 env:
   CARGO_TERM_COLOR: always
@@ -22,6 +22,8 @@ jobs:
           - '3.13.x'
         os:
           - ubuntu-latest
+          - macos-latest
+          - windows-latest
 
     runs-on: ${{ matrix.os }}
 
@@ -40,5 +42,7 @@ jobs:
           poetry install --no-interaction --no-ansi --without dev,docs
 
       - name: Run pytest
+        env:
+          HF_TOKEN: ${{ secrets.HF_TOKEN }}
         run: |
-          poetry run pytest 
+          poetry run pytest

.github/workflows/type-checkers.yml

@@ -26,8 +26,6 @@ jobs:
           python -m pip install --upgrade pip poetry
           poetry install --no-interaction --no-ansi --without dev,docs,test
 
-          poetry run pip install "numpy<2.0.0"  # https://github.com/python/mypy/issues/17396
-
       - name: mypy
         run: |
             poetry run mypy fastembed \

README.md

@@ -190,6 +190,23 @@ scores = list(encoder.rerank(query, documents))
 # [-11.48061752319336, 5.472434997558594]
 ```
 
+Text cross encoders can also be extended with models which are not in the list of supported models.
+
+```python
+from fastembed.rerank.cross_encoder import TextCrossEncoder 
+from fastembed.common.model_description import ModelSource
+
+TextCrossEncoder.add_custom_model(
+    model="Xenova/ms-marco-MiniLM-L-4-v2",
+    model_file="onnx/model.onnx",
+    sources=ModelSource(hf="Xenova/ms-marco-MiniLM-L-4-v2"),
+)
+model = TextCrossEncoder(model_name="Xenova/ms-marco-MiniLM-L-4-v2")
+scores = list(model.rerank_pairs(
+    [("What is AI?", "Artificial intelligence is ..."), ("What is ML?", "Machine learning is ..."),]
+))
+```
+
 ## ⚡️ FastEmbed on a GPU
 
 FastEmbed supports running on GPU devices.

fastembed/common/model_description.py

@@ -7,6 +7,11 @@ from typing import Optional, Any
 class ModelSource:
     hf: Optional[str] = None
     url: Optional[str] = None
+    _deprecated_tar_struct: bool = False
+
+    @property
+    def deprecated_tar_struct(self) -> bool:
+        return self._deprecated_tar_struct
 
     def __post_init__(self) -> None:
         if self.hf is None and self.url is None:

fastembed/common/model_management.py

@@ -330,9 +330,10 @@ class ModelManagement(Generic[T]):
         model_name: str,
         source_url: str,
         cache_dir: str,
+        deprecated_tar_struct: bool = False,
         local_files_only: bool = False,
     ) -> Path:
-        fast_model_name = f"fast-{model_name.split('/')[-1]}"
+        fast_model_name = f"{'fast-' if deprecated_tar_struct else ''}{model_name.split('/')[-1]}"
         cache_tmp_dir = Path(cache_dir) / "tmp"
         model_tmp_dir = cache_tmp_dir / fast_model_name
         model_dir = Path(cache_dir) / fast_model_name
@@ -438,6 +439,7 @@ class ModelManagement(Generic[T]):
                         model.model,
                         str(url_source),
                         str(cache_dir),
+                        deprecated_tar_struct=model.sources.deprecated_tar_struct,
                         local_files_only=local_files_only,
                     )
                 except Exception:

fastembed/common/onnx_model.py

@@ -68,15 +68,7 @@ class OnnxModel(Generic[T]):
             if device_id is None:
                 onnx_providers = ["CUDAExecutionProvider"]
             else:
-                # kSameAsRequested: Allocates only the requested memory, avoiding over-allocation.
-                # more precise than 'kNextPowerOfTwo', which grows memory aggressively.
-                # source: https://onnxruntime.ai/docs/get-started/with-c.html#features:~:text=Memory%20arena%20shrinkage:
-                onnx_providers = [
-                    (
-                        "CUDAExecutionProvider",
-                        {"device_id": device_id, "arena_extend_strategy": "kSameAsRequested"},
-                    )
-                ]
+                onnx_providers = [("CUDAExecutionProvider", {"device_id": device_id})]
         else:
             onnx_providers = ["CPUExecutionProvider"]
 
@@ -140,7 +132,5 @@ class EmbeddingWorker(Worker, Generic[T]):
     def start(cls, model_name: str, cache_dir: str, **kwargs: Any) -> "EmbeddingWorker[T]":
         return cls(model_name=model_name, cache_dir=cache_dir, **kwargs)
 
-    def process(
-        self, items: Iterable[tuple[int, Any]], **kwargs: Any
-    ) -> Iterable[tuple[int, Any]]:
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         raise NotImplementedError("Subclasses must implement this method")

fastembed/common/utils.py

@@ -5,12 +5,12 @@ import tempfile
 import unicodedata
 from pathlib import Path
 from itertools import islice
-from typing import Iterable, Optional, TypeVar, Sequence
+from typing import Iterable, Optional, TypeVar
 
 import numpy as np
 from numpy.typing import NDArray
 
-from fastembed.common.types import NumpyArray, OnnxProvider
+from fastembed.common.types import NumpyArray
 
 T = TypeVar("T")
 
@@ -67,18 +67,3 @@ def get_all_punctuation() -> set[str]:
 
 def remove_non_alphanumeric(text: str) -> str:
     return re.sub(r"[^\w\s]", " ", text, flags=re.UNICODE)
-
-
-def is_cuda_enabled(cuda: bool, providers: Optional[Sequence[OnnxProvider]]) -> bool:
-    """
-    Check if CUDA is enabled based on the `cuda` and `providers` parameters
-    """
-    if cuda:
-        return True
-    if not providers:
-        return False
-    if isinstance(providers, str):
-        return "CUDAExecutionProvider" in providers
-    return isinstance(providers, (list, tuple)) and any(
-        isinstance(p, str) and "CUDAExecutionProvider" in p for p in providers
-    )

fastembed/image/onnx_embedding.py

@@ -1,6 +1,5 @@
 from typing import Any, Iterable, Optional, Sequence, Type, Union
 
-import numpy as np
 
 from fastembed.common.types import NumpyArray
 from fastembed.common import ImageInput, OnnxProvider
@@ -195,7 +194,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
         return onnx_input
 
     def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
-        return normalize(output.model_output).astype(np.float32)
+        return normalize(output.model_output)
 
 
 class OnnxImageEmbeddingWorker(ImageEmbeddingWorker[NumpyArray]):

fastembed/image/onnx_image_model.py

@@ -6,14 +6,13 @@ from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import numpy as np
 from PIL import Image
-import onnxruntime as ort
 
 from fastembed.image.transform.operators import Compose
 from fastembed.common.types import NumpyArray
 from fastembed.common import ImageInput, OnnxProvider
 from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_preprocessor
-from fastembed.common.utils import iter_batch, is_cuda_enabled
+from fastembed.common.utils import iter_batch
 from fastembed.parallel_processor import ParallelWorkerPool
 
 # Holds type of the embedding result
@@ -75,21 +74,7 @@ class OnnxImageModel(OnnxModel[T]):
             encoded = np.array(self.processor(image_files))
         onnx_input = self._build_onnx_input(encoded)
         onnx_input = self._preprocess_onnx_input(onnx_input)
-
-        run_options = ort.RunOptions()
-        providers = kwargs.get("providers", None)
-        cuda = kwargs.get("cuda", False)
-        if is_cuda_enabled(cuda, providers):
-            device_id = kwargs.get("device_id", None)
-            device_id = str(device_id if isinstance(device_id, int) else 0)
-            # enables memory arena shrinkage, freeing unused memory after each Run() cycle.
-            # helps prevent excessive memory retention, especially for dynamic workloads.
-            # source: https://onnxruntime.ai/docs/get-started/with-c.html#features:~:text=Memory%20arena%20shrinkage:
-            run_options.add_run_config_entry(
-                "memory.enable_memory_arena_shrinkage", f"gpu:{device_id}"
-            )
-
-        model_output = self.model.run(None, onnx_input, run_options)  # type: ignore[union-attr]
+        model_output = self.model.run(None, onnx_input)  # type: ignore[union-attr]
         embeddings = model_output[0].reshape(len(images), -1)
         return OnnxOutputContext(model_output=embeddings)
 
@@ -119,9 +104,7 @@ class OnnxImageModel(OnnxModel[T]):
                 self.load_onnx_model()
 
             for batch in iter_batch(images, batch_size):
-                yield from self._post_process_onnx_output(
-                    self.onnx_embed(batch, cuda=cuda, providers=providers)
-                )
+                yield from self._post_process_onnx_output(self.onnx_embed(batch))
         else:
             if parallel == 0:
                 parallel = os.cpu_count()
@@ -146,9 +129,7 @@ class OnnxImageModel(OnnxModel[T]):
 
 
 class ImageEmbeddingWorker(EmbeddingWorker[T]):
-    def process(
-        self, items: Iterable[tuple[int, Any]], **kwargs: Any
-    ) -> Iterable[tuple[int, Any]]:
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         for idx, batch in items:
-            embeddings = self.model.onnx_embed(batch, **kwargs)
+            embeddings = self.model.onnx_embed(batch)
             yield idx, embeddings

fastembed/late_interaction/colbert.py

@@ -46,7 +46,7 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
         self, output: OnnxOutputContext, is_doc: bool = True
     ) -> Iterable[NumpyArray]:
         if not is_doc:
-            return output.model_output.astype(np.float32)
+            return output.model_output
 
         if output.input_ids is None or output.attention_mask is None:
             raise ValueError(
@@ -58,11 +58,11 @@ class Colbert(LateInteractionTextEmbeddingBase, OnnxTextModel[NumpyArray]):
                 if token_id in self.skip_list or token_id == self.pad_token_id:
                     output.attention_mask[i, j] = 0
 
-        output.model_output *= np.expand_dims(output.attention_mask, 2).astype(np.float32)
+        output.model_output *= np.expand_dims(output.attention_mask, 2)
         norm = np.linalg.norm(output.model_output, ord=2, axis=2, keepdims=True)
         norm_clamped = np.maximum(norm, 1e-12)
         output.model_output /= norm_clamped
-        return output.model_output.astype(np.float32)
+        return output.model_output
 
     def _preprocess_onnx_input(
         self, onnx_input: dict[str, NumpyArray], is_doc: bool = True, **kwargs: Any

fastembed/late_interaction/token_embeddings.py

@@ -0,0 +1,82 @@
+from typing import Union, Iterable, Optional, List, Dict, Any, Type
+
+import numpy as np
+
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.late_interaction.late_interaction_embedding_base import LateInteractionTextEmbeddingBase
+from fastembed.text.onnx_embedding import OnnxTextEmbedding
+from fastembed.text.onnx_text_model import TextEmbeddingWorker
+
+supported_token_embeddings_models = [
+    {
+        "model": "jinaai/jina-embeddings-v2-small-en-tokens",
+        "dim": 512,
+        "description": "Text embeddings, Unimodal (text), English, 8192 input tokens truncation,"
+                       " Prefixes for queries/documents: not necessary, 2023 year.",
+        "license": "apache-2.0",
+        "size_in_GB": 0.12,
+        "sources": {"hf": "xenova/jina-embeddings-v2-small-en"},
+        "model_file": "onnx/model.onnx",
+    },
+]
+
+
+class TokenEmbeddingsModel(OnnxTextEmbedding, LateInteractionTextEmbeddingBase):
+    @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_token_embeddings_models
+
+    @classmethod
+    def _get_worker_class(cls) -> Type[TextEmbeddingWorker]:
+        return TokensEmbeddingWorker
+
+    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[np.ndarray]:
+        # Size: (batch_size, sequence_length, hidden_size)
+        embeddings = output.model_output
+        # Size: (batch_size, sequence_length)
+        masks = output.attention_mask
+
+        # For each document we only select those embeddings that are not masked out
+
+        for i in range(embeddings.shape[0]):
+            yield embeddings[i, masks[i] == 1]
+
+    def embed(
+            self,
+            documents: Union[str, Iterable[str]],
+            batch_size: int = 256,
+            parallel: Optional[int] = None,
+            **kwargs,
+    ) -> Iterable[np.ndarray]:
+        yield from OnnxTextEmbedding.embed(self, documents, batch_size=batch_size, parallel=parallel, **kwargs)
+
+    def tokenize_docs(self, documents: List[str]) -> List[np.ndarray]:
+        encoded = self.tokenizer.encode_batch(documents)
+        return [e.ids for e in encoded]
+
+
+class TokensEmbeddingWorker(TextEmbeddingWorker):
+    def init_embedding(self, model_name: str, cache_dir: str, **kwargs) -> TokenEmbeddingsModel:
+        return TokenEmbeddingsModel(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )
+
+
+if __name__ == "__main__":
+    # Example usage
+    model = TokenEmbeddingsModel(model_name="jinaai/jina-embeddings-v2-small-en-tokens")
+    docs = ["Hello, world!", "hello", "hello hello"]
+
+    embeddings = model.embed(docs)
+    for emb in embeddings:
+        print(emb.shape)
+
+    print(model.tokenize_docs(docs))

fastembed/late_interaction_multimodal/colpali.py

@@ -142,7 +142,7 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
         assert self.model_description.dim is not None, "Model dim is not defined"
         return output.model_output.reshape(
             output.model_output.shape[0], -1, self.model_description.dim
-        ).astype(np.float32)
+        )
 
     def _post_process_onnx_text_output(
         self,
@@ -157,7 +157,7 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
         Returns:
             Iterable[NumpyArray]: Post-processed output as NumPy arrays.
         """
-        return output.model_output.astype(np.float32)
+        return output.model_output
 
     def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
         texts_query: list[str] = []

fastembed/late_interaction_multimodal/onnx_multimodal_model.py

@@ -6,14 +6,13 @@ from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import numpy as np
 from PIL import Image
-import onnxruntime as ort
 from tokenizers import Encoding, Tokenizer
 
 from fastembed.common import OnnxProvider, ImageInput
 from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_tokenizer, load_preprocessor
 from fastembed.common.types import NumpyArray
-from fastembed.common.utils import iter_batch, is_cuda_enabled
+from fastembed.common.utils import iter_batch
 from fastembed.image.transform.operators import Compose
 from fastembed.parallel_processor import ParallelWorkerPool
 
@@ -104,21 +103,7 @@ class OnnxMultimodalModel(OnnxModel[T]):
             )
 
         onnx_input = self._preprocess_onnx_text_input(onnx_input, **kwargs)
-
-        run_options = ort.RunOptions()
-        providers = kwargs.get("providers", None)
-        cuda = kwargs.get("cuda", False)
-        if is_cuda_enabled(cuda, providers):
-            device_id = kwargs.get("device_id", None)
-            device_id = str(device_id if isinstance(device_id, int) else 0)
-            # enables memory arena shrinkage, freeing unused memory after each Run() cycle.
-            # helps prevent excessive memory retention, especially for dynamic workloads.
-            # source: https://onnxruntime.ai/docs/get-started/with-c.html#features:~:text=Memory%20arena%20shrinkage:
-            run_options.add_run_config_entry(
-                "memory.enable_memory_arena_shrinkage", f"gpu:{device_id}"
-            )
-
-        model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input, run_options)  # type: ignore[union-attr]
+        model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)  # type: ignore[union-attr]
         return OnnxOutputContext(
             model_output=model_output[0],
             attention_mask=onnx_input.get("attention_mask", attention_mask),
@@ -151,9 +136,7 @@ class OnnxMultimodalModel(OnnxModel[T]):
             if not hasattr(self, "model") or self.model is None:
                 self.load_onnx_model()
             for batch in iter_batch(documents, batch_size):
-                yield from self._post_process_onnx_text_output(
-                    self.onnx_embed_text(batch, cuda=cuda, providers=providers)
-                )
+                yield from self._post_process_onnx_text_output(self.onnx_embed_text(batch))
         else:
             if parallel == 0:
                 parallel = os.cpu_count()
@@ -186,21 +169,7 @@ class OnnxMultimodalModel(OnnxModel[T]):
             encoded = np.array(self.processor(image_files))
         onnx_input = {"pixel_values": encoded}
         onnx_input = self._preprocess_onnx_image_input(onnx_input, **kwargs)
-
-        run_options = ort.RunOptions()
-        providers = kwargs.get("providers", None)
-        cuda = kwargs.get("cuda", False)
-        if is_cuda_enabled(cuda, providers):
-            device_id = kwargs.get("device_id", None)
-            device_id = str(device_id if isinstance(device_id, int) else 0)
-            # enables memory arena shrinkage, freeing unused memory after each Run() cycle.
-            # helps prevent excessive memory retention, especially for dynamic workloads.
-            # source: https://onnxruntime.ai/docs/get-started/with-c.html#features:~:text=Memory%20arena%20shrinkage:
-            run_options.add_run_config_entry(
-                "memory.enable_memory_arena_shrinkage", f"gpu:{device_id}"
-            )
-
-        model_output = self.model.run(None, onnx_input, run_options)  # type: ignore[union-attr]
+        model_output = self.model.run(None, onnx_input)  # type: ignore[union-attr]
         embeddings = model_output[0].reshape(len(images), -1)
         return OnnxOutputContext(model_output=embeddings)
 
@@ -230,9 +199,7 @@ class OnnxMultimodalModel(OnnxModel[T]):
                 self.load_onnx_model()
 
             for batch in iter_batch(images, batch_size):
-                yield from self._post_process_onnx_image_output(
-                    self.onnx_embed_image(batch, cuda=cuda, providers=providers)
-                )
+                yield from self._post_process_onnx_image_output(self.onnx_embed_image(batch))
         else:
             if parallel == 0:
                 parallel = os.cpu_count()
@@ -274,11 +241,9 @@ class TextEmbeddingWorker(EmbeddingWorker[T]):
     ) -> OnnxMultimodalModel:
         raise NotImplementedError()
 
-    def process(
-        self, items: Iterable[tuple[int, Any]], **kwargs: Any
-    ) -> Iterable[tuple[int, Any]]:
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         for idx, batch in items:
-            onnx_output = self.model.onnx_embed_text(batch, **kwargs)
+            onnx_output = self.model.onnx_embed_text(batch)
             yield idx, onnx_output
 
 
@@ -300,9 +265,7 @@ class ImageEmbeddingWorker(EmbeddingWorker[T]):
     ) -> OnnxMultimodalModel:
         raise NotImplementedError()
 
-    def process(
-        self, items: Iterable[tuple[int, Any]], **kwargs: Any
-    ) -> Iterable[tuple[int, Any]]:
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         for idx, batch in items:
-            embeddings = self.model.onnx_embed_image(batch, **kwargs)
+            embeddings = self.model.onnx_embed_image(batch)
             yield idx, embeddings

fastembed/parallel_processor.py

@@ -28,9 +28,7 @@ class Worker:
     def start(cls, *args: Any, **kwargs: Any) -> "Worker":
         raise NotImplementedError()
 
-    def process(
-        self, items: Iterable[tuple[int, Any]], **kwargs: Any
-    ) -> Iterable[tuple[int, Any]]:
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         raise NotImplementedError()
 
 
@@ -65,7 +63,7 @@ def _worker(
                     break
                 yield item
 
-        for processed_item in worker.process(input_queue_iterable(), **kwargs):
+        for processed_item in worker.process(input_queue_iterable()):
             output_queue.put(processed_item)
     except Exception as e:  # pylint: disable=broad-except
         logging.exception(e)

fastembed/rerank/cross_encoder/custom_text_cross_encoder.py

@@ -0,0 +1,46 @@
+from typing import Optional, Sequence, Any
+
+from fastembed.common import OnnxProvider
+from fastembed.common.model_description import BaseModelDescription
+from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
+
+
+class CustomTextCrossEncoder(OnnxTextCrossEncoder):
+    SUPPORTED_MODELS: list[BaseModelDescription] = []
+
+    def __init__(
+        self,
+        model_name: str,
+        cache_dir: Optional[str] = None,
+        threads: Optional[int] = None,
+        providers: Optional[Sequence[OnnxProvider]] = None,
+        cuda: bool = False,
+        device_ids: Optional[list[int]] = None,
+        lazy_load: bool = False,
+        device_id: Optional[int] = None,
+        specific_model_path: Optional[str] = None,
+        **kwargs: Any,
+    ):
+        super().__init__(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=threads,
+            providers=providers,
+            cuda=cuda,
+            device_ids=device_ids,
+            lazy_load=lazy_load,
+            device_id=device_id,
+            specific_model_path=specific_model_path,
+            **kwargs,
+        )
+
+    @classmethod
+    def _list_supported_models(cls) -> list[BaseModelDescription]:
+        return cls.SUPPORTED_MODELS
+
+    @classmethod
+    def add_model(
+        cls,
+        model_description: BaseModelDescription,
+    ) -> None:
+        cls.SUPPORTED_MODELS.append(model_description)

fastembed/rerank/cross_encoder/onnx_text_model.py

@@ -4,7 +4,6 @@ from pathlib import Path
 from typing import Any, Iterable, Optional, Sequence, Type
 
 import numpy as np
-import onnxruntime as ort
 from tokenizers import Encoding
 
 from fastembed.common.onnx_model import (
@@ -15,7 +14,7 @@ from fastembed.common.onnx_model import (
 )
 from fastembed.common.types import NumpyArray
 from fastembed.common.preprocessor_utils import load_tokenizer
-from fastembed.common.utils import iter_batch, is_cuda_enabled
+from fastembed.common.utils import iter_batch
 from fastembed.parallel_processor import ParallelWorkerPool
 
 
@@ -72,21 +71,7 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
         tokenized_input = self.tokenize(pairs, **kwargs)
         inputs = self._build_onnx_input(tokenized_input)
         onnx_input = self._preprocess_onnx_input(inputs, **kwargs)
-
-        run_options = ort.RunOptions()
-        providers = kwargs.get("providers", None)
-        cuda = kwargs.get("cuda", False)
-        if is_cuda_enabled(cuda, providers):
-            device_id = kwargs.get("device_id", None)
-            device_id = str(device_id if isinstance(device_id, int) else 0)
-            # Enables memory arena shrinkage, freeing unused memory after each Run() cycle.
-            # Helps prevent excessive memory retention, especially for dynamic workloads.
-            # Source: https://onnxruntime.ai/docs/get-started/with-c.html#features:~:text=Memory%20arena%20shrinkage:
-            run_options.add_run_config_entry(
-                "memory.enable_memory_arena_shrinkage", f"gpu:{device_id}"
-            )
-
-        outputs = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input, run_options)  # type: ignore[union-attr]
+        outputs = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)  # type: ignore[union-attr]
         relevant_output = outputs[0]
         scores: NumpyArray = relevant_output[:, 0]
         return OnnxOutputContext(model_output=scores)
@@ -125,9 +110,7 @@ class OnnxCrossEncoderModel(OnnxModel[float]):
             if not hasattr(self, "model") or self.model is None:
                 self.load_onnx_model()
             for batch in iter_batch(pairs, batch_size):
-                yield from self._post_process_onnx_output(
-                    self.onnx_embed_pairs(batch, cuda=cuda, providers=providers, **kwargs)
-                )
+                yield from self._post_process_onnx_output(self.onnx_embed_pairs(batch, **kwargs))
         else:
             if parallel == 0:
                 parallel = os.cpu_count()
@@ -180,9 +163,7 @@ class TextRerankerWorker(EmbeddingWorker[float]):
     ) -> OnnxCrossEncoderModel:
         raise NotImplementedError()
 
-    def process(
-        self, items: Iterable[tuple[int, Any]], **kwargs: Any
-    ) -> Iterable[tuple[int, Any]]:
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, Any]]:
         for idx, batch in items:
-            onnx_output = self.model.onnx_embed_pairs(batch, **kwargs)
+            onnx_output = self.model.onnx_embed_pairs(batch)
             yield idx, onnx_output

fastembed/rerank/cross_encoder/text_cross_encoder.py

@@ -3,13 +3,19 @@ from dataclasses import asdict
 
 from fastembed.common import OnnxProvider
 from fastembed.rerank.cross_encoder.onnx_text_cross_encoder import OnnxTextCrossEncoder
+from fastembed.rerank.cross_encoder.custom_text_cross_encoder import CustomTextCrossEncoder
+
 from fastembed.rerank.cross_encoder.text_cross_encoder_base import TextCrossEncoderBase
-from fastembed.common.model_description import BaseModelDescription
+from fastembed.common.model_description import (
+    ModelSource,
+    BaseModelDescription,
+)
 
 
 class TextCrossEncoder(TextCrossEncoderBase):
     CROSS_ENCODER_REGISTRY: list[Type[TextCrossEncoderBase]] = [
         OnnxTextCrossEncoder,
+        CustomTextCrossEncoder,
     ]
 
     @classmethod
@@ -124,3 +130,34 @@ class TextCrossEncoder(TextCrossEncoderBase):
         yield from self.model.rerank_pairs(
             pairs, batch_size=batch_size, parallel=parallel, **kwargs
         )
+
+    @classmethod
+    def add_custom_model(
+        cls,
+        model: str,
+        sources: ModelSource,
+        model_file: str = "onnx/model.onnx",
+        description: str = "",
+        license: str = "",
+        size_in_gb: float = 0.0,
+        additional_files: Optional[list[str]] = None,
+    ) -> None:
+        registered_models = cls._list_supported_models()
+        for registered_model in registered_models:
+            if model == registered_model.model:
+                raise ValueError(
+                    f"Model {model} is already registered in CrossEncoderModel, if you still want to add this model, "
+                    f"please use another model name"
+                )
+
+        CustomTextCrossEncoder.add_model(
+            BaseModelDescription(
+                model=model,
+                sources=sources,
+                model_file=model_file,
+                description=description,
+                license=license,
+                size_in_GB=size_in_gb,
+                additional_files=additional_files or [],
+            )
+        )

fastembed/sparse/bm25.py

@@ -21,13 +21,9 @@ from fastembed.sparse.sparse_embedding_base import (
 from fastembed.sparse.utils.tokenizer import SimpleTokenizer
 from fastembed.common.model_description import SparseModelDescription, ModelSource
 
+
 supported_languages = [
     "arabic",
-    "azerbaijani",
-    "basque",
-    "bengali",
-    "catalan",
-    "chinese",
     "danish",
     "dutch",
     "english",
@@ -35,21 +31,15 @@ supported_languages = [
     "french",
     "german",
     "greek",
-    "hebrew",
-    "hinglish",
     "hungarian",
-    "indonesian",
     "italian",
-    "kazakh",
-    "nepali",
     "norwegian",
     "portuguese",
     "romanian",
     "russian",
-    "slovene",
     "spanish",
     "swedish",
-    "tajik",
+    "tamil",
     "turkish",
 ]
 
@@ -344,7 +334,7 @@ class Bm25Worker(Worker):
         return cls(model_name=model_name, cache_dir=cache_dir, **kwargs)
 
     def process(
-        self, items: Iterable[tuple[int, Any]], **kwargs: Any
+        self, items: Iterable[tuple[int, Any]]
     ) -> Iterable[tuple[int, list[SparseEmbedding]]]:
         for idx, batch in items:
             onnx_output = self.model.raw_embed(batch)

fastembed/sparse/splade_pp.py

@@ -18,7 +18,7 @@ supported_splade_models: list[SparseModelDescription] = [
         description="Independent Implementation of SPLADE++ Model for English.",
         license="apache-2.0",
         size_in_GB=0.532,
-        sources=ModelSource(hf="Qdrant/SPLADE_PP_en_v1"),
+        sources=ModelSource(hf="Qdrant/Splade_PP_en_v1"),
         model_file="model.onnx",
     ),
     SparseModelDescription(
@@ -27,7 +27,7 @@ supported_splade_models: list[SparseModelDescription] = [
         description="Independent Implementation of SPLADE++ Model for English.",
         license="apache-2.0",
         size_in_GB=0.532,
-        sources=ModelSource(hf="Qdrant/SPLADE_PP_en_v1"),
+        sources=ModelSource(hf="Qdrant/Splade_PP_en_v1"),
         model_file="model.onnx",
     ),
 ]

fastembed/text/multitask_embedding.py

@@ -3,6 +3,7 @@ from typing import Any, Type, Iterable, Union, Optional
 
 import numpy as np
 
+from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.types import NumpyArray
 from fastembed.text.pooled_normalized_embedding import PooledNormalizedEmbedding
 from fastembed.text.onnx_embedding import OnnxTextEmbeddingWorker
@@ -44,9 +45,11 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
     PASSAGE_TASK = Task.RETRIEVAL_PASSAGE
     QUERY_TASK = Task.RETRIEVAL_QUERY
 
-    def __init__(self, *args: Any, **kwargs: Any):
+    def __init__(self, *args: Any, task_id: Optional[int] = None, **kwargs: Any):
         super().__init__(*args, **kwargs)
-        self.current_task_id: Union[Task, int] = self.PASSAGE_TASK
+        self.default_task_id: Union[Task, int] = (
+            task_id if task_id is not None else self.PASSAGE_TASK
+        )
 
     @classmethod
     def _get_worker_class(cls) -> Type[OnnxTextEmbeddingWorker]:
@@ -57,9 +60,14 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
         return supported_multitask_models
 
     def _preprocess_onnx_input(
-        self, onnx_input: dict[str, NumpyArray], **kwargs: Any
+        self,
+        onnx_input: dict[str, NumpyArray],
+        task_id: Optional[Union[int, Task]] = None,
+        **kwargs: Any,
     ) -> dict[str, NumpyArray]:
-        onnx_input["task_id"] = np.array(self.current_task_id, dtype=np.int64)
+        if task_id is None:
+            raise ValueError(f"task_id must be provided for JinaEmbeddingV3, got <{task_id}>")
+        onnx_input["task_id"] = np.array(task_id, dtype=np.int64)
         return onnx_input
 
     def embed(
@@ -67,20 +75,19 @@ class JinaEmbeddingV3(PooledNormalizedEmbedding):
         documents: Union[str, Iterable[str]],
         batch_size: int = 256,
         parallel: Optional[int] = None,
-        task_id: int = PASSAGE_TASK,
+        task_id: Optional[int] = None,
         **kwargs: Any,
     ) -> Iterable[NumpyArray]:
-        self.current_task_id = task_id
-        kwargs["task_id"] = task_id
-        yield from super().embed(documents, batch_size, parallel, **kwargs)
+        task_id = (
+            task_id if task_id is not None else self.default_task_id
+        )  # required for multiprocessing
+        yield from super().embed(documents, batch_size, parallel, task_id=task_id, **kwargs)
 
     def query_embed(self, query: Union[str, Iterable[str]], **kwargs: Any) -> Iterable[NumpyArray]:
-        self.current_task_id = self.QUERY_TASK
-        yield from super().embed(query, **kwargs)
+        yield from super().embed(query, task_id=self.QUERY_TASK, **kwargs)
 
     def passage_embed(self, texts: Iterable[str], **kwargs: Any) -> Iterable[NumpyArray]:
-        self.current_task_id = self.PASSAGE_TASK
-        yield from super().embed(texts, **kwargs)
+        yield from super().embed(texts, task_id=self.PASSAGE_TASK, **kwargs)
 
 
 class JinaEmbeddingV3Worker(OnnxTextEmbeddingWorker):
@@ -90,11 +97,15 @@ class JinaEmbeddingV3Worker(OnnxTextEmbeddingWorker):
         cache_dir: str,
         **kwargs: Any,
     ) -> JinaEmbeddingV3:
-        model = JinaEmbeddingV3(
+        return JinaEmbeddingV3(
             model_name=model_name,
             cache_dir=cache_dir,
             threads=1,
             **kwargs,
         )
-        model.current_task_id = kwargs["task_id"]
-        return model
+
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, OnnxOutputContext]]:
+        self.model: JinaEmbeddingV3  # mypy complaints `self.model` does not have `default_task_id`
+        for idx, batch in items:
+            onnx_output = self.model.onnx_embed(batch, task_id=self.model.default_task_id)
+            yield idx, onnx_output

fastembed/text/onnx_embedding.py

@@ -1,6 +1,5 @@
 from typing import Any, Iterable, Optional, Sequence, Type, Union
 
-import numpy as np
 from fastembed.common.types import NumpyArray, OnnxProvider
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import define_cache_dir, normalize
@@ -21,6 +20,7 @@ supported_onnx_models: list[DenseModelDescription] = [
         sources=ModelSource(
             hf="Qdrant/fast-bge-base-en",
             url="https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz",
+            _deprecated_tar_struct=True,
         ),
         model_file="model_optimized.onnx",
     ),
@@ -36,6 +36,7 @@ supported_onnx_models: list[DenseModelDescription] = [
         sources=ModelSource(
             hf="qdrant/bge-base-en-v1.5-onnx-q",
             url="https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en-v1.5.tar.gz",
+            _deprecated_tar_struct=True,
         ),
         model_file="model_optimized.onnx",
     ),
@@ -63,6 +64,7 @@ supported_onnx_models: list[DenseModelDescription] = [
         sources=ModelSource(
             hf="Qdrant/bge-small-en",
             url="https://storage.googleapis.com/qdrant-fastembed/BAAI-bge-small-en.tar.gz",
+            _deprecated_tar_struct=True,
         ),
         model_file="model_optimized.onnx",
     ),
@@ -90,6 +92,7 @@ supported_onnx_models: list[DenseModelDescription] = [
         sources=ModelSource(
             hf="Qdrant/bge-small-zh-v1.5",
             url="https://storage.googleapis.com/qdrant-fastembed/fast-bge-small-zh-v1.5.tar.gz",
+            _deprecated_tar_struct=True,
         ),
         model_file="model_optimized.onnx",
     ),
@@ -309,7 +312,7 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
             processed_embeddings = embeddings
         else:
             raise ValueError(f"Unsupported embedding shape: {embeddings.shape}")
-        return normalize(processed_embeddings).astype(np.float32)
+        return normalize(processed_embeddings)
 
     def load_onnx_model(self) -> None:
         self._load_onnx_model(

fastembed/text/onnx_text_model.py

@@ -4,14 +4,13 @@ from pathlib import Path
 from typing import Any, Iterable, Optional, Sequence, Type, Union
 
 import numpy as np
-import onnxruntime as ort
 from numpy.typing import NDArray
 from tokenizers import Encoding, Tokenizer
 
 from fastembed.common.types import NumpyArray, OnnxProvider
 from fastembed.common.onnx_model import EmbeddingWorker, OnnxModel, OnnxOutputContext, T
 from fastembed.common.preprocessor_utils import load_tokenizer
-from fastembed.common.utils import iter_batch, is_cuda_enabled
+from fastembed.common.utils import iter_batch
 from fastembed.parallel_processor import ParallelWorkerPool
 
 
@@ -83,21 +82,7 @@ class OnnxTextModel(OnnxModel[T]):
             )
         onnx_input = self._preprocess_onnx_input(onnx_input, **kwargs)
 
-        run_options = ort.RunOptions()
-        providers = kwargs.get("providers", None)
-        cuda = kwargs.get("cuda", False)
-
-        if is_cuda_enabled(cuda, providers):
-            device_id = kwargs.get("device_id", None)
-            device_id = str(device_id if isinstance(device_id, int) else 0)
-            # enables memory arena shrinkage, freeing unused memory after each Run() cycle.
-            # helps prevent excessive memory retention, especially for dynamic workloads.
-            # source: https://onnxruntime.ai/docs/get-started/with-c.html#features:~:text=Memory%20arena%20shrinkage:
-            run_options.add_run_config_entry(
-                "memory.enable_memory_arena_shrinkage", f"gpu:{device_id}"
-            )
-
-        model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input, run_options)  # type: ignore[union-attr]
+        model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)  # type: ignore[union-attr]
         return OnnxOutputContext(
             model_output=model_output[0],
             attention_mask=onnx_input.get("attention_mask", attention_mask),
@@ -130,9 +115,7 @@ class OnnxTextModel(OnnxModel[T]):
             if not hasattr(self, "model") or self.model is None:
                 self.load_onnx_model()
             for batch in iter_batch(documents, batch_size):
-                yield from self._post_process_onnx_output(
-                    self.onnx_embed(batch, cuda=cuda, providers=providers)
-                )
+                yield from self._post_process_onnx_output(self.onnx_embed(batch, **kwargs))
         else:
             if parallel == 0:
                 parallel = os.cpu_count()
@@ -157,9 +140,7 @@ class OnnxTextModel(OnnxModel[T]):
 
 
 class TextEmbeddingWorker(EmbeddingWorker[T]):
-    def process(
-        self, items: Iterable[tuple[int, Any]], **kwargs: Any
-    ) -> Iterable[tuple[int, OnnxOutputContext]]:
+    def process(self, items: Iterable[tuple[int, Any]]) -> Iterable[tuple[int, OnnxOutputContext]]:
         for idx, batch in items:
-            onnx_output = self.model.onnx_embed(batch, **kwargs)
+            onnx_output = self.model.onnx_embed(batch)
             yield idx, onnx_output

fastembed/text/pooled_embedding.py

@@ -82,6 +82,7 @@ supported_pooled_models: list[DenseModelDescription] = [
         sources=ModelSource(
             hf="qdrant/multilingual-e5-large-onnx",
             url="https://storage.googleapis.com/qdrant-fastembed/fast-multilingual-e5-large.tar.gz",
+            _deprecated_tar_struct=True,
         ),
         model_file="model.onnx",
         additional_files=["model.onnx_data"],
@@ -115,7 +116,7 @@ class PooledEmbedding(OnnxTextEmbedding):
 
         embeddings = output.model_output
         attn_mask = output.attention_mask
-        return self.mean_pooling(embeddings, attn_mask).astype(np.float32)
+        return self.mean_pooling(embeddings, attn_mask)
 
 
 class PooledEmbeddingWorker(OnnxTextEmbeddingWorker):

fastembed/text/pooled_normalized_embedding.py

@@ -1,6 +1,5 @@
 from typing import Any, Iterable, Type
 
-import numpy as np
 
 from fastembed.common.types import NumpyArray
 from fastembed.common.onnx_model import OnnxOutputContext
@@ -22,6 +21,7 @@ supported_pooled_normalized_models: list[DenseModelDescription] = [
         sources=ModelSource(
             url="https://storage.googleapis.com/qdrant-fastembed/sentence-transformers-all-MiniLM-L6-v2.tar.gz",
             hf="qdrant/all-MiniLM-L6-v2-onnx",
+            _deprecated_tar_struct=True,
         ),
         model_file="model.onnx",
     ),
@@ -144,7 +144,7 @@ class PooledNormalizedEmbedding(PooledEmbedding):
 
         embeddings = output.model_output
         attn_mask = output.attention_mask
-        return normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
+        return normalize(self.mean_pooling(embeddings, attn_mask))
 
 
 class PooledNormalizedEmbeddingWorker(OnnxTextEmbeddingWorker):

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
-name = "fastembed-gpu"
-version = "0.6.0"
+name = "fastembed"
+version = "0.6.1"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["Qdrant Team <info@qdrant.tech>", "NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
@@ -18,7 +18,7 @@ numpy = [
     { version = ">=2.1.0", python = ">=3.13" },
     { version = ">=1.21,<2.1.0", python = "<3.10" },
 ]
-onnxruntime-gpu = [
+onnxruntime = [
     { version = ">1.20.0", python = ">=3.13" },
     { version = ">=1.17.0,<1.20.0", python = "<3.10" },
     { version = ">=1.17.0,!=1.20.0", python = ">=3.10,<3.13" },

tests/test_custom_models.py

@@ -3,10 +3,17 @@ import os
 import numpy as np
 import pytest
 
-from fastembed.common.model_description import PoolingType, ModelSource, DenseModelDescription
+from fastembed.common.model_description import (
+    PoolingType,
+    ModelSource,
+    DenseModelDescription,
+    BaseModelDescription,
+)
 from fastembed.common.onnx_model import OnnxOutputContext
 from fastembed.common.utils import normalize, mean_pooling
 from fastembed.text.custom_text_embedding import CustomTextEmbedding, PostprocessingConfig
+from fastembed.rerank.cross_encoder.custom_text_cross_encoder import CustomTextCrossEncoder
+from fastembed.rerank.cross_encoder import TextCrossEncoder
 from fastembed.text.text_embedding import TextEmbedding
 from tests.utils import delete_model_cache
 
@@ -14,8 +21,10 @@ from tests.utils import delete_model_cache
 @pytest.fixture(autouse=True)
 def restore_custom_models_fixture():
     CustomTextEmbedding.SUPPORTED_MODELS = []
+    CustomTextCrossEncoder.SUPPORTED_MODELS = []
     yield
     CustomTextEmbedding.SUPPORTED_MODELS = []
+    CustomTextCrossEncoder.SUPPORTED_MODELS = []
 
 
 def test_text_custom_model():
@@ -65,6 +74,43 @@ def test_text_custom_model():
         delete_model_cache(model.model._model_dir)
 
 
+def test_cross_encoder_custom_model():
+    is_ci = os.getenv("CI")
+    custom_model_name = "Xenova/ms-marco-MiniLM-L-4-v2"
+    size_in_gb = 0.08
+    source = ModelSource(hf=custom_model_name)
+    canonical_vector = np.array([-5.7170815, -11.112114], dtype=np.float32)
+
+    TextCrossEncoder.add_custom_model(
+        custom_model_name,
+        model_file="onnx/model.onnx",
+        sources=source,
+        size_in_gb=size_in_gb,
+    )
+
+    assert CustomTextCrossEncoder.SUPPORTED_MODELS[0] == BaseModelDescription(
+        model=custom_model_name,
+        sources=source,
+        model_file="onnx/model.onnx",
+        description="",
+        license="",
+        size_in_GB=size_in_gb,
+    )
+
+    model = TextCrossEncoder(custom_model_name)
+    pairs = [
+        ("What is AI?", "Artificial intelligence is ..."),
+        ("What is ML?", "Machine learning is ..."),
+    ]
+    scores = list(model.rerank_pairs(pairs))
+
+    embeddings = np.stack(scores, axis=0)
+    assert embeddings.shape == (2,)
+    assert np.allclose(embeddings, canonical_vector, atol=1e-3)
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
+
+
 def test_mock_add_custom_models():
     dim = 5
     size_in_gb = 0.1
@@ -91,15 +137,11 @@ def test_mock_add_custom_models():
     expected_output = {
         f"{PoolingType.MEAN.lower()}-normalized": normalize(
             mean_pooling(dummy_token_embedding, dummy_attention_mask)
-        ).astype(np.float32),
+        ),
         f"{PoolingType.MEAN.lower()}": mean_pooling(dummy_token_embedding, dummy_attention_mask),
-        f"{PoolingType.CLS.lower()}-normalized": normalize(dummy_token_embedding[:, 0]).astype(
-            np.float32
-        ),
+        f"{PoolingType.CLS.lower()}-normalized": normalize(dummy_token_embedding[:, 0]),
         f"{PoolingType.CLS.lower()}": dummy_token_embedding[:, 0],
-        f"{PoolingType.DISABLED.lower()}-normalized": normalize(dummy_pooled_embedding).astype(
-            np.float32
-        ),
+        f"{PoolingType.DISABLED.lower()}-normalized": normalize(dummy_pooled_embedding),
         f"{PoolingType.DISABLED.lower()}": dummy_pooled_embedding,
     }
 
@@ -160,3 +202,28 @@ def test_do_not_add_existing_model():
             dim=384,
             size_in_gb=0.47,
         )
+
+
+def test_do_not_add_existing_cross_encoder():
+    existing_base_model = "Xenova/ms-marco-MiniLM-L-6-v2"
+    custom_model_name = "Xenova/ms-marco-MiniLM-L-4-v2"
+
+    with pytest.raises(ValueError, match=f"Model {existing_base_model} is already registered"):
+        TextCrossEncoder.add_custom_model(
+            existing_base_model,
+            sources=ModelSource(hf=existing_base_model),
+            size_in_gb=0.08,
+        )
+
+    TextCrossEncoder.add_custom_model(
+        custom_model_name,
+        sources=ModelSource(hf=existing_base_model),
+        size_in_gb=0.08,
+    )
+
+    with pytest.raises(ValueError, match=f"Model {custom_model_name} is already registered"):
+        TextCrossEncoder.add_custom_model(
+            custom_model_name,
+            sources=ModelSource(hf=custom_model_name),
+            size_in_gb=0.08,
+        )

tests/test_image_onnx_embeddings.py

@@ -8,7 +8,7 @@ from PIL import Image
 
 from fastembed import ImageEmbedding
 from tests.config import TEST_MISC_DIR
-from tests.utils import delete_model_cache
+from tests.utils import delete_model_cache, should_test_model
 
 CANONICAL_VECTOR_VALUES = {
     "Qdrant/clip-ViT-B-32-vision": np.array([-0.0098, 0.0128, -0.0274, 0.002, -0.0059]),
@@ -27,11 +27,13 @@ CANONICAL_VECTOR_VALUES = {
 }
 
 
-def test_embedding() -> None:
+@pytest.mark.parametrize("model_name", ["Qdrant/clip-ViT-B-32-vision"])
+def test_embedding(model_name: str) -> None:
     is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
 
     for model_desc in ImageEmbedding._list_supported_models():
-        if not is_ci and model_desc.size_in_GB > 1:
+        if not should_test_model(model_desc, model_name, is_ci, is_manual):
             continue
 
         dim = model_desc.dim
@@ -74,8 +76,12 @@ def test_batch_embedding(n_dims: int, model_name: str) -> None:
 
     embeddings = list(model.embed(images, batch_size=10))
     embeddings = np.stack(embeddings, axis=0)
+    assert np.allclose(embeddings[1], embeddings[2])
+
+    canonical_vector = CANONICAL_VECTOR_VALUES[model_name]
 
     assert embeddings.shape == (len(test_images) * n_images, n_dims)
+    assert np.allclose(embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3)
     if is_ci:
         delete_model_cache(model.model._model_dir)
 

tests/test_late_interaction_embeddings.py

@@ -6,7 +6,7 @@ import numpy as np
 from fastembed.late_interaction.late_interaction_text_embedding import (
     LateInteractionTextEmbedding,
 )
-from tests.utils import delete_model_cache
+from tests.utils import delete_model_cache, should_test_model
 
 # vectors are abridged and rounded for brevity
 CANONICAL_COLUMN_VALUES = {
@@ -153,31 +153,37 @@ CANONICAL_QUERY_VALUES = {
 docs = ["Hello World"]
 
 
-def test_batch_embedding():
+@pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
+def test_batch_embedding(model_name: str):
     is_ci = os.getenv("CI")
     docs_to_embed = docs * 10
 
-    for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
-        print("evaluating", model_name)
-        model = LateInteractionTextEmbedding(model_name=model_name)
-        result = list(model.embed(docs_to_embed, batch_size=6))
+    model = LateInteractionTextEmbedding(model_name=model_name)
+    result = list(model.embed(docs_to_embed, batch_size=6))
+    expected_result = CANONICAL_COLUMN_VALUES[model_name]
 
-        for value in result:
-            token_num, abridged_dim = expected_result.shape
-            assert np.allclose(value[:, :abridged_dim], expected_result, atol=2e-3)
+    for value in result:
+        token_num, abridged_dim = expected_result.shape
+        assert np.allclose(value[:, :abridged_dim], expected_result, atol=2e-3)
 
-        if is_ci:
-            delete_model_cache(model.model._model_dir)
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
 
 
-def test_single_embedding():
+@pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
+def test_single_embedding(model_name: str):
     is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
     docs_to_embed = docs
 
-    for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
+    for model_desc in LateInteractionTextEmbedding._list_supported_models():
+        if not should_test_model(model_desc, model_name, is_ci, is_manual):
+            continue
+
         print("evaluating", model_name)
         model = LateInteractionTextEmbedding(model_name=model_name)
         result = next(iter(model.embed(docs_to_embed, batch_size=6)))
+        expected_result = CANONICAL_COLUMN_VALUES[model_name]
         token_num, abridged_dim = expected_result.shape
         assert np.allclose(result[:, :abridged_dim], expected_result, atol=2e-3)
 
@@ -185,14 +191,20 @@ def test_single_embedding():
             delete_model_cache(model.model._model_dir)
 
 
-def test_single_embedding_query():
+@pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
+def test_single_embedding_query(model_name: str):
     is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
     queries_to_embed = docs
 
-    for model_name, expected_result in CANONICAL_QUERY_VALUES.items():
+    for model_desc in LateInteractionTextEmbedding._list_supported_models():
+        if not should_test_model(model_desc, model_name, is_ci, is_manual):
+            continue
+
         print("evaluating", model_name)
         model = LateInteractionTextEmbedding(model_name=model_name)
         result = next(iter(model.query_embed(queries_to_embed)))
+        expected_result = CANONICAL_QUERY_VALUES[model_name]
         token_num, abridged_dim = expected_result.shape
         assert np.allclose(result[:, :abridged_dim], expected_result, atol=2e-3)
 
@@ -200,10 +212,11 @@ def test_single_embedding_query():
             delete_model_cache(model.model._model_dir)
 
 
-def test_parallel_processing():
+@pytest.mark.parametrize("token_dim,model_name", [(96, "answerdotai/answerai-colbert-small-v1")])
+def test_parallel_processing(token_dim: int, model_name: str):
     is_ci = os.getenv("CI")
-    model = LateInteractionTextEmbedding(model_name="colbert-ir/colbertv2.0")
-    token_dim = 128
+    model = LateInteractionTextEmbedding(model_name=model_name)
+
     docs = ["hello world", "flag embedding"] * 100
     embeddings = list(model.embed(docs, batch_size=10, parallel=2))
     embeddings = np.stack(embeddings, axis=0)
@@ -222,10 +235,7 @@ def test_parallel_processing():
         delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize(
-    "model_name",
-    ["colbert-ir/colbertv2.0"],
-)
+@pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
 def test_lazy_load(model_name: str):
     is_ci = os.getenv("CI")
 

tests/test_late_interaction_multimodal.py

@@ -1,5 +1,6 @@
 import os
 
+import pytest
 from PIL import Image
 import numpy as np
 
@@ -45,38 +46,38 @@ images = [
 
 
 def test_batch_embedding():
-    is_ci = os.getenv("CI")
+    if os.getenv("CI"):
+        pytest.skip("Colpali is too large to test in CI")
 
-    if not is_ci:
-        for model_name, expected_result in CANONICAL_IMAGE_VALUES.items():
-            print("evaluating", model_name)
-            model = LateInteractionMultimodalEmbedding(model_name=model_name)
-            result = list(model.embed_image(images, batch_size=2))
+    for model_name, expected_result in CANONICAL_IMAGE_VALUES.items():
+        print("evaluating", model_name)
+        model = LateInteractionMultimodalEmbedding(model_name=model_name)
+        result = list(model.embed_image(images, batch_size=2))
 
-            for value in result:
-                token_num, abridged_dim = expected_result.shape
-                assert np.allclose(value[:token_num, :abridged_dim], expected_result, atol=2e-3)
+        for value in result:
+            token_num, abridged_dim = expected_result.shape
+            assert np.allclose(value[:token_num, :abridged_dim], expected_result, atol=2e-3)
 
 
 def test_single_embedding():
-    is_ci = os.getenv("CI")
-    if not is_ci:
-        for model_name, expected_result in CANONICAL_IMAGE_VALUES.items():
-            print("evaluating", model_name)
-            model = LateInteractionMultimodalEmbedding(model_name=model_name)
-            result = next(iter(model.embed_image(images, batch_size=6)))
-            token_num, abridged_dim = expected_result.shape
-            assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)
+    if os.getenv("CI"):
+        pytest.skip("Colpali is too large to test in CI")
+
+    for model_name, expected_result in CANONICAL_IMAGE_VALUES.items():
+        print("evaluating", model_name)
+        model = LateInteractionMultimodalEmbedding(model_name=model_name)
+        result = next(iter(model.embed_image(images, batch_size=6)))
+        token_num, abridged_dim = expected_result.shape
+        assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)
 
 
 def test_single_embedding_query():
-    is_ci = os.getenv("CI")
-    if not is_ci:
-        queries_to_embed = queries
+    if os.getenv("CI"):
+        pytest.skip("Colpali is too large to test in CI")
 
-        for model_name, expected_result in CANONICAL_QUERY_VALUES.items():
-            print("evaluating", model_name)
-            model = LateInteractionMultimodalEmbedding(model_name=model_name)
-            result = next(iter(model.embed_text(queries_to_embed)))
-            token_num, abridged_dim = expected_result.shape
-            assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)
+    for model_name, expected_result in CANONICAL_QUERY_VALUES.items():
+        print("evaluating", model_name)
+        model = LateInteractionMultimodalEmbedding(model_name=model_name)
+        result = next(iter(model.embed_text(queries)))
+        token_num, abridged_dim = expected_result.shape
+        assert np.allclose(result[:token_num, :abridged_dim], expected_result, atol=2e-3)

tests/test_sparse_embeddings.py

@@ -5,10 +5,10 @@ import numpy as np
 
 from fastembed.sparse.bm25 import Bm25
 from fastembed.sparse.sparse_text_embedding import SparseTextEmbedding
-from tests.utils import delete_model_cache
+from tests.utils import delete_model_cache, should_test_model
 
 CANONICAL_COLUMN_VALUES = {
-    "prithvida/Splade_PP_en_v1": {
+    "prithivida/Splade_PP_en_v1": {
         "indices": [
             2040,
             2047,
@@ -49,28 +49,41 @@ CANONICAL_COLUMN_VALUES = {
 docs = ["Hello World"]
 
 
-def test_batch_embedding() -> None:
+@pytest.mark.parametrize("model_name", ["prithivida/Splade_PP_en_v1"])
+def test_batch_embedding(model_name: str) -> None:
     is_ci = os.getenv("CI")
     docs_to_embed = docs * 10
 
-    for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
-        model = SparseTextEmbedding(model_name=model_name)
-        result = next(iter(model.embed(docs_to_embed, batch_size=6)))
-        assert result.indices.tolist() == expected_result["indices"]
+    model = SparseTextEmbedding(model_name=model_name)
+    result = next(iter(model.embed(docs_to_embed, batch_size=6)))
+    expected_result = CANONICAL_COLUMN_VALUES[model_name]
+    assert result.indices.tolist() == expected_result["indices"]
 
-        for i, value in enumerate(result.values):
-            assert pytest.approx(value, abs=0.001) == expected_result["values"][i]
-        if is_ci:
-            delete_model_cache(model.model._model_dir)
+    for i, value in enumerate(result.values):
+        assert pytest.approx(value, abs=0.001) == expected_result["values"][i]
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
 
 
-def test_single_embedding() -> None:
+@pytest.mark.parametrize("model_name", ["prithivida/Splade_PP_en_v1"])
+def test_single_embedding(model_name: str) -> None:
     is_ci = os.getenv("CI")
-    for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+
+    for model_desc in SparseTextEmbedding._list_supported_models():
+        if (
+            model_desc.model not in CANONICAL_COLUMN_VALUES
+        ):  # attention models and bm25 are also parts of
+            # SparseTextEmbedding, however, they have their own tests
+            continue
+        if not should_test_model(model_desc, model_name, is_ci, is_manual):
+            continue
+
         model = SparseTextEmbedding(model_name=model_name)
 
         passage_result = next(iter(model.embed(docs, batch_size=6)))
         query_result = next(iter(model.query_embed(docs)))
+        expected_result = CANONICAL_COLUMN_VALUES[model_name]
         for result in [passage_result, query_result]:
             assert result.indices.tolist() == expected_result["indices"]
 
@@ -80,9 +93,10 @@ def test_single_embedding() -> None:
             delete_model_cache(model.model._model_dir)
 
 
-def test_parallel_processing() -> None:
+@pytest.mark.parametrize("model_name", ["prithivida/Splade_PP_en_v1"])
+def test_parallel_processing(model_name: str) -> None:
     is_ci = os.getenv("CI")
-    model = SparseTextEmbedding(model_name="prithivida/Splade_PP_en_v1")
+    model = SparseTextEmbedding(model_name=model_name)
     docs = ["hello world", "flag embedding"] * 30
     sparse_embeddings_duo = list(model.embed(docs, batch_size=10, parallel=2))
     sparse_embeddings_all = list(model.embed(docs, batch_size=10, parallel=0))
@@ -172,10 +186,7 @@ def test_disable_stemmer_behavior(disable_stemmer: bool) -> None:
     assert result == expected, f"Expected {expected}, but got {result}"
 
 
-@pytest.mark.parametrize(
-    "model_name",
-    ["prithivida/Splade_PP_en_v1"],
-)
+@pytest.mark.parametrize("model_name", ["prithivida/Splade_PP_en_v1"])
 def test_lazy_load(model_name: str) -> None:
     is_ci = os.getenv("CI")
     model = SparseTextEmbedding(model_name=model_name, lazy_load=True)

tests/test_text_cross_encoder.py

@@ -4,7 +4,7 @@ import numpy as np
 import pytest
 
 from fastembed.rerank.cross_encoder import TextCrossEncoder
-from tests.utils import delete_model_cache
+from tests.utils import delete_model_cache, should_test_model
 
 CANONICAL_SCORE_VALUES = {
     "Xenova/ms-marco-MiniLM-L-6-v2": np.array([8.500708, -2.541011]),
@@ -15,44 +15,37 @@ CANONICAL_SCORE_VALUES = {
     "jinaai/jina-reranker-v2-base-multilingual": np.array([1.6533, -1.6455]),
 }
 
-SELECTED_MODELS = {
-    "Xenova": "Xenova/ms-marco-MiniLM-L-6-v2",
-    "BAAI": "BAAI/bge-reranker-base",
-    "jinaai": "jinaai/jina-reranker-v1-tiny-en",
-}
 
-
-@pytest.mark.parametrize(
-    "model_name",
-    [model_name for model_name in CANONICAL_SCORE_VALUES],
-)
+@pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
 def test_rerank(model_name: str) -> None:
     is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
 
-    model = TextCrossEncoder(model_name=model_name)
+    for model_desc in TextCrossEncoder._list_supported_models():
+        if not should_test_model(model_desc, model_name, is_ci, is_manual):
+            continue
 
-    query = "What is the capital of France?"
-    documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."]
-    scores = np.array(list(model.rerank(query, documents)))
+        model = TextCrossEncoder(model_name=model_name)
 
-    pairs = [(query, doc) for doc in documents]
-    scores2 = np.array(list(model.rerank_pairs(pairs)))
-    assert np.allclose(
-        scores, scores2, atol=1e-5
-    ), f"Model: {model_name}, Scores: {scores}, Scores2: {scores2}"
+        query = "What is the capital of France?"
+        documents = ["Paris is the capital of France.", "Berlin is the capital of Germany."]
+        scores = np.array(list(model.rerank(query, documents)))
 
-    canonical_scores = CANONICAL_SCORE_VALUES[model_name]
-    assert np.allclose(
-        scores, canonical_scores, atol=1e-3
-    ), f"Model: {model_name}, Scores: {scores}, Expected: {canonical_scores}"
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
+        pairs = [(query, doc) for doc in documents]
+        scores2 = np.array(list(model.rerank_pairs(pairs)))
+        assert np.allclose(
+            scores, scores2, atol=1e-5
+        ), f"Model: {model_name}, Scores: {scores}, Scores2: {scores2}"
+
+        canonical_scores = CANONICAL_SCORE_VALUES[model_name]
+        assert np.allclose(
+            scores, canonical_scores, atol=1e-3
+        ), f"Model: {model_name}, Scores: {scores}, Expected: {canonical_scores}"
+        if is_ci:
+            delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize(
-    "model_name",
-    [model_name for model_name in SELECTED_MODELS.values()],
-)
+@pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
 def test_batch_rerank(model_name: str) -> None:
     is_ci = os.getenv("CI")
 
@@ -78,10 +71,7 @@ def test_batch_rerank(model_name: str) -> None:
         delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize(
-    "model_name",
-    ["Xenova/ms-marco-MiniLM-L-6-v2"],
-)
+@pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
 def test_lazy_load(model_name: str) -> None:
     is_ci = os.getenv("CI")
     model = TextCrossEncoder(model_name=model_name, lazy_load=True)
@@ -95,10 +85,7 @@ def test_lazy_load(model_name: str) -> None:
         delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize(
-    "model_name",
-    [model_name for model_name in SELECTED_MODELS.values()],
-)
+@pytest.mark.parametrize("model_name", ["Xenova/ms-marco-MiniLM-L-6-v2"])
 def test_rerank_pairs_parallel(model_name: str) -> None:
     is_ci = os.getenv("CI")
 

tests/test_text_multitask_embeddings.py

@@ -4,7 +4,7 @@ import numpy as np
 import pytest
 
 from fastembed import TextEmbedding
-from fastembed.text.multitask_embedding import Task
+from fastembed.text.multitask_embedding import JinaEmbeddingV3, Task
 from tests.utils import delete_model_cache
 
 
@@ -60,52 +60,43 @@ CANONICAL_VECTOR_VALUES = {
 docs = ["Hello World", "Follow the white rabbit."]
 
 
-def test_batch_embedding():
+@pytest.mark.parametrize("dim,model_name", [(1024, "jinaai/jina-embeddings-v3")])
+def test_batch_embedding(dim: int, model_name: str):
     is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+    if is_ci and not is_manual:
+        pytest.skip("Skipping multitask models in CI non-manual mode")
+
     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 = TextEmbedding(model_name=model_name)
 
-        model_name = model_desc.model
-        dim = model_desc.dim
+    embeddings = list(model.embed(documents=docs_to_embed, batch_size=6))
+    embeddings = np.stack(embeddings, axis=0)
 
-        if model_name not in CANONICAL_VECTOR_VALUES.keys():
-            continue
+    assert embeddings.shape == (len(docs_to_embed), dim)
 
-        model = TextEmbedding(model_name=model_name)
+    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_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)
+    if is_ci:
+        delete_model_cache(model.model._model_dir)
 
 
 def test_single_embedding():
     is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+    if is_ci and not is_manual:
+        pytest.skip("Skipping multitask models in CI non-manual mode")
 
-    for model_desc in TextEmbedding._list_supported_models():
-        if not is_ci and model_desc.size_in_GB > 1:
-            continue
-
+    for model_desc in JinaEmbeddingV3._list_supported_models():
+        # todo: once we add more models, we should not test models >1GB size locally
         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]:
@@ -118,27 +109,42 @@ def test_single_embedding():
 
             canonical_vector = task["vectors"]
             assert np.allclose(
-                embeddings[: len(docs), : canonical_vector.shape[1]], canonical_vector, atol=1e-4
+                embeddings[:, : canonical_vector.shape[1]], canonical_vector, atol=1e-4
             ), model_desc.model
 
+        classification_embeddings = list(model.embed(documents=docs, task_id=Task.CLASSIFICATION))
+        classification_embeddings = np.stack(classification_embeddings, axis=0)
+
+        assert classification_embeddings.shape == (len(docs), dim)
+
+        model = TextEmbedding(model_name=model_name, task_id=Task.CLASSIFICATION)
+        default_embeddings = list(model.embed(documents=docs))
+        default_embeddings = np.stack(default_embeddings, axis=0)
+
+        assert default_embeddings.shape == (len(docs), dim)
+
+        assert np.allclose(
+            classification_embeddings,
+            default_embeddings,
+            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")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+    if is_ci and not is_manual:
+        pytest.skip("Skipping multitask models in CI non-manual mode")
+
     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
-
+    for model_desc in JinaEmbeddingV3._list_supported_models():
+        # todo: once we add more models, we should not test models >1GB size locally
         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}")
@@ -150,7 +156,7 @@ def test_single_embedding_query():
 
         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
+            embeddings[:, : canonical_vector.shape[1]], canonical_vector, atol=1e-4
         ), model_desc.model
 
         if is_ci:
@@ -159,18 +165,18 @@ def test_single_embedding_query():
 
 def test_single_embedding_passage():
     is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+    if is_ci and not is_manual:
+        pytest.skip("Skipping multitask models in CI non-manual mode")
+
     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
+    for model_desc in JinaEmbeddingV3._list_supported_models():
+        # todo: once we add more models, we should not test models >1GB size locally
 
         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}")
@@ -182,21 +188,22 @@ def test_single_embedding_passage():
 
         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
+            embeddings[:, : 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():
+@pytest.mark.parametrize("dim,model_name", [(1024, "jinaai/jina-embeddings-v3")])
+def test_parallel_processing(dim: int, model_name: str):
     is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+    if is_ci and not is_manual:
+        pytest.skip("Skipping in CI non-manual mode")
 
     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
@@ -216,33 +223,14 @@ def test_parallel_processing():
         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"],
-)
+@pytest.mark.parametrize("model_name", ["jinaai/jina-embeddings-v3"])
 def test_lazy_load(model_name: str):
     is_ci = os.getenv("CI")
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
+
+    if is_ci and not is_manual:
+        pytest.skip("Skipping in CI non-manual mode")
+
     model = TextEmbedding(model_name=model_name, lazy_load=True)
     assert not hasattr(model.model, "model")
 

tests/test_text_onnx_embeddings.py

@@ -5,7 +5,7 @@ import numpy as np
 import pytest
 
 from fastembed.text.text_embedding import TextEmbedding
-from tests.utils import delete_model_cache
+from tests.utils import delete_model_cache, should_test_model
 
 CANONICAL_VECTOR_VALUES = {
     "BAAI/bge-small-en": np.array([-0.0232, -0.0255, 0.0174, -0.0639, -0.0006]),
@@ -72,17 +72,19 @@ CANONICAL_VECTOR_VALUES = {
 MULTI_TASK_MODELS = ["jinaai/jina-embeddings-v3"]
 
 
-def test_embedding() -> None:
+@pytest.mark.parametrize("model_name", ["BAAI/bge-small-en-v1.5"])
+def test_embedding(model_name: str) -> None:
     is_ci = os.getenv("CI")
     is_mac = platform.system() == "Darwin"
+    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
 
     for model_desc in TextEmbedding._list_supported_models():
-        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")
+        if model_desc.model in MULTI_TASK_MODELS or (
+            is_mac and model_desc.model == "nomic-ai/nomic-embed-text-v1.5-Q"
         ):
             continue
+        if not should_test_model(model_desc, model_name, is_ci, is_manual):
+            continue
 
         dim = model_desc.dim
 
@@ -95,15 +97,12 @@ def test_embedding() -> None:
         canonical_vector = CANONICAL_VECTOR_VALUES[model_desc.model]
         assert np.allclose(
             embeddings[0, : canonical_vector.shape[0]], canonical_vector, atol=1e-3
-        ), model_desc["model"]
+        ), model_desc.model
         if is_ci:
             delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize(
-    "n_dims,model_name",
-    [(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")],
-)
+@pytest.mark.parametrize("n_dims,model_name", [(384, "BAAI/bge-small-en-v1.5")])
 def test_batch_embedding(n_dims: int, model_name: str) -> None:
     is_ci = os.getenv("CI")
     model = TextEmbedding(model_name=model_name)
@@ -112,15 +111,12 @@ def test_batch_embedding(n_dims: int, model_name: str) -> None:
     embeddings = list(model.embed(docs, batch_size=10))
     embeddings = np.stack(embeddings, axis=0)
 
-    assert embeddings.shape == (200, n_dims)
+    assert embeddings.shape == (len(docs), n_dims)
     if is_ci:
         delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize(
-    "n_dims,model_name",
-    [(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")],
-)
+@pytest.mark.parametrize("n_dims,model_name", [(384, "BAAI/bge-small-en-v1.5")])
 def test_parallel_processing(n_dims: int, model_name: str) -> None:
     is_ci = os.getenv("CI")
     model = TextEmbedding(model_name=model_name)
@@ -135,7 +131,7 @@ def test_parallel_processing(n_dims: int, model_name: str) -> None:
     embeddings_3 = list(model.embed(docs, batch_size=10, parallel=0))
     embeddings_3 = np.stack(embeddings_3, axis=0)
 
-    assert embeddings.shape == (200, n_dims)
+    assert embeddings.shape == (len(docs), n_dims)
     assert np.allclose(embeddings, embeddings_2, atol=1e-3)
     assert np.allclose(embeddings, embeddings_3, atol=1e-3)
 
@@ -143,10 +139,7 @@ def test_parallel_processing(n_dims: int, model_name: str) -> None:
         delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize(
-    "model_name",
-    ["BAAI/bge-small-en-v1.5"],
-)
+@pytest.mark.parametrize("model_name", ["BAAI/bge-small-en-v1.5"])
 def test_lazy_load(model_name: str) -> None:
     is_ci = os.getenv("CI")
     model = TextEmbedding(model_name=model_name, lazy_load=True)

tests/utils.py

@@ -3,7 +3,9 @@ import traceback
 
 from pathlib import Path
 from types import TracebackType
-from typing import Union, Callable, Any, Type
+from typing import Union, Callable, Any, Type, Optional
+
+from fastembed.common.model_description import BaseModelDescription
 
 
 def delete_model_cache(model_dir: Union[str, Path]) -> None:
@@ -35,3 +37,31 @@ def delete_model_cache(model_dir: Union[str, Path]) -> None:
     if model_dir.exists():
         # todo: PermissionDenied is raised on blobs removal in Windows, with blobs > 2GB
         shutil.rmtree(model_dir, onerror=on_error)
+
+
+def should_test_model(
+    model_desc: BaseModelDescription,
+    autotest_model_name: str,
+    is_ci: Optional[str],
+    is_manual: bool,
+):
+    """Determine if a model should be tested based on environment
+
+    Tests can be run either in ci or locally.
+    Testing all models each time in ci is too long.
+    The testing scheme in ci and on a local machine are different, therefore, there are 3 possible scenarious.
+    1) Run lightweight tests in ci:
+        - test only one model that has been manually chosen as a representative for a certain class family
+    2) Run heavyweight (manual) tests in ci:
+        - test all models
+        Running tests in ci each time is too expensive, however, it's fine to run it one time with a manual dispatch
+    3) Run tests locally:
+        - test all models, which are not too heavy, since network speed might be a bottleneck
+
+    """
+    if not is_ci:
+        if model_desc.size_in_GB > 1:
+            return False
+    elif not is_manual and model_desc.model != autotest_model_name:
+        return False
+    return True