draft

.github/workflows/python-tests.yml

@@ -1,11 +1,9 @@
 name: Tests
-run-name: Tests (gpu)
 
 on:
-  pull_request:
+  push:
     branches: [ master, main, gpu ]
-  workflow_dispatch:
-  
+  pull_request:
 
 env:
   CARGO_TERM_COLOR: always
@@ -23,6 +21,8 @@ jobs:
           - '3.13.x'
         os:
           - ubuntu-latest
+          - macos-latest
+          - windows-latest
 
     runs-on: ${{ matrix.os }}
 
@@ -42,4 +42,4 @@ jobs:
 
       - name: Run pytest
         run: |
-          poetry run pytest 
+          poetry run pytest 

README.md

@@ -63,23 +63,6 @@ embeddings = list(model.embed(documents))
 
 ```
 
-Dense text embedding can also be extended with models which are not in the list of supported models.
-
-```python
-from fastembed import TextEmbedding
-from fastembed.common.model_description import PoolingType, ModelSource
-
-TextEmbedding.add_custom_model(
-    model="intfloat/multilingual-e5-small",
-    pooling=PoolingType.MEAN,
-    normalization=True,
-    sources=ModelSource(hf="intfloat/multilingual-e5-small"),  # can be used with an `url` to load files from a private storage
-    dim=384,
-    model_file="onnx/model.onnx",  # can be used to load an already supported model with another optimization or quantization, e.g. onnx/model_O4.onnx
-)
-model = TextEmbedding(model_name="intfloat/multilingual-e5-small")
-embeddings = list(model.embed(documents))
-```
 
 
 ### 🔱 Sparse text embeddings
@@ -154,27 +137,6 @@ embeddings = list(model.embed(images))
 # ]
 ```
 
