diff --git a/fastembed/image/image_embedding.py b/fastembed/image/image_embedding.py index a36303456..0924eac0f 100644 --- a/fastembed/image/image_embedding.py +++ b/fastembed/image/image_embedding.py @@ -5,11 +5,15 @@ 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]]: diff --git a/fastembed/image/normalized_embedding.py b/fastembed/image/normalized_embedding.py new file mode 100644 index 000000000..0b83fe774 --- /dev/null +++ b/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, + ) diff --git a/fastembed/image/onnx_embedding.py b/fastembed/image/onnx_embedding.py index 3e1739951..38b79743a 100644 --- a/fastembed/image/onnx_embedding.py +++ b/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, diff --git a/tests/test_image_onnx_embeddings.py b/tests/test_image_onnx_embeddings.py index 5ac5e44f1..f9fe24952 100644 --- a/tests/test_image_onnx_embeddings.py +++ b/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 @@ "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 get_model(model_name: str): @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