diff --git a/scripts/generate_golden.py b/scripts/generate_golden.py index 95dec62e..9674efac 100644 --- a/scripts/generate_golden.py +++ b/scripts/generate_golden.py @@ -979,6 +979,55 @@ def _generate_image_classification(case: TestCase, json_path: Path, device: str) ) +def _generate_masked_lm(case: TestCase, json_path: Path, device: str) -> None: + """Generate golden data for a masked LM model (BERT, RoBERTa, ESM-2, etc.). + + Runs a forward pass with a masked input and captures the logits at the + position of the first mask token. For protein language models (ESM-2), + the mask token is ````; for BERT-style models it is ``[MASK]``. + """ + from mobius._testing.golden import save_golden_ref + from mobius._testing.torch_reference import ( + load_torch_masked_lm_model, + torch_masked_lm_forward, + ) + + model, tokenizer = load_torch_masked_lm_model( + case.model_id, + revision=case.revision, + device=device, + trust_remote_code=case.trust_remote_code, + ) + + encoded = tokenizer(case.prompts[0], return_tensors="np", padding=False) + input_ids = encoded["input_ids"] + attention_mask = encoded["attention_mask"] + token_type_ids = encoded.get("token_type_ids") + + # Forward pass → (batch, seq_len, vocab_size) + logits = torch_masked_lm_forward(model, input_ids, attention_mask, token_type_ids) + + # Find the first mask token position and extract its logits. + # Fall back to position 1 (first non-CLS token) if no mask token is found. + mask_id = tokenizer.mask_token_id + flat_ids = input_ids[0].tolist() + mask_pos = flat_ids.index(mask_id) if mask_id in flat_ids else 1 + mask_logits = logits[0, mask_pos, :] # (vocab_size,) + + golden = _extract_logits_golden(mask_logits) + + # Masked-LM is L4-only (no autoregressive generation) + save_golden_ref( + json_path, + top1_id=golden["top1_id"], + top2_id=golden["top2_id"], + top10_ids=golden["top10_ids"], + top10_logits=golden["top10_logits"], + logits_summary=golden["logits_summary"], + input_ids=input_ids, + ) + + def _detection_forced_size(model_id: str, trust_remote_code: bool) -> dict | None: """Return a fixed ``{height, width}`` size for object-detection export. @@ -1644,6 +1693,7 @@ def _hook(_module, _args, kwargs, output): _GENERATORS = { "text-generation": _generate_causal_lm, "feature-extraction": _generate_encoder, + "masked-lm": _generate_masked_lm, "seq2seq": _generate_seq2seq, "image-text-to-text": _generate_vision_language, "image-classification": _generate_image_classification, diff --git a/src/mobius/_registry.py b/src/mobius/_registry.py index a0a10410..bc273694 100644 --- a/src/mobius/_registry.py +++ b/src/mobius/_registry.py @@ -100,7 +100,7 @@ ) from mobius.models.bamba import BambaCausalLMModel from mobius.models.bart import BartForConditionalGeneration -from mobius.models.bert import BertModel +from mobius.models.bert import BertForMaskedLM, BertModel from mobius.models.blip import BlipVisionModel from mobius.models.blip2 import Blip2Model from mobius.models.clip import CLIPTextModel, CLIPVisionModel, SigLIPVisionModel @@ -682,6 +682,10 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None: "xlnet": ModelRegistration(BertModel, task="feature-extraction"), "xmod": ModelRegistration(BertModel, task="feature-extraction"), "yoso": ModelRegistration(BertModel, task="feature-extraction"), + # --- Masked LM (encoder + prediction head) --- + "bert_masked_lm": ModelRegistration(BertForMaskedLM, task="masked-lm"), + "esm_masked_lm": ModelRegistration(BertForMaskedLM, task="masked-lm"), + "roberta_masked_lm": ModelRegistration(BertForMaskedLM, task="masked-lm"), # --- Absolute positional embeddings (non-RoPE) --- "biogpt": ModelRegistration(GPT2CausalLMModel), "ctrl": ModelRegistration(CTRLCausalLMModel), @@ -1062,6 +1066,11 @@ def _create_default_registry() -> ModelRegistry: "xmod": "facebook/xmod-base", "yoso": "uw-madison/yoso-4096", + # --- Masked LM (encoder + prediction head) --- + "bert_masked_lm": "google-bert/bert-base-uncased", + "esm_masked_lm": "facebook/esm2_t6_8M_UR50D", + "roberta_masked_lm": "FacebookAI/roberta-base", + # --- Encoder-decoder --- "bart": "facebook/bart-base", "t5": "google-t5/t5-small", diff --git a/src/mobius/_testing/torch_reference.py b/src/mobius/_testing/torch_reference.py index 095e3211..ac3e5af1 100644 --- a/src/mobius/_testing/torch_reference.py +++ b/src/mobius/_testing/torch_reference.py @@ -484,6 +484,64 @@ def torch_encoder_forward( return outputs.last_hidden_state.cpu().numpy() +# --------------------------------------------------------------------------- +# Masked LM models (BERT, RoBERTa, ESM-2, etc.) +# --------------------------------------------------------------------------- + + +def load_torch_masked_lm_model( + model_id: str, + revision: str | None = None, + dtype: torch.dtype = torch.float32, + device: str = "cpu", + trust_remote_code: bool = False, +): + """Load a HuggingFace masked LM model for reference inference. + + Uses AutoModelForMaskedLM for BERT/RoBERTa/ESM-2-style architectures. + + Returns: + Tuple of (model, tokenizer). + """ + import transformers + + kwargs: dict = {"trust_remote_code": trust_remote_code} + if revision: + kwargs["revision"] = revision + tokenizer = transformers.AutoTokenizer.from_pretrained(model_id, **kwargs) + model = transformers.AutoModelForMaskedLM.from_pretrained( + model_id, + torch_dtype=dtype, + device_map=device, + **kwargs, + ) + model.eval() + return model, tokenizer + + +@torch.no_grad() +def torch_masked_lm_forward( + model, + input_ids: np.ndarray, + attention_mask: np.ndarray, + token_type_ids: np.ndarray | None = None, +) -> np.ndarray: + """Run a single forward pass on a HuggingFace masked LM model. + + Returns: + logits as numpy array [batch, seq_len, vocab_size]. + """ + device = next(model.parameters()).device + kwargs: dict = { + "input_ids": torch.from_numpy(input_ids).to(device), + "attention_mask": torch.from_numpy(attention_mask).to(device), + } + if token_type_ids is not None: + kwargs["token_type_ids"] = torch.from_numpy(token_type_ids).to(device) + outputs = model(**kwargs) + return outputs.logits.cpu().numpy() + + # --------------------------------------------------------------------------- # Seq2seq models (BART, T5, mBART, etc.) # --------------------------------------------------------------------------- diff --git a/src/mobius/models/__init__.py b/src/mobius/models/__init__.py index ec45276f..57baa476 100644 --- a/src/mobius/models/__init__.py +++ b/src/mobius/models/__init__.py @@ -10,6 +10,7 @@ "AutoencoderKLQwenImageModel", "BambaCausalLMModel", "BartForConditionalGeneration", + "BertForMaskedLM", "BertModel", "Blip2Model", "BloomCausalLMModel", @@ -157,7 +158,7 @@ from mobius.models.bamba import BambaCausalLMModel from mobius.models.bart import BartForConditionalGeneration from mobius.models.base import CausalLMModel, FusedGateUpCausalLMModel, LayerNormCausalLMModel -from mobius.models.bert import BertModel +from mobius.models.bert import BertForMaskedLM, BertModel from mobius.models.blip2 import Blip2Model from mobius.models.chatglm import ChatGLMCausalLMModel from mobius.models.clip import CLIPVisionModel, SigLIPVisionModel diff --git a/src/mobius/models/bert.py b/src/mobius/models/bert.py index 13b227d5..90b68aa2 100644 --- a/src/mobius/models/bert.py +++ b/src/mobius/models/bert.py @@ -1,7 +1,7 @@ # Copyright (c) Microsoft Corporation. # Licensed under the MIT License. -"""BERT encoder-only model with HF-aligned weight naming. +"""BERT encoder-only model and masked LM with HF-aligned weight naming. HF BERT uses deeply nested naming for encoder layers: attention.self.query / attention.self.key / attention.self.value @@ -11,8 +11,12 @@ embeddings.LayerNorm Module attributes here match HF conventions to eliminate the -rename dict entirely. Only prefix stripping (bert./roberta.) +rename dict entirely. Only prefix stripping (bert./roberta./esm.) and gamma/beta compat remain in preprocess_weights. + +BertForMaskedLM adds a masked language model head on top of the +encoder for predicting masked tokens. Supports BERT, RoBERTa, ESM-2, +and similar encoder-only architectures. """ from __future__ import annotations @@ -283,26 +287,13 @@ def forward(self, op: OpBuilder, hidden_states: ir.Value, attention_mask: ir.Val _PARAM_RENAMES = {"gamma": "weight", "beta": "bias"} -def _rename_bert_weight(name: str) -> str | None: - """Rename a single HF BERT weight to our convention. +def _rename_bert_encoder_weight(name: str) -> str: + """Collapse nested HF naming and handle gamma/beta compat. - Strips model prefixes (bert./roberta./esm./mpnet./etc.), collapses - nested HF naming (.self./.output.) to match the flat ONNX initializer - paths, and handles old-BERT gamma/beta compat. Returns None for - pooler/cls weights we don't need. + This is the shared encoder-weight rename logic used by both + ``_rename_bert_weight`` and ``_rename_masked_lm_weight``. + Assumes model prefix (bert./roberta./esm.) is already stripped. """ - # Strip model-type prefix (e.g. bert., roberta., esm., mpnet., squeezebert.) - # HF safetensors use the model class prefix before embeddings/encoder. - first_dot = name.find(".") - if first_dot > 0: - after = name[first_dot + 1 :] - if after.startswith(("embeddings.", "encoder.", "pooler.", "cls.", "lm_head.")): - name = after - - # Skip pooler, classification heads, and lm_head - if name.startswith(("pooler.", "cls.", "lm_head.")): - return None - # Collapse nested HF naming to match flat ONNX paths: # attention.self.query → attention.query # attention.output.dense → attention.dense @@ -318,3 +309,175 @@ def _rename_bert_weight(name: str) -> str | None: name = f"{parts[0]}.{_PARAM_RENAMES[parts[1]]}" return name + + +def _rename_bert_weight(name: str) -> str | None: + """Rename a single HF BERT weight to our convention. + + Strips model prefixes (bert./roberta./esm.), collapses nested HF naming + (.self./.output.) to match the flat ONNX initializer paths, and + handles old-BERT gamma/beta compat. Returns None for pooler/cls + weights we don't need. + """ + # Strip "bert." / "roberta." / "esm." prefix if present + if name.startswith("bert."): + name = name[5:] + elif name.startswith("roberta."): + name = name[8:] + elif name.startswith("esm."): + name = name[4:] + + # Skip pooler and classification heads + if name.startswith(("pooler.", "cls.")): + return None + + return _rename_bert_encoder_weight(name) + + +# --------------------------------------------------------------------------- +# Masked LM Head and BertForMaskedLM +# --------------------------------------------------------------------------- + +# BERT cls.predictions.transform.* -> our lm_head.* naming +_BERT_CLS_TO_LM_HEAD = { + "cls.predictions.transform.dense.": "lm_head.dense.", + "cls.predictions.transform.LayerNorm.": "lm_head.layer_norm.", + "cls.predictions.decoder.": "lm_head.decoder.", + # Top-level bias on the BERT predictions head + "cls.predictions.bias": "lm_head.decoder.bias", +} + + +class _MaskedLMHead(nn.Module): + """Masked LM prediction head: dense -> activation -> LayerNorm -> decoder. + + Matches HF ESM/RoBERTa naming: lm_head.dense, lm_head.layer_norm, + lm_head.decoder. BERT uses different naming (cls.predictions.*) which + is handled by preprocess_weights. + """ + + def __init__( + self, + hidden_size: int, + vocab_size: int, + hidden_act: str = "gelu", + layer_norm_eps: float = 1e-12, + ): + super().__init__() + self.dense = Linear(hidden_size, hidden_size, bias=True) + self.layer_norm = LayerNorm(hidden_size, eps=layer_norm_eps) + self.decoder = Linear(hidden_size, vocab_size, bias=True) + self._act_fn = ACT2FN[hidden_act] + + def forward(self, op: OpBuilder, hidden_states: ir.Value): + # hidden_states: (batch, seq, hidden_size) -> logits: (batch, seq, vocab) + x = self.dense(op, hidden_states) + x = self._act_fn(op, x) + x = self.layer_norm(op, x) + return self.decoder(op, x) + + +class BertForMaskedLM(nn.Module): + """BERT/ESM/RoBERTa encoder with masked language model head. + + Predicts vocabulary logits for each token position, used for + masked token prediction (fill-mask). + + Supports ESM-2 (protein masked LM), BERT, RoBERTa, and similar + encoder architectures. Output is per-token logits over the vocabulary. + + Replicates HuggingFace's ``EsmForMaskedLM`` / ``BertForMaskedLM`` / + ``RobertaForMaskedLM``. + """ + + default_task = "masked-lm" + category = "encoder" + + def __init__(self, config: ArchitectureConfig): + super().__init__() + self.config = config + self.embeddings = _BertEmbeddings( + vocab_size=config.vocab_size, + hidden_size=config.hidden_size, + max_position_embeddings=config.max_position_embeddings, + type_vocab_size=getattr(config, "type_vocab_size", 2), + layer_norm_eps=config.rms_norm_eps, + pad_token_id=config.pad_token_id or 0, + ) + self.encoder = _BertEncoder(config) + self.lm_head = _MaskedLMHead( + hidden_size=config.hidden_size, + vocab_size=config.vocab_size, + hidden_act=config.hidden_act, + layer_norm_eps=config.rms_norm_eps, + ) + + def forward( + self, + op: OpBuilder, + input_ids: ir.Value, + attention_mask: ir.Value, + token_type_ids: ir.Value, + ): + hidden_states = self.embeddings(op, input_ids, token_type_ids) + hidden_states = self.encoder(op, hidden_states, attention_mask) + logits = self.lm_head(op, hidden_states) + return logits + + def preprocess_weights( + self, state_dict: dict[str, torch.Tensor] + ) -> dict[str, torch.Tensor]: + """Rename HF weight names to match our parameter names. + + Handles ESM (esm.* prefix + lm_head.*), RoBERTa (roberta.* + + lm_head.*), and BERT (bert.* + cls.predictions.*) conventions. + """ + new_state_dict = {} + for name, tensor in state_dict.items(): + new_name = _rename_masked_lm_weight(name) + if new_name is not None: + new_state_dict[new_name] = tensor + return new_state_dict + + +def _rename_masked_lm_weight(name: str) -> str | None: + """Rename a single HF masked LM weight to our convention. + + Handles LM head weights (BERT cls.predictions.* -> lm_head.*, + ESM/RoBERTa lm_head.* passes through) and delegates encoder + weights to ``_rename_bert_weight``. + """ + # Strip "bert." / "roberta." / "esm." prefix from encoder weights + if name.startswith("bert."): + name = name[5:] + elif name.startswith("roberta."): + name = name[8:] + elif name.startswith("esm."): + name = name[4:] + + # Skip pooler (not needed for MLM) + if name.startswith("pooler."): + return None + + # Skip contact_head (ESM-specific, not needed for MLM) + if name.startswith("contact_head."): + return None + + # Map BERT cls.predictions.* -> lm_head.* + for bert_prefix, our_prefix in _BERT_CLS_TO_LM_HEAD.items(): + if name.startswith(bert_prefix): + name = our_prefix + name[len(bert_prefix) :] + return name + + # lm_head.* passes through directly (ESM/RoBERTa convention) + if name.startswith("lm_head."): + return name + + # Skip other cls.* heads (e.g. cls.seq_relationship for NSP) + if name.startswith("cls."): + return None + + # Encoder weights: delegate to shared rename logic. + # Prefix is already stripped, so pass directly to the encoder + # rename which handles HF naming collapse + gamma/beta compat. + return _rename_bert_encoder_weight(name) diff --git a/src/mobius/models/bert_test.py b/src/mobius/models/bert_test.py new file mode 100644 index 00000000..d68c6d7d --- /dev/null +++ b/src/mobius/models/bert_test.py @@ -0,0 +1,90 @@ +# Copyright (c) ONNX Project Contributors +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for BERT/ESM/RoBERTa weight-name rename helpers. + +These guard the ``preprocess_weights`` renaming used by ``BertModel`` and +``BertForMaskedLM`` so that silent mis-renames (which surface only as +mismatched initializer names during weight application) are caught early. +""" + +from __future__ import annotations + +import pytest + +from mobius.models.bert import _rename_bert_weight, _rename_masked_lm_weight + + +class TestRenameBertWeight: + @pytest.mark.parametrize( + "hf_name,expected", + [ + # Prefix stripping for the three supported encoder families. + ("bert.embeddings.word_embeddings.weight", "embeddings.word_embeddings.weight"), + ("roberta.embeddings.word_embeddings.weight", "embeddings.word_embeddings.weight"), + ("esm.embeddings.word_embeddings.weight", "embeddings.word_embeddings.weight"), + # Nested HF naming collapse (attention.self / attention.output). + ( + "encoder.layer.0.attention.self.query.weight", + "encoder.layer.0.attention.query.weight", + ), + ( + "encoder.layer.0.attention.output.dense.weight", + "encoder.layer.0.attention.dense.weight", + ), + ( + "encoder.layer.0.output.dense.weight", + "encoder.layer.0.dense.weight", + ), + # Old-BERT gamma/beta compat. + ("encoder.layer.0.output.LayerNorm.gamma", "encoder.layer.0.LayerNorm.weight"), + ("encoder.layer.0.output.LayerNorm.beta", "encoder.layer.0.LayerNorm.bias"), + ], + ) + def test_encoder_renames(self, hf_name: str, expected: str) -> None: + assert _rename_bert_weight(hf_name) == expected + + @pytest.mark.parametrize( + "hf_name", + [ + "bert.pooler.dense.weight", + "cls.predictions.decoder.weight", + "cls.seq_relationship.weight", + ], + ) + def test_skipped_weights_return_none(self, hf_name: str) -> None: + assert _rename_bert_weight(hf_name) is None + + +class TestRenameMaskedLMWeight: + @pytest.mark.parametrize( + "hf_name,expected", + [ + # BERT-style cls.predictions.* -> lm_head.* + ("cls.predictions.transform.dense.weight", "lm_head.dense.weight"), + ("cls.predictions.transform.LayerNorm.weight", "lm_head.layer_norm.weight"), + ("cls.predictions.decoder.weight", "lm_head.decoder.weight"), + ("cls.predictions.bias", "lm_head.decoder.bias"), + # ESM/RoBERTa lm_head.* passes through unchanged. + ("lm_head.dense.weight", "lm_head.dense.weight"), + ("lm_head.decoder.bias", "lm_head.decoder.bias"), + # Encoder weights still delegate to the shared rename. + ( + "esm.encoder.layer.0.attention.self.query.weight", + "encoder.layer.0.attention.query.weight", + ), + ], + ) + def test_masked_lm_renames(self, hf_name: str, expected: str) -> None: + assert _rename_masked_lm_weight(hf_name) == expected + + @pytest.mark.parametrize( + "hf_name", + [ + "esm.pooler.dense.weight", + "esm.contact_head.regression.weight", + "cls.seq_relationship.weight", + ], + ) + def test_skipped_weights_return_none(self, hf_name: str) -> None: + assert _rename_masked_lm_weight(hf_name) is None diff --git a/src/mobius/tasks/__init__.py b/src/mobius/tasks/__init__.py index 3abfc09d..a276b0b8 100644 --- a/src/mobius/tasks/__init__.py +++ b/src/mobius/tasks/__init__.py @@ -40,6 +40,7 @@ "Gemma4UnifiedTask", "Gemma4TextCausalLMTask", "HybridCausalLMTask", + "MaskedLMTask", "HybridQwenVLTask", "ImageClassificationTask", "ModelTask", @@ -101,6 +102,7 @@ from mobius.tasks._gemma4_assistant import Gemma4AssistantTask from mobius.tasks._hunyuan_vl_mot import HunYuanVLMoTTask from mobius.tasks._image_classification import ImageClassificationTask +from mobius.tasks._masked_lm import MaskedLMTask from mobius.tasks._masked_diffusion import MaskedDiffusionTask from mobius.tasks._moshi import MoshiDepformerTask, MoshiTemporalTask from mobius.tasks._multimodal import MultiModalTask @@ -138,6 +140,7 @@ "controlnet": ControlNetTask, "denoising": DenoisingTask, "feature-extraction": FeatureExtractionTask, + "masked-lm": MaskedLMTask, "masked-diffusion": MaskedDiffusionTask, "image-classification": ImageClassificationTask, "object-detection": ObjectDetectionTask, diff --git a/src/mobius/tasks/_masked_lm.py b/src/mobius/tasks/_masked_lm.py new file mode 100644 index 00000000..e22595d6 --- /dev/null +++ b/src/mobius/tasks/_masked_lm.py @@ -0,0 +1,60 @@ +# Copyright (c) ONNX Project Contributors +# SPDX-License-Identifier: Apache-2.0 + +"""Masked language modeling task for encoder-only models (BERT, ESM, RoBERTa, etc.).""" + +from __future__ import annotations + +import onnx_ir as ir +from onnxscript import nn + +from mobius._configs import BaseModelConfig +from mobius._model_package import ModelPackage +from mobius.tasks._base import ModelTask, _make_graph, _make_model + + +class MaskedLMTask(ModelTask): + """Encoder-only masked language modeling (predict masked tokens). + + The module must accept ``(op, input_ids, attention_mask, token_type_ids)`` + and return ``logits`` with shape ``[batch, sequence_len, vocab_size]``. + + Inputs: + - input_ids: [batch, sequence_len] INT64 + - attention_mask: [batch, sequence_len] INT64 + - token_type_ids: [batch, sequence_len] INT64 (always emitted; pass + zeros for models such as RoBERTa/ESM that do not use segment IDs) + + Outputs: + - logits: [batch, sequence_len, vocab_size] FLOAT + """ + + def build( + self, + module: nn.Module, + config: BaseModelConfig, + ) -> ModelPackage: + batch = ir.SymbolicDim("batch") + seq_len = ir.SymbolicDim("sequence_len") + + graph, builder = _make_graph() + op = builder.op + + input_ids = builder.input("input_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len]) + attention_mask = builder.input( + "attention_mask", dtype=ir.DataType.INT64, shape=[batch, seq_len] + ) + token_type_ids = builder.input( + "token_type_ids", dtype=ir.DataType.INT64, shape=[batch, seq_len] + ) + + logits = module( + op, + input_ids=input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + ) + + builder.add_output(logits, "logits") + + return ModelPackage({"model": _make_model(graph)}, config=config) diff --git a/testdata/cases/encoder/esm2-8m.yaml b/testdata/cases/encoder/esm2-8m.yaml index 68a73f86..8ea2a9e4 100644 --- a/testdata/cases/encoder/esm2-8m.yaml +++ b/testdata/cases/encoder/esm2-8m.yaml @@ -1,6 +1,6 @@ model_id: "facebook/esm2_t6_8M_UR50D" model_type: "esm" -revision: "main" +revision: "c731040fcd8d73dceaa04b0a8e6329b345b0f5df" task_type: "feature-extraction" dtype: "float32" diff --git a/testdata/cases/masked-lm/bert-base-uncased.yaml b/testdata/cases/masked-lm/bert-base-uncased.yaml new file mode 100644 index 00000000..cd87050a --- /dev/null +++ b/testdata/cases/masked-lm/bert-base-uncased.yaml @@ -0,0 +1,12 @@ +model_id: "google-bert/bert-base-uncased" +revision: "86b5e0934494bd15c9632b12f734a8a67f723594" +task_type: "masked-lm" +dtype: "float32" + +inputs: + prompts: + - "The quick brown [MASK] jumps over the lazy dog." + +level: "L4" + +notes: "BERT base uncased with masked LM head. Predicts masked tokens." diff --git a/testdata/cases/masked-lm/esm2-8m.yaml b/testdata/cases/masked-lm/esm2-8m.yaml new file mode 100644 index 00000000..b284fc9b --- /dev/null +++ b/testdata/cases/masked-lm/esm2-8m.yaml @@ -0,0 +1,12 @@ +model_id: "facebook/esm2_t6_8M_UR50D" +revision: "c731040fcd8d73dceaa04b0a8e6329b345b0f5df" +task_type: "masked-lm" +dtype: "float32" + +inputs: + prompts: + - "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG" + +level: "L4" + +notes: "ESM-2 8M with masked LM head. Protein language model for predicting masked amino acids." diff --git a/testdata/cases/masked-lm/roberta-base.yaml b/testdata/cases/masked-lm/roberta-base.yaml new file mode 100644 index 00000000..78a15421 --- /dev/null +++ b/testdata/cases/masked-lm/roberta-base.yaml @@ -0,0 +1,12 @@ +model_id: "FacebookAI/roberta-base" +revision: "e2da8e2f811d1448a5b465c236feacd80ffbac7b" +task_type: "masked-lm" +dtype: "float32" + +inputs: + prompts: + - "The quick brown jumps over the lazy dog." + +level: "L4" + +notes: "RoBERTa base with masked LM head. Uses token for predictions." diff --git a/testdata/cases/schema.json b/testdata/cases/schema.json index 41827882..d716c030 100644 --- a/testdata/cases/schema.json +++ b/testdata/cases/schema.json @@ -55,6 +55,7 @@ "image-text-to-text", "image-to-image", "image-to-text", + "masked-lm", "object-detection", "phi4mm-multimodal", "qwen35-mtp", diff --git a/testdata/golden/encoder/esm2-8m.json b/testdata/golden/encoder/esm2-8m.json index 4920c6af..11424bc9 100644 --- a/testdata/golden/encoder/esm2-8m.json +++ b/testdata/golden/encoder/esm2-8m.json @@ -14,22 +14,22 @@ 240 ], "top10_logits": [ - "0x1.f89c820000000p-1", - "0x1.a22c6c0000000p-1", - "0x1.8bd4140000000p-1", - "0x1.780c760000000p-1", - "0x1.74406c0000000p-1", - "0x1.6735360000000p-1", - "0x1.4d31ae0000000p-1", - "0x1.4d17340000000p-1", - "0x1.48a44c0000000p-1", - "0x1.3c190e0000000p-1" + "0x1.f89c6c0000000p-1", + "0x1.a22c5e0000000p-1", + "0x1.8bd40a0000000p-1", + "0x1.780c840000000p-1", + "0x1.7440560000000p-1", + "0x1.6735380000000p-1", + "0x1.4d31a80000000p-1", + "0x1.4d17420000000p-1", + "0x1.48a4300000000p-1", + "0x1.3c19140000000p-1" ], "logits_summary": [ - "0x1.f89c820000000p-1", - "-0x1.ec55a60000000p+1", - "-0x1.6f2c12d666666p-7", - "0x1.abcafbbbb5caep-2" + "0x1.f89c6c0000000p-1", + "-0x1.ec55a80000000p+1", + "-0x1.6f2c13bd9999ap-7", + "0x1.abcafa4ea9449p-2" ], "input_ids": [ 0, diff --git a/testdata/golden/masked-lm/bert-base-uncased.json b/testdata/golden/masked-lm/bert-base-uncased.json new file mode 100644 index 00000000..e30cb171 --- /dev/null +++ b/testdata/golden/masked-lm/bert-base-uncased.json @@ -0,0 +1,48 @@ +{ + "top1_id": 4937, + "top2_id": 3899, + "top10_ids": [ + 4937, + 3899, + 4562, + 2158, + 2666, + 2879, + 2028, + 2611, + 4702, + 3124 + ], + "top10_logits": [ + "0x1.3f00520000000p+3", + "0x1.28a7340000000p+3", + "0x1.1c69740000000p+3", + "0x1.17f0f00000000p+3", + "0x1.10a9640000000p+3", + "0x1.0da2a40000000p+3", + "0x1.015a900000000p+3", + "0x1.ea17740000000p+2", + "0x1.e409b60000000p+2", + "0x1.e0c91e0000000p+2" + ], + "logits_summary": [ + "0x1.3f00520000000p+3", + "-0x1.ec9bb60000000p+3", + "-0x1.23a32c11124d2p+2", + "0x1.426564917a4f0p+1" + ], + "input_ids": [ + 101, + 1996, + 4248, + 2829, + 103, + 14523, + 2058, + 1996, + 13971, + 3899, + 1012, + 102 + ] +} diff --git a/testdata/golden/masked-lm/esm2-8m.json b/testdata/golden/masked-lm/esm2-8m.json new file mode 100644 index 00000000..732a0700 --- /dev/null +++ b/testdata/golden/masked-lm/esm2-8m.json @@ -0,0 +1,103 @@ +{ + "top1_id": 20, + "top2_id": 15, + "top10_ids": [ + 20, + 15, + 7, + 10, + 4, + 6, + 14, + 5, + 8, + 12 + ], + "top10_logits": [ + "0x1.2e937e0000000p+3", + "0x1.75a1cc0000000p-2", + "0x1.6a99960000000p-2", + "0x1.441ec00000000p-2", + "0x1.c072400000000p-3", + "-0x1.d8add40000000p-7", + "-0x1.4dcdf80000000p-4", + "-0x1.e15ae80000000p-4", + "-0x1.1ba6180000000p-3", + "-0x1.a1bf8a0000000p-2" + ], + "logits_summary": [ + "0x1.2e937e0000000p+3", + "-0x1.fd60940000000p+3", + "-0x1.3ae86a35a2e8cp+2", + "0x1.aedc3dfcb9dbbp+2" + ], + "input_ids": [ + 0, + 20, + 15, + 11, + 7, + 10, + 16, + 9, + 10, + 4, + 15, + 8, + 12, + 7, + 10, + 12, + 4, + 9, + 10, + 8, + 15, + 9, + 14, + 7, + 8, + 6, + 5, + 16, + 4, + 5, + 9, + 9, + 4, + 8, + 7, + 8, + 10, + 16, + 7, + 12, + 7, + 16, + 13, + 12, + 5, + 19, + 4, + 10, + 8, + 4, + 6, + 19, + 17, + 12, + 7, + 5, + 11, + 14, + 10, + 6, + 19, + 7, + 4, + 5, + 6, + 6, + 2 + ] +} diff --git a/testdata/golden/masked-lm/roberta-base.json b/testdata/golden/masked-lm/roberta-base.json new file mode 100644 index 00000000..fd47a654 --- /dev/null +++ b/testdata/golden/masked-lm/roberta-base.json @@ -0,0 +1,48 @@ +{ + "top1_id": 324, + "top2_id": 4758, + "top10_ids": [ + 324, + 4758, + 23602, + 2173, + 2335, + 219, + 2143, + 19921, + 28133, + 26666 + ], + "top10_logits": [ + "0x1.ea240c0000000p+3", + "0x1.c60af00000000p+3", + "0x1.bd73d00000000p+3", + "0x1.b3ec280000000p+3", + "0x1.b02e300000000p+3", + "0x1.aa97d40000000p+3", + "0x1.9ec0b80000000p+3", + "0x1.9462de0000000p+3", + "0x1.903d360000000p+3", + "0x1.8e63320000000p+3" + ], + "logits_summary": [ + "0x1.ea240c0000000p+3", + "-0x1.2236dc0000000p+4", + "-0x1.6f00602fea530p+1", + "0x1.08999b5b0b62bp+2" + ], + "input_ids": [ + 0, + 133, + 2119, + 6219, + 50264, + 13855, + 81, + 5, + 22414, + 2335, + 4, + 2 + ] +} diff --git a/tests/_test_configs.py b/tests/_test_configs.py index 7bfcea0c..51b52dca 100644 --- a/tests/_test_configs.py +++ b/tests/_test_configs.py @@ -1464,6 +1464,16 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: ] +# --------------------------------------------------------------------------- +# Masked LM configs (task: masked-lm, module: BertForMaskedLM) +# --------------------------------------------------------------------------- +MASKED_LM_CONFIGS: list[tuple[str, dict, bool]] = [ + ("bert_masked_lm", {"hidden_act": "gelu", "type_vocab_size": 2}, True), + ("esm_masked_lm", {"hidden_act": "gelu", "type_vocab_size": 2}, True), + ("roberta_masked_lm", {"hidden_act": "gelu", "type_vocab_size": 1}, False), +] + + # --------------------------------------------------------------------------- # Seq2Seq (encoder-decoder) configs (task: seq2seq) # --------------------------------------------------------------------------- @@ -2530,6 +2540,7 @@ def _base_config(config_cls=None, **overrides) -> ArchitectureConfig: ALL_CONFIGS: list[tuple[str, dict, bool]] = ( CAUSAL_LM_CONFIGS + ENCODER_CONFIGS + + MASKED_LM_CONFIGS + SEQ2SEQ_CONFIGS + VISION_CONFIGS + DETECTION_CONFIGS diff --git a/tests/build_graph_test.py b/tests/build_graph_test.py index c63ce356..551d578b 100644 --- a/tests/build_graph_test.py +++ b/tests/build_graph_test.py @@ -29,6 +29,7 @@ DETECTION_CONFIGS, ENCODER_CONFIGS, LONGROPE_FACTORS, + MASKED_LM_CONFIGS, SEQ2SEQ_CONFIGS, SPEECH_CONFIGS, SSM_CONFIGS, @@ -459,6 +460,68 @@ def test_outputs_have_shapes_and_dtypes(self, model_type: str, config_overrides: _assert_outputs_have_shapes_and_dtypes(pkg, model_type) +# === Masked LM model configs (imported from _test_configs) === +_MASKED_LM_MODEL_PARAMS = _make_params(MASKED_LM_CONFIGS) + + +@pytest.mark.parametrize("model_type,config_overrides", _MASKED_LM_MODEL_PARAMS) +class TestBuildMaskedLMGraph: + """Verify that encoder-only models build valid masked LM ONNX graphs.""" + + def test_graph_builds_without_weights(self, model_type: str, config_overrides: dict): + config = _base_config(**config_overrides) + model_cls = registry.get(model_type) + module = model_cls(config) + task = get_task(_default_task_for_model(model_type)) + pkg = task.build(module, config) + model = pkg["model"] + + assert model.graph is not None + input_names = {inp.name for inp in model.graph.inputs} + assert "input_ids" in input_names + assert "attention_mask" in input_names + assert "token_type_ids" in input_names + + output_names = {out.name for out in model.graph.outputs} + assert "logits" in output_names + # No KV cache outputs for masked LM + assert not any(n.startswith("present.") for n in output_names) + # No hidden_state output (that's feature-extraction, not masked-lm) + assert "last_hidden_state" not in output_names + + def test_graph_has_lm_head_initializers(self, model_type: str, config_overrides: dict): + config = _base_config(**config_overrides) + model_cls = registry.get(model_type) + module = model_cls(config) + task = get_task(_default_task_for_model(model_type)) + pkg = task.build(module, config) + model = pkg["model"] + + init_names = list(model.graph.initializers) + assert len(init_names) > 0 + has_lm_head = any("lm_head" in n for n in init_names) + has_encoder = any("encoder" in n for n in init_names) + assert has_lm_head, "Should have LM head parameters" + assert has_encoder, "Should have encoder parameters" + + def test_onnx_checker_passes(self, model_type: str, config_overrides: dict): + """Run the ONNX CheckerPass to catch attribute/shape/type errors.""" + config = _base_config(**config_overrides) + model_cls = registry.get(model_type) + module = model_cls(config) + task = get_task(_default_task_for_model(model_type)) + pkg = task.build(module, config) + _run_onnx_checker(pkg, model_type) + + def test_outputs_have_shapes_and_dtypes(self, model_type: str, config_overrides: dict): + config = _base_config(**config_overrides) + model_cls = registry.get(model_type) + module = model_cls(config) + task = get_task(_default_task_for_model(model_type)) + pkg = task.build(module, config) + _assert_outputs_have_shapes_and_dtypes(pkg, model_type) + + # === Encoder-decoder model configs (imported from _test_configs) === _SEQ2SEQ_MODEL_CONFIGS: list[tuple[str, dict]] = [(mt, ov) for mt, ov, _ in SEQ2SEQ_CONFIGS] diff --git a/tests/model_coverage_test.py b/tests/model_coverage_test.py index 80a9da92..93c15fcf 100644 --- a/tests/model_coverage_test.py +++ b/tests/model_coverage_test.py @@ -59,6 +59,7 @@ ALL_CAUSAL_LM_CONFIGS, DETECTION_CONFIGS, ENCODER_CONFIGS, + MASKED_LM_CONFIGS, SEQ2SEQ_CONFIGS, SPEECH_CONFIGS, SSM_CONFIGS, @@ -77,6 +78,7 @@ def _l1_l3_model_types() -> set[str]: all_configs = ( ALL_CAUSAL_LM_CONFIGS + ENCODER_CONFIGS + + MASKED_LM_CONFIGS + SEQ2SEQ_CONFIGS + VISION_CONFIGS + DETECTION_CONFIGS