From b472d3534e452c327cb32499c11f196285d3eef3 Mon Sep 17 00:00:00 2001 From: Minamiyama Date: Wed, 7 Oct 2026 13:52:23 +0800 Subject: [PATCH 1/2] feat(embedding): add EmbeddingGemma 2 support --- .../builtin/embedding/embeddinggemma-2.rst | 21 + xinference/model/embedding/__init__.py | 8 + xinference/model/embedding/embeddinggemma2.py | 295 ++++++++++++ .../embedding/models/embeddinggemma-2.json | 55 +++ .../embedding/sentence_transformers/core.py | 2 + .../embedding/tests/test_embeddinggemma2.py | 444 ++++++++++++++++++ xinference/model/embedding/vllm/core.py | 120 ++++- 7 files changed, 941 insertions(+), 4 deletions(-) create mode 100644 doc/source/models/builtin/embedding/embeddinggemma-2.rst create mode 100644 xinference/model/embedding/embeddinggemma2.py create mode 100644 xinference/model/embedding/models/embeddinggemma-2.json create mode 100644 xinference/model/embedding/tests/test_embeddinggemma2.py diff --git a/doc/source/models/builtin/embedding/embeddinggemma-2.rst b/doc/source/models/builtin/embedding/embeddinggemma-2.rst new file mode 100644 index 0000000000..878e3352bd --- /dev/null +++ b/doc/source/models/builtin/embedding/embeddinggemma-2.rst @@ -0,0 +1,21 @@ +.. _models_builtin_embeddinggemma-2: + +================ +embeddinggemma-2 +================ + +- **Model Name:** embeddinggemma-2 +- **Languages:** multilingual +- **Abilities:** embed + +Specifications +^^^^^^^^^^^^^^ + +- **Dimensions:** 768 +- **Max Tokens:** 8192 +- **Model ID:** google/embeddinggemma-2 +- **Model Hubs**: `Hugging Face `__, `ModelScope `__ + +Execute the following command to launch the model:: + + xinference launch --model-name embeddinggemma-2 --model-type embedding \ No newline at end of file diff --git a/xinference/model/embedding/__init__.py b/xinference/model/embedding/__init__.py index fa7b06c9b7..3ac62c511c 100644 --- a/xinference/model/embedding/__init__.py +++ b/xinference/model/embedding/__init__.py @@ -210,6 +210,10 @@ def _install(): generate_embedding_description(model_spec) ) + from .embeddinggemma2 import ( + SentenceTransformerEmbeddingGemma2Model, + TransformersEmbeddingGemma2Model, + ) from .flag.core import FlagEmbeddingModel from .llama_cpp.core import XllamaCppEmbeddingModel from .sentence_transformers.core import SentenceTransformerEmbeddingModel @@ -218,6 +222,9 @@ def _install(): extend_classes_once( SENTENCE_TRANSFORMER_CLASSES, [SentenceTransformerEmbeddingModel] ) + # Specific modality-aware adapters take precedence over the generic text path. + if SentenceTransformerEmbeddingGemma2Model not in SENTENCE_TRANSFORMER_CLASSES: + SENTENCE_TRANSFORMER_CLASSES.insert(0, SentenceTransformerEmbeddingGemma2Model) extend_classes_once(FLAG_EMBEDDER_CLASSES, [FlagEmbeddingModel]) extend_classes_once(VLLM_CLASSES, [VLLMEmbeddingModel]) extend_classes_once(LLAMA_CPP_CLASSES, [XllamaCppEmbeddingModel]) @@ -226,6 +233,7 @@ def _install(): SUPPORTED_ENGINES["flag"] = FLAG_EMBEDDER_CLASSES SUPPORTED_ENGINES["vllm"] = VLLM_CLASSES SUPPORTED_ENGINES["llama.cpp"] = LLAMA_CPP_CLASSES + SUPPORTED_ENGINES["transformers"] = [TransformersEmbeddingGemma2Model] # Distribution-specific engines are appended after the built-ins. _run_engine_registration_hooks(MODEL_TYPE_EMBEDDING, SUPPORTED_ENGINES) diff --git a/xinference/model/embedding/embeddinggemma2.py b/xinference/model/embedding/embeddinggemma2.py new file mode 100644 index 0000000000..d5b39f7c7a --- /dev/null +++ b/xinference/model/embedding/embeddinggemma2.py @@ -0,0 +1,295 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import os +from importlib.metadata import PackageNotFoundError, version +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +from packaging.version import Version + +from ...device_utils import get_available_device +from ...types import Embedding, EmbeddingData, EmbeddingUsage +from ..batch import BatchMixin +from ..utils import check_dependency_available, virtual_env_allows_missing_engine +from .core import EmbeddingModel, EmbeddingModelFamilyV2, EmbeddingSpecV1 + +MODEL_NAME = "embeddinggemma-2" +MATRYOSHKA_DIMENSIONS = (128, 256, 512, 768) + + +def validate_dimensions(dimensions: Optional[int]) -> None: + if dimensions is not None and ( + not isinstance(dimensions, int) + or isinstance(dimensions, bool) + or dimensions not in MATRYOSHKA_DIMENSIONS + ): + raise ValueError( + f"EmbeddingGemma 2 dimensions must be one of {MATRYOSHKA_DIMENSIONS}, " + f"got {dimensions}." + ) + + +def resolve_dtype(value: Any, device: str) -> torch.dtype: + if value is None or value == "auto": + return ( + torch.bfloat16 + if device.startswith("cuda") and torch.cuda.is_bf16_supported() + else torch.float32 + ) + if value in ("bfloat16", "bf16", torch.bfloat16): + return torch.bfloat16 + if value in ("float32", "fp32", torch.float32): + return torch.float32 + raise ValueError( + "EmbeddingGemma 2 requires bfloat16 or float32; float16 produces invalid embeddings." + ) + + +def check_dependencies(engine: str) -> Union[bool, Tuple[bool, str]]: + requirements = [("transformers", "5.19.0")] + if engine == "sentence_transformers": + from ..utils import neutralize_broken_torchcodec + + neutralize_broken_torchcodec() + requirements.append(("sentence-transformers", "6.1.0")) + for package, minimum in requirements: + module = package.replace("-", "_") + available = check_dependency_available(module, package) + if available != True: + return available + try: + installed = version(package) + except PackageNotFoundError: + return False, f"Cannot determine the installed {package} version" + if Version(installed) < Version(minimum): + return ( + False, + f"EmbeddingGemma 2 requires {package}>={minimum}, current: {installed}", + ) + return True + + +def load_prompts(model_path: str) -> Dict[str, str]: + with open( + os.path.join(model_path, "config_sentence_transformers.json"), encoding="utf-8" + ) as file: + return json.load(file).get("prompts", {}) + + +def normalize_inputs( + inputs: Any, prompts: Dict[str, str], **kwargs: Any +) -> List[List[Dict[str, Any]]]: + """Keep each input as one ordered, possibly interleaved multimodal sample.""" + prompt = kwargs.get("prompt") + prompt_name = kwargs.get("prompt_name") + if prompt is None and prompt_name is not None: + if prompt_name not in prompts: + raise ValueError(f"Unknown EmbeddingGemma 2 prompt_name: {prompt_name}") + prompt = prompts[prompt_name] + items = [inputs] if isinstance(inputs, (str, dict)) else inputs + if not isinstance(items, list): + raise ValueError( + "EmbeddingGemma 2 input must be text, a multimodal dict, or a list of these." + ) + result = [] + for item in items: + if isinstance(item, str): + content = [{"type": "text", "text": item}] + elif isinstance(item, dict): + content = [] + for modality, value in item.items(): + if modality not in ("text", "image", "video", "audio"): + raise ValueError( + f"Unsupported EmbeddingGemma 2 input field: {modality}" + ) + if value is None: + continue + values = value if isinstance(value, list) else [value] + for part in values: + if modality == "text" and not isinstance(part, str): + raise ValueError("EmbeddingGemma 2 text must be a string.") + content.append({"type": modality, modality: part}) + if not content: + raise ValueError( + "EmbeddingGemma 2 input dict must contain text, image, video, or audio." + ) + else: + raise ValueError( + "EmbeddingGemma 2 input items must be strings or multimodal dicts." + ) + # Prefixes belong to text. Media-only inputs must not acquire a task prefix. + if prompt: + for part in content: + if part["type"] == "text": + part["text"] = prompt + part["text"] + break + result.append([{"role": "user", "content": content}]) + return result + + +class TransformersEmbeddingGemma2Model(EmbeddingModel, BatchMixin): + engine = "transformers" + _device: Optional[str] + + def __init__(self, *args, **kwargs) -> None: + EmbeddingModel.__init__(self, *args, **kwargs) + BatchMixin.__init__(self, self.create_embedding, **kwargs) # type: ignore + self._dimensions = self._kwargs.get("dimensions") + validate_dimensions(self._dimensions) + self._prompts: Dict[str, str] = {} + + @classmethod + def check_lib(cls) -> Union[bool, Tuple[bool, str]]: + return check_dependencies(cls.engine) + + @classmethod + def match_json( + cls, + model_family: EmbeddingModelFamilyV2, + model_spec: EmbeddingSpecV1, + quantization: str, + ) -> Union[bool, Tuple[bool, str]]: + if ( + model_family.model_name != MODEL_NAME + or model_spec.model_format != "pytorch" + ): + return False + return True if virtual_env_allows_missing_engine() else cls.check_lib() + + def load(self) -> None: + dependencies = self.check_lib() + if dependencies != True: + raise ImportError(dependencies[1]) # type: ignore[index] + from transformers import AutoModel, AutoProcessor + + self._device = self._device or get_available_device() + dtype = resolve_dtype(self._kwargs.get("torch_dtype"), self._device) + self._processor = AutoProcessor.from_pretrained(self._model_path) + self._tokenizer = self._processor.tokenizer + self._model = ( + AutoModel.from_pretrained( + self._model_path, + dtype=dtype, + **self._kwargs.get("config_kwargs", {}), + ) + .to(self._device) + .eval() + ) + self._prompts = load_prompts(self._model_path) + + def _encode(self, messages: List[List[Dict[str, Any]]], **kwargs: Any): + embeddings = [] + tokens = 0 + assert self._model is not None + # Media have variable token lengths. One sample at a time avoids padding + # mixed image/video/audio tensors and still uses Xinference request batching. + for sample in messages: + features = self._processor.apply_chat_template( + sample, + tokenize=True, + return_dict=True, + return_tensors="pt", + processor_kwargs={ + "truncation": True, + "max_length": self.model_family.max_tokens, + **kwargs.get("processing_kwargs", {}), + }, + ).to(device=self._device, dtype=self._model.dtype) + with torch.inference_mode(): + hidden = self._model(**features).last_hidden_state.float() + mask = features["attention_mask"].unsqueeze(-1) + pooled = (hidden * mask).sum(1) / mask.sum(1).clamp(min=1) + embeddings.append(pooled[0]) + tokens += int(features["attention_mask"].sum().item()) + return embeddings, tokens + + def _create_embedding(self, sentences: Any, **kwargs: Any) -> Embedding: + sentences = self._fix_langchain_openai_inputs(sentences) + dimensions = kwargs.pop("dimensions", self._dimensions) + validate_dimensions(dimensions) + if kwargs.pop("return_sparse", False): + raise ValueError("EmbeddingGemma 2 does not support sparse embeddings.") + messages = normalize_inputs(sentences, self._prompts, **kwargs) + embeddings, tokens = self._encode(messages, **kwargs) if messages else ([], 0) + normalize = kwargs.get( + "normalize_embeddings", kwargs.get("normalize_embedding", True) + ) + data = [] + for index, vector in enumerate(embeddings): + vector = vector.float()[:dimensions] + if normalize: + vector = torch.nn.functional.normalize(vector, p=2, dim=0) + data.append( + EmbeddingData( + index=index, object="embedding", embedding=vector.tolist() + ) + ) + self._clean_cache_if_needed(tokens) + return Embedding( + object="list", + model=kwargs.get("model_uid"), + model_replica=self._model_uid, + data=data, + usage=EmbeddingUsage(prompt_tokens=tokens, total_tokens=tokens), + ) + + +class SentenceTransformerEmbeddingGemma2Model(TransformersEmbeddingGemma2Model): + engine = "sentence_transformers" + + def load(self) -> None: + dependencies = self.check_lib() + if dependencies != True: + raise ImportError(dependencies[1]) # type: ignore[index] + from sentence_transformers import SentenceTransformer + + self._device = self._device or get_available_device() + dtype = resolve_dtype(self._kwargs.get("torch_dtype"), self._device) + self._model = SentenceTransformer( + self._model_path, + device=self._device, + model_kwargs={"dtype": dtype}, + config_kwargs=self._kwargs.get("config_kwargs"), + ) + self._model.max_seq_length = self.model_family.max_tokens + self._tokenizer = self._model.tokenizer + self._prompts = self._model.prompts + + def _encode(self, messages: List[List[Dict[str, Any]]], **kwargs: Any): + assert self._model is not None + tokens = 0 + + def count_tokens(module: Any, args: Tuple[Any, ...]) -> None: + nonlocal tokens + tokens += int(args[0]["attention_mask"].sum().item()) + + # Native encode owns modality routing and preserves batch order. Asking + # it for all features fails on mixed batches because media tensors have + # a different leading dimension; count the expanded token mask instead. + hook = self._model.register_forward_pre_hook(count_tokens) + try: + with torch.inference_mode(): + outputs = self._model.encode( + messages, + prompt="", + convert_to_tensor=True, + convert_to_numpy=False, + batch_size=kwargs.get("batch_size", 1), + processing_kwargs=kwargs.get("processing_kwargs"), + ) + finally: + hook.remove() + return list(outputs), tokens diff --git a/xinference/model/embedding/models/embeddinggemma-2.json b/xinference/model/embedding/models/embeddinggemma-2.json new file mode 100644 index 0000000000..8e4f827012 --- /dev/null +++ b/xinference/model/embedding/models/embeddinggemma-2.json @@ -0,0 +1,55 @@ +[ + { + "version": 2, + "model_name": "embeddinggemma-2", + "dimensions": 768, + "max_tokens": 8192, + "model_ability": [ + "embed_vision", + "embed_video", + "embed_audio" + ], + "language": [ + "multilingual" + ], + "model_specs": [ + { + "model_format": "pytorch", + "model_src": { + "huggingface": { + "model_id": "google/embeddinggemma-2", + "model_revision": "main", + "quantizations": [ + "none" + ] + }, + "modelscope": { + "model_id": "google/embeddinggemma-2", + "model_revision": "master", + "quantizations": [ + "none" + ] + } + } + } + ], + "virtualenv": { + "packages": [ + "transformers>=5.19.0", + "sentence-transformers>=6.1.0 ; #engine# == \"sentence_transformers\"", + "vllm>=0.32.0 ; #engine# == \"vllm\"", + "#system_torch# ; #engine# == \"sentence_transformers\"", + "#system_torchvision# ; #engine# == \"sentence_transformers\"", + "#system_torch# ; #engine# == \"transformers\"", + "#system_torchvision# ; #engine# == \"transformers\"", + "accelerate", + "pillow", + "librosa", + "soundfile", + "torchcodec" + ] + }, + "featured": true, + "updated_at": 1791324565 + } +] diff --git a/xinference/model/embedding/sentence_transformers/core.py b/xinference/model/embedding/sentence_transformers/core.py index 9b021660ad..c6b2b3ff8b 100644 --- a/xinference/model/embedding/sentence_transformers/core.py +++ b/xinference/model/embedding/sentence_transformers/core.py @@ -1035,6 +1035,8 @@ def match_json( model_spec: EmbeddingSpecV1, quantization: str, ) -> Union[bool, Tuple[bool, str]]: + if model_family.model_name == "embeddinggemma-2": + return False, "EmbeddingGemma 2 requires its modality-aware adapter" if is_wemm_model(model_family.model_name): if not virtual_env_allows_missing_engine(): dep_check = check_dependency_available( diff --git a/xinference/model/embedding/tests/test_embeddinggemma2.py b/xinference/model/embedding/tests/test_embeddinggemma2.py new file mode 100644 index 0000000000..b93ad6416f --- /dev/null +++ b/xinference/model/embedding/tests/test_embeddinggemma2.py @@ -0,0 +1,444 @@ +# Copyright 2022-2026 Xinference Holdings Pte. Ltd +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +import json +import sys +import types +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + +from xinference.model.utils import virtualenv_discovery_var + +from .. import embeddinggemma2 as gemma +from ..embed_family import match_embedding +from ..sentence_transformers.core import SentenceTransformerEmbeddingModel +from ..vllm.core import VLLMEmbeddingModel + + +def _make_model(cls=gemma.SentenceTransformerEmbeddingGemma2Model, **kwargs): + family = match_embedding("embeddinggemma-2", "pytorch", "none") + model = cls("gemma", "/unused", family, **kwargs) + model._model = MagicMock() + model._prompts = {"SearchQuery": "task: search result | query: "} + model._clean_cache_if_needed = lambda *args: None + return model + + +def test_catalog_hubs_and_engine_discovery(): + from .. import generate_engine_config_by_model_name + + for hub, revision in (("huggingface", "main"), ("modelscope", "master")): + family = match_embedding("embeddinggemma-2", "pytorch", "none", hub) + assert family.model_specs[0].model_id == "google/embeddinggemma-2" + assert family.model_specs[0].model_revision == revision + assert family.model_specs[0].model_hub == hub + assert family.model_ability == ["embed_vision", "embed_video", "embed_audio"] + engines = {} + token = virtualenv_discovery_var.set(True) + try: + family = match_embedding("embeddinggemma-2", "pytorch", "none") + generate_engine_config_by_model_name(family, engines) + assert set(engines["embeddinggemma-2"]) == { + "sentence_transformers", + "transformers", + "vllm", + } + assert ( + engines["embeddinggemma-2"]["sentence_transformers"][0]["embedding_class"] + is gemma.SentenceTransformerEmbeddingGemma2Model + ) + assert ( + SentenceTransformerEmbeddingModel.match_json( + family, family.model_specs[0], "none" + ) + is not True + ) + unsupported = SimpleNamespace(model_format="ggufv2") + assert ( + gemma.TransformersEmbeddingGemma2Model.match_json( + family, unsupported, "none" + ) + is False + ) + assert VLLMEmbeddingModel.match_json(family, unsupported, "none")[0] is False + finally: + virtualenv_discovery_var.reset(token) + + +@pytest.mark.parametrize("engine", ["sentence_transformers", "transformers", "vllm"]) +def test_virtualenv_keeps_system_torch_pins_only_for_native_engines(engine): + from xinference.core.utils import filter_virtualenv_packages_by_markers + from xinference.core.virtual_env_manager import ensure_system_torch_pin + + family = match_embedding("embeddinggemma-2", "pytorch", "none") + packages = filter_virtualenv_packages_by_markers( + family.virtualenv.packages, engine, "13.0" + ) + packages = ensure_system_torch_pin(packages) + for pin in ("#system_torch#", "#system_torchvision#"): + assert packages.count(pin) == (0 if engine == "vllm" else 1) + assert any(pkg.startswith("sentence-transformers") for pkg in packages) == ( + engine == "sentence_transformers" + ) + assert [pkg for pkg in packages if pkg.startswith("vllm")] == ( + ["vllm>=0.32.0"] if engine == "vllm" else [] + ) + + +def test_interleaved_inputs_keep_media_and_do_not_mutate(): + item = { + "text": "before <|image|> between <|image|> <|video|> <|audio|>", + "image": ["a.png", "b.png"], + "video": "clip.mp4", + "audio": "sound.wav", + } + original = copy.deepcopy(item) + messages = gemma.normalize_inputs( + [item, {"audio": "only.wav"}], + {"Document": "title: none | text: "}, + prompt_name="Document", + ) + assert item == original + parts = messages[0][0]["content"] + assert [part["type"] for part in parts] == [ + "text", + "image", + "image", + "video", + "audio", + ] + assert parts[0]["text"] == "title: none | text: " + item["text"] + assert messages[1][0]["content"] == [{"type": "audio", "audio": "only.wav"}] + + +@pytest.mark.parametrize("value", [{}, {"image_url": "x"}, {"text": 5}, [3]]) +def test_invalid_inputs_are_rejected(value): + with pytest.raises(ValueError): + gemma.normalize_inputs(value, {}) + + +@pytest.mark.parametrize("dimensions", [64, 1024, 128.0, True]) +def test_invalid_dimensions_are_rejected(dimensions): + with pytest.raises(ValueError, match="dimensions must be"): + gemma.validate_dimensions(dimensions) + + +@pytest.mark.parametrize("dtype", ["float16", "fp16", torch.float16]) +def test_float16_is_rejected(dtype): + with pytest.raises(ValueError, match="float16 produces invalid"): + gemma.resolve_dtype(dtype, "cpu") + assert gemma.resolve_dtype(None, "cpu") == torch.float32 + + +def test_sentence_transformers_uses_native_encode_and_normalizes_after_truncation(): + model = _make_model() + callbacks = [] + model._model.register_forward_pre_hook.side_effect = ( + lambda callback: callbacks.append(callback) + or model._model.register_forward_pre_hook.return_value + ) + + def encode(*args, **kwargs): + callbacks[0]( + None, ({"attention_mask": torch.cat((torch.ones(283), torch.zeros(3)))},) + ) + return torch.stack([torch.arange(1, 769, dtype=torch.float32), torch.ones(768)]) + + model._model.encode.side_effect = encode + result = model._create_embedding( + ["query", {"image": "x.png"}], + dimensions=128, + prompt_name="SearchQuery", + model_uid="public", + ) + assert result["model"] == "public" + assert result["model_replica"] == "gemma" + assert [row["index"] for row in result["data"]] == [0, 1] + assert result["usage"]["total_tokens"] == 283 + for row in result["data"]: + assert len(row["embedding"]) == 128 + assert torch.linalg.vector_norm( + torch.tensor(row["embedding"]) + ).item() == pytest.approx(1) + args, kwargs = model._model.encode.call_args + assert args[0][0][0]["content"][0]["text"].startswith( + "task: search result | query: " + ) + assert args[0][1][0]["content"] == [{"type": "image", "image": "x.png"}] + assert kwargs["convert_to_tensor"] is True + assert "output_value" not in kwargs + model._model.register_forward_pre_hook.return_value.remove.assert_called_once() + model._model.tokenize.assert_not_called() + + +def test_single_and_empty_inputs_have_correct_cardinality(): + model = _make_model(dimensions=256) + model._model.encode.return_value = torch.ones(1, 768) + assert ( + len(model._create_embedding({"audio": "x.wav"})["data"][0]["embedding"]) == 256 + ) + model._model.encode.reset_mock() + assert model._create_embedding([])["data"] == [] + model._model.encode.assert_not_called() + + +def test_sentence_transformers_removes_token_hook_after_failure(): + model = _make_model() + model._model.encode.side_effect = RuntimeError("encode failed") + with pytest.raises(RuntimeError, match="encode failed"): + model._create_embedding("text") + model._model.register_forward_pre_hook.return_value.remove.assert_called_once() + + +def test_native_transformers_and_sentence_transformers_agree_on_mixed_batch(tmp_path): + transformers = pytest.importorskip("transformers", minversion="5.19.0") + st = pytest.importorskip("sentence_transformers", minversion="6.1.0") + from PIL import Image + from sentence_transformers.base.modules.normalize import Normalize + from sentence_transformers.base.modules.transformer import Transformer + from sentence_transformers.sentence_transformer.modules.pooling import Pooling + from tokenizers import Tokenizer, models, pre_tokenizers, processors + + vocab = { + "": 0, + "": 1, + "": 2, + "": 3, + "hello": 4, + "<|image|>": 5, + "<|audio|>": 6, + "": 7, + "": 8, + "": 9, + "": 10, + } + tokenizer = Tokenizer(models.WordLevel(vocab, unk_token="")) + tokenizer.pre_tokenizer = pre_tokenizers.Whitespace() + tokenizer.post_processor = processors.TemplateProcessing( + single=" $A ", special_tokens=[("", 2), ("", 1)] + ) + tokenizer = transformers.PreTrainedTokenizerFast( + tokenizer_object=tokenizer, + pad_token="", + bos_token="", + eos_token="", + unk_token="", + model_max_length=8192, + model_specific_special_tokens={ + "image_token": "<|image|>", + "audio_token": "<|audio|>", + "boi_token": "", + "eoi_token": "", + "boa_token": "", + "eoa_token": "", + }, + ) + config = transformers.EmbeddingGemma2Config( + text_config=transformers.EmbeddingGemma2TextConfig( + hidden_size=8, + hidden_size_per_layer_input=8, + intermediate_size=16, + num_hidden_layers=1, + num_attention_heads=1, + num_key_value_heads=1, + head_dim=8, + vocab_size=32, + layer_types=["full_attention"], + per_layer_config={}, + ), + vision_config=transformers.Gemma4VisionConfig( + hidden_size=32, + intermediate_size=64, + num_hidden_layers=1, + num_attention_heads=1, + num_key_value_heads=1, + head_dim=32, + global_head_dim=32, + ), + audio_config=None, + image_token_id=5, + boi_token_id=7, + eoi_token_id=8, + ) + transformers.EmbeddingGemma2Model(config).save_pretrained(tmp_path) + processor = transformers.EmbeddingGemma2Processor( + transformers.Gemma4AudioFeatureExtractor(), + transformers.Gemma4ImageProcessor(), + tokenizer, + transformers.Gemma4VideoProcessor(), + chat_template="{% for message in messages %}{% for part in message['content'] %}" + "{% if part['type'] == 'text' %}{{ part['text'] }}{% else %}<|image|>{% endif %}" + "{% endfor %}{% endfor %}", + ) + processor.save_pretrained(tmp_path) + backbone = Transformer( + str(tmp_path), + transformer_task="feature-extraction", + modality_config={ + key: { + "method": "forward", + "method_output_name": "last_hidden_state", + **({"format": "structured"} if key == "message" else {}), + } + for key in ("text", "image", "message") + }, + module_output_name="token_embeddings", + ) + st.SentenceTransformer( + modules=[backbone, Pooling(768, pooling_mode="mean"), Normalize()] + ).save(str(tmp_path)) + image = tmp_path / "image.png" + Image.new("RGB", (32, 32), (100, 50, 0)).save(image) + results = [] + for cls in ( + gemma.TransformersEmbeddingGemma2Model, + gemma.SentenceTransformerEmbeddingGemma2Model, + ): + model = cls( + "native", str(tmp_path), match_embedding("embeddinggemma-2"), device="cpu" + ) + model.load() + # A single image in a two-item batch used to make output_value=None + # index media tensors as if they had one row for every sample. + result = model._create_embedding( + ["hello", {"image": str(image)}], batch_size=2, dimensions=128 + ) + assert len(result["data"]) == 2 + results.append(result) + assert results[0]["usage"] == results[1]["usage"] + for left, right in zip(results[0]["data"], results[1]["data"]): + torch.testing.assert_close( + torch.tensor(left["embedding"]), + torch.tensor(right["embedding"]), + atol=1e-5, + rtol=1e-4, + ) + + +def test_transformers_mean_pooling_masks_padding(): + model = _make_model(gemma.TransformersEmbeddingGemma2Model) + features = { + "input_ids": torch.tensor([[2, 3, 0]]), + "attention_mask": torch.tensor([[1, 1, 0]]), + } + batch = MagicMock() + batch.to.return_value = features + model._processor = MagicMock() + model._processor.apply_chat_template.return_value = batch + hidden = torch.ones(1, 3, 768) + hidden[0, 0, 0] = 2 + hidden[0, 1, 0] = 4 + hidden[0, 2, 0] = 1000 + model._model.return_value = SimpleNamespace(last_hidden_state=hidden) + result = model._create_embedding("text", normalize_embedding=False) + assert result["data"][0]["embedding"][0] == 3 + assert result["usage"]["prompt_tokens"] == 2 + + +def _fake_vllm(monkeypatch, supported=True): + module = types.ModuleType("vllm") + module.__version__ = "0.32.0" + module.ModelRegistry = SimpleNamespace( + get_supported_archs=lambda: ["EmbeddingGemma2Model"] if supported else [] + ) + module.PoolingParams = lambda **kwargs: SimpleNamespace(**kwargs) + module.LLM = MagicMock() + monkeypatch.setitem(sys.modules, "vllm", module) + utils = types.ModuleType("vllm.multimodal.utils") + for modality in ("image", "video", "audio"): + setattr( + utils, + f"fetch_{modality}", + lambda value, modality=modality: f"{modality}:{value}", + ) + monkeypatch.setitem(sys.modules, "vllm.multimodal.utils", utils) + return module + + +def test_vllm_rejects_missing_architecture(monkeypatch): + _fake_vllm(monkeypatch, supported=False) + family = match_embedding("embeddinggemma-2", "pytorch", "none") + token = virtualenv_discovery_var.set(False) + try: + matched = VLLMEmbeddingModel.match_json(family, family.model_specs[0], "none") + assert matched[0] is False + assert "vLLM>=0.32.0" in matched[1] + finally: + virtualenv_discovery_var.reset(token) + + +def test_vllm_load_sets_safe_dtype_and_matryoshka(monkeypatch, tmp_path): + module = _fake_vllm(monkeypatch) + (tmp_path / "chat_template.jinja").write_text("template") + (tmp_path / "config_sentence_transformers.json").write_text( + json.dumps({"prompts": {}}) + ) + model = VLLMEmbeddingModel( + "gemma", + str(tmp_path), + match_embedding("embeddinggemma-2"), + dimensions=256, + hf_overrides={"vision_config": None}, + torch_dtype="bfloat16", + ) + model.load() + kwargs = module.LLM.call_args.kwargs + assert kwargs["runner"] == "pooling" + assert kwargs["max_model_len"] == 8192 + assert kwargs["dtype"] == "bfloat16" + assert kwargs["hf_overrides"]["vision_config"] is None + assert kwargs["hf_overrides"]["matryoshka_dimensions"] == [128, 256, 512, 768] + assert "dimensions" not in kwargs + assert model._dimensions == 256 + + +def test_vllm_preserves_all_modalities_and_pooling_options(monkeypatch): + _fake_vllm(monkeypatch) + model = VLLMEmbeddingModel("gemma", "/unused", match_embedding("embeddinggemma-2")) + model._model = MagicMock() + model._tokenizer = MagicMock() + model._tokenizer.apply_chat_template.return_value = "rendered" + model._chat_template = "template" + params = object() + model._embed_embeddinggemma2( + { + "text": "caption", + "image": ["a.png", "b.png"], + "video": "clip.mp4", + "audio": "sound.wav", + }, + params, + ) + item = model._model.embed.call_args.args[0][0] + assert item["multi_modal_data"] == { + "image": [ + f"image:file://{Path('a.png').absolute()}", + f"image:file://{Path('b.png').absolute()}", + ], + "video": [f"video:file://{Path('clip.mp4').absolute()}"], + "audio": [f"audio:file://{Path('sound.wav').absolute()}"], + } + assert model._model.embed.call_args.kwargs["pooling_params"] is params + model._context_length = 8192 + model._dimensions = 128 + _, pool, _ = model._prepare_embedding( + {"audio": "x.wav"}, normalize_embeddings=False + ) + assert pool.dimensions == 128 + assert pool.use_activation is False diff --git a/xinference/model/embedding/vllm/core.py b/xinference/model/embedding/vllm/core.py index 66beefb793..4a81ded1de 100644 --- a/xinference/model/embedding/vllm/core.py +++ b/xinference/model/embedding/vllm/core.py @@ -24,6 +24,11 @@ from ...batch import BatchMixin from ...utils import check_dependency_available from ..core import EmbeddingModel, EmbeddingModelFamilyV2, EmbeddingSpecV1 +from ..embeddinggemma2 import MATRYOSHKA_DIMENSIONS +from ..embeddinggemma2 import MODEL_NAME as EMBEDDINGGEMMA2_MODEL_NAME +from ..embeddinggemma2 import load_prompts +from ..embeddinggemma2 import normalize_inputs as normalize_embeddinggemma2_inputs +from ..embeddinggemma2 import resolve_dtype, validate_dimensions from ..wemm import WeMMInput, is_wemm_model, iter_wemm_media, normalize_wemm_inputs logger = logging.getLogger(__name__) @@ -49,6 +54,8 @@ def __init__(self, *args, **kwargs): self._context_length = None self._chat_template = None self._native_pooling = False + self._prompts: Dict[str, str] = {} + self._dimensions = None def load(self): try: @@ -88,7 +95,36 @@ def load(self): is_matryoshka=True, ) - if is_wemm_model(self.model_family.model_name): + if self.model_family.model_name == EMBEDDINGGEMMA2_MODEL_NAME: + supported = self._check_embeddinggemma2_support() + if supported != True: + raise ImportError(supported[1]) + self._dimensions = self._kwargs.pop("dimensions", None) + validate_dimensions(self._dimensions) + dtype = self._kwargs.pop("torch_dtype", self._kwargs.get("dtype")) + self._kwargs["dtype"] = str( + resolve_dtype(dtype, self._device or "cuda") + ).split(".")[-1] + self._kwargs.setdefault("max_model_len", self.model_family.max_tokens) + overrides = self._kwargs.get("hf_overrides", {}) + if isinstance(overrides, str): + overrides = json.loads(overrides) + if not isinstance(overrides, dict): + raise ValueError( + "EmbeddingGemma 2 hf_overrides must be a dict or JSON object" + ) + self._kwargs["hf_overrides"] = { + **overrides, + "is_matryoshka": True, + "matryoshka_dimensions": list(MATRYOSHKA_DIMENSIONS), + } + self._prompts = load_prompts(self._model_path) + with open( + os.path.join(self._model_path, "chat_template.jinja"), encoding="utf-8" + ) as file: + self._chat_template = file.read() + self._model = LLM(model=self._model_path, runner="pooling", **self._kwargs) + elif is_wemm_model(self.model_family.model_name): if Version(vllm_version) < Version("0.27.0"): raise ValueError("WeMM-Embedding requires vLLM>=0.27.0") template_path = os.path.join( @@ -144,12 +180,22 @@ def _prepare_embedding( normalize_embedding = kwargs.get("normalize_embedding", True) dimensions = kwargs.get("dimensions", None) + if self.model_family.model_name == EMBEDDINGGEMMA2_MODEL_NAME: + normalize_embedding = kwargs.get( + "normalize_embeddings", normalize_embedding + ) + dimensions = kwargs.get("dimensions", self._dimensions) + validate_dimensions(dimensions) + if kwargs.get("return_sparse", False): + raise ValueError("EmbeddingGemma 2 does not support sparse embeddings.") assert self._model is not None # Check and truncate sentences that exceed context_length - if self._context_length is not None and not is_wemm_model( - self.model_family.model_name + if ( + self._context_length is not None + and not is_wemm_model(self.model_family.model_name) + and self.model_family.model_name != EMBEDDINGGEMMA2_MODEL_NAME ): truncated_sentences = [] for sentence in sentences if isinstance(sentences, list) else [sentences]: @@ -194,7 +240,9 @@ def _prepare_embedding( def _create_embedding(self, sentences: Any, **kwargs) -> Embedding: sentences, pool_params, model_uid = self._prepare_embedding(sentences, **kwargs) - if is_wemm_model(self.model_family.model_name): + if self.model_family.model_name == EMBEDDINGGEMMA2_MODEL_NAME: + outputs = self._embed_embeddinggemma2(sentences, pool_params, **kwargs) + elif is_wemm_model(self.model_family.model_name): outputs = self._embed_wemm(sentences, pool_params) elif self.model_family.model_name.startswith("Qwen3-VL-Embedding"): outputs = self._embed_vl(sentences) @@ -206,6 +254,62 @@ def _create_embedding(self, sentences: Any, **kwargs) -> Embedding: self._clean_cache_if_needed(result["usage"]["total_tokens"]) return result + @staticmethod + def _check_embeddinggemma2_support() -> Union[bool, Tuple[bool, str]]: + try: + from vllm import ModelRegistry + except ImportError: + return ( + False, + "EmbeddingGemma 2 requires vLLM>=0.32.0 with EmbeddingGemma2Model support.", + ) + if "EmbeddingGemma2Model" not in ModelRegistry.get_supported_archs(): + return ( + False, + "EmbeddingGemma 2 requires vLLM>=0.32.0 with EmbeddingGemma2Model support.", + ) + return True + + def _embed_embeddinggemma2(self, inputs: Any, pool_params: Any, **kwargs: Any): + from vllm.multimodal.utils import fetch_audio, fetch_image, fetch_video + + messages_batch = normalize_embeddinggemma2_inputs( + inputs, self._prompts, **kwargs + ) + loaders = {"image": fetch_image, "video": fetch_video, "audio": fetch_audio} + vllm_inputs = [] + for messages in messages_batch: + prompt = self._tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=False, + chat_template=self._chat_template, + ) + media: Dict[str, List[Any]] = {} + for part in messages[0]["content"]: + modality = part["type"] + if modality == "text": + continue + value = part[modality] + if isinstance(value, str): + if not value.startswith( + ("http://", "https://", "data:", "file://") + ): + value = f"file://{os.path.abspath(value)}" + value = loaders[modality](value) + media.setdefault(modality, []).append(value) + item: Dict[str, Any] = {"prompt": prompt} + if media: + item["multi_modal_data"] = media + if kwargs.get("processing_kwargs"): + item["mm_processor_kwargs"] = kwargs["processing_kwargs"] + vllm_inputs.append(item) + if not vllm_inputs: + return [] + return self._model.embed( + vllm_inputs, use_tqdm=False, pooling_params=pool_params + ) + async def _async_create_embedding(self, sentences: Any, **kwargs) -> Embedding: from vllm.outputs import EmbeddingRequestOutput @@ -446,6 +550,14 @@ def match_json( model_spec: EmbeddingSpecV1, quantization: str, ) -> Union[bool, Tuple[bool, str]]: + if model_family.model_name == EMBEDDINGGEMMA2_MODEL_NAME: + from ...utils import virtual_env_allows_missing_engine + + if model_spec.model_format != "pytorch": + return False, "EmbeddingGemma 2 vLLM supports pytorch format only" + if virtual_env_allows_missing_engine(): + return True + return cls._check_embeddinggemma2_support() required_vllm_version = None if is_wemm_model(model_family.model_name): required_vllm_version = "0.27.0" From 2687d3805a2b2748de8f4b237bea46063980ce75 Mon Sep 17 00:00:00 2001 From: Minamiyama Date: Wed, 7 Oct 2026 14:25:25 +0800 Subject: [PATCH 2/2] fix(embedding): correct EmbeddingGemma 2 types --- xinference/model/embedding/embeddinggemma2.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/xinference/model/embedding/embeddinggemma2.py b/xinference/model/embedding/embeddinggemma2.py index d5b39f7c7a..5258297587 100644 --- a/xinference/model/embedding/embeddinggemma2.py +++ b/xinference/model/embedding/embeddinggemma2.py @@ -143,6 +143,7 @@ def normalize_inputs( class TransformersEmbeddingGemma2Model(EmbeddingModel, BatchMixin): engine = "transformers" _device: Optional[str] + _model: Any def __init__(self, *args, **kwargs) -> None: EmbeddingModel.__init__(self, *args, **kwargs) @@ -191,7 +192,7 @@ def load(self) -> None: self._prompts = load_prompts(self._model_path) def _encode(self, messages: List[List[Dict[str, Any]]], **kwargs: Any): - embeddings = [] + embeddings: List[torch.Tensor] = [] tokens = 0 assert self._model is not None # Media have variable token lengths. One sample at a time avoids padding