diff --git a/README.md b/README.md index 3727ecc0..f897c3d1 100644 --- a/README.md +++ b/README.md @@ -56,6 +56,7 @@ Documentation sources are also available in [`docs/`](docs/). | DeBERTa | [DeBERTa: Decoding-enhanced BERT with Disentangled Attention](https://arxiv.org/abs/2006.03654) | `transformers` | | DeBERTa-v2 | [DeBERTa: Decoding-enhanced BERT with Disentangled Attention](https://arxiv.org/abs/2006.03654) | `transformers` | | DeBERTa-v3 | [DeBERTaV3: Improving DeBERTa using ELECTRA-Style Pre-Training with Gradient-Disentangled Embedding Sharing](https://arxiv.org/abs/2111.09543) | `transformers` | + | MPNet | [MPNet: Masked and Permuted Pre-training for Language Understanding](https://arxiv.org/abs/2004.09297) | `transformers` | | T5 | [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://arxiv.org/abs/1910.10683) | `transformers` | | BART | [BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension](https://arxiv.org/abs/1910.13461) | `transformers` | diff --git a/docs/mpnet.md b/docs/mpnet.md new file mode 100644 index 00000000..1070b6e9 --- /dev/null +++ b/docs/mpnet.md @@ -0,0 +1,148 @@ +# MPNet + +
+Weights: pretrained Keras weights live on Hugging Face under +zeromodels/<variant> +(each repo carries zm_config.json + model.weights.h5). +Load with from_weights("zeromodels/<variant>"). +
+ +Microsoft's MPNet in pure Keras 3: a bidirectional encoder pre-trained with **masked and +permuted** language modelling, which unifies BERT's masked-LM objective with XLNet's +permuted one. It carries a masked-LM head plus sequence / token classification, +question-answering, and multiple-choice heads. One implementation runs unmodified on +TensorFlow / Torch / JAX. + +Three things distinguish it from BERT: + +- **No token-type embeddings.** Inputs are `input_ids` + `attention_mask` only, so + `MPNetTokenizer` emits no `token_type_ids`. +- **A shared relative position bias.** Every attention layer adds the same + `(1, num_heads, L, L)` bias, gathered from `relative_attention_num_buckets` + log-spaced buckets and computed **once** per forward pass. +- **The attention output projection lives inside self-attention** (`attention.attn.o` + upstream), so the block's LayerNorm sits directly on `attention`, not on + `attention.output` as in BERT. + +Like RoBERTa, position ids are offset past the padding id. + +- Paper: [MPNet: Masked and Permuted Pre-training for Language Understanding (arXiv:2004.09297)](https://arxiv.org/abs/2004.09297) +- HF docs: [transformers/model_doc/mpnet](https://huggingface.co/docs/transformers/model_doc/mpnet) + +See also [bert.md](bert.md), [roberta.md](roberta.md), [modernbert.md](modernbert.md), +[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. + +## Variants + +Load any of these with `from_weights("zeromodels/")`. + +| Variant | Hub | layers / dim | +|---|---|---| +| `mpnet_base` | [`zeromodels/mpnet_base`](https://huggingface.co/zeromodels/mpnet_base) | 12 / 768 | + +## API + +### `MPNetModel` + +The encoder backbone plus a `tanh` pooler over the `` token. Takes a dict of +`input_ids` / `attention_mask` (both `(B, L)` int) and returns +`{"last_hidden_state": (B, L, embed_dim), "pooler_output": (B, embed_dim)}`. Pass +`add_pooler=False` to drop the pooler. + +| Arg | Default | Meaning | +|---|---|---| +| `vocab_size` | `30527` | token vocabulary size | +| `embed_dim` | `768` | model / hidden width | +| `num_layers` | `12` | transformer blocks | +| `num_heads` | `12` | attention heads | +| `mlp_dim` | `3072` | feed-forward inner width | +| `max_position_embeddings` | `512` | position-table size (padding-offset) | +| `relative_attention_num_buckets` | `32` | relative position bias buckets | +| `hidden_act` | `"gelu"` | feed-forward activation | +| `layer_norm_eps` | `1e-12` | LayerNorm epsilon | +| `pad_token_id` | `1` | padding token id | + +### Task heads + +Each composes an `MPNetModel` backbone and adds a head; all take the same backbone +constructor args, plus the extras below. The pretrained encoder + masked-LM head load +real weights; the classification / QA heads start randomly initialized (ready for +fine-tuning) and load trained weights from a `hf:` fine-tune. + +| Class | Extra args | Output | +|---|---|---| +| `MPNetMaskedLM` | `add_pooler=False` | `(B, L, vocab_size)` token logits | +| `MPNetSequenceClassify` | `num_classes=2`, `classifier_dropout=0.0`, `classifier_activation="linear"` | `(B, num_classes)` | +| `MPNetTokenClassify` | `num_classes=2`, `classifier_dropout=0.0`, `classifier_activation="linear"` | `(B, L, num_classes)` | +| `MPNetQnA` | — | `{"start_logits": (B, L), "end_logits": (B, L)}` | +| `MPNetMultipleChoice` | `num_choices=2`, `classifier_dropout=0.0` | `(B, num_choices)` | + +`MPNetSequenceClassify` pools the `` token inside its own head (dropout → `tanh` +dense → dropout → projection) rather than reading the encoder pooler, matching the +reference implementation. + +`MPNetMultipleChoice` takes `(B, num_choices, L)` inputs, folds the choice axis into the +batch, and folds the scores back out. + +### `MPNetTokenizer` + +WordPiece tokenizer (`tokenizers` Rust backend) loading a `tokenizer.json`. MPNet pairs +RoBERTa-style special tokens (``, ``, ``, ``) with a BERT-style +WordPiece vocabulary and `[UNK]`. `call` returns `input_ids` + `attention_mask` only. + +## Usage + +```python +import os + +os.environ["KERAS_BACKEND"] = "torch" # or "jax" / "tensorflow" + +from zeromodels.models.mpnet import MPNetModel, MPNetTokenizer + +model = MPNetModel.from_weights("zeromodels/mpnet_base") +tokenizer = MPNetTokenizer.from_weights("zeromodels/mpnet_base") + +out = model(tokenizer(["the quick brown fox", "jumped over the lazy dog"])) +print(out["last_hidden_state"].shape, out["pooler_output"].shape) +``` + +Fill-mask with the pretrained head: + +```python +from zeromodels.models.mpnet import MPNetMaskedLM + +mlm = MPNetMaskedLM.from_weights("zeromodels/mpnet_base") +logits = mlm(tokenizer(["the capital of France is ."])) +``` + +Or convert any MPNet checkpoint straight from the Hub: + +```python +model = MPNetModel.from_weights("hf:microsoft/mpnet-base") +``` + +## Conversion accuracy + +Converted from `microsoft/mpnet-base` and checked against the reference +implementation on a padded batch (`max|Δ|`, pad positions excluded): + +| Class | last_hidden_state | pooler_output | logits | +|---|---|---|---| +| `MPNetModel` | 2.6e-06 | 2.7e-07 | — | +| `MPNetMaskedLM` | — | — | 1.3e-05 | + +Masked-LM argmax agreement with the reference is 100%. + +Each variant repo hosts **one** file: the superset `MPNetMaskedLM(add_pooler=True)` +(encoder + pooler + masked-LM head). `MPNetModel.CHECKPOINT_SOURCE` points every class at +it, so the encoder and the task heads copy their own subset out of that single checkpoint — +the numbers above are re-measured after that round trip. + +> **Converter note.** MPNet ships both `lm_head.bias` and `lm_head.decoder.bias` with +> *different* values and does not tie them; the forward pass reads `decoder.bias`, while +> `lm_head.bias` is vestigial. The converter maps the decoder bias to +> `lm_head.decoder.bias` accordingly — using `lm_head.bias` (as is correct for BERT and +> RoBERTa, where HF does tie them) shifts every logit by up to ~7.8. diff --git a/tests/base/model_test_registry.py b/tests/base/model_test_registry.py index c68c6d5e..2eb0f352 100644 --- a/tests/base/model_test_registry.py +++ b/tests/base/model_test_registry.py @@ -2886,6 +2886,116 @@ "input_factory_kwargs": {"num_choices": 3, "seq_len": 16}, "expected_output_shape": (2, 3), }, + # ---- Text encoders (MPNet) ---- + # MPNet has no token-type ids, so it reuses the ModernBERT input factories. + "MPNetModel": { + "module": "zeromodels.models.mpnet", + "model_cls": "MPNetModel", + "model_type": "llm", + "init_kwargs": { + "vocab_size": 128, + "embed_dim": 32, + "num_layers": 2, + "num_heads": 2, + "mlp_dim": 64, + "max_position_embeddings": 64, + "relative_attention_num_buckets": 8, + }, + "input_factory": "modernbert_input", + "input_factory_kwargs": {"seq_len": 16}, + "expected_output_shape": { + "last_hidden_state": (2, 16, 32), + "pooler_output": (2, 32), + }, + }, + "MPNetMaskedLM": { + "module": "zeromodels.models.mpnet", + "model_cls": "MPNetMaskedLM", + "model_type": "llm", + "init_kwargs": { + "vocab_size": 128, + "embed_dim": 32, + "num_layers": 2, + "num_heads": 2, + "mlp_dim": 64, + "max_position_embeddings": 64, + "relative_attention_num_buckets": 8, + }, + "input_factory": "modernbert_input", + "input_factory_kwargs": {"seq_len": 16}, + "expected_output_shape": (2, 16, 128), + }, + "MPNetSequenceClassify": { + "module": "zeromodels.models.mpnet", + "model_cls": "MPNetSequenceClassify", + "model_type": "llm", + "init_kwargs": { + "vocab_size": 128, + "embed_dim": 32, + "num_layers": 2, + "num_heads": 2, + "mlp_dim": 64, + "max_position_embeddings": 64, + "relative_attention_num_buckets": 8, + "num_classes": 3, + }, + "input_factory": "modernbert_input", + "input_factory_kwargs": {"seq_len": 16}, + "expected_output_shape": (2, 3), + }, + "MPNetTokenClassify": { + "module": "zeromodels.models.mpnet", + "model_cls": "MPNetTokenClassify", + "model_type": "llm", + "init_kwargs": { + "vocab_size": 128, + "embed_dim": 32, + "num_layers": 2, + "num_heads": 2, + "mlp_dim": 64, + "max_position_embeddings": 64, + "relative_attention_num_buckets": 8, + "num_classes": 3, + }, + "input_factory": "modernbert_input", + "input_factory_kwargs": {"seq_len": 16}, + "expected_output_shape": (2, 16, 3), + }, + "MPNetQnA": { + "module": "zeromodels.models.mpnet", + "model_cls": "MPNetQnA", + "model_type": "llm", + "init_kwargs": { + "vocab_size": 128, + "embed_dim": 32, + "num_layers": 2, + "num_heads": 2, + "mlp_dim": 64, + "max_position_embeddings": 64, + "relative_attention_num_buckets": 8, + }, + "input_factory": "modernbert_input", + "input_factory_kwargs": {"seq_len": 16}, + "expected_output_shape": {"start_logits": (2, 16), "end_logits": (2, 16)}, + }, + "MPNetMultipleChoice": { + "module": "zeromodels.models.mpnet", + "model_cls": "MPNetMultipleChoice", + "model_type": "llm", + "init_kwargs": { + "vocab_size": 128, + "embed_dim": 32, + "num_layers": 2, + "num_heads": 2, + "mlp_dim": 64, + "max_position_embeddings": 64, + "relative_attention_num_buckets": 8, + "num_choices": 3, + }, + "input_factory": "modernbert_multiple_choice_input", + "input_factory_kwargs": {"num_choices": 3, "seq_len": 16}, + "expected_output_shape": (2, 3), + }, # ---- Text encoders (XLM-RoBERTa) ---- "XLMRobertaModel": { "module": "zeromodels.models.xlm_roberta", diff --git a/tests/fixtures/cross_backend_parity.json b/tests/fixtures/cross_backend_parity.json index c9be97b1..ead0d87f 100644 --- a/tests/fixtures/cross_backend_parity.json +++ b/tests/fixtures/cross_backend_parity.json @@ -13461,6 +13461,417 @@ ] } ], + "MPNetMaskedLM": [ + { + "shape": [ + 2, + 16, + 128 + ], + "sample": [ + 0.013573, + 0.019182, + -0.003003, + -0.002239, + 0.010858, + 0.002082, + -0.010376, + -0.008917, + 0.019973, + 0.000831, + 0.048757, + 0.014154, + 0.014761, + 0.004264, + -0.0039, + 0.004477, + -0.055531, + 0.024344, + -0.035723, + -0.00028, + -0.018314, + -0.004464, + -0.017497, + -0.027353, + 0.017084, + 0.011324, + -0.032609, + -0.014469, + 0.00916, + -0.00502, + -0.00857, + 0.010407, + 0.015198, + -0.015795, + 0.003306, + 0.032307, + 0.022574, + 0.007339, + 0.003671, + -0.008746, + 0.007837, + -0.001909, + 0.003813, + -0.018606, + 0.026303, + 0.038794, + 0.017043, + -0.015144, + 0.004639, + -0.003656, + -0.042075, + -0.023148, + -0.042787, + 0.001969, + 0.009279, + -0.04811, + -0.033618, + 0.002299, + 0.013592, + -0.009713, + -0.018393, + -0.010557, + 0.018521, + -0.021079 + ] + } + ], + "MPNetModel": [ + { + "shape": [ + 2, + 16, + 32 + ], + "sample": [ + -0.020645, + 0.015959, + -0.019746, + 0.015891, + -0.020546, + -0.013821, + 0.017029, + -0.015514, + 0.014505, + 0.023789, + 0.033138, + 0.024079, + 0.033855, + -0.024825, + -0.023612, + -0.024429, + -0.023467, + -0.059901, + 0.006218, + -0.062317, + 0.006591, + 0.005456, + -0.020633, + 0.005696, + -0.023892, + 0.006111, + 0.00978, + 0.006021, + 0.008071, + 0.00593, + 0.021451, + -0.01323, + 0.015734, + -0.010186, + 0.044177, + -0.061401, + 0.043231, + -0.062624, + -0.010335, + 0.026483, + -0.009721, + 0.027006, + 0.019313, + -0.004297, + 0.019373, + -0.004201, + 0.018476, + 0.015594, + -0.020873, + 0.016887, + -0.020705, + 0.032291, + -0.018332, + 0.03389, + -0.017909, + 0.023473, + -0.025739, + 0.023412, + -0.02516, + -0.053637, + -0.004469, + -0.055186, + -0.00548, + -0.015974 + ] + }, + { + "shape": [ + 2, + 32 + ], + "sample": [ + 0.00065, + 0.010326, + 0.005773, + 0.004612, + 0.007379, + -0.004441, + 0.015894, + 0.005093, + -0.010812, + 0.031536, + -0.014104, + 0.010285, + 0.006126, + 0.005762, + -0.024201, + 0.026273, + 0.010725, + 0.0098, + 0.000418, + -0.012469, + -0.008399, + -0.023222, + -0.007242, + -0.027337, + 0.016845, + 0.019998, + -0.00719, + -0.002365, + -0.011976, + 0.038827, + -0.023255, + -0.01336, + 0.000654, + 0.010103, + 0.006004, + 0.00487, + 0.007337, + -0.00421, + 0.01594, + 0.004906, + -0.010986, + 0.031251, + -0.014171, + 0.01033, + 0.006074, + 0.005645, + -0.024634, + 0.026149, + 0.010839, + 0.00979, + 0.000852, + -0.01262, + -0.008544, + -0.023639, + -0.006925, + -0.027536, + 0.016622, + 0.020242, + -0.007396, + -0.00241, + -0.012197, + 0.0386, + -0.023296, + -0.013329 + ] + } + ], + "MPNetMultipleChoice": [ + { + "shape": [ + 2, + 3 + ], + "sample": [ + -0.01615, + -0.016114, + -0.016131, + -0.016112, + -0.016114, + -0.016105 + ] + } + ], + "MPNetQnA": [ + { + "shape": [ + 2, + 16 + ], + "sample": [ + -0.040562, + -0.039599, + -0.040371, + -0.040677, + -0.039517, + -0.040737, + -0.039518, + -0.040246, + -0.040091, + -0.040533, + -0.040357, + -0.04072, + -0.039934, + -0.04053, + -0.040411, + -0.04008, + -0.039782, + -0.040645, + -0.040107, + -0.040016, + -0.04025, + -0.040043, + -0.040763, + -0.040006, + -0.040546, + -0.040007, + -0.039798, + -0.04033, + -0.040371, + -0.04012, + -0.039401, + -0.040783 + ] + }, + { + "shape": [ + 2, + 16 + ], + "sample": [ + 0.018373, + 0.0185, + 0.01844, + 0.018343, + 0.01861, + 0.018244, + 0.018584, + 0.018463, + 0.018288, + 0.018401, + 0.018483, + 0.018244, + 0.018515, + 0.018206, + 0.018466, + 0.018338, + 0.018543, + 0.018295, + 0.01841, + 0.018584, + 0.018409, + 0.018558, + 0.018376, + 0.018428, + 0.018557, + 0.018491, + 0.018478, + 0.018418, + 0.018304, + 0.018487, + 0.018556, + 0.018167 + ] + } + ], + "MPNetSequenceClassify": [ + { + "shape": [ + 2, + 3 + ], + "sample": [ + 0.005387, + 0.007261, + -0.035079, + 0.005365, + 0.007217, + -0.035066 + ] + } + ], + "MPNetTokenClassify": [ + { + "shape": [ + 2, + 16, + 3 + ], + "sample": [ + -0.024943, + -0.001819, + -0.024911, + -0.002318, + -0.024774, + -0.001945, + -0.025025, + -0.001701, + -0.024932, + -0.002333, + -0.025001, + -0.001696, + -0.024944, + -0.002243, + -0.024965, + -0.002019, + -0.025027, + -0.002081, + -0.025085, + -0.001842, + -0.024971, + -0.001967, + -0.024847, + -0.00164, + -0.024944, + -0.002156, + -0.024911, + -0.001935, + -0.024777, + -0.001788, + -0.025015, + -0.002021, + -0.02485, + -0.00203, + -0.024839, + -0.001767, + -0.02506, + -0.001828, + -0.024913, + -0.002012, + -0.025059, + -0.001892, + -0.024976, + -0.001954, + -0.024792, + -0.001738, + -0.025123, + -0.00208, + -0.024907, + -0.001628, + -0.025064, + -0.002151, + -0.024901, + -0.00207, + -0.024921, + -0.001825, + -0.024727, + -0.001856, + -0.024917, + -0.00196, + -0.025066, + -0.002352, + -0.024992, + -0.054829 + ] + } + ], "Mask2FormerUniversalSegment": [ { "shape": [ diff --git a/website/mkdocs.yml b/website/mkdocs.yml index 70b1dc7d..9da072d7 100644 --- a/website/mkdocs.yml +++ b/website/mkdocs.yml @@ -158,6 +158,7 @@ nav: - Mistral: mistral.md - Mixtral: mixtral.md - ModernBERT: modernbert.md + - MPNet: mpnet.md - Qwen: qwen.md - Qwen2: qwen2.md - Qwen2 MoE: qwen2_moe.md diff --git a/zeromodels/auto/auto_mapping_names.py b/zeromodels/auto/auto_mapping_names.py index c3ec9f08..389be02e 100644 --- a/zeromodels/auto/auto_mapping_names.py +++ b/zeromodels/auto/auto_mapping_names.py @@ -141,6 +141,7 @@ "deberta_v3": "DebertaV3MaskedLM", "electra": "ElectraMaskedLM", "modernbert": "ModernBertMaskedLM", + "mpnet": "MPNetMaskedLM", "roberta": "RobertaMaskedLM", "xlm-roberta": "XLMRobertaMaskedLM", "xlm_roberta": "XLMRobertaMaskedLM", @@ -233,6 +234,7 @@ "mobilevit": "MobileViTModel", "mobilevitv2": "MobileViTV2Model", "modernbert": "ModernBertModel", + "mpnet": "MPNetModel", "moonshine": "MoonshineModel", "nextvit": "NextViTModel", "oneformer": "OneFormerModel", @@ -298,6 +300,7 @@ "deberta_v3": "DebertaV3MultipleChoice", "electra": "ElectraMultipleChoice", "modernbert": "ModernBertMultipleChoice", + "mpnet": "MPNetMultipleChoice", "roberta": "RobertaMultipleChoice", "xlm-roberta": "XLMRobertaMultipleChoice", "xlm_roberta": "XLMRobertaMultipleChoice", @@ -316,6 +319,7 @@ "deberta_v3": "DebertaV3QnA", "electra": "ElectraQnA", "modernbert": "ModernBertQnA", + "mpnet": "MPNetQnA", "roberta": "RobertaQnA", "t5": "T5QnA", "xlm-roberta": "XLMRobertaQnA", @@ -342,6 +346,7 @@ "deberta_v3": "DebertaV3SequenceClassify", "electra": "ElectraSequenceClassify", "modernbert": "ModernBertSequenceClassify", + "mpnet": "MPNetSequenceClassify", "roberta": "RobertaSequenceClassify", "t5": "T5SequenceClassify", "xlm-roberta": "XLMRobertaSequenceClassify", @@ -424,6 +429,7 @@ "deberta_v3": "DebertaV3TokenClassify", "electra": "ElectraTokenClassify", "modernbert": "ModernBertTokenClassify", + "mpnet": "MPNetTokenClassify", "roberta": "RobertaTokenClassify", "t5": "T5TokenClassify", "xlm-roberta": "XLMRobertaTokenClassify", @@ -609,6 +615,7 @@ "mobilevit": "MobileViTConfig", "mobilevitv2": "MobileViTV2Config", "modernbert": "ModernBertConfig", + "mpnet": "MPNetConfig", "moonshine": "MoonshineConfig", "moonshine_audio": "MoonshineAudioConfig", "moonshine_text": "MoonshineTextConfig", @@ -751,6 +758,7 @@ "mistral3": "Mistral3Tokenizer", "mixtral": "MixtralTokenizer", "modernbert": "ModernBertTokenizer", + "mpnet": "MPNetTokenizer", "moonshine": "MoonshineTokenizer", "oneformer": "OneFormerTokenizer", "openai-gpt": "GptTokenizer", diff --git a/zeromodels/models/__init__.py b/zeromodels/models/__init__.py index 67346365..463e0fac 100644 --- a/zeromodels/models/__init__.py +++ b/zeromodels/models/__init__.py @@ -84,6 +84,7 @@ mobilevitv2, modernbert, moonshine, + mpnet, nextvit, oneformer, owlv2, diff --git a/zeromodels/models/mpnet/__init__.py b/zeromodels/models/mpnet/__init__.py new file mode 100644 index 00000000..4efe8724 --- /dev/null +++ b/zeromodels/models/mpnet/__init__.py @@ -0,0 +1,21 @@ +from .mpnet_config import MPNetConfig +from .mpnet_model import ( + MPNetMaskedLM, + MPNetModel, + MPNetMultipleChoice, + MPNetQnA, + MPNetSequenceClassify, + MPNetTokenClassify, +) +from .mpnet_tokenizer import MPNetTokenizer + +__all__ = [ + "MPNetConfig", + "MPNetModel", + "MPNetMaskedLM", + "MPNetSequenceClassify", + "MPNetTokenClassify", + "MPNetQnA", + "MPNetMultipleChoice", + "MPNetTokenizer", +] diff --git a/zeromodels/models/mpnet/convert_mpnet_hf_to_keras.py b/zeromodels/models/mpnet/convert_mpnet_hf_to_keras.py new file mode 100644 index 00000000..c642443d --- /dev/null +++ b/zeromodels/models/mpnet/convert_mpnet_hf_to_keras.py @@ -0,0 +1,200 @@ +import re +from typing import Dict, Optional + +import numpy as np +from tqdm import tqdm + +from zeromodels.conversion.exceptions import WeightMappingError +from zeromodels.conversion.weight_transfer_util import transfer_weights + +WEIGHT_NAME_MAPPING = { + "embeddings/word_embeddings/embeddings": "embeddings.word_embeddings.weight", + "embeddings/position_embeddings/embeddings": "embeddings.position_embeddings.weight", + "embeddings/LayerNorm/gamma": "embeddings.LayerNorm.weight", + "embeddings/LayerNorm/beta": "embeddings.LayerNorm.bias", + "encoder_relative_attention_bias/embeddings": ( + "encoder.relative_attention_bias.weight" + ), + "intermediate_dense": "intermediate.dense", + "output_dense": "output.dense", + "attention_layernorm": "attention.LayerNorm", + "output_layernorm": "output.LayerNorm", + "pooler_dense/kernel": "pooler.dense.weight", + "pooler_dense/bias": "pooler.dense.bias", + "lm_head_dense/kernel": "lm_head.dense.weight", + "lm_head_dense/bias": "lm_head.dense.bias", + "lm_head_layernorm/gamma": "lm_head.layer_norm.weight", + "lm_head_layernorm/beta": "lm_head.layer_norm.bias", + "lm_head_decoder/kernel": "embeddings.word_embeddings.weight", + "lm_head_decoder/bias": "lm_head.decoder.bias", + "classifier_dense/kernel": "classifier.dense.weight", + "classifier_dense/bias": "classifier.dense.bias", + "classifier_out_proj/kernel": "classifier.out_proj.weight", + "classifier_out_proj/bias": "classifier.out_proj.bias", + "classifier/kernel": "classifier.weight", + "classifier/bias": "classifier.bias", + "qa_outputs/kernel": "qa_outputs.weight", + "qa_outputs/bias": "qa_outputs.bias", +} + +_OPTIONAL_WEIGHTS = ("classifier", "qa_outputs", "lm_head", "pooler_dense") + +_QKVO_RE = re.compile( + r"blocks_(\d+)_attention_attn/blocks_\d+_(q|k|v|o)/(kernel|bias)$" +) +_DENSE_RE = re.compile(r"blocks_(\d+)_(intermediate_dense|output_dense)/(kernel|bias)$") +_NORM_RE = re.compile( + r"blocks_(\d+)_(attention_layernorm|output_layernorm)/(gamma|beta)$" +) + + +def hf_name_for(path: str) -> Optional[str]: + if path in WEIGHT_NAME_MAPPING: + return WEIGHT_NAME_MAPPING[path] + + m = _QKVO_RE.match(path) + if m: + idx, proj, w = m.groups() + suffix = "weight" if w == "kernel" else "bias" + return f"encoder.layer.{idx}.attention.attn.{proj}.{suffix}" + + m = _DENSE_RE.match(path) + if m: + idx, layer, w = m.groups() + suffix = "weight" if w == "kernel" else "bias" + return f"encoder.layer.{idx}.{WEIGHT_NAME_MAPPING[layer]}.{suffix}" + + m = _NORM_RE.match(path) + if m: + idx, layer, w = m.groups() + suffix = "weight" if w == "gamma" else "bias" + return f"encoder.layer.{idx}.{WEIGHT_NAME_MAPPING[layer]}.{suffix}" + + return None + + +def normalize_hf_key(key: str) -> str: + if key.startswith("mpnet."): + key = key[len("mpnet.") :] + return key.replace("LayerNorm.gamma", "LayerNorm.weight").replace( + "LayerNorm.beta", "LayerNorm.bias" + ) + + +def transfer_mpnet_weights(keras_model, hf_state_dict: Dict[str, np.ndarray]) -> None: + hf = {normalize_hf_key(k): v for k, v in hf_state_dict.items()} + for weight in tqdm(keras_model.weights, desc="Transferring weights to Keras"): + hf_name = hf_name_for(weight.path) + if hf_name is None: + continue + if hf_name not in hf: + if weight.path.startswith(_OPTIONAL_WEIGHTS): + continue + raise WeightMappingError(weight.path, hf_name) + transfer_weights(weight.path, weight, hf[hf_name]) + + +if __name__ == "__main__": + import gc + import os + + import keras + import torch + from transformers import MPNetForMaskedLM + from transformers import MPNetModel as HFMPNetModel + + from zeromodels.conversion.weight_transfer_util import ( + copy_weights_by_path_suffix, + ) + from zeromodels.models.mpnet import MPNetMaskedLM, MPNetModel + + HF_TOKEN = os.environ.get("HF_TOKEN") + MPNET_VARIANTS = {"mpnet_base": "microsoft/mpnet-base"} + + rng = np.random.default_rng(0) + + for variant, hf_id in MPNET_VARIANTS.items(): + print(f"\n{'=' * 60}\nConverting: {variant} <- {hf_id}\n{'=' * 60}") + + hf_model = HFMPNetModel.from_pretrained(hf_id, token=HF_TOKEN).eval() + hf_mlm = MPNetForMaskedLM.from_pretrained(hf_id, token=HF_TOKEN).eval() + arch = MPNetModel.config_from_hf(hf_model.config.to_dict()) + + # Pad id 1 so the pad-offset position ids exercise the masked-cumsum path. + ids = rng.integers(2, arch["vocab_size"], (2, 16)).astype("int64") + mask = np.ones((2, 16), dtype="int64") + mask[0, 12:] = 0 + ids[0, 12:] = arch["pad_token_id"] + k_inputs = { + "input_ids": ids.astype("int32"), + "attention_mask": mask.astype("int32"), + } + pt = { + "input_ids": torch.from_numpy(ids), + "attention_mask": torch.from_numpy(mask), + } + valid = mask.astype(bool) + + with torch.no_grad(): + hf_out = hf_model(**pt) + hf_logits = hf_mlm(**pt).logits.detach().cpu().numpy() + hf_seq = hf_out.last_hidden_state.detach().cpu().numpy() + hf_pool = hf_out.pooler_output.detach().cpu().numpy() + + # One superset checkpoint = encoder + pooler + MLM head. HF splits these across + # MPNetModel (pooler) and MPNetForMaskedLM (MLM head), so merge both state dicts; + # keys differ only by the "mpnet." prefix normalize_hf_key strips. Each zeromodels + # class then loads its own subset out of this single file. + merged = {**dict(hf_mlm.state_dict()), **dict(hf_model.state_dict())} + keras_full = MPNetMaskedLM(**arch, add_pooler=True) + transfer_mpnet_weights(keras_full, merged) + + full_out = keras_full(k_inputs, training=False) + full_logits = keras.ops.convert_to_numpy(full_out["logits"]) + full_pool = keras.ops.convert_to_numpy(full_out["pooler_output"]) + d_mlm = float(np.abs(hf_logits[valid] - full_logits[valid]).max()) + d_pool = float(np.abs(hf_pool - full_pool).max()) + print(f" full checkpoint mlm diff: {d_mlm:.3e} pooler diff: {d_pool:.3e}") + if max(d_mlm, d_pool) > 1e-3: + raise ValueError(f"{variant}: full-checkpoint parity failed") + + out_path = f"{variant}.weights.h5" + keras_full.save_weights(out_path) + print(f" Saved single file -> {out_path}") + + # Verify the ONE file serves every view: reload into a same-class reference, copy + # each subset out by semantic path, and re-check parity end to end. + ref = MPNetMaskedLM(**arch, add_pooler=True) + ref.load_weights(out_path) + + mm = MPNetModel(**arch) + copy_weights_by_path_suffix(ref, mm) + mm_out = mm(k_inputs, training=False) + d_seq = float( + np.abs( + hf_seq[valid] + - keras.ops.convert_to_numpy(mm_out["last_hidden_state"])[valid] + ).max() + ) + d_mmpool = float( + np.abs(hf_pool - keras.ops.convert_to_numpy(mm_out["pooler_output"])).max() + ) + + mlm2 = MPNetMaskedLM(**arch) + copy_weights_by_path_suffix(ref, mlm2) + d_mlm2 = float( + np.abs( + hf_logits[valid] + - keras.ops.convert_to_numpy(mlm2(k_inputs, training=False))[valid] + ).max() + ) + print( + f" reload MPNetModel seq: {d_seq:.3e} pooler: {d_mmpool:.3e} | " + f"MPNetMaskedLM mlm: {d_mlm2:.3e}" + ) + if max(d_seq, d_mmpool, d_mlm2) > 1e-3: + raise ValueError(f"{variant}: single-file reload parity failed") + + del hf_model, hf_mlm, keras_full, ref, mm, mlm2 + keras.backend.clear_session() + gc.collect() diff --git a/zeromodels/models/mpnet/mpnet_config.py b/zeromodels/models/mpnet/mpnet_config.py new file mode 100644 index 00000000..04dad500 --- /dev/null +++ b/zeromodels/models/mpnet/mpnet_config.py @@ -0,0 +1,60 @@ +from zeromodels.base import BaseConfig + + +class MPNetConfig(BaseConfig): + r"""Configuration for the MPNet encoder ([`MPNetModel`]) and its task heads. + + MPNet is a BERT-style bidirectional encoder pre-trained with masked **and** + permuted language modelling. It differs from BERT in three ways that matter to the + implementation: there are no token-type (segment) embeddings, position ids are + offset past the padding id as in RoBERTa, and every attention layer adds a shared + **relative position bias** gathered from `relative_attention_num_buckets` buckets. + One `zm_config.json` (declaring the canonical [`MPNetModel`]) sits on each variant's + repo; the encoder, masked-LM, and task-head classes all load from it. Fields mirror + the model constructor and serialize flat. + + Args: + vocab_size (`int`, *optional*, defaults to 30527): + Token vocabulary size. + embed_dim (`int`, *optional*, defaults to 768): + Hidden size. + num_layers (`int`, *optional*, defaults to 12): + Number of transformer encoder layers. + num_heads (`int`, *optional*, defaults to 12): + Number of attention heads. + mlp_dim (`int`, *optional*, defaults to 3072): + Feed-forward intermediate size. + max_position_embeddings (`int`, *optional*, defaults to 512): + Maximum sequence length supported by the positional embeddings. + relative_attention_num_buckets (`int`, *optional*, defaults to 32): + Number of buckets the relative position bias is gathered from. The bias is + computed once and shared by every attention layer. + hidden_act (`str`, *optional*, defaults to `"gelu"`): + Activation used in the feed-forward blocks. + layer_norm_eps (`float`, *optional*, defaults to 1e-12): + LayerNorm epsilon. + pad_token_id (`int`, *optional*, defaults to 1): + Padding token id (positions are offset past it). + + Examples: + + ```python + >>> from zeromodels.models.mpnet import MPNetConfig, MPNetModel + + >>> configuration = MPNetConfig() + >>> model = MPNetModel(configuration) + >>> configuration = model.config + ```""" + + model_type = "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 = 512 + relative_attention_num_buckets: int = 32 + hidden_act: str = "gelu" + layer_norm_eps: float = 1e-12 + pad_token_id: int = 1 diff --git a/zeromodels/models/mpnet/mpnet_layers.py b/zeromodels/models/mpnet/mpnet_layers.py new file mode 100644 index 00000000..adef3dc6 --- /dev/null +++ b/zeromodels/models/mpnet/mpnet_layers.py @@ -0,0 +1,314 @@ +import keras +from keras import layers, ops + +from zeromodels.base.base_attention import fused_attention + + +def relative_position_bucket(relative_position, num_buckets=32, max_distance=128): + """Map raw relative positions to MPNet's bidirectional bucket ids. + + Half the buckets carry each direction; within a direction the first quarter is + exact and the rest are log-spaced up to ``max_distance``. Mirrors HF + ``MPNetEncoder.relative_position_bucket``, including its ``n = -relative_position`` + sign convention (the opposite of T5's). + """ + n = -relative_position + num_buckets //= 2 + ret = ops.cast(n < 0, "int32") * num_buckets + n = ops.abs(n) + + max_exact = num_buckets // 2 + is_small = n < max_exact + val_if_large = max_exact + ops.cast( + ops.log(ops.cast(ops.maximum(n, 1), "float32") / max_exact) + / ops.log(ops.cast(max_distance / max_exact, "float32")) + * (num_buckets - max_exact), + "int32", + ) + val_if_large = ops.minimum(val_if_large, num_buckets - 1) + return ret + ops.where(is_small, ops.cast(n, "int32"), val_if_large) + + +@keras.saving.register_keras_serializable(package="zeromodels") +class MPNetEmbeddings(layers.Layer): + """Constructs MPNet's input embeddings. + + Sums learned word and absolute-position embeddings, then applies LayerNorm and + dropout. MPNet has **no token-type embeddings**. Like RoBERTa, position ids are + derived from the non-padding mask -- each non-pad token numbered sequentially from + ``pad_token_id + 1``, pad tokens mapping to ``pad_token_id`` -- computed with a + masked ``cumsum`` rather than ``arange`` so the layer stays shape-polymorphic + across the TensorFlow / JAX / PyTorch backends. + + Args: + vocab_size: Token vocabulary size. + embed_dim: Embedding / model dimension. + max_position_embeddings: Size of the position-embedding table. + pad_token_id: Padding token id; positions are offset by this value. + layer_norm_eps: Epsilon for the embedding LayerNorm. + dropout: Dropout rate applied to the summed embeddings. + """ + + def __init__( + self, + vocab_size, + embed_dim, + max_position_embeddings, + pad_token_id=1, + layer_norm_eps=1e-12, + dropout=0.0, + **kwargs, + ): + super().__init__(**kwargs) + self.vocab_size = vocab_size + self.embed_dim = embed_dim + self.max_position_embeddings = max_position_embeddings + self.pad_token_id = pad_token_id + self.layer_norm_eps = layer_norm_eps + self.dropout_rate = dropout + + self.word_embeddings = layers.Embedding( + vocab_size, embed_dim, name="word_embeddings" + ) + self.position_embeddings = layers.Embedding( + max_position_embeddings, embed_dim, name="position_embeddings" + ) + self.layer_norm = layers.LayerNormalization( + epsilon=layer_norm_eps, name="LayerNorm" + ) + self.dropout = layers.Dropout(dropout) + + def call(self, input_ids, training=None): + mask = ops.cast(ops.not_equal(input_ids, self.pad_token_id), input_ids.dtype) + position_ids = ops.cumsum(mask, axis=1) * mask + self.pad_token_id + + embeddings = self.word_embeddings(input_ids) + self.position_embeddings( + position_ids + ) + embeddings = self.layer_norm(embeddings) + return self.dropout(embeddings, training=training) + + def compute_output_shape(self, input_shape): + return tuple(input_shape) + (self.embed_dim,) + + def get_config(self): + config = super().get_config() + config.update( + { + "vocab_size": self.vocab_size, + "embed_dim": self.embed_dim, + "max_position_embeddings": self.max_position_embeddings, + "pad_token_id": self.pad_token_id, + "layer_norm_eps": self.layer_norm_eps, + "dropout": self.dropout_rate, + } + ) + return config + + +@keras.saving.register_keras_serializable(package="zeromodels") +class MPNetRelativeAttentionBias(layers.Layer): + """MPNet's shared relative position bias, ``(1, num_heads, seq, seq)``. + + Holds the ``(num_buckets, num_heads)`` bias table. The bias depends only on the + sequence length, so the encoder computes it **once** and every attention layer adds + the same tensor -- matching HF, where ``MPNetEncoder.forward`` calls + ``compute_position_bias`` before the layer loop. + + Positions are built with ``cumsum(ones_like) - 1`` instead of ``arange`` to stay + shape-polymorphic (the reference uses a plain ``arange``, which is the same values; + it does **not** use the pad-offset position ids from the embeddings). + + Args: + num_buckets: Size of the bias table (``relative_attention_num_buckets``). + num_heads: Number of attention heads (the table's second axis). + max_distance: Largest relative distance given its own bucket. + """ + + def __init__(self, num_buckets, num_heads, max_distance=128, **kwargs): + super().__init__(**kwargs) + self.num_buckets = num_buckets + self.num_heads = num_heads + self.max_distance = max_distance + + def build(self, input_shape): + self.relative_attention_bias = self.add_weight( + name="embeddings", + shape=(self.num_buckets, self.num_heads), + initializer="zeros", + ) + self.built = True + + def call(self, input_ids): + positions = ops.cumsum(ops.ones_like(input_ids), axis=1)[0] - 1 + relative_position = positions[None, :] - positions[:, None] + bucket = relative_position_bucket( + relative_position, self.num_buckets, self.max_distance + ) + # (seq, seq, num_heads) -> (1, num_heads, seq, seq) + values = ops.take(self.relative_attention_bias, bucket, axis=0) + return ops.transpose(values, (2, 0, 1))[None] + + def compute_output_shape(self, input_shape): + return (1, self.num_heads, input_shape[-1], input_shape[-1]) + + def get_config(self): + config = super().get_config() + config.update( + { + "num_buckets": self.num_buckets, + "num_heads": self.num_heads, + "max_distance": self.max_distance, + } + ) + return config + + +@keras.saving.register_keras_serializable(package="zeromodels") +class MPNetSelfAttention(layers.Layer): + """MPNet multi-head self-attention (HF's ``attention.attn`` sub-block). + + Projects the input to query/key/value, adds the shared relative position bias to + the scores alongside the additive padding mask, and -- unlike BERT, where the + output projection sits in ``attention.output`` -- applies its own ``o`` projection + before returning. The residual and LayerNorm live in the encoder layer (HF's + ``attention.LayerNorm``). + + Args: + embed_dim: Model dimension. Must be divisible by ``num_heads``. + num_heads: Number of attention heads. + attention_dropout: Dropout rate applied to the attention weights. + block_prefix: Prefix for the q/k/v/o projection names. Carries the + encoder-layer index so each layer's weights get a unique path suffix + (required for backbone weight-sharing across task heads). + """ + + def __init__( + self, + embed_dim, + num_heads, + attention_dropout=0.0, + block_prefix=None, + **kwargs, + ): + super().__init__(**kwargs) + if embed_dim % num_heads != 0: + raise ValueError( + f"embed_dim ({embed_dim}) must be divisible by num_heads ({num_heads})." + ) + self.embed_dim = embed_dim + self.num_heads = num_heads + self.attention_dropout = attention_dropout + self.block_prefix = block_prefix if block_prefix is not None else "attention" + self.head_dim = embed_dim // num_heads + self.scale = self.head_dim**-0.5 + + prefix = f"{self.block_prefix}_" + self.query = layers.Dense(embed_dim, name=prefix + "q") + self.key = layers.Dense(embed_dim, name=prefix + "k") + self.value = layers.Dense(embed_dim, name=prefix + "v") + self.output_proj = layers.Dense(embed_dim, name=prefix + "o") + self.dropout = layers.Dropout(attention_dropout) + + def build(self, input_shape): + # Explicit build: this layer takes several tensors (states + mask + bias), which + # TF cannot auto-build from a symbolic additive mask. + hidden_shape = input_shape[0] if isinstance(input_shape, list) else input_shape + input_dim = hidden_shape[-1] + self.query.build((None, input_dim)) + self.key.build((None, input_dim)) + self.value.build((None, input_dim)) + self.output_proj.build((None, self.embed_dim)) + self.built = True + + def compute_output_shape(self, input_shape, *args, **kwargs): + return input_shape[0] if isinstance(input_shape, list) else input_shape + + def transpose_for_scores(self, x): + batch_size = ops.shape(x)[0] + seq_len = ops.shape(x)[1] + x = ops.reshape(x, (batch_size, seq_len, self.num_heads, self.head_dim)) + return ops.transpose(x, (0, 2, 1, 3)) + + def call( + self, hidden_states, attention_mask=None, position_bias=None, training=None + ): + query = self.transpose_for_scores(self.query(hidden_states)) + key = self.transpose_for_scores(self.key(hidden_states)) + value = self.transpose_for_scores(self.value(hidden_states)) + + # The relative bias is part of the additive term, so fold it into the mask and + # reuse the shared attention kernel. + bias = attention_mask + if position_bias is not None: + bias = ( + position_bias + if attention_mask is None + else position_bias + attention_mask + ) + + context = fused_attention( + query, + key, + value, + self.scale, + bias, + dropout=self.dropout, + training=training, + ) + context = ops.transpose(context, (0, 2, 1, 3)) + batch_size = ops.shape(context)[0] + context = ops.reshape(context, (batch_size, -1, self.embed_dim)) + return self.output_proj(context) + + def get_config(self): + config = super().get_config() + config.update( + { + "embed_dim": self.embed_dim, + "num_heads": self.num_heads, + "attention_dropout": self.attention_dropout, + "block_prefix": self.block_prefix, + } + ) + return config + + +@keras.saving.register_keras_serializable(package="zeromodels") +class MPNetFlattenChoices(layers.Layer): + """Merge the multiple-choice axis into the batch: ``(B, C, S) -> (B*C, S)``. + + Defining ``compute_output_shape`` keeps the dynamic reshape out of the + functional-build trace, so it builds on every backend (the JAX backend rejects a + symbolic ``(-1, None)`` reshape). + """ + + def call(self, inputs): + return ops.reshape(inputs, (-1, ops.shape(inputs)[-1])) + + def compute_output_shape(self, input_shape): + return (None, input_shape[-1]) + + +@keras.saving.register_keras_serializable(package="zeromodels") +class MPNetUnflattenChoices(layers.Layer): + """Inverse of :class:`MPNetFlattenChoices` for the scores: ``(B*C, 1) -> (B, C)``. + + Args: + num_choices: Number of choices ``C`` to fold back out of the batch. + """ + + def __init__(self, num_choices, **kwargs): + super().__init__(**kwargs) + self.num_choices = num_choices + + def call(self, inputs): + return ops.reshape(inputs, (-1, self.num_choices)) + + def compute_output_shape(self, input_shape): + return (None, self.num_choices) + + def get_config(self): + config = super().get_config() + config.update({"num_choices": self.num_choices}) + return config diff --git a/zeromodels/models/mpnet/mpnet_model.py b/zeromodels/models/mpnet/mpnet_model.py new file mode 100644 index 00000000..3d529da7 --- /dev/null +++ b/zeromodels/models/mpnet/mpnet_model.py @@ -0,0 +1,779 @@ +import keras +from keras import layers, ops + +from zeromodels.base import BaseModel, CheckpointSource + +from .mpnet_config import MPNetConfig +from .mpnet_layers import ( + MPNetEmbeddings, + MPNetFlattenChoices, + MPNetRelativeAttentionBias, + MPNetSelfAttention, + MPNetUnflattenChoices, +) + +MASK_NEG = -1e9 + +# All classes (encoder + masked-LM + task heads) share the variant's weights repo, +# whose zm_config.json declares the canonical MPNetModel encoder (model.weights.h5). +MPNET_HUB_SIBLINGS = frozenset( + { + "MPNetModel", + "MPNetMaskedLM", + "MPNetSequenceClassify", + "MPNetTokenClassify", + "MPNetQnA", + "MPNetMultipleChoice", + } +) + +_BACKBONE_KW = { + "vocab_size": 30527, + "embed_dim": 768, + "num_layers": 12, + "num_heads": 12, + "mlp_dim": 3072, + "max_position_embeddings": 512, + "relative_attention_num_buckets": 32, + "hidden_act": "gelu", + "layer_norm_eps": 1e-12, + "pad_token_id": 1, + "dropout": 0.0, + "attention_dropout": 0.0, +} + + +def mpnet_encoder_layer( + x, + attention_mask, + position_bias, + *, + embed_dim, + num_heads, + mlp_dim, + hidden_act, + layer_norm_eps, + dropout, + attention_dropout, + layer_idx, +): + """One MPNet transformer block: self-attention + feed-forward. + + Both sub-blocks use post-LayerNorm residuals. The attention output projection is + inside :class:`MPNetSelfAttention` (HF's ``attention.attn.o``), so this adds only + the residual and LayerNorm (HF's ``attention.LayerNorm``). + + Args: + x: Input token states ``(B, seq, embed_dim)``. + attention_mask: Additive mask ``(B, 1, 1, seq)`` (0 keep, large-negative drop). + position_bias: Shared relative bias ``(1, num_heads, seq, seq)``. + embed_dim: Model dimension. + num_heads: Number of attention heads. + mlp_dim: Feed-forward hidden dimension. + hidden_act: Feed-forward activation. + layer_norm_eps: Epsilon for the two LayerNorms. + dropout: Hidden dropout rate. + attention_dropout: Attention-weight dropout rate. + layer_idx: Encoder-layer index (used for unique layer names). + + Returns: + Token states ``(B, seq, embed_dim)``. + """ + prefix = f"blocks_{layer_idx}" + + attn = MPNetSelfAttention( + embed_dim, + num_heads, + attention_dropout=attention_dropout, + block_prefix=prefix, + name=f"{prefix}_attention_attn", + )(x, attention_mask=attention_mask, position_bias=position_bias) + attn = layers.Dropout(dropout)(attn) + attn = layers.Add(name=f"{prefix}_attention_add")([attn, x]) + attn = layers.LayerNormalization( + epsilon=layer_norm_eps, name=f"{prefix}_attention_layernorm" + )(attn) + + inter = layers.Dense(mlp_dim, name=f"{prefix}_intermediate_dense")(attn) + inter = layers.Activation(hidden_act, name=f"{prefix}_intermediate_act")(inter) + out = layers.Dense(embed_dim, name=f"{prefix}_output_dense")(inter) + out = layers.Dropout(dropout)(out) + out = layers.Add(name=f"{prefix}_output_add")([out, attn]) + out = layers.LayerNormalization( + epsilon=layer_norm_eps, name=f"{prefix}_output_layernorm" + )(out) + return out + + +def mpnet_backbone( + input_ids, + attention_mask, + *, + vocab_size, + embed_dim, + num_layers, + num_heads, + mlp_dim, + max_position_embeddings, + relative_attention_num_buckets, + pad_token_id, + hidden_act, + layer_norm_eps, + dropout, + attention_dropout, + add_pooler, +): + """MPNet embeddings + transformer encoder (+ optional pooler). + + The relative position bias is built once here and shared by every layer, matching + HF's ``MPNetEncoder.forward``. + + Returns ``(sequence_output, pooled_output)`` where ``pooled_output`` is ``None`` + when ``add_pooler`` is False. + """ + embeddings = MPNetEmbeddings( + vocab_size=vocab_size, + embed_dim=embed_dim, + max_position_embeddings=max_position_embeddings, + pad_token_id=pad_token_id, + layer_norm_eps=layer_norm_eps, + dropout=dropout, + name="embeddings", + )(input_ids) + + position_bias = MPNetRelativeAttentionBias( + num_buckets=relative_attention_num_buckets, + num_heads=num_heads, + name="encoder_relative_attention_bias", + )(input_ids) + + mask = ops.cast(attention_mask, "float32") + mask = ops.expand_dims(ops.expand_dims(mask, 1), 1) + mask = (1.0 - mask) * MASK_NEG + + x = embeddings + for i in range(num_layers): + x = mpnet_encoder_layer( + x, + mask, + position_bias, + embed_dim=embed_dim, + num_heads=num_heads, + mlp_dim=mlp_dim, + hidden_act=hidden_act, + layer_norm_eps=layer_norm_eps, + dropout=dropout, + attention_dropout=attention_dropout, + layer_idx=i, + ) + + sequence_output = x + pooled_output = None + if add_pooler: + first_token = sequence_output[:, 0] + pooled_output = layers.Dense(embed_dim, activation="tanh", name="pooler_dense")( + first_token + ) + return sequence_output, pooled_output + + +def mpnet_classification_head(sequence_output, embed_dim, num_classes, dropout_rate): + """MPNet's sentence-level head: ```` token -> dropout -> dense/tanh -> out_proj. + + Matches HF ``MPNetClassificationHead`` (used by sequence classification), which + pools the first token itself rather than reading the encoder's pooler. + """ + x = sequence_output[:, 0] + x = layers.Dropout(dropout_rate)(x) + x = layers.Dense(embed_dim, activation="tanh", name="classifier_dense")(x) + x = layers.Dropout(dropout_rate)(x) + return x + + +@keras.saving.register_keras_serializable(package="zeromodels") +class MPNetModel(BaseModel): + """Instantiates the MPNet encoder backbone. + + MPNet embeds tokens with summed word and absolute-position embeddings (there are + **no** token-type embeddings), then applies a stack of bidirectional transformer + encoder layers. Each attention layer adds a **shared relative position bias**, + computed once from ``relative_attention_num_buckets`` log-spaced buckets, on top of + the usual scaled dot-product scores. Position ids are offset past the padding id as + in RoBERTa. An optional pooler applies a ``tanh`` dense projection to the first + (````) token. + + The model takes a dict of ``input_ids`` and ``attention_mask`` (both ``(B, seq)`` + int tensors, as produced by :class:`MPNetTokenizer`) and returns a dict with + ``last_hidden_state`` ``(B, seq, embed_dim)`` and, when ``add_pooler=True``, + ``pooler_output`` ``(B, embed_dim)``. + + References: + - [MPNet: Masked and Permuted Pre-training for Language Understanding](https://arxiv.org/abs/2004.09297) + + Args: + vocab_size: Integer, token vocabulary size. Defaults to `30527`. + embed_dim: Integer, model / embedding dimension. Defaults to `768`. + num_layers: Integer, number of transformer encoder layers. Defaults to `12`. + num_heads: Integer, number of attention heads. Defaults to `12`. + mlp_dim: Integer, feed-forward hidden dimension. Defaults to `3072`. + max_position_embeddings: Integer, size of the position-embedding table. + Defaults to `512`. + relative_attention_num_buckets: Integer, relative-bias bucket count. + Defaults to `32`. + hidden_act: String, feed-forward activation. Defaults to `"gelu"`. + layer_norm_eps: Float, LayerNorm epsilon. Defaults to `1e-12`. + pad_token_id: Integer, padding token id (also the position offset). + Defaults to `1`. + dropout: Float, hidden dropout rate. Defaults to `0.0`. + attention_dropout: Float, attention-weight dropout rate. Defaults to `0.0`. + add_pooler: Boolean, whether to add the ```` pooler. Defaults to `True`. + name: String, model name. Defaults to `"MPNetModel"`. + + Returns: + A Keras `Model` instance. + """ + + BASE_WEIGHT_CONFIG = None + HF_MODEL_TYPE = "mpnet" + config_class = MPNetConfig + HUB_REPO_SIBLINGS = MPNET_HUB_SIBLINGS + CHECKPOINT_SOURCE = CheckpointSource( + "MPNetMaskedLM", build_kwargs={"add_pooler": True} + ) + + @classmethod + def transfer_from_hf(cls, keras_model, state_dict): + from .convert_mpnet_hf_to_keras import transfer_mpnet_weights + + transfer_mpnet_weights(keras_model, state_dict) + + @classmethod + def config_from_hf(cls, hf_config): + return { + "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"], + "relative_attention_num_buckets": hf_config.get( + "relative_attention_num_buckets", 32 + ), + "hidden_act": hf_config.get("hidden_act", "gelu"), + "layer_norm_eps": hf_config.get("layer_norm_eps", 1e-12), + "pad_token_id": hf_config.get("pad_token_id", 1), + } + + def __init__( + self, + vocab_size=30527, + embed_dim=768, + num_layers=12, + num_heads=12, + mlp_dim=3072, + max_position_embeddings=512, + relative_attention_num_buckets=32, + hidden_act="gelu", + layer_norm_eps=1e-12, + pad_token_id=1, + dropout=0.0, + attention_dropout=0.0, + add_pooler=True, + name="MPNetModel", + **kwargs, + ): + for k in ("model", "hf_id", "url", "num_classes"): + kwargs.pop(k, None) + + # MPNet has no token-type embeddings, so the encoder takes these two inputs only. + inputs = { + "input_ids": layers.Input(shape=(None,), dtype="int32", name="input_ids"), + "attention_mask": layers.Input( + shape=(None,), dtype="int32", name="attention_mask" + ), + } + sequence_output, pooled_output = mpnet_backbone( + inputs["input_ids"], + inputs["attention_mask"], + 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, + pad_token_id=pad_token_id, + hidden_act=hidden_act, + layer_norm_eps=layer_norm_eps, + dropout=dropout, + attention_dropout=attention_dropout, + add_pooler=add_pooler, + ) + + outputs = {"last_hidden_state": sequence_output} + if pooled_output is not None: + outputs["pooler_output"] = pooled_output + + super().__init__(inputs=inputs, outputs=outputs, name=name, **kwargs) + + 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.dropout = dropout + self.attention_dropout = attention_dropout + self.add_pooler = add_pooler + + def get_config(self): + config = super().get_config() + config.update( + { + "vocab_size": self.vocab_size, + "embed_dim": self.embed_dim, + "num_layers": self.num_layers, + "num_heads": self.num_heads, + "mlp_dim": self.mlp_dim, + "max_position_embeddings": self.max_position_embeddings, + "relative_attention_num_buckets": self.relative_attention_num_buckets, + "hidden_act": self.hidden_act, + "layer_norm_eps": self.layer_norm_eps, + "pad_token_id": self.pad_token_id, + "dropout": self.dropout, + "attention_dropout": self.attention_dropout, + "add_pooler": self.add_pooler, + "name": self.name, + } + ) + return config + + @classmethod + def from_config(cls, config): + return cls(**config) + + +@keras.saving.register_keras_serializable(package="zeromodels") +class MPNetMaskedLM(BaseModel): + """MPNet with the masked-and-permuted language-modeling head. + + Wraps an :class:`MPNetModel` backbone (no pooler) and attaches MPNet's LM head -- a + dense transform with ``gelu`` then LayerNorm, followed by a vocabulary projection -- + producing token logits ``(B, seq, vocab_size)``. The head's weights are part of the + pretrained checkpoint, so loading restores a ready-to-use fill-mask model. + + References: + - [MPNet: Masked and Permuted Pre-training for Language Understanding](https://arxiv.org/abs/2004.09297) + + Args: + See :class:`MPNetModel` for the backbone arguments. + add_pooler: Boolean, whether to also carry the backbone's ```` pooler and + return it alongside the logits. Defaults to `False`. + name: String, model name. Defaults to `"MPNetMaskedLM"`. + + Returns: + A Keras `Model` instance. + """ + + BASE_WEIGHT_CONFIG = None + HF_MODEL_TYPE = "mpnet" + config_class = MPNetConfig + HUB_REPO_SIBLINGS = MPNET_HUB_SIBLINGS + CHECKPOINT_SOURCE = CheckpointSource( + "MPNetMaskedLM", build_kwargs={"add_pooler": True} + ) + + @classmethod + def transfer_from_hf(cls, keras_model, state_dict): + from .convert_mpnet_hf_to_keras import transfer_mpnet_weights + + transfer_mpnet_weights(keras_model, state_dict) + + @classmethod + def config_from_hf(cls, hf_config): + return MPNetModel.config_from_hf(hf_config) + + def __init__(self, add_pooler=False, name="MPNetMaskedLM", **kwargs): + for k in ("model", "hf_id", "url", "num_classes"): + kwargs.pop(k, None) + cfg = {**_BACKBONE_KW, **kwargs} + self._cfg = dict(cfg) + + backbone = MPNetModel(**cfg, add_pooler=add_pooler, name=f"{name}_backbone") + x = backbone.output["last_hidden_state"] + x = layers.Dense(cfg["embed_dim"], name="lm_head_dense")(x) + x = layers.Activation("gelu", name="lm_head_act")(x) + x = layers.LayerNormalization( + epsilon=cfg["layer_norm_eps"], name="lm_head_layernorm" + )(x) + logits = layers.Dense(cfg["vocab_size"], name="lm_head_decoder")(x) + + # add_pooler=True carries the pooler so a single hosted checkpoint serves + # MPNetModel (encoder + pooler), MPNetMaskedLM (encoder + MLM head) and the task + # heads alike; each class loads its own subset. Default False keeps the plain + # fill-mask output. + if add_pooler: + outputs = { + "logits": logits, + "pooler_output": backbone.output["pooler_output"], + } + else: + outputs = logits + + super().__init__(inputs=backbone.input, outputs=outputs, name=name) + + self.add_pooler = add_pooler + + def get_config(self): + config = super().get_config() + config.update({**self._cfg, "add_pooler": self.add_pooler, "name": self.name}) + return config + + @classmethod + def from_config(cls, config): + return cls(**config) + + +@keras.saving.register_keras_serializable(package="zeromodels") +class MPNetSequenceClassify(BaseModel): + """MPNet sentence/sequence classifier. + + Wraps an :class:`MPNetModel` backbone (no pooler) and attaches MPNet's + classification head -- dropout, a ``tanh`` dense, dropout, then a projection to + ``num_classes`` -- applied to the first (````) token, producing logits + ``(B, num_classes)``. Note MPNet's head pools the token itself rather than reading + the encoder pooler, matching Hugging Face. The pretrained checkpoint has no task + head, so the classifier stays randomly initialized and ready for fine-tuning. + + References: + - [MPNet: Masked and Permuted Pre-training for Language Understanding](https://arxiv.org/abs/2004.09297) + + Args: + See :class:`MPNetModel` for the backbone arguments. + num_classes: Integer, number of output classes. Defaults to `2`. + classifier_dropout: Float, dropout inside the head. Defaults to `0.0`. + classifier_activation: String/callable, head activation (`"linear"` for + logits). Defaults to `"linear"`. + name: String, model name. Defaults to `"MPNetSequenceClassify"`. + + Returns: + A Keras `Model` instance. + """ + + BASE_WEIGHT_CONFIG = None + HF_MODEL_TYPE = "mpnet" + config_class = MPNetConfig + HUB_REPO_SIBLINGS = MPNET_HUB_SIBLINGS + CHECKPOINT_SOURCE = CheckpointSource( + "MPNetMaskedLM", build_kwargs={"add_pooler": True} + ) + + @classmethod + def transfer_from_hf(cls, keras_model, state_dict): + from .convert_mpnet_hf_to_keras import transfer_mpnet_weights + + transfer_mpnet_weights(keras_model, state_dict) + + @classmethod + def config_from_hf(cls, hf_config): + config = MPNetModel.config_from_hf(hf_config) + config["num_classes"] = ( + len(hf_config["id2label"]) + if "id2label" in hf_config + else hf_config.get("num_labels", 2) + ) + return config + + def __init__( + self, + num_classes=2, + classifier_dropout=0.0, + classifier_activation="linear", + name="MPNetSequenceClassify", + **kwargs, + ): + for k in ("model", "hf_id", "url", "add_pooler"): + kwargs.pop(k, None) + cfg = {**_BACKBONE_KW, **kwargs} + self._cfg = dict(cfg) + + backbone = MPNetModel(**cfg, add_pooler=False, name=f"{name}_backbone") + x = mpnet_classification_head( + backbone.output["last_hidden_state"], + cfg["embed_dim"], + num_classes, + classifier_dropout, + ) + logits = layers.Dense( + num_classes, activation=classifier_activation, name="classifier_out_proj" + )(x) + + super().__init__(inputs=backbone.input, outputs=logits, name=name) + + self.num_classes = num_classes + self.classifier_dropout = classifier_dropout + self.classifier_activation = classifier_activation + + def get_config(self): + config = super().get_config() + config.update( + { + **self._cfg, + "num_classes": self.num_classes, + "classifier_dropout": self.classifier_dropout, + "classifier_activation": self.classifier_activation, + "name": self.name, + } + ) + return config + + @classmethod + def from_config(cls, config): + return cls(**config) + + +@keras.saving.register_keras_serializable(package="zeromodels") +class MPNetTokenClassify(BaseModel): + """MPNet token classifier (e.g. NER / POS tagging). + + Wraps an :class:`MPNetModel` backbone (no pooler) and attaches dropout plus a dense + head applied per token, producing logits ``(B, seq, num_classes)``. The head is + randomly initialized from the pretrained checkpoint and meant for fine-tuning. + + References: + - [MPNet: Masked and Permuted Pre-training for Language Understanding](https://arxiv.org/abs/2004.09297) + + Args: + See :class:`MPNetModel` for the backbone arguments. + num_classes: Integer, number of token classes. Defaults to `2`. + classifier_dropout: Float, dropout before the classifier. Defaults to `0.0`. + classifier_activation: String/callable, head activation. Defaults to `"linear"`. + name: String, model name. Defaults to `"MPNetTokenClassify"`. + + Returns: + A Keras `Model` instance. + """ + + BASE_WEIGHT_CONFIG = None + HF_MODEL_TYPE = "mpnet" + config_class = MPNetConfig + HUB_REPO_SIBLINGS = MPNET_HUB_SIBLINGS + CHECKPOINT_SOURCE = CheckpointSource( + "MPNetMaskedLM", build_kwargs={"add_pooler": True} + ) + + @classmethod + def transfer_from_hf(cls, keras_model, state_dict): + from .convert_mpnet_hf_to_keras import transfer_mpnet_weights + + transfer_mpnet_weights(keras_model, state_dict) + + @classmethod + def config_from_hf(cls, hf_config): + return MPNetSequenceClassify.config_from_hf(hf_config) + + def __init__( + self, + num_classes=2, + classifier_dropout=0.0, + classifier_activation="linear", + name="MPNetTokenClassify", + **kwargs, + ): + for k in ("model", "hf_id", "url", "add_pooler"): + kwargs.pop(k, None) + cfg = {**_BACKBONE_KW, **kwargs} + self._cfg = dict(cfg) + + backbone = MPNetModel(**cfg, add_pooler=False, name=f"{name}_backbone") + x = layers.Dropout(classifier_dropout)(backbone.output["last_hidden_state"]) + logits = layers.Dense( + num_classes, activation=classifier_activation, name="classifier" + )(x) + + super().__init__(inputs=backbone.input, outputs=logits, name=name) + + self.num_classes = num_classes + self.classifier_dropout = classifier_dropout + self.classifier_activation = classifier_activation + + def get_config(self): + config = super().get_config() + config.update( + { + **self._cfg, + "num_classes": self.num_classes, + "classifier_dropout": self.classifier_dropout, + "classifier_activation": self.classifier_activation, + "name": self.name, + } + ) + return config + + @classmethod + def from_config(cls, config): + return cls(**config) + + +@keras.saving.register_keras_serializable(package="zeromodels") +class MPNetQnA(BaseModel): + """MPNet extractive question answering. + + Wraps an :class:`MPNetModel` backbone (no pooler) and attaches a dense head + producing two logits per token, split into ``start_logits`` / ``end_logits`` + (each ``(B, seq)``). The head is randomly initialized from the pretrained + checkpoint and meant for fine-tuning. + + References: + - [MPNet: Masked and Permuted Pre-training for Language Understanding](https://arxiv.org/abs/2004.09297) + + Args: + See :class:`MPNetModel` for the backbone arguments. + name: String, model name. Defaults to `"MPNetQnA"`. + + Returns: + A Keras `Model` instance. + """ + + BASE_WEIGHT_CONFIG = None + HF_MODEL_TYPE = "mpnet" + config_class = MPNetConfig + HUB_REPO_SIBLINGS = MPNET_HUB_SIBLINGS + CHECKPOINT_SOURCE = CheckpointSource( + "MPNetMaskedLM", build_kwargs={"add_pooler": True} + ) + + @classmethod + def transfer_from_hf(cls, keras_model, state_dict): + from .convert_mpnet_hf_to_keras import transfer_mpnet_weights + + transfer_mpnet_weights(keras_model, state_dict) + + @classmethod + def config_from_hf(cls, hf_config): + return MPNetModel.config_from_hf(hf_config) + + def __init__(self, name="MPNetQnA", **kwargs): + for k in ("model", "hf_id", "url", "num_classes", "add_pooler"): + kwargs.pop(k, None) + cfg = {**_BACKBONE_KW, **kwargs} + self._cfg = dict(cfg) + + backbone = MPNetModel(**cfg, add_pooler=False, name=f"{name}_backbone") + span = layers.Dense(2, name="qa_outputs")(backbone.output["last_hidden_state"]) + start_logits = span[..., 0] + end_logits = span[..., 1] + + super().__init__( + inputs=backbone.input, + outputs={"start_logits": start_logits, "end_logits": end_logits}, + name=name, + ) + + def get_config(self): + config = super().get_config() + config.update({**self._cfg, "name": self.name}) + return config + + @classmethod + def from_config(cls, config): + return cls(**config) + + +@keras.saving.register_keras_serializable(package="zeromodels") +class MPNetMultipleChoice(BaseModel): + """MPNet multiple-choice classifier (e.g. SWAG / RACE). + + Takes ``(B, num_choices, seq)`` inputs, folds the choice axis into the batch, runs + the :class:`MPNetModel` backbone (with pooler), scores each choice with a single + dense unit on the pooled ```` token, then folds the choice axis back out to give + logits ``(B, num_choices)``. The head is randomly initialized from the pretrained + checkpoint and meant for fine-tuning. + + References: + - [MPNet: Masked and Permuted Pre-training for Language Understanding](https://arxiv.org/abs/2004.09297) + + Args: + See :class:`MPNetModel` for the backbone arguments. + num_choices: Integer, number of choices per example. Defaults to `2`. + classifier_dropout: Float, dropout before the scorer. Defaults to `0.0`. + name: String, model name. Defaults to `"MPNetMultipleChoice"`. + + Returns: + A Keras `Model` instance. + """ + + BASE_WEIGHT_CONFIG = None + HF_MODEL_TYPE = "mpnet" + config_class = MPNetConfig + HUB_REPO_SIBLINGS = MPNET_HUB_SIBLINGS + CHECKPOINT_SOURCE = CheckpointSource( + "MPNetMaskedLM", build_kwargs={"add_pooler": True} + ) + + @classmethod + def transfer_from_hf(cls, keras_model, state_dict): + from .convert_mpnet_hf_to_keras import transfer_mpnet_weights + + transfer_mpnet_weights(keras_model, state_dict) + + @classmethod + def config_from_hf(cls, hf_config): + return MPNetModel.config_from_hf(hf_config) + + def __init__( + self, + num_choices=2, + classifier_dropout=0.0, + name="MPNetMultipleChoice", + **kwargs, + ): + for k in ("model", "hf_id", "url", "num_classes", "add_pooler"): + kwargs.pop(k, None) + cfg = {**_BACKBONE_KW, **kwargs} + self._cfg = dict(cfg) + + inputs = { + "input_ids": layers.Input( + shape=(num_choices, None), dtype="int32", name="input_ids" + ), + "attention_mask": layers.Input( + shape=(num_choices, None), dtype="int32", name="attention_mask" + ), + } + flat = MPNetFlattenChoices(name="flatten_choices") + backbone = MPNetModel(**cfg, add_pooler=True, name=f"{name}_backbone") + pooled = backbone( + { + "input_ids": flat(inputs["input_ids"]), + "attention_mask": flat(inputs["attention_mask"]), + } + )["pooler_output"] + x = layers.Dropout(classifier_dropout)(pooled) + scores = layers.Dense(1, name="classifier")(x) + logits = MPNetUnflattenChoices(num_choices, name="unflatten_choices")(scores) + + super().__init__(inputs=inputs, outputs=logits, name=name) + + self.num_choices = num_choices + self.classifier_dropout = classifier_dropout + + def get_config(self): + config = super().get_config() + config.update( + { + **self._cfg, + "num_choices": self.num_choices, + "classifier_dropout": self.classifier_dropout, + "name": self.name, + } + ) + return config + + @classmethod + def from_config(cls, config): + return cls(**config) diff --git a/zeromodels/models/mpnet/mpnet_tokenizer.py b/zeromodels/models/mpnet/mpnet_tokenizer.py new file mode 100644 index 00000000..ec119ecd --- /dev/null +++ b/zeromodels/models/mpnet/mpnet_tokenizer.py @@ -0,0 +1,107 @@ +from typing import List, Union + +import keras +from tokenizers import Tokenizer + +from zeromodels.base import BaseTokenizer + + +@keras.saving.register_keras_serializable(package="zeromodels") +class MPNetTokenizer(BaseTokenizer): + """MPNet WordPiece tokenizer (``tokenizers`` Rust backend). + + Loads the HuggingFace fast-tokenizer ``tokenizer.json`` for ``variant`` (or an + explicit ``tokenizer_file``), with MPNet's `` A `` / `` A B `` + post-processing baked into that file, plus truncation and padding. MPNet pairs + RoBERTa-style special tokens with a BERT-style WordPiece vocabulary. + + ``call`` returns only ``input_ids`` / ``attention_mask``: MPNet has no token-type + embeddings, so unlike :class:`~zeromodels.models.bert.BertTokenizer` no + ``token_type_ids`` entry is produced. + + Args: + variant: MPNet variant key (no default; pass this or ``tokenizer_file``). + tokenizer_file: Optional explicit ``tokenizer.json`` path (overrides variant). + max_seq_len: Truncation length (default 512); batches pad to the longest. + 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 = 512, + bos_token: str = "", + eos_token: str = "", + unk_token: str = "[UNK]", + pad_token: str = "", + mask_token: str = "", + **kwargs, + ): + super().__init__(**kwargs) + self.variant = variant + tokenizer_file = self.resolve_tokenizer_json_from_hf( + (f"zeromodels/{self.variant}" if self.variant else None), tokenizer_file + ) + self.tokenizer_file = tokenizer_file + self.max_seq_len = max_seq_len + self.bos_token = bos_token + self.eos_token = eos_token + self.unk_token = unk_token + self.pad_token = pad_token + self.mask_token = mask_token + + tok = Tokenizer.from_file(tokenizer_file) + self.bos_token_id = tok.token_to_id(bos_token) + self.eos_token_id = tok.token_to_id(eos_token) + self.unk_token_id = tok.token_to_id(unk_token) + self.pad_token_id = tok.token_to_id(pad_token) + self.mask_token_id = tok.token_to_id(mask_token) + tok.enable_truncation(max_length=max_seq_len) + tok.enable_padding(pad_id=self.pad_token_id, pad_token=pad_token) + self._tok = tok + + @classmethod + def from_hf(cls, repo, **kwargs): + from huggingface_hub import hf_hub_download + + return cls(tokenizer_file=hf_hub_download(repo, "tokenizer.json"), **kwargs) + + @property + def vocab_size(self) -> int: + return self._tok.get_vocab_size() + + def tokenize( + self, text: Union[str, List[str]] + ) -> Union[List[int], List[List[int]]]: + if isinstance(text, str): + return self._tok.encode(text, add_special_tokens=False).ids + encs = self._tok.encode_batch(text, add_special_tokens=False) + return [e.ids for e in encs] + + def detokenize(self, token_ids, skip_special_tokens: bool = True) -> str: + return self.decode(token_ids, skip_special_tokens=skip_special_tokens) + + def decode(self, ids, skip_special_tokens: bool = True) -> str: + return self._tok.decode( + self.to_id_list(ids), skip_special_tokens=skip_special_tokens + ) + + def call(self, inputs: Union[str, List[str]], text_pair=None): + return self.encode_batch_to_inputs(inputs, text_pair, token_type_ids=False) + + 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, + "bos_token": self.bos_token, + "eos_token": self.eos_token, + "unk_token": self.unk_token, + "pad_token": self.pad_token, + "mask_token": self.mask_token, + } + ) + return config