From 8fea27e59177f6c57778f8115816bccdc1db7ac7 Mon Sep 17 00:00:00 2001 From: IMvision12 Date: Tue, 13 Oct 2026 16:19:33 -0700 Subject: [PATCH 01/10] Add Initial Sentence Transformers support --- docs/models.md | 1 + docs/mpnet.md | 3 +- docs/sentence_transformer.md | 88 +++ tests/base/model_test_registry.py | 49 ++ .../integration/test_sentence_transformer.py | 139 +++++ website/mkdocs.yml | 1 + zeromodels/auto/auto_mapping_names.py | 4 + zeromodels/conversion/hf_download_utils.py | 22 +- zeromodels/models/__init__.py | 1 + .../models/sentence_transformer/__init__.py | 13 + ...onvert_sentence_transformer_hf_to_keras.py | 289 +++++++++ .../sentence_transformer_backbones.py | 225 +++++++ .../sentence_transformer_config.py | 41 ++ .../sentence_transformer_layers.py | 98 +++ .../sentence_transformer_model.py | 566 ++++++++++++++++++ .../sentence_transformer_tokenizer.py | 182 ++++++ 16 files changed, 1714 insertions(+), 8 deletions(-) create mode 100644 docs/sentence_transformer.md create mode 100644 tests/integration/test_sentence_transformer.py create mode 100644 zeromodels/models/sentence_transformer/__init__.py create mode 100644 zeromodels/models/sentence_transformer/convert_sentence_transformer_hf_to_keras.py create mode 100644 zeromodels/models/sentence_transformer/sentence_transformer_backbones.py create mode 100644 zeromodels/models/sentence_transformer/sentence_transformer_config.py create mode 100644 zeromodels/models/sentence_transformer/sentence_transformer_layers.py create mode 100644 zeromodels/models/sentence_transformer/sentence_transformer_model.py create mode 100644 zeromodels/models/sentence_transformer/sentence_transformer_tokenizer.py diff --git a/docs/models.md b/docs/models.md index 342ff0f5..978541be 100644 --- a/docs/models.md +++ b/docs/models.md @@ -37,6 +37,7 @@ Encoders for embeddings and masked language modelling, and decoder LLMs, dense a - [Qwen3.5](qwen3_5.md) - [Qwen3-Next](qwen3_next.md) - [RoBERTa](roberta.md) +- [Sentence Transformer](sentence_transformer.md) - [XLM-RoBERTa](xlm_roberta.md) ## Vision models diff --git a/docs/mpnet.md b/docs/mpnet.md index 1070b6e9..6786c2e2 100644 --- a/docs/mpnet.md +++ b/docs/mpnet.md @@ -33,7 +33,8 @@ See also [bert.md](bert.md), [roberta.md](roberta.md), [modernbert.md](modernber [electra.md](electra.md). MPNet is also the backbone of `sentence-transformers/all-mpnet-base-v2`, one of the most -widely used sentence-embedding models. +widely used sentence-embedding models — see [sentence_transformer.md](sentence_transformer.md) +for the Keras port of that full ST stack (mean pool + L2 normalize). ## Variants diff --git a/docs/sentence_transformer.md b/docs/sentence_transformer.md new file mode 100644 index 00000000..2f9da7b6 --- /dev/null +++ b/docs/sentence_transformer.md @@ -0,0 +1,88 @@ +# Sentence Transformer + +ZeroModels loads the 13 original dense sentence-embedding checkpoints shown in +the [Sentence Transformers pretrained-model table](https://sbert.net/docs/sentence_transformer/pretrained_models.html). +It builds a Keras 3 encoder and applies the checkpoint's declared pooling, +optional Dense projection, and optional L2 normalization. Parity has been +verified on the Keras Torch backend; TensorFlow and JAX parity have not yet +been checked for all 13 checkpoints. + +These are **upstream** Hugging Face repos: use the `hf:` prefix below. No +conversion or upload is needed for upstream loading. + +## Supported upstream checkpoints + +| Checkpoint (`sentence-transformers/`) | Actual backbone | Embedding recipe | Dim | +|---|---|---|---:| +| `all-mpnet-base-v2` | MPNet | mean → normalize | 768 | +| `multi-qa-mpnet-base-dot-v1` | MPNet | CLS, dot-product similarity | 768 | +| `all-distilroberta-v1` | RoBERTa | mean → normalize | 768 | +| `all-MiniLM-L12-v2` | BERT | mean → normalize | 384 | +| `multi-qa-distilbert-cos-v1` | DistilBERT | mean → normalize | 768 | +| `all-MiniLM-L6-v2` | BERT | mean → normalize | 384 | +| `multi-qa-MiniLM-L6-cos-v1` | BERT | mean → normalize | 384 | +| `paraphrase-multilingual-mpnet-base-v2` | XLM-RoBERTa | mean | 768 | +| `paraphrase-albert-small-v2` | ALBERT | mean | 768 | +| `paraphrase-multilingual-MiniLM-L12-v2` | BERT | mean | 384 | +| `paraphrase-MiniLM-L3-v2` | BERT | mean | 384 | +| `distiluse-base-multilingual-cased-v1` | DistilBERT | mean → Tanh Dense | 512 | +| `distiluse-base-multilingual-cased-v2` | DistilBERT | mean → Tanh Dense | 512 | + +The name `paraphrase-multilingual-mpnet-base-v2` is historical: its actual +checkpoint `model_type` is `xlm-roberta`, and ZeroModels uses the XLM-R backbone. + +## Usage + +```python +import os + +os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" + +from zeromodels.models.sentence_transformer import ( + SentenceTransformerModel, + SentenceTransformerTokenizer, +) + +repo = "sentence-transformers/all-MiniLM-L6-v2" +model = SentenceTransformerModel.from_weights(f"hf:{repo}") +tokenizer = SentenceTransformerTokenizer.from_weights( + f"hf:{repo}", revision=model.source_revision +) + +embeddings = model.encode( + ["A man is eating food.", "A man is eating pasta."], + tokenizer=tokenizer, +) +print(embeddings.shape) # (2, 384) +print(model.similarity(embeddings[:1], embeddings[1:])) +``` + +`model(tokenizer(texts))` returns `embeddings` and `token_embeddings` tensors. +`encode()` batches raw strings and returns a NumPy array. Use the same +checkpoint and revision for the model and tokenizer. An explicit commit SHA can +be supplied as `revision=` to pin the upstream snapshot. `model.similarity()` +uses the checkpoint's declared cosine or dot-product scoring function. + +`AutoZMTextEmbed` resolves ZeroModels `sentence_transformer` configs to the +`SentenceTransformerTextEmbed` task view. For upstream `hf:` repos, load the +concrete class as shown above: their +root `model_type` describes the backbone, not the sentence-embedding task. + +Unsupported upstream module types, combinations, or pooling modes raise a +`ValueError` rather than silently changing the embedding recipe. Cross-encoders, +sparse encoders, and multi-vector encoders are separate tasks and are not +supported by this model. + +## Parity check + +The optional `sentence-transformers` dependency is needed only for comparison +against the original implementation: + +```bash +python scripts/verify_sentence_transformer_parity.py --model all +``` + +The script compares embeddings at pinned upstream commits for all 13 rows. +Use `--model ` to check one checkpoint. It downloads upstream weights +into the standard Hugging Face cache but does not create or publish converted +releases. diff --git a/tests/base/model_test_registry.py b/tests/base/model_test_registry.py index 66cf90d6..f20421ac 100644 --- a/tests/base/model_test_registry.py +++ b/tests/base/model_test_registry.py @@ -2996,6 +2996,55 @@ "input_factory_kwargs": {"num_choices": 3, "seq_len": 16}, "expected_output_shape": (2, 3), }, + # ---- Sentence Transformers (MPNet backbone) ---- + "SentenceTransformerModel": { + "module": "zeromodels.models.sentence_transformer", + "model_cls": "SentenceTransformerModel", + "model_type": "llm", + "init_kwargs": { + "backbone_type": "mpnet", + "vocab_size": 128, + "embed_dim": 32, + "num_layers": 2, + "num_heads": 2, + "mlp_dim": 64, + "max_position_embeddings": 64, + "relative_attention_num_buckets": 8, + "pooling": "mean", + "normalize_embeddings": True, + "max_seq_len": 32, + }, + "input_factory": "modernbert_input", + "input_factory_kwargs": {"seq_len": 16}, + "expected_output_shape": { + "embeddings": (2, 32), + "token_embeddings": (2, 16, 32), + }, + }, + "SentenceTransformerTextEmbed": { + "module": "zeromodels.models.sentence_transformer", + "model_cls": "SentenceTransformerTextEmbed", + "model_type": "llm", + "init_kwargs": { + "backbone_type": "mpnet", + "vocab_size": 128, + "embed_dim": 32, + "num_layers": 2, + "num_heads": 2, + "mlp_dim": 64, + "max_position_embeddings": 64, + "relative_attention_num_buckets": 8, + "pooling": "mean", + "normalize_embeddings": True, + "max_seq_len": 32, + }, + "input_factory": "modernbert_input", + "input_factory_kwargs": {"seq_len": 16}, + "expected_output_shape": { + "embeddings": (2, 32), + "token_embeddings": (2, 16, 32), + }, + }, # ---- Text encoders (XLM-RoBERTa) ---- "XLMRobertaModel": { "module": "zeromodels.models.xlm_roberta", diff --git a/tests/integration/test_sentence_transformer.py b/tests/integration/test_sentence_transformer.py new file mode 100644 index 00000000..f272f349 --- /dev/null +++ b/tests/integration/test_sentence_transformer.py @@ -0,0 +1,139 @@ +"""Offline contract tests for the multi-backbone sentence-embedding API.""" + +import json + +import numpy as np +import pytest +from keras import ops + +from zeromodels.models.sentence_transformer import ( + SentenceTransformerModel, + SentenceTransformerTextEmbed, +) + + +@pytest.mark.parametrize( + "backbone_type", + ("mpnet", "bert", "roberta", "xlm-roberta", "distilbert", "albert"), +) +def test_backbone_builds_and_returns_sentence_embeddings(backbone_type): + model = SentenceTransformerModel( + backbone_type=backbone_type, + vocab_size=32, + embed_dim=16, + embedding_size=8, + num_layers=1, + num_heads=2, + mlp_dim=32, + max_position_embeddings=32, + relative_attention_num_buckets=8, + ) + inputs = { + "input_ids": np.array([[2, 3, 4, 1]], dtype="int32"), + "attention_mask": np.array([[1, 1, 1, 0]], dtype="int32"), + } + if backbone_type in ("bert", "roberta", "xlm-roberta", "albert"): + inputs["token_type_ids"] = np.zeros((1, 4), dtype="int32") + actual = model(inputs, training=False) + assert actual["embeddings"].shape == (1, 16) + assert actual["token_embeddings"].shape == (1, 4, 16) + + +@pytest.mark.parametrize("pooling", ("mean", "cls", "max")) +def test_pooling_parser_accepts_exactly_one_supported_mode( + monkeypatch, tmp_path, pooling +): + flags = { + "pooling_mode_mean_tokens": pooling == "mean", + "pooling_mode_cls_token": pooling == "cls", + "pooling_mode_max_tokens": pooling == "max", + } + config_path = tmp_path / "config.json" + config_path.write_text(json.dumps(flags), encoding="utf-8") + monkeypatch.setattr( + "huggingface_hub.hf_hub_download", lambda *args, **kwargs: str(config_path) + ) + assert ( + SentenceTransformerModel._pooling_from_module_dir("example/model", "1_Pooling") + == pooling + ) + + +def test_pooling_parser_rejects_ambiguous_modes(monkeypatch, tmp_path): + config_path = tmp_path / "config.json" + config_path.write_text( + json.dumps({"pooling_mode_mean_tokens": True, "pooling_mode_cls_token": True}), + encoding="utf-8", + ) + monkeypatch.setattr( + "huggingface_hub.hf_hub_download", lambda *args, **kwargs: str(config_path) + ) + with pytest.raises(ValueError, match="pooling config"): + SentenceTransformerModel._pooling_from_module_dir("example/model", "1_Pooling") + + +def test_sentence_transformer_registered_for_text_embed_task(): + from zeromodels.auto import AutoZMTextEmbed + + assert ( + AutoZMTextEmbed.mapping()["sentence_transformer"] + is SentenceTransformerTextEmbed + ) + + +def test_dot_similarity_is_not_cosine_normalized(): + model = SentenceTransformerModel( + backbone_type="bert", + vocab_size=32, + embed_dim=8, + num_layers=1, + num_heads=2, + mlp_dim=16, + max_position_embeddings=16, + similarity_fn_name="dot", + ) + a = np.array([[2.0, 0.0]], dtype="float32") + b = np.array([[3.0, 0.0]], dtype="float32") + np.testing.assert_allclose(ops.convert_to_numpy(model.similarity(a, b)), [[6.0]]) + assert model.get_config()["similarity_fn_name"] == "dot" + + +def test_cosine_similarity_accepts_numpy_embeddings(): + model = SentenceTransformerModel( + backbone_type="bert", + vocab_size=32, + embed_dim=8, + num_layers=1, + num_heads=2, + mlp_dim=16, + max_position_embeddings=16, + ) + a = np.array([[2.0, 0.0]], dtype="float32") + b = np.array([[3.0, 0.0]], dtype="float32") + np.testing.assert_allclose(ops.convert_to_numpy(model.similarity(a, b)), [[1.0]]) + + +def test_sentence_transformer_h5_weights_round_trip(tmp_path): + kwargs = { + "backbone_type": "bert", + "vocab_size": 32, + "embed_dim": 8, + "num_layers": 1, + "num_heads": 2, + "mlp_dim": 16, + "max_position_embeddings": 16, + } + model = SentenceTransformerModel(**kwargs) + inputs = { + "input_ids": np.array([[2, 3, 4]], dtype="int32"), + "attention_mask": np.ones((1, 3), dtype="int32"), + "token_type_ids": np.zeros((1, 3), dtype="int32"), + } + expected = ops.convert_to_numpy(model(inputs)["embeddings"]) + path = tmp_path / "model.weights.h5" + model.save_weights(str(path)) + + reloaded = SentenceTransformerModel(**model.config.constructor_kwargs()) + reloaded.load_weights(str(path)) + actual = ops.convert_to_numpy(reloaded(inputs)["embeddings"]) + np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-5) diff --git a/website/mkdocs.yml b/website/mkdocs.yml index 06eb039b..e9f38e9d 100644 --- a/website/mkdocs.yml +++ b/website/mkdocs.yml @@ -159,6 +159,7 @@ nav: - Mixtral: mixtral.md - ModernBERT: modernbert.md - MPNet: mpnet.md + - Sentence Transformer: sentence_transformer.md - Qwen: qwen.md - Qwen2: qwen2.md - Qwen2 MoE: qwen2_moe.md diff --git a/zeromodels/auto/auto_mapping_names.py b/zeromodels/auto/auto_mapping_names.py index a659a0e0..04e98093 100644 --- a/zeromodels/auto/auto_mapping_names.py +++ b/zeromodels/auto/auto_mapping_names.py @@ -286,6 +286,7 @@ "sam3": "SAM3Model", "segformer": "SegFormerModel", "senet": "SENetModel", + "sentence_transformer": "SentenceTransformerModel", "siglip": "SigLIPModel", "siglip2": "SigLIP2Model", "stable_diffusion": "StableDiffusionModel", @@ -369,6 +370,7 @@ }, "TextEmbed": { "clip": "CLIPTextEmbed", + "sentence_transformer": "SentenceTransformerTextEmbed", }, "TextGenerate": { "deepseek_v2": "DeepseekV2TextGenerate", @@ -688,6 +690,7 @@ "sam3": "Sam3Config", "segformer": "SegformerConfig", "senet": "SENetConfig", + "sentence_transformer": "SentenceTransformerConfig", "siglip": "SigLIPConfig", "siglip2": "Siglip2Config", "siglip2_text": "Siglip2TextConfig", @@ -801,6 +804,7 @@ "qwen3_vl_moe_text": "Qwen3VLMoeTokenizer", "roberta": "RobertaTokenizer", "sam3": "SAM3CLIPTokenizer", + "sentence_transformer": "SentenceTransformerTokenizer", "siglip": "SigLIPTokenizer", "siglip2": "SigLIP2Tokenizer", "speech_to_text": "Speech2TextTokenizer", diff --git a/zeromodels/conversion/hf_download_utils.py b/zeromodels/conversion/hf_download_utils.py index 79e580b3..647d3d6d 100644 --- a/zeromodels/conversion/hf_download_utils.py +++ b/zeromodels/conversion/hf_download_utils.py @@ -97,7 +97,7 @@ def to_numpy(v): return state -def download_hf_state_dict(hf_id, token=None): +def download_hf_state_dict(hf_id, token=None, revision=None): """Download model weights and return a ``{name: numpy_array}`` mapping. ``token`` is forwarded to ``hf_hub_download`` for gated / private repos. @@ -114,41 +114,49 @@ def download_hf_state_dict(hf_id, token=None): 4. ``pytorch_model.bin.index.json`` (sharded pickle) """ try: - path = hf_hub_download(hf_id, "model.safetensors", token=token) + path = hf_hub_download( + hf_id, "model.safetensors", token=token, revision=revision + ) except EntryNotFoundError: path = None if path is not None: return LazyStateDict.from_files([path]) try: - index_path = hf_hub_download(hf_id, "model.safetensors.index.json", token=token) + index_path = hf_hub_download( + hf_id, "model.safetensors.index.json", token=token, revision=revision + ) except EntryNotFoundError: index_path = None if index_path is not None: with open(index_path, "r") as f: weight_map = json.load(f)["weight_map"] local = { - shard: hf_hub_download(hf_id, shard, token=token) + shard: hf_hub_download(hf_id, shard, token=token, revision=revision) for shard in sorted(set(weight_map.values())) } return LazyStateDict({name: local[shard] for name, shard in weight_map.items()}) try: - path = hf_hub_download(hf_id, "pytorch_model.bin", token=token) + path = hf_hub_download( + hf_id, "pytorch_model.bin", token=token, revision=revision + ) except EntryNotFoundError: path = None if path is not None: return load_bin_state_dict([path]) try: - index_path = hf_hub_download(hf_id, "pytorch_model.bin.index.json", token=token) + index_path = hf_hub_download( + hf_id, "pytorch_model.bin.index.json", token=token, revision=revision + ) except EntryNotFoundError: index_path = None if index_path is not None: with open(index_path, "r") as f: weight_map = json.load(f)["weight_map"] shards = [ - hf_hub_download(hf_id, shard, token=token) + hf_hub_download(hf_id, shard, token=token, revision=revision) for shard in sorted(set(weight_map.values())) ] return load_bin_state_dict(shards) diff --git a/zeromodels/models/__init__.py b/zeromodels/models/__init__.py index 74301652..5c282d74 100644 --- a/zeromodels/models/__init__.py +++ b/zeromodels/models/__init__.py @@ -121,6 +121,7 @@ sam3, segformer, senet, + sentence_transformer, siglip, siglip2, speech2text, diff --git a/zeromodels/models/sentence_transformer/__init__.py b/zeromodels/models/sentence_transformer/__init__.py new file mode 100644 index 00000000..a9beb6a8 --- /dev/null +++ b/zeromodels/models/sentence_transformer/__init__.py @@ -0,0 +1,13 @@ +from .sentence_transformer_config import SentenceTransformerConfig +from .sentence_transformer_model import ( + SentenceTransformerModel, + SentenceTransformerTextEmbed, +) +from .sentence_transformer_tokenizer import SentenceTransformerTokenizer + +__all__ = [ + "SentenceTransformerConfig", + "SentenceTransformerModel", + "SentenceTransformerTextEmbed", + "SentenceTransformerTokenizer", +] diff --git a/zeromodels/models/sentence_transformer/convert_sentence_transformer_hf_to_keras.py b/zeromodels/models/sentence_transformer/convert_sentence_transformer_hf_to_keras.py new file mode 100644 index 00000000..b587d161 --- /dev/null +++ b/zeromodels/models/sentence_transformer/convert_sentence_transformer_hf_to_keras.py @@ -0,0 +1,289 @@ +"""Transfer Sentence-Transformers / HF MPNet weights into Keras ST models.""" + +from __future__ import annotations + +from typing import Dict, Optional, Tuple + +import numpy as np +from huggingface_hub import hf_hub_download +from huggingface_hub.utils import EntryNotFoundError +from safetensors import safe_open + +from zeromodels.conversion.exceptions import WeightMappingError +from zeromodels.conversion.weight_transfer_util import transfer_weights +from zeromodels.models.bert.convert_bert_hf_to_keras import transfer_bert_weights +from zeromodels.models.mpnet.convert_mpnet_hf_to_keras import ( + normalize_hf_key, + transfer_mpnet_weights, +) +from zeromodels.models.roberta.convert_roberta_hf_to_keras import ( + transfer_roberta_weights, +) +from zeromodels.models.xlm_roberta.convert_xlm_roberta_hf_to_keras import ( + transfer_xlm_roberta_weights, +) + + +def strip_st_prefix(key: str) -> str: + """Normalize ST / HF key prefixes down to bare MPNet (or Dense) names. + + Handles common layouts: + + * ``0.auto_model.*`` — SentenceTransformer module stack + * ``auto_model.*`` + * ``mpnet.*`` — HF ``MPNetModel`` wrapper + """ + for prefix in ( + "0.auto_model.", + "auto_model.", + "0.transformer.", + ): + if key.startswith(prefix): + key = key[len(prefix) :] + break + return normalize_hf_key(key) + + +def _find_dense_tensors( + hf_state_dict: Dict[str, np.ndarray], +) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]: + """Locate an ST ``Dense`` module's ``linear.weight`` / ``linear.bias``.""" + kernel = bias = None + for k, v in hf_state_dict.items(): + if "auto_model" in k or ".layer." in k or k.startswith("embeddings."): + continue + if k.endswith("linear.weight"): + kernel = np.asarray(v) + elif k.endswith("linear.bias"): + bias = np.asarray(v) + return kernel, bias + + +def transfer_sentence_transformer_weights( + keras_model, hf_state_dict: Dict[str, np.ndarray] +) -> None: + """Copy backbone (+ optional Dense projection) weights into ``keras_model``. + + Pooling and L2-normalize have no parameters. Unused HF pooler / MLM head + tensors are skipped via the MPNet optional-weight rules. + """ + hf = {strip_st_prefix(k): v for k, v in hf_state_dict.items()} + kind = keras_model.backbone_type + if kind == "mpnet": + transfer_mpnet_weights(keras_model, hf) + elif kind == "bert": + transfer_bert_weights(keras_model, hf) + elif kind == "roberta": + transfer_roberta_weights(keras_model, hf) + elif kind == "xlm-roberta": + transfer_xlm_roberta_weights(keras_model, hf) + elif kind in ("distilbert", "albert"): + _transfer_custom_backbone(keras_model, hf) + else: + raise ValueError(f"Unsupported sentence embedding backbone: {kind}") + + needs_proj = any(w.path.startswith("projection/") for w in keras_model.weights) + if not needs_proj: + return + + dense_state = getattr(keras_model, "_st_dense_state", hf_state_dict) + proj_kernel, proj_bias = _find_dense_tensors(dense_state) + for weight in keras_model.weights: + path = weight.path + if path == "projection/kernel": + if proj_kernel is None: + raise WeightMappingError(path, "*.linear.weight") + transfer_weights(path, weight, proj_kernel) + elif path == "projection/bias": + if proj_bias is None: + raise WeightMappingError(path, "*.linear.bias") + transfer_weights(path, weight, proj_bias) + + +def _custom_hf_name(path, backbone_type): + """Map the two ST-only backbone families to their HF tensor names.""" + suffix = path.rsplit("/", 1)[-1] + suffix = { + "kernel": "weight", + "embeddings": "weight", + "gamma": "weight", + "beta": "bias", + }.get(suffix, suffix) + if backbone_type == "distilbert": + fixed = { + "distil_word_embeddings": "embeddings.word_embeddings", + "distil_position_embeddings": "embeddings.position_embeddings", + "distil_embeddings_norm": "embeddings.LayerNorm", + } + if "/" not in path: + return None + outer, inner = path.split("/", 1) + if outer in fixed: + return f"{fixed[outer]}.{suffix}" + if not outer.startswith("distil_layer_"): + return None + index = int(outer.removeprefix("distil_layer_")) + part = inner.split("/", 1)[0] + part = { + "q_lin": "attention.q_lin", + "k_lin": "attention.k_lin", + "v_lin": "attention.v_lin", + "out_lin": "attention.out_lin", + "ffn_lin1": "ffn.lin1", + "ffn_lin2": "ffn.lin2", + }.get(part, part) + return f"transformer.layer.{index}.{part}.{suffix}" + + fixed = { + "albert_word_embeddings": "embeddings.word_embeddings", + "albert_position_embeddings": "embeddings.position_embeddings", + "albert_token_type_embeddings": "embeddings.token_type_embeddings", + "albert_embeddings_norm": "embeddings.LayerNorm", + "albert_embedding_projection": "encoder.embedding_hidden_mapping_in", + } + if "/" not in path: + return None + outer, inner = path.split("/", 1) + if outer in fixed: + return f"{fixed[outer]}.{suffix}" + if outer != "albert_shared_block": + return None + part = inner.split("/", 1)[0] + part = { + "query": "attention.query", + "key": "attention.key", + "value": "attention.value", + "attention_dense": "attention.dense", + "attention_layer_norm": "attention.LayerNorm", + "full_layer_norm": "full_layer_layer_norm", + }.get(part, part) + return f"encoder.albert_layer_groups.0.albert_layers.0.{part}.{suffix}" + + +def _transfer_custom_backbone(keras_model, hf_state_dict): + kind = keras_model.backbone_type + prefix = "distilbert." if kind == "distilbert" else "albert." + names = {k.removeprefix(prefix): k for k in hf_state_dict} + for weight in keras_model.weights: + hf_name = _custom_hf_name(weight.path, kind) + if hf_name is None: + continue + if hf_name not in names: + raise WeightMappingError(weight.path, hf_name) + source = hf_state_dict[names[hf_name]] + if weight.path == "albert_embedding_projection/kernel": + # The generic transfer helper treats any path containing + # "embedding" as an Embedding table, but this is a Dense kernel. + weight.assign(np.asarray(source).T) + else: + transfer_weights(weight.path, weight, source) + + +def load_dense_state_dict(hf_id, rel_path, *, token=None, revision=None): + """Read the trained Dense module from its declared ST subdirectory.""" + if not rel_path or rel_path.startswith("/") or ".." in rel_path.split("/"): + raise ValueError( + f"Unsafe Sentence Transformers Dense module path: {rel_path!r}" + ) + try: + path = hf_hub_download( + hf_id, f"{rel_path}/model.safetensors", token=token, revision=revision + ) + except EntryNotFoundError: + path = hf_hub_download( + hf_id, f"{rel_path}/pytorch_model.bin", token=token, revision=revision + ) + import torch + + state = torch.load(path, map_location="cpu", weights_only=True) + return {key: value.detach().cpu().numpy() for key, value in state.items()} + with safe_open(path, framework="np") as handle: + return {key: handle.get_tensor(key) for key in handle.keys()} + + +def pooling_mode_from_config(pooling_cfg: dict) -> str: + if pooling_cfg.get("pooling_mode_cls_token"): + return "cls" + if pooling_cfg.get("pooling_mode_max_tokens"): + return "max" + return "mean" + + +if __name__ == "__main__": + import gc + import os + + import keras + import torch + from transformers import AutoModel, AutoTokenizer + + from zeromodels.conversion.zm_config import write_zm_config + from zeromodels.models.sentence_transformer import ( + SentenceTransformerModel, + SentenceTransformerTokenizer, + ) + + HF_TOKEN = os.environ.get("HF_TOKEN") + VARIANTS = { + "all-mpnet-base-v2": "sentence-transformers/all-mpnet-base-v2", + } + + texts = [ + "A man is eating food.", + "A man is eating a piece of bread.", + "The weather is lovely today.", + ] + + for variant, hf_id in VARIANTS.items(): + print(f"\n{'=' * 60}\nConverting: {variant} <- {hf_id}\n{'=' * 60}") + + keras_model = SentenceTransformerModel.from_hf(hf_id) + tok = SentenceTransformerTokenizer.from_hf( + hf_id, max_seq_len=keras_model.max_seq_len + ) + k_inputs = tok(texts) + k_emb = keras.ops.convert_to_numpy( + keras_model(k_inputs, training=False)["embeddings"] + ) + + # Reference: transformers backbone + mean pool + L2 (matches ST default). + hf_tok = AutoTokenizer.from_pretrained(hf_id, token=HF_TOKEN) + hf_enc = AutoModel.from_pretrained(hf_id, token=HF_TOKEN).eval() + enc = hf_tok( + texts, + padding=True, + truncation=True, + max_length=keras_model.max_seq_len, + return_tensors="pt", + ) + with torch.no_grad(): + out = hf_enc(**enc) + mask = enc["attention_mask"].unsqueeze(-1).float() + summed = (out.last_hidden_state * mask).sum(1) + counts = mask.sum(1).clamp(min=1e-9) + pooled = summed / counts + ref = torch.nn.functional.normalize(pooled, p=2, dim=1).cpu().numpy() + + diff = float(np.abs(k_emb - ref).max()) + cos = float(np.sum(k_emb * ref, axis=-1).mean()) + print(f" max abs diff vs mean-pool+L2 reference: {diff:.3e}") + print(f" mean cosine(keras, ref): {cos:.6f}") + if diff > 1e-3: + raise ValueError(f"{variant}: parity failed (diff={diff})") + + dest = variant + os.makedirs(dest, exist_ok=True) + out_path = os.path.join(dest, "model.weights.h5") + keras_model.save_weights(out_path) + write_zm_config( + dest, + SentenceTransformerModel, + variant, + keras_model.config, + weights_filename="model.weights.h5", + ) + print(f" Saved -> {dest}/ (model.weights.h5 + zm_config.json)") + + del keras_model, hf_enc + keras.backend.clear_session() + gc.collect() diff --git a/zeromodels/models/sentence_transformer/sentence_transformer_backbones.py b/zeromodels/models/sentence_transformer/sentence_transformer_backbones.py new file mode 100644 index 00000000..c45b7be3 --- /dev/null +++ b/zeromodels/models/sentence_transformer/sentence_transformer_backbones.py @@ -0,0 +1,225 @@ +"""Encoder backbones needed by the original Sentence Transformers checkpoints.""" + +import math + +import keras +from keras import layers, ops + + +def _attention(query, key, value, attention_mask, num_heads): + batch = ops.shape(query)[0] + length = ops.shape(query)[1] + head_dim = query.shape[-1] // num_heads + + def split_heads(x): + x = ops.reshape(x, (batch, length, num_heads, head_dim)) + return ops.transpose(x, (0, 2, 1, 3)) + + q, k, v = (split_heads(x) for x in (query, key, value)) + scores = ops.matmul(q, ops.transpose(k, (0, 1, 3, 2))) + scores = scores / math.sqrt(head_dim) + mask = ops.cast(attention_mask[:, None, None, :], "bool") + scores = ops.where(mask, scores, ops.cast(-1e4, scores.dtype)) + context = ops.matmul(ops.softmax(scores, axis=-1), v) + context = ops.transpose(context, (0, 2, 1, 3)) + return ops.reshape(context, (batch, length, num_heads * head_dim)) + + +def _gelu_new(x): + return ( + 0.5 + * x + * (1.0 + ops.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * ops.power(x, 3)))) + ) + + +@keras.saving.register_keras_serializable(package="zeromodels") +class DistilBertEncoderBlock(layers.Layer): + """One post-LayerNorm DistilBERT transformer layer.""" + + def __init__( + self, embed_dim, num_heads, mlp_dim, activation="gelu", norm_eps=1e-12, **kwargs + ): + super().__init__(**kwargs) + self.embed_dim = embed_dim + self.num_heads = num_heads + self.mlp_dim = mlp_dim + self.activation = activation + self.norm_eps = norm_eps + self.q_lin = layers.Dense(embed_dim, name="q_lin") + self.k_lin = layers.Dense(embed_dim, name="k_lin") + self.v_lin = layers.Dense(embed_dim, name="v_lin") + self.out_lin = layers.Dense(embed_dim, name="out_lin") + self.sa_layer_norm = layers.LayerNormalization( + epsilon=norm_eps, name="sa_layer_norm" + ) + self.ffn_lin1 = layers.Dense(mlp_dim, name="ffn_lin1") + self.ffn_lin2 = layers.Dense(embed_dim, name="ffn_lin2") + self.output_layer_norm = layers.LayerNormalization( + epsilon=norm_eps, name="output_layer_norm" + ) + + def call(self, inputs): + x, attention_mask = inputs + context = _attention( + self.q_lin(x), self.k_lin(x), self.v_lin(x), attention_mask, self.num_heads + ) + x = self.sa_layer_norm(x + self.out_lin(context)) + hidden = self.ffn_lin1(x) + if self.activation == "gelu": + hidden = ops.gelu(hidden) + else: + raise ValueError(f"Unsupported DistilBERT activation: {self.activation}") + return self.output_layer_norm(x + self.ffn_lin2(hidden)) + + def get_config(self): + return { + **super().get_config(), + "embed_dim": self.embed_dim, + "num_heads": self.num_heads, + "mlp_dim": self.mlp_dim, + "activation": self.activation, + "norm_eps": self.norm_eps, + } + + +@keras.saving.register_keras_serializable(package="zeromodels") +class AlbertEncoderBlock(layers.Layer): + """Shared ALBERT attention and feed-forward group.""" + + def __init__( + self, + embed_dim, + num_heads, + mlp_dim, + activation="gelu_new", + norm_eps=1e-12, + **kwargs, + ): + super().__init__(**kwargs) + self.embed_dim = embed_dim + self.num_heads = num_heads + self.mlp_dim = mlp_dim + self.activation = activation + self.norm_eps = norm_eps + self.query = layers.Dense(embed_dim, name="query") + self.key = layers.Dense(embed_dim, name="key") + self.value = layers.Dense(embed_dim, name="value") + self.attention_dense = layers.Dense(embed_dim, name="attention_dense") + self.attention_layer_norm = layers.LayerNormalization( + epsilon=norm_eps, name="attention_layer_norm" + ) + self.ffn = layers.Dense(mlp_dim, name="ffn") + self.ffn_output = layers.Dense(embed_dim, name="ffn_output") + self.full_layer_norm = layers.LayerNormalization( + epsilon=norm_eps, name="full_layer_norm" + ) + + def call(self, inputs): + x, attention_mask = inputs + context = _attention( + self.query(x), self.key(x), self.value(x), attention_mask, self.num_heads + ) + x = self.attention_layer_norm(x + self.attention_dense(context)) + hidden = self.ffn(x) + if self.activation == "gelu_new": + hidden = _gelu_new(hidden) + elif self.activation == "gelu": + hidden = ops.gelu(hidden) + else: + raise ValueError(f"Unsupported ALBERT activation: {self.activation}") + return self.full_layer_norm(x + self.ffn_output(hidden)) + + def get_config(self): + return { + **super().get_config(), + "embed_dim": self.embed_dim, + "num_heads": self.num_heads, + "mlp_dim": self.mlp_dim, + "activation": self.activation, + "norm_eps": self.norm_eps, + } + + +def distilbert_backbone( + input_ids, + attention_mask, + *, + vocab_size, + embed_dim, + num_layers, + num_heads, + mlp_dim, + max_position_embeddings, + hidden_act, + layer_norm_eps, +): + positions = ops.cumsum(ops.ones_like(input_ids), axis=1) - 1 + x = layers.Embedding(vocab_size, embed_dim, name="distil_word_embeddings")( + input_ids + ) + x = x + layers.Embedding( + max_position_embeddings, embed_dim, name="distil_position_embeddings" + )(positions) + x = layers.LayerNormalization( + epsilon=layer_norm_eps, name="distil_embeddings_norm" + )(x) + for index in range(num_layers): + x = DistilBertEncoderBlock( + embed_dim, + num_heads, + mlp_dim, + hidden_act, + layer_norm_eps, + name=f"distil_layer_{index}", + )([x, attention_mask]) + return x + + +def albert_backbone( + input_ids, + attention_mask, + token_type_ids, + *, + vocab_size, + embed_dim, + embedding_size, + num_layers, + num_heads, + mlp_dim, + max_position_embeddings, + type_vocab_size, + hidden_act, + layer_norm_eps, + num_hidden_groups=1, + inner_group_num=1, +): + if num_hidden_groups != 1 or inner_group_num != 1: + raise ValueError( + "Only ALBERT with one shared hidden group and one inner group is supported" + ) + positions = ops.cumsum(ops.ones_like(input_ids), axis=1) - 1 + x = layers.Embedding(vocab_size, embedding_size, name="albert_word_embeddings")( + input_ids + ) + x = x + layers.Embedding( + max_position_embeddings, embedding_size, name="albert_position_embeddings" + )(positions) + x = x + layers.Embedding( + type_vocab_size, embedding_size, name="albert_token_type_embeddings" + )(token_type_ids) + x = layers.LayerNormalization( + epsilon=layer_norm_eps, name="albert_embeddings_norm" + )(x) + x = layers.Dense(embed_dim, name="albert_embedding_projection")(x) + shared = AlbertEncoderBlock( + embed_dim, + num_heads, + mlp_dim, + hidden_act, + layer_norm_eps, + name="albert_shared_block", + ) + for _ in range(num_layers): + x = shared([x, attention_mask]) + return x diff --git a/zeromodels/models/sentence_transformer/sentence_transformer_config.py b/zeromodels/models/sentence_transformer/sentence_transformer_config.py new file mode 100644 index 00000000..9bb7fccf --- /dev/null +++ b/zeromodels/models/sentence_transformer/sentence_transformer_config.py @@ -0,0 +1,41 @@ +from typing import Optional + +from zeromodels.base import BaseConfig + + +class SentenceTransformerConfig(BaseConfig): + """Configuration for :class:`SentenceTransformerModel`. + + A Sentence-Transformers-style stack: encoder backbone + pooling (+ optional + dense projection) + optional L2 normalize. ``backbone_type`` selects which + ZeroModels encoder family builds the tower. + + Fields mirror the flat model constructor and serialize into ``zm_config.json``. + """ + + model_type = "sentence_transformer" + + backbone_type: str = "mpnet" + vocab_size: int = 30527 + embed_dim: int = 768 + num_layers: int = 12 + num_heads: int = 12 + mlp_dim: int = 3072 + max_position_embeddings: int = 514 + relative_attention_num_buckets: int = 32 + hidden_act: str = "gelu" + layer_norm_eps: float = 1e-5 # matches microsoft/mpnet-base / ST hubs + pad_token_id: int = 1 + type_vocab_size: int = 2 + embedding_size: Optional[int] = None + num_hidden_groups: int = 1 + inner_group_num: int = 1 + dropout: float = 0.0 + attention_dropout: float = 0.0 + + pooling: str = "mean" # mean | cls | max + normalize_embeddings: bool = True + projection_dim: Optional[int] = None + projection_activation: str = "linear" + similarity_fn_name: str = "cosine" + max_seq_len: int = 384 # all-mpnet-base-v2 sentence_bert_config.json diff --git a/zeromodels/models/sentence_transformer/sentence_transformer_layers.py b/zeromodels/models/sentence_transformer/sentence_transformer_layers.py new file mode 100644 index 00000000..fb618d54 --- /dev/null +++ b/zeromodels/models/sentence_transformer/sentence_transformer_layers.py @@ -0,0 +1,98 @@ +"""Pooling / projection helpers for Sentence-Transformers-style models.""" + +import keras +from keras import layers, ops + + +def masked_mean_pool(token_embeddings, attention_mask): + """Attention-masked mean over the sequence (ST ``pooling_mode_mean_tokens``).""" + mask = ops.cast(attention_mask, token_embeddings.dtype) + mask = ops.expand_dims(mask, axis=-1) + summed = ops.sum(token_embeddings * mask, axis=1) + counts = ops.maximum(ops.sum(mask, axis=1), 1e-9) + return summed / counts + + +def cls_pool(token_embeddings, attention_mask=None): + """First-token pool (ST ``pooling_mode_cls_token``).""" + del attention_mask + return token_embeddings[:, 0] + + +def masked_max_pool(token_embeddings, attention_mask): + """Attention-masked max over the sequence (ST ``pooling_mode_max_tokens``).""" + mask = ops.cast(attention_mask, token_embeddings.dtype) + mask = ops.expand_dims(mask, axis=-1) + # Large negative so padded positions never win the max. + masked = token_embeddings + (1.0 - mask) * (-1e9) + return ops.max(masked, axis=1) + + +POOLERS = { + "mean": masked_mean_pool, + "cls": cls_pool, + "max": masked_max_pool, +} + + +def pool_token_embeddings(token_embeddings, attention_mask, pooling="mean"): + if pooling not in POOLERS: + raise ValueError( + f"Unknown pooling={pooling!r}; expected one of {sorted(POOLERS)}" + ) + return POOLERS[pooling](token_embeddings, attention_mask) + + +def l2_normalize(embeddings, axis=-1, eps=1e-12): + norm = ops.sqrt( + ops.maximum(ops.sum(ops.square(embeddings), axis=axis, keepdims=True), eps) + ) + return embeddings / norm + + +def cosine_similarity(a, b, eps=1e-12): + """Pairwise cosine similarity matrix ``(B_a, B_b)`` for row-normalized or raw vectors.""" + a = ops.convert_to_tensor(a) + b = ops.convert_to_tensor(b) + a = l2_normalize(a, eps=eps) + b = l2_normalize(b, eps=eps) + return ops.matmul(a, ops.transpose(b)) + + +@keras.saving.register_keras_serializable(package="zeromodels") +class SentenceTransformerPooling(layers.Layer): + """Mean / CLS / max pool over token embeddings with an attention mask.""" + + def __init__(self, pooling="mean", **kwargs): + super().__init__(**kwargs) + if pooling not in POOLERS: + raise ValueError( + f"Unknown pooling={pooling!r}; expected one of {sorted(POOLERS)}" + ) + self.pooling = pooling + + def call(self, inputs): + token_embeddings, attention_mask = inputs + return pool_token_embeddings(token_embeddings, attention_mask, self.pooling) + + def get_config(self): + config = super().get_config() + config.update({"pooling": self.pooling}) + return config + + +@keras.saving.register_keras_serializable(package="zeromodels") +class SentenceTransformerNormalize(layers.Layer): + """L2-normalize embeddings along the last axis.""" + + def __init__(self, eps=1e-12, **kwargs): + super().__init__(**kwargs) + self.eps = eps + + def call(self, embeddings): + return l2_normalize(embeddings, eps=self.eps) + + def get_config(self): + config = super().get_config() + config.update({"eps": self.eps}) + return config diff --git a/zeromodels/models/sentence_transformer/sentence_transformer_model.py b/zeromodels/models/sentence_transformer/sentence_transformer_model.py new file mode 100644 index 00000000..7cb942a5 --- /dev/null +++ b/zeromodels/models/sentence_transformer/sentence_transformer_model.py @@ -0,0 +1,566 @@ +"""Sentence-Transformers-style text embedding models (Keras 3).""" + +from __future__ import annotations + +import json +import os +import re + +import keras +from keras import layers, ops + +from zeromodels.base import BaseModel +from zeromodels.base.base_mixin import warn_skipped +from zeromodels.conversion.hf_download_utils import download_hf_state_dict +from zeromodels.conversion.weight_transfer_util import skip_mismatched_weights + +from .sentence_transformer_backbones import albert_backbone, distilbert_backbone +from .sentence_transformer_config import SentenceTransformerConfig +from .sentence_transformer_layers import ( + SentenceTransformerNormalize, + SentenceTransformerPooling, + cosine_similarity, +) + +SENTENCE_TRANSFORMER_HUB_SIBLINGS = frozenset({"SentenceTransformerModel"}) + + +@keras.saving.register_keras_serializable(package="zeromodels") +class SentenceTransformerModel(BaseModel): + """Sentence embedding model: encoder + pooling (+ optional Dense) + normalize. + + Matches popular Sentence-Transformers stacks such as + ``sentence-transformers/all-mpnet-base-v2`` (MPNet encoder, mean pooling, + L2 normalize). Load an upstream checkpoint with:: + + model = SentenceTransformerModel.from_weights( + "hf:sentence-transformers/all-mpnet-base-v2" + ) + tok = SentenceTransformerTokenizer.from_weights( + "hf:sentence-transformers/all-mpnet-base-v2", + revision=model.source_revision, + ) + emb = model(tok(["a cat", "a dog"]))["embeddings"] + + Or the ST-style helper:: + + emb = model.encode(["a cat", "a dog"], tokenizer=tok) + + Args: + backbone_type: Encoder family (MPNet, BERT, RoBERTa, XLM-R, DistilBERT, ALBERT). + pooling: ``\"mean\"`` | ``\"cls\"`` | ``\"max\"``. + normalize_embeddings: If True, L2-normalize the output vectors. + projection_dim: Optional Dense width after pooling (None = identity). + max_seq_len: Default tokenizer truncation length (stored in config). + Other kwargs: forwarded to the backbone (``vocab_size``, ``embed_dim``, …). + """ + + config_class = SentenceTransformerConfig + # Accept zeromodels ``sentence_transformer`` configs and upstream ST hubs + # whose root ``config.json`` is the backbone transformers config + # (``model_type: mpnet`` for all-mpnet-base-v2). Pooling / normalize come + # from sibling ST files read in :meth:`from_hf`. + HF_MODEL_TYPE = ( + "sentence_transformer", + "mpnet", + "bert", + "roberta", + "xlm-roberta", + "distilbert", + "albert", + ) + HUB_REPO_SIBLINGS = SENTENCE_TRANSFORMER_HUB_SIBLINGS + + def __init__( + self, + backbone_type="mpnet", + vocab_size=30527, + embed_dim=768, + num_layers=12, + num_heads=12, + mlp_dim=3072, + max_position_embeddings=514, + relative_attention_num_buckets=32, + hidden_act="gelu", + layer_norm_eps=1e-5, + pad_token_id=1, + type_vocab_size=2, + embedding_size=None, + num_hidden_groups=1, + inner_group_num=1, + dropout=0.0, + attention_dropout=0.0, + pooling="mean", + normalize_embeddings=True, + projection_dim=None, + projection_activation="linear", + similarity_fn_name="cosine", + max_seq_len=384, + name="SentenceTransformerModel", + **kwargs, + ): + for k in ("model", "hf_id", "url", "num_classes"): + kwargs.pop(k, None) + + if backbone_type not in ( + "mpnet", + "bert", + "roberta", + "xlm-roberta", + "distilbert", + "albert", + ): + raise ValueError( + f"Unsupported sentence embedding backbone: {backbone_type!r}" + ) + if pooling not in ("mean", "cls", "max"): + raise ValueError(f"pooling must be mean|cls|max, got {pooling!r}") + if projection_activation not in ("linear", "tanh"): + raise ValueError( + f"Unsupported projection activation: {projection_activation!r}" + ) + if similarity_fn_name not in ("cosine", "dot"): + raise ValueError(f"Unsupported similarity function: {similarity_fn_name!r}") + + inputs = { + "input_ids": layers.Input(shape=(None,), dtype="int32", name="input_ids"), + "attention_mask": layers.Input( + shape=(None,), dtype="int32", name="attention_mask" + ), + } + if backbone_type in ("bert", "roberta", "xlm-roberta", "albert"): + inputs["token_type_ids"] = layers.Input( + shape=(None,), dtype="int32", name="token_type_ids" + ) + common = { + "vocab_size": vocab_size, + "embed_dim": embed_dim, + "num_layers": num_layers, + "num_heads": num_heads, + "mlp_dim": mlp_dim, + "max_position_embeddings": max_position_embeddings, + "hidden_act": hidden_act, + "layer_norm_eps": layer_norm_eps, + } + if backbone_type == "mpnet": + from zeromodels.models.mpnet.mpnet_model import mpnet_backbone + + sequence_output, _ = mpnet_backbone( + inputs["input_ids"], + inputs["attention_mask"], + **common, + relative_attention_num_buckets=relative_attention_num_buckets, + pad_token_id=pad_token_id, + dropout=dropout, + attention_dropout=attention_dropout, + add_pooler=False, + ) + elif backbone_type in ("roberta", "xlm-roberta"): + from zeromodels.models.roberta.roberta_model import roberta_backbone + + sequence_output, _ = roberta_backbone( + inputs["input_ids"], + inputs["attention_mask"], + inputs["token_type_ids"], + **common, + type_vocab_size=type_vocab_size, + pad_token_id=pad_token_id, + dropout=dropout, + attention_dropout=attention_dropout, + add_pooler=False, + ) + elif backbone_type == "bert": + from zeromodels.models.bert.bert_model import bert_backbone + + sequence_output, _ = bert_backbone( + inputs["input_ids"], + inputs["attention_mask"], + inputs["token_type_ids"], + **common, + type_vocab_size=type_vocab_size, + dropout=dropout, + attention_dropout=attention_dropout, + add_pooler=False, + ) + elif backbone_type == "distilbert": + sequence_output = distilbert_backbone( + inputs["input_ids"], + inputs["attention_mask"], + **common, + ) + else: + sequence_output = albert_backbone( + inputs["input_ids"], + inputs["attention_mask"], + inputs["token_type_ids"], + **common, + embedding_size=embedding_size or embed_dim, + type_vocab_size=type_vocab_size, + num_hidden_groups=num_hidden_groups, + inner_group_num=inner_group_num, + ) + + pooled = SentenceTransformerPooling(pooling=pooling, name="pooling")( + [sequence_output, inputs["attention_mask"]] + ) + + if projection_dim is not None: + pooled = layers.Dense( + int(projection_dim), activation=projection_activation, name="projection" + )(pooled) + out_dim = int(projection_dim) + else: + out_dim = embed_dim + + if normalize_embeddings: + embeddings = SentenceTransformerNormalize(name="normalize")(pooled) + else: + embeddings = pooled + + super().__init__( + inputs=inputs, + outputs={ + "embeddings": embeddings, + "token_embeddings": sequence_output, + }, + name=name, + **kwargs, + ) + self.backbone_type = backbone_type + self.vocab_size = vocab_size + self.embed_dim = embed_dim + self.num_layers = num_layers + self.num_heads = num_heads + self.mlp_dim = mlp_dim + self.max_position_embeddings = max_position_embeddings + self.relative_attention_num_buckets = relative_attention_num_buckets + self.hidden_act = hidden_act + self.layer_norm_eps = layer_norm_eps + self.pad_token_id = pad_token_id + self.type_vocab_size = type_vocab_size + self.embedding_size = embedding_size + self.num_hidden_groups = num_hidden_groups + self.inner_group_num = inner_group_num + self.dropout = dropout + self.attention_dropout = attention_dropout + self.pooling = pooling + self.normalize_embeddings = normalize_embeddings + self.projection_dim = projection_dim + self.projection_activation = projection_activation + self.similarity_fn_name = similarity_fn_name + self.max_seq_len = max_seq_len + self.output_embedding_dim = out_dim + self.config = SentenceTransformerConfig( + backbone_type=backbone_type, + vocab_size=vocab_size, + embed_dim=embed_dim, + num_layers=num_layers, + num_heads=num_heads, + mlp_dim=mlp_dim, + max_position_embeddings=max_position_embeddings, + relative_attention_num_buckets=relative_attention_num_buckets, + hidden_act=hidden_act, + layer_norm_eps=layer_norm_eps, + pad_token_id=pad_token_id, + type_vocab_size=type_vocab_size, + embedding_size=embedding_size, + num_hidden_groups=num_hidden_groups, + inner_group_num=inner_group_num, + dropout=dropout, + attention_dropout=attention_dropout, + pooling=pooling, + normalize_embeddings=normalize_embeddings, + projection_dim=projection_dim, + projection_activation=projection_activation, + similarity_fn_name=similarity_fn_name, + max_seq_len=max_seq_len, + ) + + @classmethod + def config_from_hf(cls, hf_config): + """Map a transformers backbone ``config.json`` to constructor kwargs. + + Pooling / normalize / max_seq_len are filled by :meth:`from_hf` from the + Sentence-Transformers sidecar files when present. + """ + kind = hf_config["model_type"] + if kind == "distilbert": + return { + "backbone_type": kind, + "vocab_size": hf_config["vocab_size"], + "embed_dim": hf_config["dim"], + "num_layers": hf_config["n_layers"], + "num_heads": hf_config["n_heads"], + "mlp_dim": hf_config["hidden_dim"], + "max_position_embeddings": hf_config["max_position_embeddings"], + "hidden_act": hf_config.get("activation", "gelu"), + "layer_norm_eps": 1e-12, + "pad_token_id": hf_config.get("pad_token_id", 0), + } + result = { + "backbone_type": kind, + "vocab_size": hf_config["vocab_size"], + "embed_dim": hf_config["hidden_size"], + "num_layers": hf_config["num_hidden_layers"], + "num_heads": hf_config["num_attention_heads"], + "mlp_dim": hf_config["intermediate_size"], + "max_position_embeddings": hf_config["max_position_embeddings"], + "hidden_act": hf_config.get("hidden_act", "gelu"), + "layer_norm_eps": hf_config.get("layer_norm_eps", 1e-5), + "pad_token_id": hf_config.get("pad_token_id", 1), + "type_vocab_size": hf_config.get("type_vocab_size", 2), + } + if kind == "mpnet": + result["relative_attention_num_buckets"] = hf_config.get( + "relative_attention_num_buckets", 32 + ) + if kind == "albert": + result.update( + embedding_size=hf_config["embedding_size"], + num_hidden_groups=hf_config.get("num_hidden_groups", 1), + inner_group_num=hf_config.get("inner_group_num", 1), + ) + return result + + @classmethod + def transfer_from_hf(cls, keras_model, state_dict): + from .convert_sentence_transformer_hf_to_keras import ( + transfer_sentence_transformer_weights, + ) + + transfer_sentence_transformer_weights(keras_model, state_dict) + + @classmethod + def from_hf( + cls, + hf_id, + load_weights=True, + variant=None, + skip_mismatch=False, + quantization=None, + **kwargs, + ): + """Load from a Sentence-Transformers (or plain MPNet) Hub repo. + + Reads the backbone ``config.json`` via the base loader, then overlays + pooling / normalize / ``max_seq_len`` from ``modules.json``, + ``*_Pooling/config.json``, and ``sentence_bert_config.json`` when present. + """ + from huggingface_hub import HfApi, hf_hub_download + + token = os.environ.get("HF_TOKEN") + revision = kwargs.pop("revision", None) + if variant is not None: + raise ValueError("variant= is not used for Sentence Transformers Hub repos") + # Resolve once: config, module metadata, tokenizer and weight files must + # describe the same commit even if the upstream default branch moves. + if not (isinstance(revision, str) and re.fullmatch(r"[0-9a-f]{40}", revision)): + revision = HfApi(token=token).model_info(hf_id, revision=revision).sha + with open( + hf_hub_download(hf_id, "config.json", token=token, revision=revision) + ) as f: + hf_config = json.load(f) + cls.assert_hf_model_type(hf_id, hf_config) + if hf_config["model_type"] == "sentence_transformer": + raise ValueError( + "Use the ZeroModels repo loader for a converted checkpoint" + ) + sidecar, dense_path = cls._st_sidecar_kwargs( + hf_id, token=token, revision=revision + ) + model_kwargs = cls.config_from_hf(hf_config) + model_kwargs.update(sidecar) + model_kwargs.update(kwargs) + model = cls(**model_kwargs) + model.source_revision = revision + if not load_weights: + return model + state_dict = download_hf_state_dict(hf_id, token=token, revision=revision) + completed = False + try: + if dense_path: + from .convert_sentence_transformer_hf_to_keras import ( + load_dense_state_dict, + ) + + model._st_dense_state = load_dense_state_dict( + hf_id, dense_path, token=token, revision=revision + ) + with skip_mismatched_weights(skip_mismatch) as skipped: + cls._quantized_transfer(model, state_dict, quantization, skip_mismatch) + completed = True + finally: + close = getattr(state_dict, "close", None) + if callable(close): + close(completed=completed) + if hasattr(model, "_st_dense_state"): + del model._st_dense_state + warn_skipped(skipped) + return model + + @classmethod + def _st_sidecar_kwargs(cls, hf_id, token=None, revision=None): + """Parse Sentence-Transformers module configs into constructor kwargs.""" + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import EntryNotFoundError + + out = {} + try: + path = hf_hub_download( + hf_id, + "config_sentence_transformers.json", + token=token, + revision=revision, + ) + with open(path, "r") as f: + st_config = json.load(f) + if st_config.get("similarity_fn_name") is not None: + out["similarity_fn_name"] = st_config["similarity_fn_name"] + except EntryNotFoundError: + pass + try: + path = hf_hub_download( + hf_id, "sentence_bert_config.json", token=token, revision=revision + ) + with open(path, "r") as f: + sbert = json.load(f) + if "max_seq_length" in sbert: + out["max_seq_len"] = int(sbert["max_seq_length"]) + except EntryNotFoundError: + pass + + modules = None + try: + path = hf_hub_download( + hf_id, "modules.json", token=token, revision=revision + ) + with open(path, "r") as f: + modules = json.load(f) + except EntryNotFoundError: + modules = None + + if modules is None: + raise ValueError( + f"{hf_id} has no modules.json; it is not a Sentence Transformers checkpoint" + ) + types = [mod.get("type") for mod in modules] + allowed = [ + "sentence_transformers.models.Transformer", + "sentence_transformers.models.Pooling", + ] + if types[:2] != allowed or types[2:] not in ( + [], + ["sentence_transformers.models.Normalize"], + ["sentence_transformers.models.Dense"], + [ + "sentence_transformers.models.Dense", + "sentence_transformers.models.Normalize", + ], + ): + raise ValueError( + f"Unsupported Sentence Transformers module recipe in {hf_id}: {types}" + ) + out["pooling"] = cls._pooling_from_module_dir( + hf_id, modules[1].get("path"), token=token, revision=revision + ) + out["normalize_embeddings"] = ( + types[-1] == "sentence_transformers.models.Normalize" + ) + dense_path = None + if "sentence_transformers.models.Dense" in types: + dense_path = modules[2].get("path") + dense = cls._dense_config(hf_id, dense_path, token=token, revision=revision) + out["projection_dim"] = int(dense["out_features"]) + activation = dense.get("activation_function") + if activation in (None, "torch.nn.modules.linear.Identity"): + out["projection_activation"] = "linear" + elif activation == "torch.nn.modules.activation.Tanh": + out["projection_activation"] = "tanh" + else: + raise ValueError( + f"Unsupported Sentence Transformers Dense activation: {activation}" + ) + return out, dense_path + + @classmethod + def _pooling_from_module_dir(cls, hf_id, rel_path, token=None, revision=None): + from huggingface_hub import hf_hub_download + + filename = "config.json" if not rel_path else f"{rel_path}/config.json" + path = hf_hub_download(hf_id, filename, token=token, revision=revision) + with open(path, "r") as f: + cfg = json.load(f) + flags = { + "cls": cfg.get("pooling_mode_cls_token", False), + "mean": cfg.get("pooling_mode_mean_tokens", False), + "max": cfg.get("pooling_mode_max_tokens", False), + } + unsupported = ( + cfg.get("pooling_mode_mean_sqrt_len_tokens") + or cfg.get("pooling_mode_weightedmean_tokens") + or cfg.get("pooling_mode_lasttoken") + or cfg.get("include_prompt") is False + ) + selected = [name for name, enabled in flags.items() if enabled] + if unsupported or len(selected) != 1: + raise ValueError(f"Unsupported Sentence Transformers pooling config: {cfg}") + return selected[0] + + @classmethod + def _dense_config(cls, hf_id, rel_path, token=None, revision=None): + from huggingface_hub import hf_hub_download + + filename = "config.json" if not rel_path else f"{rel_path}/config.json" + path = hf_hub_download(hf_id, filename, token=token, revision=revision) + with open(path, "r") as f: + return json.load(f) + + def encode(self, texts, tokenizer, normalize=None, batch_size=32): + """Embed ``texts`` and return a numpy array ``(N, D)``. + + Args: + texts: String or list of strings. + tokenizer: :class:`SentenceTransformerTokenizer` (or any callable that + returns ``input_ids`` / ``attention_mask``). + normalize: Override config L2-normalize. ``None`` keeps model output. + batch_size: Encode in chunks of this size. + """ + import numpy as np + + if isinstance(texts, str): + texts = [texts] + texts = list(texts) + chunks = [] + for start in range(0, len(texts), batch_size): + batch = texts[start : start + batch_size] + inputs = tokenizer(batch) + out = self(inputs, training=False) + emb = ops.convert_to_numpy(out["embeddings"]) + if normalize is False and self.normalize_embeddings: + raise ValueError( + "This model has a Normalize module; encode(normalize=False) cannot undo it" + ) + if normalize is True: + emb = emb / np.clip( + np.linalg.norm(emb, axis=-1, keepdims=True), 1e-12, None + ) + chunks.append(emb) + if not chunks: + return np.zeros((0, self.output_embedding_dim), dtype="float32") + return np.concatenate(chunks, axis=0) + + def similarity(self, a, b): + """Pairwise scores using this checkpoint's declared similarity function.""" + if self.similarity_fn_name == "dot": + return ops.matmul(a, ops.transpose(b)) + return cosine_similarity(a, b) + + def get_config(self): + config = super().get_config() + config.update(self.config.constructor_kwargs()) + return config + + +@keras.saving.register_keras_serializable(package="zeromodels") +class SentenceTransformerTextEmbed(SentenceTransformerModel): + """Text-embedding task view of :class:`SentenceTransformerModel`.""" diff --git a/zeromodels/models/sentence_transformer/sentence_transformer_tokenizer.py b/zeromodels/models/sentence_transformer/sentence_transformer_tokenizer.py new file mode 100644 index 00000000..83015d2e --- /dev/null +++ b/zeromodels/models/sentence_transformer/sentence_transformer_tokenizer.py @@ -0,0 +1,182 @@ +"""Fast-tokenizer adapter for supported Sentence Transformers encoders.""" + +import re +from typing import List, Union + +import keras +from tokenizers import Tokenizer + +from zeromodels.base import BaseTokenizer + + +@keras.saving.register_keras_serializable(package="zeromodels") +class SentenceTransformerTokenizer(BaseTokenizer): + """Tokenizer for supported ST models, preserving their tokenizer.json rules. + + Loads ``tokenizer.json`` from a Hub repo (via ``from_weights("org/repo")`` / + ``from_hub_repo``) or an explicit ``tokenizer_file``. Defaults match + ``sentence-transformers/all-mpnet-base-v2`` (``max_seq_len=384``, MPNet + special tokens). ``call`` returns ``input_ids`` / ``attention_mask`` only. + + Args: + variant: Optional Hub variant key; resolves + ``zeromodels//tokenizer.json`` when ``tokenizer_file`` is + omitted. + tokenizer_file: Explicit ``tokenizer.json`` path (overrides variant). + max_seq_len: Truncation length (default 384). + bos_token / eos_token / unk_token / pad_token / mask_token: Special tokens. + """ + + def __init__( + self, + variant: str = None, + tokenizer_file: str = None, + max_seq_len: int = 384, + backbone_type: str = "mpnet", + bos_token: str = "", + eos_token: str = "", + unk_token: str = "[UNK]", + pad_token: str = None, + mask_token: str = "", + **kwargs, + ): + super().__init__(**kwargs) + self.variant = variant + tokenizer_file = self.resolve_tokenizer_json_from_hf( + (f"zeromodels/{self.variant}" if self.variant else None), tokenizer_file + ) + self.tokenizer_file = tokenizer_file + self.max_seq_len = max_seq_len + self.backbone_type = backbone_type + self.bos_token = bos_token + self.eos_token = eos_token + self.unk_token = unk_token + self.mask_token = mask_token + + tok = Tokenizer.from_file(tokenizer_file) + if pad_token is None: + pad_token = next( + ( + candidate + for candidate in ("", "[PAD]", "") + if tok.token_to_id(candidate) is not None + ), + None, + ) + if pad_token is None: + raise ValueError("tokenizer.json has no supported padding token") + self.pad_token = pad_token + self.bos_token_id = tok.token_to_id(bos_token) + self.eos_token_id = tok.token_to_id(eos_token) + self.unk_token_id = tok.token_to_id(unk_token) + self.pad_token_id = tok.token_to_id(pad_token) + self.mask_token_id = tok.token_to_id(mask_token) + tok.enable_truncation(max_length=max_seq_len) + tok.enable_padding(pad_id=self.pad_token_id, pad_token=pad_token) + self._tok = tok + + @classmethod + def from_hf(cls, repo, **kwargs): + import json + import os + + from huggingface_hub import HfApi, hf_hub_download + from huggingface_hub.utils import EntryNotFoundError + + token = os.environ.get("HF_TOKEN") + revision = kwargs.pop("revision", None) + if not (isinstance(revision, str) and re.fullmatch(r"[0-9a-f]{40}", revision)): + revision = HfApi(token=token).model_info(repo, revision=revision).sha + with open( + hf_hub_download(repo, "config.json", token=token, revision=revision) + ) as f: + hf_config = json.load(f) + kwargs.setdefault("backbone_type", hf_config["model_type"]) + if "max_seq_len" not in kwargs: + try: + path = hf_hub_download( + repo, "sentence_bert_config.json", token=token, revision=revision + ) + except EntryNotFoundError: + pass + else: + with open(path) as f: + kwargs["max_seq_len"] = int(json.load(f)["max_seq_length"]) + path = hf_hub_download(repo, "tokenizer.json", token=token, revision=revision) + return cls(tokenizer_file=path, **kwargs) + + @classmethod + def from_hub_repo(cls, repo_id, **kwargs): + """Load ``tokenizer.json`` and honour ``max_seq_len`` from zm / ST configs.""" + import json + import os + + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import EntryNotFoundError + + token = os.environ.get("HF_TOKEN") + path = hf_hub_download(repo_id.rstrip("/"), "tokenizer.json", token=token) + if "max_seq_len" not in kwargs: + for filename in ("zm_config.json", "sentence_bert_config.json"): + try: + cfg_path = hf_hub_download( + repo_id.rstrip("/"), filename, token=token + ) + except EntryNotFoundError: + continue + with open(cfg_path, "r") as f: + cfg = json.load(f) + if "max_seq_len" in cfg: + kwargs["max_seq_len"] = int(cfg["max_seq_len"]) + elif "max_seq_length" in cfg: + kwargs["max_seq_len"] = int(cfg["max_seq_length"]) + if "backbone_type" in cfg: + kwargs.setdefault("backbone_type", cfg["backbone_type"]) + if "max_seq_len" in kwargs: + break + return cls(tokenizer_file=path, **kwargs) + + @property + def vocab_size(self) -> int: + return self._tok.get_vocab_size() + + def tokenize( + self, text: Union[str, List[str]] + ) -> Union[List[int], List[List[int]]]: + if isinstance(text, str): + return self._tok.encode(text, add_special_tokens=False).ids + encs = self._tok.encode_batch(text, add_special_tokens=False) + return [e.ids for e in encs] + + def detokenize(self, token_ids, skip_special_tokens: bool = True) -> str: + return self.decode(token_ids, skip_special_tokens=skip_special_tokens) + + def decode(self, ids, skip_special_tokens: bool = True) -> str: + return self._tok.decode( + self.to_id_list(ids), skip_special_tokens=skip_special_tokens + ) + + def call(self, inputs: Union[str, List[str]], text_pair=None): + return self.encode_batch_to_inputs( + inputs, + text_pair, + token_type_ids=self.backbone_type + in ("bert", "roberta", "xlm-roberta", "albert"), + ) + + def get_config(self): + config = super().get_config() + config.update( + { + "variant": self.variant, + "tokenizer_file": self.tokenizer_file, + "max_seq_len": self.max_seq_len, + "backbone_type": self.backbone_type, + "bos_token": self.bos_token, + "eos_token": self.eos_token, + "unk_token": self.unk_token, + "pad_token": self.pad_token, + "mask_token": self.mask_token, + } + ) + return config From 977fd15d8bb0f532032bafd905491f1151e97334 Mon Sep 17 00:00:00 2001 From: IMvision12 Date: Tue, 13 Oct 2026 19:24:11 -0700 Subject: [PATCH 02/10] refactor --- docs/mpnet.md | 2 +- docs/sentence_transformer.md | 55 +- pyproject.toml | 1 + tests/base/model_test_registry.py | 77 +- tests/fixtures/cross_backend_parity.json | 1757 +++++++++++++++++ tests/integration/test_auto_registry.py | 14 +- .../integration/test_sentence_transformer.py | 63 +- zeromodels/__init__.py | 3 +- zeromodels/auto/auto_factory.py | 14 +- zeromodels/models/__init__.py | 1 - .../models/sentence_transformer/__init__.py | 40 +- zeromodels/st_models/__init__.py | 20 + .../sentence_transformer/__init__.py | 15 + ...onvert_sentence_transformer_hf_to_keras.py | 4 +- .../sentence_transformer_backbones.py | 2 +- .../sentence_transformer_config.py | 2 + .../sentence_transformer_layers.py | 2 +- .../sentence_transformer_model.py | 2 +- .../sentence_transformer_tokenizer.py | 2 +- 19 files changed, 2028 insertions(+), 48 deletions(-) create mode 100644 zeromodels/st_models/__init__.py create mode 100644 zeromodels/st_models/sentence_transformer/__init__.py rename zeromodels/{models => st_models}/sentence_transformer/convert_sentence_transformer_hf_to_keras.py (98%) rename zeromodels/{models => st_models}/sentence_transformer/sentence_transformer_backbones.py (98%) rename zeromodels/{models => st_models}/sentence_transformer/sentence_transformer_config.py (95%) rename zeromodels/{models => st_models}/sentence_transformer/sentence_transformer_layers.py (97%) rename zeromodels/{models => st_models}/sentence_transformer/sentence_transformer_model.py (99%) rename zeromodels/{models => st_models}/sentence_transformer/sentence_transformer_tokenizer.py (98%) diff --git a/docs/mpnet.md b/docs/mpnet.md index 6786c2e2..c84a5751 100644 --- a/docs/mpnet.md +++ b/docs/mpnet.md @@ -34,7 +34,7 @@ See also [bert.md](bert.md), [roberta.md](roberta.md), [modernbert.md](modernber MPNet is also the backbone of `sentence-transformers/all-mpnet-base-v2`, one of the most widely used sentence-embedding models — see [sentence_transformer.md](sentence_transformer.md) -for the Keras port of that full ST stack (mean pool + L2 normalize). +for the separate embedding API. ## Variants diff --git a/docs/sentence_transformer.md b/docs/sentence_transformer.md index 2f9da7b6..b592dae2 100644 --- a/docs/sentence_transformer.md +++ b/docs/sentence_transformer.md @@ -1,5 +1,15 @@ # Sentence Transformer +
+Weights: Pre-converted Keras weights for all 13 checkpoints are hosted under the +zeromodels organization. Load one with +from_weights("zeromodels/all-MiniLM-L6-v2"); no conversion is needed. +Original sentence-transformers checkpoints can also be converted on the fly +with the hf: prefix, e.g. +from_weights("hf:sentence-transformers/all-MiniLM-L6-v2"). +The backbone and embedding recipe are read from the checkpoint metadata. +
+ ZeroModels loads the 13 original dense sentence-embedding checkpoints shown in the [Sentence Transformers pretrained-model table](https://sbert.net/docs/sentence_transformer/pretrained_models.html). It builds a Keras 3 encoder and applies the checkpoint's declared pooling, @@ -7,12 +17,9 @@ optional Dense projection, and optional L2 normalization. Parity has been verified on the Keras Torch backend; TensorFlow and JAX parity have not yet been checked for all 13 checkpoints. -These are **upstream** Hugging Face repos: use the `hf:` prefix below. No -conversion or upload is needed for upstream loading. - -## Supported upstream checkpoints +## Available ZeroModels checkpoints -| Checkpoint (`sentence-transformers/`) | Actual backbone | Embedding recipe | Dim | +| Checkpoint (`zeromodels/`) | Actual backbone | Embedding recipe | Dim | |---|---|---|---:| | `all-mpnet-base-v2` | MPNet | mean → normalize | 768 | | `multi-qa-mpnet-base-dot-v1` | MPNet | CLS, dot-product similarity | 768 | @@ -38,16 +45,14 @@ import os os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" -from zeromodels.models.sentence_transformer import ( - SentenceTransformerModel, +from zeromodels.st_models import ( + SentenceTransformer, SentenceTransformerTokenizer, ) -repo = "sentence-transformers/all-MiniLM-L6-v2" -model = SentenceTransformerModel.from_weights(f"hf:{repo}") -tokenizer = SentenceTransformerTokenizer.from_weights( - f"hf:{repo}", revision=model.source_revision -) +repo = "zeromodels/all-MiniLM-L6-v2" +model = SentenceTransformer.from_weights(repo) +tokenizer = SentenceTransformerTokenizer.from_weights(repo) embeddings = model.encode( ["A man is eating food.", "A man is eating pasta."], @@ -58,15 +63,29 @@ print(model.similarity(embeddings[:1], embeddings[1:])) ``` `model(tokenizer(texts))` returns `embeddings` and `token_embeddings` tensors. -`encode()` batches raw strings and returns a NumPy array. Use the same -checkpoint and revision for the model and tokenizer. An explicit commit SHA can -be supplied as `revision=` to pin the upstream snapshot. `model.similarity()` -uses the checkpoint's declared cosine or dot-product scoring function. +`encode()` batches raw strings and returns a NumPy array. Use the same repo for +the model and tokenizer. `model.similarity()` uses the checkpoint's declared +cosine or dot-product scoring function. + +To load directly from an original Sentence Transformers repo instead, use the +`hf:` prefix for both components and pass the model's resolved revision to the +tokenizer: + +```python +upstream = "hf:sentence-transformers/all-MiniLM-L6-v2" +model = SentenceTransformer.from_weights(upstream) +tokenizer = SentenceTransformerTokenizer.from_weights( + upstream, revision=model.source_revision +) +``` + +An explicit commit SHA can be supplied as `revision=` to pin the upstream +snapshot. `AutoZMTextEmbed` resolves ZeroModels `sentence_transformer` configs to the `SentenceTransformerTextEmbed` task view. For upstream `hf:` repos, load the -concrete class as shown above: their -root `model_type` describes the backbone, not the sentence-embedding task. +concrete class as shown above: their root `model_type` describes the backbone, +not the sentence-embedding task. Unsupported upstream module types, combinations, or pooling modes raise a `ValueError` rather than silently changing the embedding recipe. Cross-encoders, diff --git a/pyproject.toml b/pyproject.toml index a405f764..d87a7b03 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -104,6 +104,7 @@ version = { attr = "zeromodels.__version__" } source = ["zeromodels", "tests"] omit = [ "zeromodels/models/*/convert_*.py", + "zeromodels/st_models/*/convert_*.py", ] [tool.coverage.report] diff --git a/tests/base/model_test_registry.py b/tests/base/model_test_registry.py index f20421ac..41703485 100644 --- a/tests/base/model_test_registry.py +++ b/tests/base/model_test_registry.py @@ -2998,7 +2998,7 @@ }, # ---- Sentence Transformers (MPNet backbone) ---- "SentenceTransformerModel": { - "module": "zeromodels.models.sentence_transformer", + "module": "zeromodels.st_models", "model_cls": "SentenceTransformerModel", "model_type": "llm", "init_kwargs": { @@ -3022,7 +3022,7 @@ }, }, "SentenceTransformerTextEmbed": { - "module": "zeromodels.models.sentence_transformer", + "module": "zeromodels.st_models", "model_cls": "SentenceTransformerTextEmbed", "model_type": "llm", "init_kwargs": { @@ -4434,6 +4434,73 @@ }, } +_SENTENCE_TRANSFORMER_VARIANTS = { + "all-mpnet-base-v2": ("mpnet", "mean", True, None, "cosine"), + "multi-qa-mpnet-base-dot-v1": ("mpnet", "cls", False, None, "dot"), + "all-distilroberta-v1": ("roberta", "mean", True, None, "cosine"), + "all-MiniLM-L12-v2": ("bert", "mean", True, None, "cosine"), + "multi-qa-distilbert-cos-v1": ("distilbert", "mean", True, None, "cosine"), + "all-MiniLM-L6-v2": ("bert", "mean", True, None, "cosine"), + "multi-qa-MiniLM-L6-cos-v1": ("bert", "mean", True, None, "cosine"), + "paraphrase-multilingual-mpnet-base-v2": ( + "xlm-roberta", + "mean", + False, + None, + "cosine", + ), + "paraphrase-albert-small-v2": ("albert", "mean", False, None, "cosine"), + "paraphrase-multilingual-MiniLM-L12-v2": ("bert", "mean", False, None, "cosine"), + "paraphrase-MiniLM-L3-v2": ("bert", "mean", False, None, "cosine"), + "distiluse-base-multilingual-cased-v1": ("distilbert", "mean", False, 8, "cosine"), + "distiluse-base-multilingual-cased-v2": ("distilbert", "mean", False, 8, "cosine"), +} +for _variant, ( + _backbone, + _pooling, + _normalize, + _projection, + _similarity, +) in _SENTENCE_TRANSFORMER_VARIANTS.items(): + _kwargs = { + "backbone_type": _backbone, + "vocab_size": 128, + "embed_dim": 16, + "num_layers": 1, + "num_heads": 2, + "mlp_dim": 32, + "max_position_embeddings": 64, + "relative_attention_num_buckets": 8, + "pooling": _pooling, + "normalize_embeddings": _normalize, + "projection_dim": _projection, + "similarity_fn_name": _similarity, + "max_seq_len": 32, + } + if _projection is not None: + _kwargs["projection_activation"] = "tanh" + if _backbone == "albert": + _kwargs["embedding_size"] = 8 + _factory = ( + "modernbert_input" + if _backbone in ("mpnet", "distilbert") + else "roberta_input" + if _backbone in ("roberta", "xlm-roberta") + else "bert_input" + ) + MODEL_TEST_CONFIGS[f"SentenceTransformer_{_variant.replace('-', '_')}"] = { + "module": "zeromodels.st_models", + "model_cls": "SentenceTransformerModel", + "model_type": "llm", + "init_kwargs": _kwargs, + "input_factory": _factory, + "input_factory_kwargs": {"seq_len": 16}, + "expected_output_shape": { + "embeddings": (2, _projection or 16), + "token_embeddings": (2, 16, 16), + }, + } + # The Qwen-VL `*Generate` classes share their base model's config + input but # add an LM head, so derive their test entries (logits output) from the bases. @@ -5002,8 +5069,8 @@ def instantiate_model(config): def get_cached_model(config): """Return a shared, build-once model for ``config``'s READ-ONLY tests. - Keyed by (class, active data_format, quantized) so a channels_first build or - a quantized variant never aliases the plain one. Use this only for tests that + Keyed by (registry entry, active data_format, quantized) so two entries using + the same class with different architectures never alias one another. Use this only for tests that do not mutate the model (forward / shape / NaN / ``get_config`` / ``save*``); a test that assigns weights, builds a different data format, or calls ``clear_session`` must call :func:`instantiate_model` for a fresh instance. @@ -5016,7 +5083,7 @@ def get_cached_model(config): import keras key = ( - config["model_cls"], + id(config), keras.config.image_data_format(), bool(config.get("quantization_config")), ) diff --git a/tests/fixtures/cross_backend_parity.json b/tests/fixtures/cross_backend_parity.json index 9cd433b3..5f98aa08 100644 --- a/tests/fixtures/cross_backend_parity.json +++ b/tests/fixtures/cross_backend_parity.json @@ -24042,6 +24042,1763 @@ ] } ], + "SentenceTransformerModel": [ + { + "shape": [ + 2, + 32 + ], + "sample": [ + -0.138292, + 0.107383, + 0.227017, + -0.157133, + 0.043625, + -0.147396, + 0.077896, + 0.145626, + 0.28461, + -0.071275, + 0.122478, + -0.137024, + -0.125705, + -0.16926, + -0.036518, + 0.066965, + 0.107061, + -0.087836, + 0.160158, + -0.163558, + -0.408743, + 0.039037, + 0.045357, + -0.055598, + -0.408058, + 0.181733, + -0.026227, + 0.110355, + 0.224607, + 0.15685, + -0.357952, + -0.10573, + -0.136634, + 0.107497, + 0.226582, + -0.15824, + 0.043457, + -0.147285, + 0.080659, + 0.144457, + 0.282996, + -0.068748, + 0.124357, + -0.137633, + -0.122886, + -0.167857, + -0.037022, + 0.067155, + 0.108349, + -0.086218, + 0.160484, + -0.164189, + -0.406857, + 0.036848, + 0.046163, + -0.051938, + -0.41074, + 0.180494, + -0.025904, + 0.110552, + 0.223734, + 0.157155, + -0.360933, + -0.106077 + ] + }, + { + "shape": [ + 2, + 16, + 32 + ], + "sample": [ + -0.020572, + 0.015957, + -0.019703, + 0.015892, + -0.020483, + -0.01368, + 0.016936, + -0.015326, + 0.014478, + 0.023795, + 0.033209, + 0.02408, + 0.033903, + -0.024844, + -0.023571, + -0.024457, + -0.02343, + -0.059766, + 0.006253, + -0.062119, + 0.006618, + 0.005499, + -0.020619, + 0.005742, + -0.023827, + 0.006142, + 0.009706, + 0.006107, + 0.008017, + 0.00602, + 0.021343, + -0.013266, + 0.015743, + -0.010255, + 0.043917, + -0.061277, + 0.042983, + -0.062477, + -0.010429, + 0.026378, + -0.009833, + 0.02688, + 0.019112, + -0.004345, + 0.019161, + -0.004257, + 0.018308, + 0.015562, + -0.020819, + 0.01682, + -0.02066, + 0.032244, + -0.018348, + 0.033819, + -0.017939, + 0.02347, + -0.025788, + 0.023409, + -0.025221, + -0.053642, + -0.004411, + -0.055167, + -0.005387, + -0.015949 + ] + } + ], + "SentenceTransformerTextEmbed": [ + { + "shape": [ + 2, + 32 + ], + "sample": [ + -0.138292, + 0.107383, + 0.227017, + -0.157133, + 0.043625, + -0.147396, + 0.077896, + 0.145626, + 0.28461, + -0.071275, + 0.122478, + -0.137024, + -0.125705, + -0.16926, + -0.036518, + 0.066965, + 0.107061, + -0.087836, + 0.160158, + -0.163558, + -0.408743, + 0.039037, + 0.045357, + -0.055598, + -0.408058, + 0.181733, + -0.026227, + 0.110355, + 0.224607, + 0.15685, + -0.357952, + -0.10573, + -0.136634, + 0.107497, + 0.226582, + -0.15824, + 0.043457, + -0.147285, + 0.080659, + 0.144457, + 0.282996, + -0.068748, + 0.124357, + -0.137633, + -0.122886, + -0.167857, + -0.037022, + 0.067155, + 0.108349, + -0.086218, + 0.160484, + -0.164189, + -0.406857, + 0.036848, + 0.046163, + -0.051938, + -0.41074, + 0.180494, + -0.025904, + 0.110552, + 0.223734, + 0.157155, + -0.360933, + -0.106077 + ] + }, + { + "shape": [ + 2, + 16, + 32 + ], + "sample": [ + -0.020572, + 0.015957, + -0.019703, + 0.015892, + -0.020483, + -0.01368, + 0.016936, + -0.015326, + 0.014478, + 0.023795, + 0.033209, + 0.02408, + 0.033903, + -0.024844, + -0.023571, + -0.024457, + -0.02343, + -0.059766, + 0.006253, + -0.062119, + 0.006618, + 0.005499, + -0.020619, + 0.005742, + -0.023827, + 0.006142, + 0.009706, + 0.006107, + 0.008017, + 0.00602, + 0.021343, + -0.013266, + 0.015743, + -0.010255, + 0.043917, + -0.061277, + 0.042983, + -0.062477, + -0.010429, + 0.026378, + -0.009833, + 0.02688, + 0.019112, + -0.004345, + 0.019161, + -0.004257, + 0.018308, + 0.015562, + -0.020819, + 0.01682, + -0.02066, + 0.032244, + -0.018348, + 0.033819, + -0.017939, + 0.02347, + -0.025788, + 0.023409, + -0.025221, + -0.053642, + -0.004411, + -0.055167, + -0.005387, + -0.015949 + ] + } + ], + "SentenceTransformer_all_MiniLM_L12_v2": [ + { + "shape": [ + 2, + 16 + ], + "sample": [ + -0.006283, + -0.28829, + -0.510357, + -0.181786, + 0.267572, + -0.204303, + 0.065134, + -0.036725, + 0.221611, + -0.276, + 0.044667, + -0.319701, + -0.356479, + 0.162738, + 0.068047, + -0.341659, + -0.007111, + -0.286077, + -0.509385, + -0.181498, + 0.266028, + -0.204473, + 0.059981, + -0.032958, + 0.218527, + -0.278198, + 0.045874, + -0.318675, + -0.36582, + 0.169479, + 0.064059, + -0.336116 + ] + }, + { + "shape": [ + 2, + 16, + 16 + ], + "sample": [ + -0.001388, + 0.02555, + 0.001283, + 0.024534, + 0.000669, + 0.0236, + -0.000574, + 0.024204, + -9.2e-05, + -0.031526, + -0.032391, + -0.030826, + -0.032146, + -0.031856, + -0.032201, + -0.030922, + -0.032239, + -0.032286, + -0.05965, + 0.004568, + -0.064821, + 0.002639, + -0.055885, + 0.007288, + -0.060506, + 0.002031, + -0.053601, + -0.035833, + -0.021838, + -0.035347, + -0.020152, + -0.036121, + -0.020518, + -0.036075, + -0.020536, + -0.036254, + 0.030103, + -0.039094, + 0.030121, + -0.035787, + 0.030136, + -0.047276, + 0.030123, + -0.041288, + 0.030123, + 0.014357, + -0.020862, + 0.01455, + -0.023323, + 0.019685, + -0.02363, + 0.00853, + -0.020957, + 0.031693, + 0.003902, + 0.005723, + 0.007174, + 0.006236, + 0.003883, + 0.006472, + 0.015475, + 0.010045, + -0.001431, + -0.036287 + ] + } + ], + "SentenceTransformer_all_MiniLM_L6_v2": [ + { + "shape": [ + 2, + 16 + ], + "sample": [ + -0.006283, + -0.28829, + -0.510357, + -0.181786, + 0.267572, + -0.204303, + 0.065134, + -0.036725, + 0.221611, + -0.276, + 0.044667, + -0.319701, + -0.356479, + 0.162738, + 0.068047, + -0.341659, + -0.007111, + -0.286077, + -0.509385, + -0.181498, + 0.266028, + -0.204473, + 0.059981, + -0.032958, + 0.218527, + -0.278198, + 0.045874, + -0.318675, + -0.36582, + 0.169479, + 0.064059, + -0.336116 + ] + }, + { + "shape": [ + 2, + 16, + 16 + ], + "sample": [ + -0.001388, + 0.02555, + 0.001283, + 0.024534, + 0.000669, + 0.0236, + -0.000574, + 0.024204, + -9.2e-05, + -0.031526, + -0.032391, + -0.030826, + -0.032146, + -0.031856, + -0.032201, + -0.030922, + -0.032239, + -0.032286, + -0.05965, + 0.004568, + -0.064821, + 0.002639, + -0.055885, + 0.007288, + -0.060506, + 0.002031, + -0.053601, + -0.035833, + -0.021838, + -0.035347, + -0.020152, + -0.036121, + -0.020518, + -0.036075, + -0.020536, + -0.036254, + 0.030103, + -0.039094, + 0.030121, + -0.035787, + 0.030136, + -0.047276, + 0.030123, + -0.041288, + 0.030123, + 0.014357, + -0.020862, + 0.01455, + -0.023323, + 0.019685, + -0.02363, + 0.00853, + -0.020957, + 0.031693, + 0.003902, + 0.005723, + 0.007174, + 0.006236, + 0.003883, + 0.006472, + 0.015475, + 0.010045, + -0.001431, + -0.036287 + ] + } + ], + "SentenceTransformer_all_distilroberta_v1": [ + { + "shape": [ + 2, + 16 + ], + "sample": [ + -0.006428, + -0.288801, + -0.508576, + -0.181893, + 0.268083, + -0.204161, + 0.058898, + -0.038747, + 0.221928, + -0.27647, + 0.045056, + -0.320287, + -0.360028, + 0.155845, + 0.068021, + -0.342725, + -0.003668, + -0.288044, + -0.509014, + -0.183166, + 0.268015, + -0.205943, + 0.054317, + -0.032845, + 0.218373, + -0.281978, + 0.046054, + -0.321448, + -0.358621, + 0.167182, + 0.063274, + -0.335987 + ] + }, + { + "shape": [ + 2, + 16, + 16 + ], + "sample": [ + 0.00054, + 0.024421, + 0.00022, + 0.024225, + 0.000921, + 0.023616, + -0.00275, + 0.026056, + 8.3e-05, + -0.030825, + -0.032213, + -0.031122, + -0.032074, + -0.0312, + -0.032084, + -0.033028, + -0.032695, + -0.030283, + -0.064057, + 0.005521, + -0.059847, + 0.003993, + -0.058114, + 0.005001, + -0.053613, + 0.000932, + -0.055822, + -0.035447, + -0.021108, + -0.036471, + -0.020006, + -0.035757, + -0.020632, + -0.036374, + -0.020734, + -0.03649, + 0.030099, + -0.045783, + 0.030133, + -0.03234, + 0.030137, + -0.041763, + 0.030125, + -0.045943, + 0.030106, + 0.020362, + -0.021803, + 0.011972, + -0.022738, + 0.021537, + -0.023277, + 0.009515, + -0.022329, + 0.040238, + 0.014673, + 0.006805, + 0.01218, + 0.008585, + -0.000855, + 0.006342, + 0.002824, + 0.008836, + -0.000799, + -0.036171 + ] + } + ], + "SentenceTransformer_all_mpnet_base_v2": [ + { + "shape": [ + 2, + 16 + ], + "sample": [ + 0.131721, + -0.261829, + -0.385196, + -0.173886, + 0.281246, + 0.001545, + 0.245415, + -0.072511, + -0.050302, + -0.224904, + 0.388975, + -0.262667, + -0.358547, + 0.329789, + 0.2014, + -0.199019, + 0.130397, + -0.259231, + -0.389688, + -0.17159, + 0.279488, + 0.000753, + 0.247705, + -0.083166, + -0.051451, + -0.230933, + 0.382909, + -0.258658, + -0.349881, + 0.325492, + 0.221085, + -0.203389 + ] + }, + { + "shape": [ + 2, + 16, + 16 + ], + "sample": [ + 0.012509, + -0.007654, + 0.012005, + -0.005402, + 0.012474, + -0.005269, + 0.013393, + -0.00231, + 0.01668, + -0.023257, + -0.027662, + -0.027021, + -0.028156, + -0.025185, + -0.027928, + -0.026351, + -0.02774, + -0.028348, + -0.031596, + 0.042527, + -0.038723, + 0.039582, + -0.045372, + 0.041086, + -0.044355, + 0.041554, + -0.03575, + -0.027495, + -0.018154, + -0.024365, + -0.019006, + -0.026141, + -0.018336, + -0.038645, + -0.018094, + -0.020396, + 0.030104, + -0.045994, + 0.030238, + -0.01341, + 0.030433, + -0.053719, + 0.03042, + -0.077768, + 0.03017, + 0.034797, + -0.002746, + 0.033777, + -0.006548, + 0.037556, + 0.012106, + 0.030717, + -0.003646, + 0.039079, + 0.038353, + 0.012978, + 0.040298, + 0.028337, + 0.011519, + 0.025773, + 0.027593, + 0.011926, + 0.014531, + -0.022591 + ] + } + ], + "SentenceTransformer_distiluse_base_multilingual_cased_v1": [ + { + "shape": [ + 2, + 8 + ], + "sample": [ + -0.005696, + 0.011375, + -0.000598, + -0.005064, + -0.008866, + 0.021272, + 0.007137, + -0.015385, + -0.005726, + 0.011304, + -0.000321, + -0.005126, + -0.008742, + 0.021266, + 0.00725, + -0.015222 + ] + }, + { + "shape": [ + 2, + 16, + 16 + ], + "sample": [ + 0.00223, + -0.007174, + 0.002873, + -0.009289, + 0.003751, + -0.007166, + 0.003752, + -0.007555, + 0.001998, + 0.011331, + -0.02211, + -0.001803, + -0.023278, + 0.00989, + -0.0181, + -0.00481, + -0.031739, + 0.000239, + 0.015857, + 0.06128, + 0.015085, + 0.06705, + 0.017578, + 0.05266, + 0.013209, + 0.060971, + 0.015753, + 0.049514, + -0.001151, + 0.049367, + -0.001122, + 0.033595, + -0.000648, + 0.041646, + -0.000901, + 0.048977, + -0.010922, + -0.039426, + -0.004766, + -0.042949, + -0.007115, + -0.037974, + -0.003507, + -0.042364, + -0.006889, + -0.005795, + -0.003808, + -0.007997, + -0.003384, + 0.000805, + -0.004321, + -0.008732, + -0.005165, + -0.006883, + -0.002073, + 0.050016, + -0.005411, + 0.052228, + -0.006485, + 0.053975, + -0.00512, + 0.0522, + -0.004789, + 0.017832 + ] + } + ], + "SentenceTransformer_distiluse_base_multilingual_cased_v2": [ + { + "shape": [ + 2, + 8 + ], + "sample": [ + -0.005696, + 0.011375, + -0.000598, + -0.005064, + -0.008866, + 0.021272, + 0.007137, + -0.015385, + -0.005726, + 0.011304, + -0.000321, + -0.005126, + -0.008742, + 0.021266, + 0.00725, + -0.015222 + ] + }, + { + "shape": [ + 2, + 16, + 16 + ], + "sample": [ + 0.00223, + -0.007174, + 0.002873, + -0.009289, + 0.003751, + -0.007166, + 0.003752, + -0.007555, + 0.001998, + 0.011331, + -0.02211, + -0.001803, + -0.023278, + 0.00989, + -0.0181, + -0.00481, + -0.031739, + 0.000239, + 0.015857, + 0.06128, + 0.015085, + 0.06705, + 0.017578, + 0.05266, + 0.013209, + 0.060971, + 0.015753, + 0.049514, + -0.001151, + 0.049367, + -0.001122, + 0.033595, + -0.000648, + 0.041646, + -0.000901, + 0.048977, + -0.010922, + -0.039426, + -0.004766, + -0.042949, + -0.007115, + -0.037974, + -0.003507, + -0.042364, + -0.006889, + -0.005795, + -0.003808, + -0.007997, + -0.003384, + 0.000805, + -0.004321, + -0.008732, + -0.005165, + -0.006883, + -0.002073, + 0.050016, + -0.005411, + 0.052228, + -0.006485, + 0.053975, + -0.00512, + 0.0522, + -0.004789, + 0.017832 + ] + } + ], + "SentenceTransformer_multi_qa_MiniLM_L6_cos_v1": [ + { + "shape": [ + 2, + 16 + ], + "sample": [ + -0.006283, + -0.28829, + -0.510357, + -0.181786, + 0.267572, + -0.204303, + 0.065134, + -0.036725, + 0.221611, + -0.276, + 0.044667, + -0.319701, + -0.356479, + 0.162738, + 0.068047, + -0.341659, + -0.007111, + -0.286077, + -0.509385, + -0.181498, + 0.266028, + -0.204473, + 0.059981, + -0.032958, + 0.218527, + -0.278198, + 0.045874, + -0.318675, + -0.36582, + 0.169479, + 0.064059, + -0.336116 + ] + }, + { + "shape": [ + 2, + 16, + 16 + ], + "sample": [ + -0.001388, + 0.02555, + 0.001283, + 0.024534, + 0.000669, + 0.0236, + -0.000574, + 0.024204, + -9.2e-05, + -0.031526, + -0.032391, + -0.030826, + -0.032146, + -0.031856, + -0.032201, + -0.030922, + -0.032239, + -0.032286, + -0.05965, + 0.004568, + -0.064821, + 0.002639, + -0.055885, + 0.007288, + -0.060506, + 0.002031, + -0.053601, + -0.035833, + -0.021838, + -0.035347, + -0.020152, + -0.036121, + -0.020518, + -0.036075, + -0.020536, + -0.036254, + 0.030103, + -0.039094, + 0.030121, + -0.035787, + 0.030136, + -0.047276, + 0.030123, + -0.041288, + 0.030123, + 0.014357, + -0.020862, + 0.01455, + -0.023323, + 0.019685, + -0.02363, + 0.00853, + -0.020957, + 0.031693, + 0.003902, + 0.005723, + 0.007174, + 0.006236, + 0.003883, + 0.006472, + 0.015475, + 0.010045, + -0.001431, + -0.036287 + ] + } + ], + "SentenceTransformer_multi_qa_distilbert_cos_v1": [ + { + "shape": [ + 2, + 16 + ], + "sample": [ + 0.025973, + -0.206354, + 0.137715, + -0.007853, + -0.064003, + -0.047599, + -0.033647, + 0.035637, + -0.075466, + 0.007383, + 0.548441, + 0.416146, + -0.376682, + -0.05115, + 0.495402, + 0.243534, + 0.025587, + -0.19556, + 0.143738, + -0.007995, + -0.066312, + -0.045355, + -0.03309, + 0.023793, + -0.075691, + -0.024764, + 0.571854, + 0.411226, + -0.383017, + -0.0612, + 0.497094, + 0.179573 + ] + }, + { + "shape": [ + 2, + 16, + 16 + ], + "sample": [ + 0.00223, + -0.007174, + 0.002873, + -0.009289, + 0.003751, + -0.007166, + 0.003752, + -0.007555, + 0.001998, + 0.011331, + -0.02211, + -0.001803, + -0.023278, + 0.00989, + -0.0181, + -0.00481, + -0.031739, + 0.000239, + 0.015857, + 0.06128, + 0.015085, + 0.06705, + 0.017578, + 0.05266, + 0.013209, + 0.060971, + 0.015753, + 0.049514, + -0.001151, + 0.049367, + -0.001122, + 0.033595, + -0.000648, + 0.041646, + -0.000901, + 0.048977, + -0.010922, + -0.039426, + -0.004766, + -0.042949, + -0.007115, + -0.037974, + -0.003507, + -0.042364, + -0.006889, + -0.005795, + -0.003808, + -0.007997, + -0.003384, + 0.000805, + -0.004321, + -0.008732, + -0.005165, + -0.006883, + -0.002073, + 0.050016, + -0.005411, + 0.052228, + -0.006485, + 0.053975, + -0.00512, + 0.0522, + -0.004789, + 0.017832 + ] + } + ], + "SentenceTransformer_multi_qa_mpnet_base_dot_v1": [ + { + "shape": [ + 2, + 16 + ], + "sample": [ + 0.012509, + -0.027962, + -0.041582, + -0.018233, + 0.030212, + -0.002694, + 0.010153, + 0.004949, + -0.007654, + -0.025851, + 0.042763, + -0.02836, + -0.03605, + 0.034484, + 0.022787, + -0.014465, + 0.007949, + -0.027739, + -0.047459, + -0.018336, + 0.030239, + -0.001236, + 0.020397, + -0.011969, + -0.007566, + -0.027303, + 0.041396, + -0.038645, + -0.02535, + 0.034228, + 0.041315, + -0.024879 + ] + }, + { + "shape": [ + 2, + 16, + 16 + ], + "sample": [ + 0.012509, + -0.007654, + 0.012005, + -0.005402, + 0.012474, + -0.005269, + 0.013393, + -0.00231, + 0.01668, + -0.023257, + -0.027662, + -0.027021, + -0.028156, + -0.025185, + -0.027928, + -0.026351, + -0.02774, + -0.028348, + -0.031596, + 0.042527, + -0.038723, + 0.039582, + -0.045372, + 0.041086, + -0.044355, + 0.041554, + -0.03575, + -0.027495, + -0.018154, + -0.024365, + -0.019006, + -0.026141, + -0.018336, + -0.038645, + -0.018094, + -0.020396, + 0.030104, + -0.045994, + 0.030238, + -0.01341, + 0.030433, + -0.053719, + 0.03042, + -0.077768, + 0.03017, + 0.034797, + -0.002746, + 0.033777, + -0.006548, + 0.037556, + 0.012106, + 0.030717, + -0.003646, + 0.039079, + 0.038353, + 0.012978, + 0.040298, + 0.028337, + 0.011519, + 0.025773, + 0.027593, + 0.011926, + 0.014531, + -0.022591 + ] + } + ], + "SentenceTransformer_paraphrase_MiniLM_L3_v2": [ + { + "shape": [ + 2, + 16 + ], + "sample": [ + -0.000708, + -0.032463, + -0.05747, + -0.02047, + 0.03013, + -0.023006, + 0.007335, + -0.004135, + 0.024955, + -0.03108, + 0.00503, + -0.036001, + -0.040142, + 0.018325, + 0.007663, + -0.038473, + -0.000805, + -0.032396, + -0.057684, + -0.020553, + 0.030126, + -0.023155, + 0.006792, + -0.003732, + 0.024747, + -0.031504, + 0.005195, + -0.036088, + -0.041427, + 0.019192, + 0.007254, + -0.038063 + ] + }, + { + "shape": [ + 2, + 16, + 16 + ], + "sample": [ + -0.001388, + 0.02555, + 0.001283, + 0.024534, + 0.000669, + 0.0236, + -0.000574, + 0.024204, + -9.2e-05, + -0.031526, + -0.032391, + -0.030826, + -0.032146, + -0.031856, + -0.032201, + -0.030922, + -0.032239, + -0.032286, + -0.05965, + 0.004568, + -0.064821, + 0.002639, + -0.055885, + 0.007288, + -0.060506, + 0.002031, + -0.053601, + -0.035833, + -0.021838, + -0.035347, + -0.020152, + -0.036121, + -0.020518, + -0.036075, + -0.020536, + -0.036254, + 0.030103, + -0.039094, + 0.030121, + -0.035787, + 0.030136, + -0.047276, + 0.030123, + -0.041288, + 0.030123, + 0.014357, + -0.020862, + 0.01455, + -0.023323, + 0.019685, + -0.02363, + 0.00853, + -0.020957, + 0.031693, + 0.003902, + 0.005723, + 0.007174, + 0.006236, + 0.003883, + 0.006472, + 0.015475, + 0.010045, + -0.001431, + -0.036287 + ] + } + ], + "SentenceTransformer_paraphrase_albert_small_v2": [ + { + "shape": [ + 2, + 16 + ], + "sample": [ + 0.019625, + -0.000452, + -0.065652, + 0.038684, + -0.016475, + -0.077698, + 0.066542, + 0.033501, + -0.000845, + 0.026137, + -0.01556, + 0.010658, + -0.019583, + -0.014868, + -0.019154, + 0.056296, + 0.019622, + -0.0004, + -0.065506, + 0.038478, + -0.016457, + -0.077675, + 0.066516, + 0.033927, + -0.000842, + 0.026165, + -0.015646, + 0.010543, + -0.019583, + -0.014792, + -0.018992, + 0.055857 + ] + }, + { + "shape": [ + 2, + 16, + 16 + ], + "sample": [ + 0.019678, + -0.000842, + 0.019711, + -0.000838, + 0.01957, + -0.000845, + 0.019591, + -0.000842, + 0.019633, + 0.026074, + -0.000336, + 0.026192, + -0.000561, + 0.026158, + -0.000532, + 0.026085, + -0.000782, + 0.026046, + -0.066363, + -0.015558, + -0.065007, + -0.015725, + -0.066967, + -0.015305, + -0.065715, + -0.015685, + -0.064612, + 0.0103, + 0.038897, + 0.010929, + 0.039194, + 0.010795, + 0.038024, + 0.010261, + 0.03857, + 0.010324, + -0.016849, + -0.019532, + -0.016213, + -0.019658, + -0.016232, + -0.019634, + -0.017388, + -0.019531, + -0.015906, + -0.015022, + -0.077507, + -0.014948, + -0.077572, + -0.014514, + -0.07773, + -0.014773, + -0.077588, + -0.014602, + 0.066499, + -0.019658, + 0.066418, + -0.018982, + 0.066457, + -0.018579, + 0.066452, + -0.018653, + 0.066585, + 0.057405 + ] + } + ], + "SentenceTransformer_paraphrase_multilingual_MiniLM_L12_v2": [ + { + "shape": [ + 2, + 16 + ], + "sample": [ + -0.000708, + -0.032463, + -0.05747, + -0.02047, + 0.03013, + -0.023006, + 0.007335, + -0.004135, + 0.024955, + -0.03108, + 0.00503, + -0.036001, + -0.040142, + 0.018325, + 0.007663, + -0.038473, + -0.000805, + -0.032396, + -0.057684, + -0.020553, + 0.030126, + -0.023155, + 0.006792, + -0.003732, + 0.024747, + -0.031504, + 0.005195, + -0.036088, + -0.041427, + 0.019192, + 0.007254, + -0.038063 + ] + }, + { + "shape": [ + 2, + 16, + 16 + ], + "sample": [ + -0.001388, + 0.02555, + 0.001283, + 0.024534, + 0.000669, + 0.0236, + -0.000574, + 0.024204, + -9.2e-05, + -0.031526, + -0.032391, + -0.030826, + -0.032146, + -0.031856, + -0.032201, + -0.030922, + -0.032239, + -0.032286, + -0.05965, + 0.004568, + -0.064821, + 0.002639, + -0.055885, + 0.007288, + -0.060506, + 0.002031, + -0.053601, + -0.035833, + -0.021838, + -0.035347, + -0.020152, + -0.036121, + -0.020518, + -0.036075, + -0.020536, + -0.036254, + 0.030103, + -0.039094, + 0.030121, + -0.035787, + 0.030136, + -0.047276, + 0.030123, + -0.041288, + 0.030123, + 0.014357, + -0.020862, + 0.01455, + -0.023323, + 0.019685, + -0.02363, + 0.00853, + -0.020957, + 0.031693, + 0.003902, + 0.005723, + 0.007174, + 0.006236, + 0.003883, + 0.006472, + 0.015475, + 0.010045, + -0.001431, + -0.036287 + ] + } + ], + "SentenceTransformer_paraphrase_multilingual_mpnet_base_v2": [ + { + "shape": [ + 2, + 16 + ], + "sample": [ + -0.000722, + -0.032459, + -0.05716, + -0.020443, + 0.03013, + -0.022946, + 0.00662, + -0.004355, + 0.024943, + -0.031073, + 0.005064, + -0.035998, + -0.040464, + 0.017516, + 0.007645, + -0.038519, + -0.000412, + -0.032375, + -0.057211, + -0.020587, + 0.030124, + -0.023147, + 0.006105, + -0.003692, + 0.024544, + -0.031693, + 0.005176, + -0.036129, + -0.040307, + 0.01879, + 0.007112, + -0.037763 + ] + }, + { + "shape": [ + 2, + 16, + 16 + ], + "sample": [ + 0.00054, + 0.024421, + 0.00022, + 0.024225, + 0.000921, + 0.023616, + -0.00275, + 0.026056, + 8.3e-05, + -0.030825, + -0.032213, + -0.031122, + -0.032074, + -0.0312, + -0.032084, + -0.033028, + -0.032695, + -0.030283, + -0.064057, + 0.005521, + -0.059847, + 0.003993, + -0.058114, + 0.005001, + -0.053613, + 0.000932, + -0.055822, + -0.035447, + -0.021108, + -0.036471, + -0.020006, + -0.035757, + -0.020632, + -0.036374, + -0.020734, + -0.03649, + 0.030099, + -0.045783, + 0.030133, + -0.03234, + 0.030137, + -0.041763, + 0.030125, + -0.045943, + 0.030106, + 0.020362, + -0.021803, + 0.011972, + -0.022738, + 0.021537, + -0.023277, + 0.009515, + -0.022329, + 0.040238, + 0.014673, + 0.006805, + 0.01218, + 0.008585, + -0.000855, + 0.006342, + 0.002824, + 0.008836, + -0.000799, + -0.036171 + ] + } + ], "SigLIP2ImageClassify": [ { "shape": [ diff --git a/tests/integration/test_auto_registry.py b/tests/integration/test_auto_registry.py index c1d92066..01011d23 100644 --- a/tests/integration/test_auto_registry.py +++ b/tests/integration/test_auto_registry.py @@ -85,14 +85,16 @@ def _hf_types(cls): def _iter_model_classes(): import zeromodels.models as models + import zeromodels.st_models as st_models from zeromodels.base.base_model import BaseModel - for family in sorted(n for n in dir(models) if not n.startswith("_")): - package = getattr(models, family) - for name in getattr(package, "__all__", []): - obj = getattr(package, name, None) - if inspect.isclass(obj) and issubclass(obj, BaseModel): - yield name, obj + for namespace in (models, st_models): + for family in sorted(n for n in dir(namespace) if not n.startswith("_")): + package = getattr(namespace, family) + for name in getattr(package, "__all__", []): + obj = getattr(package, name, None) + if inspect.isclass(obj) and issubclass(obj, BaseModel): + yield name, obj def test_top_level_and_task_autos_exported(): diff --git a/tests/integration/test_sentence_transformer.py b/tests/integration/test_sentence_transformer.py index f272f349..bb454100 100644 --- a/tests/integration/test_sentence_transformer.py +++ b/tests/integration/test_sentence_transformer.py @@ -1,17 +1,78 @@ """Offline contract tests for the multi-backbone sentence-embedding API.""" +import importlib import json import numpy as np import pytest from keras import ops -from zeromodels.models.sentence_transformer import ( +from tests.base.model_test_registry import MODEL_TEST_CONFIGS +from zeromodels.st_models import ( + SentenceTransformer, SentenceTransformerModel, SentenceTransformerTextEmbed, ) +def test_top_level_st_api_and_legacy_imports_resolve_to_same_classes(): + legacy = importlib.import_module("zeromodels.models.sentence_transformer") + legacy_model_module = importlib.import_module( + "zeromodels.models.sentence_transformer.sentence_transformer_model" + ) + assert SentenceTransformer is SentenceTransformerModel + assert legacy.SentenceTransformerModel is SentenceTransformerModel + assert legacy_model_module.SentenceTransformerModel is SentenceTransformerModel + + +def test_existing_hub_config_with_old_module_path_still_builds(monkeypatch): + from zeromodels.conversion import zm_config + + spec = { + "model_module": "zeromodels.models.sentence_transformer", + "model_class": "SentenceTransformerModel", + "model_type": "sentence_transformer", + "backbone_type": "bert", + "vocab_size": 32, + "embed_dim": 8, + "num_layers": 1, + "num_heads": 2, + "mlp_dim": 16, + "max_position_embeddings": 16, + } + monkeypatch.setattr(zm_config, "load_zm_config", lambda repo_id: spec) + + model = SentenceTransformer.from_hub_repo( + "zeromodels/example-sentence-transformer", load_weights=False + ) + assert isinstance(model, SentenceTransformerModel) + assert model.config.backbone_type == "bert" + + +def test_all_published_sentence_transformer_variants_are_in_offline_registry(): + published = { + "all-mpnet-base-v2", + "multi-qa-mpnet-base-dot-v1", + "all-distilroberta-v1", + "all-MiniLM-L12-v2", + "multi-qa-distilbert-cos-v1", + "all-MiniLM-L6-v2", + "multi-qa-MiniLM-L6-cos-v1", + "paraphrase-multilingual-mpnet-base-v2", + "paraphrase-albert-small-v2", + "paraphrase-multilingual-MiniLM-L12-v2", + "paraphrase-MiniLM-L3-v2", + "distiluse-base-multilingual-cased-v1", + "distiluse-base-multilingual-cased-v2", + } + registered = { + name.removeprefix("SentenceTransformer_") + for name in MODEL_TEST_CONFIGS + if name.startswith("SentenceTransformer_") + } + assert registered == {name.replace("-", "_") for name in published} + + @pytest.mark.parametrize( "backbone_type", ("mpnet", "bert", "roberta", "xlm-roberta", "distilbert", "albert"), diff --git a/zeromodels/__init__.py b/zeromodels/__init__.py index 3a5387a8..0069ac5d 100644 --- a/zeromodels/__init__.py +++ b/zeromodels/__init__.py @@ -1,4 +1,4 @@ -from zeromodels import auto, conversion, models, samplers, utils +from zeromodels import auto, conversion, models, samplers, st_models, utils from zeromodels.auto import ( AutoZMConfig, AutoZMImageProcessor, @@ -14,6 +14,7 @@ "auto", "models", "samplers", + "st_models", "utils", "conversion", "version", diff --git a/zeromodels/auto/auto_factory.py b/zeromodels/auto/auto_factory.py index e733f338..328e6b2e 100644 --- a/zeromodels/auto/auto_factory.py +++ b/zeromodels/auto/auto_factory.py @@ -104,14 +104,16 @@ def _class_index(): global _CLASS_INDEX if _CLASS_INDEX is None: import zeromodels.models as models + import zeromodels.st_models as st_models index = {} - for family in sorted(n for n in dir(models) if not n.startswith("_")): - package = getattr(models, family) - for name in getattr(package, "__all__", []): - obj = getattr(package, name, None) - if inspect.isclass(obj): - index[name] = obj + for namespace in (models, st_models): + for family in sorted(n for n in dir(namespace) if not n.startswith("_")): + package = getattr(namespace, family) + for name in getattr(package, "__all__", []): + obj = getattr(package, name, None) + if inspect.isclass(obj): + index[name] = obj _CLASS_INDEX = index return _CLASS_INDEX diff --git a/zeromodels/models/__init__.py b/zeromodels/models/__init__.py index 5c282d74..74301652 100644 --- a/zeromodels/models/__init__.py +++ b/zeromodels/models/__init__.py @@ -121,7 +121,6 @@ sam3, segformer, senet, - sentence_transformer, siglip, siglip2, speech2text, diff --git a/zeromodels/models/sentence_transformer/__init__.py b/zeromodels/models/sentence_transformer/__init__.py index a9beb6a8..48653db8 100644 --- a/zeromodels/models/sentence_transformer/__init__.py +++ b/zeromodels/models/sentence_transformer/__init__.py @@ -1,9 +1,43 @@ -from .sentence_transformer_config import SentenceTransformerConfig -from .sentence_transformer_model import ( +"""Compatibility imports for Sentence Transformer models moved to ``st_models``.""" + +import sys + +from zeromodels.st_models import ( + SentenceTransformerConfig, SentenceTransformerModel, SentenceTransformerTextEmbed, + SentenceTransformerTokenizer, +) +from zeromodels.st_models.sentence_transformer import ( + convert_sentence_transformer_hf_to_keras as _converter, +) +from zeromodels.st_models.sentence_transformer import ( + sentence_transformer_backbones as _backbones, +) +from zeromodels.st_models.sentence_transformer import ( + sentence_transformer_config as _config, ) -from .sentence_transformer_tokenizer import SentenceTransformerTokenizer +from zeromodels.st_models.sentence_transformer import ( + sentence_transformer_layers as _layers, +) +from zeromodels.st_models.sentence_transformer import ( + sentence_transformer_model as _model, +) +from zeromodels.st_models.sentence_transformer import ( + sentence_transformer_tokenizer as _tokenizer, +) + +# Old deep import paths may appear in user code and saved Keras objects. Keep +# them importable without retaining a second copy of the implementation. +for _name, _module in { + "convert_sentence_transformer_hf_to_keras": _converter, + "sentence_transformer_backbones": _backbones, + "sentence_transformer_config": _config, + "sentence_transformer_layers": _layers, + "sentence_transformer_model": _model, + "sentence_transformer_tokenizer": _tokenizer, +}.items(): + sys.modules[f"{__name__}.{_name}"] = _module __all__ = [ "SentenceTransformerConfig", diff --git a/zeromodels/st_models/__init__.py b/zeromodels/st_models/__init__.py new file mode 100644 index 00000000..e08d52a5 --- /dev/null +++ b/zeromodels/st_models/__init__.py @@ -0,0 +1,20 @@ +"""Sentence-embedding models, separate from the general model zoo.""" + +from .sentence_transformer import ( + SentenceTransformerConfig, + SentenceTransformerModel, + SentenceTransformerTextEmbed, + SentenceTransformerTokenizer, +) + +# The compact public API mirrors sentence-transformers' main class name while +# preserving the existing ZeroModels class and Hub model_class metadata. +SentenceTransformer = SentenceTransformerModel + +__all__ = [ + "SentenceTransformer", + "SentenceTransformerConfig", + "SentenceTransformerModel", + "SentenceTransformerTextEmbed", + "SentenceTransformerTokenizer", +] diff --git a/zeromodels/st_models/sentence_transformer/__init__.py b/zeromodels/st_models/sentence_transformer/__init__.py new file mode 100644 index 00000000..e75d298a --- /dev/null +++ b/zeromodels/st_models/sentence_transformer/__init__.py @@ -0,0 +1,15 @@ +"""Implementation of the Sentence Transformer model family.""" + +from .sentence_transformer_config import SentenceTransformerConfig +from .sentence_transformer_model import ( + SentenceTransformerModel, + SentenceTransformerTextEmbed, +) +from .sentence_transformer_tokenizer import SentenceTransformerTokenizer + +__all__ = [ + "SentenceTransformerConfig", + "SentenceTransformerModel", + "SentenceTransformerTextEmbed", + "SentenceTransformerTokenizer", +] diff --git a/zeromodels/models/sentence_transformer/convert_sentence_transformer_hf_to_keras.py b/zeromodels/st_models/sentence_transformer/convert_sentence_transformer_hf_to_keras.py similarity index 98% rename from zeromodels/models/sentence_transformer/convert_sentence_transformer_hf_to_keras.py rename to zeromodels/st_models/sentence_transformer/convert_sentence_transformer_hf_to_keras.py index b587d161..7fc958e1 100644 --- a/zeromodels/models/sentence_transformer/convert_sentence_transformer_hf_to_keras.py +++ b/zeromodels/st_models/sentence_transformer/convert_sentence_transformer_hf_to_keras.py @@ -1,4 +1,4 @@ -"""Transfer Sentence-Transformers / HF MPNet weights into Keras ST models.""" +"""Transfer upstream Sentence Transformer weights into Keras models.""" from __future__ import annotations @@ -218,7 +218,7 @@ def pooling_mode_from_config(pooling_cfg: dict) -> str: from transformers import AutoModel, AutoTokenizer from zeromodels.conversion.zm_config import write_zm_config - from zeromodels.models.sentence_transformer import ( + from zeromodels.st_models import ( SentenceTransformerModel, SentenceTransformerTokenizer, ) diff --git a/zeromodels/models/sentence_transformer/sentence_transformer_backbones.py b/zeromodels/st_models/sentence_transformer/sentence_transformer_backbones.py similarity index 98% rename from zeromodels/models/sentence_transformer/sentence_transformer_backbones.py rename to zeromodels/st_models/sentence_transformer/sentence_transformer_backbones.py index c45b7be3..e24020a9 100644 --- a/zeromodels/models/sentence_transformer/sentence_transformer_backbones.py +++ b/zeromodels/st_models/sentence_transformer/sentence_transformer_backbones.py @@ -1,4 +1,4 @@ -"""Encoder backbones needed by the original Sentence Transformers checkpoints.""" +"""Encoder backbones for the supported Sentence Transformer checkpoints.""" import math diff --git a/zeromodels/models/sentence_transformer/sentence_transformer_config.py b/zeromodels/st_models/sentence_transformer/sentence_transformer_config.py similarity index 95% rename from zeromodels/models/sentence_transformer/sentence_transformer_config.py rename to zeromodels/st_models/sentence_transformer/sentence_transformer_config.py index 9bb7fccf..1deca5be 100644 --- a/zeromodels/models/sentence_transformer/sentence_transformer_config.py +++ b/zeromodels/st_models/sentence_transformer/sentence_transformer_config.py @@ -1,3 +1,5 @@ +"""Configuration for Sentence Transformer embedding models.""" + from typing import Optional from zeromodels.base import BaseConfig diff --git a/zeromodels/models/sentence_transformer/sentence_transformer_layers.py b/zeromodels/st_models/sentence_transformer/sentence_transformer_layers.py similarity index 97% rename from zeromodels/models/sentence_transformer/sentence_transformer_layers.py rename to zeromodels/st_models/sentence_transformer/sentence_transformer_layers.py index fb618d54..405f354e 100644 --- a/zeromodels/models/sentence_transformer/sentence_transformer_layers.py +++ b/zeromodels/st_models/sentence_transformer/sentence_transformer_layers.py @@ -1,4 +1,4 @@ -"""Pooling / projection helpers for Sentence-Transformers-style models.""" +"""Pooling and normalization layers for Sentence Transformer models.""" import keras from keras import layers, ops diff --git a/zeromodels/models/sentence_transformer/sentence_transformer_model.py b/zeromodels/st_models/sentence_transformer/sentence_transformer_model.py similarity index 99% rename from zeromodels/models/sentence_transformer/sentence_transformer_model.py rename to zeromodels/st_models/sentence_transformer/sentence_transformer_model.py index 7cb942a5..240d1325 100644 --- a/zeromodels/models/sentence_transformer/sentence_transformer_model.py +++ b/zeromodels/st_models/sentence_transformer/sentence_transformer_model.py @@ -1,4 +1,4 @@ -"""Sentence-Transformers-style text embedding models (Keras 3).""" +"""Sentence Transformer text embedding models (Keras 3).""" from __future__ import annotations diff --git a/zeromodels/models/sentence_transformer/sentence_transformer_tokenizer.py b/zeromodels/st_models/sentence_transformer/sentence_transformer_tokenizer.py similarity index 98% rename from zeromodels/models/sentence_transformer/sentence_transformer_tokenizer.py rename to zeromodels/st_models/sentence_transformer/sentence_transformer_tokenizer.py index 83015d2e..953d264c 100644 --- a/zeromodels/models/sentence_transformer/sentence_transformer_tokenizer.py +++ b/zeromodels/st_models/sentence_transformer/sentence_transformer_tokenizer.py @@ -1,4 +1,4 @@ -"""Fast-tokenizer adapter for supported Sentence Transformers encoders.""" +"""Fast-tokenizer adapter for supported Sentence Transformer encoders.""" import re from typing import List, Union From e437a744034e0a4acd9529bce0187aa190d5b014 Mon Sep 17 00:00:00 2001 From: IMvision12 Date: Thu, 15 Oct 2026 16:34:00 -0700 Subject: [PATCH 03/10] fix-qwen --- zeromodels/models/qwen_image_21/qwen_image_21_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/zeromodels/models/qwen_image_21/qwen_image_21_model.py b/zeromodels/models/qwen_image_21/qwen_image_21_model.py index b8643806..d919310e 100644 --- a/zeromodels/models/qwen_image_21/qwen_image_21_model.py +++ b/zeromodels/models/qwen_image_21/qwen_image_21_model.py @@ -387,7 +387,7 @@ def __init__( ) is_text = ops.convert_to_tensor((~image_pad_mask).astype(np.float32)) key_ok = text_valid_on_joint * is_text[None, :] + (1.0 - is_text[None, :]) - attn_mask = attn_mask + (1.0 - key_ok)[:, None, None, :] * MASK_NEG + attn_mask = ops.add(attn_mask, (1.0 - key_ok)[:, None, None, :] * MASK_NEG) if causal_condition: t_all = ops.concatenate( From 349479bf0d4d16a226d779b234349d43549611bf Mon Sep 17 00:00:00 2001 From: IMvision12 Date: Thu, 15 Oct 2026 16:35:51 -0700 Subject: [PATCH 04/10] fix-docs --- docs/models.md | 8 +++++- docs/sentence_transformer.md | 26 +++----------------- docs/sentence_transformers.md | 46 +++++++++++++++++++++++++++++++++++ website/mkdocs.yml | 4 ++- 4 files changed, 59 insertions(+), 25 deletions(-) create mode 100644 docs/sentence_transformers.md diff --git a/docs/models.md b/docs/models.md index 978541be..5c06b6ab 100644 --- a/docs/models.md +++ b/docs/models.md @@ -37,9 +37,15 @@ Encoders for embeddings and masked language modelling, and decoder LLMs, dense a - [Qwen3.5](qwen3_5.md) - [Qwen3-Next](qwen3_next.md) - [RoBERTa](roberta.md) -- [Sentence Transformer](sentence_transformer.md) - [XLM-RoBERTa](xlm_roberta.md) +## Sentence Transformer models + +Dense text embeddings with checkpoint-specific pooling and normalization. + +- [All 13 Sentence Transformer checkpoints](sentence_transformers.md) +- [Sentence Transformer API and usage](sentence_transformer.md) + ## Vision models Detection, segmentation, monocular depth, and self-supervised backbones. diff --git a/docs/sentence_transformer.md b/docs/sentence_transformer.md index b592dae2..36b8c2e7 100644 --- a/docs/sentence_transformer.md +++ b/docs/sentence_transformer.md @@ -10,34 +10,14 @@ with the hf: prefix, e.g. The backbone and embedding recipe are read from the checkpoint metadata. -ZeroModels loads the 13 original dense sentence-embedding checkpoints shown in -the [Sentence Transformers pretrained-model table](https://sbert.net/docs/sentence_transformer/pretrained_models.html). +ZeroModels loads 13 original dense sentence-embedding checkpoints. See the +[Sentence Transformer models](sentence_transformers.md) page for the available +ZeroModels-hosted checkpoint IDs and Hub links. It builds a Keras 3 encoder and applies the checkpoint's declared pooling, optional Dense projection, and optional L2 normalization. Parity has been verified on the Keras Torch backend; TensorFlow and JAX parity have not yet been checked for all 13 checkpoints. -## Available ZeroModels checkpoints - -| Checkpoint (`zeromodels/`) | Actual backbone | Embedding recipe | Dim | -|---|---|---|---:| -| `all-mpnet-base-v2` | MPNet | mean → normalize | 768 | -| `multi-qa-mpnet-base-dot-v1` | MPNet | CLS, dot-product similarity | 768 | -| `all-distilroberta-v1` | RoBERTa | mean → normalize | 768 | -| `all-MiniLM-L12-v2` | BERT | mean → normalize | 384 | -| `multi-qa-distilbert-cos-v1` | DistilBERT | mean → normalize | 768 | -| `all-MiniLM-L6-v2` | BERT | mean → normalize | 384 | -| `multi-qa-MiniLM-L6-cos-v1` | BERT | mean → normalize | 384 | -| `paraphrase-multilingual-mpnet-base-v2` | XLM-RoBERTa | mean | 768 | -| `paraphrase-albert-small-v2` | ALBERT | mean | 768 | -| `paraphrase-multilingual-MiniLM-L12-v2` | BERT | mean | 384 | -| `paraphrase-MiniLM-L3-v2` | BERT | mean | 384 | -| `distiluse-base-multilingual-cased-v1` | DistilBERT | mean → Tanh Dense | 512 | -| `distiluse-base-multilingual-cased-v2` | DistilBERT | mean → Tanh Dense | 512 | - -The name `paraphrase-multilingual-mpnet-base-v2` is historical: its actual -checkpoint `model_type` is `xlm-roberta`, and ZeroModels uses the XLM-R backbone. - ## Usage ```python diff --git a/docs/sentence_transformers.md b/docs/sentence_transformers.md new file mode 100644 index 00000000..84536f69 --- /dev/null +++ b/docs/sentence_transformers.md @@ -0,0 +1,46 @@ +# Sentence Transformer models + +ZeroModels provides 13 converted dense sentence-embedding checkpoints. Each +checkpoint combines a text encoder with its original pooling, optional Dense +projection, and optional L2 normalization. Load them through the +[`zeromodels.st_models` API](sentence_transformer.md); no conversion is needed +for the ZeroModels-hosted weights. + +| Checkpoint | Backbone | Embedding recipe | Dimensions | +|---|---|---|---:| +| [`all-mpnet-base-v2`](https://huggingface.co/zeromodels/all-mpnet-base-v2) | MPNet | Mean pooling → L2 normalization | 768 | +| [`multi-qa-mpnet-base-dot-v1`](https://huggingface.co/zeromodels/multi-qa-mpnet-base-dot-v1) | MPNet | CLS pooling; dot-product similarity | 768 | +| [`all-distilroberta-v1`](https://huggingface.co/zeromodels/all-distilroberta-v1) | RoBERTa | Mean pooling → L2 normalization | 768 | +| [`all-MiniLM-L12-v2`](https://huggingface.co/zeromodels/all-MiniLM-L12-v2) | BERT | Mean pooling → L2 normalization | 384 | +| [`multi-qa-distilbert-cos-v1`](https://huggingface.co/zeromodels/multi-qa-distilbert-cos-v1) | DistilBERT | Mean pooling → L2 normalization | 768 | +| [`all-MiniLM-L6-v2`](https://huggingface.co/zeromodels/all-MiniLM-L6-v2) | BERT | Mean pooling → L2 normalization | 384 | +| [`multi-qa-MiniLM-L6-cos-v1`](https://huggingface.co/zeromodels/multi-qa-MiniLM-L6-cos-v1) | BERT | Mean pooling → L2 normalization | 384 | +| [`paraphrase-multilingual-mpnet-base-v2`](https://huggingface.co/zeromodels/paraphrase-multilingual-mpnet-base-v2) | XLM-RoBERTa | Mean pooling | 768 | +| [`paraphrase-albert-small-v2`](https://huggingface.co/zeromodels/paraphrase-albert-small-v2) | ALBERT | Mean pooling | 768 | +| [`paraphrase-multilingual-MiniLM-L12-v2`](https://huggingface.co/zeromodels/paraphrase-multilingual-MiniLM-L12-v2) | BERT | Mean pooling | 384 | +| [`paraphrase-MiniLM-L3-v2`](https://huggingface.co/zeromodels/paraphrase-MiniLM-L3-v2) | BERT | Mean pooling | 384 | +| [`distiluse-base-multilingual-cased-v1`](https://huggingface.co/zeromodels/distiluse-base-multilingual-cased-v1) | DistilBERT | Mean pooling → Tanh Dense | 512 | +| [`distiluse-base-multilingual-cased-v2`](https://huggingface.co/zeromodels/distiluse-base-multilingual-cased-v2) | DistilBERT | Mean pooling → Tanh Dense | 512 | + +The `paraphrase-multilingual-mpnet-base-v2` name is historical; the checkpoint +actually uses an XLM-RoBERTa backbone. + +## Quick start + +```python +import os + +os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" + +from zeromodels.st_models import SentenceTransformer, SentenceTransformerTokenizer + +repo = "zeromodels/all-MiniLM-L6-v2" +model = SentenceTransformer.from_weights(repo) +tokenizer = SentenceTransformerTokenizer.from_weights(repo) + +embeddings = model.encode(["A cat on a mat.", "A dog on a mat."], tokenizer=tokenizer) +print(embeddings.shape) # (2, 384) +``` + +See [Sentence Transformer](sentence_transformer.md) for the complete API, +upstream `hf:` loading, similarity scoring, and parity-check instructions. diff --git a/website/mkdocs.yml b/website/mkdocs.yml index e9f38e9d..365e80aa 100644 --- a/website/mkdocs.yml +++ b/website/mkdocs.yml @@ -159,7 +159,6 @@ nav: - Mixtral: mixtral.md - ModernBERT: modernbert.md - MPNet: mpnet.md - - Sentence Transformer: sentence_transformer.md - Qwen: qwen.md - Qwen2: qwen2.md - Qwen2 MoE: qwen2_moe.md @@ -170,6 +169,9 @@ nav: - RoBERTa: roberta.md - T5: t5.md - XLM-RoBERTa: xlm_roberta.md + - Sentence Transformer models: + - sentence_transformers.md + - Sentence Transformer: sentence_transformer.md - Vision models: - BEiT: beit.md - Classification backbones: classification_backbones.md From cb364b2941e8fe4afb0f7d8d7259aa02c0b7e872 Mon Sep 17 00:00:00 2001 From: IMvision12 Date: Thu, 15 Oct 2026 17:49:58 -0700 Subject: [PATCH 05/10] Update token logic --- docs/sentence_transformer.md | 47 +++ .../test_sentence_transformer_tokenizer.py | 250 ++++++++++++++ ...t_sentence_transformer_tokenizers_vs_hf.py | 54 ++++ zeromodels/conversion/zm_config.py | 12 +- zeromodels/st_models/__init__.py | 4 - .../sentence_transformer_tokenizer.py | 306 +++++++++++++----- 6 files changed, 592 insertions(+), 81 deletions(-) create mode 100644 tests/integration/test_sentence_transformer_tokenizer.py create mode 100644 tests/processors/test_sentence_transformer_tokenizers_vs_hf.py diff --git a/docs/sentence_transformer.md b/docs/sentence_transformer.md index 36b8c2e7..29a8e6f5 100644 --- a/docs/sentence_transformer.md +++ b/docs/sentence_transformer.md @@ -47,6 +47,53 @@ print(model.similarity(embeddings[:1], embeddings[1:])) the model and tokenizer. `model.similarity()` uses the checkpoint's declared cosine or dot-product scoring function. +## Tokenizer preprocessing metadata + +Each checkpoint uses its own `tokenizer.json` for the vocabulary, normalization, +token splitting, and special-token insertion. Optional `zm_preprocessor.json` +records the surrounding official configuration in two separate blocks: + +```json +{ + "library_name": "zeromodels", + "preprocessor_module": "zeromodels.st_models", + "preprocessor_class": "SentenceTransformerTokenizer", + "tokenizer_file": "tokenizer.json", + "backbone_type": "bert", + "tokenizer_config": { + "do_lower_case": true, + "model_max_length": 512, + "cls_token": "[CLS]", + "sep_token": "[SEP]", + "pad_token": "[PAD]" + }, + "sentence_transformer_config": { + "max_seq_length": 256, + "do_lower_case": false + } +} +``` + +This abbreviated example represents `all-MiniLM-L6-v2`. The tokenizer's own +lowercasing is already encoded in `tokenizer.json`; the Sentence Transformers +flag requests an additional normalization step only when true. The effective +sequence limit is 256, including special tokens. Padding and truncation sides +come from `tokenizer_config`, defaulting to right; batches pad to their longest +sequence and pairs use longest-first truncation. + +Both metadata and tokenizer assets load from one resolved Hub commit, available +as `tokenizer.source_revision`. Pass `revision=` to select a particular snapshot. +An explicit `max_seq_len=` overrides the stored limit. Legacy repos without +`zm_preprocessor.json` continue to read the backbone and sequence limit from +`zm_config.json`, including its nested `text_config` form. + +`write_zm_preprocessor(directory, tokenizer, variant)` exports a portable recipe +referencing `tokenizer.json` in the same repo, rather than a local cache path. +Text preprocessing uses the `tokenizers` library; `transformers` is not needed +to load these converted checkpoints. + +## Original Sentence Transformers checkpoints + To load directly from an original Sentence Transformers repo instead, use the `hf:` prefix for both components and pass the model's resolved revision to the tokenizer: diff --git a/tests/integration/test_sentence_transformer_tokenizer.py b/tests/integration/test_sentence_transformer_tokenizer.py new file mode 100644 index 00000000..91e12c53 --- /dev/null +++ b/tests/integration/test_sentence_transformer_tokenizer.py @@ -0,0 +1,250 @@ +"""Checkpoint metadata must control batching without replacing the JSON pipeline.""" + +import json +from types import SimpleNamespace + +import keras +import numpy as np +import pytest +from huggingface_hub.utils import EntryNotFoundError +from keras import ops +from tokenizers import Tokenizer, models, normalizers, pre_tokenizers, processors + +from zeromodels.st_models import SentenceTransformerTokenizer + +SHA = "a" * 40 + + +@pytest.fixture +def tokenizer_path(tmp_path): + vocab = { + token: i + for i, token in enumerate( + ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]", "hello", "world", "Hello"] + ) + } + tok = Tokenizer(models.WordPiece(vocab, unk_token="[UNK]")) + tok.normalizer = normalizers.BertNormalizer(lowercase=False) + tok.pre_tokenizer = pre_tokenizers.BertPreTokenizer() + tok.post_processor = processors.TemplateProcessing( + single="[CLS] $A [SEP]", + pair="[CLS] $A [SEP] $B:1 [SEP]:1", + special_tokens=[("[CLS]", vocab["[CLS]"]), ("[SEP]", vocab["[SEP]"])], + ) + path = tmp_path / "tokenizer.json" + tok.save(str(path)) + return path + + +def metadata(): + return { + "preprocessor_module": "zeromodels.st_models", + "preprocessor_class": "SentenceTransformerTokenizer", + "tokenizer_file": "assets/tokenizer.json", + "backbone_type": "bert", + "tokenizer_config": { + "model_max_length": 512, + "do_lower_case": False, + "cls_token": "[CLS]", + "sep_token": "[SEP]", + "pad_token": {"content": "[PAD]", "special": True}, + }, + "sentence_transformer_config": { + "max_seq_length": 6, + "do_lower_case": False, + }, + } + + +def mock_repo(monkeypatch, tmp_path, files): + calls = [] + + def download(repo_id, filename, token=None, revision=None): + calls.append((filename, revision)) + if filename not in files: + raise EntryNotFoundError(f"Missing {filename}") + value = files[filename] + if not isinstance(value, dict) and not isinstance(value, list): + return str(value) + path = tmp_path / filename + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(value), encoding="utf-8") + return str(path) + + monkeypatch.setattr("huggingface_hub.hf_hub_download", download) + monkeypatch.setattr( + "huggingface_hub.HfApi.model_info", + lambda *args, **kwargs: SimpleNamespace(sha=SHA), + ) + return calls + + +def test_hub_metadata_uses_one_snapshot_and_official_sequence_limit( + monkeypatch, tmp_path, tokenizer_path +): + calls = mock_repo( + monkeypatch, + tmp_path, + { + "zm_preprocessor.json": metadata(), + "assets/tokenizer.json": tokenizer_path, + }, + ) + tok = SentenceTransformerTokenizer.from_weights( + "zeromodels/example", revision="release" + ) + encoded = tok(["hello world " * 10, "hello"]) + assert tok.max_seq_len == 6 + assert tok.cls_token_id == 2 and tok.sep_token_id == 3 + assert tok.pad_token_id == 0 + assert encoded["input_ids"].shape == (2, 6) + assert ops.convert_to_numpy(encoded["input_ids"])[0, -1] == 3 + np.testing.assert_array_equal( + ops.convert_to_numpy(encoded["attention_mask"])[1], [1, 1, 1, 0, 0, 0] + ) + assert tok.source_revision == SHA + assert all(revision == SHA for _, revision in calls) + assert {name for name, _ in calls} == { + "zm_preprocessor.json", + "assets/tokenizer.json", + } + + +def test_portable_preprocessor_export_loads_again( + monkeypatch, tmp_path, tokenizer_path +): + from zeromodels.conversion.zm_config import write_zm_preprocessor + + original = SentenceTransformerTokenizer( + tokenizer_file=str(tokenizer_path), + backbone_type="bert", + max_seq_len=6, + tokenizer_config={"padding_side": "left"}, + sentence_transformer_config={"do_lower_case": True}, + ) + path = write_zm_preprocessor(tmp_path, original, "example") + with open(path, encoding="utf-8") as stream: + spec = json.load(stream) + assert spec["tokenizer_file"] == "tokenizer.json" + mock_repo( + monkeypatch, + tmp_path, + {"zm_preprocessor.json": spec, "tokenizer.json": tokenizer_path}, + ) + reloaded = SentenceTransformerTokenizer.from_weights("example", revision=SHA) + for key, expected in original(["Hello world", "hello"]).items(): + np.testing.assert_array_equal( + ops.convert_to_numpy(reloaded(["Hello world", "hello"])[key]), + ops.convert_to_numpy(expected), + ) + + +def test_legacy_nested_config_and_explicit_limit_keep_backbone( + monkeypatch, tmp_path, tokenizer_path +): + mock_repo( + monkeypatch, + tmp_path, + { + "zm_config.json": { + "model_type": "sentence_transformer", + "text_config": {"max_seq_len": 8, "backbone_type": "bert"}, + }, + "tokenizer.json": tokenizer_path, + }, + ) + tok = SentenceTransformerTokenizer.from_weights("zeromodels/example", max_seq_len=5) + assert tok.max_seq_len == 5 + assert tok.backbone_type == "bert" + assert "token_type_ids" in tok(["hello"]) + + +def test_sentence_transformer_lowercase_is_separate_from_tokenizer_normalization( + tokenizer_path, +): + tc = {"do_lower_case": False, "pad_token": "[PAD]"} + plain = SentenceTransformerTokenizer( + tokenizer_file=str(tokenizer_path), + tokenizer_config=tc, + sentence_transformer_config={"do_lower_case": False}, + ) + extra = SentenceTransformerTokenizer( + tokenizer_file=str(tokenizer_path), + tokenizer_config=tc, + sentence_transformer_config={"do_lower_case": True}, + ) + assert plain.tokenize("Hello") == [7] + assert extra.tokenize("Hello") == [5] + raw = Tokenizer.from_file(str(tokenizer_path)) + raw.normalizer = normalizers.BertNormalizer(lowercase=True) + raw.save(str(tokenizer_path)) + inherited = SentenceTransformerTokenizer( + tokenizer_file=str(tokenizer_path), + tokenizer_config={"do_lower_case": True}, + sentence_transformer_config={"do_lower_case": False}, + ) + assert inherited.tokenize("Hello") == [5] + + +def test_padding_truncation_pairs_and_serialization(tokenizer_path): + tok = SentenceTransformerTokenizer( + tokenizer_file=str(tokenizer_path), + backbone_type="bert", + tokenizer_config={"padding_side": "left", "truncation_side": "left"}, + sentence_transformer_config={"max_seq_length": 6}, + ) + clone = keras.saving.deserialize_keras_object( + keras.saving.serialize_keras_object(tok) + ) + inputs = ["hello world hello world hello", "hello"] + for current in (tok, clone): + encoded = current(inputs, text_pair=["world", "world"]) + ids = ops.convert_to_numpy(encoded["input_ids"]) + np.testing.assert_array_equal(ids[1], [0, 2, 5, 3, 6, 3]) + np.testing.assert_array_equal( + ops.convert_to_numpy(encoded["token_type_ids"])[1], [0, 0, 0, 0, 1, 1] + ) + assert ids[0, 0] == 2 and ids[0, -1] == 3 + + +def test_upstream_module_path_and_sidecars(monkeypatch, tmp_path, tokenizer_path): + calls = mock_repo( + monkeypatch, + tmp_path, + { + "modules.json": [{"path": "0_Transformer"}], + "0_Transformer/config.json": {"model_type": "bert"}, + "0_Transformer/tokenizer_config.json": {"model_max_length": 512}, + "0_Transformer/special_tokens_map.json": { + "cls_token": "[CLS]", + "sep_token": "[SEP]", + "pad_token": "[PAD]", + }, + "0_Transformer/sentence_bert_config.json": { + "max_seq_length": 5, + "do_lower_case": True, + }, + "0_Transformer/tokenizer.json": tokenizer_path, + }, + ) + tok = SentenceTransformerTokenizer.from_weights("hf:sentence-transformers/example") + assert tok.max_seq_len == 5 and tok.cls_token_id == 2 + assert tok.tokenize("Hello") == [5] + assert all(revision == SHA for _, revision in calls) + + +@pytest.mark.parametrize( + "filename", + [ + "../tokenizer.json", + "/tmp/tokenizer.json", + "https://example.com/tokenizer.json", + "C:\\tokenizer.json", + ], +) +def test_hub_rejects_non_repository_tokenizer_paths(monkeypatch, tmp_path, filename): + spec = metadata() + spec["tokenizer_file"] = filename + mock_repo(monkeypatch, tmp_path, {"zm_preprocessor.json": spec}) + with pytest.raises(ValueError, match="repository-relative"): + SentenceTransformerTokenizer.from_weights("zeromodels/example") diff --git a/tests/processors/test_sentence_transformer_tokenizers_vs_hf.py b/tests/processors/test_sentence_transformer_tokenizers_vs_hf.py new file mode 100644 index 00000000..94f48bee --- /dev/null +++ b/tests/processors/test_sentence_transformer_tokenizers_vs_hf.py @@ -0,0 +1,54 @@ +"""Optional official-tokenizer parity for all 13 ST checkpoints; no weights. + +Run with transformers installed: + KERAS_BACKEND=torch pytest tests/processors/test_sentence_transformer_tokenizers_vs_hf.py +""" + +import numpy as np +import pytest +from keras import ops + +from tests.base.model_test_registry import _SENTENCE_TRANSFORMER_VARIANTS +from zeromodels.st_models import SentenceTransformerTokenizer + +transformers = pytest.importorskip("transformers") + +TEXTS = [ + "", + " Hello, WORLD! ", + "Café naïve résumé", + "你好，世界！ مرحبا بالعالم", + "line one\nline two\t123", + "[CLS] hello [SEP]", + "hello world " * 600, +] +PAIRS = ["a pair", "", "deuxième phrase", "हिन्दी पाठ", "test", "[MASK]", "world " * 500] + + +@pytest.mark.parametrize("name", list(_SENTENCE_TRANSFORMER_VARIANTS)) +def test_official_tokenizer_parity_without_weights(name): + repo = f"sentence-transformers/{name}" + ours = SentenceTransformerTokenizer.from_hf(repo) + reference = transformers.AutoTokenizer.from_pretrained( + repo, revision=ours.source_revision, use_fast=True, trust_remote_code=False + ) + # The 13 supported checkpoints all disable extra ST lowercasing; their own + # tokenizer.json normalizers still apply. The offline suite covers True. + assert ours.sentence_transformer_config["do_lower_case"] is False + for pairs in (None, PAIRS): + actual = ours(TEXTS, text_pair=pairs) + expected = reference( + TEXTS, + text_pair=pairs, + padding=True, + truncation="longest_first", + max_length=ours.max_seq_len, + return_tensors="np", + return_token_type_ids="token_type_ids" in actual, + ) + for key, value in actual.items(): + np.testing.assert_array_equal( + ops.convert_to_numpy(value), + expected[key], + err_msg=f"{name}: {key}, pairs={pairs is not None}", + ) diff --git a/zeromodels/conversion/zm_config.py b/zeromodels/conversion/zm_config.py index d62e2f39..67e0ef32 100644 --- a/zeromodels/conversion/zm_config.py +++ b/zeromodels/conversion/zm_config.py @@ -4,7 +4,7 @@ root, next to model.weights.h5: zm_config.json the model: class + flat hyperparameters + weights filename. - zm_preprocessor.json the image processor / feature extractor (when present). + zm_preprocessor.json preprocessing settings (image/audio or tokenizer). They let a model be rebuilt and loaded from the repo alone, so a caller loads by repo id (zeromodels/detr-resnet-50) instead of a variant hardcoded in the @@ -13,7 +13,8 @@ (``model_module`` / ``model_class`` / ``variant`` / ``weights``). The zm_ prefix avoids the config.json / preprocessor_config.json names that would make the Hub mis-detect these Keras repos as transformers or keras-hub. The tokenizer stays a -real tokenizer.json and is not described here. +real tokenizer.json; optional tokenizer and Sentence Transformers settings are +recorded as separate blocks in zm_preprocessor.json. """ import inspect @@ -87,12 +88,17 @@ def model_config_dict(config): def preprocessor_config(preprocessor): """Serialize a preprocessor's constructor state. + A preprocessor can provide ``get_hub_config`` for a portable Hub recipe + (e.g. tokenizer metadata referencing repo-relative assets). zeromodels image processors do not all implement ``get_config``, so read the values straight off the instance by walking the constructor signature: every named parameter maps to a same-named attribute (``self.size``, ``self.resample``, ...). Missing attributes fall back to the parameter default. ``self`` and ``*args`` / ``**kwargs`` are skipped. """ + hub_config = getattr(preprocessor, "get_hub_config", None) + if callable(hub_config): + return _jsonable(hub_config()) config = {} for name, param in inspect.signature( type(preprocessor).__init__ @@ -181,7 +187,7 @@ def write_zm_config( def write_zm_preprocessor(dest_dir, preprocessor, variant): - """Write a flat zm_preprocessor.json for a built ``preprocessor`` instance.""" + """Write zm_preprocessor.json for a built ``preprocessor`` instance.""" payload = { "library_name": "zeromodels", "zeromodels_version": _zm_version(), diff --git a/zeromodels/st_models/__init__.py b/zeromodels/st_models/__init__.py index e08d52a5..fe44e6ae 100644 --- a/zeromodels/st_models/__init__.py +++ b/zeromodels/st_models/__init__.py @@ -1,5 +1,3 @@ -"""Sentence-embedding models, separate from the general model zoo.""" - from .sentence_transformer import ( SentenceTransformerConfig, SentenceTransformerModel, @@ -7,8 +5,6 @@ SentenceTransformerTokenizer, ) -# The compact public API mirrors sentence-transformers' main class name while -# preserving the existing ZeroModels class and Hub model_class metadata. SentenceTransformer = SentenceTransformerModel __all__ = [ diff --git a/zeromodels/st_models/sentence_transformer/sentence_transformer_tokenizer.py b/zeromodels/st_models/sentence_transformer/sentence_transformer_tokenizer.py index 953d264c..67612f92 100644 --- a/zeromodels/st_models/sentence_transformer/sentence_transformer_tokenizer.py +++ b/zeromodels/st_models/sentence_transformer/sentence_transformer_tokenizer.py @@ -1,10 +1,11 @@ -"""Fast-tokenizer adapter for supported Sentence Transformer encoders.""" - +import json +import os import re +from pathlib import PurePosixPath from typing import List, Union import keras -from tokenizers import Tokenizer +from tokenizers import Tokenizer, normalizers from zeromodels.base import BaseTokenizer @@ -14,16 +15,18 @@ class SentenceTransformerTokenizer(BaseTokenizer): """Tokenizer for supported ST models, preserving their tokenizer.json rules. Loads ``tokenizer.json`` from a Hub repo (via ``from_weights("org/repo")`` / - ``from_hub_repo``) or an explicit ``tokenizer_file``. Defaults match - ``sentence-transformers/all-mpnet-base-v2`` (``max_seq_len=384``, MPNet - special tokens). ``call`` returns ``input_ids`` / ``attention_mask`` only. + ``from_hub_repo``) or an explicit ``tokenizer_file``. ``zm_preprocessor.json`` + keeps official tokenizer settings separate from Sentence Transformers' + additional preprocessing. The JSON pipeline owns normalization and special + token insertion; metadata supplies special token identities and batching. + Repos without preprocessing metadata retain the legacy loading path. Args: variant: Optional Hub variant key; resolves ``zeromodels//tokenizer.json`` when ``tokenizer_file`` is omitted. tokenizer_file: Explicit ``tokenizer.json`` path (overrides variant). - max_seq_len: Truncation length (default 384). + max_seq_len: Explicit truncation override; otherwise read from metadata. bos_token / eos_token / unk_token / pad_token / mask_token: Special tokens. """ @@ -31,13 +34,15 @@ def __init__( self, variant: str = None, tokenizer_file: str = None, - max_seq_len: int = 384, + max_seq_len: int = None, backbone_type: str = "mpnet", - bos_token: str = "", - eos_token: str = "", - unk_token: str = "[UNK]", + bos_token: str = None, + eos_token: str = None, + unk_token: str = None, pad_token: str = None, - mask_token: str = "", + mask_token: str = None, + tokenizer_config: dict = None, + sentence_transformer_config: dict = None, **kwargs, ): super().__init__(**kwargs) @@ -46,14 +51,53 @@ def __init__( (f"zeromodels/{self.variant}" if self.variant else None), tokenizer_file ) self.tokenizer_file = tokenizer_file - self.max_seq_len = max_seq_len self.backbone_type = backbone_type - self.bos_token = bos_token - self.eos_token = eos_token - self.unk_token = unk_token - self.mask_token = mask_token + self.tokenizer_config = dict(tokenizer_config or {}) + self.sentence_transformer_config = dict(sentence_transformer_config or {}) + tc, sc = self.tokenizer_config, self.sentence_transformer_config + if max_seq_len is None: + max_seq_len = sc.get("max_seq_length") + if max_seq_len is None: + candidate = tc.get("model_max_length") + max_seq_len = ( + candidate + if isinstance(candidate, int) and 0 < candidate < 10**9 + else 384 + ) + if ( + isinstance(max_seq_len, bool) + or not isinstance(max_seq_len, int) + or max_seq_len < 1 + ): + raise ValueError("max_seq_len must be a positive integer") + self.max_seq_len = max_seq_len tok = Tokenizer.from_file(tokenizer_file) + if sc.get("do_lower_case", False) and not self._has_lowercase( + json.loads(tok.to_str()).get("normalizer") + ): + tok.normalizer = normalizers.Sequence( + [normalizers.Lowercase()] + ([tok.normalizer] if tok.normalizer else []) + ) + for name, explicit, fallback in ( + ("bos_token", bos_token, ""), + ("eos_token", eos_token, ""), + ("unk_token", unk_token, "[UNK]"), + ("pad_token", pad_token, None), + ("mask_token", mask_token, ""), + ("cls_token", None, None), + ("sep_token", None, None), + ): + value = explicit if explicit is not None else tc.get(name, fallback) + if isinstance(value, dict): + value = value["content"] + setattr(self, name, value) + setattr( + self, + f"{name}_id", + tok.token_to_id(value) if value is not None else None, + ) + pad_token = self.pad_token if pad_token is None: pad_token = next( ( @@ -66,75 +110,114 @@ def __init__( if pad_token is None: raise ValueError("tokenizer.json has no supported padding token") self.pad_token = pad_token - self.bos_token_id = tok.token_to_id(bos_token) - self.eos_token_id = tok.token_to_id(eos_token) - self.unk_token_id = tok.token_to_id(unk_token) self.pad_token_id = tok.token_to_id(pad_token) - self.mask_token_id = tok.token_to_id(mask_token) - tok.enable_truncation(max_length=max_seq_len) - tok.enable_padding(pad_id=self.pad_token_id, pad_token=pad_token) + if self.pad_token_id is None: + raise ValueError( + f"Padding token {pad_token!r} is absent from tokenizer.json" + ) + tok.enable_truncation( + max_length=max_seq_len, + direction=tc.get("truncation_side", "right"), + strategy="longest_first", + ) + tok.enable_padding( + direction=tc.get("padding_side", "right"), + pad_id=self.pad_token_id, + pad_token=pad_token, + ) self._tok = tok @classmethod - def from_hf(cls, repo, **kwargs): - import json - import os + def from_variant(cls, variant, /, **kwargs): + if kwargs.get("tokenizer_file") is not None: + return cls(variant=variant, **kwargs) + return cls.from_hub_repo(f"zeromodels/{variant}", **kwargs) - from huggingface_hub import HfApi, hf_hub_download - from huggingface_hub.utils import EntryNotFoundError + @classmethod + def from_hf(cls, repo, **kwargs): + from huggingface_hub import hf_hub_download - token = os.environ.get("HF_TOKEN") - revision = kwargs.pop("revision", None) - if not (isinstance(revision, str) and re.fullmatch(r"[0-9a-f]{40}", revision)): - revision = HfApi(token=token).model_info(repo, revision=revision).sha - with open( - hf_hub_download(repo, "config.json", token=token, revision=revision) - ) as f: - hf_config = json.load(f) + token = kwargs.pop("token", os.environ.get("HF_TOKEN")) + revision = cls._snapshot(repo, kwargs.pop("revision", None), token) + modules = cls._optional_json(repo, "modules.json", token, revision) + prefix = "" + if modules and modules[0].get("path"): + prefix = cls._repo_file(modules[0]["path"]) + "/" + hf_config = cls._optional_json(repo, prefix + "config.json", token, revision) + if hf_config is None: + raise ValueError(f"{repo}: missing Transformer config.json") kwargs.setdefault("backbone_type", hf_config["model_type"]) - if "max_seq_len" not in kwargs: - try: - path = hf_hub_download( - repo, "sentence_bert_config.json", token=token, revision=revision - ) - except EntryNotFoundError: - pass - else: - with open(path) as f: - kwargs["max_seq_len"] = int(json.load(f)["max_seq_length"]) - path = hf_hub_download(repo, "tokenizer.json", token=token, revision=revision) - return cls(tokenizer_file=path, **kwargs) + tc = ( + cls._optional_json(repo, prefix + "tokenizer_config.json", token, revision) + or {} + ) + tc.update( + cls._optional_json( + repo, prefix + "special_tokens_map.json", token, revision + ) + or {} + ) + kwargs.setdefault("tokenizer_config", tc) + kwargs.setdefault( + "sentence_transformer_config", + cls._optional_json( + repo, prefix + "sentence_bert_config.json", token, revision + ) + or {}, + ) + path = hf_hub_download( + repo, prefix + "tokenizer.json", token=token, revision=revision + ) + model = cls(tokenizer_file=path, **kwargs) + model.source_revision = revision + return model @classmethod def from_hub_repo(cls, repo_id, **kwargs): - """Load ``tokenizer.json`` and honour ``max_seq_len`` from zm / ST configs.""" - import json - import os - + """Load official preprocessing metadata, with legacy config fallback.""" from huggingface_hub import hf_hub_download - from huggingface_hub.utils import EntryNotFoundError - token = os.environ.get("HF_TOKEN") - path = hf_hub_download(repo_id.rstrip("/"), "tokenizer.json", token=token) - if "max_seq_len" not in kwargs: - for filename in ("zm_config.json", "sentence_bert_config.json"): - try: - cfg_path = hf_hub_download( - repo_id.rstrip("/"), filename, token=token - ) - except EntryNotFoundError: - continue - with open(cfg_path, "r") as f: - cfg = json.load(f) - if "max_seq_len" in cfg: - kwargs["max_seq_len"] = int(cfg["max_seq_len"]) - elif "max_seq_length" in cfg: - kwargs["max_seq_len"] = int(cfg["max_seq_length"]) - if "backbone_type" in cfg: - kwargs.setdefault("backbone_type", cfg["backbone_type"]) - if "max_seq_len" in kwargs: - break - return cls(tokenizer_file=path, **kwargs) + repo_id = repo_id.rstrip("/") + token = kwargs.pop("token", os.environ.get("HF_TOKEN")) + revision = cls._snapshot(repo_id, kwargs.pop("revision", None), token) + spec = cls._optional_json(repo_id, "zm_preprocessor.json", token, revision) + filename = "tokenizer.json" + if spec is not None: + if spec.get("preprocessor_class") != cls.__name__: + raise ValueError( + f"{repo_id}: zm_preprocessor.json must declare {cls.__name__}" + ) + filename = cls._repo_file(spec.get("tokenizer_file", filename)) + for field in ( + "tokenizer_config", + "sentence_transformer_config", + "backbone_type", + "max_seq_len", + ): + if field in spec: + kwargs.setdefault(field, spec[field]) + else: + spec = cls._optional_json(repo_id, "zm_config.json", token, revision) or {} + config = {**spec.get("text_config", {}), **spec} + for field in ("backbone_type", "max_seq_len"): + if field in config: + kwargs.setdefault(field, config[field]) + kwargs.setdefault( + "tokenizer_config", + cls._optional_json(repo_id, "tokenizer_config.json", token, revision) + or {}, + ) + kwargs.setdefault( + "sentence_transformer_config", + cls._optional_json( + repo_id, "sentence_bert_config.json", token, revision + ) + or {}, + ) + path = hf_hub_download(repo_id, filename, token=token, revision=revision) + model = cls(tokenizer_file=path, **kwargs) + model.source_revision = revision + return model @property def vocab_size(self) -> int: @@ -177,6 +260,81 @@ def get_config(self): "unk_token": self.unk_token, "pad_token": self.pad_token, "mask_token": self.mask_token, + "tokenizer_config": self.tokenizer_config, + "sentence_transformer_config": self.sentence_transformer_config, } ) return config + + def get_hub_config(self): + """Portable preprocessing recipe; tokenizer.json is stored alongside it.""" + tc = dict(self.tokenizer_config) + for name in ( + "bos_token", + "eos_token", + "unk_token", + "pad_token", + "mask_token", + "cls_token", + "sep_token", + ): + if getattr(self, f"{name}_id") is not None: + tc[name] = getattr(self, name) + sc = dict(self.sentence_transformer_config) + sc["max_seq_length"] = self.max_seq_len + return { + "tokenizer_file": "tokenizer.json", + "backbone_type": self.backbone_type, + "tokenizer_config": tc, + "sentence_transformer_config": sc, + } + + @staticmethod + def _repo_file(filename): + """Accept only filenames inside the selected Hub repository.""" + if ( + not isinstance(filename, str) + or not filename + or "\\" in filename + or ":" in filename + or filename.startswith("/") + or any(part in ("", ".", "..") for part in filename.split("/")) + ): + raise ValueError( + f"Expected a repository-relative filename, got {filename!r}" + ) + return str(PurePosixPath(filename)) + + @staticmethod + def _optional_json(repo, filename, token, revision): + from huggingface_hub import hf_hub_download + from huggingface_hub.utils import EntryNotFoundError + + try: + path = hf_hub_download(repo, filename, token=token, revision=revision) + except EntryNotFoundError: + return None + with open(path, encoding="utf-8") as stream: + return json.load(stream) + + @staticmethod + def _snapshot(repo, revision, token): + from huggingface_hub import HfApi + + if isinstance(revision, str) and re.fullmatch(r"[0-9a-f]{40}", revision): + return revision + sha = HfApi(token=token).model_info(repo, revision=revision).sha + if not sha: + raise OSError(f"{repo}: Hub returned no commit SHA for {revision!r}") + return sha + + @staticmethod + def _has_lowercase(config): + return config is not None and ( + config.get("type") == "Lowercase" + or config.get("lowercase") is True + or any( + SentenceTransformerTokenizer._has_lowercase(item) + for item in config.get("normalizers", []) + ) + ) From 390cf599574d8675b44cd372398bce4d23a7afd6 Mon Sep 17 00:00:00 2001 From: IMvision12 Date: Thu, 15 Oct 2026 17:57:36 -0700 Subject: [PATCH 06/10] Update Readme.md --- README.md | 25 ++++++++++++++++++++++++- 1 file changed, 24 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index 900108b9..4d0006fc 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,7 @@ ## 📖 Introduction -ZeroModels is a collection of pretrained models built entirely in Keras 3. It spans a broad range of tasks, including image classification, object detection, segmentation, monocular depth estimation, feature extraction, vision-language modeling (VLMs), speech recognition, speech-aware language modeling, text encoding and masked language modeling, large language models (LLMs), text-to-text encoder-decoder modeling, multimodal vision-language generation, text-to-image diffusion, and more. +ZeroModels is a collection of pretrained models built entirely in Keras 3. It spans a broad range of tasks, including image classification, object detection, segmentation, monocular depth estimation, feature extraction, vision-language modeling (VLMs), speech recognition, speech-aware language modeling, text encoding and masked language modeling, sentence embeddings, large language models (LLMs), text-to-text encoder-decoder modeling, multimodal vision-language generation, text-to-image diffusion, and more. ## ⚡ Installation @@ -40,6 +40,28 @@ Classification backbones share a single [documentation page](https://imvision12. Documentation sources are also available in [`docs/`](docs/). +## 💬 Sentence Transformers Models (st_models) + +All 13 checkpoints below have pre-converted Keras weights under the `zeromodels` Hugging Face organization. Load them through `zeromodels.st_models` for sentence similarity, semantic search, or clustering. + +| 🏷️ Model Name | 📜 Reference Paper | 📦 Source of Weights | +|---------------|-------------------|---------------------| +| [all-mpnet-base-v2](https://huggingface.co/zeromodels/all-mpnet-base-v2) | [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://arxiv.org/abs/1908.10084) | `sentence-transformers` | +| [multi-qa-mpnet-base-dot-v1](https://huggingface.co/zeromodels/multi-qa-mpnet-base-dot-v1) | [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://arxiv.org/abs/1908.10084) | `sentence-transformers` | +| [all-distilroberta-v1](https://huggingface.co/zeromodels/all-distilroberta-v1) | [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://arxiv.org/abs/1908.10084) | `sentence-transformers` | +| [all-MiniLM-L12-v2](https://huggingface.co/zeromodels/all-MiniLM-L12-v2) | [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://arxiv.org/abs/1908.10084) | `sentence-transformers` | +| [multi-qa-distilbert-cos-v1](https://huggingface.co/zeromodels/multi-qa-distilbert-cos-v1) | [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://arxiv.org/abs/1908.10084) | `sentence-transformers` | +| [all-MiniLM-L6-v2](https://huggingface.co/zeromodels/all-MiniLM-L6-v2) | [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://arxiv.org/abs/1908.10084) | `sentence-transformers` | +| [multi-qa-MiniLM-L6-cos-v1](https://huggingface.co/zeromodels/multi-qa-MiniLM-L6-cos-v1) | [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://arxiv.org/abs/1908.10084) | `sentence-transformers` | +| [paraphrase-multilingual-mpnet-base-v2](https://huggingface.co/zeromodels/paraphrase-multilingual-mpnet-base-v2) | [Making Monolingual Sentence Embeddings Multilingual using Knowledge Distillation](https://arxiv.org/abs/2004.09813) | `sentence-transformers` | +| [paraphrase-albert-small-v2](https://huggingface.co/zeromodels/paraphrase-albert-small-v2) | [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://arxiv.org/abs/1908.10084) | `sentence-transformers` | +| [paraphrase-multilingual-MiniLM-L12-v2](https://huggingface.co/zeromodels/paraphrase-multilingual-MiniLM-L12-v2) | [Making Monolingual Sentence Embeddings Multilingual using Knowledge Distillation](https://arxiv.org/abs/2004.09813) | `sentence-transformers` | +| [paraphrase-MiniLM-L3-v2](https://huggingface.co/zeromodels/paraphrase-MiniLM-L3-v2) | [Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks](https://arxiv.org/abs/1908.10084) | `sentence-transformers` | +| [distiluse-base-multilingual-cased-v1](https://huggingface.co/zeromodels/distiluse-base-multilingual-cased-v1) | [Making Monolingual Sentence Embeddings Multilingual using Knowledge Distillation](https://arxiv.org/abs/2004.09813) | `sentence-transformers` | +| [distiluse-base-multilingual-cased-v2](https://huggingface.co/zeromodels/distiluse-base-multilingual-cased-v2) | [Making Monolingual Sentence Embeddings Multilingual using Knowledge Distillation](https://arxiv.org/abs/2004.09813) | `sentence-transformers` | + +
+ ## 📑 Models ### 📝 Text Models @@ -295,6 +317,7 @@ This project leverages [timm](https://github.com/huggingface/pytorch-image-model - The [Transformers](https://github.com/huggingface/transformers) library for its robust tools for loading and adapting pretrained models - The [pytorch-image-models (timm)](https://github.com/huggingface/pytorch-image-models) project for pioneering many computer vision model implementations - The [Diffusers](https://github.com/huggingface/diffusers) library, the reference the diffusion models are converted from and verified against +- The [Sentence Transformers](https://github.com/huggingface/sentence-transformers) project for its sentence-embedding checkpoints and preprocessing recipes - All contributors to the original papers and architectures implemented in this library ## Citing From 5047a60212c1fcfe9d422ce89b9a45aae1902fed Mon Sep 17 00:00:00 2001 From: IMvision12 Date: Thu, 15 Oct 2026 18:29:59 -0700 Subject: [PATCH 07/10] fix-conversion --- ...onvert_sentence_transformer_hf_to_keras.py | 198 +++++++----------- 1 file changed, 79 insertions(+), 119 deletions(-) diff --git a/zeromodels/st_models/sentence_transformer/convert_sentence_transformer_hf_to_keras.py b/zeromodels/st_models/sentence_transformer/convert_sentence_transformer_hf_to_keras.py index 7fc958e1..d500ac7d 100644 --- a/zeromodels/st_models/sentence_transformer/convert_sentence_transformer_hf_to_keras.py +++ b/zeromodels/st_models/sentence_transformer/convert_sentence_transformer_hf_to_keras.py @@ -1,8 +1,6 @@ -"""Transfer upstream Sentence Transformer weights into Keras models.""" - from __future__ import annotations -from typing import Dict, Optional, Tuple +from typing import Dict import numpy as np from huggingface_hub import hf_hub_download @@ -23,40 +21,43 @@ transfer_xlm_roberta_weights, ) - -def strip_st_prefix(key: str) -> str: - """Normalize ST / HF key prefixes down to bare MPNet (or Dense) names. - - Handles common layouts: - - * ``0.auto_model.*`` — SentenceTransformer module stack - * ``auto_model.*`` - * ``mpnet.*`` — HF ``MPNetModel`` wrapper - """ - for prefix in ( - "0.auto_model.", - "auto_model.", - "0.transformer.", - ): - if key.startswith(prefix): - key = key[len(prefix) :] - break - return normalize_hf_key(key) - - -def _find_dense_tensors( - hf_state_dict: Dict[str, np.ndarray], -) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]: - """Locate an ST ``Dense`` module's ``linear.weight`` / ``linear.bias``.""" - kernel = bias = None - for k, v in hf_state_dict.items(): - if "auto_model" in k or ".layer." in k or k.startswith("embeddings."): - continue - if k.endswith("linear.weight"): - kernel = np.asarray(v) - elif k.endswith("linear.bias"): - bias = np.asarray(v) - return kernel, bias +WEIGHT_NAME_MAPPING = { + "distilbert": { + "distil_word_embeddings": "embeddings.word_embeddings", + "distil_position_embeddings": "embeddings.position_embeddings", + "distil_embeddings_norm": "embeddings.LayerNorm", + "distil_layer_{index}/q_lin": "transformer.layer.{index}.attention.q_lin", + "distil_layer_{index}/k_lin": "transformer.layer.{index}.attention.k_lin", + "distil_layer_{index}/v_lin": "transformer.layer.{index}.attention.v_lin", + "distil_layer_{index}/out_lin": "transformer.layer.{index}.attention.out_lin", + "distil_layer_{index}/sa_layer_norm": "transformer.layer.{index}.sa_layer_norm", + "distil_layer_{index}/ffn_lin1": "transformer.layer.{index}.ffn.lin1", + "distil_layer_{index}/ffn_lin2": "transformer.layer.{index}.ffn.lin2", + "distil_layer_{index}/output_layer_norm": "transformer.layer.{index}.output_layer_norm", + }, + "albert": { + "albert_word_embeddings": "embeddings.word_embeddings", + "albert_position_embeddings": "embeddings.position_embeddings", + "albert_token_type_embeddings": "embeddings.token_type_embeddings", + "albert_embeddings_norm": "embeddings.LayerNorm", + "albert_embedding_projection": "encoder.embedding_hidden_mapping_in", + "albert_shared_block/query": "encoder.albert_layer_groups.0.albert_layers.0.attention.query", + "albert_shared_block/key": "encoder.albert_layer_groups.0.albert_layers.0.attention.key", + "albert_shared_block/value": "encoder.albert_layer_groups.0.albert_layers.0.attention.value", + "albert_shared_block/attention_dense": "encoder.albert_layer_groups.0.albert_layers.0.attention.dense", + "albert_shared_block/attention_layer_norm": "encoder.albert_layer_groups.0.albert_layers.0.attention.LayerNorm", + "albert_shared_block/ffn": "encoder.albert_layer_groups.0.albert_layers.0.ffn", + "albert_shared_block/ffn_output": "encoder.albert_layer_groups.0.albert_layers.0.ffn_output", + "albert_shared_block/full_layer_norm": "encoder.albert_layer_groups.0.albert_layers.0.full_layer_layer_norm", + }, +} + +WEIGHT_SUFFIX_MAPPING = { + "kernel": "weight", + "embeddings": "weight", + "gamma": "weight", + "beta": "bias", +} def transfer_sentence_transformer_weights( @@ -67,7 +68,13 @@ def transfer_sentence_transformer_weights( Pooling and L2-normalize have no parameters. Unused HF pooler / MLM head tensors are skipped via the MPNet optional-weight rules. """ - hf = {strip_st_prefix(k): v for k, v in hf_state_dict.items()} + hf = {} + for key, value in hf_state_dict.items(): + for prefix in ("0.auto_model.", "auto_model.", "0.transformer."): + if key.startswith(prefix): + key = key[len(prefix) :] + break + hf[normalize_hf_key(key)] = value kind = keras_model.backbone_type if kind == "mpnet": transfer_mpnet_weights(keras_model, hf) @@ -78,7 +85,32 @@ def transfer_sentence_transformer_weights( elif kind == "xlm-roberta": transfer_xlm_roberta_weights(keras_model, hf) elif kind in ("distilbert", "albert"): - _transfer_custom_backbone(keras_model, hf) + prefix = f"{kind}." + names = {key.removeprefix(prefix): key for key in hf} + mapping = WEIGHT_NAME_MAPPING[kind] + for weight in keras_model.weights: + if "/" not in weight.path: + continue + layer_path, suffix = weight.path.rsplit("/", 1) + index = None + if kind == "distilbert" and layer_path.startswith("distil_layer_"): + block, component = layer_path.split("/", 1) + index = int(block.removeprefix("distil_layer_")) + layer_path = f"distil_layer_{{index}}/{component}" + hf_name = mapping.get(layer_path) + if hf_name is None: + if layer_path.startswith(("distil_", "albert_")): + raise WeightMappingError(weight.path, "WEIGHT_NAME_MAPPING entry") + continue + suffix = WEIGHT_SUFFIX_MAPPING.get(suffix, suffix) + hf_name = f"{hf_name.format(index=index)}.{suffix}" + if hf_name not in names: + raise WeightMappingError(weight.path, hf_name) + source = hf[names[hf_name]] + if weight.path == "albert_embedding_projection/kernel": + weight.assign(np.asarray(source).T) + else: + transfer_weights(weight.path, weight, source) else: raise ValueError(f"Unsupported sentence embedding backbone: {kind}") @@ -87,7 +119,14 @@ def transfer_sentence_transformer_weights( return dense_state = getattr(keras_model, "_st_dense_state", hf_state_dict) - proj_kernel, proj_bias = _find_dense_tensors(dense_state) + proj_kernel = proj_bias = None + for key, value in dense_state.items(): + if "auto_model" in key or ".layer." in key or key.startswith("embeddings."): + continue + if key.endswith("linear.weight"): + proj_kernel = np.asarray(value) + elif key.endswith("linear.bias"): + proj_bias = np.asarray(value) for weight in keras_model.weights: path = weight.path if path == "projection/kernel": @@ -100,85 +139,6 @@ def transfer_sentence_transformer_weights( transfer_weights(path, weight, proj_bias) -def _custom_hf_name(path, backbone_type): - """Map the two ST-only backbone families to their HF tensor names.""" - suffix = path.rsplit("/", 1)[-1] - suffix = { - "kernel": "weight", - "embeddings": "weight", - "gamma": "weight", - "beta": "bias", - }.get(suffix, suffix) - if backbone_type == "distilbert": - fixed = { - "distil_word_embeddings": "embeddings.word_embeddings", - "distil_position_embeddings": "embeddings.position_embeddings", - "distil_embeddings_norm": "embeddings.LayerNorm", - } - if "/" not in path: - return None - outer, inner = path.split("/", 1) - if outer in fixed: - return f"{fixed[outer]}.{suffix}" - if not outer.startswith("distil_layer_"): - return None - index = int(outer.removeprefix("distil_layer_")) - part = inner.split("/", 1)[0] - part = { - "q_lin": "attention.q_lin", - "k_lin": "attention.k_lin", - "v_lin": "attention.v_lin", - "out_lin": "attention.out_lin", - "ffn_lin1": "ffn.lin1", - "ffn_lin2": "ffn.lin2", - }.get(part, part) - return f"transformer.layer.{index}.{part}.{suffix}" - - fixed = { - "albert_word_embeddings": "embeddings.word_embeddings", - "albert_position_embeddings": "embeddings.position_embeddings", - "albert_token_type_embeddings": "embeddings.token_type_embeddings", - "albert_embeddings_norm": "embeddings.LayerNorm", - "albert_embedding_projection": "encoder.embedding_hidden_mapping_in", - } - if "/" not in path: - return None - outer, inner = path.split("/", 1) - if outer in fixed: - return f"{fixed[outer]}.{suffix}" - if outer != "albert_shared_block": - return None - part = inner.split("/", 1)[0] - part = { - "query": "attention.query", - "key": "attention.key", - "value": "attention.value", - "attention_dense": "attention.dense", - "attention_layer_norm": "attention.LayerNorm", - "full_layer_norm": "full_layer_layer_norm", - }.get(part, part) - return f"encoder.albert_layer_groups.0.albert_layers.0.{part}.{suffix}" - - -def _transfer_custom_backbone(keras_model, hf_state_dict): - kind = keras_model.backbone_type - prefix = "distilbert." if kind == "distilbert" else "albert." - names = {k.removeprefix(prefix): k for k in hf_state_dict} - for weight in keras_model.weights: - hf_name = _custom_hf_name(weight.path, kind) - if hf_name is None: - continue - if hf_name not in names: - raise WeightMappingError(weight.path, hf_name) - source = hf_state_dict[names[hf_name]] - if weight.path == "albert_embedding_projection/kernel": - # The generic transfer helper treats any path containing - # "embedding" as an Embedding table, but this is a Dense kernel. - weight.assign(np.asarray(source).T) - else: - transfer_weights(weight.path, weight, source) - - def load_dense_state_dict(hf_id, rel_path, *, token=None, revision=None): """Read the trained Dense module from its declared ST subdirectory.""" if not rel_path or rel_path.startswith("/") or ".." in rel_path.split("/"): From b178e4bd9d268c8b6e1aceb4bcbca148f269201d Mon Sep 17 00:00:00 2001 From: IMvision12 Date: Thu, 15 Oct 2026 18:33:59 -0700 Subject: [PATCH 08/10] Fix act --- .../sentence_transformer/sentence_transformer_backbones.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/zeromodels/st_models/sentence_transformer/sentence_transformer_backbones.py b/zeromodels/st_models/sentence_transformer/sentence_transformer_backbones.py index e24020a9..5d6ddfdc 100644 --- a/zeromodels/st_models/sentence_transformer/sentence_transformer_backbones.py +++ b/zeromodels/st_models/sentence_transformer/sentence_transformer_backbones.py @@ -67,7 +67,7 @@ def call(self, inputs): x = self.sa_layer_norm(x + self.out_lin(context)) hidden = self.ffn_lin1(x) if self.activation == "gelu": - hidden = ops.gelu(hidden) + hidden = ops.gelu(hidden, approximate=False) else: raise ValueError(f"Unsupported DistilBERT activation: {self.activation}") return self.output_layer_norm(x + self.ffn_lin2(hidden)) @@ -125,7 +125,7 @@ def call(self, inputs): if self.activation == "gelu_new": hidden = _gelu_new(hidden) elif self.activation == "gelu": - hidden = ops.gelu(hidden) + hidden = ops.gelu(hidden, approximate=False) else: raise ValueError(f"Unsupported ALBERT activation: {self.activation}") return self.full_layer_norm(x + self.ffn_output(hidden)) From 05721dd45f68d254334fd968a688764e564c6bbd Mon Sep 17 00:00:00 2001 From: IMvision12 Date: Thu, 15 Oct 2026 18:59:20 -0700 Subject: [PATCH 09/10] refactor --- README.md | 4 +- docs/sentence_transformer.md | 49 ++++++---- docs/sentence_transformers.md | 11 ++- pyproject.toml | 2 +- tests/base/model_test_registry.py | 14 +-- tests/integration/test_auto_registry.py | 33 +++++-- .../integration/test_sentence_transformer.py | 93 ++++++++++++++----- .../test_sentence_transformer_tokenizer.py | 41 +++++--- ...t_sentence_transformer_tokenizers_vs_hf.py | 4 +- zeromodels/__init__.py | 11 ++- zeromodels/auto/auto_factory.py | 12 ++- zeromodels/auto/auto_mapping_names.py | 4 +- .../models/sentence_transformer/__init__.py | 20 ++-- zeromodels/sentence_transformers/__init__.py | 37 ++++++++ ...onvert_sentence_transformer_hf_to_keras.py | 12 +-- .../sentence_transformer_backbones.py | 22 ++--- .../sentence_transformer_config.py | 10 +- .../sentence_transformer_layers.py | 39 +++----- .../sentence_transformer_model.py | 37 ++++++-- .../sentence_transformer_tokenizer.py | 15 ++- zeromodels/st_models/__init__.py | 28 +++++- .../sentence_transformer/__init__.py | 15 --- 22 files changed, 339 insertions(+), 174 deletions(-) create mode 100644 zeromodels/sentence_transformers/__init__.py rename zeromodels/{st_models/sentence_transformer => sentence_transformers}/convert_sentence_transformer_hf_to_keras.py (97%) rename zeromodels/{st_models/sentence_transformer => sentence_transformers}/sentence_transformer_backbones.py (95%) rename zeromodels/{st_models/sentence_transformer => sentence_transformers}/sentence_transformer_config.py (77%) rename zeromodels/{st_models/sentence_transformer => sentence_transformers}/sentence_transformer_layers.py (69%) rename zeromodels/{st_models/sentence_transformer => sentence_transformers}/sentence_transformer_model.py (94%) rename zeromodels/{st_models/sentence_transformer => sentence_transformers}/sentence_transformer_tokenizer.py (96%) delete mode 100644 zeromodels/st_models/sentence_transformer/__init__.py diff --git a/README.md b/README.md index 4d0006fc..ae038a46 100644 --- a/README.md +++ b/README.md @@ -40,9 +40,9 @@ Classification backbones share a single [documentation page](https://imvision12. Documentation sources are also available in [`docs/`](docs/). -## 💬 Sentence Transformers Models (st_models) +## 💬 Sentence Transformers Models (sentence_transformers) -All 13 checkpoints below have pre-converted Keras weights under the `zeromodels` Hugging Face organization. Load them through `zeromodels.st_models` for sentence similarity, semantic search, or clustering. +All 13 checkpoints below have pre-converted Keras weights under the `zeromodels` Hugging Face organization. Use `ZMSentenceTransformer` and `ZMSentenceTransformerTokenizer` from `zeromodels.sentence_transformers` for sentence similarity, semantic search, or clustering. | 🏷️ Model Name | 📜 Reference Paper | 📦 Source of Weights | |---------------|-------------------|---------------------| diff --git a/docs/sentence_transformer.md b/docs/sentence_transformer.md index 29a8e6f5..cf418b52 100644 --- a/docs/sentence_transformer.md +++ b/docs/sentence_transformer.md @@ -18,6 +18,14 @@ optional Dense projection, and optional L2 normalization. Parity has been verified on the Keras Torch backend; TensorFlow and JAX parity have not yet been checked for all 13 checkpoints. +The public API is `zeromodels.sentence_transformers`, separate from the upstream +`sentence_transformers` reference library. The former `zeromodels.st_models` +and `zeromodels.models.sentence_transformer` paths remain compatibility aliases +for existing code and saved metadata; new examples and exports use the new namespace. +Use `ZMSentenceTransformer` and `ZMSentenceTransformerTokenizer`; the former +`SentenceTransformer`, `SentenceTransformerModel`, and `SentenceTransformerTokenizer` +names remain compatibility aliases. + ## Usage ```python @@ -25,14 +33,14 @@ import os os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" -from zeromodels.st_models import ( - SentenceTransformer, - SentenceTransformerTokenizer, +from zeromodels.sentence_transformers import ( + ZMSentenceTransformer, + ZMSentenceTransformerTokenizer, ) repo = "zeromodels/all-MiniLM-L6-v2" -model = SentenceTransformer.from_weights(repo) -tokenizer = SentenceTransformerTokenizer.from_weights(repo) +model = ZMSentenceTransformer.from_weights(repo) +tokenizer = ZMSentenceTransformerTokenizer.from_weights(repo) embeddings = model.encode( ["A man is eating food.", "A man is eating pasta."], @@ -56,8 +64,8 @@ records the surrounding official configuration in two separate blocks: ```json { "library_name": "zeromodels", - "preprocessor_module": "zeromodels.st_models", - "preprocessor_class": "SentenceTransformerTokenizer", + "preprocessor_module": "zeromodels.sentence_transformers", + "preprocessor_class": "ZMSentenceTransformerTokenizer", "tokenizer_file": "tokenizer.json", "backbone_type": "bert", "tokenizer_config": { @@ -89,6 +97,8 @@ An explicit `max_seq_len=` overrides the stored limit. Legacy repos without `write_zm_preprocessor(directory, tokenizer, variant)` exports a portable recipe referencing `tokenizer.json` in the same repo, rather than a local cache path. +New exports set both `model_module` in `zm_config.json` and `preprocessor_module` +in `zm_preprocessor.json` to `zeromodels.sentence_transformers`. Text preprocessing uses the `tokenizers` library; `transformers` is not needed to load these converted checkpoints. @@ -100,8 +110,8 @@ tokenizer: ```python upstream = "hf:sentence-transformers/all-MiniLM-L6-v2" -model = SentenceTransformer.from_weights(upstream) -tokenizer = SentenceTransformerTokenizer.from_weights( +model = ZMSentenceTransformer.from_weights(upstream) +tokenizer = ZMSentenceTransformerTokenizer.from_weights( upstream, revision=model.source_revision ) ``` @@ -122,13 +132,20 @@ supported by this model. ## Parity check The optional `sentence-transformers` dependency is needed only for comparison -against the original implementation: +against the original implementation. After the Usage example above, compare +the published ZeroModels weights with the matching original checkpoint: -```bash -python scripts/verify_sentence_transformer_parity.py --model all +```python +import numpy as np +from sentence_transformers import SentenceTransformer as ReferenceSentenceTransformer + +texts = ["A man is eating food.", "A man is eating pasta."] +reference = ReferenceSentenceTransformer("sentence-transformers/all-MiniLM-L6-v2") +expected = reference.encode(texts, normalize_embeddings=False) +actual = model.encode(texts, tokenizer=tokenizer) +print("Maximum absolute difference:", np.max(np.abs(actual - expected))) ``` -The script compares embeddings at pinned upstream commits for all 13 rows. -Use `--model ` to check one checkpoint. It downloads upstream weights -into the standard Hugging Face cache but does not create or publish converted -releases. +Use matching checkpoint names from the models table, compare in float32, and +record Hub commit SHAs for reproducible results. This downloads original +weights for comparison; it does not convert or publish a new checkpoint. diff --git a/docs/sentence_transformers.md b/docs/sentence_transformers.md index 84536f69..5a37ee8b 100644 --- a/docs/sentence_transformers.md +++ b/docs/sentence_transformers.md @@ -3,7 +3,7 @@ ZeroModels provides 13 converted dense sentence-embedding checkpoints. Each checkpoint combines a text encoder with its original pooling, optional Dense projection, and optional L2 normalization. Load them through the -[`zeromodels.st_models` API](sentence_transformer.md); no conversion is needed +[`zeromodels.sentence_transformers` API](sentence_transformer.md); no conversion is needed for the ZeroModels-hosted weights. | Checkpoint | Backbone | Embedding recipe | Dimensions | @@ -32,11 +32,14 @@ import os os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" -from zeromodels.st_models import SentenceTransformer, SentenceTransformerTokenizer +from zeromodels.sentence_transformers import ( + ZMSentenceTransformer, + ZMSentenceTransformerTokenizer, +) repo = "zeromodels/all-MiniLM-L6-v2" -model = SentenceTransformer.from_weights(repo) -tokenizer = SentenceTransformerTokenizer.from_weights(repo) +model = ZMSentenceTransformer.from_weights(repo) +tokenizer = ZMSentenceTransformerTokenizer.from_weights(repo) embeddings = model.encode(["A cat on a mat.", "A dog on a mat."], tokenizer=tokenizer) print(embeddings.shape) # (2, 384) diff --git a/pyproject.toml b/pyproject.toml index d87a7b03..019f301d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -104,7 +104,7 @@ version = { attr = "zeromodels.__version__" } source = ["zeromodels", "tests"] omit = [ "zeromodels/models/*/convert_*.py", - "zeromodels/st_models/*/convert_*.py", + "zeromodels/sentence_transformers/convert_*.py", ] [tool.coverage.report] diff --git a/tests/base/model_test_registry.py b/tests/base/model_test_registry.py index 41703485..673c0490 100644 --- a/tests/base/model_test_registry.py +++ b/tests/base/model_test_registry.py @@ -2997,9 +2997,9 @@ "expected_output_shape": (2, 3), }, # ---- Sentence Transformers (MPNet backbone) ---- - "SentenceTransformerModel": { - "module": "zeromodels.st_models", - "model_cls": "SentenceTransformerModel", + "ZMSentenceTransformer": { + "module": "zeromodels.sentence_transformers", + "model_cls": "ZMSentenceTransformer", "model_type": "llm", "init_kwargs": { "backbone_type": "mpnet", @@ -3022,7 +3022,7 @@ }, }, "SentenceTransformerTextEmbed": { - "module": "zeromodels.st_models", + "module": "zeromodels.sentence_transformers", "model_cls": "SentenceTransformerTextEmbed", "model_type": "llm", "init_kwargs": { @@ -4488,9 +4488,9 @@ if _backbone in ("roberta", "xlm-roberta") else "bert_input" ) - MODEL_TEST_CONFIGS[f"SentenceTransformer_{_variant.replace('-', '_')}"] = { - "module": "zeromodels.st_models", - "model_cls": "SentenceTransformerModel", + MODEL_TEST_CONFIGS[f"ZMSentenceTransformer_{_variant.replace('-', '_')}"] = { + "module": "zeromodels.sentence_transformers", + "model_cls": "ZMSentenceTransformer", "model_type": "llm", "init_kwargs": _kwargs, "input_factory": _factory, diff --git a/tests/integration/test_auto_registry.py b/tests/integration/test_auto_registry.py index 01011d23..43c41e6c 100644 --- a/tests/integration/test_auto_registry.py +++ b/tests/integration/test_auto_registry.py @@ -26,6 +26,7 @@ "DptSemanticSegment": "SemanticSegment", "DptDensePredict": "DensePredict", "UnifiedTextGenerate": "TextGenerate", + "SentenceTransformer": "Model", } _TASK_SUFFIXES = sorted( { @@ -62,6 +63,7 @@ "TextModel", "VisionModel", "EncoderModel", + "SentenceTransformer", "Model", }, key=len, @@ -85,16 +87,27 @@ def _hf_types(cls): def _iter_model_classes(): import zeromodels.models as models - import zeromodels.st_models as st_models + import zeromodels.sentence_transformers as sentence_transformers from zeromodels.base.base_model import BaseModel - for namespace in (models, st_models): - for family in sorted(n for n in dir(namespace) if not n.startswith("_")): - package = getattr(namespace, family) + seen = set() + for namespace in (models, sentence_transformers): + packages = [namespace] + packages.extend( + getattr(namespace, family) + for family in sorted(n for n in dir(namespace) if not n.startswith("_")) + ) + for package in packages: for name in getattr(package, "__all__", []): obj = getattr(package, name, None) - if inspect.isclass(obj) and issubclass(obj, BaseModel): - yield name, obj + if ( + inspect.isclass(obj) + and issubclass(obj, BaseModel) + and obj not in seen + ): + seen.add(obj) + # Compatibility aliases are not separate model classes. + yield obj.__name__, obj def test_top_level_and_task_autos_exported(): @@ -142,6 +155,7 @@ def test_no_transformers_name_collision(): # Versioned families are disambiguated by the distinct config model_type. ("AutoZModel", "deberta_v2", "DebertaV2Model"), ("AutoZModel", "deberta_v3", "DebertaV3Model"), + ("AutoZModel", "sentence_transformer", "ZMSentenceTransformer"), ], ) def test_key_resolutions(auto_name, model_type, expected_cls): @@ -197,6 +211,9 @@ def test_no_key_in_both_single_and_ambiguous(): def test_preprocessor_and_config_registries(): assert zm.AutoZMConfig.for_model_type("detr").__name__ == "DetrConfig" assert zm.AutoZMTokenizer.mapping()["bert"].__name__ == "BertTokenizer" + assert zm.AutoZMTokenizer.mapping()["sentence_transformer"].__name__ == ( + "ZMSentenceTransformerTokenizer" + ) assert zm.AutoZMImageProcessor.mapping()["detr"].__name__ == "DETRImageProcessor" assert zm.AutoZMProcessor.mapping()["clip"].__name__ == "CLIPProcessor" @@ -430,7 +447,9 @@ def test_no_cross_family_model_type(): model. The genuinely-shared model_types are allowlisted with their exact family set.""" def family(class_name): - return A._resolve_class(class_name).__module__.split(".")[2] + parts = A._resolve_class(class_name).__module__.split(".") + # Model-zoo families are nested; standalone APIs are a top-level package. + return parts[2] if parts[1] == "models" else parts[1] per_type = {} for task, table in names.MODEL_TASK_MAPPING_NAMES.items(): diff --git a/tests/integration/test_sentence_transformer.py b/tests/integration/test_sentence_transformer.py index bb454100..7a03d2d6 100644 --- a/tests/integration/test_sentence_transformer.py +++ b/tests/integration/test_sentence_transformer.py @@ -5,32 +5,63 @@ import numpy as np import pytest -from keras import ops +from keras import ops, saving from tests.base.model_test_registry import MODEL_TEST_CONFIGS -from zeromodels.st_models import ( - SentenceTransformer, - SentenceTransformerModel, +from zeromodels.sentence_transformers import ( SentenceTransformerTextEmbed, + ZMSentenceTransformer, + ZMSentenceTransformerTokenizer, ) -def test_top_level_st_api_and_legacy_imports_resolve_to_same_classes(): - legacy = importlib.import_module("zeromodels.models.sentence_transformer") +@pytest.mark.parametrize( + "module", + ( + "zeromodels.models.sentence_transformer", + "zeromodels.st_models", + "zeromodels.st_models.sentence_transformer", + "zeromodels.sentence_transformers.sentence_transformer", + ), +) +def test_top_level_st_api_and_legacy_imports_resolve_to_same_classes(module): + legacy = importlib.import_module(module) + package = ( + module + ".sentence_transformer" if module.endswith("st_models") else module + ) legacy_model_module = importlib.import_module( - "zeromodels.models.sentence_transformer.sentence_transformer_model" + package + ".sentence_transformer_model" ) - assert SentenceTransformer is SentenceTransformerModel - assert legacy.SentenceTransformerModel is SentenceTransformerModel - assert legacy_model_module.SentenceTransformerModel is SentenceTransformerModel + assert ZMSentenceTransformer.__name__ == "ZMSentenceTransformer" + assert ZMSentenceTransformerTokenizer.__name__ == "ZMSentenceTransformerTokenizer" + assert legacy.ZMSentenceTransformer is ZMSentenceTransformer + assert legacy_model_module.ZMSentenceTransformer is ZMSentenceTransformer + assert legacy.SentenceTransformerModel is ZMSentenceTransformer + assert legacy_model_module.SentenceTransformerModel is ZMSentenceTransformer + assert legacy.SentenceTransformerTokenizer is ZMSentenceTransformerTokenizer -def test_existing_hub_config_with_old_module_path_still_builds(monkeypatch): +@pytest.mark.parametrize( + "module", + ( + "zeromodels.models.sentence_transformer", + "zeromodels.st_models.sentence_transformer", + "zeromodels.sentence_transformers.sentence_transformer", + "zeromodels.sentence_transformers", + ), +) +@pytest.mark.parametrize( + "declared", + ("ZMSentenceTransformer", "SentenceTransformerModel", "SentenceTransformer"), +) +def test_existing_hub_config_with_old_module_path_still_builds( + monkeypatch, module, declared +): from zeromodels.conversion import zm_config spec = { - "model_module": "zeromodels.models.sentence_transformer", - "model_class": "SentenceTransformerModel", + "model_module": module, + "model_class": declared, "model_type": "sentence_transformer", "backbone_type": "bert", "vocab_size": 32, @@ -42,13 +73,25 @@ def test_existing_hub_config_with_old_module_path_still_builds(monkeypatch): } monkeypatch.setattr(zm_config, "load_zm_config", lambda repo_id: spec) - model = SentenceTransformer.from_hub_repo( + model = ZMSentenceTransformer.from_hub_repo( "zeromodels/example-sentence-transformer", load_weights=False ) - assert isinstance(model, SentenceTransformerModel) + assert isinstance(model, ZMSentenceTransformer) assert model.config.backbone_type == "bert" +@pytest.mark.parametrize( + "canonical,legacy", + ( + (ZMSentenceTransformer, "SentenceTransformerModel"), + (ZMSentenceTransformerTokenizer, "SentenceTransformerTokenizer"), + ), +) +def test_keras_registration_uses_new_name_and_preserves_legacy_name(canonical, legacy): + assert saving.get_registered_name(canonical) == f"zeromodels>{canonical.__name__}" + assert saving.get_registered_object(f"zeromodels>{legacy}") is canonical + + def test_all_published_sentence_transformer_variants_are_in_offline_registry(): published = { "all-mpnet-base-v2", @@ -66,9 +109,9 @@ def test_all_published_sentence_transformer_variants_are_in_offline_registry(): "distiluse-base-multilingual-cased-v2", } registered = { - name.removeprefix("SentenceTransformer_") + name.removeprefix("ZMSentenceTransformer_") for name in MODEL_TEST_CONFIGS - if name.startswith("SentenceTransformer_") + if name.startswith("ZMSentenceTransformer_") } assert registered == {name.replace("-", "_") for name in published} @@ -78,7 +121,7 @@ def test_all_published_sentence_transformer_variants_are_in_offline_registry(): ("mpnet", "bert", "roberta", "xlm-roberta", "distilbert", "albert"), ) def test_backbone_builds_and_returns_sentence_embeddings(backbone_type): - model = SentenceTransformerModel( + model = ZMSentenceTransformer( backbone_type=backbone_type, vocab_size=32, embed_dim=16, @@ -115,7 +158,7 @@ def test_pooling_parser_accepts_exactly_one_supported_mode( "huggingface_hub.hf_hub_download", lambda *args, **kwargs: str(config_path) ) assert ( - SentenceTransformerModel._pooling_from_module_dir("example/model", "1_Pooling") + ZMSentenceTransformer._pooling_from_module_dir("example/model", "1_Pooling") == pooling ) @@ -130,7 +173,7 @@ def test_pooling_parser_rejects_ambiguous_modes(monkeypatch, tmp_path): "huggingface_hub.hf_hub_download", lambda *args, **kwargs: str(config_path) ) with pytest.raises(ValueError, match="pooling config"): - SentenceTransformerModel._pooling_from_module_dir("example/model", "1_Pooling") + ZMSentenceTransformer._pooling_from_module_dir("example/model", "1_Pooling") def test_sentence_transformer_registered_for_text_embed_task(): @@ -143,7 +186,7 @@ def test_sentence_transformer_registered_for_text_embed_task(): def test_dot_similarity_is_not_cosine_normalized(): - model = SentenceTransformerModel( + model = ZMSentenceTransformer( backbone_type="bert", vocab_size=32, embed_dim=8, @@ -160,7 +203,7 @@ def test_dot_similarity_is_not_cosine_normalized(): def test_cosine_similarity_accepts_numpy_embeddings(): - model = SentenceTransformerModel( + model = ZMSentenceTransformer( backbone_type="bert", vocab_size=32, embed_dim=8, @@ -184,7 +227,7 @@ def test_sentence_transformer_h5_weights_round_trip(tmp_path): "mlp_dim": 16, "max_position_embeddings": 16, } - model = SentenceTransformerModel(**kwargs) + model = ZMSentenceTransformer(name="SentenceTransformerModel", **kwargs) inputs = { "input_ids": np.array([[2, 3, 4]], dtype="int32"), "attention_mask": np.ones((1, 3), dtype="int32"), @@ -194,7 +237,9 @@ def test_sentence_transformer_h5_weights_round_trip(tmp_path): path = tmp_path / "model.weights.h5" model.save_weights(str(path)) - reloaded = SentenceTransformerModel(**model.config.constructor_kwargs()) + reloaded = ZMSentenceTransformer(**model.config.constructor_kwargs()) + assert model.name == "SentenceTransformerModel" + assert reloaded.name == "ZMSentenceTransformer" reloaded.load_weights(str(path)) actual = ops.convert_to_numpy(reloaded(inputs)["embeddings"]) np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-5) diff --git a/tests/integration/test_sentence_transformer_tokenizer.py b/tests/integration/test_sentence_transformer_tokenizer.py index 91e12c53..f4b615e0 100644 --- a/tests/integration/test_sentence_transformer_tokenizer.py +++ b/tests/integration/test_sentence_transformer_tokenizer.py @@ -10,7 +10,7 @@ from keras import ops from tokenizers import Tokenizer, models, normalizers, pre_tokenizers, processors -from zeromodels.st_models import SentenceTransformerTokenizer +from zeromodels.sentence_transformers import ZMSentenceTransformerTokenizer SHA = "a" * 40 @@ -38,8 +38,8 @@ def tokenizer_path(tmp_path): def metadata(): return { - "preprocessor_module": "zeromodels.st_models", - "preprocessor_class": "SentenceTransformerTokenizer", + "preprocessor_module": "zeromodels.sentence_transformers", + "preprocessor_class": "ZMSentenceTransformerTokenizer", "tokenizer_file": "assets/tokenizer.json", "backbone_type": "bert", "tokenizer_config": { @@ -79,18 +79,23 @@ def download(repo_id, filename, token=None, revision=None): return calls +@pytest.mark.parametrize( + "declared", ("ZMSentenceTransformerTokenizer", "SentenceTransformerTokenizer") +) def test_hub_metadata_uses_one_snapshot_and_official_sequence_limit( - monkeypatch, tmp_path, tokenizer_path + monkeypatch, tmp_path, tokenizer_path, declared ): + spec = metadata() + spec["preprocessor_class"] = declared calls = mock_repo( monkeypatch, tmp_path, { - "zm_preprocessor.json": metadata(), + "zm_preprocessor.json": spec, "assets/tokenizer.json": tokenizer_path, }, ) - tok = SentenceTransformerTokenizer.from_weights( + tok = ZMSentenceTransformerTokenizer.from_weights( "zeromodels/example", revision="release" ) encoded = tok(["hello world " * 10, "hello"]) @@ -115,7 +120,7 @@ def test_portable_preprocessor_export_loads_again( ): from zeromodels.conversion.zm_config import write_zm_preprocessor - original = SentenceTransformerTokenizer( + original = ZMSentenceTransformerTokenizer( tokenizer_file=str(tokenizer_path), backbone_type="bert", max_seq_len=6, @@ -125,13 +130,15 @@ def test_portable_preprocessor_export_loads_again( path = write_zm_preprocessor(tmp_path, original, "example") with open(path, encoding="utf-8") as stream: spec = json.load(stream) + assert spec["preprocessor_module"] == "zeromodels.sentence_transformers" assert spec["tokenizer_file"] == "tokenizer.json" + assert spec["preprocessor_class"] == "ZMSentenceTransformerTokenizer" mock_repo( monkeypatch, tmp_path, {"zm_preprocessor.json": spec, "tokenizer.json": tokenizer_path}, ) - reloaded = SentenceTransformerTokenizer.from_weights("example", revision=SHA) + reloaded = ZMSentenceTransformerTokenizer.from_weights("example", revision=SHA) for key, expected in original(["Hello world", "hello"]).items(): np.testing.assert_array_equal( ops.convert_to_numpy(reloaded(["Hello world", "hello"])[key]), @@ -153,7 +160,9 @@ def test_legacy_nested_config_and_explicit_limit_keep_backbone( "tokenizer.json": tokenizer_path, }, ) - tok = SentenceTransformerTokenizer.from_weights("zeromodels/example", max_seq_len=5) + tok = ZMSentenceTransformerTokenizer.from_weights( + "zeromodels/example", max_seq_len=5 + ) assert tok.max_seq_len == 5 assert tok.backbone_type == "bert" assert "token_type_ids" in tok(["hello"]) @@ -163,12 +172,12 @@ def test_sentence_transformer_lowercase_is_separate_from_tokenizer_normalization tokenizer_path, ): tc = {"do_lower_case": False, "pad_token": "[PAD]"} - plain = SentenceTransformerTokenizer( + plain = ZMSentenceTransformerTokenizer( tokenizer_file=str(tokenizer_path), tokenizer_config=tc, sentence_transformer_config={"do_lower_case": False}, ) - extra = SentenceTransformerTokenizer( + extra = ZMSentenceTransformerTokenizer( tokenizer_file=str(tokenizer_path), tokenizer_config=tc, sentence_transformer_config={"do_lower_case": True}, @@ -178,7 +187,7 @@ def test_sentence_transformer_lowercase_is_separate_from_tokenizer_normalization raw = Tokenizer.from_file(str(tokenizer_path)) raw.normalizer = normalizers.BertNormalizer(lowercase=True) raw.save(str(tokenizer_path)) - inherited = SentenceTransformerTokenizer( + inherited = ZMSentenceTransformerTokenizer( tokenizer_file=str(tokenizer_path), tokenizer_config={"do_lower_case": True}, sentence_transformer_config={"do_lower_case": False}, @@ -187,7 +196,7 @@ def test_sentence_transformer_lowercase_is_separate_from_tokenizer_normalization def test_padding_truncation_pairs_and_serialization(tokenizer_path): - tok = SentenceTransformerTokenizer( + tok = ZMSentenceTransformerTokenizer( tokenizer_file=str(tokenizer_path), backbone_type="bert", tokenizer_config={"padding_side": "left", "truncation_side": "left"}, @@ -227,7 +236,9 @@ def test_upstream_module_path_and_sidecars(monkeypatch, tmp_path, tokenizer_path "0_Transformer/tokenizer.json": tokenizer_path, }, ) - tok = SentenceTransformerTokenizer.from_weights("hf:sentence-transformers/example") + tok = ZMSentenceTransformerTokenizer.from_weights( + "hf:sentence-transformers/example" + ) assert tok.max_seq_len == 5 and tok.cls_token_id == 2 assert tok.tokenize("Hello") == [5] assert all(revision == SHA for _, revision in calls) @@ -247,4 +258,4 @@ def test_hub_rejects_non_repository_tokenizer_paths(monkeypatch, tmp_path, filen spec["tokenizer_file"] = filename mock_repo(monkeypatch, tmp_path, {"zm_preprocessor.json": spec}) with pytest.raises(ValueError, match="repository-relative"): - SentenceTransformerTokenizer.from_weights("zeromodels/example") + ZMSentenceTransformerTokenizer.from_weights("zeromodels/example") diff --git a/tests/processors/test_sentence_transformer_tokenizers_vs_hf.py b/tests/processors/test_sentence_transformer_tokenizers_vs_hf.py index 94f48bee..9645ca0c 100644 --- a/tests/processors/test_sentence_transformer_tokenizers_vs_hf.py +++ b/tests/processors/test_sentence_transformer_tokenizers_vs_hf.py @@ -9,7 +9,7 @@ from keras import ops from tests.base.model_test_registry import _SENTENCE_TRANSFORMER_VARIANTS -from zeromodels.st_models import SentenceTransformerTokenizer +from zeromodels.sentence_transformers import ZMSentenceTransformerTokenizer transformers = pytest.importorskip("transformers") @@ -28,7 +28,7 @@ @pytest.mark.parametrize("name", list(_SENTENCE_TRANSFORMER_VARIANTS)) def test_official_tokenizer_parity_without_weights(name): repo = f"sentence-transformers/{name}" - ours = SentenceTransformerTokenizer.from_hf(repo) + ours = ZMSentenceTransformerTokenizer.from_hf(repo) reference = transformers.AutoTokenizer.from_pretrained( repo, revision=ours.source_revision, use_fast=True, trust_remote_code=False ) diff --git a/zeromodels/__init__.py b/zeromodels/__init__.py index 0069ac5d..d85bd976 100644 --- a/zeromodels/__init__.py +++ b/zeromodels/__init__.py @@ -1,4 +1,12 @@ -from zeromodels import auto, conversion, models, samplers, st_models, utils +from zeromodels import ( + auto, + conversion, + models, + samplers, + sentence_transformers, + st_models, + utils, +) from zeromodels.auto import ( AutoZMConfig, AutoZMImageProcessor, @@ -14,6 +22,7 @@ "auto", "models", "samplers", + "sentence_transformers", "st_models", "utils", "conversion", diff --git a/zeromodels/auto/auto_factory.py b/zeromodels/auto/auto_factory.py index 328e6b2e..45f71086 100644 --- a/zeromodels/auto/auto_factory.py +++ b/zeromodels/auto/auto_factory.py @@ -104,12 +104,16 @@ def _class_index(): global _CLASS_INDEX if _CLASS_INDEX is None: import zeromodels.models as models - import zeromodels.st_models as st_models + import zeromodels.sentence_transformers as sentence_transformers index = {} - for namespace in (models, st_models): - for family in sorted(n for n in dir(namespace) if not n.startswith("_")): - package = getattr(namespace, family) + for namespace in (models, sentence_transformers): + packages = [namespace] + packages.extend( + getattr(namespace, family) + for family in sorted(n for n in dir(namespace) if not n.startswith("_")) + ) + for package in packages: for name in getattr(package, "__all__", []): obj = getattr(package, name, None) if inspect.isclass(obj): diff --git a/zeromodels/auto/auto_mapping_names.py b/zeromodels/auto/auto_mapping_names.py index 04e98093..edfee13d 100644 --- a/zeromodels/auto/auto_mapping_names.py +++ b/zeromodels/auto/auto_mapping_names.py @@ -286,7 +286,7 @@ "sam3": "SAM3Model", "segformer": "SegFormerModel", "senet": "SENetModel", - "sentence_transformer": "SentenceTransformerModel", + "sentence_transformer": "ZMSentenceTransformer", "siglip": "SigLIPModel", "siglip2": "SigLIP2Model", "stable_diffusion": "StableDiffusionModel", @@ -804,7 +804,7 @@ "qwen3_vl_moe_text": "Qwen3VLMoeTokenizer", "roberta": "RobertaTokenizer", "sam3": "SAM3CLIPTokenizer", - "sentence_transformer": "SentenceTransformerTokenizer", + "sentence_transformer": "ZMSentenceTransformerTokenizer", "siglip": "SigLIPTokenizer", "siglip2": "SigLIP2Tokenizer", "speech_to_text": "Speech2TextTokenizer", diff --git a/zeromodels/models/sentence_transformer/__init__.py b/zeromodels/models/sentence_transformer/__init__.py index 48653db8..a9f861bc 100644 --- a/zeromodels/models/sentence_transformer/__init__.py +++ b/zeromodels/models/sentence_transformer/__init__.py @@ -1,29 +1,31 @@ -"""Compatibility imports for Sentence Transformer models moved to ``st_models``.""" +"""Compatibility imports for models moved to ``zeromodels.sentence_transformers``.""" import sys -from zeromodels.st_models import ( +from zeromodels.sentence_transformers import ( SentenceTransformerConfig, SentenceTransformerModel, SentenceTransformerTextEmbed, SentenceTransformerTokenizer, + ZMSentenceTransformer, + ZMSentenceTransformerTokenizer, ) -from zeromodels.st_models.sentence_transformer import ( +from zeromodels.sentence_transformers import ( convert_sentence_transformer_hf_to_keras as _converter, ) -from zeromodels.st_models.sentence_transformer import ( +from zeromodels.sentence_transformers import ( sentence_transformer_backbones as _backbones, ) -from zeromodels.st_models.sentence_transformer import ( +from zeromodels.sentence_transformers import ( sentence_transformer_config as _config, ) -from zeromodels.st_models.sentence_transformer import ( +from zeromodels.sentence_transformers import ( sentence_transformer_layers as _layers, ) -from zeromodels.st_models.sentence_transformer import ( +from zeromodels.sentence_transformers import ( sentence_transformer_model as _model, ) -from zeromodels.st_models.sentence_transformer import ( +from zeromodels.sentence_transformers import ( sentence_transformer_tokenizer as _tokenizer, ) @@ -44,4 +46,6 @@ "SentenceTransformerModel", "SentenceTransformerTextEmbed", "SentenceTransformerTokenizer", + "ZMSentenceTransformer", + "ZMSentenceTransformerTokenizer", ] diff --git a/zeromodels/sentence_transformers/__init__.py b/zeromodels/sentence_transformers/__init__.py new file mode 100644 index 00000000..d55870ec --- /dev/null +++ b/zeromodels/sentence_transformers/__init__.py @@ -0,0 +1,37 @@ +"""Sentence Transformer models and tokenizers for the ZeroModels embedding API.""" + +import importlib +import sys + +from .sentence_transformer_config import SentenceTransformerConfig +from .sentence_transformer_model import ( + SentenceTransformerTextEmbed, + ZMSentenceTransformer, +) +from .sentence_transformer_tokenizer import ZMSentenceTransformerTokenizer + +SentenceTransformer = ZMSentenceTransformer +SentenceTransformerModel = ZMSentenceTransformer +SentenceTransformerTokenizer = ZMSentenceTransformerTokenizer + +# Keep the former nested import paths usable without a nested implementation. +sentence_transformer = sys.modules[__name__] +sys.modules[f"{__name__}.sentence_transformer"] = sentence_transformer +for _name in ( + "convert_sentence_transformer_hf_to_keras", + "sentence_transformer_backbones", + "sentence_transformer_config", + "sentence_transformer_layers", + "sentence_transformer_model", + "sentence_transformer_tokenizer", +): + sys.modules[f"{__name__}.sentence_transformer.{_name}"] = importlib.import_module( + f"{__name__}.{_name}" + ) + +__all__ = [ + "SentenceTransformerConfig", + "SentenceTransformerTextEmbed", + "ZMSentenceTransformer", + "ZMSentenceTransformerTokenizer", +] diff --git a/zeromodels/st_models/sentence_transformer/convert_sentence_transformer_hf_to_keras.py b/zeromodels/sentence_transformers/convert_sentence_transformer_hf_to_keras.py similarity index 97% rename from zeromodels/st_models/sentence_transformer/convert_sentence_transformer_hf_to_keras.py rename to zeromodels/sentence_transformers/convert_sentence_transformer_hf_to_keras.py index d500ac7d..4a41e88f 100644 --- a/zeromodels/st_models/sentence_transformer/convert_sentence_transformer_hf_to_keras.py +++ b/zeromodels/sentence_transformers/convert_sentence_transformer_hf_to_keras.py @@ -178,9 +178,9 @@ def pooling_mode_from_config(pooling_cfg: dict) -> str: from transformers import AutoModel, AutoTokenizer from zeromodels.conversion.zm_config import write_zm_config - from zeromodels.st_models import ( - SentenceTransformerModel, - SentenceTransformerTokenizer, + from zeromodels.sentence_transformers import ( + ZMSentenceTransformer, + ZMSentenceTransformerTokenizer, ) HF_TOKEN = os.environ.get("HF_TOKEN") @@ -197,8 +197,8 @@ def pooling_mode_from_config(pooling_cfg: dict) -> str: for variant, hf_id in VARIANTS.items(): print(f"\n{'=' * 60}\nConverting: {variant} <- {hf_id}\n{'=' * 60}") - keras_model = SentenceTransformerModel.from_hf(hf_id) - tok = SentenceTransformerTokenizer.from_hf( + keras_model = ZMSentenceTransformer.from_hf(hf_id) + tok = ZMSentenceTransformerTokenizer.from_hf( hf_id, max_seq_len=keras_model.max_seq_len ) k_inputs = tok(texts) @@ -237,7 +237,7 @@ def pooling_mode_from_config(pooling_cfg: dict) -> str: keras_model.save_weights(out_path) write_zm_config( dest, - SentenceTransformerModel, + ZMSentenceTransformer, variant, keras_model.config, weights_filename="model.weights.h5", diff --git a/zeromodels/st_models/sentence_transformer/sentence_transformer_backbones.py b/zeromodels/sentence_transformers/sentence_transformer_backbones.py similarity index 95% rename from zeromodels/st_models/sentence_transformer/sentence_transformer_backbones.py rename to zeromodels/sentence_transformers/sentence_transformer_backbones.py index 5d6ddfdc..6bd10e5e 100644 --- a/zeromodels/st_models/sentence_transformer/sentence_transformer_backbones.py +++ b/zeromodels/sentence_transformers/sentence_transformer_backbones.py @@ -1,5 +1,3 @@ -"""Encoder backbones for the supported Sentence Transformer checkpoints.""" - import math import keras @@ -25,14 +23,6 @@ def split_heads(x): return ops.reshape(context, (batch, length, num_heads * head_dim)) -def _gelu_new(x): - return ( - 0.5 - * x - * (1.0 + ops.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * ops.power(x, 3)))) - ) - - @keras.saving.register_keras_serializable(package="zeromodels") class DistilBertEncoderBlock(layers.Layer): """One post-LayerNorm DistilBERT transformer layer.""" @@ -123,7 +113,17 @@ def call(self, inputs): x = self.attention_layer_norm(x + self.attention_dense(context)) hidden = self.ffn(x) if self.activation == "gelu_new": - hidden = _gelu_new(hidden) + hidden = ( + 0.5 + * hidden + * ( + 1.0 + + ops.tanh( + math.sqrt(2.0 / math.pi) + * (hidden + 0.044715 * ops.power(hidden, 3)) + ) + ) + ) elif self.activation == "gelu": hidden = ops.gelu(hidden, approximate=False) else: diff --git a/zeromodels/st_models/sentence_transformer/sentence_transformer_config.py b/zeromodels/sentence_transformers/sentence_transformer_config.py similarity index 77% rename from zeromodels/st_models/sentence_transformer/sentence_transformer_config.py rename to zeromodels/sentence_transformers/sentence_transformer_config.py index 1deca5be..89be4331 100644 --- a/zeromodels/st_models/sentence_transformer/sentence_transformer_config.py +++ b/zeromodels/sentence_transformers/sentence_transformer_config.py @@ -1,12 +1,10 @@ -"""Configuration for Sentence Transformer embedding models.""" - from typing import Optional from zeromodels.base import BaseConfig class SentenceTransformerConfig(BaseConfig): - """Configuration for :class:`SentenceTransformerModel`. + """Configuration for :class:`ZMSentenceTransformer`. A Sentence-Transformers-style stack: encoder backbone + pooling (+ optional dense projection) + optional L2 normalize. ``backbone_type`` selects which @@ -26,7 +24,7 @@ class SentenceTransformerConfig(BaseConfig): max_position_embeddings: int = 514 relative_attention_num_buckets: int = 32 hidden_act: str = "gelu" - layer_norm_eps: float = 1e-5 # matches microsoft/mpnet-base / ST hubs + layer_norm_eps: float = 1e-5 pad_token_id: int = 1 type_vocab_size: int = 2 embedding_size: Optional[int] = None @@ -35,9 +33,9 @@ class SentenceTransformerConfig(BaseConfig): dropout: float = 0.0 attention_dropout: float = 0.0 - pooling: str = "mean" # mean | cls | max + pooling: str = "mean" normalize_embeddings: bool = True projection_dim: Optional[int] = None projection_activation: str = "linear" similarity_fn_name: str = "cosine" - max_seq_len: int = 384 # all-mpnet-base-v2 sentence_bert_config.json + max_seq_len: int = 384 diff --git a/zeromodels/st_models/sentence_transformer/sentence_transformer_layers.py b/zeromodels/sentence_transformers/sentence_transformer_layers.py similarity index 69% rename from zeromodels/st_models/sentence_transformer/sentence_transformer_layers.py rename to zeromodels/sentence_transformers/sentence_transformer_layers.py index 405f354e..2806f63f 100644 --- a/zeromodels/st_models/sentence_transformer/sentence_transformer_layers.py +++ b/zeromodels/sentence_transformers/sentence_transformer_layers.py @@ -1,5 +1,3 @@ -"""Pooling and normalization layers for Sentence Transformer models.""" - import keras from keras import layers, ops @@ -35,30 +33,6 @@ def masked_max_pool(token_embeddings, attention_mask): } -def pool_token_embeddings(token_embeddings, attention_mask, pooling="mean"): - if pooling not in POOLERS: - raise ValueError( - f"Unknown pooling={pooling!r}; expected one of {sorted(POOLERS)}" - ) - return POOLERS[pooling](token_embeddings, attention_mask) - - -def l2_normalize(embeddings, axis=-1, eps=1e-12): - norm = ops.sqrt( - ops.maximum(ops.sum(ops.square(embeddings), axis=axis, keepdims=True), eps) - ) - return embeddings / norm - - -def cosine_similarity(a, b, eps=1e-12): - """Pairwise cosine similarity matrix ``(B_a, B_b)`` for row-normalized or raw vectors.""" - a = ops.convert_to_tensor(a) - b = ops.convert_to_tensor(b) - a = l2_normalize(a, eps=eps) - b = l2_normalize(b, eps=eps) - return ops.matmul(a, ops.transpose(b)) - - @keras.saving.register_keras_serializable(package="zeromodels") class SentenceTransformerPooling(layers.Layer): """Mean / CLS / max pool over token embeddings with an attention mask.""" @@ -73,7 +47,11 @@ def __init__(self, pooling="mean", **kwargs): def call(self, inputs): token_embeddings, attention_mask = inputs - return pool_token_embeddings(token_embeddings, attention_mask, self.pooling) + if self.pooling not in POOLERS: + raise ValueError( + f"Unknown pooling={self.pooling!r}; expected one of {sorted(POOLERS)}" + ) + return POOLERS[self.pooling](token_embeddings, attention_mask) def get_config(self): config = super().get_config() @@ -90,7 +68,12 @@ def __init__(self, eps=1e-12, **kwargs): self.eps = eps def call(self, embeddings): - return l2_normalize(embeddings, eps=self.eps) + norm = ops.sqrt( + ops.maximum( + ops.sum(ops.square(embeddings), axis=-1, keepdims=True), self.eps + ) + ) + return embeddings / norm def get_config(self): config = super().get_config() diff --git a/zeromodels/st_models/sentence_transformer/sentence_transformer_model.py b/zeromodels/sentence_transformers/sentence_transformer_model.py similarity index 94% rename from zeromodels/st_models/sentence_transformer/sentence_transformer_model.py rename to zeromodels/sentence_transformers/sentence_transformer_model.py index 240d1325..66998dd9 100644 --- a/zeromodels/st_models/sentence_transformer/sentence_transformer_model.py +++ b/zeromodels/sentence_transformers/sentence_transformer_model.py @@ -19,24 +19,28 @@ from .sentence_transformer_layers import ( SentenceTransformerNormalize, SentenceTransformerPooling, - cosine_similarity, ) -SENTENCE_TRANSFORMER_HUB_SIBLINGS = frozenset({"SentenceTransformerModel"}) +SENTENCE_TRANSFORMER_HUB_SIBLINGS = frozenset( + {"ZMSentenceTransformer", "SentenceTransformerModel", "SentenceTransformer"} +) @keras.saving.register_keras_serializable(package="zeromodels") -class SentenceTransformerModel(BaseModel): +@keras.saving.register_keras_serializable( + package="zeromodels", name="SentenceTransformerModel" +) +class ZMSentenceTransformer(BaseModel): """Sentence embedding model: encoder + pooling (+ optional Dense) + normalize. Matches popular Sentence-Transformers stacks such as ``sentence-transformers/all-mpnet-base-v2`` (MPNet encoder, mean pooling, L2 normalize). Load an upstream checkpoint with:: - model = SentenceTransformerModel.from_weights( + model = ZMSentenceTransformer.from_weights( "hf:sentence-transformers/all-mpnet-base-v2" ) - tok = SentenceTransformerTokenizer.from_weights( + tok = ZMSentenceTransformerTokenizer.from_weights( "hf:sentence-transformers/all-mpnet-base-v2", revision=model.source_revision, ) @@ -96,7 +100,7 @@ def __init__( projection_activation="linear", similarity_fn_name="cosine", max_seq_len=384, - name="SentenceTransformerModel", + name="ZMSentenceTransformer", **kwargs, ): for k in ("model", "hf_id", "url", "num_classes"): @@ -520,7 +524,7 @@ def encode(self, texts, tokenizer, normalize=None, batch_size=32): Args: texts: String or list of strings. - tokenizer: :class:`SentenceTransformerTokenizer` (or any callable that + tokenizer: :class:`ZMSentenceTransformerTokenizer` (or any callable that returns ``input_ids`` / ``attention_mask``). normalize: Override config L2-normalize. ``None`` keeps model output. batch_size: Encode in chunks of this size. @@ -553,7 +557,15 @@ def similarity(self, a, b): """Pairwise scores using this checkpoint's declared similarity function.""" if self.similarity_fn_name == "dot": return ops.matmul(a, ops.transpose(b)) - return cosine_similarity(a, b) + a = ops.convert_to_tensor(a) + b = ops.convert_to_tensor(b) + a = a / ops.sqrt( + ops.maximum(ops.sum(ops.square(a), axis=-1, keepdims=True), 1e-12) + ) + b = b / ops.sqrt( + ops.maximum(ops.sum(ops.square(b), axis=-1, keepdims=True), 1e-12) + ) + return ops.matmul(a, ops.transpose(b)) def get_config(self): config = super().get_config() @@ -562,5 +574,10 @@ def get_config(self): @keras.saving.register_keras_serializable(package="zeromodels") -class SentenceTransformerTextEmbed(SentenceTransformerModel): - """Text-embedding task view of :class:`SentenceTransformerModel`.""" +class SentenceTransformerTextEmbed(ZMSentenceTransformer): + """Text-embedding task view of :class:`ZMSentenceTransformer`.""" + + +# Preserve old deep imports and serialized objects without duplicating classes. +SentenceTransformerModel = ZMSentenceTransformer +SentenceTransformer = ZMSentenceTransformer diff --git a/zeromodels/st_models/sentence_transformer/sentence_transformer_tokenizer.py b/zeromodels/sentence_transformers/sentence_transformer_tokenizer.py similarity index 96% rename from zeromodels/st_models/sentence_transformer/sentence_transformer_tokenizer.py rename to zeromodels/sentence_transformers/sentence_transformer_tokenizer.py index 67612f92..de792fc3 100644 --- a/zeromodels/st_models/sentence_transformer/sentence_transformer_tokenizer.py +++ b/zeromodels/sentence_transformers/sentence_transformer_tokenizer.py @@ -11,7 +11,10 @@ @keras.saving.register_keras_serializable(package="zeromodels") -class SentenceTransformerTokenizer(BaseTokenizer): +@keras.saving.register_keras_serializable( + package="zeromodels", name="SentenceTransformerTokenizer" +) +class ZMSentenceTransformerTokenizer(BaseTokenizer): """Tokenizer for supported ST models, preserving their tokenizer.json rules. Loads ``tokenizer.json`` from a Hub repo (via ``from_weights("org/repo")`` / @@ -183,7 +186,10 @@ def from_hub_repo(cls, repo_id, **kwargs): spec = cls._optional_json(repo_id, "zm_preprocessor.json", token, revision) filename = "tokenizer.json" if spec is not None: - if spec.get("preprocessor_class") != cls.__name__: + if spec.get("preprocessor_class") not in ( + cls.__name__, + "SentenceTransformerTokenizer", + ): raise ValueError( f"{repo_id}: zm_preprocessor.json must declare {cls.__name__}" ) @@ -334,7 +340,10 @@ def _has_lowercase(config): config.get("type") == "Lowercase" or config.get("lowercase") is True or any( - SentenceTransformerTokenizer._has_lowercase(item) + ZMSentenceTransformerTokenizer._has_lowercase(item) for item in config.get("normalizers", []) ) ) + + +SentenceTransformerTokenizer = ZMSentenceTransformerTokenizer diff --git a/zeromodels/st_models/__init__.py b/zeromodels/st_models/__init__.py index fe44e6ae..90b6eba0 100644 --- a/zeromodels/st_models/__init__.py +++ b/zeromodels/st_models/__init__.py @@ -1,11 +1,33 @@ -from .sentence_transformer import ( +"""Compatibility aliases for the renamed ``zeromodels.sentence_transformers`` API.""" + +import importlib +import sys + +from zeromodels.sentence_transformers import ( + SentenceTransformer, SentenceTransformerConfig, SentenceTransformerModel, SentenceTransformerTextEmbed, SentenceTransformerTokenizer, + ZMSentenceTransformer, + ZMSentenceTransformerTokenizer, + sentence_transformer, ) -SentenceTransformer = SentenceTransformerModel +# Saved configs and deep imports must resolve to the canonical classes, not +# load a second copy with a different module name or Keras registration. +sys.modules[f"{__name__}.sentence_transformer"] = sentence_transformer +for _name in ( + "convert_sentence_transformer_hf_to_keras", + "sentence_transformer_backbones", + "sentence_transformer_config", + "sentence_transformer_layers", + "sentence_transformer_model", + "sentence_transformer_tokenizer", +): + sys.modules[f"{__name__}.sentence_transformer.{_name}"] = importlib.import_module( + f"{sentence_transformer.__name__}.{_name}" + ) __all__ = [ "SentenceTransformer", @@ -13,4 +35,6 @@ "SentenceTransformerModel", "SentenceTransformerTextEmbed", "SentenceTransformerTokenizer", + "ZMSentenceTransformer", + "ZMSentenceTransformerTokenizer", ] diff --git a/zeromodels/st_models/sentence_transformer/__init__.py b/zeromodels/st_models/sentence_transformer/__init__.py deleted file mode 100644 index e75d298a..00000000 --- a/zeromodels/st_models/sentence_transformer/__init__.py +++ /dev/null @@ -1,15 +0,0 @@ -"""Implementation of the Sentence Transformer model family.""" - -from .sentence_transformer_config import SentenceTransformerConfig -from .sentence_transformer_model import ( - SentenceTransformerModel, - SentenceTransformerTextEmbed, -) -from .sentence_transformer_tokenizer import SentenceTransformerTokenizer - -__all__ = [ - "SentenceTransformerConfig", - "SentenceTransformerModel", - "SentenceTransformerTextEmbed", - "SentenceTransformerTokenizer", -] From fe31e0e128df5976f9b8abcfa867c15b7e48107e Mon Sep 17 00:00:00 2001 From: IMvision12 Date: Thu, 15 Oct 2026 19:04:22 -0700 Subject: [PATCH 10/10] fix --- docs/auto.md | 23 +++++++++++++++++ docs/sentence_transformer.md | 16 ++++++++++++ .../integration/test_sentence_transformer.py | 25 ++++++++++++++++--- .../test_sentence_transformer_tokenizer.py | 13 +++++++--- zeromodels/auto/__init__.py | 5 ++-- 5 files changed, 74 insertions(+), 8 deletions(-) diff --git a/docs/auto.md b/docs/auto.md index 9e169dfb..94b7304f 100644 --- a/docs/auto.md +++ b/docs/auto.md @@ -44,6 +44,29 @@ fine-tune of the same architecture loads exactly like the official weights. inspect, so pass it to the concrete class instead (`ResNetImageClassify.from_weights("resnet50_a1_in1k")`). +## Sentence Transformer models + +All 13 ZeroModels-hosted Sentence Transformer checkpoints use +`model_type: "sentence_transformer"`. The Auto loaders resolve them to +`ZMSentenceTransformer` and `ZMSentenceTransformerTokenizer`: + +```python +from zeromodels import AutoZModel, AutoZMTokenizer + +repo = "zeromodels/all-MiniLM-L6-v2" +model = AutoZModel.from_weights(repo) # -> ZMSentenceTransformer +tokenizer = AutoZMTokenizer.from_weights(repo) # -> ZMSentenceTransformerTokenizer +embeddings = model.encode(["A cat on a mat.", "A dog on a mat."], tokenizer=tokenizer) +print(embeddings.shape) # (2, 384) +``` + +Here `AutoZModel` loads the complete sentence-embedding stack, including its +pooling, optional Dense projection, and normalization—not just the text encoder. +The tokenizer preserves the checkpoint's own preprocessing recipe. See +[Sentence Transformer](sentence_transformer.md) for the concrete API. Original +upstream `hf:sentence-transformers/...` configs describe their backbone type; +use the concrete classes to load their sentence-embedding recipe. + ## Task classes ### Pipeline components diff --git a/docs/sentence_transformer.md b/docs/sentence_transformer.md index cf418b52..073442df 100644 --- a/docs/sentence_transformer.md +++ b/docs/sentence_transformer.md @@ -55,6 +55,22 @@ print(model.similarity(embeddings[:1], embeddings[1:])) the model and tokenizer. `model.similarity()` uses the checkpoint's declared cosine or dot-product scoring function. +## Auto loading + +The same ZeroModels-hosted repositories work with the Auto API: + +```python +from zeromodels import AutoZModel, AutoZMTokenizer + +repo = "zeromodels/all-MiniLM-L6-v2" +model = AutoZModel.from_weights(repo) # ZMSentenceTransformer +tokenizer = AutoZMTokenizer.from_weights(repo) # ZMSentenceTransformerTokenizer +``` + +Detection uses `model_type: "sentence_transformer"` in `zm_config.json`, so this +works for all 13 hosted checkpoints and compatible fine-tunes. The model's full +embedding recipe and the tokenizer's checkpoint-specific settings are preserved. + ## Tokenizer preprocessing metadata Each checkpoint uses its own `tokenizer.json` for the vocabulary, normalization, diff --git a/tests/integration/test_sentence_transformer.py b/tests/integration/test_sentence_transformer.py index 7a03d2d6..0edd9c53 100644 --- a/tests/integration/test_sentence_transformer.py +++ b/tests/integration/test_sentence_transformer.py @@ -217,7 +217,16 @@ def test_cosine_similarity_accepts_numpy_embeddings(): np.testing.assert_allclose(ops.convert_to_numpy(model.similarity(a, b)), [[1.0]]) -def test_sentence_transformer_h5_weights_round_trip(tmp_path): +@pytest.mark.parametrize( + "declared", ("ZMSentenceTransformer", "SentenceTransformerModel") +) +def test_sentence_transformer_h5_weights_round_trip_through_auto( + monkeypatch, tmp_path, declared +): + from zeromodels import AutoZModel + from zeromodels.base import base_mixin + from zeromodels.conversion import zm_config + kwargs = { "backbone_type": "bert", "vocab_size": 32, @@ -237,9 +246,19 @@ def test_sentence_transformer_h5_weights_round_trip(tmp_path): path = tmp_path / "model.weights.h5" model.save_weights(str(path)) - reloaded = ZMSentenceTransformer(**model.config.constructor_kwargs()) + config_path = zm_config.write_zm_config( + tmp_path, ZMSentenceTransformer, "example", model.config, weight_dtype="float32" + ) + with open(config_path, encoding="utf-8") as stream: + spec = json.load(stream) + spec["model_class"] = declared + monkeypatch.setattr(zm_config, "load_zm_config", lambda repo_id: spec) + monkeypatch.setattr(base_mixin, "resolve_weights_url", lambda url: url) + monkeypatch.setattr(base_mixin, "download_weights", lambda url: str(path)) + + reloaded = AutoZModel.from_weights("zeromodels/example") + assert isinstance(reloaded, ZMSentenceTransformer) assert model.name == "SentenceTransformerModel" assert reloaded.name == "ZMSentenceTransformer" - reloaded.load_weights(str(path)) actual = ops.convert_to_numpy(reloaded(inputs)["embeddings"]) np.testing.assert_allclose(actual, expected, rtol=1e-5, atol=1e-5) diff --git a/tests/integration/test_sentence_transformer_tokenizer.py b/tests/integration/test_sentence_transformer_tokenizer.py index f4b615e0..99757143 100644 --- a/tests/integration/test_sentence_transformer_tokenizer.py +++ b/tests/integration/test_sentence_transformer_tokenizer.py @@ -85,6 +85,14 @@ def download(repo_id, filename, token=None, revision=None): def test_hub_metadata_uses_one_snapshot_and_official_sequence_limit( monkeypatch, tmp_path, tokenizer_path, declared ): + from zeromodels import AutoZMTokenizer + from zeromodels.conversion import zm_config + + monkeypatch.setattr( + zm_config, + "load_zm_config", + lambda repo_id: {"model_type": "sentence_transformer"}, + ) spec = metadata() spec["preprocessor_class"] = declared calls = mock_repo( @@ -95,9 +103,8 @@ def test_hub_metadata_uses_one_snapshot_and_official_sequence_limit( "assets/tokenizer.json": tokenizer_path, }, ) - tok = ZMSentenceTransformerTokenizer.from_weights( - "zeromodels/example", revision="release" - ) + tok = AutoZMTokenizer.from_weights("zeromodels/example", revision="release") + assert isinstance(tok, ZMSentenceTransformerTokenizer) encoded = tok(["hello world " * 10, "hello"]) assert tok.max_seq_len == 6 assert tok.cls_token_id == 2 and tok.sep_token_id == 3 diff --git a/zeromodels/auto/__init__.py b/zeromodels/auto/__init__.py index a9427528..f57401fa 100644 --- a/zeromodels/auto/__init__.py +++ b/zeromodels/auto/__init__.py @@ -5,8 +5,9 @@ ``AutoTokenizer`` / ... when both libraries are imported. The ``AutoZM`` model classes (``AutoZModel``, ``AutoZMDetect``, ``AutoZMImageClassify``, ``AutoZMTextGenerate``, ...) plus ``AutoZMConfig`` / ``AutoZMTokenizer`` / ``AutoZMProcessor`` / -``AutoZMImageProcessor`` are generated from :mod:`zeromodels.models`, so this package -re-exports whatever the factory built. +``AutoZMImageProcessor`` are generated from the committed task tables. Class discovery +includes :mod:`zeromodels.models` and :mod:`zeromodels.sentence_transformers`; this +package re-exports whatever the factory built. """ from zeromodels.auto import auto_factory