-### Late interaction multimodal models (ColPali)
-
-```python
-from fastembed import LateInteractionMultimodalEmbedding
-
-doc_images = [
-    "./path/to/qdrant_pdf_doc_1_screenshot.jpg",
-    "./path/to/colpali_pdf_doc_2_screenshot.jpg",
-]
-
-query = "What is Qdrant?"
-
-model = LateInteractionMultimodalEmbedding(model_name="Qdrant/colpali-v1.3-fp16")
-doc_images_embeddings = list(model.embed_image(doc_images))
-# shape (2, 1030, 128)
-# [array([[-0.03353882, -0.02090454, ..., -0.15576172, -0.07678223]], dtype=float32)]
-query_embedding = model.embed_text(query)
-# shape (1, 20, 128)
-# [array([[-0.00218201,  0.14758301, ...,  -0.02207947,  0.16833496]], dtype=float32)]
-```
-
 ### 🔄 Rerankers
 ```python
 from fastembed.rerank.cross_encoder import TextCrossEncoder

fastembed/common/model_description.py

@@ -1,5 +1,4 @@
 from dataclasses import dataclass, field
-from enum import Enum
 from typing import Optional, Any
 
 
@@ -7,11 +6,6 @@ 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:
@@ -34,7 +28,7 @@ class BaseModelDescription:
 @dataclass(frozen=True)
 class DenseModelDescription(BaseModelDescription):
     dim: Optional[int] = None
-    tasks: Optional[dict[str, Any]] = field(default_factory=dict)
+    tasks: Optional[dict[str, Any]] = None
 
     def __post_init__(self) -> None:
         assert self.dim is not None, "dim is required for dense model description"
@@ -44,9 +38,3 @@ class DenseModelDescription(BaseModelDescription):
 class SparseModelDescription(BaseModelDescription):
     requires_idf: Optional[bool] = None
     vocab_size: Optional[int] = None
-
-
-class PoolingType(str, Enum):
-    CLS = "CLS"
-    MEAN = "MEAN"
-    DISABLED = "DISABLED"

fastembed/common/model_management.py

@@ -33,31 +33,6 @@ class ModelManagement(Generic[T]):
         """
         raise NotImplementedError()
 
-    @classmethod
-    def add_custom_model(
-        cls,
-        *args: Any,
-        **kwargs: Any,
-    ) -> None:
-        """Add a custom model to the existing embedding classes based on the passed model descriptions
-
-        Model description dict should contain the fields same as in one of the model descriptions presented
-         in fastembed.common.model_description
-
-         E.g. for BaseModelDescription:
-              model: str
-              sources: ModelSource
-              model_file: str
-              description: str
-              license: str
-              size_in_GB: float
-              additional_files: list[str]
-
-        Returns:
-            None
-        """
-        raise NotImplementedError()
-
     @classmethod
     def _list_supported_models(cls) -> list[T]:
         raise NotImplementedError()
@@ -330,10 +305,9 @@ 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-' if deprecated_tar_struct else ''}{model_name.split('/')[-1]}"
+        fast_model_name = f"fast-{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
@@ -439,7 +413,6 @@ 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/utils.py

@@ -5,12 +5,11 @@ 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")
 
@@ -23,15 +22,6 @@ def normalize(input_array: NumpyArray, p: int = 2, dim: int = 1, eps: float = 1e
     return normalized_array
 
 
-def mean_pooling(input_array: NumpyArray, attention_mask: NDArray[np.int64]) -> NumpyArray:
-    input_mask_expanded = np.expand_dims(attention_mask, axis=-1).astype(np.int64)
-    input_mask_expanded = np.tile(input_mask_expanded, (1, 1, input_array.shape[-1]))
-    sum_embeddings = np.sum(input_array * input_mask_expanded, axis=1)
-    sum_mask = np.sum(input_mask_expanded, axis=1)
-    pooled_embeddings = sum_embeddings / np.maximum(sum_mask, 1e-9)
-    return pooled_embeddings
-
-
 def iter_batch(iterable: Iterable[T], size: int) -> Iterable[list[T]]:
     """
     >>> list(iter_batch([1,2,3,4,5], 3))
@@ -67,18 +57,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/config/dense_models.json

@@ -0,0 +1,16 @@
+{
+  "models": [
+    {
+      "model": "BAAI/bge-base-en",
+      "dim": 768,
+      "description": "Text embeddings, Unimodal (text), English...",
+      "license": "mit",
+      "size_in_GB": 0.42,
+      "sources": {
+        "hf": "Qdrant/fast-bge-base-en",
+        "url": "https://storage.googleapis.com/qdrant-fastembed/fast-bge-base-en.tar.gz"
+      },
+      "model_file": "model_optimized.onnx"
+    }
+  ]
+}

fastembed/image/onnx_embedding.py

@@ -1,5 +1,6 @@
 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
@@ -194,7 +195,7 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
         return onnx_input
 
     def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
-        return normalize(output.model_output)
+        return normalize(output.model_output).astype(np.float32)
 
 
 class OnnxImageEmbeddingWorker(ImageEmbeddingWorker[NumpyArray]):

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
+            return output.model_output.astype(np.float32)
 
         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)
+        output.model_output *= np.expand_dims(output.attention_mask, 2).astype(np.float32)
         norm = np.linalg.norm(output.model_output, ord=2, axis=2, keepdims=True)
         norm_clamped = np.maximum(norm, 1e-12)
         output.model_output /= norm_clamped
-        return output.model_output
+        return output.model_output.astype(np.float32)
 
     def _preprocess_onnx_input(
         self, onnx_input: dict[str, NumpyArray], is_doc: bool = True, **kwargs: Any

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
+        return output.model_output.astype(np.float32)
 
     def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
         texts_query: list[str] = []
@@ -197,11 +197,12 @@ class ColPali(LateInteractionMultimodalEmbeddingBase, OnnxMultimodalModel[NumpyA
         Returns:
             Dict[str, NumpyArray]: ONNX input with text placeholders.
         """
+
         onnx_input["input_ids"] = np.array(
-            [self.EMPTY_TEXT_PLACEHOLDER for _ in onnx_input["pixel_values"]]
+            [self.EMPTY_TEXT_PLACEHOLDER for _ in onnx_input["input_ids"]]
         )
         onnx_input["attention_mask"] = np.array(
-            [self.EVEN_ATTENTION_MASK for _ in onnx_input["pixel_values"]]
+            [self.EVEN_ATTENTION_MASK for _ in onnx_input["input_ids"]]
         )
         return onnx_input
 

fastembed/late_interaction_multimodal/onnx_multimodal_model.py

@@ -73,8 +73,8 @@ class OnnxMultimodalModel(OnnxModel[T]):
             cuda=cuda,
             device_id=device_id,
         )
-        self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
         assert self.tokenizer is not None
+        self.tokenizer, self.special_token_to_id = load_tokenizer(model_dir=model_dir)
         self.processor = load_preprocessor(model_dir=model_dir)
 
     def load_onnx_model(self) -> None:
@@ -159,6 +159,10 @@ class OnnxMultimodalModel(OnnxModel[T]):
             for batch in pool.ordered_map(iter_batch(documents, batch_size), **params):
                 yield from self._post_process_onnx_text_output(batch)  # type: ignore
 
+    def _build_onnx_image_input(self, encoded: NumpyArray) -> dict[str, NumpyArray]:
+        input_name = self.model.get_inputs()[0].name  # type: ignore[union-attr]
+        return {input_name: encoded}
+
     def onnx_embed_image(self, images: list[ImageInput], **kwargs: Any) -> OnnxOutputContext:
         with contextlib.ExitStack():
             image_files = [
@@ -167,7 +171,7 @@ class OnnxMultimodalModel(OnnxModel[T]):
             ]
             assert self.processor is not None, "Processor is not initialized"
             encoded = np.array(self.processor(image_files))
-        onnx_input = {"pixel_values": encoded}
+        onnx_input = self._build_onnx_image_input(encoded)
         onnx_input = self._preprocess_onnx_image_input(onnx_input, **kwargs)
         model_output = self.model.run(None, onnx_input)  # type: ignore[union-attr]
         embeddings = model_output[0].reshape(len(images), -1)

fastembed/model_loader.py

@@ -0,0 +1,16 @@
+from pathlib import Path
+import json
+from typing import Dict, List
+
+
+class ModelLoader:
+    def __init__(self):
+        self.config_dir = Path(__file__).parent / "configs"
+        self._models: Dict[str, List[Dict]] = {}
+
+    def load_models(self, model_type: str) -> List[Dict]:
+        if model_type not in self._models:
+            config_path = self.config_dir / f"{model_type}_models.json"
+            with open(config_path) as f:
+                self._models[model_type] = json.load(f)["models"]
+        return self._models[model_type]

fastembed/text/custom_text_embedding.py

@@ -1,91 +0,0 @@
-from typing import Optional, Sequence, Any, Iterable
-
-from dataclasses import dataclass
-
-import numpy as np
-from numpy.typing import NDArray
-
-from fastembed.common import OnnxProvider
-from fastembed.common.model_description import (
-    PoolingType,
-    DenseModelDescription,
-)
-from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.types import NumpyArray
-from fastembed.common.utils import normalize, mean_pooling
-from fastembed.text.onnx_embedding import OnnxTextEmbedding
-
-
-@dataclass(frozen=True)
-class PostprocessingConfig:
-    pooling: PoolingType
-    normalization: bool
-
-
-class CustomTextEmbedding(OnnxTextEmbedding):
-    SUPPORTED_MODELS: list[DenseModelDescription] = []
-    POSTPROCESSING_MAPPING: dict[str, PostprocessingConfig] = {}
-
-    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,
-        )
-        self._pooling = self.POSTPROCESSING_MAPPING[model_name].pooling
-        self._normalization = self.POSTPROCESSING_MAPPING[model_name].normalization
-
-    @classmethod
-    def _list_supported_models(cls) -> list[DenseModelDescription]:
-        return cls.SUPPORTED_MODELS
-
-    def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
-        return self._normalize(self._pool(output.model_output, output.attention_mask))
-
-    def _pool(
-        self, embeddings: NumpyArray, attention_mask: Optional[NDArray[np.int64]] = None
-    ) -> NumpyArray:
-        if self._pooling == PoolingType.CLS:
-            return embeddings[:, 0]
-
-        if self._pooling == PoolingType.MEAN:
-            if attention_mask is None:
-                raise ValueError("attention_mask must be provided for mean pooling")
-            return mean_pooling(embeddings, attention_mask)
-
-        if self._pooling == PoolingType.DISABLED:
-            return embeddings
-
-    def _normalize(self, embeddings: NumpyArray) -> NumpyArray:
-        return normalize(embeddings) if self._normalization else embeddings
-
-    @classmethod
-    def add_model(
-        cls,
-        model_description: DenseModelDescription,
-        pooling: PoolingType,
-        normalization: bool,
-    ) -> None:
-        cls.SUPPORTED_MODELS.append(model_description)
-        cls.POSTPROCESSING_MAPPING[model_description.model] = PostprocessingConfig(
-            pooling=pooling, normalization=normalization
-        )

fastembed/text/onnx_embedding.py

@@ -1,5 +1,6 @@
 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
@@ -20,7 +21,6 @@ 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,7 +36,6 @@ 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",
     ),
@@ -64,7 +63,6 @@ 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",
     ),
@@ -92,10 +90,21 @@ 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",
     ),
+    DenseModelDescription(
+        model="thenlper/gte-large",
+        dim=1024,
+        description=(
+            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
+            "Prefixes for queries/documents: not necessary, 2023 year."
+        ),
+        license="mit",
+        size_in_GB=1.20,
+        sources=ModelSource(hf="qdrant/gte-large-onnx"),
+        model_file="model.onnx",
+    ),
     DenseModelDescription(
         model="mixedbread-ai/mxbai-embed-large-v1",
         dim=1024,
@@ -305,14 +314,13 @@ class OnnxTextEmbedding(TextEmbeddingBase, OnnxTextModel[NumpyArray]):
 
     def _post_process_onnx_output(self, output: OnnxOutputContext) -> Iterable[NumpyArray]:
         embeddings = output.model_output
-
         if embeddings.ndim == 3:  # (batch_size, seq_len, embedding_dim)
             processed_embeddings = embeddings[:, 0]
         elif embeddings.ndim == 2:  # (batch_size, embedding_dim)
             processed_embeddings = embeddings
         else:
             raise ValueError(f"Unsupported embedding shape: {embeddings.shape}")
-        return normalize(processed_embeddings)
+        return normalize(processed_embeddings).astype(np.float32)
 
     def load_onnx_model(self) -> None:
         self._load_onnx_model(

fastembed/text/onnx_text_model.py

@@ -6,12 +6,11 @@ from typing import Any, Iterable, Optional, Sequence, Type, Union
 import numpy as np
 from numpy.typing import NDArray
 from tokenizers import Encoding, Tokenizer
-import onnxruntime as ort
 
 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
 
 
@@ -63,30 +62,6 @@ class OnnxTextModel(OnnxModel[T]):
     def tokenize(self, documents: list[str], **kwargs: Any) -> list[Encoding]:
         return self.tokenizer.encode_batch(documents)  # type: ignore[union-attr]
 
-    def _run_with_io_binding(
-        self, onnx_input: dict[str, NumpyArray], device_id: int = 0
-    ) -> NumpyArray:
-        """Run inference with IO Binding for optimized memory transfer for non-CPU execution providers."""
-        io_binding = self.model.io_binding()
-
-        for name, value in onnx_input.items():
-            ort_value = ort.OrtValue.ortvalue_from_numpy(
-                numpy_obj=value, device_type="cuda", device_id=device_id
-            )
-            io_binding.bind_ortvalue_input(name=name, ortvalue=ort_value)
-
-        output_names = [output.name for output in self.model.get_outputs()]
-        for output_name in output_names:
-            io_binding.bind_output(
-                name=output_name,
-                device_type="cuda",
-                device_id=device_id,
-            )
-
-        self.model.run_with_iobinding(io_binding)
-
-        return io_binding.copy_outputs_to_cpu()
-
     def onnx_embed(
         self,
         documents: list[str],
@@ -107,13 +82,7 @@ class OnnxTextModel(OnnxModel[T]):
             )
         onnx_input = self._preprocess_onnx_input(onnx_input, **kwargs)
 
-        providers = kwargs.get("providers", None)
-        cuda = kwargs.get("cuda", False)
-
-        if is_cuda_enabled(cuda, providers):
-            model_output = self._run_with_io_binding(onnx_input)
-        else:
-            model_output = self.model.run(self.ONNX_OUTPUT_NAMES, onnx_input)  # 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),
@@ -146,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))
         else:
             if parallel == 0:
                 parallel = os.cpu_count()

fastembed/text/pooled_embedding.py

@@ -1,11 +1,9 @@
 from typing import Any, Iterable, Type
 
 import numpy as np
-from numpy.typing import NDArray
 
 from fastembed.common.types import NumpyArray
 from fastembed.common.onnx_model import OnnxOutputContext
-from fastembed.common.utils import mean_pooling
 from fastembed.text.onnx_embedding import OnnxTextEmbedding, OnnxTextEmbeddingWorker
 from fastembed.common.model_description import DenseModelDescription, ModelSource
 
@@ -82,7 +80,6 @@ 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"],
@@ -96,10 +93,16 @@ class PooledEmbedding(OnnxTextEmbedding):
         return PooledEmbeddingWorker
 
     @classmethod
-    def mean_pooling(
-        cls, model_output: NumpyArray, attention_mask: NDArray[np.int64]
-    ) -> NumpyArray:
-        return mean_pooling(model_output, attention_mask)
+    def mean_pooling(cls, model_output: NumpyArray, attention_mask: NumpyArray) -> NumpyArray:
+        token_embeddings = model_output.astype(np.float32)
+        attention_mask = attention_mask.astype(np.float32)
+        input_mask_expanded = np.expand_dims(attention_mask, axis=-1)
+        input_mask_expanded = np.tile(input_mask_expanded, (1, 1, token_embeddings.shape[-1]))
+        input_mask_expanded = input_mask_expanded.astype(np.float32)
+        sum_embeddings = np.sum(token_embeddings * input_mask_expanded, axis=1)
+        sum_mask = np.sum(input_mask_expanded, axis=1)
+        pooled_embeddings = sum_embeddings / np.maximum(sum_mask, 1e-9)
+        return pooled_embeddings
 
     @classmethod
     def _list_supported_models(cls) -> list[DenseModelDescription]:
@@ -116,7 +119,7 @@ class PooledEmbedding(OnnxTextEmbedding):
 
         embeddings = output.model_output
         attn_mask = output.attention_mask
-        return self.mean_pooling(embeddings, attn_mask)
+        return self.mean_pooling(embeddings, attn_mask).astype(np.float32)
 
 
 class PooledEmbeddingWorker(OnnxTextEmbeddingWorker):

fastembed/text/pooled_normalized_embedding.py

@@ -1,5 +1,6 @@
 from typing import Any, Iterable, Type
 
+import numpy as np
 
 from fastembed.common.types import NumpyArray
 from fastembed.common.onnx_model import OnnxOutputContext
@@ -21,7 +22,6 @@ 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",
     ),
@@ -109,18 +109,6 @@ supported_pooled_normalized_models: list[DenseModelDescription] = [
         sources=ModelSource(hf="thenlper/gte-base"),
         model_file="onnx/model.onnx",
     ),
-    DenseModelDescription(
-        model="thenlper/gte-large",
-        dim=1024,
-        description=(
-            "Text embeddings, Unimodal (text), English, 512 input tokens truncation, "
-            "Prefixes for queries/documents: not necessary, 2023 year."
-        ),
-        license="mit",
-        size_in_GB=1.20,
-        sources=ModelSource(hf="qdrant/gte-large-onnx"),
-        model_file="model.onnx",
-    ),
 ]
 
 
@@ -144,7 +132,7 @@ class PooledNormalizedEmbedding(PooledEmbedding):
 
         embeddings = output.model_output
         attn_mask = output.attention_mask
-        return normalize(self.mean_pooling(embeddings, attn_mask))
+        return normalize(self.mean_pooling(embeddings, attn_mask)).astype(np.float32)
 
 
 class PooledNormalizedEmbeddingWorker(OnnxTextEmbeddingWorker):

fastembed/text/text_embedding.py

@@ -4,13 +4,12 @@ from dataclasses import asdict
 
 from fastembed.common.types import NumpyArray, OnnxProvider
 from fastembed.text.clip_embedding import CLIPOnnxEmbedding
-from fastembed.text.custom_text_embedding import CustomTextEmbedding
 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
-from fastembed.common.model_description import DenseModelDescription, ModelSource, PoolingType
+from fastembed.common.model_description import DenseModelDescription
 
 
 class TextEmbedding(TextEmbeddingBase):
@@ -20,7 +19,6 @@ class TextEmbedding(TextEmbeddingBase):
         PooledNormalizedEmbedding,
         PooledEmbedding,
         JinaEmbeddingV3,
-        CustomTextEmbedding,
     ]
 
     @classmethod
@@ -39,43 +37,6 @@ class TextEmbedding(TextEmbeddingBase):
             result.extend(embedding._list_supported_models())
         return result
 
-    @classmethod
-    def add_custom_model(
-        cls,
-        model: str,
-        pooling: PoolingType,
-        normalization: bool,
-        sources: ModelSource,
-        dim: int,
-        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 TextEmbedding, if you still want to add this model, "
-                    f"please use another model name"
-                )
-
-        CustomTextEmbedding.add_model(
-            DenseModelDescription(
-                model=model,
-                sources=sources,
-                dim=dim,
-                model_file=model_file,
-                description=description,
-                license=license,
-                size_in_GB=size_in_gb,
-                additional_files=additional_files or [],
-            ),
-            pooling=pooling,
-            normalization=normalization,
-        )
-
     def __init__(
         self,
         model_name: str = "BAAI/bge-small-en-v1.5",
@@ -90,24 +51,29 @@ class TextEmbedding(TextEmbeddingBase):
         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 on HF and release notes to ensure compatibility with your workflow. ",
+                "The model 'nomic-ai/nomic-embed-text-v1.5-Q' has been updated on HuggingFace. "
+                "Please review the latest documentation and release notes to ensure compatibility with your workflow. ",
+                UserWarning,
+                stacklevel=2,
+            )
+        if model_name == "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2":
+            warnings.warn(
+                "The model 'sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2' has been updated to "
+                "include a mean pooling layer. Please ensure your usage aligns with the new functionality. "
+                "Support for the previous version without mean pooling will be removed as of version 0.5.2.",
                 UserWarning,
                 stacklevel=2,
             )
         if model_name in {
-            "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
-            "thenlper/gte-large",
-            "intfloat/multilingual-e5-large",
             "sentence-transformers/paraphrase-multilingual-mpnet-base-v2",
+            "intfloat/multilingual-e5-large",
         }:
             warnings.warn(
-                f"The model {model_name} now uses mean pooling instead of CLS embedding. "
-                f"In order to preserve the previous behaviour, consider either pinning fastembed version to 0.5.1 or "
-                "using `add_custom_model` functionality.",
+                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):

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
-name = "fastembed-gpu"
-version = "0.6.0"
+name = "fastembed"
+version = "0.5.1"
 description = "Fast, light, accurate library built for retrieval embedding generation"
 authors = ["Qdrant Team <info@qdrant.tech>", "NirantK <nirant.bits@gmail.com>"]
 license = "Apache License"
@@ -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" },
@@ -29,7 +29,7 @@ tokenizers = ">=0.15,<1.0"
 huggingface-hub = ">=0.20,<1.0"
 loguru = "^0.7.2"
 pillow = ">=10.3.0,<12.0.0"
-mmh3 = ">=4.1.0,<6.0.0"
+mmh3 = "^4.1.0"
 py-rust-stemmers = "^0.1.0"
 
 [tool.poetry.group.test.dependencies]

tests/test_custom_models.py

@@ -1,158 +0,0 @@
-import itertools
-import os
-import numpy as np
-import pytest
-
-from fastembed.common.model_description import PoolingType, ModelSource, DenseModelDescription
-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.text.text_embedding import TextEmbedding
-from tests.utils import delete_model_cache
-
-
-@pytest.fixture(autouse=True)
-def restore_custom_models_fixture():
-    CustomTextEmbedding.SUPPORTED_MODELS = []
-    yield
-    CustomTextEmbedding.SUPPORTED_MODELS = []
-
-
-def test_text_custom_model():
-    is_ci = os.getenv("CI")
-    custom_model_name = "intfloat/multilingual-e5-small"
-    canonical_vector = np.array(
-        [3.1317e-02, 3.0939e-02, -3.5117e-02, -6.7274e-02, 8.5084e-02], dtype=np.float32
-    )
-    pooling = PoolingType.MEAN
-    normalization = True
-    dim = 384
-    size_in_gb = 0.47
-    source = ModelSource(hf=custom_model_name)
-
-    TextEmbedding.add_custom_model(
-        custom_model_name,
-        pooling=pooling,
-        normalization=normalization,
-        sources=source,
-        dim=dim,
-        size_in_gb=size_in_gb,
-    )
-
-    assert CustomTextEmbedding.SUPPORTED_MODELS[0] == DenseModelDescription(
-        model=custom_model_name,
-        sources=source,
-        model_file="onnx/model.onnx",
-        description="",
-        license="",
-        size_in_GB=size_in_gb,
-        additional_files=[],
-        dim=dim,
-        tasks={},
-    )
-    assert CustomTextEmbedding.POSTPROCESSING_MAPPING[custom_model_name] == PostprocessingConfig(
-        pooling=pooling, normalization=normalization
-    )
-
-    model = TextEmbedding(custom_model_name)
-    docs = ["hello world", "flag embedding"]
-    embeddings = list(model.embed(docs))
-    embeddings = np.stack(embeddings, axis=0)
-    assert embeddings.shape == (2, dim)
-
-    assert np.allclose(embeddings[0, : canonical_vector.shape[0]], 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
-    source = ModelSource(hf="artificial")
-
-    num_tokens = 10
-    dummy_pooled_embedding = np.random.random((1, dim)).astype(np.float32)
-    dummy_token_embedding = np.random.random((1, num_tokens, dim)).astype(np.float32)
-    dummy_attention_mask = np.ones((1, num_tokens)).astype(np.int64)
-
-    dummy_token_output = OnnxOutputContext(
-        model_output=dummy_token_embedding, attention_mask=dummy_attention_mask
-    )
-    dummy_pooled_output = OnnxOutputContext(model_output=dummy_pooled_embedding)
-    input_data = {
-        f"{PoolingType.MEAN.lower()}-normalized": dummy_token_output,
-        f"{PoolingType.MEAN.lower()}": dummy_token_output,
-        f"{PoolingType.CLS.lower()}-normalized": dummy_token_output,
-        f"{PoolingType.CLS.lower()}": dummy_token_output,
-        f"{PoolingType.DISABLED.lower()}-normalized": dummy_pooled_output,
-        f"{PoolingType.DISABLED.lower()}": dummy_pooled_output,
-    }
-
-    expected_output = {
-        f"{PoolingType.MEAN.lower()}-normalized": normalize(
-            mean_pooling(dummy_token_embedding, dummy_attention_mask)
-        ),
-        f"{PoolingType.MEAN.lower()}": mean_pooling(dummy_token_embedding, dummy_attention_mask),
-        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),
-        f"{PoolingType.DISABLED.lower()}": dummy_pooled_embedding,
-    }
-
-    for pooling, normalization in itertools.product(
-        (PoolingType.MEAN, PoolingType.CLS, PoolingType.DISABLED), (True, False)
-    ):
-        model_name = f"{pooling.name.lower()}{'-normalized' if normalization else ''}"
-        TextEmbedding.add_custom_model(
-            model_name,
-            pooling=pooling,
-            normalization=normalization,
-            sources=source,
-            dim=dim,
-            size_in_gb=size_in_gb,
-        )
-
-        custom_text_embedding = CustomTextEmbedding(
-            model_name,
-            lazy_load=True,
-            specific_model_path="./",  # disable model downloading and loading
-        )
-
-        post_processed_output = next(
-            iter(custom_text_embedding._post_process_onnx_output(input_data[model_name]))
-        )
-        assert np.allclose(post_processed_output, expected_output[model_name], atol=1e-3)
-
-
-def test_do_not_add_existing_model():
-    existing_base_model = "sentence-transformers/all-MiniLM-L6-v2"
-    custom_model_name = "intfloat/multilingual-e5-small"
-
-    with pytest.raises(ValueError, match=f"Model {existing_base_model} is already registered"):
-        TextEmbedding.add_custom_model(
-            existing_base_model,
-            pooling=PoolingType.MEAN,
-            normalization=True,
-            sources=ModelSource(hf=existing_base_model),
-            dim=384,
-            size_in_gb=0.47,
-        )
-
-    TextEmbedding.add_custom_model(
-        custom_model_name,
-        pooling=PoolingType.MEAN,
-        normalization=False,
-        sources=ModelSource(hf=existing_base_model),
-        dim=384,
-        size_in_gb=0.47,
-    )
-
-    with pytest.raises(ValueError, match=f"Model {custom_model_name} is already registered"):
-        TextEmbedding.add_custom_model(
-            custom_model_name,
-            pooling=PoolingType.MEAN,
-            normalization=True,
-            sources=ModelSource(hf=custom_model_name),
-            dim=384,
-            size_in_gb=0.47,
-        )

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, should_test_model
+from tests.utils import delete_model_cache
 
 CANONICAL_VECTOR_VALUES = {
     "Qdrant/clip-ViT-B-32-vision": np.array([-0.0098, 0.0128, -0.0274, 0.002, -0.0059]),
@@ -27,13 +27,11 @@ CANONICAL_VECTOR_VALUES = {
 }
 
 
-@pytest.mark.parametrize("model_name", ["Qdrant/clip-ViT-B-32-vision"])
-def test_embedding(model_name: str) -> None:
+def test_embedding() -> 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 should_test_model(model_desc, model_name, is_ci, is_manual):
+        if not is_ci and model_desc.size_in_GB > 1:
             continue
 
         dim = model_desc.dim
@@ -76,12 +74,8 @@ 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, should_test_model
+from tests.utils import delete_model_cache
 
 # vectors are abridged and rounded for brevity
 CANONICAL_COLUMN_VALUES = {
@@ -153,37 +153,31 @@ CANONICAL_QUERY_VALUES = {
 docs = ["Hello World"]
 
 
-@pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
-def test_batch_embedding(model_name: str):
+def test_batch_embedding():
     is_ci = os.getenv("CI")
     docs_to_embed = docs * 10
 
-    model = LateInteractionTextEmbedding(model_name=model_name)
-    result = list(model.embed(docs_to_embed, batch_size=6))
-    expected_result = CANONICAL_COLUMN_VALUES[model_name]
+    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))
 
-    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)
 
 
-@pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
-def test_single_embedding(model_name: str):
+def test_single_embedding():
     is_ci = os.getenv("CI")
-    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
     docs_to_embed = docs
 
-    for model_desc in LateInteractionTextEmbedding._list_supported_models():
-        if not should_test_model(model_desc, model_name, is_ci, is_manual):
-            continue
-
+    for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
         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)
 
@@ -191,20 +185,14 @@ def test_single_embedding(model_name: str):
             delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
-def test_single_embedding_query(model_name: str):
+def test_single_embedding_query():
     is_ci = os.getenv("CI")
-    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
     queries_to_embed = docs
 
-    for model_desc in LateInteractionTextEmbedding._list_supported_models():
-        if not should_test_model(model_desc, model_name, is_ci, is_manual):
-            continue
-
+    for model_name, expected_result in CANONICAL_QUERY_VALUES.items():
         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)
 
@@ -212,11 +200,10 @@ def test_single_embedding_query(model_name: str):
             delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize("token_dim,model_name", [(96, "answerdotai/answerai-colbert-small-v1")])
-def test_parallel_processing(token_dim: int, model_name: str):
+def test_parallel_processing():
     is_ci = os.getenv("CI")
-    model = LateInteractionTextEmbedding(model_name=model_name)
-
+    model = LateInteractionTextEmbedding(model_name="colbert-ir/colbertv2.0")
+    token_dim = 128
     docs = ["hello world", "flag embedding"] * 100
     embeddings = list(model.embed(docs, batch_size=10, parallel=2))
     embeddings = np.stack(embeddings, axis=0)
@@ -235,7 +222,10 @@ def test_parallel_processing(token_dim: int, model_name: str):
         delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize("model_name", ["answerdotai/answerai-colbert-small-v1"])
+@pytest.mark.parametrize(
+    "model_name",
+    ["colbert-ir/colbertv2.0"],
+)
 def test_lazy_load(model_name: str):
     is_ci = os.getenv("CI")
 

tests/test_late_interaction_multimodal.py

@@ -1,6 +1,5 @@
 import os
 
-import pytest
 from PIL import Image
 import numpy as np
 
@@ -12,13 +11,15 @@ from tests.config import TEST_MISC_DIR
 CANONICAL_IMAGE_VALUES = {
     "Qdrant/colpali-v1.3-fp16": np.array(
         [
-            [-0.0345, -0.022, 0.0567, -0.0518, -0.0782, 0.1714, -0.1738],
-            [-0.1181, -0.099, 0.0268, 0.0774, 0.0228, 0.0563, -0.1021],
-            [-0.117, -0.0683, 0.0371, 0.0921, 0.0107, 0.0659, -0.0666],
-            [-0.1393, -0.0948, 0.037, 0.0951, -0.0126, 0.0678, -0.087],
-            [-0.0957, -0.081, 0.0404, 0.052, 0.0409, 0.0335, -0.064],
-            [-0.0626, -0.0445, 0.056, 0.0592, -0.0229, 0.0409, -0.0301],
-            [-0.1299, -0.0691, 0.1097, 0.0728, 0.0123, 0.0519, 0.0122],
+            [
+                [-0.0345, -0.022, 0.0567, -0.0518, -0.0782, 0.1714, -0.1738],
+                [-0.1181, -0.099, 0.0268, 0.0774, 0.0228, 0.0563, -0.1021],
+                [-0.117, -0.0683, 0.0371, 0.0921, 0.0107, 0.0659, -0.0666],
+                [-0.1393, -0.0948, 0.037, 0.0951, -0.0126, 0.0678, -0.087],
+                [-0.0957, -0.081, 0.0404, 0.052, 0.0409, 0.0335, -0.064],
+                [-0.0626, -0.0445, 0.056, 0.0592, -0.0229, 0.0409, -0.0301],
+                [-0.1299, -0.0691, 0.1097, 0.0728, 0.0123, 0.0519, 0.0122],
+            ]
         ]
     ),
 }
@@ -46,38 +47,38 @@ images = [
 
 
 def test_batch_embedding():
-    if os.getenv("CI"):
-        pytest.skip("Colpali is too large to test in CI")
+    is_ci = os.getenv("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))
+    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 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:
+                batch_size, token_num, abridged_dim = expected_result.shape
+                assert np.allclose(value[:token_num, :abridged_dim], expected_result, atol=1e-3)
 
 
 def test_single_embedding():
-    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)
+    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)))
+            batch_size, 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():
-    if os.getenv("CI"):
-        pytest.skip("Colpali is too large to test in CI")
+    is_ci = os.getenv("CI")
+    if not is_ci:
+        queries_to_embed = queries
 
-    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)
+        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)

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, should_test_model
+from tests.utils import delete_model_cache
 
 CANONICAL_COLUMN_VALUES = {
-    "prithivida/Splade_PP_en_v1": {
+    "prithvida/Splade_PP_en_v1": {
         "indices": [
             2040,
             2047,
@@ -49,41 +49,28 @@ CANONICAL_COLUMN_VALUES = {
 docs = ["Hello World"]
 
 
-@pytest.mark.parametrize("model_name", ["prithivida/Splade_PP_en_v1"])
-def test_batch_embedding(model_name: str) -> None:
+def test_batch_embedding() -> None:
     is_ci = os.getenv("CI")
     docs_to_embed = docs * 10
 
-    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 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"]
 
-    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)
 
 
-@pytest.mark.parametrize("model_name", ["prithivida/Splade_PP_en_v1"])
-def test_single_embedding(model_name: str) -> None:
+def test_single_embedding() -> None:
     is_ci = os.getenv("CI")
-    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
-
+    for model_name, expected_result in CANONICAL_COLUMN_VALUES.items():
         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"]
 
@@ -93,10 +80,9 @@ def test_single_embedding(model_name: str) -> None:
             delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize("model_name", ["prithivida/Splade_PP_en_v1"])
-def test_parallel_processing(model_name: str) -> None:
+def test_parallel_processing() -> None:
     is_ci = os.getenv("CI")
-    model = SparseTextEmbedding(model_name=model_name)
+    model = SparseTextEmbedding(model_name="prithivida/Splade_PP_en_v1")
     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))
@@ -186,7 +172,10 @@ 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, should_test_model
+from tests.utils import delete_model_cache
 
 CANONICAL_SCORE_VALUES = {
     "Xenova/ms-marco-MiniLM-L-6-v2": np.array([8.500708, -2.541011]),
@@ -15,37 +15,44 @@ 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", ["Xenova/ms-marco-MiniLM-L-6-v2"])
+
+@pytest.mark.parametrize(
+    "model_name",
+    [model_name for model_name in CANONICAL_SCORE_VALUES],
+)
 def test_rerank(model_name: str) -> None:
     is_ci = os.getenv("CI")
-    is_manual = os.getenv("GITHUB_EVENT_NAME") == "workflow_dispatch"
 
-    for model_desc in TextCrossEncoder._list_supported_models():
-        if not should_test_model(model_desc, model_name, is_ci, is_manual):
-            continue
+    model = TextCrossEncoder(model_name=model_name)
 
-        model = TextCrossEncoder(model_name=model_name)
+    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)))
 
-        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)))
+    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}"
 
-        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)
+    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", ["Xenova/ms-marco-MiniLM-L-6-v2"])
+@pytest.mark.parametrize(
+    "model_name",
+    [model_name for model_name in SELECTED_MODELS.values()],
+)
 def test_batch_rerank(model_name: str) -> None:
     is_ci = os.getenv("CI")
 
@@ -71,7 +78,10 @@ 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)
@@ -85,7 +95,10 @@ def test_lazy_load(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",
+    [model_name for model_name in SELECTED_MODELS.values()],
+)
 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 JinaEmbeddingV3, Task
+from fastembed.text.multitask_embedding import Task
 from tests.utils import delete_model_cache
 
 
@@ -60,43 +60,52 @@ CANONICAL_VECTOR_VALUES = {
 docs = ["Hello World", "Follow the white rabbit."]
 
 
-@pytest.mark.parametrize("dim,model_name", [(1024, "jinaai/jina-embeddings-v3")])
-def test_batch_embedding(dim: int, model_name: str):
+def test_batch_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")
-
     docs_to_embed = docs * 10
     default_task = Task.RETRIEVAL_PASSAGE
 
-    model = TextEmbedding(model_name=model_name)
+    for model_desc in TextEmbedding._list_supported_models():
+        if not is_ci and model_desc.size_in_GB > 1:
+            continue
 
-    embeddings = list(model.embed(documents=docs_to_embed, batch_size=6))
-    embeddings = np.stack(embeddings, axis=0)
+        model_name = model_desc.model
+        dim = model_desc.dim
 
-    assert embeddings.shape == (len(docs_to_embed), dim)
+        if model_name not in CANONICAL_VECTOR_VALUES.keys():
+            continue
 
-    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
+        model = TextEmbedding(model_name=model_name)
 
-    if is_ci:
-        delete_model_cache(model.model._model_dir)
+        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")
-    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 JinaEmbeddingV3._list_supported_models():
-        # todo: once we add more models, we should not test models >1GB size locally
+    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]:
@@ -118,17 +127,18 @@ def test_single_embedding():
 
 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 JinaEmbeddingV3._list_supported_models():
-        # todo: once we add more models, we should not test models >1GB size locally
+    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}")
@@ -149,18 +159,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 JinaEmbeddingV3._list_supported_models():
-        # todo: once we add more models, we should not test models >1GB size locally
+    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}")
@@ -179,15 +189,14 @@ def test_single_embedding_passage():
             delete_model_cache(model.model._model_dir)
 
 
-@pytest.mark.parametrize("dim,model_name", [(1024, "jinaai/jina-embeddings-v3")])
-def test_parallel_processing(dim: int, model_name: str):
+def test_parallel_processing():
     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
@@ -209,14 +218,14 @@ def test_parallel_processing(dim: int, model_name: str):
 
 def test_task_assignment():
     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")
+    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
+        if model_name not in CANONICAL_VECTOR_VALUES.keys():
+            continue
 
         model = TextEmbedding(model_name=model_name)
 
@@ -228,14 +237,12 @@ def test_task_assignment():
             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, should_test_model
+from tests.utils import delete_model_cache
 
 CANONICAL_VECTOR_VALUES = {
     "BAAI/bge-small-en": np.array([-0.0232, -0.0255, 0.0174, -0.0639, -0.0006]),
@@ -52,7 +52,7 @@ CANONICAL_VECTOR_VALUES = {
         [0.0802303, 0.3700881, -4.3053818, 0.4431803, -0.271572]
     ),
     "thenlper/gte-large": np.array(
-        [-0.00986551, -0.00018734, 0.00605892, -0.03289612, -0.0387564],
+        [-0.01920587, 0.00113156, -0.00708992, -0.00632304, -0.04025577]
     ),
     "mixedbread-ai/mxbai-embed-large-v1": np.array(
         [0.02295546, 0.03196154, 0.016512, -0.04031524, -0.0219634]
@@ -72,19 +72,17 @@ CANONICAL_VECTOR_VALUES = {
 MULTI_TASK_MODELS = ["jinaai/jina-embeddings-v3"]
 
 
-@pytest.mark.parametrize("model_name", ["BAAI/bge-small-en-v1.5"])
-def test_embedding(model_name: str) -> None:
+def test_embedding() -> 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 model_desc.model in MULTI_TASK_MODELS 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
-        if not should_test_model(model_desc, model_name, is_ci, is_manual):
-            continue
 
         dim = model_desc.dim
 
@@ -97,12 +95,15 @@ def test_embedding(model_name: str) -> 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")])
+@pytest.mark.parametrize(
+    "n_dims,model_name",
+    [(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")],
+)
 def test_batch_embedding(n_dims: int, model_name: str) -> None:
     is_ci = os.getenv("CI")
     model = TextEmbedding(model_name=model_name)
@@ -111,12 +112,15 @@ 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 == (len(docs), n_dims)
+    assert embeddings.shape == (200, 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")])
+@pytest.mark.parametrize(
+    "n_dims,model_name",
+    [(384, "BAAI/bge-small-en-v1.5"), (768, "jinaai/jina-embeddings-v2-base-en")],
+)
 def test_parallel_processing(n_dims: int, model_name: str) -> None:
     is_ci = os.getenv("CI")
     model = TextEmbedding(model_name=model_name)
@@ -131,7 +135,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 == (len(docs), n_dims)
+    assert embeddings.shape == (200, n_dims)
     assert np.allclose(embeddings, embeddings_2, atol=1e-3)
     assert np.allclose(embeddings, embeddings_3, atol=1e-3)
 
@@ -139,7 +143,10 @@ 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,9 +3,7 @@ import traceback
 
 from pathlib import Path
 from types import TracebackType
-from typing import Union, Callable, Any, Type, Optional
-
-from fastembed.common.model_description import BaseModelDescription
+from typing import Union, Callable, Any, Type
 
 
 def delete_model_cache(model_dir: Union[str, Path]) -> None:
@@ -37,31 +35,3 @@ 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