diff --git a/README.md b/README.md index 900108b9..ae038a46 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 (sentence_transformers) + +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 | +|---------------|-------------------|---------------------| +| [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 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/models.md b/docs/models.md index 342ff0f5..5c06b6ab 100644 --- a/docs/models.md +++ b/docs/models.md @@ -39,6 +39,13 @@ Encoders for embeddings and masked language modelling, and decoder LLMs, dense a - [RoBERTa](roberta.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/mpnet.md b/docs/mpnet.md index 1070b6e9..c84a5751 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 separate embedding API. ## Variants diff --git a/docs/sentence_transformer.md b/docs/sentence_transformer.md new file mode 100644 index 00000000..073442df --- /dev/null +++ b/docs/sentence_transformer.md @@ -0,0 +1,167 @@ +# 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 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. + +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 +import os + +os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" + +from zeromodels.sentence_transformers import ( + ZMSentenceTransformer, + ZMSentenceTransformerTokenizer, +) + +repo = "zeromodels/all-MiniLM-L6-v2" +model = ZMSentenceTransformer.from_weights(repo) +tokenizer = ZMSentenceTransformerTokenizer.from_weights(repo) + +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 repo for +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, +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.sentence_transformers", + "preprocessor_class": "ZMSentenceTransformerTokenizer", + "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. +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. + +## 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: + +```python +upstream = "hf:sentence-transformers/all-MiniLM-L6-v2" +model = ZMSentenceTransformer.from_weights(upstream) +tokenizer = ZMSentenceTransformerTokenizer.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. + +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. After the Usage example above, compare +the published ZeroModels weights with the matching original checkpoint: + +```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))) +``` + +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 new file mode 100644 index 00000000..5a37ee8b --- /dev/null +++ b/docs/sentence_transformers.md @@ -0,0 +1,49 @@ +# 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.sentence_transformers` 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.sentence_transformers import ( + ZMSentenceTransformer, + ZMSentenceTransformerTokenizer, +) + +repo = "zeromodels/all-MiniLM-L6-v2" +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) +``` + +See [Sentence Transformer](sentence_transformer.md) for the complete API, +upstream `hf:` loading, similarity scoring, and parity-check instructions. diff --git a/pyproject.toml b/pyproject.toml index a405f764..019f301d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -104,6 +104,7 @@ version = { attr = "zeromodels.__version__" } source = ["zeromodels", "tests"] omit = [ "zeromodels/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 66cf90d6..673c0490 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) ---- + "ZMSentenceTransformer": { + "module": "zeromodels.sentence_transformers", + "model_cls": "ZMSentenceTransformer", + "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.sentence_transformers", + "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", @@ -4385,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"ZMSentenceTransformer_{_variant.replace('-', '_')}"] = { + "module": "zeromodels.sentence_transformers", + "model_cls": "ZMSentenceTransformer", + "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. @@ -4953,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. @@ -4967,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..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,14 +87,27 @@ def _hf_types(cls): def _iter_model_classes(): import zeromodels.models as models + import zeromodels.sentence_transformers as sentence_transformers 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 + 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) + 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(): @@ -140,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): @@ -195,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" @@ -428,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 new file mode 100644 index 00000000..0edd9c53 --- /dev/null +++ b/tests/integration/test_sentence_transformer.py @@ -0,0 +1,264 @@ +"""Offline contract tests for the multi-backbone sentence-embedding API.""" + +import importlib +import json + +import numpy as np +import pytest +from keras import ops, saving + +from tests.base.model_test_registry import MODEL_TEST_CONFIGS +from zeromodels.sentence_transformers import ( + SentenceTransformerTextEmbed, + ZMSentenceTransformer, + ZMSentenceTransformerTokenizer, +) + + +@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( + package + ".sentence_transformer_model" + ) + 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 + + +@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": module, + "model_class": declared, + "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 = ZMSentenceTransformer.from_hub_repo( + "zeromodels/example-sentence-transformer", load_weights=False + ) + 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", + "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("ZMSentenceTransformer_") + for name in MODEL_TEST_CONFIGS + if name.startswith("ZMSentenceTransformer_") + } + assert registered == {name.replace("-", "_") for name in published} + + +@pytest.mark.parametrize( + "backbone_type", + ("mpnet", "bert", "roberta", "xlm-roberta", "distilbert", "albert"), +) +def test_backbone_builds_and_returns_sentence_embeddings(backbone_type): + model = ZMSentenceTransformer( + 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 ( + ZMSentenceTransformer._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"): + ZMSentenceTransformer._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 = ZMSentenceTransformer( + 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 = ZMSentenceTransformer( + 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]]) + + +@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, + "embed_dim": 8, + "num_layers": 1, + "num_heads": 2, + "mlp_dim": 16, + "max_position_embeddings": 16, + } + model = ZMSentenceTransformer(name="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)) + + 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" + 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 new file mode 100644 index 00000000..99757143 --- /dev/null +++ b/tests/integration/test_sentence_transformer_tokenizer.py @@ -0,0 +1,268 @@ +"""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.sentence_transformers import ZMSentenceTransformerTokenizer + +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.sentence_transformers", + "preprocessor_class": "ZMSentenceTransformerTokenizer", + "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 + + +@pytest.mark.parametrize( + "declared", ("ZMSentenceTransformerTokenizer", "SentenceTransformerTokenizer") +) +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( + monkeypatch, + tmp_path, + { + "zm_preprocessor.json": spec, + "assets/tokenizer.json": tokenizer_path, + }, + ) + 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 + 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 = ZMSentenceTransformerTokenizer( + 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["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 = 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]), + 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 = 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"]) + + +def test_sentence_transformer_lowercase_is_separate_from_tokenizer_normalization( + tokenizer_path, +): + tc = {"do_lower_case": False, "pad_token": "[PAD]"} + plain = ZMSentenceTransformerTokenizer( + tokenizer_file=str(tokenizer_path), + tokenizer_config=tc, + sentence_transformer_config={"do_lower_case": False}, + ) + extra = ZMSentenceTransformerTokenizer( + 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 = ZMSentenceTransformerTokenizer( + 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 = ZMSentenceTransformerTokenizer( + 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 = 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) + + +@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"): + 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 new file mode 100644 index 00000000..9645ca0c --- /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.sentence_transformers import ZMSentenceTransformerTokenizer + +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 = ZMSentenceTransformerTokenizer.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/website/mkdocs.yml b/website/mkdocs.yml index 06eb039b..365e80aa 100644 --- a/website/mkdocs.yml +++ b/website/mkdocs.yml @@ -169,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 diff --git a/zeromodels/__init__.py b/zeromodels/__init__.py index 3a5387a8..d85bd976 100644 --- a/zeromodels/__init__.py +++ b/zeromodels/__init__.py @@ -1,4 +1,12 @@ -from zeromodels import auto, conversion, models, samplers, utils +from zeromodels import ( + auto, + conversion, + models, + samplers, + sentence_transformers, + st_models, + utils, +) from zeromodels.auto import ( AutoZMConfig, AutoZMImageProcessor, @@ -14,6 +22,8 @@ "auto", "models", "samplers", + "sentence_transformers", + "st_models", "utils", "conversion", "version", 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 diff --git a/zeromodels/auto/auto_factory.py b/zeromodels/auto/auto_factory.py index e733f338..45f71086 100644 --- a/zeromodels/auto/auto_factory.py +++ b/zeromodels/auto/auto_factory.py @@ -104,14 +104,20 @@ def _class_index(): global _CLASS_INDEX if _CLASS_INDEX is None: import zeromodels.models as models + import zeromodels.sentence_transformers as sentence_transformers 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, 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): + index[name] = obj _CLASS_INDEX = index return _CLASS_INDEX diff --git a/zeromodels/auto/auto_mapping_names.py b/zeromodels/auto/auto_mapping_names.py index a659a0e0..edfee13d 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": "ZMSentenceTransformer", "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": "ZMSentenceTransformerTokenizer", "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/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/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( diff --git a/zeromodels/models/sentence_transformer/__init__.py b/zeromodels/models/sentence_transformer/__init__.py new file mode 100644 index 00000000..a9f861bc --- /dev/null +++ b/zeromodels/models/sentence_transformer/__init__.py @@ -0,0 +1,51 @@ +"""Compatibility imports for models moved to ``zeromodels.sentence_transformers``.""" + +import sys + +from zeromodels.sentence_transformers import ( + SentenceTransformerConfig, + SentenceTransformerModel, + SentenceTransformerTextEmbed, + SentenceTransformerTokenizer, + ZMSentenceTransformer, + ZMSentenceTransformerTokenizer, +) +from zeromodels.sentence_transformers import ( + convert_sentence_transformer_hf_to_keras as _converter, +) +from zeromodels.sentence_transformers import ( + sentence_transformer_backbones as _backbones, +) +from zeromodels.sentence_transformers import ( + sentence_transformer_config as _config, +) +from zeromodels.sentence_transformers import ( + sentence_transformer_layers as _layers, +) +from zeromodels.sentence_transformers import ( + sentence_transformer_model as _model, +) +from zeromodels.sentence_transformers 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", + "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/sentence_transformers/convert_sentence_transformer_hf_to_keras.py b/zeromodels/sentence_transformers/convert_sentence_transformer_hf_to_keras.py new file mode 100644 index 00000000..4a41e88f --- /dev/null +++ b/zeromodels/sentence_transformers/convert_sentence_transformer_hf_to_keras.py @@ -0,0 +1,249 @@ +from __future__ import annotations + +from typing import Dict + +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, +) + +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( + 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 = {} + 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) + 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"): + 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}") + + 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 = 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": + 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 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.sentence_transformers import ( + ZMSentenceTransformer, + ZMSentenceTransformerTokenizer, + ) + + 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 = ZMSentenceTransformer.from_hf(hf_id) + tok = ZMSentenceTransformerTokenizer.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, + ZMSentenceTransformer, + 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/sentence_transformers/sentence_transformer_backbones.py b/zeromodels/sentence_transformers/sentence_transformer_backbones.py new file mode 100644 index 00000000..6bd10e5e --- /dev/null +++ b/zeromodels/sentence_transformers/sentence_transformer_backbones.py @@ -0,0 +1,225 @@ +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)) + + +@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, approximate=False) + 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 = ( + 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: + 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/sentence_transformers/sentence_transformer_config.py b/zeromodels/sentence_transformers/sentence_transformer_config.py new file mode 100644 index 00000000..89be4331 --- /dev/null +++ b/zeromodels/sentence_transformers/sentence_transformer_config.py @@ -0,0 +1,41 @@ +from typing import Optional + +from zeromodels.base import BaseConfig + + +class SentenceTransformerConfig(BaseConfig): + """Configuration for :class:`ZMSentenceTransformer`. + + 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 + 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" + normalize_embeddings: bool = True + projection_dim: Optional[int] = None + projection_activation: str = "linear" + similarity_fn_name: str = "cosine" + max_seq_len: int = 384 diff --git a/zeromodels/sentence_transformers/sentence_transformer_layers.py b/zeromodels/sentence_transformers/sentence_transformer_layers.py new file mode 100644 index 00000000..2806f63f --- /dev/null +++ b/zeromodels/sentence_transformers/sentence_transformer_layers.py @@ -0,0 +1,81 @@ +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, +} + + +@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 + 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() + 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): + 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() + config.update({"eps": self.eps}) + return config diff --git a/zeromodels/sentence_transformers/sentence_transformer_model.py b/zeromodels/sentence_transformers/sentence_transformer_model.py new file mode 100644 index 00000000..66998dd9 --- /dev/null +++ b/zeromodels/sentence_transformers/sentence_transformer_model.py @@ -0,0 +1,583 @@ +"""Sentence Transformer 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, +) + +SENTENCE_TRANSFORMER_HUB_SIBLINGS = frozenset( + {"ZMSentenceTransformer", "SentenceTransformerModel", "SentenceTransformer"} +) + + +@keras.saving.register_keras_serializable(package="zeromodels") +@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 = ZMSentenceTransformer.from_weights( + "hf:sentence-transformers/all-mpnet-base-v2" + ) + tok = ZMSentenceTransformerTokenizer.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="ZMSentenceTransformer", + **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:`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. + """ + 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)) + 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() + config.update(self.config.constructor_kwargs()) + return config + + +@keras.saving.register_keras_serializable(package="zeromodels") +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/sentence_transformers/sentence_transformer_tokenizer.py b/zeromodels/sentence_transformers/sentence_transformer_tokenizer.py new file mode 100644 index 00000000..de792fc3 --- /dev/null +++ b/zeromodels/sentence_transformers/sentence_transformer_tokenizer.py @@ -0,0 +1,349 @@ +import json +import os +import re +from pathlib import PurePosixPath +from typing import List, Union + +import keras +from tokenizers import Tokenizer, normalizers + +from zeromodels.base import BaseTokenizer + + +@keras.saving.register_keras_serializable(package="zeromodels") +@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")`` / + ``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: Explicit truncation override; otherwise read from metadata. + 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 = None, + backbone_type: str = "mpnet", + bos_token: str = None, + eos_token: str = None, + unk_token: str = None, + pad_token: str = None, + mask_token: str = None, + tokenizer_config: dict = None, + sentence_transformer_config: dict = None, + **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.backbone_type = backbone_type + 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( + ( + 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.pad_token_id = tok.token_to_id(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_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) + + @classmethod + def from_hf(cls, repo, **kwargs): + from huggingface_hub import hf_hub_download + + 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"]) + 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 official preprocessing metadata, with legacy config fallback.""" + from huggingface_hub import hf_hub_download + + 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") not in ( + cls.__name__, + "SentenceTransformerTokenizer", + ): + 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: + 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, + "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( + 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 new file mode 100644 index 00000000..90b6eba0 --- /dev/null +++ b/zeromodels/st_models/__init__.py @@ -0,0 +1,40 @@ +"""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, +) + +# 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", + "SentenceTransformerConfig", + "SentenceTransformerModel", + "SentenceTransformerTextEmbed", + "SentenceTransformerTokenizer", + "ZMSentenceTransformer", + "ZMSentenceTransformerTokenizer", +]