Add MaskedLMTask for protein/bio models (ESM-2, ProtBert) - #86
Draft
justinchuby wants to merge 4 commits into
Draft
Add MaskedLMTask for protein/bio models (ESM-2, ProtBert)#86justinchuby wants to merge 4 commits into
justinchuby wants to merge 4 commits into
Conversation
Add masked language modeling support for encoder-only models (ESM-2, BERT, RoBERTa, etc.). This is the highest-ROI task for protein language models like Meta's ESM-2 (3.3M+ downloads). New components: - MaskedLMTask: task defining masked-lm I/O contract (input_ids, attention_mask, token_type_ids -> logits) - BertForMaskedLM: encoder + MLM head (dense -> GELU -> LayerNorm -> decoder) - _MaskedLMHead: reusable MLM prediction head component - _rename_masked_lm_weight: weight renaming for ESM/BERT/RoBERTa conventions Testing: - MASKED_LM_CONFIGS in _test_configs.py (bert, esm, roberta) - TestBuildMaskedLMGraph with 4 test methods x 3 model types = 12 tests - YAML golden test cases in testdata/cases/masked-lm/ (3 models) - model_coverage_test.py updated with MASKED_LM_CONFIGS - 'masked-lm' added to schema.json task_type enum Documentation: - .github/skills/adding-new-task-types/SKILL.md: comprehensive guide for adding new task types, using MaskedLMTask as worked example Also: - Add esm. prefix stripping to _rename_bert_weight (was missing) - Register masked-lm task in TASK_REGISTRY - Export BertForMaskedLM from models package Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Performance Comparison
|
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
🏗️ Architecture Diff
No architecture changes detected. ✅ Legend: ⚪ No change · 🔵 Minor (attrs/inits) · 🟡 Moderate (nodes added/removed) · 🔴 Major (interface changed) |
Contributor
There was a problem hiding this comment.
Pull request overview
Adds a new masked language modeling (MLM) inference path to mobius by introducing a dedicated MaskedLMTask, an encoder+MLM-head module (BertForMaskedLM), and corresponding tests / golden cases so encoder-only models (BERT/RoBERTa/ESM) can produce per-token vocabulary logits.
Changes:
- Introduces
MaskedLMTaskand registers it astask_type: masked-lm(including YAML schema + new golden cases). - Adds
BertForMaskedLMand_MaskedLMHead, plus MLM-specific HF→mobius weight renaming (includingesm.prefix handling). - Extends test config/coverage lists and adds graph-build tests for masked-lm graphs.
Reviewed changes
Copilot reviewed 11 out of 11 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/model_coverage_test.py | Includes MASKED_LM_CONFIGS in coverage computation. |
| tests/build_graph_test.py | Adds parametrized build/checker tests for masked-lm graphs. |
| tests/_test_configs.py | Adds MASKED_LM_CONFIGS and includes them in ALL_CONFIGS. |
| testdata/cases/schema.json | Extends task_type enum with "masked-lm". |
| testdata/cases/masked-lm/bert-base-uncased.yaml | Adds masked-lm golden case for BERT. |
| testdata/cases/masked-lm/esm2-8m.yaml | Adds masked-lm golden case for ESM-2. |
| testdata/cases/masked-lm/roberta-base.yaml | Adds masked-lm golden case for RoBERTa. |
| src/mobius/tasks/_masked_lm.py | New task wiring encoder-only MLM modules to logits output. |
| src/mobius/tasks/init.py | Exports MaskedLMTask and registers "masked-lm" in TASK_REGISTRY. |
| src/mobius/models/bert.py | Adds BertForMaskedLM, _MaskedLMHead, and MLM weight renaming; extends _rename_bert_weight to strip esm.. |
| src/mobius/models/init.py | Exports BertForMaskedLM. |
Comment on lines
+399
to
+402
| _MASKED_LM_MODEL_CONFIGS: list[tuple[str, dict]] = [ | ||
| (mt, ov) for mt, ov, _ in MASKED_LM_CONFIGS | ||
| ] | ||
|
|
There was a problem hiding this comment.
_MASKED_LM_MODEL_CONFIGS is defined but never used. This will likely trigger an unused-variable lint failure (e.g., Ruff F841). Remove it or use it in the parametrization (e.g., in place of re-deriving params).
Suggested change
| _MASKED_LM_MODEL_CONFIGS: list[tuple[str, dict]] = [ | |
| (mt, ov) for mt, ov, _ in MASKED_LM_CONFIGS | |
| ] |
Comment on lines
+341
to
+342
| # Top-level bias on the BERT predictions head | ||
| "cls.predictions.bias": "lm_head.bias", |
There was a problem hiding this comment.
The MLM weight renaming maps cls.predictions.bias to lm_head.bias, and _rename_masked_lm_weight also passes through lm_head.bias unchanged. But _MaskedLMHead only defines lm_head.decoder.bias (via Linear(..., bias=True)) and has no lm_head.bias parameter, so this bias will be skipped during weight application and logits will be wrong. Map HF bias keys (cls.predictions.bias and/or lm_head.bias) to the actual parameter name used by the module (likely lm_head.decoder.bias), or add an explicit lm_head.bias parameter and ensure it’s used/tied correctly.
Suggested change
| # Top-level bias on the BERT predictions head | |
| "cls.predictions.bias": "lm_head.bias", | |
| # Top-level bias on the BERT predictions head maps to decoder bias | |
| "cls.predictions.bias": "lm_head.decoder.bias", |
| Inputs: | ||
| - input_ids: [batch, sequence_len] INT64 | ||
| - attention_mask: [batch, sequence_len] INT64 | ||
| - token_type_ids: [batch, sequence_len] INT64 (optional, for BERT) |
There was a problem hiding this comment.
The docstring says token_type_ids is optional, but build() always creates it as a required graph input and always passes it to the module. Either update the docstring to reflect that token_type_ids is always present (and should be zeros / type_vocab_size=1 for models that don’t use it), or implement an actually-optional input contract.
Suggested change
| - token_type_ids: [batch, sequence_len] INT64 (optional, for BERT) | |
| - token_type_ids: [batch, sequence_len] INT64 | |
| Always provided as an input. Models that do not use token types | |
| should configure ``type_vocab_size=1`` and/or ignore this tensor | |
| (e.g., pass or store all zeros). |
Comment on lines
295
to
+309
| def _rename_bert_weight(name: str) -> str | None: | ||
| """Rename a single HF BERT weight to our convention. | ||
|
|
||
| Strips model prefixes (bert./roberta.), collapses nested HF naming | ||
| 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." or "roberta." prefix if present | ||
| # 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:] |
There was a problem hiding this comment.
New weight-renaming logic (_rename_masked_lm_weight and the updated _rename_bert_weight handling esm.) isn’t covered by any targeted unit tests, and silent mis-renames will only surface as runtime warnings/incorrect outputs when weights are applied. Add a small _test.py (similar to src/mobius/models/t5_test.py) that asserts key HF→Mobius name mappings for BERT/ESM/RoBERTa MLM heads and encoder weights.
justinchuby
added a commit
that referenced
this pull request
Apr 2, 2026
Bug fixes for PR #86 review: 1. (HIGH) Register BertForMaskedLM in the registry as bert_masked_lm, esm_masked_lm, roberta_masked_lm with task='masked-lm' and test_model_ids. Tests now use registry.get() instead of hardcoding. 2. (MEDIUM) Fix cls.predictions.bias mapping from 'lm_head.bias' to 'lm_head.decoder.bias' — the HF BERT output bias is shared with the decoder layer bias. 3. (LOW) Extract _rename_bert_encoder_weight() helper to eliminate duplicated encoder rename logic between _rename_bert_weight and _rename_masked_lm_weight. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Bug fixes for PR #86 review: 1. (HIGH) Register BertForMaskedLM in the registry as bert_masked_lm, esm_masked_lm, roberta_masked_lm with task='masked-lm' and test_model_ids. Tests now use registry.get() instead of hardcoding. 2. (MEDIUM) Fix cls.predictions.bias mapping from 'lm_head.bias' to 'lm_head.decoder.bias' — the HF BERT output bias is shared with the decoder layer bias. 3. (LOW) Extract _rename_bert_encoder_weight() helper to eliminate duplicated encoder rename logic between _rename_bert_weight and _rename_masked_lm_weight. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
justinchuby
force-pushed
the
justinchu/bio-models
branch
from
April 2, 2026 01:27
99de4a0 to
755b34d
Compare
- Add cls.* skip in _rename_masked_lm_weight for NSP head weights - Pin ESM-2 YAML revision to SHA c731040f (was 'main') Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
Add _generate_masked_lm() to scripts/generate_golden.py and load_torch_masked_lm_model()/torch_masked_lm_forward() to torch_reference.py. The generator uses AutoModelForMaskedLM and extracts logits at the first mask token position (falls back to position 1 if no mask token is found). Generate golden JSON files for all 4 cases lacking them: - testdata/golden/encoder/esm2-8m.json (feature-extraction) - testdata/golden/masked-lm/bert-base-uncased.json - testdata/golden/masked-lm/esm2-8m.json - testdata/golden/masked-lm/roberta-base.json Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds MaskedLMTask — a new task type for masked language model inference. This unlocks the #1 protein model use case (masked amino acid prediction) for ESM-2, ProtBert, and all encoder-only models.
Changes
MaskedLMTaskinsrc/mobius/tasks/_masked_lm.pyBertForMaskedLMmodel class with MLM head (dense → GELU → LayerNorm → decoder)_MaskedLMHeadreusable prediction head component_rename_masked_lm_weight)esm.prefix stripping fix for_rename_bert_weightmasked-lmtask inTASK_REGISTRYtestdata/cases/masked-lm/(3 models)masked-lmadded totask_typeenummodel_coverage_test.pyupdated withMASKED_LM_CONFIGSTesting
Impact
Unlocks MLM inference for ESM-2 (1.44M+ monthly downloads), ProtBert, and any encoder-only model.