add nomic-embed-vision-v1.5 (#651)

fastembed/image/image_embedding.py

@@ -5,11 +5,15 @@ from fastembed.common.types import NumpyArray, Device
 from fastembed.common import ImageInput, OnnxProvider
 from fastembed.image.image_embedding_base import ImageEmbeddingBase
 from fastembed.image.onnx_embedding import OnnxImageEmbedding
+from fastembed.image.normalized_embedding import NormalizedEmbedding
 from fastembed.common.model_description import DenseModelDescription
 
 
 class ImageEmbedding(ImageEmbeddingBase):
-    EMBEDDINGS_REGISTRY: list[Type[ImageEmbeddingBase]] = [OnnxImageEmbedding]
+    EMBEDDINGS_REGISTRY: list[Type[ImageEmbeddingBase]] = [
+        OnnxImageEmbedding,
+        NormalizedEmbedding,
+    ]
 
     @classmethod
     def list_supported_models(cls) -> list[dict[str, Any]]:

fastembed/image/normalized_embedding.py

@@ -0,0 +1,69 @@
+from typing import Any, Iterable, Type
+
+
+from fastembed.common.types import NumpyArray
+from fastembed.common.onnx_model import OnnxOutputContext
+from fastembed.common.utils import normalize
+from fastembed.image.onnx_embedding import OnnxImageEmbedding
+from fastembed.image.onnx_image_model import ImageEmbeddingWorker
+from fastembed.common.model_description import DenseModelDescription, ModelSource
+
+supported_normalized_models: list[DenseModelDescription] = [
+    DenseModelDescription(
+        model="nomic-ai/nomic-embed-vision-v1.5",
+        dim=768,
+        description="Image embeddings, Multimodal (text&image), 2024 year",
+        license="apache-2.0",
+        size_in_GB=0.37,
+        sources=ModelSource(hf="nomic-ai/nomic-embed-vision-v1.5"),
+        model_file="onnx/model.onnx",
+    ),
+    DenseModelDescription(
+        model="nomic-ai/nomic-embed-vision-v1.5-Q",
+        dim=768,
+        description="Image embeddings, Multimodal (text&image), 2024 year",
+        license="apache-2.0",
+        size_in_GB=0.1,
+        sources=ModelSource(hf="nomic-ai/nomic-embed-vision-v1.5"),
+        model_file="onnx/model_quantized.onnx",
+    ),
+]
+
+
+class NormalizedEmbedding(OnnxImageEmbedding):
+    @classmethod
+    def _list_supported_models(cls) -> list[DenseModelDescription]:
+        """
+        Lists the supported models.
+
+        Returns:
+            list[DenseModelDescription]: A list of DenseModelDescription objects containing the model information.
+        """
+        return supported_normalized_models
+
+    @classmethod
+    def _get_worker_class(cls) -> Type["ImageEmbeddingWorker[NumpyArray]"]:
+        return NormalizedEmbeddingWorker
+
+    def _post_process_onnx_output(
+        self, output: OnnxOutputContext, **kwargs: Any
+    ) -> Iterable[NumpyArray]:
+        # The model emits last_hidden_state, which onnx_embed flattens to (batch, tokens * dim).
+        # Recover the token axis, take the CLS token (index 0) and normalize, matching the reference
+        # F.normalize(last_hidden_state[:, 0], p=2, dim=1).
+        dim = self.model_description.dim
+        assert dim is not None, "Model description is missing the embedding dim"
+        hidden_states = output.model_output.reshape(output.model_output.shape[0], -1, dim)
+        return normalize(hidden_states[:, 0])
+
+
+class NormalizedEmbeddingWorker(ImageEmbeddingWorker[NumpyArray]):
+    def init_embedding(
+        self, model_name: str, cache_dir: str, **kwargs: Any
+    ) -> NormalizedEmbedding:
+        return NormalizedEmbedding(
+            model_name=model_name,
+            cache_dir=cache_dir,
+            threads=1,
+            **kwargs,
+        )

fastembed/image/onnx_embedding.py

@@ -63,7 +63,6 @@ class OnnxImageEmbedding(ImageEmbeddingBase, OnnxImageModel[NumpyArray]):
     def __init__(
         self,
         model_name: str,
-
         cache_dir: str | None = None,
         threads: int | None = None,
         providers: Sequence[OnnxProvider] | None = None,

tests/test_image_onnx_embeddings.py

@@ -1,4 +1,5 @@
 import os
+import platform
 from contextlib import contextmanager
 from io import BytesIO
 
@@ -25,6 +26,12 @@ CANONICAL_VECTOR_VALUES = {
     "jinaai/jina-clip-v1": np.array(
         [-0.029, 0.0216, 0.0396, 0.0283, -0.0023, 0.0151, 0.011, -0.0235, 0.0251, -0.0343]
     ),
+    "nomic-ai/nomic-embed-vision-v1.5": np.array(
+        [0.0048, -0.0254, 0.0067, -0.0296, -0.0435, -0.0123, 0.0024, -0.0361, -0.0703, -0.0186]
+    ),
+    "nomic-ai/nomic-embed-vision-v1.5-Q": np.array(
+        [-0.0011, -0.0477, 0.0024, -0.049, -0.0458, -0.0314, 0.017, -0.0383, -0.0537, -0.021]
+    ),
 }
 
 _MODELS_TO_CACHE = ("Qdrant/clip-ViT-B-32-vision",)
@@ -59,9 +66,13 @@ def model_cache():
 @pytest.mark.parametrize("model_name", ["Qdrant/clip-ViT-B-32-vision"])
 def test_embedding(model_cache, 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 ImageEmbedding._list_supported_models():
+        # quantized int8 ops diverge on macOS; canonical vector is generated on linux/amd64 (CI)
+        if is_mac and model_desc.model == "nomic-ai/nomic-embed-vision-v1.5-Q":
+            continue
         if not should_test_model(model_desc, model_name, is_ci, is_manual):
             continue