Add MaskedLMTask for protein/bio models (ESM-2, ProtBert)

scripts/generate_golden.py

@@ -555,12 +555,62 @@ 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 ``<mask>``; 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,
+    )
+
+
 # ---- Dispatcher ----
 
 # Map task_type strings to generator functions.
 _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,

src/mobius/_registry.py

@@ -80,7 +80,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
@@ -589,6 +589,14 @@ def _create_default_registry() -> ModelRegistry:
     reg.register("layoutlmv3", LayoutLMv3Model, task="feature-extraction")
     reg.register("modernbert", ModernBertModel, task="feature-extraction")
 
+    # --- Masked LM (encoder + prediction head) ---
+    for name in (
+        "bert_masked_lm",
+        "esm_masked_lm",
+        "roberta_masked_lm",
+    ):
+        reg.register(name, BertForMaskedLM, task="masked-lm")
+
     # --- Absolute positional embeddings (non-RoPE) ---
     reg.register("gpt2", GPT2CausalLMModel)
     for name in (
@@ -899,6 +907,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",

src/mobius/_testing/torch_reference.py

@@ -263,6 +263,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.)
 # ---------------------------------------------------------------------------

src/mobius/models/__init__.py

@@ -10,6 +10,7 @@
     "AutoencoderKLQwenImageModel",
     "BambaCausalLMModel",
     "BartForConditionalGeneration",
+    "BertForMaskedLM",
     "BertModel",
     "Blip2Model",
     "BloomCausalLMModel",
@@ -126,7 +127,7 @@
 from mobius.models.bamba import BambaCausalLMModel
 from mobius.models.bart import BartForConditionalGeneration
 from mobius.models.base import CausalLMModel, 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