Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion fastembed/image/image_embedding.py
Original file line number Diff line number Diff line change
Expand Up @@ -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]]:
Expand Down
69 changes: 69 additions & 0 deletions fastembed/image/normalized_embedding.py
Original file line number Diff line number Diff line change
@@ -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,
)
1 change: 0 additions & 1 deletion fastembed/image/onnx_embedding.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down
11 changes: 11 additions & 0 deletions tests/test_image_onnx_embeddings.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
import os
import platform
from contextlib import contextmanager
from io import BytesIO

Expand All @@ -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",)
Expand Down Expand Up @@ -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

Expand Down