diff --git a/packages/sie_server/README.md b/packages/sie_server/README.md index da7e99bfe..f3d5c76ee 100644 --- a/packages/sie_server/README.md +++ b/packages/sie_server/README.md @@ -69,7 +69,8 @@ pip install sie-server ``` - **Transformers 5 bundle** (LightOnOCR, GLM-OCR, GLiGuard, GLiNER2.5 entity and Decide models, Privacy Filter, and the TopK-Embed-V1 - multi-vector models, and Perplexity decisions) — requires `transformers` 5.x, and is served with `-b transformers5`. The GLiNER2.5-Decide + multi-vector models, Perplexity decisions, and EmbeddingGemma 2) — requires `transformers` 5.x, and is served with `-b transformers5`. + `google/embeddinggemma-2` needs `EmbeddingGemma2Model`, which first shipped in `transformers` 5.19.0. The GLiNER2.5-Decide and entity models also need `gliner2` 2.x. `sie-server` itself asks for `gliner2<2`, which the default bundle's GLiNER2 models need, so pip reports that conflict when the second command below installs 2.x; the transformers5 bundle's GLiNER2 models are pinned to 2.0.0: diff --git a/packages/sie_server/bundles/transformers5.yaml b/packages/sie_server/bundles/transformers5.yaml index 05de663ae..4a9536d5c 100644 --- a/packages/sie_server/bundles/transformers5.yaml +++ b/packages/sie_server/bundles/transformers5.yaml @@ -21,6 +21,9 @@ adapters: - sie_server.adapters.pplx_decider.runtime_lora # SAM 3 open-vocabulary detection (Sam3Model / Sam3Processor exist only in transformers 5.x). - sie_server.adapters.sam3.adapter +# EmbeddingGemma 2 text embeddings (EmbeddingGemma2Model). The class first +# shipped in transformers 5.19.0; see models/google__embeddinggemma-2.yaml. +- sie_server.adapters.embedding_gemma2.adapter deps: # Transformers 5.x for models requiring newer architecture support # (e.g., LightOnOCR-2-1B needs LightOnOcrForConditionalGeneration, @@ -28,8 +31,9 @@ deps: # for the visual-SPLADE family (custom modeling code built on 5.13+), # then 5.17 for Perplexity's pinned native Qwen decision checkpoint, then # 5.19: Parakeet-TDT refuses to load below 5.18 (padded rows turn into NaN - # under eager attention), and 5.19.0 is the version the v0.10.x images - # already ship. The saved locks pin this floor exactly. + # under eager attention), EmbeddingGemma2Model first shipped in 5.19.0, and + # 5.19.0 is the version the v0.10.x images already ship. The saved locks + # pin this floor exactly. transformers: '>=5.19,<6' # Includes multimodal SparseEncoder (image documents), with the local-code # trust boundary validated against this exact release. diff --git a/packages/sie_server/models/google__embeddinggemma-2.yaml b/packages/sie_server/models/google__embeddinggemma-2.yaml new file mode 100644 index 000000000..029446ef5 --- /dev/null +++ b/packages/sie_server/models/google__embeddinggemma-2.yaml @@ -0,0 +1,58 @@ +# google/embeddinggemma-2 — Apache-2.0, ungated, 744,371,512 parameters. +# Released 2026-09-14. Native transformers class EmbeddingGemma2Model +# (model_type embedding_gemma2); the repo has no custom modeling code. +# +# The class first shipped in transformers 5.19.0 (2026-10-06). 5.18.x does +# not contain it. This profile is on the transformers5 bundle only, whose +# floor is >=5.19,<6 and whose saved locks pin 5.19.0. The root pin stays +# transformers 4.57.6. Do not serve this model from the default bundle. +# +# Sentence Transformers >=6.1.0 is the upstream wrapper. The bundle pins +# sentence-transformers==5.7.0, so the adapter applies the text recipe itself: +# mean pooling with the prompt included, then L2 normalization in float32. +# Query prompt: "task: search result | query: " +# Document prompt: "title: none | text: " +# Matryoshka prefixes, re-normalized after the slice: 128, 256, 512, and the +# native 768. Pass options.truncate_dim to request a shorter prefix. The +# catalog width stays 768. +# +# Text only. Image, audio, and video need EmbeddingGemma2Processor's soft +# tokens, video-token sampling, and Gemma 4 audio features. SIE's image, frame, +# and speech-PCM plumbing does not produce those inputs, and this profile +# drops both towers at load (vision_config=None, audio_config=None). +# +# Float16 overflows this model (NaN activations). The profile asks for +# bfloat16; off CUDA the adapter uses float32. +sie_id: google/embeddinggemma-2 +hf_id: google/embeddinggemma-2 +# Pinned to an immutable commit SHA. Resolved 2026-11-01 from the Hub +# model API (ungated, sha of the default revision). +hf_revision: f9972ac22e696a105c482202646979e33e4b5a1e +inputs: + text: true + image: false + audio: false + video: false +tasks: + encode: + dense: + dim: 768 + sparse: null + multivector: null + score: null + extract: null +# Model card: one shared 8,192-token context. The tokenizer's +# model_max_length is an unbounded sentinel, so the cap lives here. +max_sequence_length: 8192 +profiles: + default: + max_batch_tokens: 16384 + compute_precision: bfloat16 + adapter_path: sie_server.adapters.embedding_gemma2.adapter:EmbeddingGemma2Adapter + adapter_options: + loadtime: {} + runtime: + normalize: true + pooling: mean + query_template: 'task: search result | query: {text}' + doc_template: 'title: none | text: {text}' diff --git a/packages/sie_server/src/sie_server/adapters/embedding_gemma2/__init__.py b/packages/sie_server/src/sie_server/adapters/embedding_gemma2/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/packages/sie_server/src/sie_server/adapters/embedding_gemma2/adapter.py b/packages/sie_server/src/sie_server/adapters/embedding_gemma2/adapter.py new file mode 100644 index 000000000..92060648b --- /dev/null +++ b/packages/sie_server/src/sie_server/adapters/embedding_gemma2/adapter.py @@ -0,0 +1,340 @@ +"""Text embeddings for google/embeddinggemma-2 on the native transformers class. + +EmbeddingGemma 2 is an omni embedding model (text, image, video, audio) whose +checkpoint uses ``EmbeddingGemma2Model`` with no custom code. That class first +shipped in transformers 5.19.0. Sentence Transformers 6.1 applies the official +prompt, pooling, and normalization recipe, but the transformers5 bundle pins +sentence-transformers 5.7.0, so this adapter reproduces the text profile itself: + +* task prompts from ``config_sentence_transformers.json``, prepended and kept + inside the pool (``1_Pooling`` sets ``include_prompt: true``); +* mask-aware mean pooling of ``last_hidden_state``; +* L2 normalization in float32, after any Matryoshka slice. + +The text-only load sets ``vision_config`` and ``audio_config`` to ``None``, +which the model card documents as the ~270M deployment of the 744M checkpoint. +Image, audio, and video are refused: SIE's media plumbing does not produce the +processor inputs those towers consume (see :func:`_reject_non_text`). +""" + +from __future__ import annotations + +import logging +import os +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from packaging.version import Version +from torch.nn import functional as F + +from sie_server.adapters._base_adapter import BaseAdapter +from sie_server.adapters._spec import AdapterSpec +from sie_server.adapters._transformers_compat import install_transformers_security_fixes +from sie_server.adapters._types import ERR_NOT_LOADED, ERR_REQUIRES_TEXT, ComputePrecision +from sie_server.core.inference_output import EncodeOutput +from sie_server.types.inputs import Item + +logger = logging.getLogger(__name__) + +# First transformers release that defines EmbeddingGemma2Model / embedding_gemma2. +# 5.18.0 and 5.18.1 do not; 5.19.0 (2026-10-06) does. +MIN_TRANSFORMERS = "5.19.0" + +NATIVE_DIM = 768 +# Card and config_sentence_transformers.json. 768 is the untruncated head. +MATRYOSHKA_DIMS = (128, 256, 512, 768) +# Shared context window from the model card. The tokenizer's model_max_length +# is an unbounded sentinel, so the catalog cap is what truncation uses. +CONTEXT_LENGTH = 8192 + +# Trailing space is part of the published prompt. Prompts are included in the +# mean pool; they are not a prefix that pooling strips. +QUERY_PROMPT = "task: search result | query: " +DOCUMENT_PROMPT = "title: none | text: " +QUERY_TEMPLATE = QUERY_PROMPT + "{text}" +DOCUMENT_TEMPLATE = DOCUMENT_PROMPT + "{text}" + +_MODALITY_GAP = ( + "Image, audio, and video are not wired. SIE decodes images to PIL, video to " + "sampled frames, and audio to speech PCM. EmbeddingGemma2Processor instead " + "owns image soft-tokens (max_soft_tokens), video frames as video tokens " + "(fps, timestamps, pixel_values_videos), and Gemma4 audio features. Those " + "are different inputs, and this profile does not load the vision or audio towers." +) + + +class EmbeddingGemma2Adapter(BaseAdapter): + """Dense text encoder for ``google/embeddinggemma-2``. + + Queries use the retrieval query prompt and documents use the document + prompt. Both are included in mean pooling. Callers can pass + ``options.truncate_dim`` of 128, 256, 512, or 768; the slice is + L2-normalized again when normalization is on. Float16 is not used: the + card reports NaN activations in that dtype. + """ + + spec = AdapterSpec( + inputs=("text",), + outputs=("dense",), + dense_dim=NATIVE_DIM, + unload_fields=("_model", "_tokenizer", "_dense_dim"), + ) + + def __init__( + self, + model_name_or_path: str | Path, + *, + normalize: bool = True, + max_seq_length: int | None = None, + compute_precision: ComputePrecision | None = None, + trust_remote_code: bool = False, + revision: str | None = None, + query_template: str | None = QUERY_TEMPLATE, + doc_template: str | None = DOCUMENT_TEMPLATE, + pooling: str = "mean", + dense_dim: int | None = None, + **kwargs: Any, + ) -> None: + if trust_remote_code: + msg = "EmbeddingGemma2Adapter does not execute Hub code; trust_remote_code must be false" + raise ValueError(msg) + if pooling != "mean": + msg = f"EmbeddingGemma 2 pools with mean, not {pooling!r}" + raise ValueError(msg) + _ = kwargs + self._model_name_or_path = str(model_name_or_path) + self._normalize = normalize + self._max_seq_length = max_seq_length or CONTEXT_LENGTH + self._compute_precision = compute_precision + self._revision = revision + self._query_template = query_template + self._doc_template = doc_template + self._configured_dense_dim = dense_dim + + self._model: Any = None + self._tokenizer: Any = None + self._device: str | None = None + self._dense_dim: int | None = dense_dim or NATIVE_DIM + + def load(self, device: str) -> None: + import transformers + + if Version(transformers.__version__) < Version(MIN_TRANSFORMERS): + msg = ( + f"EmbeddingGemma2Model requires transformers>={MIN_TRANSFORMERS}, " + f"found {transformers.__version__}. The class first shipped in transformers " + "5.19.0. This adapter is on the transformers5 bundle; transformers 4.57.6 " + "cannot load it." + ) + raise RuntimeError(msg) + + from transformers import AutoConfig, AutoTokenizer, EmbeddingGemma2Model # ty: ignore[unresolved-import] + + self._device = device + dtype = self._resolve_dtype() + shared_kwargs: dict[str, Any] = { + "trust_remote_code": False, + "token": os.environ.get("HF_TOKEN"), + } + if self._revision is not None: + shared_kwargs["revision"] = self._revision + + logger.info( + "Loading EmbeddingGemma 2 %s on device=%s dtype=%s revision=%s (text towers only)", + self._model_name_or_path, + device, + dtype, + self._revision, + ) + install_transformers_security_fixes() + config = AutoConfig.from_pretrained(self._model_name_or_path, **shared_kwargs) + # Card: unused towers are skipped and their checkpoint keys ignored. + config.vision_config = None + config.audio_config = None + + self._tokenizer = AutoTokenizer.from_pretrained( + self._model_name_or_path, + padding_side="right", + **shared_kwargs, + ) + self._model = EmbeddingGemma2Model.from_pretrained( + self._model_name_or_path, + config=config, + torch_dtype=dtype, + **shared_kwargs, + ) + self._model.to(device) + self._model.eval() + + text_config = self._model.config.text_config + observed = int(getattr(text_config, "embedding_dim", NATIVE_DIM)) + self._dense_dim = self._validate_dense_dim(observed) + + def _validate_dense_dim(self, observed_dim: int) -> int: + if self._configured_dense_dim is not None and observed_dim != self._configured_dense_dim: + msg = ( + "EmbeddingGemma 2 dimension mismatch: " + f"configured dense_dim={self._configured_dense_dim}, embedding_dim={observed_dim}" + ) + raise ValueError(msg) + return observed_dim + + def _resolve_dtype(self) -> torch.dtype: + """bfloat16 on CUDA, float32 elsewhere. Never float16.""" + precision = self._compute_precision or "float32" + cuda = bool(self._device and str(self._device).startswith("cuda")) + if precision == "float16": + logger.warning( + "EmbeddingGemma 2 overflows float16. Using %s instead.", + "bfloat16" if cuda else "float32", + ) + precision = "bfloat16" if cuda else "float32" + if not cuda: + return torch.float32 + if precision == "bfloat16": + return torch.bfloat16 + return torch.float32 + + def encode( + self, + items: list[Item], + output_types: list[str], + *, + instruction: str | None = None, + is_query: bool = False, + prepared_items: Any = None, + options: dict[str, Any] | None = None, + ) -> EncodeOutput: + _ = instruction, prepared_items + self._check_loaded() + if self._model is None or self._tokenizer is None: + raise RuntimeError(ERR_NOT_LOADED) + + unsupported = set(output_types) - {"dense"} + if unsupported: + msg = f"Unsupported output types: {unsupported}. This model only supports 'dense'." + raise ValueError(msg) + + opts = options or {} + pooling = opts.get("pooling", "mean") + if pooling != "mean": + msg = f"EmbeddingGemma 2 pools with mean, not {pooling!r}" + raise ValueError(msg) + normalize = bool(opts.get("normalize", self._normalize)) + truncate_dim = _resolve_truncate_dim(opts.get("truncate_dim", None), self._dense_dim or NATIVE_DIM) + query_template = _option_template(opts, "query_template", self._query_template) + doc_template = _option_template(opts, "doc_template", self._doc_template) + template = query_template if is_query else doc_template + + for item in items: + _reject_non_text(item) + + texts = [_format_text(template, _require_text(item)) for item in items] + if not texts: + width = truncate_dim or self._dense_dim or NATIVE_DIM + return EncodeOutput( + dense=np.zeros((0, width), dtype=np.float32), + batch_size=0, + is_query=is_query, + dense_dim=width, + ) + + with self._tokenizer_guard(), torch.inference_mode(): + encoded = self._tokenizer( + texts, + max_length=self._max_seq_length, + padding=True, + truncation=True, + return_tensors="pt", + ) + input_ids = encoded["input_ids"].to(self._device) + attention_mask = encoded["attention_mask"].to(self._device) + outputs = self._model( + input_ids=input_ids, + attention_mask=attention_mask, + return_dict=True, + ) + pooled = _mean_pool(outputs.last_hidden_state, attention_mask) + if truncate_dim is not None: + pooled = pooled[..., :truncate_dim] + if normalize: + pooled = F.normalize(pooled.float(), p=2, dim=-1) + else: + pooled = pooled.float() + token_counts = [int(count) for count in attention_mask.sum(dim=1).tolist()] + + embeddings = pooled.cpu().numpy() + width = int(embeddings.shape[1]) + return EncodeOutput( + dense=embeddings, + batch_size=len(items), + is_query=is_query, + dense_dim=width, + extra={"input_token_counts": token_counts}, + ) + + +def _require_text(item: Item) -> str: + if item.text is None: + raise ValueError(ERR_REQUIRES_TEXT.format(adapter_name="EmbeddingGemma2Adapter")) + return item.text + + +def _reject_non_text(item: Item) -> None: + present: list[str] = [] + if item.images: + present.append("image") + if item.audio is not None: + present.append("audio") + if item.video is not None: + present.append("video") + if item.document is not None: + present.append("document") + if present: + msg = f"EmbeddingGemma2Adapter encodes text only in this release; got {', '.join(present)}. {_MODALITY_GAP}" + raise ValueError(msg) + + +def _format_text(template: str | None, text: str) -> str: + if not template: + return text + if "{text}" not in template: + msg = f"EmbeddingGemma 2 prompt template must contain '{{text}}', got {template!r}" + raise ValueError(msg) + return template.format(text=text) + + +def _option_template(options: dict[str, Any], key: str, default: str | None) -> str | None: + if key not in options: + return default + value = options[key] + if value is None: + return None + if not isinstance(value, str): + msg = f"options.{key} must be a string or null, got {type(value).__name__}" + raise ValueError(msg) + return value + + +def _resolve_truncate_dim(value: object, native_dim: int) -> int | None: + if value is None: + return None + if isinstance(value, bool) or not isinstance(value, int): + msg = f"truncate_dim must be an int or null, got {value!r}" + raise ValueError(msg) + if value == native_dim: + return None + if value not in MATRYOSHKA_DIMS or value > native_dim: + msg = f"truncate_dim must be one of {MATRYOSHKA_DIMS}, got {value}" + raise ValueError(msg) + return value + + +def _mean_pool(last_hidden: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: + """Mean of non-padding tokens, in float32, matching the card's recipe.""" + mask = attention_mask.to(dtype=torch.float32).unsqueeze(-1) + summed = (last_hidden.float() * mask).sum(dim=1) + counts = mask.sum(dim=1).clamp(min=1e-9) + return summed / counts diff --git a/packages/sie_server/tests/adapters/test_embedding_gemma2.py b/packages/sie_server/tests/adapters/test_embedding_gemma2.py new file mode 100644 index 000000000..74cca673e --- /dev/null +++ b/packages/sie_server/tests/adapters/test_embedding_gemma2.py @@ -0,0 +1,242 @@ +"""Text recipe for EmbeddingGemma2Adapter: prompts, mean pool, Matryoshka.""" + +from __future__ import annotations + +import re +from types import SimpleNamespace +from typing import Any +from unittest.mock import MagicMock + +import pytest +import torch +from sie_server.adapters._types import ERR_NOT_LOADED +from sie_server.adapters.embedding_gemma2.adapter import ( + DOCUMENT_TEMPLATE, + QUERY_TEMPLATE, + EmbeddingGemma2Adapter, + _mean_pool, +) +from sie_server.types.inputs import AudioInput, Item + + +def _adapter(**kwargs: Any) -> EmbeddingGemma2Adapter: + dense_dim = int(kwargs.pop("dense_dim", 4)) + adapter = EmbeddingGemma2Adapter("google/embeddinggemma-2", revision="abc", dense_dim=dense_dim, **kwargs) + adapter._device = "cpu" + adapter._model = object() + adapter._tokenizer = _FakeTokenizer() + adapter._dense_dim = dense_dim + return adapter + + +class _FakeTokenizer: + def __init__(self) -> None: + self.texts: list[str] | None = None + self.kwargs: dict[str, Any] | None = None + + def __call__(self, texts: list[str], **kwargs: Any) -> dict[str, torch.Tensor]: + self.texts = list(texts) + self.kwargs = kwargs + # One content token per word, plus a padding column on the shorter row. + lengths = [max(1, len(text.split())) for text in texts] + width = max(lengths) + input_ids = [] + mask = [] + for length in lengths: + input_ids.append([1] * length + [0] * (width - length)) + mask.append([1] * length + [0] * (width - length)) + return { + "input_ids": torch.tensor(input_ids), + "attention_mask": torch.tensor(mask), + } + + +class _FakeModel: + def __init__(self, last_hidden: torch.Tensor) -> None: + self._last_hidden = last_hidden + self.calls: list[dict[str, Any]] = [] + + def __call__(self, **kwargs: Any) -> SimpleNamespace: + self.calls.append(kwargs) + batch, seq = kwargs["input_ids"].shape + hidden = self._last_hidden + if hidden.shape[0] != batch or hidden.shape[1] != seq: + hidden = hidden[:, :1].expand(batch, seq, hidden.shape[-1]).clone() + return SimpleNamespace(last_hidden_state=hidden) + + +def test_mean_pool_skips_padding() -> None: + hidden = torch.tensor( + [ + [[2.0, 0.0], [6.0, 0.0], [100.0, 0.0]], + [[0.0, 3.0], [0.0, 3.0], [0.0, 3.0]], + ] + ) + mask = torch.tensor([[1, 1, 0], [1, 1, 1]]) + pooled = _mean_pool(hidden, mask) + assert torch.allclose(pooled[0], torch.tensor([4.0, 0.0])) + assert torch.allclose(pooled[1], torch.tensor([0.0, 3.0])) + + +def test_query_and_document_prompts_are_included_in_the_pool() -> None: + adapter = _adapter() + hidden = torch.tensor([[[3.0, 4.0]]]) + adapter._model = _FakeModel(hidden) # ty: ignore[invalid-assignment] + + query = adapter.encode([Item(text="mars")], ["dense"], is_query=True) + document = adapter.encode([Item(text="mars")], ["dense"], is_query=False) + + assert adapter._tokenizer.texts == [DOCUMENT_TEMPLATE.format(text="mars")] # ty: ignore[unresolved-attribute] + adapter.encode([Item(text="mars")], ["dense"], is_query=True) + assert adapter._tokenizer.texts == [QUERY_TEMPLATE.format(text="mars")] # ty: ignore[unresolved-attribute] + # [3, 4] normalized is [0.6, 0.8]. Prompt tokens are part of the string the + # tokenizer saw, so they are inside the pooled span. + assert query.dense is not None + assert document.dense is not None + assert pytest.approx(query.dense[0].tolist(), abs=1e-5) == [0.6, 0.8] + assert pytest.approx(document.dense[0].tolist(), abs=1e-5) == [0.6, 0.8] + assert query.extra["input_token_counts"] == [len(QUERY_TEMPLATE.format(text="mars").split())] + + +def test_empty_template_disables_the_task_prompt() -> None: + adapter = _adapter() + adapter._model = _FakeModel(torch.tensor([[[1.0, 0.0]]])) # ty: ignore[invalid-assignment] + adapter.encode( + [Item(text="mars")], + ["dense"], + is_query=True, + options={"query_template": None}, + ) + assert adapter._tokenizer.texts == ["mars"] # ty: ignore[unresolved-attribute] + + +def test_truncate_dim_slices_then_renormalizes() -> None: + adapter = _adapter(dense_dim=768) + adapter._dense_dim = 768 + hidden = torch.zeros(1, 1, 768) + hidden[0, 0, 0] = 3.0 + hidden[0, 0, 1] = 4.0 + adapter._model = _FakeModel(hidden) # ty: ignore[invalid-assignment] + output = adapter.encode([Item(text="mars")], ["dense"], options={"truncate_dim": 128}) + assert output.dense is not None + assert output.dense.shape == (1, 128) + assert output.dense_dim == 128 + # The leading [3, 4] pair is re-normalized after the slice; the rest is zero + # and does not change the direction of that pair. + assert pytest.approx(output.dense[0, 0].item(), abs=1e-5) == 0.6 + assert pytest.approx(output.dense[0, 1].item(), abs=1e-5) == 0.8 + + +def test_native_truncate_dim_keeps_the_full_vector() -> None: + adapter = _adapter(dense_dim=768) + adapter._dense_dim = 768 + adapter._model = _FakeModel(torch.ones(1, 1, 768)) # ty: ignore[invalid-assignment] + output = adapter.encode([Item(text="mars")], ["dense"], options={"truncate_dim": 768}) + assert output.dense is not None + assert output.dense.shape == (1, 768) + + +@pytest.mark.parametrize("value", [64, 384, 0, True, "256"]) +def test_truncate_dim_rejects_untrained_widths(value: object) -> None: + adapter = _adapter() + adapter._model = _FakeModel(torch.ones(1, 1, 4)) # ty: ignore[invalid-assignment] + with pytest.raises(ValueError, match="truncate_dim"): + adapter.encode([Item(text="mars")], ["dense"], options={"truncate_dim": value}) + + +def test_non_mean_pooling_is_rejected() -> None: + with pytest.raises(ValueError, match="mean"): + EmbeddingGemma2Adapter("google/embeddinggemma-2", pooling="last") + adapter = _adapter() + adapter._model = _FakeModel(torch.ones(1, 1, 4)) # ty: ignore[invalid-assignment] + with pytest.raises(ValueError, match="mean"): + adapter.encode([Item(text="mars")], ["dense"], options={"pooling": "cls"}) + + +def test_non_text_modalities_are_refused() -> None: + adapter = _adapter() + image = Item(text="a cat", images=[{"data": b"not-an-image"}]) + audio = Item(audio=AudioInput(data=b"abc", format="wav")) + video = Item(video={"data": b"mp4"}) + for item, modality in ((image, "image"), (audio, "audio"), (video, "video")): + with pytest.raises(ValueError, match=modality): + adapter.encode([item], ["dense"]) + + +def test_encode_requires_a_loaded_model() -> None: + adapter = EmbeddingGemma2Adapter("google/embeddinggemma-2") + with pytest.raises(RuntimeError, match=re.escape(ERR_NOT_LOADED)): + adapter.encode([Item(text="mars")], ["dense"]) + + +def test_trust_remote_code_is_refused() -> None: + with pytest.raises(ValueError, match="trust_remote_code"): + EmbeddingGemma2Adapter("google/embeddinggemma-2", trust_remote_code=True) + + +def test_load_refuses_transformers_before_5_19(monkeypatch: pytest.MonkeyPatch) -> None: + import transformers + + monkeypatch.setattr(transformers, "__version__", "5.17.0", raising=False) + adapter = EmbeddingGemma2Adapter("google/embeddinggemma-2", revision="f" * 40) + with pytest.raises(RuntimeError, match=r"5\.19\.0"): + adapter.load("cpu") + + +def test_load_drops_modality_towers_and_pins_revision(monkeypatch: pytest.MonkeyPatch) -> None: + import transformers + + monkeypatch.setattr(transformers, "__version__", "5.19.0", raising=False) + config = SimpleNamespace(vision_config={"keep": False}, audio_config={"keep": False}, text_config=None) + model = MagicMock() + model.config.text_config.embedding_dim = 768 + auto_config = MagicMock() + auto_config.from_pretrained.return_value = config + auto_tokenizer = MagicMock() + embedding_cls = MagicMock() + embedding_cls.from_pretrained.return_value = model + monkeypatch.setattr(transformers, "AutoConfig", auto_config, raising=False) + monkeypatch.setattr(transformers, "AutoTokenizer", auto_tokenizer, raising=False) + monkeypatch.setattr(transformers, "EmbeddingGemma2Model", embedding_cls, raising=False) + + revision = "f9972ac22e696a105c482202646979e33e4b5a1e" + adapter = EmbeddingGemma2Adapter( + "google/embeddinggemma-2", + revision=revision, + dense_dim=768, + compute_precision="float16", + ) + adapter.load("cpu") + + assert config.vision_config is None + assert config.audio_config is None + assert adapter._dense_dim == 768 + for loader in (auto_config, auto_tokenizer, embedding_cls): + kwargs = loader.from_pretrained.call_args.kwargs + assert kwargs["trust_remote_code"] is False + assert kwargs["revision"] == revision + model_kwargs = embedding_cls.from_pretrained.call_args.kwargs + assert model_kwargs["config"] is config + assert model_kwargs["torch_dtype"] == torch.float32 + model.to.assert_called_once_with("cpu") + model.eval.assert_called_once() + + +def test_float16_on_cuda_upgrades_to_bfloat16(monkeypatch: pytest.MonkeyPatch) -> None: + import transformers + + monkeypatch.setattr(transformers, "__version__", "5.22.1", raising=False) + config = SimpleNamespace(vision_config=object(), audio_config=object()) + model = MagicMock() + model.config.text_config.embedding_dim = 768 + auto_config = MagicMock() + auto_config.from_pretrained.return_value = config + monkeypatch.setattr(transformers, "AutoConfig", auto_config, raising=False) + monkeypatch.setattr(transformers, "AutoTokenizer", MagicMock(), raising=False) + embedding_cls = MagicMock() + embedding_cls.from_pretrained.return_value = model + monkeypatch.setattr(transformers, "EmbeddingGemma2Model", embedding_cls, raising=False) + + adapter = EmbeddingGemma2Adapter("google/embeddinggemma-2", compute_precision="float16", dense_dim=768) + adapter.load("cuda:0") + assert embedding_cls.from_pretrained.call_args.kwargs["torch_dtype"] == torch.bfloat16 diff --git a/packages/sie_server/tests/config/test_embedding_gemma2.py b/packages/sie_server/tests/config/test_embedding_gemma2.py new file mode 100644 index 000000000..297f40ae8 --- /dev/null +++ b/packages/sie_server/tests/config/test_embedding_gemma2.py @@ -0,0 +1,80 @@ +"""Catalog identity and transformers5 routing for google/embeddinggemma-2.""" + +from pathlib import Path + +import yaml +from packaging.specifiers import SpecifierSet +from packaging.version import Version +from sie_sdk.bundle_utils import match_bundle_models +from sie_server.adapters.embedding_gemma2.adapter import ( + DOCUMENT_TEMPLATE, + MATRYOSHKA_DIMS, + MIN_TRANSFORMERS, + NATIVE_DIM, + QUERY_TEMPLATE, +) +from sie_server.core.loader import load_model_configs, resolve_adapter_class + +SERVER = Path(__file__).resolve().parents[2] +MODELS = SERVER / "models" +BUNDLES = SERVER / "bundles" +MODEL_ID = "google/embeddinggemma-2" +REVISION = "f9972ac22e696a105c482202646979e33e4b5a1e" + + +def test_embeddinggemma2_descriptor_is_pinned_text_encode() -> None: + config = load_model_configs(MODELS)[MODEL_ID] + assert config.sie_id == config.hf_id == MODEL_ID + assert config.hf_revision == REVISION + assert config.inputs.text + assert not config.inputs.image + assert not config.inputs.audio + assert not config.inputs.video + assert config.max_sequence_length == 8192 + assert config.tasks.encode is not None + assert config.tasks.encode.dense is not None + assert config.tasks.encode.dense.dim == NATIVE_DIM + assert config.tasks.score is None + assert config.tasks.extract is None + profile = config.resolve_profile("default") + assert profile.compute_precision == "bfloat16" + assert profile.max_batch_tokens == 16384 + assert profile.adapter_path == "sie_server.adapters.embedding_gemma2.adapter:EmbeddingGemma2Adapter" + assert resolve_adapter_class(config, MODELS).__name__ == "EmbeddingGemma2Adapter" + assert profile.loadtime == {} + assert dict(profile.runtime) == { + "normalize": True, + "pooling": "mean", + "query_template": QUERY_TEMPLATE, + "doc_template": DOCUMENT_TEMPLATE, + } + + +def test_embeddinggemma2_prompts_and_matryoshka_match_the_card() -> None: + assert QUERY_TEMPLATE == "task: search result | query: {text}" + assert DOCUMENT_TEMPLATE == "title: none | text: {text}" + assert MATRYOSHKA_DIMS == (128, 256, 512, 768) + assert MIN_TRANSFORMERS == "5.19.0" + assert Version("5.18.1") < Version(MIN_TRANSFORMERS) + assert Version("5.19.0") >= Version(MIN_TRANSFORMERS) + + +def test_embeddinggemma2_is_served_only_from_the_transformers5_bundle() -> None: + bundle = yaml.safe_load((BUNDLES / "transformers5.yaml").read_text()) + assert "sie_server.adapters.embedding_gemma2.adapter" in bundle["adapters"] + # Every transformers release the bundle admits ships EmbeddingGemma2Model. + spec = SpecifierSet(bundle["deps"]["transformers"]) + assert Version("5.19.0") in spec + assert Version("5.18.1") not in spec + assert Version("4.57.6") not in spec + assert MODEL_ID in match_bundle_models(BUNDLES / "transformers5.yaml", MODELS) + + default = yaml.safe_load((BUNDLES / "default.yaml").read_text()) + assert "sie_server.adapters.embedding_gemma2.adapter" not in default["adapters"] + assert default["deps"]["transformers"] == "==4.57.6" + assert MODEL_ID not in match_bundle_models(BUNDLES / "default.yaml", MODELS) + + for bundle_path in sorted(BUNDLES.glob("*.yaml")): + if bundle_path.name == "transformers5.yaml": + continue + assert MODEL_ID not in match_bundle_models(bundle_path, MODELS)