Update FastPLMs runtime files

Add-only FastPLMs files-only publication. Checkpoint weights and complete-artifact attestations are unchanged.

LICENSES/FastPLMs-Apache-2.0.txt

@@ -0,0 +1,203 @@
+PLEASE NOTE THE APACHE LICENSE ONLY APPLIES TO THE CODE IN THE FastPLMs GITHUB AND ASSOCIATED HUGGINGFACE REPOSITORIES, NOT NECESSARILY THE MODEL WEIGHTS. THOSE LICENSES CAN BE FOUND HERE https://github.com/Synthyra/FastPLMs/tree/main/LICENSES
+
+                                Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

LICENSES/dplm/LICENSE

@@ -0,0 +1,202 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

LICENSES/dplm/PROVENANCE.md

@@ -0,0 +1,23 @@
+# DPLM checkpoint license provenance
+
+FastPLMs uses the ByteDance DPLM repository at immutable revision
+`8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d` as the official source for both
+DPLM1 and DPLM2.
+
+At that revision:
+
+- the repository contains the complete [Apache License 2.0](https://github.com/bytedance/dplm/blob/8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d/LICENSE); and
+- the [official README](https://github.com/bytedance/dplm/blob/8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d/README.md#overview)
+  defines the repository release as including the pretrained weights for the
+  DPLM family, specifically DPLM1 and DPLM2, alongside training and inference
+  implementations.
+
+FastPLMs therefore records the official DPLM1 and DPLM2 checkpoint weights as
+Apache-2.0. Converted Synthyra checkpoints retain that license and include the
+verbatim upstream `LICENSE`. The deterministic conversion identifiers are
+`dplm_to_fastplms_v1` and `dplm2_to_fastplms_v1`; neither adds restrictions to
+the upstream terms.
+
+Complete publication is permitted only after the ordinary FastPLMs artifact,
+state-parity, legal-inventory, and atomic-publication checks pass. This record
+does not change the terms of third-party training data or downstream outputs.

README.md

@@ -1,127 +1,200 @@
----
-library_name: transformers
-tags: []
----
-
-# NOTE
-The GitHub with the implementation and requirements can be found [here](https://github.com/Synthyra/FastPLMs.git).
-
-# DPLM2
-Synthyra DPLM2 checkpoints are HuggingFace AutoModel compatible and include FastPLMs embedding helpers.
-
-## Supported models
-```python
-model_dict = {
-    "Synthyra/DPLM2-150M": "airkingbd/dplm2_150m",
-    "Synthyra/DPLM2-650M": "airkingbd/dplm2_650m",
-    "Synthyra/DPLM2-3B": "airkingbd/dplm2_3b",
-}
-```
-
-## Use with transformers
-```python
-import torch
-from transformers import AutoModel, AutoModelForMaskedLM
-
-model_path = "Synthyra/DPLM2-150M"
-model = AutoModel.from_pretrained(model_path, trust_remote_code=True, dtype=torch.float16).eval()
-tokenizer = model.tokenizer
-
-batch = tokenizer(["MPRTEIN", "MSEQWENCE"], padding=True, return_tensors="pt")
-with torch.no_grad():
-    hidden = model(**batch).last_hidden_state
-
-mlm = AutoModelForMaskedLM.from_pretrained(model_path, trust_remote_code=True, dtype=torch.float16).eval()
-with torch.no_grad():
-    logits = mlm(**batch).logits
-```
-
-## Experimental test-time training
-
-TTT is disabled by default. Normal DPLM2 inference, embeddings, logits, and
-`state_dict()` keys are unchanged unless you explicitly call `model.ttt(...)`.
-The current implementation is experimental and trains only local LoRA adapters
-on the PLM backbone with masked language modeling on the test protein. It can
-help some difficult proteins, but it adds test-time compute and can degrade
-already confident predictions.
-
-```python
-metrics = mlm.ttt(
-    seq="MSTNPKPQRKTKRNT",
-    ttt_config={"steps": 3, "ags": 1, "batch_size": 1},
-)
-mlm.ttt_reset()
-print(metrics["losses"])
-```
-
-## DPLM2 modality types
-DPLM2 infers `type_ids` automatically from `input_ids` and `attention_mask` when they are not provided.
-
-## Attention backends
-
-`sdpa` (PyTorch Scaled Dot Product Attention) is the default.
-
-| Backend | Key | Notes |
-| :--- | :--- | :--- |
-| PyTorch SDPA | `"sdpa"` | Default. Exact numerics, stable on all hardware. |
-| Flash Attention | `"kernels_flash"` | Fastest on Ampere/Hopper GPUs. Requires `pip install kernels` (pre-built — no hours-long compilation). Outputs are not bitwise identical to SDPA due to online softmax reordering; differences are often small but not guaranteed to be inconsequential — use `"sdpa"` if exact numerics matter. |
-| Flex Attention | `"flex"` | Skips padding tokens via block mask — faster on variable-length batches. Near-exact numerics. First use compiles a Triton kernel (30–120 s). Best combined with `torch.compile`. |
-| Auto | `"auto"` | Picks the best available: `kernels_flash` → `flex` → `sdpa`. |
-
-Set via config before loading, or change on the model after loading (DPLM2 propagates the change to all attention layers immediately):
-
-```python
-from transformers import AutoConfig, AutoModel
-
-# Option 1: set before loading
-config = AutoConfig.from_pretrained("Synthyra/DPLM2-150M", trust_remote_code=True)
-config.attn_backend = "flex"
-model = AutoModel.from_pretrained("Synthyra/DPLM2-150M", config=config, trust_remote_code=True)
-
-# Option 2: set after loading
-model = AutoModel.from_pretrained("Synthyra/DPLM2-150M", trust_remote_code=True)
-model.attn_backend = "flex"  # propagates to all attention layers in-place
-```
-
-## Embed datasets
-All DPLM2 models inherit `EmbeddingMixin`, so you can call `model.embed_dataset(...)` directly.
-
-## Citations
-
-```bibtex
-@article{wang2024dplm2,
-  title={DPLM-2: A Multimodal Diffusion Protein Language Model},
-  author={Wang, Xinyou and Ye, Zaixiang and Huang, Fei and Cao, Dongyan and Liang, Shujian and Huang, Liang},
-  journal={arXiv preprint arXiv:2410.13782},
-  year={2024}
-}
-```
-
-```bibtex
-@misc{FastPLMs,
-  author={Hallee, Logan and Bichara, David and Gleghorn, Jason P.},
-  title={FastPLMs: Fast, efficient, protein language model inference from Huggingface AutoModel.},
-  year={2024},
-  url={https://huggingface.co/Synthyra/ESMplusplus_small},
-  DOI={10.57967/hf/3726},
-  publisher={Hugging Face}
-}
-```
-
-```bibtex
-@article{dong2024flexattention,
-  title={Flex Attention: A Programming Model for Generating Optimized Attention Kernels},
-  author={Dong, Juechu and Feng, Boyuan and Guessous, Driss and Liang, Yanbo and He, Horace},
-  journal={arXiv preprint arXiv:2412.05496},
-  year={2024}
-}
-```
-
-```bibtex
-@inproceedings{paszke2019pytorch,
-  title={PyTorch: An Imperative Style, High-Performance Deep Learning Library},
-  author={Paszke, Adam and Gross, Sam and Massa, Francisco and Lerer, Adam and Bradbury, James and Chanan, Gregory and Killeen, Trevor and Lin, Zeming and Gimelshein, Natalia and Antiga, Luca and Desmaison, Alban and K{\"o}pf, Andreas and Yang, Edward and DeVito, Zach and Raison, Martin and Tejani, Alykhan and Chilamkurthy, Sasank and Steiner, Benoit and Fang, Lu and Bai, Junjie and Chintala, Soumith},
-  booktitle={Advances in Neural Information Processing Systems 32},
-  year={2019}
-}
-```
+---
+library_name: transformers
+license: "apache-2.0"
+tags:
+  - protein-language-model
+  - fastplms
+---
+
+<!-- Generated from src/fastplms/models.toml. Do not edit. -->
+
+# Synthyra/DPLM2-150M
+
+This checkpoint packages the FastPLMs `DPLM2` implementation.
+
+Accepted inputs are tokenized amino-acid and structure tracks with explicit
+modality boundaries.
+Supported Transformers entry points are `AutoConfig`, `AutoModel`,
+`AutoModelForMaskedLM`, `AutoModelForSequenceClassification`,
+`AutoModelForTokenClassification`.
+
+## Install and platform requirements
+
+Install FastPLMs from the exact source revision paired with this model card:
+
+```bash
+python -m pip install \
+  "fastplms @ git+https://github.com/Synthyra/FastPLMs.git@1b9ce023f1e06571cf3e6324be0610ffa53e0a4a"
+```
+
+Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The declared CPU gate covers tiny offline contracts; published checkpoint throughput and parity require the documented device tier. The Hub quick start below requires network
+access on first download. For an air-gapped run, first build the manifest-pinned
+local artifact and use the offline form shown in the example.
+
+## Quick start
+
+```python
+from transformers import AutoModel
+
+model_id = "Synthyra/DPLM2-150M"
+model = AutoModel.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+).eval()
+```
+
+This example uses the published Hub repository. For offline validation, build
+the manifest-pinned artifact and replace `model_id` with its local
+`dist/hub/DPLM2-150M` path, then pass `local_files_only=True`.
+
+Leave attention unspecified for the Transformers default. Supported explicit
+choices are `sdpa`.
+Pass the selected name through `attn_implementation`.
+When an optimized backend cannot return full attention tensors,
+`output_attentions=True` emits one explicit runtime warning and uses a correctly
+masked eager implementation for that call only. The warning identifies the
+configured backend, effective backend, and reason. Configuration and later
+calls are unchanged.
+For BF16 execution, this family uses FP32 parameters with CUDA BF16 autocast.
+
+## Dataset embeddings
+
+The shared embedding API accepts sequences, `(id, sequence)` pairs,
+`EmbeddingInput` records, insertion-ordered `{id: sequence}` mappings, or a
+FASTA path. Results preserve order and duplicate identifiers:
+
+```python
+result = model.embed_dataset(
+    ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"],
+    batch_size=2,
+    pooling=("mean", "std"),
+)
+
+for record in result:
+    print(record.id, record.sequence, record.tensor.shape)
+```
+
+Set `full_embeddings=True` for one residue tensor with shape `(l, d)` per
+sequence. Set `output` to a directory for bounded-memory, transactional
+safetensors with ordered-prefix resume, or choose `format="sqlite"` for
+batch-level database commits and exact resume. Pooling excludes boundary,
+padding, and other non-biological positions.
+
+For a long FASTA run, stream completed batches into SQLite:
+
+```python
+persisted = model.embed_dataset(
+    "proteins.fasta",
+    batch_size=64,
+    pooling=("mean",),
+    output="protein-embeddings.sqlite",
+    format="sqlite",
+    resume=True,
+)
+```
+
+Resume verifies the input order, model state, tokenizer policy, backend, dtype,
+and pooling configuration. It never appends incompatible records to an
+existing run.
+
+## Amino-acid and structure co-generation
+
+DPLM2 uses separate structure and amino-acid tracks with modality-specific
+boundary and mask tokens:
+
+```python
+import torch
+from transformers import AutoModelForMaskedLM, AutoTokenizer
+
+model_id = "Synthyra/DPLM2-150M"
+tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
+generator = AutoModelForMaskedLM.from_pretrained(
+    model_id,
+    trust_remote_code=True,
+).cuda().eval()
+vocab = tokenizer.get_vocab()
+l = 64
+structure = [
+    vocab["<cls_struct>"],
+    *([vocab["<mask_struct>"]] * l),
+    vocab["<eos_struct>"],
+]
+amino_acids = [
+    vocab["<cls_aa>"],
+    *([vocab["<mask_aa>"]] * l),
+    vocab["<eos_aa>"],
+]
+input_ids = torch.tensor([structure + amino_acids], device="cuda")
+
+with torch.inference_mode():
+    generated = generator.generate(input_ids, max_iter=100)["output_tokens"]
+print(generated.shape)
+```
+
+Generic `cls_token`, `eos_token`, `mask_token`, and `unk_token` aliases are
+intentionally unset. Callers constructing multimodal tensors must choose the
+amino-acid or structure token explicitly. Raw amino-acid sequences remain
+supported by `model.embed_dataset(...)`.
+
+Plain `AutoModel` omits the optional ESM pooler because this co-generation
+checkpoint contains no trained pooler weights. Pass `add_pooling_layer=True`
+only when intentionally initializing and training that head.
+
+The checkpoint weights are Apache-2.0. The pinned ByteDance
+[LICENSE](https://github.com/bytedance/dplm/blob/8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d/LICENSE) and [README](https://github.com/bytedance/dplm/blob/8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d/README.md#overview) provide the immutable license
+basis for the pretrained DPLM1 and DPLM2 weights. Complete publication remains
+subject to all artifact, legal, parity, and atomic-publication preflights.
+
+## Runtime contract
+
+- Public input: Tokenized amino-acid and structure tracks with explicit modality boundaries
+- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification`
+- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head`
+- Attention implementations: `sdpa`
+- Precision policies: `default`
+- BF16 execution: `fp32_parameters_autocast`
+- Generation contract: `required`
+- Optional dependency group: `core`
+- Weight publication allowed: `true`
+- Weight license status: `resolved`
+- Redistributable: `true`
+- Complete weight publication required: `false`
+
+## Provenance
+
+- FastPLMs weights: `Synthyra/DPLM2-150M@182745b8dc5661f898481a4fa60a7af9d53385c4`
+- Runtime revision: `1b9ce023f1e06571cf3e6324be0610ffa53e0a4a`
+- Runtime source-tree SHA-256: `5de51d1c5c03b023e3db0af2d6a8849140af8011164caf84722faa356428efe1`
+- Runtime bundle SHA-256: `3ca371517961576b5af3c51ca4f07e486d91a11c8da3b0d7916d0bbc8f76be72`
+- Generator/schema version and complete/runtime-only attestations: recorded in `provenance.json`
+- Canonical transformed state SHA-256: `82e1751f59052b8de72b082517557db47947e8d9b4ac2f11278369e6c0cbf001`
+- Conversion equality attestation: recorded in `provenance.json`
+- Official checkpoint: `airkingbd/dplm2_150m@3451d984d06497f835ed49634bd68c9dfb54d730`
+- Artifact source: `official`
+- State transform: `dplm2_to_fastplms_v1`
+- BF16 execution: `fp32_parameters_autocast`
+- Tokenizer class: `fastplms.models.dplm2.tokenization_dplm2.DPLM2Tokenizer`
+- Pinned upstreams: `dplm`
+- Reference container: `reference-dplm`
+- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark`
+- Unresolved required file identities: `0`
+
+The local artifact records exact file identities, conversion provenance, source
+revisions, and legal texts in `provenance.json`. A nonzero unresolved count is a
+release blocker.
+
+## Validation boundary
+
+For tiers declared by the manifest, the release contract compares applicable
+semantic configuration, tokenizer behavior, state keys, shapes, dtypes,
+values, aliases, and representative inference with the pinned official
+implementation. This metadata does not by itself claim that a particular build
+passed, that one backend is faster, or that an output has biological or
+therapeutic validity.
+
+## License
+
+Checkpoint terms: Apache-2.0. The Hub model-card identifier is
+`apache-2.0`. Applicable source licenses, notices, attribution,
+and conversion records are distributed with the local artifact. Review them
+before use.

THIRD_PARTY_NOTICES.md

@@ -0,0 +1,99 @@
+# Third-party notices
+
+FastPLMs implements interfaces and checkpoint mappings for independently
+released protein models. The pinned repositories under `vendor/upstream/` are
+parity oracles. Production code does not import them, and runtime images do not
+contain them.
+
+This notice is informational and is not legal advice. A checkpoint license can
+differ from the license covering its source implementation. The typed inventory
+in `src/fastplms/models.toml` and the verbatim files under `LICENSES/` are the
+distribution record.
+
+## ANKH
+
+The pinned ANKH implementation and the mirrored ANKH checkpoints are identified
+as CC BY-NC-SA 4.0. FastPLMs displays those terms but does not enforce them in
+software. Users are responsible for determining whether their use and
+redistribution comply. The complete text is in `LICENSES/ankh/LICENSE.md`.
+
+## Profluent-E1
+
+Profluent identifies its E1 model code as Apache-2.0. The E1 weights and full
+release are subject to the Profluent-E1 Clickthrough License Agreement and the
+incorporated attribution requirements. Any E1 distribution must retain all of
+the following files:
+
+- `LICENSES/e1/LICENSE`, the Profluent-E1 agreement
+- `LICENSES/e1/ATTRIBUTION`, the attribution guidelines
+- `LICENSES/e1/NOTICE`, the required notice
+- `LICENSES/e1/Apache-2.0.txt`, the code license
+- `LICENSES/e1/BSD-3-Clause.txt`, covering the FlashAttention-derived padding
+  utility identified by the official E1 source
+- `LICENSES/e1/MODIFICATIONS.md`, the FastPLMs modified-file notice
+
+The exact text `Profluent-E1` must remain prominently displayed in E1
+documentation and at each launch of an executable E1 workflow, as required by
+the upstream attribution guidelines. Certain commercial outputs, including
+specified pharmaceutical and target-related outputs, can require the separate
+`Built with Profluent-E1` statement described in `ATTRIBUTION`.
+
+## DPLM
+
+The pinned ByteDance DPLM repository is Apache-2.0. Its
+[README](https://github.com/bytedance/dplm/blob/8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d/README.md#overview)
+explicitly defines the repository release as including pretrained DPLM1 and
+DPLM2 weights, and the same revision carries the complete
+[Apache-2.0 license](https://github.com/bytedance/dplm/blob/8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d/LICENSE).
+FastPLMs records both checkpoint families as Apache-2.0 and distributes the
+verbatim license plus `LICENSES/dplm/PROVENANCE.md`. Converted weights retain
+those terms and remain subject to the ordinary artifact and publication gates.
+
+## Biohub
+
+The pinned Biohub ESM implementation is MIT and includes a separate
+`THIRD_PARTY_NOTICE.md`; both files are distributed under
+`LICENSES/biohub-esm/`. The pinned Biohub Transformers fork is Apache-2.0, with
+its complete text under `LICENSES/biohub-transformers/`.
+
+## Boltz
+
+The pinned Boltz source is MIT. The verbatim notice is in
+`LICENSES/boltz/LICENSE`.
+
+## Meta ESM and OpenFold
+
+The pinned Meta ESM source is MIT. The pinned OpenFold source is Apache-2.0.
+Their verbatim texts and revision-specific provenance notices are under
+`LICENSES/fair-esm/` and `LICENSES/openfold/`.
+
+The native H100 ESMFold reference image applies the tracked
+`docker/constraints/openfold-sm90.patch` to the copied OpenFold `setup.py`.
+This build-only change restricts the CUDA extension to `sm90` and selects the
+C++17 standard required by the reference PyTorch version. It leaves the pinned
+submodule, extension source, model classes, checkpoint data, and public API
+unchanged. The complete modified-file record is in
+`LICENSES/openfold/MODIFICATIONS.md`.
+
+The isolated reference image also includes Apache-2.0 PyTorch Lightning,
+TorchMetrics, Lightning Utilities, and NVIDIA DLLogger. Their exact versions or
+revision are pinned in `docker/constraints/esmfold.txt`; OpenFold imports them
+eagerly, and FastPLMs production code does not depend on them. DLLogger's exact
+source identity and installed-license handling are recorded in
+`LICENSES/dllogger/PROVENANCE.md`.
+
+## ProteinTTT
+
+The optional test-time training workflow is validated against the pinned
+ProteinTTT repository under its MIT license. Its verbatim license and
+revision-specific provenance are under `LICENSES/protein-ttt/`.
+
+## Conversion and packaging record
+
+For every supported family, `src/fastplms/models.toml` records an immutable
+official checkpoint revision, an immutable FastPLMs checkpoint revision, file
+digests, a named state transformation, and a mechanism-level conversion record.
+Generated artifacts reproduce that record in `provenance.json`. A release or
+artifact build must fail when a required file identity, legal text, attribution
+notice, modified-file notice, upstream revision, or conversion record is absent
+or differs from its manifest digest.

config.json

@@ -1,42 +1,50 @@
 {
-  "aa_type": 1,
+  "_name_or_path": "facebook/esm2_t30_150M_UR50D",
+  "add_cross_attention": false,
   "architectures": [
     "EsmForDPLM2"
   ],
   "attention_probs_dropout_prob": 0.0,
-  "attn_backend": "sdpa",
   "auto_map": {
-    "AutoConfig": "modeling_dplm2.DPLM2Config",
-    "AutoModel": "modeling_dplm2.DPLM2Model",
-    "AutoModelForMaskedLM": "modeling_dplm2.DPLM2ForMaskedLM",
-    "AutoModelForSequenceClassification": "modeling_dplm2.DPLM2ForSequenceClassification",
-    "AutoModelForTokenClassification": "modeling_dplm2.DPLM2ForTokenClassification"
+    "AutoConfig": "modeling_fastplms.DPLM2Config",
+    "AutoModel": "modeling_fastplms.DPLM2Model",
+    "AutoModelForMaskedLM": "modeling_fastplms.DPLM2ForMaskedLM",
+    "AutoModelForSequenceClassification": "modeling_fastplms.DPLM2ForSequenceClassification",
+    "AutoModelForTokenClassification": "modeling_fastplms.DPLM2ForTokenClassification"
   },
   "classifier_dropout": null,
   "dplm_type": "dplm2_esm",
-  "dtype": "float32",
   "emb_layer_norm_before": false,
   "esmfold_config": null,
+  "fastplms_checkpoint_hash": "f86bc2df968f20de59aa45cbcfbeabdc6100ebe41ed587271d30cb1dcefdb1ee",
+  "fastplms_checkpoint_repo_id": "airkingbd/dplm2_150m",
+  "fastplms_checkpoint_revision": "3451d984d06497f835ed49634bd68c9dfb54d730",
+  "fastplms_model_id": "dplm2_150m",
+  "fastplms_release_tool_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "fastplms_release_tool_sha256": "1459b5d7d13d9b07bd97b3eee764f2ce73623e15e32d07ddf6825c2a9509afb9",
+  "fastplms_runtime_bundle_sha256": "3ca371517961576b5af3c51ca4f07e486d91a11c8da3b0d7916d0bbc8f76be72",
+  "fastplms_runtime_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "fastplms_source_tree_sha256": "5de51d1c5c03b023e3db0af2d6a8849140af8011164caf84722faa356428efe1",
+  "fastplms_weights_revision": "3451d984d06497f835ed49634bd68c9dfb54d730",
   "hidden_act": "gelu",
   "hidden_dropout_prob": 0.1,
   "hidden_size": 640,
   "initializer_range": 0.02,
   "intermediate_size": 2560,
+  "is_decoder": false,
   "is_folding_model": false,
   "layer_norm_eps": 1e-05,
   "mask_token_id": 32,
   "max_position_embeddings": 1026,
-  "model_type": "dplm2",
+  "model_type": "esm",
   "num_attention_heads": 20,
   "num_hidden_layers": 30,
   "pad_token_id": 1,
-  "pad_type": 2,
   "position_embedding_type": "rotary",
-  "struct_type": 0,
-  "tie_word_embeddings": false,
   "token_dropout": true,
-  "transformers_version": "4.57.6",
-  "use_cache": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.39.2",
+  "use_cache": false,
   "vocab_list": null,
   "vocab_size": 8229
 }

fastplms/__init__.py

@@ -0,0 +1,48 @@
+"""FastPLMs public package interface.
+
+The module uses lazy exports so importing :mod:`fastplms` does not initialize
+Torch, download checkpoints, construct tokenizers, or compile kernels.
+"""
+
+from __future__ import annotations
+
+from importlib import import_module
+from typing import Any
+
+__version__ = "1.0.0"
+
+_LAZY_EXPORTS = {
+    "CheckpointSource": ("fastplms.registry", "CheckpointSource"),
+    "EmbeddingInput": ("fastplms.embeddings", "EmbeddingInput"),
+    "EmbeddingRecord": ("fastplms.embeddings", "EmbeddingRecord"),
+    "EmbeddingResult": ("fastplms.embeddings", "EmbeddingResult"),
+    "FileDigest": ("fastplms.registry", "FileDigest"),
+    "ModelFamily": ("fastplms.registry", "ModelFamily"),
+    "ModelRegistry": ("fastplms.registry", "ModelRegistry"),
+    "ModelSpec": ("fastplms.registry", "ModelSpec"),
+    "OracleAsset": ("fastplms.registry", "OracleAsset"),
+    "RegistryError": ("fastplms.registry", "RegistryError"),
+    "RuntimeProfile": ("fastplms.runtime", "RuntimeProfile"),
+    "UpstreamSource": ("fastplms.registry", "UpstreamSource"),
+    "embed_dataset": ("fastplms.embeddings", "embed_dataset"),
+    "get_model_registry": ("fastplms.registry", "get_model_registry"),
+    "get_model_spec": ("fastplms.registry", "get_model_spec"),
+    "load_model_registry": ("fastplms.registry", "load_model_registry"),
+    "runtime_profile": ("fastplms.runtime", "runtime_profile"),
+}
+
+__all__ = ["__version__", *_LAZY_EXPORTS]
+
+
+def __getattr__(name: str) -> Any:
+    try:
+        module_name, attribute_name = _LAZY_EXPORTS[name]
+    except KeyError as error:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}") from error
+    value = getattr(import_module(module_name), attribute_name)
+    globals()[name] = value
+    return value
+
+
+def __dir__() -> list[str]:
+    return sorted(set(globals()).union(__all__))

fastplms/attention/__init__.py

@@ -0,0 +1,63 @@
+"""Shared attention backends, masks, and optional optimized kernels."""
+
+from ._core import (
+    VALID_ATTENTION_BACKENDS,
+    AttentionBackend,
+    BlockMask,
+    _ensure_flash_kernels_loaded,
+    _get_flex_attention_fn,
+    _get_flex_block_mask,
+    _kernels_flash_forward,
+    _kernels_flash_varlen_forward,
+    _unpad_input,
+    bool_to_additive_mask,
+    clear_flex_attention_caches,
+    create_block_mask,
+    flex_attention,
+    get_attention_mask,
+    get_attn_implementation,
+    index_first_axis,
+    index_put_first_axis,
+    kernels_flash_attention_func,
+    pad_input,
+    resolve_attention_backend,
+    resolve_attention_backend_for_call,
+    set_config_attn_implementation,
+    warn_attention_backend_fallback,
+)
+from .interfaces import (
+    FASTPLMS_ATTENTION_FUNCTIONS,
+    FASTPLMS_ATTENTION_MASKS,
+    FastPLMsAttentionMixin,
+    validate_transformers_attention_interfaces,
+)
+
+__all__ = [
+    "FASTPLMS_ATTENTION_FUNCTIONS",
+    "FASTPLMS_ATTENTION_MASKS",
+    "VALID_ATTENTION_BACKENDS",
+    "AttentionBackend",
+    "BlockMask",
+    "FastPLMsAttentionMixin",
+    "_ensure_flash_kernels_loaded",
+    "_get_flex_attention_fn",
+    "_get_flex_block_mask",
+    "_kernels_flash_forward",
+    "_kernels_flash_varlen_forward",
+    "_unpad_input",
+    "bool_to_additive_mask",
+    "clear_flex_attention_caches",
+    "create_block_mask",
+    "flex_attention",
+    "get_attention_mask",
+    "get_attn_implementation",
+    "index_first_axis",
+    "index_put_first_axis",
+    "kernels_flash_attention_func",
+    "pad_input",
+    "resolve_attention_backend",
+    "resolve_attention_backend_for_call",
+    "set_config_attn_implementation",
+    "validate_transformers_attention_interfaces",
+    "warn_attention_backend_fallback",
+]

fastplms/attention/_core.py

@@ -0,0 +1,779 @@
+"""Low-level attention kernels and mask construction.
+
+The public backend contract lives in :mod:`fastplms.attention`. Optional
+kernels are resolved only after a caller explicitly requests them, so importing
+FastPLMs never downloads or compiles code.
+"""
+
+from __future__ import annotations
+
+import warnings
+from collections import OrderedDict
+from collections.abc import Callable
+from enum import Enum
+from threading import RLock
+
+import torch
+from einops import rearrange
+from torch.nn import functional as F
+
+from ._kernel_lock import load_locked_kernel
+
+try:
+    from torch.nn.attention.flex_attention import BlockMask, create_block_mask, flex_attention
+except ImportError:
+    create_block_mask = None
+    flex_attention = None
+    BlockMask = None
+
+_MAX_FLEX_CACHE_ENTRIES = 128
+_compiled_flex_attention: OrderedDict[tuple, object] = OrderedDict()
+_flex_block_masks: OrderedDict[tuple, BlockMask] = OrderedDict()
+_flex_cache_lock = RLock()
+
+
+def _remember(cache: OrderedDict, key: tuple, value):
+    """Insert an item into a bounded least-recently-used cache."""
+    cache[key] = value
+    cache.move_to_end(key)
+    while len(cache) > _MAX_FLEX_CACHE_ENTRIES:
+        cache.popitem(last=False)
+    return value
+
+
+def clear_flex_attention_caches() -> None:
+    """Drop FastPLMs-owned compiled Flex callables and block masks.
+
+    This deliberately does not call :func:`torch.compiler.reset`, which would
+    clear process-global Torch compilation state owned by unrelated models.
+    Active forwards retain their local references and can complete safely.
+    """
+
+    with _flex_cache_lock:
+        _compiled_flex_attention.clear()
+        _flex_block_masks.clear()
+
+
+def _get_flex_attention_fn(
+    *,
+    device: torch.device | None = None,
+    dtype: torch.dtype | None = None,
+    shape: tuple[int, ...] | None = None,
+    sequence_lengths: tuple[int, ...] | None = None,
+    mask_semantics: str = "padding",
+):
+    """Return a compiled Flex callable for an explicit execution signature.
+
+    Compilation depends on execution shape, device, dtype, and mask semantics.
+    Per-example padding lengths are represented by the ``BlockMask`` argument
+    and must not create a new compiled graph for every batch composition.
+    """
+    if flex_attention is None:
+        return None
+    # Retain the keyword for compatibility with remote-code artifacts while
+    # deliberately excluding data-dependent lengths from the compile key.
+    del sequence_lengths
+    flex_mod = torch.nn.attention.flex_attention
+    if getattr(flex_mod, "_FLEX_ATTENTION_DISABLE_COMPILE_DEBUG", False):
+        return flex_attention
+    key = (
+        None if device is None else str(device),
+        None if dtype is None else str(dtype),
+        shape,
+        mask_semantics,
+    )
+    with _flex_cache_lock:
+        compiled = _compiled_flex_attention.get(key)
+        if compiled is None:
+            compiled = torch.compile(flex_attention, dynamic=False)
+            _remember(_compiled_flex_attention, key, compiled)
+        else:
+            _compiled_flex_attention.move_to_end(key)
+    return compiled
+
+
+def _get_flex_block_mask(
+    *,
+    mask_pattern: torch.Tensor,
+    batch_size: int,
+    query_length: int,
+    key_value_length: int,
+    device: torch.device,
+    dtype: torch.dtype | None,
+    mask_semantics: str,
+    mask_mod: Callable[
+        [torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor],
+        torch.Tensor,
+    ],
+) -> BlockMask:
+    """Return a bounded, exact-pattern cached Flex ``BlockMask``.
+
+    The complete pattern is transferred to the host once to avoid a CUDA
+    synchronization per batch row. Execution dtype remains part of the key
+    because compiled Flex plans can specialize on it even though the pattern
+    tensor itself is boolean or integer.
+    """
+    if create_block_mask is None:
+        raise RuntimeError(
+            "'flex_attention' was requested, but torch.create_block_mask is unavailable."
+        )
+    pattern = mask_pattern.detach().to(device=device).contiguous()
+    # One device-to-host transfer is required for an exact cache identity. Use
+    # the contiguous buffer directly instead of materializing one Python int
+    # per byte, which is prohibitively expensive for long batched sequences.
+    host_pattern = pattern.to(device="cpu").contiguous()
+    pattern_bytes = host_pattern.view(torch.uint8).numpy().tobytes(order="C")
+    cache_key = (
+        str(device),
+        None if dtype is None else str(dtype),
+        (batch_size, query_length, key_value_length),
+        str(pattern.dtype),
+        pattern_bytes,
+        mask_semantics,
+    )
+    with _flex_cache_lock:
+        flex_block_mask = _flex_block_masks.get(cache_key)
+        if flex_block_mask is None:
+            flex_block_mask = create_block_mask(
+                mask_mod,
+                batch_size,
+                1,
+                query_length,
+                key_value_length,
+                device=device,
+            )
+            _remember(_flex_block_masks, cache_key, flex_block_mask)
+        else:
+            _flex_block_masks.move_to_end(cache_key)
+    return flex_block_mask
+
+
+# Hugging Face `kernels` exposes slightly different APIs for FlashAttention 2
+# and 3. Detect the loaded variant once so every caller uses the same dispatch.
+def _infer_kernels_flash_variant(kernel) -> str | None:
+    if hasattr(kernel, "fwd") and hasattr(kernel, "varlen_fwd"):
+        return "flash_attn2"
+    if hasattr(kernel, "flash_attn_func") and hasattr(kernel, "flash_attn_varlen_func"):
+        return "flash_attn3"
+    return None
+
+
+def _load_kernels_flash(implementation: str) -> tuple[object, str]:
+    """Load exactly the requested FlashAttention kernel.
+
+    Loading is deferred until backend selection. A FlashAttention-2 request
+    never falls through to FlashAttention-3, or vice versa.
+    """
+    from fastplms.registry import get_model_registry
+
+    kernel_spec = get_model_registry().attention_kernels[implementation]
+    repository = kernel_spec.repository
+    try:
+        flash_kernel = load_locked_kernel(repository, kernel_spec.revision)
+    except Exception as error:
+        raise RuntimeError(
+            f"Unable to load the manifest-pinned kernel "
+            f"{repository}@{kernel_spec.revision} for {implementation!r}."
+        ) from error
+    flash_kernel_variant = _infer_kernels_flash_variant(flash_kernel)
+    if flash_kernel_variant != kernel_spec.expected_variant:
+        raise RuntimeError(
+            f"{repository}@{kernel_spec.revision} exposed {flash_kernel_variant!r}; "
+            f"expected {kernel_spec.expected_variant!r}."
+        )
+    if not all(
+        callable(getattr(flash_kernel, name, None))
+        for name in ("flash_attn_func", "flash_attn_varlen_func")
+    ):
+        raise RuntimeError(
+            f"{repository}@{kernel_spec.revision} does not expose the "
+            "autograd-enabled flash_attn_func and flash_attn_varlen_func APIs."
+        )
+    return flash_kernel, flash_kernel_variant
+
+
+_FLASH_KERNELS: dict[str, tuple[object, str]] = {}
+
+
+def _validate_kernels_flash_dtype(
+    query_states: torch.Tensor,
+    key_states: torch.Tensor,
+    value_states: torch.Tensor,
+    implementation: str,
+) -> torch.dtype:
+    """Reject dtypes outside the immutable kernel manifest before dispatch."""
+
+    tensor_dtypes = {query_states.dtype, key_states.dtype, value_states.dtype}
+    if len(tensor_dtypes) != 1:
+        observed = ", ".join(sorted(str(dtype) for dtype in tensor_dtypes))
+        raise RuntimeError(
+            f"{implementation!r} requires Q, K, and V to share one dtype; received {observed}."
+        )
+    runtime_dtype = query_states.dtype
+    if (
+        runtime_dtype == torch.float32
+        and query_states.is_cuda
+        and torch.is_autocast_enabled("cuda")
+    ):
+        runtime_dtype = torch.get_autocast_dtype("cuda")
+    dtype_names = {
+        torch.float32: "float32",
+        torch.bfloat16: "bfloat16",
+        torch.float16: "float16",
+    }
+    runtime_dtype_name = dtype_names.get(runtime_dtype, str(runtime_dtype))
+    from fastplms.registry import get_model_registry
+
+    supported = get_model_registry().attention_kernels[implementation].dtypes
+    if runtime_dtype_name not in supported:
+        expected = ", ".join(supported)
+        raise RuntimeError(
+            f"{implementation!r} supports only manifest-declared dtype(s) {expected}; "
+            f"received {runtime_dtype_name}. Use CUDA BF16 autocast for FP32-resident "
+            "models."
+        )
+    return runtime_dtype
+
+
+def _validate_kernels_flash_device(
+    query_states: torch.Tensor,
+    key_states: torch.Tensor,
+    value_states: torch.Tensor,
+    implementation: str,
+) -> torch.device:
+    """Require Q, K, and V on one CUDA device before loading a kernel."""
+
+    devices = (query_states.device, key_states.device, value_states.device)
+    if len(set(devices)) != 1:
+        observed = ", ".join(str(device) for device in devices)
+        raise RuntimeError(
+            f"{implementation!r} requires Q, K, and V on one device; received {observed}."
+        )
+    device = devices[0]
+    if device.type != "cuda" or not all(
+        tensor.is_cuda for tensor in (query_states, key_states, value_states)
+    ):
+        raise RuntimeError(
+            f"{implementation!r} requires CUDA Q, K, and V; received device {device}."
+        )
+    return device
+
+
+def _ensure_flash_kernels_loaded(implementation: str) -> tuple[object, str]:
+    cached = _FLASH_KERNELS.get(implementation)
+    if cached is not None:
+        return cached
+    loaded = _load_kernels_flash(implementation)
+    _FLASH_KERNELS[implementation] = loaded
+    return loaded
+
+
+def _kernels_flash_forward(
+    query_states: torch.Tensor,
+    key_states: torch.Tensor,
+    value_states: torch.Tensor,
+    causal: bool = False,
+    softmax_scale: float | None = None,
+    implementation: str = "flash_attention_3",
+) -> torch.Tensor:
+    """Flash-attention forward, optionally overriding the softmax scale.
+
+    When `softmax_scale is None`, the flash kernel applies its default
+    `1 / sqrt(head_dim)`. Pass `softmax_scale=1.0` if the caller has already
+    pre-scaled Q (the convention used by ESM2, DPLM, DPLM2, E1, ESMFold).
+    Failing to override when Q is pre-scaled applies the scale twice and breaks
+    parity with eager attention and SDPA.
+    """
+    flash_kernel, flash_kernel_variant = _ensure_flash_kernels_loaded(implementation)
+    if flash_kernel_variant == "flash_attn2":
+        output = flash_kernel.flash_attn_func(
+            q=query_states,
+            k=key_states,
+            v=value_states,
+            dropout_p=0.0,
+            softmax_scale=softmax_scale,
+            causal=causal,
+        )
+        return output[0] if isinstance(output, tuple) else output
+    if flash_kernel_variant == "flash_attn3":
+        output = flash_kernel.flash_attn_func(
+            q=query_states,
+            k=key_states,
+            v=value_states,
+            softmax_scale=softmax_scale,
+            causal=causal,
+        )
+        if isinstance(output, tuple):
+            return output[0]
+        return output
+    raise RuntimeError(f"Unsupported FlashAttention kernel variant: {flash_kernel_variant}")
+
+
+def _kernels_flash_varlen_forward(
+    query_states: torch.Tensor,
+    key_states: torch.Tensor,
+    value_states: torch.Tensor,
+    cu_seqlens_q: torch.Tensor,
+    cu_seqlens_k: torch.Tensor,
+    max_seqlen_in_batch_q: int,
+    max_seqlen_in_batch_k: int,
+    causal: bool = False,
+    softmax_scale: float | None = None,
+    implementation: str = "flash_attention_3",
+) -> torch.Tensor:
+    """Varlen flash-attention forward, optionally overriding the softmax scale.
+
+    See `_kernels_flash_forward` docstring for why `softmax_scale=1.0` must be
+    passed when Q has been pre-scaled by the caller.
+    """
+    flash_kernel, flash_kernel_variant = _ensure_flash_kernels_loaded(implementation)
+    if flash_kernel_variant == "flash_attn2":
+        output = flash_kernel.flash_attn_varlen_func(
+            q=query_states,
+            k=key_states,
+            v=value_states,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=max_seqlen_in_batch_q,
+            max_seqlen_k=max_seqlen_in_batch_k,
+            dropout_p=0.0,
+            softmax_scale=softmax_scale,
+            causal=causal,
+        )
+        return output[0] if isinstance(output, tuple) else output
+    if flash_kernel_variant == "flash_attn3":
+        output = flash_kernel.flash_attn_varlen_func(
+            q=query_states,
+            k=key_states,
+            v=value_states,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_q=max_seqlen_in_batch_q,
+            max_seqlen_k=max_seqlen_in_batch_k,
+            softmax_scale=softmax_scale,
+            causal=causal,
+        )
+        if isinstance(output, tuple):
+            return output[0]
+        return output
+    raise RuntimeError(f"Unsupported FlashAttention kernel variant: {flash_kernel_variant}")
+
+
+# Varlen flash attention runs only on real tokens. These helpers remove padding
+# before the kernel call and restore the original padded batch shape afterward.
+class IndexFirstAxis(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input, indices) -> torch.Tensor:
+        ctx.save_for_backward(indices)
+        if input.ndim < 2:
+            raise ValueError(
+                "index_first_axis input must have at least two dimensions; "
+                f"received shape {tuple(input.shape)}."
+            )
+        if indices.ndim != 1:
+            raise ValueError(
+                "index_first_axis indices must be one-dimensional; "
+                f"received shape {tuple(indices.shape)}."
+            )
+        ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:]
+        second_dim = other_shape.numel()
+        return torch.gather(
+            rearrange(input, "b ... -> b (...)"), 0, indices.unsqueeze(1).expand(-1, second_dim)
+        ).reshape(-1, *other_shape)
+
+    @staticmethod
+    def backward(ctx, grad_output) -> tuple[torch.Tensor, None]:
+        (indices,) = ctx.saved_tensors
+        if grad_output.ndim < 2:
+            raise RuntimeError(
+                "index_first_axis received an invalid gradient with fewer than "
+                "two dimensions."
+            )
+        other_shape = grad_output.shape[1:]
+        grad_output = rearrange(grad_output, "b ... -> b (...)")
+        grad_input = torch.zeros(
+            [ctx.first_axis_dim, grad_output.shape[1]],
+            device=grad_output.device,
+            dtype=grad_output.dtype,
+        )
+        grad_input.scatter_(0, indices.unsqueeze(1).expand(-1, grad_output.shape[1]), grad_output)
+        return grad_input.reshape(ctx.first_axis_dim, *other_shape), None
+
+
+class IndexPutFirstAxis(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, values, indices, first_axis_dim) -> torch.Tensor:
+        ctx.save_for_backward(indices)
+        if indices.ndim != 1:
+            raise ValueError(
+                "index_put_first_axis indices must be one-dimensional; "
+                f"received shape {tuple(indices.shape)}."
+            )
+        if values.ndim < 2:
+            raise ValueError(
+                "index_put_first_axis values must have at least two dimensions; "
+                f"received shape {tuple(values.shape)}."
+            )
+        output = torch.zeros(
+            first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype
+        )
+        output[indices] = values
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output) -> tuple[torch.Tensor, None, None]:
+        (indices,) = ctx.saved_tensors
+        return grad_output[indices], None, None
+
+
+index_first_axis = IndexFirstAxis.apply
+index_put_first_axis = IndexPutFirstAxis.apply
+
+
+def pad_input(
+    hidden_states: torch.Tensor, indices: torch.Tensor, batch: int, seqlen: int
+) -> torch.Tensor:
+    output = index_put_first_axis(hidden_states, indices, batch * seqlen)
+    return rearrange(output, "(b s) ... -> b s ...", b=batch)
+
+
+def _unpad_input(
+    query_layer: torch.Tensor,
+    key_layer: torch.Tensor,
+    value_layer: torch.Tensor,
+    attention_mask_2d: torch.Tensor,
+) -> tuple[
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    tuple[torch.Tensor, torch.Tensor],
+    tuple[int, int],
+]:
+    batch_size, seq_len, num_heads, head_dim = query_layer.shape
+    seqlens = attention_mask_2d.sum(dim=1).int()
+    cu_seqlens = F.pad(seqlens.cumsum(0, dtype=torch.int32), (1, 0))
+    max_seqlen = int(seqlens.max().item())
+    indices = attention_mask_2d.flatten().nonzero(as_tuple=False).flatten()
+    query_layer = index_first_axis(
+        query_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices
+    )
+    key_layer = index_first_axis(
+        key_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices
+    )
+    value_layer = index_first_axis(
+        value_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices
+    )
+    return (
+        query_layer,
+        key_layer,
+        value_layer,
+        indices,
+        (cu_seqlens, cu_seqlens),
+        (max_seqlen, max_seqlen),
+    )
+
+
+def _validate_flash_padding_mask(
+    query_states: torch.Tensor,
+    key_states: torch.Tensor,
+    value_states: torch.Tensor,
+    attention_mask_2d: torch.Tensor,
+) -> torch.Tensor:
+    """Validate the self-attention padding mask used by the varlen kernels."""
+
+    if attention_mask_2d.ndim != 2:
+        raise ValueError("FlashAttention padding masks must have shape (batch, sequence_length).")
+    expected_shape = query_states.shape[:2]
+    if tuple(attention_mask_2d.shape) != tuple(expected_shape):
+        raise ValueError(
+            "FlashAttention padding mask shape must match the query batch and "
+            f"sequence dimensions; expected {tuple(expected_shape)}, received "
+            f"{tuple(attention_mask_2d.shape)}."
+        )
+    if key_states.shape[:2] != expected_shape or value_states.shape[:2] != expected_shape:
+        raise ValueError(
+            "Masked FlashAttention requires Q, K, and V to share batch and sequence dimensions."
+        )
+    if attention_mask_2d.device != query_states.device:
+        raise ValueError("FlashAttention padding mask and Q, K, and V must be on the same device.")
+    return attention_mask_2d.to(dtype=torch.bool)
+
+
+def kernels_flash_attention_func(
+    query_states: torch.Tensor,
+    key_states: torch.Tensor,
+    value_states: torch.Tensor,
+    attention_mask_2d: torch.Tensor | None = None,
+    causal: bool = False,
+    softmax_scale: float | None = None,
+    implementation: str = "flash_attention_3",
+) -> torch.Tensor:
+    """Public flash-attention entry point with optional padding handling.
+
+    `softmax_scale`:
+        None -> kernel applies its default `1 / sqrt(head_dim)`.
+        float -> kernel uses the given scale (pass 1.0 when Q is pre-scaled
+        by the caller).
+
+    Caller contract: if a model family pre-scales Q by `1/sqrt(head_dim)`
+    before calling this function (ESM2, DPLM, DPLM2, E1, and ESMFold do), pass
+    `softmax_scale=1.0`. Otherwise the flash kernel applies its default scale
+    again, yielding an effective `1/head_dim` scale that drifts across layers.
+    """
+    _validate_kernels_flash_device(
+        query_states,
+        key_states,
+        value_states,
+        implementation,
+    )
+    runtime_dtype = _validate_kernels_flash_dtype(
+        query_states,
+        key_states,
+        value_states,
+        implementation,
+    )
+    if query_states.dtype != runtime_dtype:
+        query_states = query_states.to(dtype=runtime_dtype)
+        key_states = key_states.to(dtype=runtime_dtype)
+        value_states = value_states.to(dtype=runtime_dtype)
+    if attention_mask_2d is not None:
+        attention_mask_2d = _validate_flash_padding_mask(
+            query_states,
+            key_states,
+            value_states,
+            attention_mask_2d,
+        )
+    _ensure_flash_kernels_loaded(implementation)
+    if attention_mask_2d is not None:
+        batch_size, q_len = query_states.shape[:2]
+        (
+            query_states,
+            key_states,
+            value_states,
+            indices_q,
+            (cu_seqlens_q, cu_seqlens_k),
+            (max_seqlen_q, max_seqlen_k),
+        ) = _unpad_input(query_states, key_states, value_states, attention_mask_2d)
+        attn_output_unpad = _kernels_flash_varlen_forward(
+            query_states=query_states,
+            key_states=key_states,
+            value_states=value_states,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_in_batch_q=max_seqlen_q,
+            max_seqlen_in_batch_k=max_seqlen_k,
+            causal=causal,
+            softmax_scale=softmax_scale,
+            implementation=implementation,
+        )
+        output = pad_input(attn_output_unpad, indices_q, batch_size, q_len)
+        return output.masked_fill(~attention_mask_2d[:, :, None, None], 0)
+    else:
+        return _kernels_flash_forward(
+            query_states=query_states,
+            key_states=key_states,
+            value_states=value_states,
+            causal=causal,
+            softmax_scale=softmax_scale,
+            implementation=implementation,
+        )
+
+
+# User-facing backend strings follow the Transformers attention interface.
+# Keep ``str`` plus ``Enum`` so stringification stays compatible with existing
+# configuration serialization rather than adopting ``StrEnum.__str__``.
+class AttentionBackend(str, Enum):  # noqa: UP042
+    EAGER = "eager"
+    SDPA = "sdpa"
+    FLEX_ATTENTION = "flex_attention"
+    FLASH_ATTENTION_2 = "flash_attention_2"
+    FLASH_ATTENTION_3 = "flash_attention_3"
+
+    # Internal spelling retained to keep attention modules concise. It is an
+    # enum alias, not an accepted public backend string.
+    FLEX = FLEX_ATTENTION
+
+    @property
+    def is_flash(self) -> bool:
+        return self in {
+            AttentionBackend.FLASH_ATTENTION_2,
+            AttentionBackend.FLASH_ATTENTION_3,
+        }
+
+
+VALID_ATTENTION_BACKENDS = tuple(b.value for b in AttentionBackend)
+
+
+def warn_attention_backend_fallback(
+    requested_backend: str | AttentionBackend,
+    *,
+    effective_backend: str | AttentionBackend,
+    reason: str,
+) -> None:
+    """Warn when one forward call cannot honor the configured backend."""
+
+    requested = resolve_attention_backend(requested_backend).value
+    effective = resolve_attention_backend(effective_backend).value
+    if requested == effective:
+        return
+    warnings.warn(
+        f"{reason} The requested {requested!r} attention implementation cannot "
+        f"satisfy this call, so FastPLMs is using {effective!r} attention for this "
+        "call only. This can change performance and memory use; the configured "
+        "backend remains unchanged for subsequent calls.",
+        RuntimeWarning,
+        stacklevel=3,
+    )
+
+
+def resolve_attention_backend_for_call(
+    requested_backend: str | AttentionBackend,
+    *,
+    output_attentions: bool,
+) -> AttentionBackend:
+    """Resolve the effective backend for one call and report substitutions once."""
+
+    requested = resolve_attention_backend(requested_backend)
+    if not output_attentions or requested == AttentionBackend.EAGER:
+        return requested
+    warn_attention_backend_fallback(
+        requested,
+        effective_backend=AttentionBackend.EAGER,
+        reason=(
+            "output_attentions=True requires the full materialized attention probability "
+            "matrix, which optimized PyTorch attention APIs do not return."
+        ),
+    )
+    return AttentionBackend.EAGER
+
+
+def resolve_attention_backend(
+    requested_backend: str | AttentionBackend | None,
+) -> AttentionBackend:
+    """Validate a backend without silently substituting another implementation."""
+    if requested_backend is None:
+        requested_backend = AttentionBackend.SDPA.value
+    if isinstance(requested_backend, AttentionBackend):
+        resolved = requested_backend
+    else:
+        try:
+            resolved = AttentionBackend(requested_backend)
+        except ValueError as error:
+            raise ValueError(
+                f"Unsupported attention implementation {requested_backend!r}; "
+                f"expected one of {VALID_ATTENTION_BACKENDS}."
+            ) from error
+    if resolved == AttentionBackend.FLEX_ATTENTION and flex_attention is None:
+        raise RuntimeError(
+            "'flex_attention' was requested, but this PyTorch build does not provide it."
+        )
+    return resolved
+
+
+def get_attn_implementation(config) -> str:
+    """Read the Transformers attention setting, defaulting to SDPA."""
+    requested = getattr(config, "_attn_implementation", None)
+    if requested is None:
+        requested = getattr(config, "attn_backend", None)
+    return resolve_attention_backend(requested).value
+
+
+def set_config_attn_implementation(config, implementation: str) -> str:
+    """Set both the Transformers field and the internal dispatch field."""
+    resolved = resolve_attention_backend(implementation).value
+    if hasattr(config, "_attn_implementation_internal"):
+        config._attn_implementation_internal = resolved
+    else:
+        config._attn_implementation = resolved
+    # Existing checkpoint configs contain this field. Keeping it synchronized
+    # preserves their state schema while the public API uses attn_implementation.
+    config.attn_backend = resolved
+    return resolved
+
+
+@torch.compiler.disable
+def get_attention_mask(
+    effective_backend: AttentionBackend,
+    batch_size: int,
+    seq_len: int,
+    device: torch.device,
+    attention_mask: torch.Tensor | None = None,
+    dtype: torch.dtype | None = None,
+    mask_semantics: str = "padding",
+) -> tuple[torch.Tensor | None, torch.Tensor | None, BlockMask | None]:
+    """Build padding masks once for all encoder layers.
+
+    Returns (attention_mask_2d, attention_mask_4d, flex_block_mask).
+    """
+    if attention_mask is None:
+        return None, None, None
+
+    if attention_mask.ndim != 2:
+        raise ValueError(
+            "attention_mask must have shape (batch, sequence_length); "
+            f"received rank {attention_mask.ndim} with shape {tuple(attention_mask.shape)}."
+        )
+    expected_shape = (batch_size, seq_len)
+    if tuple(attention_mask.shape) != expected_shape:
+        raise ValueError(
+            "attention_mask shape must match the input batch and sequence dimensions; "
+            f"expected {expected_shape}, received {tuple(attention_mask.shape)}."
+        )
+    attention_mask_2d = attention_mask.to(device=device, dtype=torch.bool)
+    if not bool(attention_mask_2d.any(dim=1).all()):
+        raise ValueError("attention_mask must keep at least one valid key per batch row.")
+
+    effective_backend = resolve_attention_backend(effective_backend)
+
+    if effective_backend.is_flash:
+        return attention_mask_2d, None, None
+
+    if effective_backend == AttentionBackend.FLEX_ATTENTION:
+        if create_block_mask is None:
+            raise RuntimeError(
+                "'flex_attention' was requested, but torch.create_block_mask is unavailable."
+            )
+        def mask_mod(batch_idx, head_idx, q_idx, kv_idx):
+            del head_idx, q_idx
+            # Match eager and SDPA: padding masks suppress invalid keys only.
+            # Invalid queries still attend to real keys and therefore remain
+            # finite; downstream residue masks exclude their outputs.
+            return attention_mask_2d[batch_idx, kv_idx]
+
+        flex_block_mask = _get_flex_block_mask(
+            mask_pattern=attention_mask_2d,
+            batch_size=batch_size,
+            query_length=seq_len,
+            key_value_length=seq_len,
+            device=device,
+            dtype=dtype,
+            mask_semantics=mask_semantics,
+            mask_mod=mask_mod,
+        )
+        return attention_mask_2d, None, flex_block_mask
+
+    # SDPA/manual masks only keys. Padding queries still attend to real keys, so
+    # their outputs stay finite instead of softmaxing over all -inf scores.
+    attention_mask_4d = attention_mask_2d[:, None, None, :]
+    return attention_mask_2d, attention_mask_4d, None
+
+
+def bool_to_additive_mask(
+    bool_mask: torch.Tensor,
+    dtype: torch.dtype,
+) -> torch.Tensor:
+    """Convert a bool mask (True = valid) to a float additive mask (0.0 valid, -inf invalid).
+
+    Why this exists: calling `bool_mask.masked_fill(bool_mask.logical_not(), float('-inf'))`
+    directly on a bool tensor returns a bool tensor because `-inf` casts to `True`.
+    That silently drops the mask. Always allocate a float tensor first, then fill it.
+    This helper is the sanctioned way to build an SDPA additive mask from a bool validity mask.
+    """
+    if bool_mask.dtype != torch.bool:
+        raise TypeError(
+            f"bool_to_additive_mask requires a bool tensor, got dtype={bool_mask.dtype}"
+        )
+    additive = torch.zeros_like(bool_mask, dtype=dtype)
+    additive.masked_fill_(bool_mask.logical_not(), float("-inf"))
+    return additive

fastplms/attention/_kernel_lock.py

@@ -0,0 +1,191 @@
+"""Resolve and validate Hugging Face kernels before importing their binaries."""
+
+from __future__ import annotations
+
+import importlib.metadata
+import json
+import os
+from pathlib import Path
+from typing import Any
+
+
+def require_kernels_package() -> None:
+    """Fail early when the precompiled-kernel runtime is not installed."""
+    try:
+        import kernels  # noqa: F401
+    except ImportError as error:
+        raise RuntimeError(
+            "Precompiled FlashAttention requires the FastPLMs 'flash' extra."
+        ) from error
+
+
+def _kernel_lock_path() -> Path:
+    """Return the lock from an artifact, checkout, or installed distribution."""
+    source_path = Path(__file__).resolve()
+    candidates = [
+        source_path.parents[1] / "kernels.lock",
+        source_path.parents[3] / "kernels.lock",
+    ]
+    try:
+        import fastplms
+
+        candidates.extend(Path(root) / "kernels.lock" for root in fastplms.__path__)
+    except (ImportError, AttributeError):
+        pass
+    for candidate in candidates:
+        if candidate.is_file():
+            return candidate
+
+    try:
+        distribution = importlib.metadata.distribution("fastplms")
+    except importlib.metadata.PackageNotFoundError as error:
+        raise RuntimeError("FastPLMs was installed without kernels.lock.") from error
+    for relative in distribution.files or ():
+        if relative.name != "kernels.lock":
+            continue
+        candidate = Path(distribution.locate_file(relative))
+        if candidate.is_file():
+            return candidate
+    raise RuntimeError("The installed FastPLMs distribution does not contain kernels.lock.")
+
+
+def _locked_entry(lock_path: Path, repository: str) -> dict[str, Any]:
+    try:
+        data = json.loads(lock_path.read_text(encoding="utf-8"))
+    except (OSError, json.JSONDecodeError) as error:
+        raise RuntimeError(f"Unable to read the packaged kernel lock: {lock_path}") from error
+    if not isinstance(data, list):
+        raise RuntimeError("kernels.lock must contain a JSON list.")
+    if any(not isinstance(entry, dict) for entry in data):
+        raise RuntimeError("Every kernels.lock entry must be a JSON object.")
+    matches = [entry for entry in data if entry.get("repo_id") == repository]
+    if len(matches) != 1:
+        raise RuntimeError(
+            f"kernels.lock must contain exactly one entry for {repository!r}; found {len(matches)}."
+        )
+    return matches[0]
+
+
+def _offline_mode() -> bool:
+    """Return whether Hub access was explicitly disabled for this process."""
+
+    enabled_values = {"1", "on", "true", "yes"}
+    return any(
+        os.environ.get(name, "").strip().lower() in enabled_values
+        for name in ("HF_HUB_OFFLINE", "TRANSFORMERS_OFFLINE")
+    )
+
+
+def _offline_snapshot_path(repository: str, revision: str) -> Path:
+    """Locate one exact, possibly sparse, kernel snapshot without using Hub APIs."""
+
+    try:
+        from huggingface_hub import constants
+        from huggingface_hub.file_download import repo_folder_name
+    except ImportError as error:
+        raise RuntimeError("Offline kernel loading requires huggingface-hub.") from error
+
+    cache_root = Path(os.environ.get("KERNELS_CACHE") or constants.HF_HUB_CACHE).resolve()
+    repository_root = (
+        cache_root / repo_folder_name(repo_id=repository, repo_type="kernel")
+    ).resolve()
+    snapshot = repository_root / "snapshots" / revision
+    if not snapshot.is_dir():
+        raise RuntimeError(
+            f"The exact offline kernel snapshot {repository}@{revision} is not cached under "
+            f"{cache_root}. Run `kernels download` before enabling offline mode."
+        )
+    if repository_root not in snapshot.resolve().parents:
+        raise RuntimeError(f"Refusing kernel snapshot outside its cache repository: {snapshot}")
+    return snapshot
+
+
+def _load_offline_locked_kernel(
+    repository: str,
+    revision: str,
+    variant_locks: dict[str, object],
+) -> object:
+    """Validate and import the one compatible variant from a sparse Hub snapshot."""
+    snapshot = _offline_snapshot_path(repository, revision)
+    build_root = snapshot / "build"
+    if not build_root.is_dir():
+        raise RuntimeError(f"The cached kernel snapshot has no build directory: {snapshot}")
+
+    cached_names = sorted(entry.name for entry in build_root.iterdir() if entry.is_dir())
+    unexpected = sorted(set(cached_names).difference(variant_locks))
+    if unexpected:
+        raise RuntimeError(
+            f"The cached {repository}@{revision} snapshot contains unlocked variants: "
+            f"{', '.join(unexpected)}"
+        )
+
+    try:
+        from kernels import get_local_kernel
+        from kernels.utils import validate_kernel
+        from kernels.variants import get_variants_local, resolve_variants
+    except ImportError as error:
+        raise RuntimeError(
+            "Precompiled FlashAttention requires the FastPLMs 'flash' extra."
+        ) from error
+
+    parsed = get_variants_local(build_root)
+    parsed_names = {variant.variant_str for variant in parsed}
+    invalid = sorted(set(cached_names).difference(parsed_names))
+    if invalid:
+        raise RuntimeError(
+            f"The cached {repository}@{revision} snapshot contains invalid variants: "
+            f"{', '.join(invalid)}"
+        )
+
+    compatible, _ = resolve_variants(parsed)
+    if len(compatible) != 1:
+        names = ", ".join(variant.variant_str for variant in compatible) or "none"
+        raise RuntimeError(
+            f"Expected exactly one compatible cached variant for {repository}@{revision}; "
+            f"found {names}."
+        )
+    variant_name = compatible[0].variant_str
+    variant_lock = variant_locks.get(variant_name)
+    expected_hash = getattr(variant_lock, "hash", None)
+    if not isinstance(expected_hash, str) or not expected_hash.startswith("sha256-"):
+        raise RuntimeError(f"The kernel lock for {variant_name} has no valid SHA-256 digest.")
+
+    # Hash validation deliberately happens before import. This operates on the
+    # sparse snapshot produced by `kernels download` and avoids Hub 1.23's
+    # full-snapshot completeness check in offline mode.
+    validate_kernel(repo_path=snapshot, variant=variant_name, hash=expected_hash)
+    return get_local_kernel(build_root / variant_name)
+
+
+def load_locked_kernel(repository: str, revision: str) -> object:
+    """Download, hash-validate, then import one immutable precompiled kernel."""
+    require_kernels_package()
+    try:
+        from kernels import get_local_kernel, install_kernel
+        from kernels.lockfile import KernelLock
+    except ImportError as error:
+        raise RuntimeError(
+            "Precompiled FlashAttention requires the FastPLMs 'flash' extra."
+        ) from error
+
+    lock_path = _kernel_lock_path()
+    kernel_lock = KernelLock.from_json(_locked_entry(lock_path, repository))
+    if kernel_lock.sha != revision:
+        raise RuntimeError(
+            f"The typed manifest pins {repository}@{revision}, but kernels.lock pins "
+            f"{kernel_lock.sha}."
+        )
+
+    if _offline_mode():
+        return _load_offline_locked_kernel(repository, revision, kernel_lock.variants)
+
+    # `install_kernel` downloads data without importing it and validates the
+    # selected build against the tracked variant hash. Only then is the exact
+    # validated path imported directly. Offline mode uses the sparse-cache
+    # resolver above because Hub 1.23 rejects partial snapshots as incomplete.
+    validated_path = install_kernel(
+        repository,
+        revision=kernel_lock.sha,
+        variant_locks=kernel_lock.variants,
+    )
+    return get_local_kernel(validated_path)

fastplms/attention/interfaces.py

@@ -0,0 +1,242 @@
+"""Transformers-compatible attention selection for FastPLMs models."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from functools import partial
+from typing import Any
+
+import torch
+from transformers import AttentionInterface, AttentionMaskInterface
+
+from ._core import (
+    AttentionBackend,
+    get_attn_implementation,
+    kernels_flash_attention_func,
+    resolve_attention_backend,
+    set_config_attn_implementation,
+)
+from ._kernel_lock import require_kernels_package
+
+
+def _kernels_attention_forward(
+    module: torch.nn.Module,
+    query: torch.Tensor,
+    key: torch.Tensor,
+    value: torch.Tensor,
+    attention_mask: torch.Tensor | None,
+    *,
+    implementation: str,
+    **kwargs: Any,
+) -> tuple[torch.Tensor, None]:
+    """Run one canonical FlashAttention backend through Hugging Face kernels.
+
+    Transformers attention functions receive Q, K, and V with shape
+    (b, h, l, d) and return an output with shape (b, l, h, d). The shared
+    FastPLMs kernel adapter uses the latter layout internally.
+    """
+
+    dropout = float(kwargs.get("dropout", 0.0) or 0.0)
+    if module.training and dropout:
+        raise RuntimeError(
+            "Hugging Face kernels FlashAttention is inference-only when attention dropout "
+            "is nonzero. Use SDPA for this training configuration."
+        )
+    causal = bool(kwargs.get("is_causal", getattr(module, "is_causal", False)))
+    softmax_scale = kwargs.get("scaling")
+    output = kernels_flash_attention_func(
+        query_states=query.transpose(1, 2).contiguous(),
+        key_states=key.transpose(1, 2).contiguous(),
+        value_states=value.transpose(1, 2).contiguous(),
+        attention_mask_2d=attention_mask,
+        causal=causal,
+        softmax_scale=softmax_scale,
+        implementation=implementation,
+    )
+    return output, None
+
+
+# Keep FastPLMs' kernels-only adapters local to this registry instance.
+# ``GeneralInterface.register`` updates Transformers' class-wide mapping, so
+# using it here would replace the canonical FlashAttention handlers for every
+# model in the process, including models unrelated to FastPLMs.
+FASTPLMS_ATTENTION_FUNCTIONS = AttentionInterface()
+FASTPLMS_ATTENTION_MASKS = AttentionMaskInterface()
+FASTPLMS_ATTENTION_FUNCTIONS["flash_attention_2"] = partial(
+    _kernels_attention_forward,
+    implementation="flash_attention_2",
+)
+FASTPLMS_ATTENTION_FUNCTIONS["flash_attention_3"] = partial(
+    _kernels_attention_forward,
+    implementation="flash_attention_3",
+)
+for _flash_name in ("flash_attention_2", "flash_attention_3"):
+    FASTPLMS_ATTENTION_MASKS[_flash_name] = FASTPLMS_ATTENTION_MASKS[_flash_name]
+
+
+class FastPLMsAttentionMixin:
+    """Synchronize Transformers attention selection with custom model layers.
+
+    Model families retain their checkpoint parameter names. Only runtime
+    attributes are updated when ``set_attn_implementation`` is called.
+    """
+
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    # Transformers 5.13 uses the singular flag during model construction. A
+    # family opts in only when its manifest entry advertises at least one of
+    # the two FastPLMs kernels-only FlashAttention implementations.
+    _supports_flash_attn = False
+    _supports_flash_attn_2 = False
+    _supports_flash_attn_3 = False
+    _fastplms_attention_implementations = (
+        "eager",
+        "sdpa",
+        "flex_attention",
+    )
+
+    def _validate_attention_name(self, implementation: str) -> None:
+        if implementation not in self._fastplms_attention_implementations:
+            raise ValueError(
+                f"{type(self).__name__} does not support {implementation!r}; expected one of "
+                f"{self._fastplms_attention_implementations}."
+            )
+
+    def _check_and_adjust_attn_implementation(
+        self,
+        attn_implementation: str | None,
+        is_init_check: bool = False,
+        allow_all_kernels: bool = False,
+    ) -> str:
+        """Resolve attention without invoking Transformers' source-Flash probe.
+
+        The standard ``flash_attention_2`` and ``flash_attention_3`` names are
+        retained for the Transformers API, but FastPLMs resolves them only
+        through the exact Hugging Face ``kernels`` artifacts pinned by
+        ``models.toml``. Repository-qualified or otherwise external kernels
+        are never accepted through this model hook.
+        """
+
+        if allow_all_kernels:
+            raise ValueError("FastPLMs does not load external attention kernels.")
+        if attn_implementation is None:
+            return super()._check_and_adjust_attn_implementation(
+                None,
+                is_init_check=is_init_check,
+                allow_all_kernels=False,
+            )
+
+        self._validate_attention_name(attn_implementation)
+        if attn_implementation in {"flash_attention_2", "flash_attention_3"}:
+            if not self._supports_flash_attn:
+                raise ValueError(
+                    f"{type(self).__name__} does not advertise kernels-only FlashAttention."
+                )
+            # Validate the lightweight Python dependency here, but defer binary
+            # download and import until Q, K, and V have passed the CUDA gate.
+            require_kernels_package()
+            return attn_implementation
+
+        return super()._check_and_adjust_attn_implementation(
+            attn_implementation,
+            is_init_check=is_init_check,
+            allow_all_kernels=False,
+        )
+
+    def __init__(self, config, *args: Any, **kwargs: Any) -> None:
+        sentinel = object()
+        internal = getattr(config, "_attn_implementation_internal", sentinel)
+        canonical = (
+            getattr(config, "_attn_implementation", None) if internal is sentinel else internal
+        )
+        legacy = getattr(config, "attn_backend", None)
+        requested = canonical if canonical is not None else legacy
+        if requested is not None:
+            if not isinstance(requested, str):
+                raise TypeError(
+                    "The configured attention implementation must be a string or None; "
+                    f"received {type(requested).__name__}."
+                )
+            self._validate_attention_name(requested)
+            # ``PreTrainedModel.__init__`` resolves a missing Transformers
+            # implementation to the family default.  Legacy FastPLMs configs
+            # persist their explicit choice in ``attn_backend``, so forward it
+            # into the canonical Transformers field before the base class can
+            # replace it with SDPA.  A non-None canonical value still wins,
+            # including an explicit ``attn_implementation=...`` load override.
+            if canonical is None and legacy is not None:
+                set_config_attn_implementation(config, legacy)
+        super().__init__(config, *args, **kwargs)
+        # Transformers resolves an unspecified implementation during the base
+        # model initialization. Synchronize that choice before family layers
+        # are constructed.
+        resolved = get_attn_implementation(config)
+        self._validate_attention_name(resolved)
+        set_config_attn_implementation(config, resolved)
+
+    def set_attn_implementation(
+        self,
+        attn_implementation: str | Mapping[str, str],
+        allow_all_kernels: bool = False,
+    ) -> None:
+        """Select an advertised backend and update every instantiated layer."""
+        if isinstance(attn_implementation, Mapping):
+            if set(attn_implementation) == {""}:
+                attn_implementation = attn_implementation[""]
+            else:
+                raise ValueError(
+                    "FastPLMs models have one attention backbone; pass a string or {'': name}."
+                )
+        resolved_name = self._check_and_adjust_attn_implementation(
+            attn_implementation,
+            is_init_check=False,
+            allow_all_kernels=allow_all_kernels,
+        )
+        set_config_attn_implementation(self.config, resolved_name)
+        resolved = resolve_attention_backend(resolved_name)
+        for module in self.modules():
+            if module is self:
+                continue
+            for attribute in ("attn_backend", "attention_backend", "_attn_backend"):
+                if attribute not in module.__dict__:
+                    continue
+                current = module.__dict__[attribute]
+                module.__dict__[attribute] = (
+                    resolved if isinstance(current, AttentionBackend) else resolved_name
+                )
+
+
+def validate_transformers_attention_interfaces() -> None:
+    """Verify that Transformers exposes functions and masks for every backend.
+
+    Transformers 5.13 registers these canonical names. The FastPLMs function
+    overrides remain instance-local and do not replace process-global handlers.
+    """
+    function_registry = FASTPLMS_ATTENTION_FUNCTIONS
+    mask_registry = FASTPLMS_ATTENTION_MASKS
+    missing_functions = [
+        name
+        for name in (
+            "sdpa",
+            "flex_attention",
+            "flash_attention_2",
+            "flash_attention_3",
+        )
+        if name not in function_registry
+    ]
+    missing_masks = [
+        name
+        for name in (
+            "eager",
+            "sdpa",
+            "flex_attention",
+            "flash_attention_2",
+            "flash_attention_3",
+        )
+        if name not in mask_registry
+    ]
+    if missing_functions or missing_masks:
+        raise RuntimeError(
+            "Transformers attention registry is incomplete: "
+            f"functions={missing_functions}, masks={missing_masks}."
+        )

fastplms/embeddings/__init__.py

@@ -0,0 +1,65 @@
+"""Ordered, residue-aware protein embedding utilities."""
+
+from .pooling import POOLING_NAMES, Pooler, pagerank_weights
+from .runner import (
+    EmbeddingMixin,
+    embed_dataset,
+    iter_fasta,
+    parse_fasta,
+    select_hidden_state_embeddings,
+)
+from .storage import (
+    DEFAULT_SHARD_SIZE,
+    append_sqlite_records,
+    convert_legacy_sqlite,
+    garbage_collect_safetensors_generations,
+    initialize_sqlite_run,
+    load_legacy_pth,
+    load_result,
+    load_safetensors_result,
+    load_sqlite_result,
+    save_result,
+    save_safetensors_result,
+    save_sqlite_result,
+    tensor_sha256,
+    update_sqlite_run_metadata,
+)
+from .types import (
+    EmbeddingBatch,
+    EmbeddingInput,
+    EmbeddingRecord,
+    EmbeddingResult,
+    LazyTensorReference,
+    TensorValue,
+)
+
+__all__ = [
+    "DEFAULT_SHARD_SIZE",
+    "POOLING_NAMES",
+    "EmbeddingBatch",
+    "EmbeddingInput",
+    "EmbeddingMixin",
+    "EmbeddingRecord",
+    "EmbeddingResult",
+    "LazyTensorReference",
+    "Pooler",
+    "TensorValue",
+    "append_sqlite_records",
+    "convert_legacy_sqlite",
+    "embed_dataset",
+    "garbage_collect_safetensors_generations",
+    "initialize_sqlite_run",
+    "iter_fasta",
+    "load_legacy_pth",
+    "load_result",
+    "load_safetensors_result",
+    "load_sqlite_result",
+    "pagerank_weights",
+    "parse_fasta",
+    "save_result",
+    "save_safetensors_result",
+    "save_sqlite_result",
+    "select_hidden_state_embeddings",
+    "tensor_sha256",
+    "update_sqlite_run_metadata",
+]

fastplms/embeddings/pooling.py

@@ -0,0 +1,210 @@
+"""Residue-aware pooling implemented entirely with PyTorch."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Sequence
+
+import torch
+from torch import Tensor
+
+POOLING_NAMES = frozenset({"mean", "max", "norm", "median", "std", "var", "cls", "parti"})
+
+
+def _validate_inputs(X: Tensor, M: Tensor) -> Tensor:
+    if not isinstance(X, Tensor) or not isinstance(M, Tensor):
+        raise TypeError("X and M must be tensors.")
+    if X.ndim != 3:
+        raise ValueError(f"X must have shape (b, l, d), got {tuple(X.shape)}.")
+    if not X.is_floating_point():
+        raise TypeError("X must use a floating-point embedding dtype.")
+    if M.shape != X.shape[:2]:
+        raise ValueError(f"M must have shape (b, l)={tuple(X.shape[:2])}, got {tuple(M.shape)}.")
+    if M.is_complex():
+        raise TypeError("M must be a boolean or binary numeric residue mask.")
+    if not bool(torch.isfinite(M).all()) or not bool(((M == 0) | (M == 1)).all()):
+        raise ValueError("M must contain only finite binary mask values.")
+    M = M.to(device=X.device, dtype=torch.bool)
+    if not bool(M.any(dim=1).all()):
+        raise ValueError("Every sample must contain at least one biological residue.")
+    if not bool((torch.isfinite(X) | ~M.unsqueeze(-1)).all()):
+        raise ValueError("Biological residue embeddings produced non-finite output.")
+    return M
+
+
+def _pooled_attention(attentions: Tensor | Sequence[Tensor], *, batch_size: int) -> Tensor:
+    """Max-pool layer/head attention A to shape ``(b, l, l)``.
+
+    ``parti`` historically keeps the strongest directed edge across the
+    available attention maps before PageRank. Replacing NetworkX with Torch
+    must not change that reduction.
+    """
+
+    if isinstance(attentions, Sequence):
+        if not attentions:
+            raise ValueError("parti received an empty attention sequence.")
+        # Each A_i has shape (b, h, l, l).
+        A = torch.stack(tuple(attentions), dim=1)
+    else:
+        A = attentions
+
+    if A.ndim == 5:
+        if A.shape[0] != batch_size and A.shape[1] == batch_size:
+            A = A.transpose(0, 1)
+        if A.shape[0] != batch_size:
+            raise ValueError("Five-dimensional attentions must use (b, n, h, l, l).")
+        A = A.flatten(1, 2).amax(dim=1)
+    elif A.ndim == 4:
+        if A.shape[0] != batch_size:
+            raise ValueError("Four-dimensional attentions must use (b, h, l, l).")
+        A = A.amax(dim=1)
+    elif A.ndim == 3:
+        if A.shape[0] != batch_size:
+            raise ValueError("Three-dimensional attentions must use (b, l, l).")
+    else:
+        raise ValueError("Attentions must have shape (b, l, l), (b, h, l, l), or (b, n, h, l, l).")
+    return A
+
+
+def pagerank_weights(
+    A: Tensor,
+    *,
+    damping: float = 0.85,
+    tolerance: float = 1e-6,
+    max_iterations: int = 100,
+) -> Tensor:
+    """Compute PageRank weights for a non-negative attention matrix A.
+
+    A has shape ``(l, l)``. Rows are normalized into transition
+    probabilities; dangling rows transition uniformly.
+    """
+
+    if not isinstance(A, Tensor):
+        raise TypeError("A must be a tensor.")
+    if A.ndim != 2 or A.shape[0] != A.shape[1]:
+        raise ValueError(f"A must be square, got shape {tuple(A.shape)}.")
+    if not A.is_floating_point():
+        raise TypeError("A must use a floating-point attention dtype.")
+    if not isinstance(damping, (int, float)) or isinstance(damping, bool):
+        raise TypeError("damping must be a finite float in [0, 1).")
+    if not math.isfinite(float(damping)) or not 0 <= damping < 1:
+        raise ValueError("damping must be a finite float in [0, 1).")
+    if not isinstance(tolerance, (int, float)) or isinstance(tolerance, bool):
+        raise TypeError("tolerance must be a positive finite float.")
+    if not math.isfinite(float(tolerance)) or tolerance <= 0:
+        raise ValueError("tolerance must be a positive finite float.")
+    if not isinstance(max_iterations, int) or isinstance(max_iterations, bool):
+        raise TypeError("max_iterations must be a positive integer.")
+    if max_iterations <= 0:
+        raise ValueError("max_iterations must be a positive integer.")
+    length = A.shape[0]
+    if length == 0:
+        raise ValueError("PageRank requires at least one residue.")
+    if not bool(torch.isfinite(A).all()):
+        raise ValueError("A must contain only finite attention values.")
+    work_dtype = torch.float64 if A.dtype == torch.float64 else torch.float32
+    P = A.detach().to(dtype=work_dtype).clamp_min(0)
+    row_sum = P.sum(dim=-1, keepdim=True)
+    uniform = torch.full_like(P, 1.0 / length)
+    P = torch.where(row_sum > 0, P / row_sum.clamp_min(torch.finfo(work_dtype).tiny), uniform)
+    p = torch.full((length,), 1.0 / length, device=P.device, dtype=work_dtype)
+    teleport = (1.0 - damping) / length
+    for _ in range(max_iterations):
+        p_next = teleport + damping * (P.transpose(0, 1) @ p)
+        if torch.linalg.vector_norm(p_next - p, ord=1) <= tolerance:
+            p = p_next
+            break
+        p = p_next
+    return p / p.sum()
+
+
+class Pooler:
+    """Apply one or more pooling operations to biological residue rows."""
+
+    def __init__(self, pooling: str | Sequence[str] = ("mean",)) -> None:
+        pooling_value: object = pooling
+        if isinstance(pooling_value, (bytes, bytearray)) or not isinstance(
+            pooling_value, (str, Sequence)
+        ):
+            raise TypeError("pooling must be a name or a sequence of names.")
+        names = (pooling_value,) if isinstance(pooling_value, str) else tuple(pooling_value)
+        if not all(isinstance(name, str) for name in names):
+            raise TypeError("pooling names must be strings.")
+        if not names:
+            raise ValueError("At least one pooling operation is required.")
+        unknown = set(names) - POOLING_NAMES
+        if unknown:
+            raise ValueError(f"Unknown pooling operations: {sorted(unknown)}.")
+        duplicates = sorted({name for name in names if names.count(name) > 1})
+        if duplicates:
+            raise ValueError(f"Duplicate pooling operations are not supported: {duplicates}.")
+        self.names = names
+
+    def output_slices(self, d: int) -> dict[str, tuple[int, int]]:
+        """Return the output interval assigned to each pooler."""
+
+        if not isinstance(d, int) or isinstance(d, bool):
+            raise TypeError("d must be a positive integer.")
+        if d <= 0:
+            raise ValueError("d must be a positive integer.")
+        return {name: (i * d, (i + 1) * d) for i, name in enumerate(self.names)}
+
+    def __call__(
+        self,
+        X: Tensor,
+        residue_mask: Tensor,
+        *,
+        attentions: Tensor | Sequence[Tensor] | None = None,
+        attention_backend: str | None = None,
+    ) -> Tensor:
+        M = _validate_inputs(X, residue_mask)
+        M_expanded = M.unsqueeze(-1)
+        count = M_expanded.sum(dim=1).clamp_min(1)
+        X_residues = X.masked_fill(~M_expanded, 0)
+        outputs: list[Tensor] = []
+
+        for name in self.names:
+            if name == "mean":
+                Y = X_residues.sum(dim=1) / count
+            elif name == "max":
+                Y = X.masked_fill(~M_expanded, -torch.inf).max(dim=1).values
+            elif name == "norm":
+                Y = torch.linalg.vector_norm(X_residues, ord=2, dim=1)
+            elif name == "median":
+                Y = X.masked_fill(~M_expanded, torch.nan).nanmedian(dim=1).values
+            elif name in {"var", "std"}:
+                mean = X_residues.sum(dim=1, keepdim=True) / count.unsqueeze(1)
+                centered = (X - mean).masked_fill(~M_expanded, 0)
+                variance = (centered**2).sum(dim=1) / count
+                Y = variance.sqrt() if name == "std" else variance
+            elif name == "cls":
+                Y = X[:, 0]
+            else:
+                if attention_backend != "eager":
+                    raise ValueError(
+                        "parti requires attn_implementation='eager' so full "
+                        "attention matrices are available."
+                    )
+                if attentions is None:
+                    raise ValueError("parti requires model attention matrices.")
+                if int(M.sum(dim=1).max().item()) > 2048:
+                    raise ValueError("parti supports at most 2,048 biological residues.")
+                A = _pooled_attention(attentions, batch_size=X.shape[0]).to(X.device)
+                pooled: list[Tensor] = []
+                for X_i, M_i, A_i in zip(X, M, A, strict=True):
+                    indices = M_i.nonzero(as_tuple=True)[0]
+                    A_residue = A_i.index_select(0, indices).index_select(1, indices)
+                    w = pagerank_weights(A_residue).to(dtype=X.dtype)
+                    pooled.append(w @ X_i.index_select(0, indices))
+                Y = torch.stack(pooled)
+            if not bool(torch.isfinite(Y).all()):
+                raise ValueError(
+                    f"Pooling operation {name!r} produced non-finite output from "
+                    "biological residue embeddings."
+                )
+            outputs.append(Y)
+
+        return torch.cat(outputs, dim=-1)
+
+
+__all__ = ["POOLING_NAMES", "Pooler", "pagerank_weights"]

fastplms/embeddings/runner.py

@@ -0,0 +1,1559 @@
+"""Model-independent dataset embedding orchestration."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import platform
+import sqlite3
+import tempfile
+from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Any, overload
+
+import torch
+from torch import Tensor
+
+from .pooling import Pooler
+from .storage import (
+    SafetensorsStreamWriter,
+    append_sqlite_records,
+    initialize_sqlite_run,
+    load_result,
+    load_sqlite_result,
+    safetensors_result_exists,
+    save_result,
+    tensor_sha256,
+    update_sqlite_run_metadata,
+)
+from .types import (
+    EmbeddingBatch,
+    EmbeddingInput,
+    EmbeddingRecord,
+    EmbeddingResult,
+    LazyTensorReference,
+)
+
+_MAX_PARTI_RESIDUES = 2_048
+_RUN_FINGERPRINT_SCHEMA_VERSION = 3
+_MODEL_STATE_HASH_CHUNK_BYTES = 16 * 1024**2
+_DEFAULT_BATCH_WINDOW_MULTIPLIER = 16
+_SUPPORTED_STORAGE_FORMATS = frozenset({"safetensors", "sqlite"})
+
+
+def _validate_parti_length(M: Tensor) -> None:
+    """Reject an oversized attention graph before model inference."""
+
+    n_residues = int(M.to(dtype=torch.int64).sum(dim=1).max().item())
+    if n_residues > _MAX_PARTI_RESIDUES:
+        raise ValueError(f"parti supports at most {_MAX_PARTI_RESIDUES:,} biological residues.")
+
+
+def select_hidden_state_embeddings(
+    last_hidden_state: Tensor,
+    hidden_states: tuple[Tensor, ...] | None,
+    *,
+    hidden_state_index: int = -1,
+    store_all_hidden_states: bool = False,
+) -> Tensor:
+    """Select one hidden state or stack every state without changing values."""
+    if store_all_hidden_states:
+        if not hidden_states:
+            raise ValueError("store_all_hidden_states requires model hidden states.")
+        # H has shape (b, n, l, d), where n follows the model's output order.
+        return torch.stack(hidden_states, dim=1)
+    if hidden_state_index == -1:
+        return last_hidden_state
+    if not hidden_states:
+        raise ValueError("hidden_state_index requires model hidden states.")
+    return hidden_states[hidden_state_index]
+
+
+def iter_fasta(path: str | Path) -> Iterator[EmbeddingInput]:
+    """Yield FASTA records in source order without reading the file into memory."""
+
+    identifier: str | None = None
+    sequence_parts: list[str] = []
+    found_record = False
+    with Path(path).open("r", encoding="utf-8") as handle:
+        for line_number, raw_line in enumerate(handle, start=1):
+            line = raw_line.strip()
+            if not line:
+                continue
+            if line.startswith(">"):
+                if identifier is not None:
+                    found_record = True
+                    yield EmbeddingInput(identifier, "".join(sequence_parts))
+                identifier = line[1:].strip().split(maxsplit=1)[0]
+                if not identifier:
+                    raise ValueError(f"Missing FASTA identifier on line {line_number}.")
+                sequence_parts = []
+            else:
+                if identifier is None:
+                    raise ValueError(
+                        f"Sequence data precedes the first FASTA header on line {line_number}."
+                    )
+                sequence_parts.append("".join(line.split()))
+    if identifier is not None:
+        found_record = True
+        yield EmbeddingInput(identifier, "".join(sequence_parts))
+    if not found_record:
+        raise ValueError(f"No FASTA records found in {path}.")
+
+
+def parse_fasta(path: str | Path) -> list[EmbeddingInput]:
+    """Parse FASTA records while preserving identifiers, order, and duplicates."""
+
+    return list(iter_fasta(path))
+
+
+def _normalize_input_item(
+    position: int,
+    item: str | EmbeddingInput | tuple[str, str],
+) -> EmbeddingInput:
+    if isinstance(item, EmbeddingInput):
+        return item
+    if isinstance(item, str):
+        return EmbeddingInput(str(position), item)
+    if isinstance(item, tuple) and len(item) == 2:
+        return EmbeddingInput(str(item[0]), str(item[1]))
+    raise TypeError(
+        "inputs must contain sequences, EmbeddingInput values, or (id, sequence) tuples."
+    )
+
+
+class _InputSpool(Sequence[EmbeddingInput]):
+    """Immutable disk-backed normalized inputs with an incremental digest."""
+
+    def __init__(
+        self,
+        values: Iterable[str | EmbeddingInput | tuple[str, str]],
+    ) -> None:
+        self._temporary: tempfile.TemporaryDirectory[str] | None = tempfile.TemporaryDirectory(
+            prefix="fastplms-inputs-"
+        )
+        self.path = Path(self._temporary.name) / "inputs.sqlite"
+        self._connection: sqlite3.Connection | None = sqlite3.connect(self.path)
+        self._connection.execute(
+            "CREATE TABLE inputs ("
+            "position INTEGER PRIMARY KEY, input_id TEXT NOT NULL, sequence TEXT NOT NULL)"
+        )
+        digest = hashlib.sha256()
+        count = 0
+        pending: list[tuple[int, str, str]] = []
+        try:
+            for position, item in enumerate(values):
+                record = _normalize_input_item(position, item)
+                for value in (record.id, record.sequence):
+                    encoded = value.encode("utf-8")
+                    digest.update(len(encoded).to_bytes(8, "big"))
+                    digest.update(encoded)
+                pending.append((position, record.id, record.sequence))
+                count += 1
+                if len(pending) == 1_024:
+                    self._connection.executemany("INSERT INTO inputs VALUES (?, ?, ?)", pending)
+                    pending.clear()
+            if pending:
+                self._connection.executemany("INSERT INTO inputs VALUES (?, ?, ?)", pending)
+            if count == 0:
+                raise ValueError("inputs must contain at least one sequence.")
+            self._connection.commit()
+            self._connection.close()
+            self._connection = sqlite3.connect(
+                f"{self.path.resolve().as_uri()}?mode=ro",
+                uri=True,
+            )
+        except BaseException:
+            self.close()
+            raise
+        digest.update(count.to_bytes(8, "big"))
+        self.input_fingerprint = digest.hexdigest()
+        self._count = count
+
+    def _require_connection(self) -> sqlite3.Connection:
+        if self._connection is None:
+            raise RuntimeError("Input spool is closed.")
+        return self._connection
+
+    def __len__(self) -> int:
+        return self._count
+
+    def __iter__(self) -> Iterator[EmbeddingInput]:
+        cursor = self._require_connection().execute(
+            "SELECT input_id, sequence FROM inputs ORDER BY position"
+        )
+        while rows := cursor.fetchmany(1_024):
+            for input_id, sequence in rows:
+                yield EmbeddingInput(input_id, sequence)
+
+    @overload
+    def __getitem__(self, index: int, /) -> EmbeddingInput: ...
+
+    @overload
+    def __getitem__(self, index: slice, /) -> list[EmbeddingInput]: ...
+
+    def __getitem__(self, index: int | slice) -> EmbeddingInput | list[EmbeddingInput]:
+        connection = self._require_connection()
+
+        if isinstance(index, slice):
+            start, stop, step = index.indices(self._count)
+            if step != 1:
+                return [self[position] for position in range(start, stop, step)]
+            rows = connection.execute(
+                "SELECT input_id, sequence FROM inputs "
+                "WHERE position >= ? AND position < ? ORDER BY position",
+                (start, stop),
+            ).fetchall()
+            return [EmbeddingInput(input_id, sequence) for input_id, sequence in rows]
+        position = index + self._count if index < 0 else index
+        if position < 0 or position >= self._count:
+            raise IndexError(index)
+        row = connection.execute(
+            "SELECT input_id, sequence FROM inputs WHERE position = ?", (position,)
+        ).fetchone()
+        if row is None:
+            raise IndexError(index)
+        return EmbeddingInput(row[0], row[1])
+
+    def close(self) -> None:
+        connection = getattr(self, "_connection", None)
+        if connection is not None:
+            connection.close()
+            self._connection = None
+        temporary = getattr(self, "_temporary", None)
+        if temporary is not None:
+            temporary.cleanup()
+            self._temporary = None
+
+    def __del__(self) -> None:
+        self.close()
+
+
+def _normalize_inputs(
+    inputs: (Iterable[str | EmbeddingInput | tuple[str, str]] | Mapping[str, str] | str | Path),
+    *,
+    disk_backed: bool,
+) -> Sequence[EmbeddingInput]:
+    is_fasta_path = isinstance(inputs, Path)
+    if isinstance(inputs, str):
+        try:
+            is_fasta_path = Path(inputs).is_file()
+        except OSError:
+            is_fasta_path = False
+    should_spool = disk_backed or is_fasta_path or not isinstance(inputs, (str, Sequence, Mapping))
+    values: Iterable[str | EmbeddingInput | tuple[str, str]]
+    if isinstance(inputs, Path):
+        values = iter_fasta(inputs)
+    elif isinstance(inputs, str):
+        values = iter_fasta(inputs) if is_fasta_path else [inputs]
+    elif isinstance(inputs, Mapping):
+        values = inputs.items()
+    else:
+        values = inputs
+    if should_spool:
+        return _InputSpool(values)
+    records: list[EmbeddingInput] = []
+    for position, item in enumerate(values):
+        records.append(_normalize_input_item(position, item))
+    if not records:
+        raise ValueError("inputs must contain at least one sequence.")
+    return records
+
+
+def _validate_untruncated_lengths(
+    records: Sequence[EmbeddingInput],
+    *,
+    max_length: int | None,
+    truncate: bool,
+) -> None:
+    """Fail before inference when a biological-residue limit would be exceeded."""
+
+    if max_length is None or truncate:
+        return
+    for position, record in enumerate(records):
+        residue_count = len(record.sequence)
+        if residue_count > max_length:
+            raise ValueError(
+                f"Input at position {position} with id {record.id!r} has "
+                f"{residue_count} biological residues, exceeding max_length={max_length} "
+                "while truncate=False."
+            )
+
+
+def _model_device(model: Any) -> torch.device:
+    try:
+        return torch.device(next(model.parameters()).device)
+    except (AttributeError, StopIteration):
+        return torch.device("cpu")
+
+
+def _attention_backend(model: Any) -> str | None:
+    config = getattr(model, "config", None)
+    for name in ("_attn_implementation", "attn_implementation", "attn_backend"):
+        value = getattr(config, name, None)
+        if value:
+            return str(value)
+    return None
+
+
+def _attention_kernel_metadata(backend: str | None) -> dict[str, Any] | None:
+    if backend not in {"flash_attention_2", "flash_attention_3"}:
+        return None
+    from fastplms.registry import get_model_registry
+
+    spec = get_model_registry().attention_kernels[backend]
+    return {
+        "repository": spec.repository,
+        "revision": spec.revision,
+        "version": spec.version,
+        "expected_variant": spec.expected_variant,
+        "dtypes": list(spec.dtypes),
+    }
+
+
+def _fingerprint_jsonable(value: Any) -> Any:
+    if isinstance(value, Mapping):
+        return {str(key): _fingerprint_jsonable(item) for key, item in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [_fingerprint_jsonable(item) for item in value]
+    if isinstance(value, (set, frozenset)):
+        return sorted((_fingerprint_jsonable(item) for item in value), key=repr)
+    if isinstance(value, Path):
+        return str(value)
+    if isinstance(value, Tensor):
+        return {
+            "dtype": str(value.dtype).removeprefix("torch."),
+            "shape": list(value.shape),
+            "sha256": tensor_sha256(value),
+        }
+    if isinstance(value, torch.dtype):
+        return str(value).removeprefix("torch.")
+    if isinstance(value, torch.device):
+        return str(value)
+    if value is None or isinstance(value, (str, int, float, bool)):
+        return value
+    return {
+        "class": f"{value.__class__.__module__}.{value.__class__.__qualname__}",
+        "value": str(value),
+    }
+
+
+def _tokenizer_content_sha256(tokenizer: Any) -> str:
+    content: dict[str, Any] = {
+        "init_kwargs": getattr(tokenizer, "init_kwargs", None),
+        "special_tokens_map": getattr(tokenizer, "special_tokens_map", None),
+        "model_max_length": getattr(tokenizer, "model_max_length", None),
+        "padding_side": getattr(tokenizer, "padding_side", None),
+        "truncation_side": getattr(tokenizer, "truncation_side", None),
+    }
+    get_vocab = getattr(tokenizer, "get_vocab", None)
+    if callable(get_vocab):
+        content["vocabulary"] = get_vocab()
+    get_added_vocab = getattr(tokenizer, "get_added_vocab", None)
+    if callable(get_added_vocab):
+        content["added_vocabulary"] = get_added_vocab()
+    backend = getattr(tokenizer, "backend_tokenizer", None)
+    backend_to_str = getattr(backend, "to_str", None)
+    if callable(backend_to_str):
+        content["backend"] = backend_to_str()
+    serialized = json.dumps(
+        _fingerprint_jsonable(content),
+        sort_keys=True,
+        separators=(",", ":"),
+        ensure_ascii=False,
+    ).encode()
+    return hashlib.sha256(serialized).hexdigest()
+
+
+def _tokenizer_metadata(model: Any, tokenizer: Any | None) -> dict[str, Any]:
+    resolved = tokenizer if tokenizer is not None else getattr(model, "tokenizer", None)
+    if resolved is None:
+        # Raw-sequence families such as E1 retain their loader context on the
+        # model/encoder rather than exposing a Transformers tokenizer. Bind the
+        # non-secret source policy to resume identity without serializing a Hub
+        # token or forcing lazy tokenizer initialization.
+        for candidate in (model, getattr(model, "model", None)):
+            settings = getattr(candidate, "__dict__", {}).get("_fastplms_tokenizer_kwargs")
+            if isinstance(settings, Mapping):
+                token_value = settings.get("token")
+                return {
+                    "mode": "native-sequence",
+                    "source": (
+                        str(settings.get("tokenizer_source"))
+                        if settings.get("tokenizer_source") is not None
+                        else None
+                    ),
+                    "revision": settings.get("revision"),
+                    "cache_dir": (
+                        str(settings.get("cache_dir"))
+                        if settings.get("cache_dir") is not None
+                        else None
+                    ),
+                    "local_files_only": bool(settings.get("local_files_only", False)),
+                    "token_policy": (
+                        "disabled"
+                        if token_value is False
+                        else "provided"
+                        if token_value is not None
+                        else "default"
+                    ),
+                }
+        return {"mode": "native-sequence"}
+    return {
+        "mode": "tokenizer",
+        "class": f"{resolved.__class__.__module__}.{resolved.__class__.__qualname__}",
+        "name_or_path": getattr(resolved, "name_or_path", None),
+        "vocab_size": getattr(resolved, "vocab_size", None),
+        "special_token_ids": list(getattr(resolved, "all_special_ids", ())),
+        "content_sha256": _tokenizer_content_sha256(resolved),
+    }
+
+
+@contextmanager
+def _temporary_eval(model: Any) -> Iterator[None]:
+    was_training = getattr(model, "training", None)
+    eval_method = getattr(model, "eval", None)
+    train_method = getattr(model, "train", None)
+    if (
+        not isinstance(was_training, bool)
+        or not callable(eval_method)
+        or not callable(train_method)
+    ):
+        yield
+        return
+    eval_method()
+    try:
+        yield
+    finally:
+        train_method(was_training)
+
+
+def _software_versions() -> dict[str, str | None]:
+    try:
+        import fastplms
+
+        fastplms_version = fastplms.__version__
+    except (AttributeError, ImportError):
+        fastplms_version = None
+    try:
+        import safetensors
+
+        safetensors_version = safetensors.__version__
+    except ImportError:
+        safetensors_version = None
+    try:
+        import transformers
+
+        transformers_version = transformers.__version__
+    except ImportError:
+        transformers_version = None
+    return {
+        "fastplms": fastplms_version,
+        "python": platform.python_version(),
+        "safetensors": safetensors_version,
+        "torch": torch.__version__,
+        "torch_cuda": torch.version.cuda,
+        "transformers": transformers_version,
+    }
+
+
+def _adapter_identity_metadata(model: Any) -> dict[str, Any] | None:
+    """Return deterministic PEFT/adapter identity without tensor payloads."""
+
+    peft_config = getattr(model, "peft_config", None)
+    if not isinstance(peft_config, Mapping) or not peft_config:
+        return None
+    configurations: dict[str, Any] = {}
+    for name, config in sorted(peft_config.items(), key=lambda item: str(item[0])):
+        to_dict = getattr(config, "to_dict", None)
+        if callable(to_dict):
+            value = to_dict()
+        else:
+            try:
+                value = vars(config)
+            except TypeError:
+                value = config
+        configurations[str(name)] = _fingerprint_jsonable(value)
+    active_adapters = getattr(model, "active_adapters", None)
+    if callable(active_adapters):
+        active_adapters = active_adapters()
+    return {
+        "active": _fingerprint_jsonable(active_adapters),
+        "configurations": configurations,
+    }
+
+
+def _execution_identity_metadata(model: Any) -> dict[str, Any]:
+    """Capture runtime policy that can change persisted numerical results."""
+
+    parameter_dtypes = sorted(
+        {
+            str(parameter.dtype).removeprefix("torch.")
+            for parameter in getattr(model, "parameters", lambda: ())()
+        }
+    )
+    return {
+        "device": _model_device(model).type,
+        "hf_device_map": _fingerprint_jsonable(getattr(model, "hf_device_map", None)),
+        "parameter_dtypes": parameter_dtypes,
+        "software": _software_versions(),
+    }
+
+
+def _biological_residue_mask(
+    input_ids: Tensor,
+    attention_mask: Tensor,
+    tokenizer: Any,
+) -> Tensor:
+    """Remove padding and tokenizer-declared special tokens from M."""
+
+    M = attention_mask.to(dtype=torch.bool)
+    special_ids = tuple(int(token_id) for token_id in getattr(tokenizer, "all_special_ids", ()))
+    if special_ids:
+        specials = torch.tensor(special_ids, device=input_ids.device, dtype=input_ids.dtype)
+        M = M & ~torch.isin(input_ids, specials)
+    return M
+
+
+def _generic_embedding_batch(
+    model: Any,
+    sequences: list[str],
+    *,
+    tokenizer: Any | None,
+    max_length: int | None,
+    truncate: bool,
+    need_attentions: bool,
+    model_kwargs: dict[str, Any],
+) -> EmbeddingBatch:
+    config = getattr(model, "config", None)
+    model_type = str(getattr(config, "model_type", "")).lower()
+    if tokenizer is None:
+        tokenizer = getattr(model, "tokenizer", None)
+
+    if tokenizer is None and model_type == "e1":
+        output = model._embed(sequences, return_attention_mask=True, **model_kwargs)
+        if not isinstance(output, tuple) or len(output) != 2:
+            raise TypeError("E1 _embed must return (X, residue_mask).")
+        X, M = output
+        preparer = getattr(model, "prep_tokens", None)
+        if preparer is not None and hasattr(preparer, "get_batch_kwargs"):
+            prepared = preparer.get_batch_kwargs(sequences, device=X.device)
+            input_ids = prepared["input_ids"]
+            boundary_ids = preparer.boundary_token_ids.to(
+                device=input_ids.device, dtype=input_ids.dtype
+            )
+            # E1 wraps each raw sequence in BOS, context-label, terminal-label,
+            # and EOS tokens. Only amino-acid rows are biological residues.
+            M = M.to(dtype=torch.bool) & ~torch.isin(input_ids, boundary_ids)
+        if need_attentions:
+            raise ValueError("parti is not available for tokenizer-free E1 embedding.")
+        return EmbeddingBatch(X=X, residue_mask=M.to(dtype=torch.bool))
+    if tokenizer is None:
+        raise ValueError("A tokenizer is required for this model's embedding path.")
+
+    tokenize_kwargs: dict[str, Any] = {
+        "return_tensors": "pt",
+        "padding": True,
+        "truncation": truncate,
+    }
+    if max_length is not None and truncate:
+        # ``max_length`` is a biological-residue limit. Tokenizer limits include
+        # boundary tokens, so reserve their declared width instead of dropping
+        # residues at the exact boundary.
+        special_token_count = 0
+        num_special_tokens_to_add = getattr(tokenizer, "num_special_tokens_to_add", None)
+        if callable(num_special_tokens_to_add):
+            special_token_count = int(num_special_tokens_to_add(pair=False))
+        tokenize_kwargs["max_length"] = max_length + special_token_count
+    sequence_tokenizer = getattr(model, "_tokenize_sequence_batch", None)
+    if callable(sequence_tokenizer):
+        encoded = sequence_tokenizer(sequences, tokenizer=tokenizer, **tokenize_kwargs)
+    else:
+        encoded = tokenizer(sequences, **tokenize_kwargs)
+    device = _model_device(model)
+    input_ids = encoded["input_ids"].to(device)
+    attention_mask = encoded.get("attention_mask", input_ids.new_ones(input_ids.shape)).to(device)
+    M = _biological_residue_mask(input_ids, attention_mask, tokenizer)
+    if need_attentions:
+        # Validate l before either the backbone or its quadratic attention graph
+        # is materialized. M has shape (b, l).
+        _validate_parti_length(M)
+    X = model._embed(input_ids, attention_mask, **model_kwargs)
+    attentions = None
+    if need_attentions:
+        output = model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_attentions=True,
+            return_dict=True,
+        )
+        attentions = getattr(output, "attentions", None)
+        if attentions is None:
+            raise ValueError("The model did not return attentions required by parti.")
+    return EmbeddingBatch(X=X, residue_mask=M, attentions=attentions)
+
+
+def _first_metadata_value(*values: Any) -> Any:
+    for value in values:
+        if isinstance(value, str):
+            if value.strip():
+                return value
+        elif value is not None:
+            return value
+    return None
+
+
+def _model_identity_metadata(model: Any) -> dict[str, Any]:
+    """Resolve model and checkpoint identity, including local artifact fallbacks."""
+
+    config = getattr(model, "config", None)
+    checkpoint_revision = _first_metadata_value(
+        getattr(config, "fastplms_checkpoint_revision", None),
+        getattr(config, "_commit_hash", None),
+    )
+    return {
+        "model_id": _first_metadata_value(
+            getattr(config, "fastplms_model_id", None),
+            getattr(config, "_name_or_path", None),
+        ),
+        "model_revision": _first_metadata_value(
+            getattr(config, "_commit_hash", None),
+            checkpoint_revision,
+        ),
+        "checkpoint_repo_id": getattr(config, "fastplms_checkpoint_repo_id", None),
+        "checkpoint_revision": checkpoint_revision,
+        "checkpoint_hash": _first_metadata_value(
+            getattr(model, "checkpoint_hash", None),
+            getattr(config, "checkpoint_hash", None),
+            getattr(config, "fastplms_checkpoint_hash", None),
+        ),
+        "weights_revision": getattr(config, "fastplms_weights_revision", None),
+        "runtime_revision": getattr(config, "fastplms_runtime_revision", None),
+        "source_tree_sha256": getattr(config, "fastplms_source_tree_sha256", None),
+        "runtime_bundle_sha256": getattr(config, "fastplms_runtime_bundle_sha256", None),
+    }
+
+
+def _bounded_tensor_chunks(X: Tensor, max_elements: int) -> Iterable[Tensor]:
+    """Yield X in logical row-major order without materializing a full copy."""
+
+    if X.numel() == 0:
+        return
+    if X.ndim == 0:
+        yield X
+        return
+    trailing_elements = 1
+    for size in X.shape[1:]:
+        trailing_elements *= int(size)
+    if trailing_elements <= max_elements:
+        rows_per_chunk = max(1, max_elements // trailing_elements)
+        for start in range(0, X.shape[0], rows_per_chunk):
+            yield X[start : start + rows_per_chunk]
+        return
+    for row in X:
+        yield from _bounded_tensor_chunks(row, max_elements)
+
+
+def _model_state_sha256(model: Any) -> str:
+    """Hash named parameters and persistent buffers using bounded CPU copies."""
+
+    # Never cache this digest from tensor identity or ``Tensor._version``.
+    # ``Parameter.data`` and independent tensor aliases can mutate shared storage
+    # without changing either signal, while persisted resume identity must bind
+    # the authoritative bytes visible at the start of this run.
+    state = model.state_dict(keep_vars=True)
+    digest = hashlib.sha256()
+    for name, value in sorted(state.items()):
+        if not isinstance(value, Tensor):
+            raise TypeError(f"Model state entry {name!r} is not a tensor.")
+        if value.is_meta:
+            raise ValueError(
+                f"Cannot fingerprint meta-device model state entry {name!r}; pass "
+                "model_state_fingerprint with a caller-owned state identity."
+            )
+        header = json.dumps(
+            {
+                "name": name,
+                "dtype": str(value.dtype).removeprefix("torch."),
+                "shape": list(value.shape),
+            },
+            sort_keys=True,
+            separators=(",", ":"),
+        ).encode()
+        digest.update(len(header).to_bytes(8, "big"))
+        digest.update(header)
+        max_elements = max(1, _MODEL_STATE_HASH_CHUNK_BYTES // value.element_size())
+        for chunk in _bounded_tensor_chunks(value.detach(), max_elements):
+            cpu_chunk = chunk.to(device="cpu").contiguous()
+            digest.update(cpu_chunk.reshape(-1).view(torch.uint8).numpy().tobytes())
+    return digest.hexdigest()
+
+
+def _input_sha256(records: Iterable[EmbeddingInput]) -> str:
+    """Hash an ordered input stream without constructing a duplicate JSON payload."""
+
+    precomputed = getattr(records, "input_fingerprint", None)
+    if isinstance(precomputed, str):
+        return precomputed
+    digest = hashlib.sha256()
+    count = 0
+    for record in records:
+        count += 1
+        for value in (record.id, record.sequence):
+            encoded = value.encode("utf-8")
+            digest.update(len(encoded).to_bytes(8, "big"))
+            digest.update(encoded)
+    digest.update(count.to_bytes(8, "big"))
+    return digest.hexdigest()
+
+
+def _run_fingerprint(
+    model: Any,
+    records: Sequence[EmbeddingInput],
+    *,
+    pooling: Sequence[str],
+    full_embeddings: bool,
+    max_length: int | None,
+    truncate: bool,
+    dtype: torch.dtype | None,
+    model_kwargs: dict[str, Any],
+    tokenizer_metadata: dict[str, Any],
+    model_state_fingerprint: str | None,
+    persist_output: bool,
+    embedding_context: Mapping[str, Any],
+    batch_size: int,
+    batch_window_size: int,
+    max_tokens_per_batch: int | None,
+) -> tuple[str, str, str | None, str]:
+    input_fingerprint = _input_sha256(records)
+    attention_backend = _attention_backend(model)
+    model_identity = _model_identity_metadata(model)
+    if model_state_fingerprint is None and persist_output:
+        resolved_model_state_fingerprint = _model_state_sha256(model)
+        model_state_fingerprint_source = "computed"
+    elif model_state_fingerprint is not None:
+        resolved_model_state_fingerprint = model_state_fingerprint.strip()
+        if not resolved_model_state_fingerprint:
+            raise ValueError("model_state_fingerprint must not be empty.")
+        model_state_fingerprint_source = "caller"
+    else:
+        resolved_model_state_fingerprint = None
+        model_state_fingerprint_source = "not-computed"
+    payload = {
+        "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
+        "input_fingerprint": input_fingerprint,
+        "model_state_fingerprint": resolved_model_state_fingerprint,
+        "model_state_fingerprint_source": model_state_fingerprint_source,
+        "model_class": f"{model.__class__.__module__}.{model.__class__.__qualname__}",
+        **model_identity,
+        "attention_backend": attention_backend,
+        "attention_kernel": _attention_kernel_metadata(attention_backend),
+        "layer": repr(
+            getattr(model, "embedding_layer", model_kwargs.get("hidden_state_index", -1))
+        ),
+        "projection": getattr(model, "embedding_projection", None),
+        "esmc_source": getattr(model, "_esmc_source", None),
+        "esmc_revision": getattr(model, "_esmc_source_revision", None),
+        "esmc_files": getattr(model, "_esmc_source_files", None),
+        "token_policy": getattr(model, "embedding_token_policy", None),
+        "tokenizer": tokenizer_metadata,
+        "adapter": _adapter_identity_metadata(model),
+        "execution": _execution_identity_metadata(model),
+        "embedding_context": _fingerprint_jsonable(embedding_context),
+        "pooling": list(pooling),
+        "full_embeddings": full_embeddings,
+        "max_length": max_length,
+        "truncate": truncate,
+        "dtype": str(dtype) if dtype is not None else None,
+        "batching": {
+            "batch_size": batch_size,
+            "batch_window_size": batch_window_size,
+            "max_tokens_per_batch": max_tokens_per_batch,
+            "input_storage": ("disk-spool" if isinstance(records, _InputSpool) else "memory"),
+        },
+        "model_kwargs": {
+            key: _fingerprint_jsonable(value) for key, value in sorted(model_kwargs.items())
+        },
+        "residue_mask_policy": "attention-mask-minus-special-tokens",
+    }
+    run_fingerprint = hashlib.sha256(
+        json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
+    ).hexdigest()
+    return (
+        input_fingerprint,
+        run_fingerprint,
+        resolved_model_state_fingerprint,
+        model_state_fingerprint_source,
+    )
+
+
+def _output_exists(path: str | Path, format: str) -> bool:
+    path = Path(path)
+    if format == "sqlite":
+        return path.is_file()
+    return safetensors_result_exists(path)
+
+
+def _output_descriptor(position: int, record: EmbeddingRecord) -> dict[str, Any]:
+    tensor = record.tensor
+    if isinstance(tensor, LazyTensorReference):
+        dtype = tensor.dtype
+        shape = tensor.shape
+        digest = tensor.sha256
+    else:
+        dtype = str(tensor.dtype).removeprefix("torch.")
+        shape = tuple(tensor.shape)
+        digest = tensor_sha256(tensor)
+    return {
+        "position": position,
+        "id": record.id,
+        "dtype": dtype,
+        "shape": shape,
+        "sha256": digest,
+    }
+
+
+def _ordered_string_sha256(values: Sequence[str]) -> str:
+    digest = hashlib.sha256()
+    for value in values:
+        encoded = value.encode("utf-8")
+        digest.update(len(encoded).to_bytes(8, "big"))
+        digest.update(encoded)
+    digest.update(len(values).to_bytes(8, "big"))
+    return digest.hexdigest()
+
+
+def _embedding_context(
+    model: Any,
+    records: Sequence[EmbeddingInput],
+    *,
+    hidden_state_source: str,
+    decoder_inputs: Sequence[str] | None,
+    decoder_input_ids: Tensor | None,
+    decoder_attention_mask: Tensor | None,
+    model_kwargs: Mapping[str, Any],
+) -> tuple[dict[str, Any], tuple[str, ...] | None]:
+    if hidden_state_source not in {"encoder", "decoder"}:
+        raise ValueError("hidden_state_source must be 'encoder' or 'decoder'.")
+    hidden_state_index = model_kwargs.get("hidden_state_index", -1)
+    if not isinstance(hidden_state_index, int) or isinstance(hidden_state_index, bool):
+        raise TypeError("hidden_state_index must be an integer.")
+    store_all_hidden_states = model_kwargs.get("store_all_hidden_states", False)
+    if not isinstance(store_all_hidden_states, bool):
+        raise TypeError("store_all_hidden_states must be a boolean.")
+    normalized_decoder_inputs: tuple[str, ...] | None = None
+    has_decoder_inputs = decoder_inputs is not None
+    has_decoder_ids = decoder_input_ids is not None
+    if hidden_state_source == "encoder":
+        if has_decoder_inputs or has_decoder_ids or decoder_attention_mask is not None:
+            raise ValueError("Decoder inputs are only valid when hidden_state_source='decoder'.")
+    else:
+        if has_decoder_inputs == has_decoder_ids:
+            raise ValueError(
+                "Decoder embedding requires exactly one of decoder_inputs or decoder_input_ids."
+            )
+    decoder_input_fingerprint: str | None = None
+    if decoder_inputs is not None:
+        if isinstance(decoder_inputs, (str, bytes)) or not isinstance(decoder_inputs, Sequence):
+            raise TypeError("decoder_inputs must be an aligned sequence of strings.")
+        normalized_decoder_inputs = tuple(decoder_inputs)
+        if not all(isinstance(value, str) and value for value in normalized_decoder_inputs):
+            raise ValueError("decoder_inputs must contain non-empty strings.")
+        if len(normalized_decoder_inputs) != len(records):
+            raise ValueError("decoder_inputs must align one-to-one with embedding inputs.")
+        decoder_input_fingerprint = _ordered_string_sha256(normalized_decoder_inputs)
+        if decoder_attention_mask is not None:
+            raise ValueError("decoder_attention_mask requires decoder_input_ids.")
+    if decoder_input_ids is not None:
+        if not isinstance(decoder_input_ids, Tensor) or decoder_input_ids.ndim != 2:
+            raise ValueError("decoder_input_ids must have shape (batch, sequence).")
+        if decoder_input_ids.shape[0] != len(records):
+            raise ValueError("decoder_input_ids must align one-to-one with embedding inputs.")
+        if decoder_input_ids.dtype == torch.bool or decoder_input_ids.is_floating_point():
+            raise TypeError("decoder_input_ids must use an integer token dtype.")
+        decoder_input_fingerprint = tensor_sha256(decoder_input_ids)
+    decoder_mask_fingerprint: str | None = None
+    if decoder_attention_mask is not None:
+        if not isinstance(decoder_attention_mask, Tensor):
+            raise TypeError("decoder_attention_mask must be a tensor.")
+        if decoder_input_ids is None or decoder_attention_mask.shape != decoder_input_ids.shape:
+            raise ValueError("decoder_attention_mask must match decoder_input_ids shape.")
+        decoder_mask_fingerprint = tensor_sha256(decoder_attention_mask)
+
+    context: dict[str, Any] = {
+        "hidden_state_source": hidden_state_source,
+        "hidden_state_index": hidden_state_index,
+        "store_all_hidden_states": store_all_hidden_states,
+        "decoder_input_fingerprint": decoder_input_fingerprint,
+        "decoder_attention_mask_fingerprint": decoder_mask_fingerprint,
+        "decoder_alignment": "input-position" if hidden_state_source == "decoder" else None,
+    }
+    metadata_hook = getattr(model, "_embedding_metadata", None)
+    model_metadata: Mapping[str, Any] | None = None
+    if callable(metadata_hook):
+        model_metadata = metadata_hook(**context)
+        if not isinstance(model_metadata, Mapping):
+            raise TypeError("_embedding_metadata must return a mapping.")
+        context["model_embedding"] = _fingerprint_jsonable(model_metadata)
+    if hidden_state_source == "decoder":
+        has_decoder_batch = callable(getattr(model, "_embedding_batch", None))
+        declares_decoder_stack = (
+            model_metadata is not None and model_metadata.get("hidden_state_stack") == "decoder"
+        )
+        if not has_decoder_batch or not declares_decoder_stack:
+            raise ValueError(
+                f"{model.__class__.__name__} does not declare decoder embedding support."
+            )
+    return context, normalized_decoder_inputs
+
+
+def _planned_batches(
+    records: Sequence[EmbeddingInput],
+    positions: range,
+    *,
+    batch_size: int,
+    max_tokens_per_batch: int | None,
+    max_length: int | None,
+    truncate: bool,
+) -> Iterator[list[int]]:
+    """Length-bucket one bounded window while retaining stable output positions."""
+
+    def effective_length(position: int) -> int:
+        length = len(records[position].sequence)
+        return min(length, max_length) if truncate and max_length is not None else length
+
+    ordered = sorted(positions, key=lambda position: (-effective_length(position), position))
+    batch: list[int] = []
+    longest = 0
+    for position in ordered:
+        length = effective_length(position)
+        if max_tokens_per_batch is not None and length > max_tokens_per_batch:
+            raise ValueError(
+                f"Input at position {position} has {length} residues, exceeding "
+                f"max_tokens_per_batch={max_tokens_per_batch}."
+            )
+        candidate_longest = max(longest, length)
+        exceeds_tokens = (
+            max_tokens_per_batch is not None
+            and candidate_longest * (len(batch) + 1) > max_tokens_per_batch
+        )
+        if batch and (len(batch) >= batch_size or exceeds_tokens):
+            yield batch
+            batch = []
+            longest = 0
+        batch.append(position)
+        longest = max(longest, length)
+    if batch:
+        yield batch
+
+
+def embed_dataset(
+    model: Any,
+    inputs: (Iterable[str | EmbeddingInput | tuple[str, str]] | Mapping[str, str] | str | Path),
+    *,
+    batch_size: int = 2,
+    pooling: str | Sequence[str] | None = None,
+    full_embeddings: bool = False,
+    output: str | Path | None = None,
+    format: str = "safetensors",
+    resume: bool = True,
+    tokenizer: Any | None = None,
+    max_length: int | None = None,
+    truncate: bool = True,
+    dtype: torch.dtype | None = torch.float32,
+    shard_size: int = 2 * 1024**3,
+    model_state_fingerprint: str | None = None,
+    batch_window_size: int | None = None,
+    max_tokens_per_batch: int | None = None,
+    hidden_state_source: str = "encoder",
+    decoder_inputs: Sequence[str] | None = None,
+    decoder_input_ids: Tensor | None = None,
+    decoder_attention_mask: Tensor | None = None,
+    _embedding_batch_fn: Callable[..., EmbeddingBatch] | None = None,
+    _embedding_batch_identity: Mapping[str, Any] | None = None,
+    _allowed_unsupported_pooling: Sequence[str] = (),
+    **model_kwargs: Any,
+) -> EmbeddingResult:
+    """Embed protein sequences with stable ordering and residue-only pooling."""
+
+    for name, value in (
+        ("batch_size", batch_size),
+        ("shard_size", shard_size),
+    ):
+        if not isinstance(value, int) or isinstance(value, bool):
+            raise TypeError(f"{name} must be a positive integer.")
+        if value <= 0:
+            raise ValueError(f"{name} must be a positive integer.")
+    for optional_name, optional_value in (
+        ("max_length", max_length),
+        ("max_tokens_per_batch", max_tokens_per_batch),
+        ("batch_window_size", batch_window_size),
+    ):
+        if optional_value is not None and (
+            not isinstance(optional_value, int) or isinstance(optional_value, bool)
+        ):
+            raise TypeError(f"{optional_name} must be a positive integer when provided.")
+        if optional_value is not None and optional_value <= 0:
+            raise ValueError(f"{optional_name} must be a positive integer when provided.")
+    for name, value in (
+        ("full_embeddings", full_embeddings),
+        ("resume", resume),
+        ("truncate", truncate),
+    ):
+        if not isinstance(value, bool):
+            raise TypeError(f"{name} must be a boolean.")
+    if not isinstance(format, str):
+        raise TypeError("format must be a string.")
+    if output is not None and not isinstance(output, (str, Path)):
+        raise TypeError("output must be a path or None.")
+    if model_state_fingerprint is not None and (
+        not isinstance(model_state_fingerprint, str) or not model_state_fingerprint
+    ):
+        raise ValueError("model_state_fingerprint must be a non-empty string when provided.")
+    if hidden_state_source not in {"encoder", "decoder"}:
+        raise ValueError("hidden_state_source must be 'encoder' or 'decoder'.")
+    hidden_state_index = model_kwargs.get("hidden_state_index", -1)
+    if not isinstance(hidden_state_index, int) or isinstance(hidden_state_index, bool):
+        raise TypeError("hidden_state_index must be an integer.")
+    store_all_hidden_states = model_kwargs.get("store_all_hidden_states", False)
+    if not isinstance(store_all_hidden_states, bool):
+        raise TypeError("store_all_hidden_states must be a boolean.")
+    if decoder_input_ids is not None:
+        if not isinstance(decoder_input_ids, Tensor):
+            raise TypeError("decoder_input_ids must be a tensor.")
+        if decoder_input_ids.is_meta:
+            raise ValueError("decoder_input_ids cannot be a meta tensor.")
+        if decoder_input_ids.ndim != 2 or decoder_input_ids.shape[1] == 0:
+            raise ValueError("decoder_input_ids must have non-empty shape (batch, sequence).")
+        if decoder_input_ids.dtype not in {torch.int32, torch.int64}:
+            raise TypeError("decoder_input_ids must use torch.int32 or torch.int64.")
+    if decoder_attention_mask is not None:
+        if not isinstance(decoder_attention_mask, Tensor):
+            raise TypeError("decoder_attention_mask must be a tensor.")
+        if decoder_attention_mask.is_meta:
+            raise ValueError("decoder_attention_mask cannot be a meta tensor.")
+        if decoder_attention_mask.is_complex() or not bool(
+            torch.isfinite(decoder_attention_mask).all()
+        ):
+            raise ValueError("decoder_attention_mask must contain finite binary values.")
+        if not bool(((decoder_attention_mask == 0) | (decoder_attention_mask == 1)).all()):
+            raise ValueError("decoder_attention_mask must contain finite binary values.")
+    pooling_names = (
+        (("mean",) if not full_embeddings else ())
+        if pooling is None
+        else ((pooling,) if isinstance(pooling, str) else tuple(pooling))
+    )
+    if full_embeddings and pooling is not None:
+        raise ValueError("full_embeddings=True cannot be combined with pooling.")
+    if not full_embeddings and not pooling_names:
+        raise ValueError("pooling is required unless full_embeddings=True.")
+    pooler = Pooler(pooling_names) if pooling_names else None
+
+    if batch_size <= 0:
+        raise ValueError("batch_size must be positive.")
+    if format == "pth" or (output is not None and Path(output).suffix.lower() == ".pth"):
+        raise ValueError("Writing pickle-based .pth embeddings is not supported.")
+    if format not in _SUPPORTED_STORAGE_FORMATS:
+        raise ValueError("format must be 'safetensors' or 'sqlite'.")
+    if max_length is not None and max_length <= 0:
+        raise ValueError("max_length must be positive when provided.")
+    if max_tokens_per_batch is not None and max_tokens_per_batch <= 0:
+        raise ValueError("max_tokens_per_batch must be positive when provided.")
+    if not isinstance(dtype, (torch.dtype, type(None))):
+        raise TypeError("dtype must be a torch.dtype or None.")
+    if batch_window_size is not None and batch_window_size <= 0:
+        raise ValueError("batch_window_size must be positive when provided.")
+    if _embedding_batch_fn is not None and not callable(_embedding_batch_fn):
+        raise TypeError("_embedding_batch_fn must be callable when provided.")
+    if _embedding_batch_fn is not None and _embedding_batch_identity is None:
+        raise ValueError(
+            "_embedding_batch_identity is required with _embedding_batch_fn so persisted "
+            "runs bind the family-specific embedding behavior."
+        )
+    if _embedding_batch_identity is not None and not isinstance(_embedding_batch_identity, Mapping):
+        raise TypeError("_embedding_batch_identity must be a mapping when provided.")
+    if isinstance(_allowed_unsupported_pooling, (str, bytes)) or not isinstance(
+        _allowed_unsupported_pooling, Sequence
+    ):
+        raise TypeError("_allowed_unsupported_pooling must be a sequence of pooler names.")
+    if not all(isinstance(name, str) for name in _allowed_unsupported_pooling):
+        raise TypeError("_allowed_unsupported_pooling must contain only strings.")
+    allowed_unsupported_pooling = frozenset(_allowed_unsupported_pooling)
+    if allowed_unsupported_pooling and _embedding_batch_fn is None:
+        raise ValueError(
+            "_allowed_unsupported_pooling is only valid with a family-specific _embedding_batch_fn."
+        )
+    resolved_batch_window_size = (
+        batch_size * _DEFAULT_BATCH_WINDOW_MULTIPLIER
+        if batch_window_size is None
+        else batch_window_size
+    )
+    if resolved_batch_window_size < batch_size:
+        raise ValueError("batch_window_size must be at least batch_size.")
+    records = _normalize_inputs(inputs, disk_backed=output is not None)
+    _validate_untruncated_lengths(
+        records,
+        max_length=max_length,
+        truncate=truncate,
+    )
+    pooling_names = (
+        (("mean",) if not full_embeddings else ())
+        if pooling is None
+        else ((pooling,) if isinstance(pooling, str) else tuple(pooling))
+    )
+    if full_embeddings:
+        if pooling is not None:
+            raise ValueError("full_embeddings=True cannot be combined with pooling.")
+    elif not pooling_names:
+        raise ValueError("pooling is required unless full_embeddings=True.")
+    store_all_hidden_states = bool(model_kwargs.get("store_all_hidden_states", False))
+    if store_all_hidden_states and not full_embeddings:
+        raise ValueError("store_all_hidden_states=True requires full_embeddings=True.")
+
+    unsupported = set(getattr(model, "embedding_unsupported_pooling", ()))
+    unknown_pooling_overrides = allowed_unsupported_pooling.difference(unsupported)
+    if unknown_pooling_overrides:
+        raise ValueError(
+            "_allowed_unsupported_pooling may only override poolers declared unsupported "
+            f"by the model; unknown overrides: {sorted(unknown_pooling_overrides)}."
+        )
+    unsupported.difference_update(allowed_unsupported_pooling)
+    requested_unsupported = unsupported.intersection(pooling_names)
+    if requested_unsupported:
+        raise ValueError(
+            f"{model.__class__.__name__} does not support pooling operations "
+            f"{sorted(requested_unsupported)}."
+        )
+
+    # Constructing the pooler validates names and duplicate operations before
+    # any checkpoint hashing, tokenization, or inference occurs.
+    pooler = Pooler(pooling_names) if pooling_names else None
+    embedding_context, normalized_decoder_inputs = _embedding_context(
+        model,
+        records,
+        hidden_state_source=hidden_state_source,
+        decoder_inputs=decoder_inputs,
+        decoder_input_ids=decoder_input_ids,
+        decoder_attention_mask=decoder_attention_mask,
+        model_kwargs=model_kwargs,
+    )
+    if _embedding_batch_identity is not None:
+        embedding_context["family_adapter"] = _fingerprint_jsonable(_embedding_batch_identity)
+        if allowed_unsupported_pooling:
+            embedding_context["family_adapter_pooling_override"] = sorted(
+                allowed_unsupported_pooling
+            )
+
+    tokenizer_metadata = _tokenizer_metadata(model, tokenizer)
+    (
+        input_fingerprint,
+        run_fingerprint,
+        resolved_model_state_fingerprint,
+        model_state_fingerprint_source,
+    ) = _run_fingerprint(
+        model,
+        records,
+        pooling=pooling_names,
+        full_embeddings=full_embeddings,
+        max_length=max_length,
+        truncate=truncate,
+        dtype=dtype,
+        model_kwargs=model_kwargs,
+        tokenizer_metadata=tokenizer_metadata,
+        model_state_fingerprint=model_state_fingerprint,
+        persist_output=output is not None,
+        embedding_context=embedding_context,
+        batch_size=batch_size,
+        batch_window_size=resolved_batch_window_size,
+        max_tokens_per_batch=max_tokens_per_batch,
+    )
+    output_already_exists = output is not None and _output_exists(output, format)
+    existing: EmbeddingResult | None = None
+    start_position = 0
+    if output is not None and resume and output_already_exists:
+        if format == "sqlite":
+            try:
+                existing = load_sqlite_result(output, run_id=run_fingerprint)
+            except KeyError:
+                existing = load_result(output, format=format)
+        else:
+            existing = load_result(output, format=format)
+        if existing.metadata.get("fingerprint_schema_version") != (_RUN_FINGERPRINT_SCHEMA_VERSION):
+            raise ValueError(
+                "Existing embeddings use an incompatible run fingerprint schema; "
+                "choose another output or set resume=False."
+            )
+        if existing.metadata.get("run_fingerprint") != run_fingerprint:
+            raise ValueError(
+                "Existing embeddings were produced by a different run fingerprint; "
+                "choose another output or set resume=False."
+            )
+        if len(existing) > len(records):
+            raise ValueError(
+                "Existing embeddings are not an ordered prefix of the requested inputs."
+            )
+        prefix_matches = all(
+            (observed.id, observed.sequence) == (expected.id, expected.sequence)
+            for expected, observed in zip(records, existing, strict=False)
+        )
+        if not prefix_matches:
+            raise ValueError(
+                "Existing embeddings are not an ordered prefix of the requested inputs."
+            )
+        if len(existing) == len(records) and existing.metadata.get("complete", True):
+            return existing
+        start_position = len(existing)
+
+    sqlite_run_id: str | None = None
+    sqlite_replace_on_first_commit = False
+    sqlite_initial_metadata: dict[str, Any] | None = None
+    if output is not None and format == "sqlite":
+        sqlite_initial_metadata = {
+            "format_version": 1,
+            "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
+            "run_fingerprint": run_fingerprint,
+            "input_fingerprint": input_fingerprint,
+            "model_state_fingerprint": resolved_model_state_fingerprint,
+            "model_state_fingerprint_source": model_state_fingerprint_source,
+            "complete": False,
+        }
+        sqlite_run_id = run_fingerprint
+        if not resume and output_already_exists:
+            try:
+                load_sqlite_result(output, run_id=run_fingerprint)
+            except KeyError:
+                pass
+            else:
+                # Keep an exact prior run readable until replacement inference
+                # has produced the first complete commit window.
+                sqlite_replace_on_first_commit = True
+        if not sqlite_replace_on_first_commit:
+            initialize_sqlite_run(
+                output,
+                sqlite_initial_metadata,
+                resume=resume,
+            )
+
+    stream_safetensors = output is not None and format == "safetensors"
+    attention_backend = _attention_backend(model)
+    output_records: list[EmbeddingRecord] = (
+        [] if sqlite_run_id is not None or stream_safetensors else list(existing or ())
+    )
+    output_descriptors: list[dict[str, Any]] | None = [] if output is None else None
+    pool_slices: dict[str, tuple[int, int]] = {}
+    if existing and pooler is not None:
+        pooled_width = existing[0].load_tensor().shape[-1]
+        if pooled_width % len(pooling_names) != 0:
+            raise ValueError("Stored pooled width is inconsistent with pooling metadata.")
+        pool_slices = pooler.output_slices(pooled_width // len(pooling_names))
+
+    safetensors_writer: SafetensorsStreamWriter | None = None
+    if stream_safetensors:
+        if output is None:
+            raise RuntimeError("Safetensors streaming was enabled without an output destination.")
+        transactional_overwrite = output_already_exists and not resume
+        safetensors_writer = SafetensorsStreamWriter(
+            output,
+            {
+                "format_version": 1,
+                "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
+                "run_fingerprint": run_fingerprint,
+                "input_fingerprint": input_fingerprint,
+                "model_state_fingerprint": resolved_model_state_fingerprint,
+                "model_state_fingerprint_source": model_state_fingerprint_source,
+                "complete": False,
+            },
+            shard_size=shard_size,
+            existing=existing or (),
+            reuse_existing=bool(resume and existing is not None),
+            publish_initial=not transactional_overwrite,
+            publish_incremental=not transactional_overwrite,
+        )
+    need_attentions = "parti" in pooling_names
+
+    config = getattr(model, "config", None)
+    model_type = str(getattr(config, "model_type", "")).lower()
+    resolved_tokenizer = tokenizer if tokenizer is not None else getattr(model, "tokenizer", None)
+    with _temporary_eval(model), torch.inference_mode():
+        for window_start in range(start_position, len(records), resolved_batch_window_size):
+            window_stop = min(window_start + resolved_batch_window_size, len(records))
+            window_records = records[window_start:window_stop]
+            if not isinstance(window_records, Sequence):
+                raise RuntimeError("The immutable embedding spool returned a non-sequence window.")
+            window_results: dict[int, EmbeddingRecord] = {}
+            for local_positions in _planned_batches(
+                window_records,
+                range(len(window_records)),
+                batch_size=batch_size,
+                max_tokens_per_batch=max_tokens_per_batch,
+                max_length=max_length,
+                truncate=truncate,
+            ):
+                batch_positions = [window_start + position for position in local_positions]
+                batch_records = [window_records[position] for position in local_positions]
+                sequences = [
+                    record.sequence[:max_length]
+                    if truncate and max_length is not None
+                    else record.sequence
+                    for record in batch_records
+                ]
+                batch_model_kwargs = dict(model_kwargs)
+                if model_type == "fast_ankh" or hidden_state_source == "decoder":
+                    batch_model_kwargs["hidden_state_source"] = hidden_state_source
+                    if normalized_decoder_inputs is not None:
+                        batch_model_kwargs["decoder_inputs"] = [
+                            normalized_decoder_inputs[position] for position in batch_positions
+                        ]
+                    if decoder_input_ids is not None:
+                        indices = torch.tensor(
+                            batch_positions,
+                            device=decoder_input_ids.device,
+                            dtype=torch.long,
+                        )
+                        batch_model_kwargs["decoder_input_ids"] = decoder_input_ids.index_select(
+                            0, indices
+                        )
+                    if decoder_attention_mask is not None:
+                        indices = torch.tensor(
+                            batch_positions,
+                            device=decoder_attention_mask.device,
+                            dtype=torch.long,
+                        )
+                        batch_model_kwargs["decoder_attention_mask"] = (
+                            decoder_attention_mask.index_select(0, indices)
+                        )
+                custom_batch = _embedding_batch_fn or getattr(model, "_embedding_batch", None)
+                if custom_batch is not None:
+                    if model_type == "fast_ankh":
+                        batch = custom_batch(
+                            sequences,
+                            tokenizer=resolved_tokenizer,
+                            max_length=max_length,
+                            truncate=truncate,
+                            need_attentions=need_attentions,
+                            **batch_model_kwargs,
+                        )
+                    else:
+                        batch = custom_batch(sequences, **batch_model_kwargs)
+                    if not isinstance(batch, EmbeddingBatch):
+                        raise TypeError("_embedding_batch must return EmbeddingBatch.")
+                else:
+                    batch = _generic_embedding_batch(
+                        model,
+                        sequences,
+                        tokenizer=tokenizer,
+                        max_length=max_length,
+                        truncate=truncate,
+                        need_attentions=need_attentions,
+                        model_kwargs=batch_model_kwargs,
+                    )
+                X = batch.X
+                raw_mask = batch.residue_mask
+                if not isinstance(X, Tensor) or not isinstance(raw_mask, Tensor):
+                    raise TypeError("Embedding batches must provide Tensor X and residue_mask.")
+                if X.is_meta or raw_mask.is_meta:
+                    raise ValueError("Embedding batches cannot contain meta tensors.")
+                if not X.is_floating_point():
+                    raise TypeError("Embedding batches must use a floating-point X dtype.")
+                if raw_mask.is_complex() or not bool(torch.isfinite(raw_mask).all()):
+                    raise ValueError("Embedding residue_mask must contain finite binary values.")
+                if not bool(((raw_mask == 0) | (raw_mask == 1)).all()):
+                    raise ValueError("Embedding residue_mask must contain finite binary values.")
+                M = raw_mask.to(device=X.device, dtype=torch.bool)
+                valid_X_shape = (
+                    X.ndim == 3
+                    and X.shape[0] == len(batch_records)
+                    and X.shape[-1] > 0
+                    and M.shape == X.shape[:2]
+                )
+                valid_all_states_shape = (
+                    X.ndim == 4
+                    and store_all_hidden_states
+                    and full_embeddings
+                    and X.shape[0] == len(batch_records)
+                    and X.shape[1] > 0
+                    and X.shape[-1] > 0
+                    and M.shape == (X.shape[0], X.shape[2])
+                )
+                if not (valid_X_shape or valid_all_states_shape):
+                    raise ValueError(
+                        "Embedding batches must provide X with shape (b, l, d), or "
+                        "(b, states, l, d) when storing all hidden states, and "
+                        "residue_mask with shape (b, l)."
+                    )
+                if not bool(M.any(dim=1).all()):
+                    raise ValueError("Every embedding sample must contain a biological residue.")
+                finite_selected = (
+                    torch.isfinite(X) | ~M.unsqueeze(-1)
+                    if X.ndim == 3
+                    else torch.isfinite(X) | ~M[:, None, :, None]
+                )
+                if not bool(finite_selected.all()):
+                    raise ValueError("Biological residue embeddings produced non-finite output.")
+                if need_attentions:
+                    # Validate the biological graph only after mask integrity is established.
+                    _validate_parti_length(M)
+                if dtype is not None:
+                    X = X.to(dtype=dtype)
+
+                if full_embeddings:
+                    if X.ndim == 4:
+                        values = [
+                            X_i[:, M_i, :].detach().cpu() for X_i, M_i in zip(X, M, strict=True)
+                        ]
+                    else:
+                        values = [X_i[M_i].detach().cpu() for X_i, M_i in zip(X, M, strict=True)]
+                else:
+                    if pooler is None:
+                        raise RuntimeError(
+                            "Pooled embedding output was requested without an initialized pooler."
+                        )
+                    Y = pooler(
+                        X,
+                        M,
+                        attentions=batch.attentions,
+                        attention_backend=attention_backend,
+                    )
+                    pool_slices = pooler.output_slices(X.shape[-1])
+                    values = list(Y.detach().cpu().unbind(0))
+                for position, record, value in zip(
+                    batch_positions, batch_records, values, strict=True
+                ):
+                    window_results[position] = EmbeddingRecord(record.id, record.sequence, value)
+
+            new_records = [
+                window_results[position] for position in range(window_start, window_stop)
+            ]
+            if output_descriptors is not None:
+                output_descriptors.extend(
+                    _output_descriptor(window_start + offset, record)
+                    for offset, record in enumerate(new_records)
+                )
+            if output is not None and sqlite_run_id is not None:
+                append_sqlite_records(
+                    output,
+                    sqlite_run_id,
+                    window_start,
+                    new_records,
+                    replace_metadata=(
+                        sqlite_initial_metadata if sqlite_replace_on_first_commit else None
+                    ),
+                )
+                sqlite_replace_on_first_commit = False
+            elif safetensors_writer is not None:
+                safetensors_writer.append(new_records)
+            else:
+                output_records.extend(new_records)
+
+    software_versions = _software_versions()
+    projection = getattr(model, "embedding_projection", None)
+    resolved_layer = getattr(
+        model,
+        "embedding_layer",
+        model_kwargs.get("hidden_state_index", -1),
+    )
+    token_policy = getattr(
+        model,
+        "embedding_token_policy",
+        {
+            "unit": "residue",
+            "include": ["biological residues"],
+            "exclude": [
+                "BOS",
+                "EOS",
+                "padding",
+                "chain delimiters",
+                "non-protein tokens",
+            ],
+        },
+    )
+    model_identity = _model_identity_metadata(model)
+    metadata: dict[str, Any] = {
+        "format_version": 1,
+        "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
+        "run_fingerprint": run_fingerprint,
+        "input_fingerprint": input_fingerprint,
+        "model_state_fingerprint": resolved_model_state_fingerprint,
+        "model_state_fingerprint_source": model_state_fingerprint_source,
+        "model_class": f"{model.__class__.__module__}.{model.__class__.__qualname__}",
+        **model_identity,
+        "dtype": str(dtype).removeprefix("torch.") if dtype is not None else "model",
+        "attention_backend": attention_backend,
+        "attention_kernel": _attention_kernel_metadata(attention_backend),
+        "layer": resolved_layer,
+        "projection": projection,
+        "esmc_source": getattr(model, "_esmc_source", None),
+        "esmc_revision": getattr(model, "_esmc_source_revision", None),
+        "esmc_files": getattr(model, "_esmc_source_files", None),
+        "token_policy": token_policy,
+        "tokenizer": tokenizer_metadata,
+        **embedding_context,
+        "pooling": list(pooling_names),
+        "pool_slices": pool_slices,
+        "full_embeddings": full_embeddings,
+        "max_length": max_length,
+        "truncate": truncate,
+        "truncation": {"enabled": truncate, "max_length": max_length},
+        "batching": {
+            "batch_size": batch_size,
+            "batch_window_size": resolved_batch_window_size,
+            "max_tokens_per_batch": max_tokens_per_batch,
+            "input_storage": ("disk-spool" if isinstance(records, _InputSpool) else "memory"),
+            "ordering": "bounded-length-bucketed-stable-output",
+            "resume_commit_granularity": (
+                "not-applicable"
+                if output is None
+                else "batch-window"
+                if format == "sqlite"
+                else "shard-flush"
+            ),
+        },
+        "residue_mask_policy": "biological-residues-only",
+        "record_count": len(records),
+        "descriptor_index": (
+            "memory-metadata"
+            if output is None
+            else "sqlite-records"
+            if format == "sqlite"
+            else "safetensors-generation-index"
+        ),
+        "storage_format": format if output is not None else "memory",
+        "software": software_versions,
+        "execution": _execution_identity_metadata(model),
+        "adapter": _adapter_identity_metadata(model),
+        "torch_version": software_versions["torch"],
+        "transformers_version": software_versions["transformers"],
+        "complete": True,
+    }
+    if output_descriptors is not None:
+        metadata["outputs"] = output_descriptors
+        metadata["tensor_hashes"] = [item["sha256"] for item in output_descriptors]
+    status = getattr(model, "esmc_precision_status", None)
+    if status is not None:
+        metadata["esmc_precision"] = status.as_dict() if hasattr(status, "as_dict") else status
+    if output is not None and sqlite_run_id is not None:
+        update_sqlite_run_metadata(output, sqlite_run_id, metadata)
+        return load_sqlite_result(output, run_id=sqlite_run_id)
+    if safetensors_writer is not None:
+        return safetensors_writer.publish(complete=True, metadata=metadata)
+    result = EmbeddingResult(output_records, metadata)
+    if output is not None:
+        return save_result(result, output, format=format, shard_size=shard_size)
+    return result
+
+
+class EmbeddingMixin:
+    """Small delegation mixin shared by FastPLMs model classes."""
+
+    def embed_dataset(self, inputs: Any, **kwargs: Any) -> EmbeddingResult:
+        return embed_dataset(self, inputs, **kwargs)
+
+
+__all__ = [
+    "EmbeddingMixin",
+    "embed_dataset",
+    "iter_fasta",
+    "parse_fasta",
+    "select_hidden_state_embeddings",
+]

fastplms/embeddings/storage.py

@@ -0,0 +1,1594 @@
+"""Lossless, reproducible storage for :mod:`fastplms.embeddings`."""
+
+from __future__ import annotations
+
+import hashlib
+import io
+import json
+import sqlite3
+import struct
+from bisect import bisect_right
+from collections.abc import Iterable, Iterator, Sequence
+from pathlib import Path
+from typing import Any, cast, overload
+from uuid import uuid4
+
+import numpy as np
+import torch
+from torch import Tensor
+
+from .types import (
+    EmbeddingRecord,
+    EmbeddingResult,
+    LazyTensorReference,
+)
+
+_DTYPE_NAMES: dict[torch.dtype, str] = {
+    torch.float16: "float16",
+    torch.bfloat16: "bfloat16",
+    torch.float32: "float32",
+    torch.float64: "float64",
+    torch.int64: "int64",
+    torch.int32: "int32",
+    torch.int16: "int16",
+    torch.int8: "int8",
+    torch.uint8: "uint8",
+    torch.bool: "bool",
+}
+_NAME_DTYPES = {name: dtype for dtype, name in _DTYPE_NAMES.items()}
+DEFAULT_SHARD_SIZE = 2 * 1024**3
+_MAX_RECORDS_PER_DESCRIPTOR_SHARD = 1_024
+_TENSOR_HASH_CHUNK_BYTES = 16 * 1024**2
+
+
+def _jsonable(value: Any) -> Any:
+    if isinstance(value, dict):
+        return {str(key): _jsonable(item) for key, item in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [_jsonable(item) for item in value]
+    if isinstance(value, Path):
+        return str(value)
+    if isinstance(value, torch.dtype):
+        return str(value).removeprefix("torch.")
+    if isinstance(value, torch.device):
+        return str(value)
+    if value is None or isinstance(value, (str, int, float, bool)):
+        return value
+    return repr(value)
+
+
+def _persistent_metadata(
+    metadata: dict[str, Any],
+    *,
+    descriptor_index: str,
+    record_count: int | None = None,
+) -> dict[str, Any]:
+    """Remove per-record copies from metadata and identify the authoritative index."""
+
+    cleaned_value = _jsonable(metadata)
+    if not isinstance(cleaned_value, dict):
+        raise TypeError("Embedding metadata must serialize to a JSON object.")
+    cleaned: dict[str, Any] = cleaned_value
+    cleaned.pop("outputs", None)
+    cleaned.pop("tensor_hashes", None)
+    cleaned["descriptor_index"] = descriptor_index
+    if record_count is not None:
+        cleaned["record_count"] = record_count
+    return cleaned
+
+
+def _tensor_bytes(X: Tensor) -> bytes:
+    """Return the exact contiguous byte representation of X."""
+
+    X = X.detach().cpu().contiguous()
+    return X.view(torch.uint8).numpy().tobytes()
+
+
+def _bounded_tensor_chunks(X: Tensor, max_bytes: int) -> Iterator[Tensor]:
+    """Yield row-major CPU chunks without materializing one full byte string."""
+
+    flattened = X.detach().to(device="cpu").reshape(-1)
+    if flattened.numel() == 0:
+        return
+    chunk_elements = max(1, max_bytes // flattened.element_size())
+    for start in range(0, flattened.numel(), chunk_elements):
+        chunk = flattened[start : start + chunk_elements]
+        if chunk.stride(0) != 1:
+            chunk = chunk.clone(memory_format=torch.contiguous_format)
+        yield chunk
+
+
+def _tensor_hash_chunks(X: Tensor) -> Iterator[bytes]:
+    for chunk in _bounded_tensor_chunks(X, _TENSOR_HASH_CHUNK_BYTES):
+        yield chunk.view(torch.uint8).numpy().tobytes()
+
+
+def tensor_sha256(X: Tensor) -> str:
+    """Hash dtype, shape, and exact tensor bytes."""
+
+    if not isinstance(X, Tensor):
+        raise TypeError("X must be a tensor.")
+    if X.dtype not in _DTYPE_NAMES:
+        raise TypeError(f"Unsupported tensor dtype {X.dtype}.")
+    if X.is_meta:
+        raise ValueError("Cannot hash a meta tensor without storage.")
+    if X.layout != torch.strided:
+        raise TypeError("Only strided tensors can be hashed.")
+    digest = hashlib.sha256()
+    digest.update(_DTYPE_NAMES[X.dtype].encode())
+    digest.update(json.dumps(tuple(X.shape)).encode())
+    for chunk in _tensor_hash_chunks(X):
+        digest.update(chunk)
+    return digest.hexdigest()
+
+
+def _encode_tensor(X: Tensor) -> tuple[str, str, bytes]:
+    if X.dtype not in _DTYPE_NAMES:
+        raise TypeError(f"Unsupported tensor dtype {X.dtype}.")
+    shape = json.dumps(tuple(X.shape), separators=(",", ":"))
+    return _DTYPE_NAMES[X.dtype], shape, _tensor_bytes(X)
+
+
+def _decode_tensor(dtype_name: str, shape_json: str, data: bytes) -> Tensor:
+    try:
+        dtype = _NAME_DTYPES[dtype_name]
+    except KeyError as error:
+        raise ValueError(f"Unsupported stored dtype {dtype_name!r}.") from error
+    shape = tuple(json.loads(shape_json))
+    # uint8 is used only as a byte-level carrier, preserving BF16 bits exactly.
+    byte_array = np.frombuffer(data, dtype=np.uint8).copy()
+    X = torch.from_numpy(byte_array).view(dtype)
+    return X.reshape(shape).clone()
+
+
+def _index_path(path: str | Path) -> Path:
+    path = Path(path)
+    if path.suffix == ".json":
+        return path
+    if path.suffix == ".safetensors":
+        return path.with_suffix(".json")
+    return path / "index.json"
+
+
+def _run_manifest_path(path: str | Path) -> Path:
+    path = Path(path)
+    if path.name == "index.json":
+        return path.with_name("run.json")
+    if path.suffix == ".json":
+        return path.with_name(f"{path.stem}.run.json")
+    if path.suffix == ".safetensors":
+        return path.with_suffix(".run.json")
+    return path / "run.json"
+
+
+def _resolve_index_child(root: Path, relative: str, *, label: str) -> Path:
+    relative_path = Path(relative)
+    candidate = (root / relative_path).resolve()
+    if relative_path.is_absolute() or candidate.parent != root.resolve():
+        raise ValueError(f"Safetensors {label} references a file outside its output directory.")
+    return candidate
+
+
+def _canonical_json_bytes(payload: dict[str, Any]) -> bytes:
+    return (json.dumps(payload, indent=2, sort_keys=True) + "\n").encode("utf-8")
+
+
+def _load_authoritative_index(
+    path: str | Path,
+) -> tuple[dict[str, Any], Path, dict[str, Any]]:
+    """Load the index selected by the atomic run-manifest commit record."""
+
+    stable_index_path = _index_path(path)
+    run_manifest_path = _run_manifest_path(path)
+    if not run_manifest_path.is_file():
+        raise ValueError(f"Missing safetensors run manifest: {run_manifest_path}.")
+    run_manifest = json.loads(run_manifest_path.read_text(encoding="utf-8"))
+    if not isinstance(run_manifest, dict):
+        raise ValueError("Safetensors run manifest must contain a JSON object.")
+    if run_manifest.get("format") != "fastplms-embedding-run":
+        raise ValueError(f"Not a FastPLMs embedding run manifest: {run_manifest_path}.")
+    version = run_manifest.get("version")
+    index_reference = run_manifest.get("index")
+    if not isinstance(index_reference, dict):
+        raise ValueError("Safetensors run manifest contains an invalid index reference.")
+    if version == 1:
+        snapshot = run_manifest.get("index_payload")
+        if isinstance(snapshot, dict):
+            payload = snapshot
+            index_bytes = _canonical_json_bytes(payload)
+        elif snapshot is None:
+            index_bytes = stable_index_path.read_bytes()
+            payload = json.loads(index_bytes.decode("utf-8"))
+            if not isinstance(payload, dict):
+                raise ValueError("Safetensors index must contain a JSON object.")
+        else:
+            raise ValueError("Safetensors run manifest contains an invalid index snapshot.")
+        expected = {
+            "file": stable_index_path.name,
+            "sha256": hashlib.sha256(index_bytes).hexdigest(),
+        }
+        index_path = stable_index_path
+    elif version == 2:
+        relative = index_reference.get("file")
+        if not isinstance(relative, str):
+            raise ValueError("Safetensors run manifest index file is invalid.")
+        index_path = _resolve_index_child(stable_index_path.parent, relative, label="run manifest")
+        index_bytes = index_path.read_bytes()
+        payload = json.loads(index_bytes.decode("utf-8"))
+        if not isinstance(payload, dict):
+            raise ValueError("Safetensors generation index must contain a JSON object.")
+        if payload.get("version") != 2:
+            raise ValueError("Safetensors v2 run manifest must reference a v2 generation index.")
+        expected = {
+            "file": relative,
+            "sha256": hashlib.sha256(index_bytes).hexdigest(),
+        }
+    else:
+        raise ValueError(f"Unsupported safetensors run manifest version {version!r}.")
+    if index_reference != expected:
+        raise ValueError("Safetensors run manifest does not match its index.")
+    if payload.get("format") != "fastplms-embedding-safetensors":
+        raise ValueError(f"Not a FastPLMs embedding index: {index_path}.")
+    record_count = payload.get("record_count")
+    if record_count is None:
+        legacy_records = payload.get("records", ())
+        if not isinstance(legacy_records, list):
+            raise ValueError("Safetensors index contains invalid records.")
+        record_count = len(legacy_records)
+    if not isinstance(record_count, int) or isinstance(record_count, bool) or record_count < 0:
+        raise ValueError("Safetensors record count must be a non-negative integer.")
+    if run_manifest.get("record_count") != record_count:
+        raise ValueError("Safetensors run manifest record count does not match its index.")
+    metadata = payload.get("metadata", {})
+    if not isinstance(metadata, dict):
+        raise ValueError("Safetensors index metadata must contain a JSON object.")
+    if metadata.get("record_count", record_count) != record_count:
+        raise ValueError("Safetensors metadata record count does not match its index.")
+    if version == 1 and run_manifest.get("metadata") != payload.get("metadata"):
+        raise ValueError("Safetensors run manifest metadata does not match its index.")
+    return payload, index_path, run_manifest
+
+
+def safetensors_result_exists(path: str | Path) -> bool:
+    """Return whether an authoritative committed safetensors run exists."""
+
+    try:
+        _load_authoritative_index(path)
+    except (OSError, ValueError, json.JSONDecodeError):
+        return False
+    return True
+
+
+def _load_safetensor(path: Path, key: str) -> Tensor:
+    try:
+        from safetensors import safe_open
+    except ImportError as error:
+        raise ImportError("Loading embeddings requires the 'safetensors' package.") from error
+    with safe_open(path, framework="pt", device="cpu") as handle:
+        return cast(Tensor, handle.get_tensor(key))
+
+
+def _safetensors_shard_prefix(path: str | Path) -> str:
+    requested_path = Path(path)
+    if requested_path.suffix in {".json", ".safetensors"}:
+        return f"{requested_path.stem}-embeddings"
+    return "embeddings"
+
+
+def _authoritative_index_payload(path: str | Path) -> dict[str, Any] | None:
+    """Return the last atomically committed generation index when available."""
+
+    try:
+        payload, _, _ = _load_authoritative_index(path)
+    except (OSError, ValueError, json.JSONDecodeError):
+        return None
+    return payload
+
+
+def _referenced_shards(
+    index_path: Path,
+    payload: dict[str, Any] | None = None,
+) -> set[Path]:
+    if payload is None:
+        payload = _authoritative_index_payload(index_path)
+    if payload is None:
+        return set()
+    shards: set[Path] = set()
+    for descriptor_shard in payload.get("descriptor_shards", ()):
+        tensor_file = descriptor_shard.get("tensor_file")
+        if isinstance(tensor_file, str):
+            candidate = _resolve_index_child(
+                index_path.parent, tensor_file, label="descriptor index"
+            )
+            shards.add(candidate)
+    for item in payload.get("records", ()):
+        relative = item.get("tensor", {}).get("file")
+        if not isinstance(relative, str):
+            continue
+        candidate = (index_path.parent / relative).resolve()
+        if candidate.parent == index_path.parent.resolve():
+            shards.add(candidate)
+    return shards
+
+
+def _validate_tensor_descriptor(
+    tensor: dict[str, Any],
+) -> tuple[str, str, tuple[int, ...], str]:
+    key = tensor.get("key")
+    if not isinstance(key, str) or not key:
+        raise ValueError("Safetensors descriptor tensor key is invalid.")
+    dtype = tensor.get("dtype")
+    if not isinstance(dtype, str) or dtype not in _NAME_DTYPES:
+        raise ValueError("Safetensors descriptor tensor dtype is invalid.")
+    raw_shape = tensor.get("shape")
+    if not isinstance(raw_shape, (list, tuple)) or not all(
+        isinstance(dimension, int) and not isinstance(dimension, bool) and dimension >= 0
+        for dimension in raw_shape
+    ):
+        raise ValueError("Safetensors descriptor tensor shape is invalid.")
+    sha256 = tensor.get("sha256")
+    if (
+        not isinstance(sha256, str)
+        or len(sha256) != 64
+        or sha256 != sha256.lower()
+        or any(character not in "0123456789abcdef" for character in sha256)
+    ):
+        raise ValueError("Safetensors descriptor tensor SHA-256 is invalid.")
+    return key, dtype, tuple(raw_shape), sha256
+
+
+def _record_from_safetensors_descriptor(root: Path, item: dict[str, Any]) -> EmbeddingRecord:
+    if not isinstance(item, dict):
+        raise ValueError("Safetensors record descriptor must contain a JSON object.")
+    record_id = item.get("id")
+    sequence = item.get("sequence")
+    if not isinstance(record_id, str) or not record_id:
+        raise ValueError("Safetensors descriptor record ID is invalid.")
+    if not isinstance(sequence, str) or not sequence:
+        raise ValueError("Safetensors descriptor sequence is invalid.")
+    tensor = item.get("tensor")
+    if not isinstance(tensor, dict):
+        raise ValueError("Safetensors descriptor is missing tensor metadata.")
+    relative = tensor.get("file")
+    if not isinstance(relative, str) or not relative:
+        raise ValueError("Safetensors descriptor tensor file is invalid.")
+    key, dtype, shape, sha256 = _validate_tensor_descriptor(tensor)
+    tensor_path = _resolve_index_child(root, relative, label="descriptor")
+    if not tensor_path.is_file():
+        raise ValueError(f"Safetensors tensor shard is missing: {relative}.")
+
+    def load_tensor() -> Tensor:
+        return _load_safetensor(tensor_path, key)
+
+    reference = LazyTensorReference(
+        source=str(tensor_path),
+        key=key,
+        dtype=dtype,
+        shape=shape,
+        sha256=sha256,
+        _loader=load_tensor,
+    )
+    return EmbeddingRecord(record_id, sequence, reference)
+
+
+class _SafetensorsRecordSequence(Sequence[EmbeddingRecord]):
+    """Lazy immutable view over bounded descriptor JSONL shards."""
+
+    _fastplms_immutable_sequence = True
+
+    def __init__(self, root: Path, descriptor_shards: Sequence[dict[str, Any]]) -> None:
+        if not isinstance(descriptor_shards, (list, tuple)):
+            raise ValueError("Safetensors generation index has invalid descriptor shards.")
+        self.root = root
+        self.shards = tuple(descriptor_shards)
+        cumulative: list[int] = []
+        total = 0
+        for shard in self.shards:
+            if not isinstance(shard, dict):
+                raise ValueError("Safetensors descriptor shard entry is invalid.")
+            relative = shard.get("file")
+            declared_count = shard.get("count")
+            if (
+                not isinstance(declared_count, int)
+                or isinstance(declared_count, bool)
+                or declared_count < 0
+            ):
+                raise ValueError("Safetensors descriptor shard count is invalid.")
+            declared_sha256 = shard.get("sha256")
+            if not isinstance(declared_sha256, str) or len(declared_sha256) != 64:
+                raise ValueError("Safetensors descriptor shard SHA-256 is invalid.")
+            if not isinstance(relative, str):
+                raise ValueError("Safetensors descriptor index file is invalid.")
+            descriptor_path = _resolve_index_child(root, relative, label="index")
+            tensor_file = shard.get("tensor_file")
+            if not isinstance(tensor_file, str):
+                raise ValueError("Safetensors descriptor tensor file is invalid.")
+            tensor_path = _resolve_index_child(root, tensor_file, label="index")
+            if not tensor_path.is_file():
+                raise ValueError(f"Safetensors tensor shard is missing: {tensor_file}.")
+            digest = hashlib.sha256()
+            count = 0
+            with descriptor_path.open("rb") as handle:
+                for line in handle:
+                    digest.update(line)
+                    if line.strip():
+                        item = json.loads(line)
+                        if not isinstance(item, dict):
+                            raise ValueError("Safetensors record descriptor must be a JSON object.")
+                        item_tensor = item.get("tensor")
+                        if not isinstance(item_tensor, dict):
+                            raise ValueError("Safetensors descriptor is missing tensor metadata.")
+                        item_tensor_file = item_tensor.get("file")
+                        if not isinstance(item_tensor_file, str):
+                            raise ValueError("Safetensors descriptor tensor file is invalid.")
+                        _resolve_index_child(root, item_tensor_file, label="descriptor")
+                        if item_tensor_file != tensor_file:
+                            raise ValueError(
+                                "Safetensors descriptor tensor file does not match its shard."
+                            )
+                        count += 1
+                        _validate_tensor_descriptor(item_tensor)
+            if digest.hexdigest() != declared_sha256 or count != declared_count:
+                raise ValueError(
+                    f"Safetensors descriptor shard failed integrity validation: {relative}."
+                )
+            total += count
+            cumulative.append(total)
+        self._cumulative = tuple(cumulative)
+        self._count = total
+
+    def __len__(self) -> int:
+        return self._count
+
+    def _iter_shard(self, shard_index: int) -> Iterator[EmbeddingRecord]:
+        descriptor_path = _resolve_index_child(
+            self.root, str(self.shards[shard_index]["file"]), label="index"
+        )
+        with descriptor_path.open("r", encoding="utf-8") as handle:
+            for line in handle:
+                if line.strip():
+                    yield _record_from_safetensors_descriptor(self.root, json.loads(line))
+
+    def __iter__(self) -> Iterator[EmbeddingRecord]:
+        for shard_index in range(len(self.shards)):
+            yield from self._iter_shard(shard_index)
+
+    @overload
+    def __getitem__(self, index: int, /) -> EmbeddingRecord: ...
+
+    @overload
+    def __getitem__(self, index: slice, /) -> Sequence[EmbeddingRecord]: ...
+
+    def __getitem__(self, index: int | slice) -> EmbeddingRecord | Sequence[EmbeddingRecord]:
+        if isinstance(index, slice):
+            start, stop, step = index.indices(self._count)
+            return [self[position] for position in range(start, stop, step)]
+        position = index + self._count if index < 0 else index
+        if position < 0 or position >= self._count:
+            raise IndexError(index)
+        shard_index = bisect_right(self._cumulative, position)
+        previous = self._cumulative[shard_index - 1] if shard_index else 0
+        local_position = position - previous
+        for offset, record in enumerate(self._iter_shard(shard_index)):
+            if offset == local_position:
+                return record
+        raise IndexError(index)
+
+
+class SafetensorsStreamWriter:
+    """Bounded-memory, resumable publisher with immutable retained generations."""
+
+    def __init__(
+        self,
+        path: str | Path,
+        metadata: dict[str, Any],
+        *,
+        shard_size: int = DEFAULT_SHARD_SIZE,
+        existing: Iterable[EmbeddingRecord] = (),
+        reuse_existing: bool = False,
+        publish_initial: bool = True,
+        publish_incremental: bool = True,
+    ) -> None:
+        try:
+            from safetensors.torch import save_file
+        except ImportError as error:
+            raise ImportError("Saving embeddings requires the 'safetensors' package.") from error
+        if shard_size <= 0:
+            raise ValueError("shard_size must be positive.")
+
+        self.path = Path(path)
+        self.index_path = _index_path(path)
+        self.run_manifest_path = _run_manifest_path(path)
+        self.index_path.parent.mkdir(parents=True, exist_ok=True)
+        self.metadata = _persistent_metadata(
+            metadata,
+            descriptor_index="safetensors-generation-index",
+            record_count=0,
+        )
+        self.shard_size = shard_size
+        self.publish_incremental = publish_incremental
+        self._save_file = save_file
+        authoritative_payload = _authoritative_index_payload(path)
+        prefix = _safetensors_shard_prefix(path)
+        # A random generation identity prevents a new writer from reusing a
+        # previously published or interrupted generation name. Published files
+        # are immutable and remain available to lazy readers until explicit GC.
+        self._generation = uuid4().hex
+        self._prefix = prefix
+        self._shard_index = 0
+        self._seed_index = 0
+        self._commit_index = 0
+        self._descriptor_shards: list[dict[str, Any]] = []
+        self._record_count = 0
+        self._current: dict[str, Tensor] = {}
+        self._pending: list[tuple[EmbeddingRecord, str, str, tuple[int, ...], str]] = []
+        self._current_size = 0
+        if reuse_existing:
+            if authoritative_payload is None:
+                raise ValueError("Cannot resume without an authoritative safetensors index.")
+            authoritative_metadata = authoritative_payload.get("metadata")
+            if not isinstance(authoritative_metadata, dict) or authoritative_metadata.get(
+                "run_fingerprint"
+            ) != self.metadata.get("run_fingerprint"):
+                raise ValueError("Cannot resume a safetensors run with a different fingerprint.")
+            expected_prefix_length = (
+                len(existing) if isinstance(existing, Sequence) else sum(1 for _ in existing)
+            )
+            if authoritative_payload.get("version") == 2:
+                self._descriptor_shards = list(authoritative_payload.get("descriptor_shards", ()))
+                self._record_count = int(authoritative_payload.get("record_count", 0))
+            else:
+                legacy_records = list(authoritative_payload.get("records", ()))
+                self._record_count = len(legacy_records)
+                if legacy_records:
+                    self._descriptor_shards.extend(self._write_descriptor_seed(legacy_records))
+            if expected_prefix_length != self._record_count:
+                raise ValueError(
+                    "The resumable safetensors prefix does not match the validated "
+                    "embedding records."
+                )
+
+        if publish_initial:
+            self._publish_metadata(complete=False)
+
+    def _write_descriptor_file(
+        self,
+        name: str,
+        descriptors: Sequence[dict[str, Any]],
+        *,
+        tensor_file: str,
+    ) -> dict[str, Any]:
+        temporary = self.index_path.parent / f".{name}.tmp"
+        destination = self.index_path.parent / name
+        if temporary.exists() or destination.exists():
+            raise FileExistsError(
+                f"Refusing to reuse immutable safetensors generation path {destination}."
+            )
+        digest = hashlib.sha256()
+        with temporary.open("wb") as handle:
+            for item in descriptors:
+                encoded = (
+                    json.dumps(item, sort_keys=True, separators=(",", ":")).encode("utf-8") + b"\n"
+                )
+                handle.write(encoded)
+                digest.update(encoded)
+        temporary.replace(destination)
+        return {
+            "file": name,
+            "sha256": digest.hexdigest(),
+            "count": len(descriptors),
+            "tensor_file": tensor_file,
+        }
+
+    def _write_descriptor_seed(self, records: Sequence[dict[str, Any]]) -> list[dict[str, Any]]:
+        groups: list[tuple[str, list[dict[str, Any]]]] = []
+        for record in records:
+            tensor_file = str(record["tensor"]["file"])
+            if (
+                not groups
+                or groups[-1][0] != tensor_file
+                or len(groups[-1][1]) == _MAX_RECORDS_PER_DESCRIPTOR_SHARD
+            ):
+                groups.append((tensor_file, []))
+            groups[-1][1].append(record)
+        descriptor_shards: list[dict[str, Any]] = []
+        for tensor_file, descriptors in groups:
+            self._seed_index += 1
+            name = (
+                f"{self._prefix}-records-run-{self._generation}-seed-{self._seed_index:05d}.jsonl"
+            )
+            descriptor_shards.append(
+                self._write_descriptor_file(name, descriptors, tensor_file=tensor_file)
+            )
+        return descriptor_shards
+
+    def _write_shard(self) -> None:
+        if not self._current:
+            return
+        self._shard_index += 1
+        name = f"{self._prefix}-run-{self._generation}-{self._shard_index:05d}.safetensors"
+        temporary = self.index_path.parent / f".{name}.tmp"
+        destination = self.index_path.parent / name
+        if temporary.exists() or destination.exists():
+            raise FileExistsError(
+                f"Refusing to reuse immutable safetensors generation path {destination}."
+            )
+        self._save_file(self._current, temporary)
+        temporary.replace(destination)
+        descriptors: list[dict[str, Any]] = []
+        for record, key, dtype_name, shape, digest in self._pending:
+            descriptors.append(
+                {
+                    "id": record.id,
+                    "sequence": record.sequence,
+                    "tensor": {
+                        "file": name,
+                        "key": key,
+                        "dtype": dtype_name,
+                        "shape": list(shape),
+                        "sha256": digest,
+                    },
+                }
+            )
+        descriptor_name = (
+            f"{self._prefix}-records-run-{self._generation}-{self._shard_index:05d}.jsonl"
+        )
+        self._descriptor_shards.append(
+            self._write_descriptor_file(descriptor_name, descriptors, tensor_file=name)
+        )
+        self._record_count += len(descriptors)
+        self._current = {}
+        self._pending = []
+        self._current_size = 0
+
+    def append(
+        self,
+        records: Iterable[EmbeddingRecord],
+        *,
+        publish: bool | None = None,
+    ) -> None:
+        """Persist records while retaining at most one shard of tensors."""
+
+        for record in records:
+            position = self._record_count + len(self._pending)
+            tensor = record.load_tensor().detach().cpu().contiguous()
+            if tensor.dtype not in _DTYPE_NAMES:
+                raise TypeError(f"Unsupported tensor dtype {tensor.dtype}.")
+            nbytes = tensor.numel() * tensor.element_size()
+            if nbytes > self.shard_size:
+                raise ValueError(
+                    f"Embedding {position} requires {nbytes} bytes and cannot fit in a "
+                    f"{self.shard_size}-byte safetensors shard."
+                )
+            if self._current and (
+                self._current_size + nbytes > self.shard_size
+                or len(self._pending) == _MAX_RECORDS_PER_DESCRIPTOR_SHARD
+            ):
+                self._write_shard()
+                if self.publish_incremental:
+                    self._publish_metadata(complete=False)
+                position = self._record_count
+            key = f"embedding_{position:08d}"
+            self._current[key] = tensor
+            self._current_size += nbytes
+            self._pending.append(
+                (
+                    record,
+                    key,
+                    _DTYPE_NAMES[tensor.dtype],
+                    tuple(tensor.shape),
+                    tensor_sha256(tensor),
+                )
+            )
+        if publish:
+            self.publish(complete=False)
+
+    def _publish_metadata(
+        self,
+        *,
+        complete: bool,
+        metadata: dict[str, Any] | None = None,
+    ) -> EmbeddingResult:
+        """Atomically expose one self-consistent metadata generation."""
+
+        if metadata is not None:
+            self.metadata = _persistent_metadata(
+                metadata,
+                descriptor_index="safetensors-generation-index",
+            )
+        self.metadata["complete"] = complete
+        self.metadata["record_count"] = self._record_count
+        self._commit_index += 1
+        payload = {
+            "version": 2,
+            "format": "fastplms-embedding-safetensors",
+            "metadata": self.metadata,
+            "record_count": self._record_count,
+            "descriptor_shards": self._descriptor_shards,
+        }
+        generation_index_name = (
+            f"{self._prefix}-index-run-{self._generation}-{self._commit_index:05d}.json"
+        )
+        generation_index_path = self.index_path.parent / generation_index_name
+        temporary_generation_index = generation_index_path.with_name(
+            f".{generation_index_path.name}.tmp"
+        )
+        if temporary_generation_index.exists() or generation_index_path.exists():
+            raise FileExistsError(
+                f"Refusing to reuse immutable safetensors generation index {generation_index_path}."
+            )
+        encoded_index = _canonical_json_bytes(payload)
+        temporary_generation_index.write_bytes(encoded_index)
+        temporary_generation_index.replace(generation_index_path)
+
+        index_sha256 = hashlib.sha256(encoded_index).hexdigest()
+        index_reference = {
+            "file": generation_index_name,
+            "sha256": index_sha256,
+        }
+        run_manifest = {
+            "version": 2,
+            "format": "fastplms-embedding-run",
+            "index": index_reference,
+            "record_count": self._record_count,
+        }
+        pointer_identity = f"{self._generation}-{self._commit_index:05d}"
+        temporary_manifest = self.run_manifest_path.with_name(
+            f".{self.run_manifest_path.name}.{pointer_identity}.tmp"
+        )
+        temporary_manifest.write_bytes(_canonical_json_bytes(run_manifest))
+        temporary_manifest.replace(self.run_manifest_path)
+
+        # ``index.json`` is a non-authoritative convenience pointer. The run
+        # manifest is committed first, so interruption here cannot invalidate
+        # the newly committed generation.
+        stable_pointer = {
+            "version": 2,
+            "format": "fastplms-embedding-index-pointer",
+            "index": index_reference,
+        }
+        temporary_index = self.index_path.with_name(
+            f".{self.index_path.name}.{pointer_identity}.tmp"
+        )
+        temporary_index.write_bytes(_canonical_json_bytes(stable_pointer))
+        temporary_index.replace(self.index_path)
+
+        return load_safetensors_result(self.index_path)
+
+    def publish(
+        self,
+        *,
+        complete: bool,
+        metadata: dict[str, Any] | None = None,
+    ) -> EmbeddingResult:
+        """Flush the current shard and atomically expose a consistent generation."""
+
+        self._write_shard()
+        return self._publish_metadata(complete=complete, metadata=metadata)
+
+
+def save_safetensors_result(
+    result: EmbeddingResult,
+    path: str | Path,
+    *,
+    shard_size: int = DEFAULT_SHARD_SIZE,
+) -> EmbeddingResult:
+    """Write sharded safetensors without materializing the full result."""
+
+    writer = SafetensorsStreamWriter(
+        path,
+        result.metadata,
+        shard_size=shard_size,
+        publish_initial=False,
+        publish_incremental=False,
+    )
+    writer.append(result, publish=False)
+    return writer.publish(complete=bool(result.metadata.get("complete", True)))
+
+
+def load_safetensors_result(path: str | Path) -> EmbeddingResult:
+    """Load an indexed safetensors result without loading tensor payloads."""
+
+    payload, index_path, _ = _load_authoritative_index(path)
+    if payload.get("version") == 2:
+        lazy_records = _SafetensorsRecordSequence(
+            index_path.parent, payload.get("descriptor_shards", ())
+        )
+        if len(lazy_records) != payload.get("record_count"):
+            raise ValueError("Safetensors descriptor count does not match its generation index.")
+        return EmbeddingResult(lazy_records, payload.get("metadata", {}))
+
+    records: list[EmbeddingRecord] = []
+    for item in payload["records"]:
+        records.append(_record_from_safetensors_descriptor(index_path.parent, item))
+    return EmbeddingResult(records, payload.get("metadata", {}))
+
+
+def garbage_collect_safetensors_generations(
+    path: str | Path,
+    *,
+    dry_run: bool = True,
+    confirm_no_active_readers_or_writers: bool = False,
+) -> tuple[Path, ...]:
+    """Remove non-authoritative generations after an explicit exclusivity check.
+
+    Safetensors results retain immutable historical generations because an
+    already-open :class:`EmbeddingResult` resolves tensors through those exact
+    descriptor and shard paths. Destructive collection is therefore safe only
+    when the caller guarantees that no reader or writer for ``path`` remains
+    active. ``dry_run=True`` is the default and returns the paths that would be
+    removed without changing the output directory.
+    """
+
+    if not isinstance(dry_run, bool):
+        raise TypeError("dry_run must be a bool.")
+    if not isinstance(confirm_no_active_readers_or_writers, bool):
+        raise TypeError("confirm_no_active_readers_or_writers must be a bool.")
+    if not dry_run and not confirm_no_active_readers_or_writers:
+        raise ValueError(
+            "Destructive safetensors generation collection requires "
+            "confirm_no_active_readers_or_writers=True."
+        )
+
+    # Validate the full descriptor graph before identifying anything as stale.
+    load_safetensors_result(path)
+    payload, authoritative_index_path, _ = _load_authoritative_index(path)
+    stable_index_path = _index_path(path)
+    run_manifest_path = _run_manifest_path(path)
+    root = stable_index_path.parent
+    prefix = _safetensors_shard_prefix(path)
+    protected = {
+        stable_index_path.resolve(),
+        run_manifest_path.resolve(),
+        authoritative_index_path.resolve(),
+        *_referenced_shards(stable_index_path, payload),
+    }
+    for descriptor_shard in payload.get("descriptor_shards", ()):
+        relative = descriptor_shard.get("file")
+        if isinstance(relative, str):
+            protected.add(_resolve_index_child(root, relative, label="index").resolve())
+
+    candidates: set[Path] = set()
+    for pattern in (
+        f"{prefix}-run-*-*.safetensors",
+        f"{prefix}-records-run-*.jsonl",
+        f"{prefix}-index-run-*.json",
+        f".{prefix}-*.tmp",
+    ):
+        candidates.update(root.glob(pattern))
+    candidates.update(root.glob(f".{stable_index_path.name}.*.tmp"))
+    candidates.update(root.glob(f".{run_manifest_path.name}.*.tmp"))
+
+    stale = tuple(
+        sorted(
+            (candidate for candidate in candidates if candidate.resolve() not in protected),
+            key=lambda candidate: candidate.name,
+        )
+    )
+    if not dry_run:
+        for candidate in stale:
+            candidate.unlink(missing_ok=True)
+    return stale
+
+
+def _ensure_sqlite_schema(connection: sqlite3.Connection) -> None:
+    connection.executescript(
+        """
+        PRAGMA foreign_keys = ON;
+        CREATE TABLE IF NOT EXISTS runs (
+            run_id TEXT PRIMARY KEY,
+            metadata_json TEXT NOT NULL,
+            created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
+            published_order INTEGER
+        );
+        CREATE TABLE IF NOT EXISTS tensors (
+            run_id TEXT NOT NULL,
+            position INTEGER NOT NULL,
+            dtype TEXT NOT NULL,
+            shape_json TEXT NOT NULL,
+            data BLOB NOT NULL,
+            sha256 TEXT NOT NULL,
+            PRIMARY KEY (run_id, position),
+            FOREIGN KEY (run_id) REFERENCES runs(run_id) ON DELETE CASCADE
+        );
+        CREATE TABLE IF NOT EXISTS records (
+            run_id TEXT NOT NULL,
+            position INTEGER NOT NULL,
+            record_id TEXT NOT NULL,
+            sequence TEXT NOT NULL,
+            PRIMARY KEY (run_id, position),
+            FOREIGN KEY (run_id, position) REFERENCES tensors(run_id, position)
+                ON DELETE CASCADE
+        );
+        """
+    )
+    run_columns = {str(row[1]) for row in connection.execute("PRAGMA table_info(runs)").fetchall()}
+    if "published_order" not in run_columns:
+        connection.execute("ALTER TABLE runs ADD COLUMN published_order INTEGER")
+        # Databases created before staged publication exposed every stored run.
+        # Preserve that view for historical runs containing committed records.
+        connection.execute(
+            "UPDATE runs SET published_order = rowid "
+            "WHERE published_order IS NULL AND EXISTS ("
+            "SELECT 1 FROM records WHERE records.run_id = runs.run_id)"
+        )
+    connection.execute(
+        "CREATE INDEX IF NOT EXISTS runs_published_order_idx ON runs(published_order)"
+    )
+    if "published_order" not in run_columns:
+        # Schema upgrades run before callers open their data transaction.
+        # End the migration transaction explicitly so BEGIN IMMEDIATE below
+        # remains valid on existing databases.
+        connection.commit()
+
+
+def save_sqlite_result(result: EmbeddingResult, path: str | Path) -> EmbeddingResult:
+    """Transactionally store an ordered result in normalized SQLite tables."""
+
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    run_id = str(result.metadata.get("run_fingerprint", ""))
+    if not run_id:
+        raise ValueError("SQLite results require metadata['run_fingerprint'].")
+    metadata_json = json.dumps(
+        _persistent_metadata(
+            result.metadata,
+            descriptor_index="sqlite-records",
+            record_count=len(result),
+        ),
+        sort_keys=True,
+    )
+    with sqlite3.connect(path, timeout=30) as connection:
+        _ensure_sqlite_schema(connection)
+        connection.execute("PRAGMA journal_mode = WAL")
+        connection.execute("BEGIN IMMEDIATE")
+        connection.execute("DELETE FROM runs WHERE run_id = ?", (run_id,))
+        connection.execute(
+            "INSERT INTO runs(run_id, metadata_json, published_order) "
+            "SELECT ?, ?, COALESCE(MAX(published_order), 0) + 1 FROM runs",
+            (run_id, metadata_json),
+        )
+        for position, record in enumerate(result):
+            X = record.load_tensor().detach().cpu().contiguous()
+            dtype_name, shape_json, data = _encode_tensor(X)
+            digest = tensor_sha256(X)
+            connection.execute(
+                "INSERT INTO tensors VALUES (?, ?, ?, ?, ?, ?)",
+                (run_id, position, dtype_name, shape_json, data, digest),
+            )
+            connection.execute(
+                "INSERT INTO records VALUES (?, ?, ?, ?)",
+                (run_id, position, record.id, record.sequence),
+            )
+        connection.commit()
+    return load_sqlite_result(path, run_id=run_id)
+
+
+def initialize_sqlite_run(
+    path: str | Path,
+    metadata: dict[str, Any],
+    *,
+    resume: bool,
+) -> str:
+    """Create a resumable SQLite run without buffering tensor results."""
+
+    path = Path(path)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    run_id = str(metadata.get("run_fingerprint", ""))
+    if not run_id:
+        raise ValueError("SQLite runs require metadata['run_fingerprint'].")
+    with sqlite3.connect(path, timeout=30) as connection:
+        _ensure_sqlite_schema(connection)
+        connection.execute("PRAGMA journal_mode = WAL")
+        connection.execute("BEGIN IMMEDIATE")
+        exists = connection.execute("SELECT 1 FROM runs WHERE run_id = ?", (run_id,)).fetchone()
+        if exists and not resume:
+            connection.execute("DELETE FROM runs WHERE run_id = ?", (run_id,))
+            exists = None
+        if exists is None:
+            initial_metadata = _persistent_metadata(
+                metadata,
+                descriptor_index="sqlite-records",
+                record_count=0,
+            )
+            connection.execute(
+                "INSERT INTO runs(run_id, metadata_json) VALUES (?, ?)",
+                (run_id, json.dumps(initial_metadata, sort_keys=True)),
+            )
+        connection.commit()
+    return run_id
+
+
+def append_sqlite_records(
+    path: str | Path,
+    run_id: str,
+    start_position: int,
+    records: list[EmbeddingRecord],
+    *,
+    replace_metadata: dict[str, Any] | None = None,
+) -> None:
+    """Commit one ordered embedding batch so an interrupted run can resume."""
+
+    if not isinstance(run_id, str) or not run_id:
+        raise ValueError("run_id must be a non-empty string.")
+    if not isinstance(start_position, int) or isinstance(start_position, bool):
+        raise TypeError("start_position must be a non-negative integer.")
+    if start_position < 0:
+        raise ValueError("start_position must be a non-negative integer.")
+    if not isinstance(records, list) or not all(
+        isinstance(record, EmbeddingRecord) for record in records
+    ):
+        raise TypeError("records must be a list of EmbeddingRecord values.")
+
+    with sqlite3.connect(Path(path), timeout=30) as connection:
+        _ensure_sqlite_schema(connection)
+        connection.execute("PRAGMA journal_mode = WAL")
+        connection.execute("BEGIN IMMEDIATE")
+        if replace_metadata is not None:
+            replacement_run_id = str(replace_metadata.get("run_fingerprint", ""))
+            if replacement_run_id != run_id:
+                raise ValueError("Replacement metadata must match the SQLite run ID.")
+            initial_metadata = _persistent_metadata(
+                replace_metadata,
+                descriptor_index="sqlite-records",
+                record_count=0,
+            )
+            connection.execute("DELETE FROM runs WHERE run_id = ?", (run_id,))
+            connection.execute(
+                "INSERT INTO runs(run_id, metadata_json) VALUES (?, ?)",
+                (run_id, json.dumps(initial_metadata, sort_keys=True)),
+            )
+        if connection.execute("SELECT 1 FROM runs WHERE run_id = ?", (run_id,)).fetchone() is None:
+            raise KeyError(f"Missing SQLite embedding run {run_id}.")
+        current_count, minimum_position, maximum_position = connection.execute(
+            "SELECT COUNT(*), MIN(position), MAX(position) FROM records WHERE run_id = ?",
+            (run_id,),
+        ).fetchone()
+        if current_count and (minimum_position != 0 or maximum_position != current_count - 1):
+            raise ValueError("SQLite embedding run has a non-contiguous record prefix.")
+        if start_position != current_count:
+            raise ValueError(
+                f"start_position={start_position} does not match the contiguous "
+                f"SQLite prefix length {current_count}."
+            )
+        for offset, record in enumerate(records):
+            position = start_position + offset
+            X = record.load_tensor().detach().cpu().contiguous()
+            dtype_name, shape_json, data = _encode_tensor(X)
+            digest = tensor_sha256(X)
+            connection.execute(
+                "INSERT INTO tensors VALUES (?, ?, ?, ?, ?, ?)",
+                (run_id, position, dtype_name, shape_json, data, digest),
+            )
+            connection.execute(
+                "INSERT INTO records VALUES (?, ?, ?, ?)",
+                (run_id, position, record.id, record.sequence),
+            )
+        row = connection.execute(
+            "SELECT metadata_json FROM runs WHERE run_id = ?", (run_id,)
+        ).fetchone()
+        if row is None:
+            raise KeyError(f"Missing SQLite embedding run {run_id}.")
+        metadata = json.loads(row[0])
+        if not isinstance(metadata, dict):
+            raise ValueError("SQLite run metadata must contain a JSON object.")
+        metadata["record_count"] = start_position + len(records)
+        metadata["descriptor_index"] = "sqlite-records"
+        connection.execute(
+            "UPDATE runs SET metadata_json = ? WHERE run_id = ?",
+            (json.dumps(metadata, sort_keys=True), run_id),
+        )
+        if records:
+            connection.execute(
+                "UPDATE runs SET published_order = ("
+                "SELECT COALESCE(MAX(published_order), 0) + 1 FROM runs"
+                ") WHERE run_id = ? AND published_order IS NULL",
+                (run_id,),
+            )
+        connection.commit()
+
+
+def update_sqlite_run_metadata(path: str | Path, run_id: str, metadata: dict[str, Any]) -> None:
+    """Finalize reproducibility metadata after the last streamed batch."""
+
+    with sqlite3.connect(Path(path), timeout=30) as connection:
+        row = connection.execute(
+            "SELECT COUNT(*) FROM records WHERE run_id = ?", (run_id,)
+        ).fetchone()
+        record_count = int(row[0]) if row is not None else 0
+        cleaned_metadata = _persistent_metadata(
+            metadata,
+            descriptor_index="sqlite-records",
+            record_count=record_count,
+        )
+        updated = connection.execute(
+            "UPDATE runs SET metadata_json = ? WHERE run_id = ?",
+            (json.dumps(cleaned_metadata, sort_keys=True), run_id),
+        ).rowcount
+        if updated != 1:
+            raise KeyError(f"Missing SQLite embedding run {run_id}.")
+        connection.commit()
+
+
+def _connect_sqlite_read_only(path: Path) -> sqlite3.Connection:
+    if not path.is_file():
+        raise FileNotFoundError(path)
+    return sqlite3.connect(f"{path.resolve().as_uri()}?mode=ro", uri=True, timeout=30)
+
+
+def _validate_sqlite_result_schema(connection: sqlite3.Connection, path: Path) -> None:
+    tables = {
+        str(row[0])
+        for row in connection.execute(
+            "SELECT name FROM sqlite_master WHERE type = 'table'"
+        ).fetchall()
+    }
+    required = {"runs", "records", "tensors"}
+    if not required.issubset(tables):
+        raise ValueError(
+            f"Not a FastPLMs embedding SQLite database: {path}. "
+            "Use convert_legacy_sqlite() for a legacy embeddings table."
+        )
+
+
+def _load_sqlite_tensor(path: Path, run_id: str, position: int) -> Tensor:
+    with _connect_sqlite_read_only(path) as connection:
+        row = connection.execute(
+            "SELECT dtype, shape_json, data FROM tensors WHERE run_id = ? AND position = ?",
+            (run_id, position),
+        ).fetchone()
+    if row is None:
+        raise KeyError(f"Missing SQLite tensor {run_id}:{position}.")
+    return _decode_tensor(*row)
+
+
+def _validate_sqlite_descriptor_row(
+    row: Sequence[Any],
+) -> tuple[int, str, str, str, str, str]:
+    if len(row) != 6:
+        raise ValueError("SQLite embedding descriptor has an invalid column count.")
+    position, record_id, sequence, dtype_name, shape_json, digest = row
+    if not isinstance(position, int) or isinstance(position, bool) or position < 0:
+        raise ValueError("SQLite embedding position is invalid.")
+    if not isinstance(record_id, str) or not record_id:
+        raise ValueError("SQLite embedding record ID is invalid.")
+    if not isinstance(sequence, str) or not sequence:
+        raise ValueError("SQLite embedding sequence is invalid.")
+    if not isinstance(shape_json, str):
+        raise ValueError("SQLite embedding tensor shape is invalid.")
+    try:
+        shape = json.loads(shape_json)
+    except json.JSONDecodeError as error:
+        raise ValueError("SQLite embedding tensor shape is invalid.") from error
+    _validate_tensor_descriptor(
+        {
+            "key": f"embedding_{position}",
+            "dtype": dtype_name,
+            "shape": shape,
+            "sha256": digest,
+        }
+    )
+    return position, record_id, sequence, dtype_name, shape_json, digest
+
+
+def _sqlite_record_from_row(path: Path, run_id: str, row: Sequence[Any]) -> EmbeddingRecord:
+    position, record_id, sequence, dtype_name, shape_json, digest = _validate_sqlite_descriptor_row(
+        row
+    )
+
+    def load_tensor() -> Tensor:
+        return _load_sqlite_tensor(path, run_id, position)
+
+    reference = LazyTensorReference(
+        source=str(path),
+        key=f"{run_id}:{position}",
+        dtype=dtype_name,
+        shape=tuple(json.loads(shape_json)),
+        sha256=digest,
+        _loader=load_tensor,
+    )
+    return EmbeddingRecord(record_id, sequence, reference)
+
+
+class _SQLiteRecordSequence(Sequence[EmbeddingRecord]):
+    """Lazy immutable descriptor view over one SQLite embedding run."""
+
+    _fastplms_immutable_sequence = True
+
+    def __init__(self, path: Path, run_id: str, count: int) -> None:
+        self.path = path
+        self.run_id = run_id
+        self._count = count
+
+    @staticmethod
+    def _row_query() -> str:
+        return (
+            "SELECT r.position, r.record_id, r.sequence, t.dtype, t.shape_json, t.sha256 "
+            "FROM records r JOIN tensors t USING (run_id, position) "
+            "WHERE r.run_id = ?"
+        )
+
+    def __len__(self) -> int:
+        return self._count
+
+    def __iter__(self) -> Iterator[EmbeddingRecord]:
+        with _connect_sqlite_read_only(self.path) as connection:
+            cursor = connection.execute(f"{self._row_query()} ORDER BY r.position", (self.run_id,))
+            while rows := cursor.fetchmany(1_024):
+                for row in rows:
+                    yield _sqlite_record_from_row(self.path, self.run_id, row)
+
+    @overload
+    def __getitem__(self, index: int, /) -> EmbeddingRecord: ...
+
+    @overload
+    def __getitem__(self, index: slice, /) -> Sequence[EmbeddingRecord]: ...
+
+    def __getitem__(self, index: int | slice) -> EmbeddingRecord | Sequence[EmbeddingRecord]:
+        if isinstance(index, slice):
+            start, stop, step = index.indices(self._count)
+            return [self[position] for position in range(start, stop, step)]
+        position = index + self._count if index < 0 else index
+        if position < 0 or position >= self._count:
+            raise IndexError(index)
+        with _connect_sqlite_read_only(self.path) as connection:
+            row = connection.execute(
+                f"{self._row_query()} AND r.position = ?",
+                (self.run_id, position),
+            ).fetchone()
+        if row is None:
+            raise IndexError(index)
+        return _sqlite_record_from_row(self.path, self.run_id, row)
+
+
+def load_sqlite_result(
+    path: str | Path,
+    *,
+    run_id: str | None = None,
+    positions: Iterable[int] | None = None,
+    record_ids: Iterable[str] | None = None,
+    sequences: Iterable[str] | None = None,
+) -> EmbeddingResult:
+    """Load one SQLite run read-only, optionally in explicit selector order.
+
+    Exactly one selector may be supplied. Repeated selectors are retained. An
+    ID or sequence selector that matches multiple stored rows returns those
+    rows in their original order for every occurrence of that selector.
+    """
+
+    path = Path(path).resolve()
+    supplied_selectors = sum(
+        selector is not None for selector in (positions, record_ids, sequences)
+    )
+    if supplied_selectors > 1:
+        raise ValueError("Choose at most one of positions, record_ids, or sequences.")
+    normalized_positions = tuple(positions) if positions is not None else None
+    normalized_ids = tuple(record_ids) if record_ids is not None else None
+    normalized_sequences = tuple(sequences) if sequences is not None else None
+    if normalized_positions is not None and not all(
+        isinstance(position, int) and not isinstance(position, bool) and position >= 0
+        for position in normalized_positions
+    ):
+        raise ValueError("positions must contain non-negative integers.")
+    for name, values in (
+        ("record_ids", normalized_ids),
+        ("sequences", normalized_sequences),
+    ):
+        if values is not None and not all(isinstance(value, str) for value in values):
+            raise TypeError(f"{name} must contain strings.")
+
+    with _connect_sqlite_read_only(path) as connection:
+        _validate_sqlite_result_schema(connection, path)
+        if run_id is None:
+            run_columns = {
+                str(info[1]) for info in connection.execute("PRAGMA table_info(runs)").fetchall()
+            }
+            if "published_order" in run_columns:
+                row = connection.execute(
+                    "SELECT run_id, metadata_json FROM runs "
+                    "WHERE published_order IS NOT NULL "
+                    "ORDER BY published_order DESC, rowid DESC LIMIT 1"
+                ).fetchone()
+            else:
+                row = connection.execute(
+                    "SELECT run_id, metadata_json FROM runs "
+                    "ORDER BY created_at DESC, rowid DESC LIMIT 1"
+                ).fetchone()
+        else:
+            row = connection.execute(
+                "SELECT run_id, metadata_json FROM runs WHERE run_id = ?", (run_id,)
+            ).fetchone()
+        if row is None:
+            raise KeyError(f"No embedding run found in {path}.")
+        selected_run, metadata_json = row
+        metadata = json.loads(metadata_json)
+        if not isinstance(metadata, dict):
+            raise ValueError("SQLite run metadata must contain a JSON object.")
+        row_prefix = (
+            "SELECT r.position, r.record_id, r.sequence, t.dtype, t.shape_json, t.sha256 "
+            "FROM records r JOIN tensors t USING (run_id, position) "
+            "WHERE r.run_id = ?"
+        )
+        record_count, minimum_position, maximum_position = connection.execute(
+            "SELECT COUNT(*), MIN(position), MAX(position) FROM records WHERE run_id = ?",
+            (selected_run,),
+        ).fetchone()
+        (tensor_count,) = connection.execute(
+            "SELECT COUNT(*) FROM tensors WHERE run_id = ?", (selected_run,)
+        ).fetchone()
+        (joined_count,) = connection.execute(
+            "SELECT COUNT(*) FROM records r JOIN tensors t USING (run_id, position) "
+            "WHERE r.run_id = ?",
+            (selected_run,),
+        ).fetchone()
+        if (
+            tensor_count != record_count
+            or joined_count != record_count
+            or (record_count and (minimum_position != 0 or maximum_position != record_count - 1))
+        ):
+            raise ValueError("SQLite embedding run has inconsistent or non-contiguous records.")
+        metadata_count = metadata.get("record_count")
+        if (
+            not isinstance(metadata_count, int)
+            or isinstance(metadata_count, bool)
+            or metadata_count != record_count
+        ):
+            raise ValueError("SQLite metadata record count does not match stored records.")
+        descriptor_cursor = connection.execute(f"{row_prefix} ORDER BY r.position", (selected_run,))
+        while descriptor_rows := descriptor_cursor.fetchmany(1_024):
+            for descriptor_row in descriptor_rows:
+                _validate_sqlite_descriptor_row(descriptor_row)
+        if supplied_selectors == 0:
+            rows: list[tuple[Any, ...]] | None = None
+        else:
+            selector_values: tuple[Any, ...]
+            selector_column: str
+            if normalized_positions is not None:
+                selector_values = normalized_positions
+                selector_column = "r.position"
+            elif normalized_ids is not None:
+                selector_values = normalized_ids
+                selector_column = "r.record_id"
+            else:
+                if normalized_sequences is None:
+                    raise RuntimeError("Filtered SQLite retrieval resolved no selector values.")
+                selector_values = normalized_sequences
+                selector_column = "r.sequence"
+            fetched: list[tuple[Any, ...]] = []
+            unique_values = tuple(dict.fromkeys(selector_values))
+            for start in range(0, len(unique_values), 900):
+                chunk = unique_values[start : start + 900]
+                placeholders = ",".join("?" for _ in chunk)
+                fetched.extend(
+                    connection.execute(
+                        f"{row_prefix} AND {selector_column} IN ({placeholders}) "
+                        "ORDER BY r.position",
+                        (selected_run, *chunk),
+                    ).fetchall()
+                )
+            value_index = (
+                0 if normalized_positions is not None else (1 if normalized_ids is not None else 2)
+            )
+            matched: dict[Any, list[tuple[Any, ...]]] = {}
+            for fetched_row in sorted(fetched, key=lambda item: int(item[0])):
+                matched.setdefault(fetched_row[value_index], []).append(fetched_row)
+            missing = [value for value in selector_values if value not in matched]
+            if missing:
+                raise KeyError(f"SQLite embedding selectors were not found: {missing!r}.")
+            rows = [
+                fetched_row for value in selector_values for fetched_row in matched.get(value, ())
+            ]
+
+    if rows is None:
+        return EmbeddingResult(
+            _SQLiteRecordSequence(path, selected_run, int(record_count)),
+            metadata,
+        )
+    records = [_sqlite_record_from_row(path, selected_run, selected_row) for selected_row in rows]
+    if supplied_selectors:
+        metadata = dict(metadata)
+        metadata["selection"] = {
+            "kind": (
+                "positions"
+                if normalized_positions is not None
+                else "record_ids"
+                if normalized_ids is not None
+                else "sequences"
+            ),
+            "count": len(rows),
+            "duplicate_policy": "preserve-request-order",
+        }
+    return EmbeddingResult(records, metadata)
+
+
+def load_legacy_pth(path: str | Path, *, allow_unsafe_pickle: bool = False) -> EmbeddingResult:
+    """Import a legacy mapping-only ``.pth`` file after explicit opt-in."""
+
+    if not allow_unsafe_pickle:
+        raise ValueError(
+            "Legacy .pth loading can execute pickle payloads. Pass "
+            "allow_unsafe_pickle=True only for a trusted file."
+        )
+    payload = torch.load(Path(path), map_location="cpu", weights_only=False)
+    if not isinstance(payload, dict):
+        raise ValueError("A legacy .pth embedding file must contain a mapping.")
+    records: list[EmbeddingRecord] = []
+    for position, (sequence, X) in enumerate(payload.items()):
+        if not isinstance(sequence, str) or not isinstance(X, Tensor):
+            raise ValueError("Legacy embedding mappings must use str keys and Tensor values.")
+        records.append(EmbeddingRecord(str(position), sequence, X.detach().cpu()))
+    return EmbeddingResult(records, {"format": "legacy-pth", "unsafe_pickle": True})
+
+
+_LEGACY_COMPACT_VERSION = 0x01
+_LEGACY_CODE_DTYPES: dict[int, tuple[np.dtype[Any], torch.dtype]] = {
+    0: (np.dtype(np.float16), torch.float16),
+    # Legacy BF16 blobs stored FP16 payload bytes and converted back to BF16.
+    1: (np.dtype(np.float16), torch.bfloat16),
+    2: (np.dtype(np.float32), torch.float32),
+}
+
+
+def _decode_legacy_sqlite_blob(
+    data: bytes,
+    *,
+    fallback_shape: tuple[int, ...] | None,
+    allow_unsafe_pickle: bool,
+) -> Tensor:
+    if len(data) >= 6 and data[0] == _LEGACY_COMPACT_VERSION:
+        dtype_code = int(data[1])
+        if dtype_code not in _LEGACY_CODE_DTYPES:
+            raise ValueError(f"Unsupported legacy compact dtype code {dtype_code}.")
+        (ndim,) = struct.unpack_from("<i", data, 2)
+        if ndim < 0 or ndim > 16 or len(data) < 6 + 4 * ndim:
+            raise ValueError("Malformed legacy compact embedding header.")
+        shape = tuple(int(value) for value in struct.unpack_from(f"<{ndim}i", data, 6))
+        if any(size < 0 for size in shape):
+            raise ValueError("Malformed negative legacy embedding dimension.")
+        numpy_dtype, target_dtype = _LEGACY_CODE_DTYPES[dtype_code]
+        offset = 6 + 4 * ndim
+        expected = int(np.prod(shape, dtype=np.int64)) * numpy_dtype.itemsize
+        if len(data) - offset != expected:
+            raise ValueError("Legacy compact embedding payload length does not match shape.")
+        array = np.frombuffer(data, dtype=numpy_dtype, offset=offset).copy().reshape(shape)
+        return torch.from_numpy(array).to(dtype=target_dtype)
+
+    try:
+        loaded = torch.load(io.BytesIO(data), map_location="cpu", weights_only=True)
+    except Exception as safe_error:
+        if allow_unsafe_pickle:
+            loaded = torch.load(io.BytesIO(data), map_location="cpu", weights_only=False)
+        elif fallback_shape is None:
+            raise ValueError(
+                "Legacy embedding blob is neither compact nor safely loadable. "
+                "Provide fallback_shape for raw FP32 bytes, or set "
+                "allow_unsafe_pickle=True only for a trusted database."
+            ) from safe_error
+        else:
+            expected = int(np.prod(fallback_shape, dtype=np.int64)) * 4
+            if len(data) != expected:
+                raise ValueError(
+                    "Legacy raw FP32 payload length does not match fallback_shape."
+                ) from safe_error
+            array = np.frombuffer(data, dtype=np.float32).copy().reshape(fallback_shape)
+            return torch.from_numpy(array)
+    if not isinstance(loaded, Tensor):
+        raise ValueError("Legacy serialized embedding payload must contain one tensor.")
+    return loaded.detach().cpu()
+
+
+def convert_legacy_sqlite(
+    source: str | Path,
+    output: str | Path,
+    *,
+    fallback_shape: tuple[int, ...] | None = None,
+    allow_unsafe_pickle: bool = False,
+    metadata: dict[str, Any] | None = None,
+) -> EmbeddingResult:
+    """Convert the v0 ``embeddings(sequence, embedding)`` database safely.
+
+    The source is opened read-only. Compact blobs and ``weights_only`` Torch
+    tensors are accepted by default. Unsafe general pickle deserialization
+    remains an explicit opt-in.
+    """
+
+    source_path = Path(source)
+    output_path = Path(output)
+    if source_path.resolve() == output_path.resolve():
+        raise ValueError("Legacy SQLite conversion requires a different output path.")
+    if fallback_shape is not None and (
+        not fallback_shape or any(not isinstance(size, int) or size < 0 for size in fallback_shape)
+    ):
+        raise ValueError("fallback_shape must contain non-negative integer dimensions.")
+    with _connect_sqlite_read_only(source_path) as connection:
+        columns = {
+            str(row[1]) for row in connection.execute("PRAGMA table_info(embeddings)").fetchall()
+        }
+        if not {"sequence", "embedding"}.issubset(columns):
+            raise ValueError("Legacy SQLite database must contain embeddings(sequence, embedding).")
+        rows = connection.execute(
+            "SELECT sequence, embedding FROM embeddings ORDER BY rowid"
+        ).fetchall()
+    if not rows:
+        raise ValueError("Legacy SQLite database contains no embeddings.")
+
+    records: list[EmbeddingRecord] = []
+    content_digest = hashlib.sha256()
+    for position, (sequence, data) in enumerate(rows):
+        if not isinstance(sequence, str) or not sequence:
+            raise ValueError("Legacy embedding sequences must be non-empty strings.")
+        if not isinstance(data, bytes):
+            data = bytes(data)
+        tensor = _decode_legacy_sqlite_blob(
+            data,
+            fallback_shape=fallback_shape,
+            allow_unsafe_pickle=allow_unsafe_pickle,
+        )
+        tensor_digest = tensor_sha256(tensor)
+        for value in (sequence.encode("utf-8"), tensor_digest.encode("ascii")):
+            content_digest.update(len(value).to_bytes(8, "big"))
+            content_digest.update(value)
+        records.append(EmbeddingRecord(str(position), sequence, tensor))
+
+    content_sha256 = content_digest.hexdigest()
+    run_fingerprint = hashlib.sha256(
+        f"fastplms-legacy-sqlite-v1:{content_sha256}".encode("ascii")
+    ).hexdigest()
+    converted_metadata: dict[str, Any] = {
+        "format_version": 1,
+        "run_fingerprint": run_fingerprint,
+        "source_format": "legacy-fastplms-sqlite-v0",
+        "source_content_sha256": content_sha256,
+        "unsafe_pickle": allow_unsafe_pickle,
+        "complete": True,
+    }
+    if metadata:
+        converted_metadata["conversion_metadata"] = _jsonable(metadata)
+    return save_sqlite_result(
+        EmbeddingResult(records, converted_metadata),
+        output_path,
+    )
+
+
+def save_result(
+    result: EmbeddingResult,
+    path: str | Path,
+    *,
+    format: str = "safetensors",
+    shard_size: int = DEFAULT_SHARD_SIZE,
+) -> EmbeddingResult:
+    if format == "safetensors":
+        return save_safetensors_result(result, path, shard_size=shard_size)
+    if format == "sqlite":
+        return save_sqlite_result(result, path)
+    if format == "pth":
+        raise ValueError("Writing pickle-based .pth embeddings is not supported.")
+    raise ValueError("format must be 'safetensors' or 'sqlite'.")
+
+
+def load_result(path: str | Path, *, format: str = "safetensors") -> EmbeddingResult:
+    if format == "safetensors":
+        return load_safetensors_result(path)
+    if format == "sqlite":
+        return load_sqlite_result(path)
+    raise ValueError("format must be 'safetensors' or 'sqlite'.")
+
+
+__all__ = [
+    "DEFAULT_SHARD_SIZE",
+    "SafetensorsStreamWriter",
+    "append_sqlite_records",
+    "convert_legacy_sqlite",
+    "garbage_collect_safetensors_generations",
+    "initialize_sqlite_run",
+    "load_legacy_pth",
+    "load_result",
+    "load_safetensors_result",
+    "load_sqlite_result",
+    "safetensors_result_exists",
+    "save_result",
+    "save_safetensors_result",
+    "save_sqlite_result",
+    "tensor_sha256",
+    "update_sqlite_run_metadata",
+]

fastplms/embeddings/types.py

@@ -0,0 +1,187 @@
+"""Public value types for dataset embedding."""
+
+from __future__ import annotations
+
+from collections.abc import Callable, Iterator, Mapping, Sequence
+from dataclasses import dataclass, field
+from typing import Any, Literal, overload
+
+from torch import Tensor
+
+
+@dataclass(frozen=True, slots=True)
+class EmbeddingInput:
+    """One named protein sequence supplied to :func:`embed_dataset`."""
+
+    id: str
+    sequence: str
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.id, str) or not self.id:
+            raise ValueError("EmbeddingInput.id must be a non-empty string.")
+        if not isinstance(self.sequence, str) or not self.sequence:
+            raise ValueError("EmbeddingInput.sequence must be a non-empty string.")
+
+
+@dataclass(frozen=True, slots=True)
+class LazyTensorReference:
+    """A tensor stored outside memory and loaded only when requested."""
+
+    source: str
+    key: str
+    dtype: str
+    shape: tuple[int, ...]
+    sha256: str
+    _loader: Callable[[], Tensor] = field(repr=False, compare=False)
+
+    def load(self, *, verify: bool = True) -> Tensor:
+        """Load X and optionally verify its content digest."""
+
+        if not isinstance(verify, bool):
+            raise TypeError("verify must be a boolean.")
+        X = self._loader()
+        if not isinstance(X, Tensor):
+            raise TypeError(f"Stored tensor loader for {self.key!r} must return a Tensor.")
+        if tuple(X.shape) != self.shape:
+            raise ValueError(
+                f"Stored tensor {self.key!r} has shape {tuple(X.shape)}, expected {self.shape}."
+            )
+        dtype = str(X.dtype).removeprefix("torch.")
+        if dtype != self.dtype:
+            raise ValueError(
+                f"Stored tensor {self.key!r} has dtype {dtype!r}, expected {self.dtype!r}."
+            )
+        if verify:
+            from .storage import tensor_sha256
+
+            digest = tensor_sha256(X)
+            if digest != self.sha256:
+                raise ValueError(f"Stored tensor {self.key!r} failed SHA-256 verification.")
+        return X
+
+
+TensorValue = Tensor | LazyTensorReference
+
+
+@dataclass(frozen=True, slots=True)
+class EmbeddingRecord:
+    """One ordered embedding result."""
+
+    id: str
+    sequence: str
+    tensor: TensorValue
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.id, str) or not self.id:
+            raise ValueError("EmbeddingRecord.id must be a non-empty string.")
+        if not isinstance(self.sequence, str) or not self.sequence:
+            raise ValueError("EmbeddingRecord.sequence must be a non-empty string.")
+        if not isinstance(self.tensor, (Tensor, LazyTensorReference)):
+            raise TypeError("EmbeddingRecord.tensor must be a Tensor or LazyTensorReference.")
+
+    def load_tensor(self, *, verify: bool = True) -> Tensor:
+        """Return X regardless of whether this record is memory-backed or lazy."""
+
+        if not isinstance(verify, bool):
+            raise TypeError("verify must be a boolean.")
+        if isinstance(self.tensor, LazyTensorReference):
+            return self.tensor.load(verify=verify)
+        return self.tensor
+
+
+class EmbeddingResult(Sequence[EmbeddingRecord]):
+    """Ordered embedding records and the metadata needed to reproduce them."""
+
+    def __init__(
+        self,
+        records: Sequence[EmbeddingRecord],
+        metadata: Mapping[str, Any] | None = None,
+    ) -> None:
+        self.records: Sequence[EmbeddingRecord] = (
+            records if getattr(records, "_fastplms_immutable_sequence", False) else tuple(records)
+        )
+        self.metadata = dict(metadata or {})
+
+    def __len__(self) -> int:
+        return len(self.records)
+
+    def __iter__(self) -> Iterator[EmbeddingRecord]:
+        return iter(self.records)
+
+    @overload
+    def __getitem__(self, index: int, /) -> EmbeddingRecord: ...
+
+    @overload
+    def __getitem__(self, index: slice, /) -> Sequence[EmbeddingRecord]: ...
+
+    def __getitem__(self, index: int | slice) -> EmbeddingRecord | Sequence[EmbeddingRecord]:
+        return self.records[index]
+
+    def as_dict(
+        self,
+        *,
+        key: Literal["id", "sequence"] = "id",
+        duplicates: Literal["error", "first", "last"] = "error",
+        materialize: bool = True,
+    ) -> dict[str, TensorValue]:
+        """Convert records to a mapping under an explicit duplicate policy."""
+
+        if key not in {"id", "sequence"}:
+            raise ValueError("key must be 'id' or 'sequence'.")
+        if duplicates not in {"error", "first", "last"}:
+            raise ValueError("duplicates must be 'error', 'first', or 'last'.")
+        if not isinstance(materialize, bool):
+            raise TypeError("materialize must be a boolean.")
+        output: dict[str, TensorValue] = {}
+        for record in self.records:
+            record_key = getattr(record, key)
+            if record_key in output:
+                if duplicates == "error":
+                    raise ValueError(
+                        f"Duplicate {key} {record_key!r}; choose duplicates='first' "
+                        "or duplicates='last' explicitly."
+                    )
+                if duplicates == "first":
+                    continue
+            output[record_key] = record.load_tensor() if materialize else record.tensor
+        return output
+
+    def materialize(self, *, verify: bool = True) -> EmbeddingResult:
+        """Return an equivalent result with every X loaded into CPU memory."""
+
+        if not isinstance(verify, bool):
+            raise TypeError("verify must be a boolean.")
+        return EmbeddingResult(
+            [
+                EmbeddingRecord(
+                    id=record.id,
+                    sequence=record.sequence,
+                    tensor=record.load_tensor(verify=verify),
+                )
+                for record in self.records
+            ],
+            self.metadata,
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class EmbeddingBatch:
+    """Internal model-to-runner contract.
+
+    ``X`` has shape ``(b, l, d)`` and ``residue_mask`` has shape ``(b, l)``.
+    ``attentions`` may contain layer/head attention matrices for ``parti``.
+    """
+
+    X: Tensor
+    residue_mask: Tensor
+    attentions: Tensor | tuple[Tensor, ...] | None = None
+
+
+__all__ = [
+    "EmbeddingBatch",
+    "EmbeddingInput",
+    "EmbeddingRecord",
+    "EmbeddingResult",
+    "LazyTensorReference",
+    "TensorValue",
+]

fastplms/models.toml

@@ -0,0 +1,1223 @@
+schema_version = 1
+legal_files = [
+  "LICENSE=sha256:2d2b50c7b1414bff1189a1db1f0cfb92e3e064b50f4c2b1019827b683e1b629a",
+  "THIRD_PARTY_NOTICES.md=sha256:25704b3c76404696cae52e7fca13088d329f70f412687340351259e86cd62baa",
+]
+
+[[attention_kernels]]
+implementation = "flash_attention_2"
+repository = "kernels-community/flash-attn2"
+revision = "db6b51744f0cd7061386442c09df890fc6d9f47e"
+version = 2
+expected_variant = "flash_attn2"
+dtypes = ["bfloat16"]
+
+[[attention_kernels]]
+implementation = "flash_attention_3"
+repository = "kernels-community/flash-attn3"
+revision = "43f0bd269777115d94ff826e0d113ce9c1c9087b"
+version = 1
+expected_variant = "flash_attn3"
+dtypes = ["bfloat16"]
+
+[[runtime_assets]]
+id = "esmfold2_ccd"
+repository = "biohub/ESMFold2"
+revision = "1ebf0e3481a5184eb6171d40615c79e384b48796"
+path = "ccd.pkl"
+sha256 = "9ff44b1927c6b9198e38ffe0928706827a09a350c15530beeeabebfa88038fc5"
+size = 417306584
+consumer_family = "esmfold2"
+trust_kind = "hash_pinned_pickle"
+license = "MIT"
+offline_behavior = "requires_cached_verified_file"
+
+[[upstreams]]
+id = "ankh"
+path = "vendor/upstream/ankh"
+url = "https://github.com/agemagician/Ankh.git"
+revision = "02b4e25ce5389b9e771c9df6e546c62af1216f8e"
+license = "CC-BY-NC-SA-4.0"
+license_files = ["LICENSE.md"]
+license_digests = ["LICENSE.md=sha256:cd041d7f9f52936e8824ac3f754e9c67410763205fc8a7020ba74fc8b6edc088"]
+distribution_files = ["LICENSE.md=sha256:cd041d7f9f52936e8824ac3f754e9c67410763205fc8a7020ba74fc8b6edc088"]
+
+[[upstreams]]
+id = "biohub-esm"
+path = "vendor/upstream/biohub-esm"
+url = "https://github.com/Biohub/esm.git"
+revision = "82ee35553d39169d678f784c8d3f8712ffd7d2c4"
+license = "MIT"
+license_files = ["LICENSE.md", "THIRD_PARTY_NOTICE.md"]
+license_digests = [
+  "LICENSE.md=sha256:b63df9ca1dd96b3b21eec226b51b236d0bd152ac20eafc43aad46bf832b48d8a",
+  "THIRD_PARTY_NOTICE.md=sha256:5bff8515ba4e0f53abdc43714c180b79c5b606160497d98de741a369cb9b6a23",
+]
+distribution_files = [
+  "LICENSE.md=sha256:b63df9ca1dd96b3b21eec226b51b236d0bd152ac20eafc43aad46bf832b48d8a",
+  "THIRD_PARTY_NOTICE.md=sha256:5bff8515ba4e0f53abdc43714c180b79c5b606160497d98de741a369cb9b6a23",
+]
+
+[[upstreams]]
+id = "biohub-transformers"
+path = "vendor/upstream/biohub-transformers"
+url = "https://github.com/Biohub/transformers.git"
+revision = "3a8956fb4d4ea16b0ec8e71deef2c2909b6a5cbf"
+license = "Apache-2.0"
+license_files = ["LICENSE"]
+license_digests = ["LICENSE=sha256:77fd4710def9ec3c0f6225800e0235f15a425abd4a8b03559127fcd782612049"]
+distribution_files = ["LICENSE=sha256:77fd4710def9ec3c0f6225800e0235f15a425abd4a8b03559127fcd782612049"]
+
+[[upstreams]]
+id = "boltz"
+path = "vendor/upstream/boltz"
+url = "https://github.com/jwohlwend/boltz.git"
+revision = "b1ebfc46ecf57f5414e0d1a6f9027bbb122c53bc"
+license = "MIT"
+license_files = ["LICENSE"]
+license_digests = ["LICENSE=sha256:f0667fd5e66c51e1ba8ddaa0249c6d7225b30037e02c45782d8f2c2943ac2617"]
+distribution_files = ["LICENSE=sha256:f0667fd5e66c51e1ba8ddaa0249c6d7225b30037e02c45782d8f2c2943ac2617"]
+
+[[upstreams]]
+id = "dplm"
+path = "vendor/upstream/dplm"
+url = "https://github.com/bytedance/dplm.git"
+revision = "8a2e15e53416b4536f03f79ad1f6f6a9cbd5e19d"
+license = "Apache-2.0"
+license_files = ["LICENSE"]
+license_digests = ["LICENSE=sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30"]
+distribution_files = [
+  "LICENSE=sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30",
+  "PROVENANCE.md=sha256:a659f74be9073cf1ad2d2f7071531ca56959b421f111152cf4c41184ace5970e",
+]
+
+[[upstreams]]
+id = "e1"
+path = "vendor/upstream/e1"
+url = "https://github.com/Profluent-AI/E1.git"
+revision = "bfd2620a602248499f3d2583d85a7ecddf0b6e02"
+license = "Apache-2.0 AND Profluent-E1-Agreement"
+license_files = ["LICENSE", "ATTRIBUTION", "NOTICE"]
+license_digests = [
+  "LICENSE=sha256:8ef1dd556091544db3044164a8015424a3dcb3450fb3765a81b88463551bbe81",
+  "ATTRIBUTION=sha256:deb22b250f6491b649eda5c63e080dd56486b8d2736cea6a52ef875436214367",
+  "NOTICE=sha256:6de9db0320b4ee82f665c0951d8fd4cd53701a659c9dbce9bc3e3ea6afc4c6b3",
+]
+distribution_files = [
+  "LICENSE=sha256:8ef1dd556091544db3044164a8015424a3dcb3450fb3765a81b88463551bbe81",
+  "ATTRIBUTION=sha256:deb22b250f6491b649eda5c63e080dd56486b8d2736cea6a52ef875436214367",
+  "NOTICE=sha256:6de9db0320b4ee82f665c0951d8fd4cd53701a659c9dbce9bc3e3ea6afc4c6b3",
+  "Apache-2.0.txt=sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30",
+  "BSD-3-Clause.txt=sha256:36e1987f2f17db7f8ad36cd7a37dbb7aeaaf0ab68b97ab4b9d3556f3a7a76ae8",
+  "MODIFICATIONS.md=sha256:2506f47c0f5475af8e8ff2cff13eb8b79e8e25a08a054cdd617bf336536750ca",
+]
+
+[[upstreams]]
+id = "fair-esm"
+path = "vendor/upstream/fair-esm"
+url = "https://github.com/facebookresearch/esm.git"
+revision = "2b369911bb5b4b0dda914521b9475cad1656b2ac"
+license = "MIT"
+license_files = ["LICENSE"]
+license_digests = ["LICENSE=sha256:da6d3703ed11cbe42bd212c725957c98da23cbff1998c05fa4b3d976d1a58e93"]
+distribution_files = [
+  "LICENSE=sha256:da6d3703ed11cbe42bd212c725957c98da23cbff1998c05fa4b3d976d1a58e93",
+  "PROVENANCE.md=sha256:950adb94daf15e646ddf226dacfe2a8e77801aa0793e439a9a3490a48eb666e7",
+]
+
+[[upstreams]]
+id = "openfold"
+path = "vendor/upstream/openfold"
+url = "https://github.com/aqlaboratory/openfold.git"
+revision = "4b41059694619831a7db195b7e0988fc4ff3a307"
+license = "Apache-2.0"
+license_files = ["LICENSE"]
+license_digests = ["LICENSE=sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30"]
+distribution_files = [
+  "LICENSE=sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30",
+  "MODIFICATIONS.md=sha256:fd6f0aa1086a0c996cf967b326d18e965660cda0ad5c7f36a3474a8490720da3",
+  "PROVENANCE.md=sha256:48c903db43a217a3126afaefbac60b7ddac7efda2dfcc0cbff0bffc7d6c30081",
+]
+
+[[upstreams]]
+id = "protein-ttt"
+path = "vendor/upstream/protein-ttt"
+url = "https://github.com/anton-bushuiev/ProteinTTT.git"
+revision = "fde2817cd84b936167cc76ccabf31e5c0fe49962"
+license = "MIT"
+license_files = ["LICENSE"]
+license_digests = ["LICENSE=sha256:bb01e7d5554f9e2e117172e56551452f68a7818df7bc8e71cd7a776a1d4ba3df"]
+distribution_files = [
+  "LICENSE=sha256:bb01e7d5554f9e2e117172e56551452f68a7818df7bc8e71cd7a776a1d4ba3df",
+  "PROVENANCE.md=sha256:dc641c37353c2efd50ccbdb316ca4aae495ec02c1563e0e15bac92f75fc482e5",
+]
+
+[families.esm2]
+architecture = "ESM2"
+upstreams = ["fair-esm"]
+tokenizer_mode = "tokenizer"
+public_input = "Amino-acid sequences tokenized to residue IDs"
+extra = "core"
+reference_container = "reference-esm2"
+reference_adapter = "tests.parity.support.reference_adapters.esm2"
+attention = ["eager", "sdpa", "flex_attention", "flash_attention_2", "flash_attention_3"]
+dtypes = ["float32", "bfloat16"]
+bf16_execution = "fp32_parameters_autocast"
+precisions = ["default"]
+vram_tier = "sequence"
+checkpoint_license = "MIT"
+hub_license = "mit"
+weights_publication_allowed = true
+state_transform = "esm2_hf_to_fastplms_v1"
+conversion_provenance = "Input: the pinned official ESM2 state dictionary. Transformation: apply the deterministic esm2_hf_to_fastplms_v1 key map while preserving tensor values and materializing the tied input/output embedding values as independent tensors. Output: the pinned Synthyra FastPLMs checkpoint. Validation: release parity compares exact keys and values after the declared non-aliasing transform, tokenizer behavior, and inference. Limitation: any numerical rewrite requires a new transform identifier and exact conversion test."
+representative = "esm2_8m"
+documentation = "docs/models.md#esm2"
+test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/_esm_rotary.py", "models/esm2", "models/ttt.py"]
+auto_map = { AutoConfig = "fastplms.models.esm2.modeling_fastesm.FastEsmConfig", AutoModel = "fastplms.models.esm2.modeling_fastesm.FastEsmModel", AutoModelForMaskedLM = "fastplms.models.esm2.modeling_fastesm.FastEsmForMaskedLM", AutoModelForSequenceClassification = "fastplms.models.esm2.modeling_fastesm.FastEsmForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esm2.modeling_fastesm.FastEsmForTokenClassification" }
+
+[families.esm_plusplus]
+architecture = "ESMC"
+upstreams = ["biohub-esm", "biohub-transformers"]
+tokenizer_mode = "tokenizer"
+public_input = "Amino-acid sequences tokenized to residue IDs"
+extra = "core"
+reference_container = "reference-biohub-esm"
+reference_adapter = "tests.parity.support.reference_adapters.esm_plusplus"
+attention = ["eager", "sdpa", "flex_attention", "flash_attention_2", "flash_attention_3"]
+dtypes = ["float32", "bfloat16"]
+bf16_execution = "static_parameters"
+precisions = ["default"]
+vram_tier = "sequence"
+checkpoint_license = "MIT"
+hub_license = "mit"
+weights_publication_allowed = true
+state_transform = "esmc_to_fastplms_v1"
+conversion_provenance = "Input: the pinned Biohub ESMC checkpoint. Transformation: apply the deterministic esmc_to_fastplms_v1 parameter map into the FastPLMs ESMC modules. Output: the pinned Synthyra ESMplusplus checkpoint. Validation: release parity compares keys, shapes, dtypes, values, aliases, and live inference. Limitation: runtime attention and precision selection are not serialized weight transforms."
+representative = "esmc_small"
+documentation = "docs/models.md#esm-and-esmc"
+test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/esm_plusplus", "models/ttt.py"]
+auto_map = { AutoConfig = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusConfig", AutoModel = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusModel", AutoModelForMaskedLM = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusForMaskedLM" }
+
+[families.esm3]
+architecture = "ESM3"
+upstreams = ["biohub-esm", "biohub-transformers"]
+tokenizer_mode = "tokenizer"
+public_input = "Sequence, structure, and function tracks prepared through the multimodal helpers"
+extra = "core"
+reference_container = "reference-biohub-esm"
+reference_adapter = "tests.parity.support.reference_adapters.esm3"
+attention = ["eager", "sdpa", "flex_attention"]
+dtypes = ["float32", "bfloat16"]
+bf16_execution = "fp32_parameters_autocast"
+precisions = ["default"]
+vram_tier = "large-sequence"
+checkpoint_license = "MIT"
+hub_license = "mit"
+weights_publication_allowed = true
+state_transform = "esm3_to_fastplms_v1"
+conversion_provenance = "Input: the pinned Biohub ESM3 checkpoint. Transformation: apply the deterministic esm3_to_fastplms_v1 parameter map for the supported sequence and multimodal modules and expand BF16 checkpoint tensors to FP32 storage. Output: the pinned Synthyra ESM3 checkpoint. Validation: release parity compares exact state identity after the declared map and live feature behavior. Limitation: unsupported upstream modalities may not be inferred from this record."
+representative = "esm3_small"
+documentation = "docs/models.md#esm3"
+test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/esm3", "models/ttt.py"]
+auto_map = { AutoConfig = "fastplms.models.esm3.modeling_esm3.FastESM3Config", AutoModel = "fastplms.models.esm3.modeling_esm3.FastESM3Model" }
+
+[families.e1]
+architecture = "E1"
+upstreams = ["e1"]
+tokenizer_mode = "sequence"
+public_input = "Raw amino-acid sequences prepared by the native E1 adapter"
+extra = "core"
+reference_container = "reference-e1"
+reference_adapter = "tests.parity.support.reference_adapters.e1"
+attention = ["sdpa", "flex_attention"]
+dtypes = ["float32", "bfloat16"]
+bf16_execution = "static_parameters"
+precisions = ["default"]
+vram_tier = "sequence"
+checkpoint_license = "Profluent-E1-Agreement"
+hub_license = "other"
+hub_license_name = "Profluent-E1 Clickthrough License Agreement"
+hub_license_link = "https://github.com/Profluent-AI/E1/blob/bfd2620a602248499f3d2583d85a7ecddf0b6e02/LICENSE"
+weights_publication_allowed = true
+state_transform = "e1_to_fastplms_v1"
+conversion_provenance = "Input: the pinned Profluent-E1 checkpoint and tokenizer-free sequence contract. Transformation: apply e1_to_fastplms_v1 to the FastPLMs encoder and official task heads, storing floating tensors in BF16. Output: the pinned Synthyra Profluent-E1 checkpoint. Validation: release parity covers state identity after the declared cast, sequence and RAG preparation, aliases, and inference. Limitation: the FastPLMs scoring extension is not represented as an official E1 head."
+representative = "e1_150m"
+documentation = "docs/models.md#e1"
+test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/e1", "models/ttt.py"]
+auto_map = { AutoConfig = "fastplms.models.e1.modeling_e1.E1Config", AutoModel = "fastplms.models.e1.modeling_e1.E1Model", AutoModelForMaskedLM = "fastplms.models.e1.modeling_e1.E1ForMaskedLM", AutoModelForSequenceClassification = "fastplms.models.e1.modeling_e1.E1ForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.e1.modeling_e1.E1ForTokenClassification" }
+
+[families.dplm]
+architecture = "DPLM"
+upstreams = ["dplm"]
+tokenizer_mode = "tokenizer"
+public_input = "Amino-acid sequences tokenized to masked or partially masked residue IDs"
+extra = "core"
+reference_container = "reference-dplm"
+reference_adapter = "tests.parity.support.reference_adapters.dplm"
+attention = ["eager", "sdpa", "flex_attention", "flash_attention_3"]
+dtypes = ["float32", "bfloat16"]
+bf16_execution = "fp32_parameters_autocast"
+precisions = ["default"]
+vram_tier = "sequence"
+checkpoint_license = "Apache-2.0"
+hub_license = "apache-2.0"
+weights_publication_allowed = true
+state_transform = "dplm_to_fastplms_v1"
+conversion_provenance = "Input: the pinned official DPLM1 checkpoint. Transformation: apply dplm_to_fastplms_v1, omitting the unused absolute-position table for rotary checkpoints and materializing the tied input/output embedding values as independent tensors. Output: the pinned Synthyra DPLM checkpoint. Validation: release parity compares exact state identity after the declared transform, tokenizer behavior, generation, and inference. License basis: the pinned ByteDance DPLM Apache-2.0 LICENSE and README explicitly scope the repository release to the pretrained DPLM1 and DPLM2 weights; immutable evidence is recorded in LICENSES/dplm/PROVENANCE.md. Limitation: redistribution remains subject to Apache-2.0 and the pinned provenance record; no broader rights are inferred."
+representative = "dplm_150m"
+documentation = "docs/models.md#dplm"
+test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/_diffusion_generation.py", "models/_esm_rotary.py", "models/dplm", "models/ttt.py"]
+auto_map = { AutoConfig = "fastplms.models.dplm.modeling_dplm.DPLMConfig", AutoModel = "fastplms.models.dplm.modeling_dplm.DPLMModel", AutoModelForMaskedLM = "fastplms.models.dplm.modeling_dplm.DPLMForMaskedLM", AutoModelForSequenceClassification = "fastplms.models.dplm.modeling_dplm.DPLMForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.dplm.modeling_dplm.DPLMForTokenClassification" }
+
+[families.dplm2]
+architecture = "DPLM2"
+upstreams = ["dplm"]
+tokenizer_mode = "tokenizer"
+public_input = "Tokenized amino-acid and structure tracks with explicit modality boundaries"
+extra = "core"
+reference_container = "reference-dplm"
+reference_adapter = "tests.parity.support.reference_adapters.dplm2"
+attention = ["sdpa"]
+dtypes = ["float32", "bfloat16"]
+bf16_execution = "fp32_parameters_autocast"
+precisions = ["default"]
+vram_tier = "sequence"
+checkpoint_license = "Apache-2.0"
+hub_license = "apache-2.0"
+weights_publication_allowed = true
+state_transform = "dplm2_to_fastplms_v1"
+conversion_provenance = "Input: the pinned official DPLM2 checkpoint. Transformation: apply dplm2_to_fastplms_v1, retaining the independent language-model head and trained encoder contact head while omitting the unused absolute-position table for rotary checkpoints. Output: the pinned Synthyra DPLM2 checkpoint. Validation: release parity compares exact keys and values after the declared omission, non-aliasing, tokenizer behavior, generation, and inference. License basis: the pinned ByteDance DPLM Apache-2.0 LICENSE and README explicitly scope the repository release to the pretrained DPLM1 and DPLM2 weights; immutable evidence is recorded in LICENSES/dplm/PROVENANCE.md. Limitation: no head exception is permitted by this record, and redistribution remains subject to Apache-2.0."
+representative = "dplm2_150m"
+documentation = "docs/models.md#dplm2"
+test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/_diffusion_generation.py", "models/_esm_rotary.py", "models/dplm2", "models/ttt.py"]
+auto_map = { AutoConfig = "fastplms.models.dplm2.modeling_dplm2.DPLM2Config", AutoModel = "fastplms.models.dplm2.modeling_dplm2.DPLM2Model", AutoModelForMaskedLM = "fastplms.models.dplm2.modeling_dplm2.DPLM2ForMaskedLM", AutoModelForSequenceClassification = "fastplms.models.dplm2.modeling_dplm2.DPLM2ForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.dplm2.modeling_dplm2.DPLM2ForTokenClassification" }
+tokenizer_class = "fastplms.models.dplm2.tokenization_dplm2.DPLM2Tokenizer"
+
+[families.ankh]
+architecture = "ANKH"
+upstreams = ["ankh"]
+tokenizer_mode = "tokenizer"
+public_input = "Amino-acid sequences tokenized for encoder or sequence-to-sequence use"
+extra = "core"
+reference_container = "reference-ankh"
+reference_adapter = "tests.parity.support.reference_adapters.ankh"
+attention = ["eager", "sdpa"]
+dtypes = ["float32", "bfloat16"]
+bf16_execution = "static_parameters"
+precisions = ["default"]
+vram_tier = "large-sequence"
+checkpoint_license = "CC-BY-NC-SA-4.0"
+hub_license = "cc-by-nc-sa-4.0"
+weights_publication_allowed = true
+state_transform = "ankh_t5_to_fastplms_v1"
+conversion_provenance = "Input: the pinned official ANKH T5 checkpoint. Transformation: apply ankh_t5_to_fastplms_v1 to the official encoder and sequence-to-sequence heads. Output: the pinned Synthyra ANKH checkpoint. Validation: release parity compares exact mapped state, tokenizer behavior, official heads, and inference. Limitation: the separately named FastPLMs masked-language-model extension is not an official ANKH head."
+representative = "ankh_base"
+documentation = "docs/models.md#ankh"
+test_tiers = ["check", "compliance", "feature", "artifact", "benchmark"]
+requires_complete_weight_publication = true
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/ankh", "models/ttt.py"]
+auto_map = { AutoConfig = "fastplms.models.ankh.modeling_ankh.FastAnkhConfig", AutoModel = "fastplms.models.ankh.modeling_ankh.FastAnkhModel", AutoModelForMaskedLM = "fastplms.models.ankh.modeling_ankh.FastAnkhForMaskedLMExtension", AutoModelForSeq2SeqLM = "fastplms.models.ankh.modeling_ankh.FastAnkhForConditionalGeneration", AutoModelForSequenceClassification = "fastplms.models.ankh.modeling_ankh.FastAnkhForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.ankh.modeling_ankh.FastAnkhForTokenClassification" }
+
+[families.boltz2]
+architecture = "Boltz2"
+upstreams = ["boltz"]
+tokenizer_mode = "structure"
+public_input = "Raw amino-acid sequences through the convenience API, or prepared model features"
+extra = "structure"
+reference_container = "reference-boltz2"
+reference_adapter = "tests.parity.support.reference_adapters.boltz"
+attention = ["eager"]
+dtypes = ["float32", "bfloat16"]
+bf16_execution = "fp32_parameters_autocast"
+precisions = ["default"]
+vram_tier = "structure"
+checkpoint_license = "MIT"
+hub_license = "mit"
+weights_publication_allowed = true
+state_transform = "boltz2_inference_core_v1"
+conversion_provenance = "Input: the pinned official Boltz2 checkpoint. Transformation: select and map the supported Boltz2 inference-core parameters with boltz2_inference_core_v1. Output: the pinned Synthyra Boltz2 checkpoint. Validation: release parity covers state identity for the declared subset, feature preparation, seeded inference, and structure outputs. Limitation: this record does not claim support for undeclared upstream training components."
+representative = "boltz2"
+documentation = "docs/models.md#boltz2"
+test_tiers = ["structure", "artifact", "benchmark"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "models/boltz"]
+auto_map = { AutoConfig = "fastplms.models.boltz.modeling_boltz2.Boltz2Config", AutoModel = "fastplms.models.boltz.modeling_boltz2.Boltz2Model" }
+
+[families.esmfold]
+architecture = "ESMFold"
+upstreams = ["fair-esm", "openfold"]
+tokenizer_mode = "structure"
+public_input = "Raw amino-acid sequences through folding helpers, or prepared residue tensors"
+extra = "structure"
+reference_container = "reference-esmfold"
+reference_adapter = "tests.parity.support.reference_adapters.esmfold"
+attention = ["eager", "sdpa", "flex_attention"]
+dtypes = ["float32", "bfloat16"]
+bf16_execution = "fp32_parameters_autocast"
+precisions = ["default"]
+vram_tier = "structure"
+checkpoint_license = "MIT"
+hub_license = "mit"
+weights_publication_allowed = true
+state_transform = "esmfold_meta_to_fastplms_v1"
+conversion_provenance = "Input: the pinned native Meta ESMFold checkpoint plus its pinned ESM2 backbone. Transformation: apply esmfold_meta_to_fastplms_v1 to map native ESM2 names into the structure-only FastPLMs backbone, retain folding tensors, omit five deterministically reconstructed geometry buffers, omit the folding-unused ESM2 masked-LM and contact-regression heads, and remove the obsolete random FastPLMs TTT head from earlier mirrors. Output: canonical FP32 FastPLMs ESMFold state with an explicit CUDA BF16-autocast execution path. Validation: release parity compares exact mapped keys, shapes, dtypes, values, aliases, semantic configuration, FP32 and BF16-compute seeded inference, and structure metrics with pLDDT normalized to (0, 1). Limitation: ESMFold TTT is rejected because the official checkpoint contains no trained masked-language-model head."
+representative = "esmfold"
+documentation = "docs/models.md#esmfold"
+test_tiers = ["check", "compliance", "structure", "feature", "artifact", "benchmark"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/_esm_rotary.py", "models/esmfold"]
+auto_map = { AutoConfig = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmFoldConfig", AutoModel = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmForProteinFolding" }
+
+[families.esmfold2]
+architecture = "ESMFold2"
+upstreams = ["biohub-esm", "biohub-transformers", "protein-ttt"]
+backbone_model = "esmc_6b"
+tokenizer_mode = "structure"
+public_input = "Raw amino-acid sequences or typed molecular-complex specifications; low-level forward accepts prepared feature tensors"
+extra = "structure"
+reference_container = "reference-esmfold2"
+reference_adapter = "tests.parity.support.reference_adapters.esmfold2"
+attention = ["eager", "sdpa", "flex_attention"]
+dtypes = ["float32", "bfloat16"]
+bf16_execution = "fp32_parameters_autocast"
+precisions = ["auto", "fp32", "bf16", "fp8"]
+experimental_precisions = ["fp8"]
+vram_tier = "structure-6b"
+checkpoint_license = "MIT"
+hub_license = "mit"
+weights_publication_allowed = true
+state_transform = "identity"
+conversion_provenance = "Input: each pinned Biohub ESMFold2 checkpoint and its separately pinned ESMC checkpoint. Transformation: apply identity to preserve the folding checkpoint exactly, load its parameters in FP32 for CUDA BF16-autocast execution, retain canonical BF16 ESMC weights, and optionally rebuild exactly 80 ESMC attention output projections as transient Transformer Engine linears. Output: the corresponding pinned Synthyra ESMFold2 checkpoint plus its declared ESMC precision policy. Validation: release parity covers exact canonical state, learned projection, prepared features, and seeded BF16 folding; experimental FP8 validation covers strict unavailable-device behavior, all four variants, and three BF16-to-FP8 reload cycles on the standard variant. Limitation: only the four manifest-listed ESMFold2 variants are supported; FP8 is experimental, applies only to inference-time ESMC execution, and requires direct CUDA loading with Transformer Engine availability."
+representative = "esmfold2"
+documentation = "docs/esmfold2.md"
+test_tiers = ["check", "compliance", "structure", "feature", "artifact", "benchmark"]
+runtime_paths = ["__init__.py", "registry.py", "runtime.py", "models.toml", "models/__init__.py", "attention", "embeddings", "models/esmfold2", "models/esm_plusplus", "models/ttt.py"]
+auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2.ESMFold2Model" }
+
+[[models]]
+id = "esm2_8m"
+family = "esm2"
+size_category = "small"
+generation_contract = "not_applicable"
+official_golden = { metadata = "tests/goldens/esm2_8m.json=sha256:6975e86d1d8f27488bf2a676551feaa48cc19254c9d24b6acb09198122745609", tensors = "tests/goldens/esm2_8m.safetensors=sha256:b40217566c33c71988d28869de353be54a3b3ebfc21fdfd29056e88cf7e99f4c" }
+fast_repo = "Synthyra/ESM2-8M"
+fast_revision = "185ecbd45665d050a8dae326d91886d330c5f9d0"
+fast_files = [
+  "config.json=git-sha1:46d0a7b517f59123c6ebc6d1011585731cbab259",
+  "model.safetensors=sha256:c824e6ded5fb71c72bc5ac05300699947819023cb26cdaf6897665e6b2645e1b",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:3cfc5db0c6790859a3bc2a4dc053a813acd65295",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+official_repo = "facebook/esm2_t6_8M_UR50D"
+official_revision = "c731040fcd8d73dceaa04b0a8e6329b345b0f5df"
+official_files = [
+  "config.json=git-sha1:c2c6e65a87d9d20d47699ae236d605b80c741dd3",
+  "model.safetensors=sha256:24c5fa474c48f3b754b86efe752d5f189d2bcd88190fa2270fc92b2ef3034189",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:3f0d47e841e1cb75257aeaf76d156802899a217e",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+
+[[models.oracle_assets]]
+role = "weights"
+path = "models/esm2_t6_8M_UR50D.pt"
+url = "https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t6_8M_UR50D.pt"
+sha256 = "46f002a9870c9bdecd0ea887acb1f9a38a6b561e8f8bf8a6990b679b9d31b928"
+size = 30099493
+
+[[models.oracle_assets]]
+role = "contact_regression"
+path = "regression/esm2_t6_8M_UR50D-contact-regression.pt"
+url = "https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t6_8M_UR50D-contact-regression.pt"
+sha256 = "8f7a4557d57713b97ba0e484303007efb7230d25299c0ac47a0a1b12a87bbb9d"
+size = 1511
+
+[[models]]
+id = "esm2_35m"
+family = "esm2"
+size_category = "small"
+generation_contract = "not_applicable"
+official_golden = { metadata = "tests/goldens/esm2_35m.json=sha256:e919d3ce6d20b6a942d27d92323814ae7594a0129dc9c4de27c5053e96675bcd", tensors = "tests/goldens/esm2_35m.safetensors=sha256:c9b8bb616cf884fb7744521a2fcc6eed23586342d11241e6c9ef16454ec31e17" }
+fast_repo = "Synthyra/ESM2-35M"
+fast_revision = "37ab9f56b41e365b3bd9e25d6fefe9150fd910f0"
+fast_files = [
+  "config.json=git-sha1:4d428c9934572f39e2a00db162249971f37c88e4",
+  "model.safetensors=sha256:21d95ab6bb9aa91bfec87eff11da61a657b732f2df279cbddbae6a7f1f0bba9c",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:3cfc5db0c6790859a3bc2a4dc053a813acd65295",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+official_repo = "facebook/esm2_t12_35M_UR50D"
+official_revision = "6fbf070e65b0b7291e7bbcd451118c216cff79d8"
+official_files = [
+  "config.json=git-sha1:3f64131bb610ed1ce482c4b5421fc358c785278f",
+  "model.safetensors=sha256:e35647818e0e064351d4531ed480d225a002567b4b2b93ad3a9246d753150fc0",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:3f0d47e841e1cb75257aeaf76d156802899a217e",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+
+[[models.oracle_assets]]
+role = "weights"
+path = "models/esm2_t12_35M_UR50D.pt"
+url = "https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t12_35M_UR50D.pt"
+sha256 = "7f21e80e61d16a71735163ef555d3009afb0c98da74c48e29df08606973cc55e"
+size = 134095705
+
+[[models.oracle_assets]]
+role = "contact_regression"
+path = "regression/esm2_t12_35M_UR50D-contact-regression.pt"
+url = "https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t12_35M_UR50D-contact-regression.pt"
+sha256 = "16641e05d830d0ce863dd152dbb8c2f3ddfa3c3ec2a66080152c8abad01d8585"
+size = 1959
+
+[[models]]
+id = "esm2_150m"
+family = "esm2"
+size_category = "medium"
+generation_contract = "not_applicable"
+official_golden = { metadata = "tests/goldens/esm2_150m.json=sha256:c04c93486024ba0fa1c81fbfbe92ee79d1d4c7f1cfcc2c9886728522f752feab", tensors = "tests/goldens/esm2_150m.safetensors=sha256:c03fe9916dba137b452a6bbe944c7dc414db4019a6f0921e87b92d4bb6a8a42f" }
+fast_repo = "Synthyra/ESM2-150M"
+fast_revision = "979e0880dfc9e0c0080839b83d9d2dc05b92786a"
+fast_files = [
+  "config.json=git-sha1:efeae2af182b7d34dc35740a45f157661e7acdf4",
+  "model.safetensors=sha256:d1f7c60f98c31af328381519a750972b6a31b13b97aa7cca2e71b5ae1b3f8f53",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:3cfc5db0c6790859a3bc2a4dc053a813acd65295",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+official_repo = "facebook/esm2_t30_150M_UR50D"
+official_revision = "a695f6045e2e32885fa60af20c13cb35398ce30c"
+official_files = [
+  "config.json=git-sha1:52e04179e6fbad6663a94ea5cc44f09d764c5cd4",
+  "model.safetensors=sha256:c3f1da8aea53bddd32c246c86168c23b9fd72341fb9db9a94436f855f5053566",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:3f0d47e841e1cb75257aeaf76d156802899a217e",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+
+[[models.oracle_assets]]
+role = "weights"
+path = "models/esm2_t30_150M_UR50D.pt"
+url = "https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t30_150M_UR50D.pt"
+sha256 = "881c7176cf198ef8dec26a3c375d40eb58d0c33df95c22562ca6cc6d3f812c62"
+size = 592774773
+
+[[models.oracle_assets]]
+role = "contact_regression"
+path = "regression/esm2_t30_150M_UR50D-contact-regression.pt"
+url = "https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t30_150M_UR50D-contact-regression.pt"
+sha256 = "6a604b96722ed052eef8a094ad90b275ba2e987d406315dbed0bdc6b3c4238a7"
+size = 3431
+
+[[models]]
+id = "esm2_650m"
+family = "esm2"
+size_category = "large"
+generation_contract = "not_applicable"
+official_golden = { metadata = "tests/goldens/esm2_650m.json=sha256:f18332172fcb3abf5dd2485fd55f5b0d193ad3b93a44cc744e0d02817c927477", tensors = "tests/goldens/esm2_650m.safetensors=sha256:c3a66b75add03628e62e238cb63da6a9e4d321f8160e84bdf2a131c096977f86" }
+fast_repo = "Synthyra/ESM2-650M"
+fast_revision = "ca0718a5d52b80d5c60dd76860e55e061a95fb0a"
+fast_files = [
+  "config.json=git-sha1:88f6bd240680b29c3244df8292246048401f5caf",
+  "model.safetensors=sha256:a15142e94ecf36f0edde9b37796f591e609ebe1694ca411e93640f0ee384994a",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:3cfc5db0c6790859a3bc2a4dc053a813acd65295",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+official_repo = "facebook/esm2_t33_650M_UR50D"
+official_revision = "08e4846e537177426273712802403f7ba8261b6c"
+official_files = [
+  "config.json=git-sha1:a956a25d277f30bd870d3760b9a116f19ead885e",
+  "model.safetensors=sha256:a08adabb949fa67ad3c14b509d04fd60368b35007b0095e3358f81200c4f4db0",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:3f0d47e841e1cb75257aeaf76d156802899a217e",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+
+[[models.oracle_assets]]
+role = "weights"
+path = "models/esm2_t33_650M_UR50D.pt"
+url = "https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t33_650M_UR50D.pt"
+sha256 = "ea9d0522b335a8778dea6535a65301f10208dece28cd5865482b0b1fc446168c"
+size = 2604537549
+
+[[models.oracle_assets]]
+role = "contact_regression"
+path = "regression/esm2_t33_650M_UR50D-contact-regression.pt"
+url = "https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t33_650M_UR50D-contact-regression.pt"
+sha256 = "8ffe6edbd4173dc8d45c2cd5cb27d43aad77ec26b4c768200c58ae1f96693575"
+size = 3687
+
+[[models]]
+id = "esm2_3b"
+family = "esm2"
+size_category = "xlarge"
+generation_contract = "not_applicable"
+official_golden = { metadata = "tests/goldens/esm2_3b.json=sha256:5043b2333c57a34d54fac53916722d1acb4b6fd50395b9abafa805435b184a48", tensors = "tests/goldens/esm2_3b.safetensors=sha256:dfd5a8cb05d3e814a080185c4808c8e7ec2277f070f395562fcfbe4376789e4e" }
+notes = "The pinned default SDPA BF16 path uses a checkpoint-specific numeric calibration: relative L2 target/hard limit 0.06/0.07, relative Q99.9 0.15/0.18, first-percentile residue cosine 0.994/0.992, and pooled cosine 0.998/0.997. Exact state identity and the global logits-distribution contract remain required."
+fast_repo = "Synthyra/ESM2-3B"
+fast_revision = "ff89d0180f414ab9c677219a25da79bf09185456"
+fast_files = [
+  "config.json=git-sha1:94944ad6cabaa40a3ce1cbe6699cf464fdc1b2c0",
+  "model-00001-of-00003.safetensors=sha256:04b57854545c23779b562ee2ae22f10021ba0f4d586ba0ad482ee6eda187d562",
+  "model-00002-of-00003.safetensors=sha256:34954aaa05bc91635776ba6672946da5822626753d80db97b38c0538e9525102",
+  "model-00003-of-00003.safetensors=sha256:a6b3a55b9e3b2e1778de34c665c3dd17bdfdf6da9d6d5c97730c57168709ccae",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:3cfc5db0c6790859a3bc2a4dc053a813acd65295",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+official_repo = "facebook/esm2_t36_3B_UR50D"
+official_revision = "476b639933c8baad5ad09a60ac1a87f987b656fc"
+official_files = [
+  "config.json=git-sha1:69e7563923f87d2d7439bfb83e5a19b44b46d71b",
+  "pytorch_model-00001-of-00002.bin=sha256:0f971f11c449d21422aa982b791619c10351972992c735f4c3cd43fe09790412",
+  "pytorch_model-00002-of-00002.bin=sha256:7560b46fc383c691fb74b915b7d4bcef40d3df181447f16ba4b298845e308d0c",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:3f0d47e841e1cb75257aeaf76d156802899a217e",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+
+[[models.oracle_assets]]
+role = "weights"
+path = "models/esm2_t36_3B_UR50D.pt"
+url = "https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t36_3B_UR50D.pt"
+sha256 = "7de8b4082ba15891959ab368b77ce3886697af1efb16d3c9e9e7b0c5d3f07500"
+size = 5678116398
+
+[[models.oracle_assets]]
+role = "contact_regression"
+path = "regression/esm2_t36_3B_UR50D-contact-regression.pt"
+url = "https://dl.fbaipublicfiles.com/fair-esm/regression/esm2_t36_3B_UR50D-contact-regression.pt"
+sha256 = "4da500eab246481dc9c8c95bc7b1d02f2803d761c380b0e95186d4a07d0fc84e"
+size = 6759
+
+[[models]]
+id = "esmc_small"
+family = "esm_plusplus"
+size_category = "medium"
+generation_contract = "not_applicable"
+official_golden = { metadata = "tests/goldens/esmc_small.json=sha256:bb02652cf3cc484756b98ffa4ba55ed4c55870d2cea3342adb1d920ba9dfe10a", tensors = "tests/goldens/esmc_small.safetensors=sha256:03378d0f0fdd8161178ebb2c1f0da1b9776a726c8e8d3a10c009808a24de5654" }
+notes = "Release contract: SDPA must match the pinned Biohub implementation bit-for-bit across every hidden state, last hidden state, logits, special token, and padding position. Eager and FlashAttention 2 are release-gated in BF16 against the pinned boundary-length and biological panels with a relative-L2 engineering target of 0.029, hard limit of 0.03, relative-Q99.9 target of 0.049, first-percentile residue-cosine target of 0.997, and Jensen-Shannon target of 0.0004. The global pooled-cosine and top-1 thresholds remain unchanged. Flex Attention and FlashAttention 3 remain selectable as opt-in alternatives, but they are not strict-parity choices: on the locked H100 BF16 generated-boundary panel, ESMC-6B Flex Attention exceeds the 0.03 relative-L2 hard limit and FlashAttention 3 falls below the 0.995 residue-cosine hard limit. The deviation is consistent with backend-specific BF16 kernel arithmetic; it is not a weight-conversion difference or silent fallback. Use SDPA for exact Biohub parity or FlashAttention 2 for release-gated acceleration."
+fast_repo = "Synthyra/ESMplusplus_small"
+fast_revision = "46c5f7d562e47d4c14165b424c71ab7db008e6fb"
+fast_files = [
+  "config.json=git-sha1:df2f44187157b0cc371c48c887b77b1783679201",
+  "model.safetensors=sha256:d099223765bc4f1ae8d6c7e18561ce41df1d54073fdc5327ef0a229235a8f52a",
+  "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b",
+  "tokenizer.json=git-sha1:f49735e56cebab0e791aeaae777757b7fd114f71",
+  "tokenizer_config.json=git-sha1:2985ed2b8aa8ecfb1d12f53f47d2b8a44cc21756",
+]
+official_repo = "biohub/ESMC-300M"
+official_revision = "a59b831785f907e96e6a246b1d142bfb76df31ee"
+official_files = [
+  "config.json=git-sha1:9a49eacf4e65c39f74381f0f0d240e3b89ef43d7",
+  "model.safetensors=sha256:0772d8fe64bb25e14fe6f23b80e3c9a7d215d0da3c6cba5bd356d7c0e0bb22cc",
+  "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b",
+  "tokenizer.json=git-sha1:81c797f56768b22dec0301fa771f018b7e43e98c",
+  "tokenizer_config.json=git-sha1:2238856624f8d39f03af53a2576c2d9b18c82f61",
+]
+
+[[models]]
+id = "esmc_large"
+family = "esm_plusplus"
+size_category = "large"
+generation_contract = "not_applicable"
+official_golden = { metadata = "tests/goldens/esmc_large.json=sha256:7a4d614f67b6fde417f3fd89f61e7ec442ae284769734b2b73e14945a816a8fd", tensors = "tests/goldens/esmc_large.safetensors=sha256:e13302df4cf7e8381552f1043a8fd0f31f3e0d50b2ab6009fb86b7940ae8ff79" }
+notes = "Release contract: SDPA must match the pinned Biohub implementation bit-for-bit across every hidden state, last hidden state, logits, special token, and padding position. Eager and FlashAttention 2 are release-gated in BF16 against the pinned boundary-length and biological panels with a relative-L2 engineering target of 0.029, hard limit of 0.03, relative-Q99.9 target of 0.049, first-percentile residue-cosine target of 0.997, and Jensen-Shannon target of 0.0004. The global pooled-cosine and top-1 thresholds remain unchanged. Flex Attention and FlashAttention 3 remain selectable as opt-in alternatives, but they are not strict-parity choices: on the locked H100 BF16 generated-boundary panel, ESMC-6B Flex Attention exceeds the 0.03 relative-L2 hard limit and FlashAttention 3 falls below the 0.995 residue-cosine hard limit. The deviation is consistent with backend-specific BF16 kernel arithmetic; it is not a weight-conversion difference or silent fallback. Use SDPA for exact Biohub parity or FlashAttention 2 for release-gated acceleration."
+fast_repo = "Synthyra/ESMplusplus_large"
+fast_revision = "f813401638b3fddab09748aec1ad2bf537aa4208"
+fast_files = [
+  "config.json=git-sha1:5736371902fe5d04e2859be30ac7dbd31b271b25",
+  "model.safetensors=sha256:4aff3f8c5de68c4d3e3824eb2c478e4a47355d3f849f3c745e5c8a5ee6cff851",
+  "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b",
+  "tokenizer.json=git-sha1:f49735e56cebab0e791aeaae777757b7fd114f71",
+  "tokenizer_config.json=git-sha1:2985ed2b8aa8ecfb1d12f53f47d2b8a44cc21756",
+]
+official_repo = "biohub/ESMC-600M"
+official_revision = "a7e82012c83126b9eedb055fea9fa84b6c02f094"
+official_files = [
+  "config.json=git-sha1:71c8241dc28a5fb636248267a0927c0242b264c1",
+  "model.safetensors=sha256:e4232c30fd35fe2f57051ec88a703996ac94520580b4b836894207a3d45d9ff8",
+  "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b",
+  "tokenizer.json=git-sha1:81c797f56768b22dec0301fa771f018b7e43e98c",
+  "tokenizer_config.json=git-sha1:2238856624f8d39f03af53a2576c2d9b18c82f61",
+]
+
+[[models]]
+id = "esmc_6b"
+family = "esm_plusplus"
+size_category = "xlarge"
+generation_contract = "not_applicable"
+official_golden = { metadata = "tests/goldens/esmc_6b.json=sha256:e229d938719782f280fab22dfc4c43e86109fdb0cc523631168c5a491afaace3", tensors = "tests/goldens/esmc_6b.safetensors=sha256:a948945e985c7deaca7be8b7eed09c0a9521a2af3f2b10fc2ec7a7d2a0f99ada" }
+notes = "Release contract: SDPA must match the pinned Biohub implementation bit-for-bit across every hidden state, last hidden state, logits, special token, and padding position. Eager and FlashAttention 2 are release-gated in BF16 against the pinned boundary-length and biological panels with a relative-L2 engineering target of 0.029, hard limit of 0.03, relative-Q99.9 target of 0.049, first-percentile residue-cosine target of 0.997, and Jensen-Shannon target of 0.0004. The global pooled-cosine and top-1 thresholds remain unchanged. Flex Attention and FlashAttention 3 remain selectable as opt-in alternatives, but they are not strict-parity choices: on the locked H100 BF16 generated-boundary panel, ESMC-6B Flex Attention exceeds the 0.03 relative-L2 hard limit and FlashAttention 3 falls below the 0.995 residue-cosine hard limit. The deviation is consistent with backend-specific BF16 kernel arithmetic; it is not a weight-conversion difference or silent fallback. Use SDPA for exact Biohub parity or FlashAttention 2 for release-gated acceleration."
+fast_repo = "Synthyra/ESMplusplus_6B"
+fast_revision = "0d579cce3b0f09efa6b3baddf6cc3fd8c9b616c8"
+fast_files = [
+  "config.json=git-sha1:e740cbcf211f2511c70c25a1ff6017a757ba7a69",
+  "model-00001-of-00006.safetensors=sha256:d30d18703453019f2d2d050866309888720c28eebc9a10307d1ddf3799e85a65",
+  "model-00002-of-00006.safetensors=sha256:b3d85378ab5023f4160a96e9c8cbd4cc6f78a771a83c856e88d48112f555bc13",
+  "model-00003-of-00006.safetensors=sha256:52595519b59349c5c6e373e6f5ca4a3d48ea6dde345f7e61e24766df5fab0e5b",
+  "model-00004-of-00006.safetensors=sha256:e46c6113c89c6f3e9b072c1bef02d763a625c37bcd8f9da2ed9363891c9a0758",
+  "model-00005-of-00006.safetensors=sha256:6d92cb2bf9791de644de2ae86f8523d802ac3b4aaabfff0716ab6c2b97f6fb14",
+  "model-00006-of-00006.safetensors=sha256:5fc1a8632490bb34162823c35d0d591337b9e4195b22cc0560741397a6e9d0b3",
+  "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b",
+  "tokenizer.json=git-sha1:f49735e56cebab0e791aeaae777757b7fd114f71",
+  "tokenizer_config.json=git-sha1:2985ed2b8aa8ecfb1d12f53f47d2b8a44cc21756",
+]
+official_repo = "biohub/ESMC-6B"
+official_revision = "45b0fa5d7fb06faefbd5e3b89bdcef35d564e79a"
+official_files = [
+  "config.json=git-sha1:19f5fb09e4f630fb5b748a497183c22a87ec5102",
+  "model-00001-of-00006.safetensors=sha256:bd90149ff223e6ac1a0cac6147a5ae0df20d3a21df4f65356a1f19cd14f4aa8a",
+  "model-00002-of-00006.safetensors=sha256:f75e2144d8269fe2eb4b3e0823fb089b94f176d8024153e85b8fb573a42294fa",
+  "model-00003-of-00006.safetensors=sha256:f699f01ecc9691d9c6470492765fe54b8b5d2e9f277c139e89427433ffdfe0b2",
+  "model-00004-of-00006.safetensors=sha256:46add1b7be098bbfdc3073884851ba3057f1b33ea23a158b650a37007dabd13d",
+  "model-00005-of-00006.safetensors=sha256:1e1cb62f060a34e18f54a31a76683ef888b8cec59e73315f5b31d25d45a1f88c",
+  "model-00006-of-00006.safetensors=sha256:56c73e13ae96e777ce65eee99364056069ef93b646470f352f83c5f1037b1b18",
+  "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b",
+  "tokenizer.json=git-sha1:81c797f56768b22dec0301fa771f018b7e43e98c",
+  "tokenizer_config.json=git-sha1:2238856624f8d39f03af53a2576c2d9b18c82f61",
+]
+
+[[models]]
+id = "esm3_small"
+family = "esm3"
+tokenizer_source = "esmc_small"
+size_category = "large"
+generation_contract = "not_applicable"
+official_golden = { metadata = "tests/goldens/esm3_small.json=sha256:5470e8596cbba0e2882647eccbc53c36d8b48b0f3947d1fe0bcea68da1078c32", tensors = "tests/goldens/esm3_small.safetensors=sha256:d957922f810c9ab4c557d80d5aaaf6a3aab79a5a45e4638012a634a4134803b1" }
+fast_repo = "Synthyra/ESM3_small"
+fast_revision = "7ddb5a740f9e5f93933eb6410c0ee8684bc63ec1"
+fast_files = [
+  "config.json=git-sha1:60526e2fdd8af9d4fba17f323775458ef5a1a1f9",
+  "model-00001-of-00002.safetensors=sha256:a4c9b736c4c59d51180e966005a164859b47d5cd36e1f8ecdea619fbd34a0e92",
+  "model-00002-of-00002.safetensors=sha256:bea60e4e91b03bb00b6cedd29b07606b8543f0869fb74454af7b26e216d80d2b",
+  "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b",
+  "tokenizer.json=git-sha1:f49735e56cebab0e791aeaae777757b7fd114f71",
+  "tokenizer_config.json=git-sha1:2985ed2b8aa8ecfb1d12f53f47d2b8a44cc21756",
+]
+official_repo = "biohub/esm3-sm-open-v1"
+official_revision = "47f0545b2b6daf26a93439a3cd610f4f7f3d5478"
+official_files = [
+  "config.json=git-sha1:0967ef424bce6791893e9a57bb952f80fd536e93",
+  "data/weights/esm3_function_decoder_v0.pth=sha256:f76d074efcaccfe21365a4fa96f212dadd66798e1e49d809ab7ffbe025d227c9",
+  "data/weights/esm3_sm_open_v1.pth=sha256:5ead5a135c658068db6a4f1b933e72d6110992c4668822e1c0e2dcc53e38acd9",
+  "data/weights/esm3_structure_decoder_v0.pth=sha256:3b726258a44274792b40ce7ea307e10c5da09936368a4ffa2970264d909da65b",
+  "data/weights/esm3_structure_encoder_v0.pth=sha256:467acbaee703ba3ccde6e75241a912a316952e5ff071355f85c1d33c68704f40",
+]
+
+[[models]]
+id = "e1_150m"
+family = "e1"
+size_category = "small"
+generation_contract = "not_applicable"
+official_golden = { metadata = "tests/goldens/e1_150m.json=sha256:701a64a6ab1a2fec5a427555b6af96232526c15cb3d5b4dc7fb253ac8f20b922", tensors = "tests/goldens/e1_150m.safetensors=sha256:6558bc8f1a7b20629eaaaa6f72601d0c2cdb859a5dc13595549b1773b6e2de41" }
+fast_repo = "Synthyra/Profluent-E1-150M"
+fast_revision = "7c5f3bbf697226a2e0900db7a100f9201774a907"
+fast_files = [
+  "config.json=git-sha1:562ef21e722ca708064fc3d54d25b731d4ac8171",
+  "model.safetensors=sha256:d779ed3a4e23799aafc932dc09c9963428d10aa7075999b5f8851b39c76b67f6",
+]
+official_repo = "Profluent-Bio/E1-150m"
+official_revision = "c4dbfe827e4aa6ed7f95eaef50dc1e084f4d77dc"
+official_files = [
+  "config.json=git-sha1:485e649199b46fe6ee7456bebf7aae9b3d4baeab",
+  "model.safetensors=sha256:ba2656339005e6598642836acfdafde480fecc7e145ce0058eb54adf572c3484",
+]
+
+[[models]]
+id = "e1_300m"
+family = "e1"
+size_category = "medium"
+generation_contract = "not_applicable"
+official_golden = { metadata = "tests/goldens/e1_300m.json=sha256:d3478f3f5957a0e0377864074dde0107de890019f96cb63548ee17ffb8f3ec3a", tensors = "tests/goldens/e1_300m.safetensors=sha256:92778b9ef95a803ddc84b3e3ca764c59e045872a94bcff0eb0cd47647732c188" }
+fast_repo = "Synthyra/Profluent-E1-300M"
+fast_revision = "5ef52c0ad2ae2578f40622696b763523810e8e26"
+fast_files = [
+  "config.json=git-sha1:f5c91498b76a3e3282a0d716d87738abb1a1b6c1",
+  "model.safetensors=sha256:9271c4176a8a2e0905a0bb769570ba1c2978fb999a87da92db4cf2b041224864",
+]
+official_repo = "Profluent-Bio/E1-300m"
+official_revision = "5a2871c587eadbcc9237bc686ea45e5b4d28dfb3"
+official_files = [
+  "config.json=git-sha1:918cb09e6e96d4719ed85951f38c693360f9cdb8",
+  "model.safetensors=sha256:31e09a2542f45b04e6ce4adafb3b657f21e2d56d12bf68fd2266b1576a80bc9b",
+]
+
+[[models]]
+id = "e1_600m"
+family = "e1"
+size_category = "large"
+generation_contract = "not_applicable"
+official_golden = { metadata = "tests/goldens/e1_600m.json=sha256:914be191c28141c1f84535cdb69ead0588a2057bb19d46c5bc7f3891a3d6739e", tensors = "tests/goldens/e1_600m.safetensors=sha256:22ed8417a4651ded255099f6d15c63c2c40552e700d2b0470d1adfde3a39c513" }
+fast_repo = "Synthyra/Profluent-E1-600M"
+fast_revision = "6c8bf0ec83b0e0178677c528b101efffd0677742"
+fast_files = [
+  "config.json=git-sha1:1d35c0b35b473259875fd29ee80167487a0d6afe",
+  "model.safetensors=sha256:793483b1b3411eab73fe5214b94d1424ca0545992dfac6889cfc0186af472363",
+]
+official_repo = "Profluent-Bio/E1-600m"
+official_revision = "52d959fb87a609d15cf223a485127b29ed5c382a"
+official_files = [
+  "config.json=git-sha1:8a0a439ed4201462bc01189c9f8b43523b257b5c",
+  "model.safetensors=sha256:cfc108d4b98baaa62932331b40be265eae39dc382595bc3cde4a5ab55db1bf7a",
+]
+
+[[models]]
+id = "dplm_150m"
+family = "dplm"
+size_category = "small"
+generation_contract = "required"
+official_golden = { metadata = "tests/goldens/dplm_150m.json=sha256:3228551fe3bed951db9ec97347143ec4462ce7c221ac240b7ce7730948c1dc1f", tensors = "tests/goldens/dplm_150m.safetensors=sha256:392992235195beed97ab8359b90a2e11e52f4326606f99a471447bed81d146bd" }
+fast_repo = "Synthyra/DPLM-150M"
+fast_revision = "90ba742754151a774f3b7ed580170d0a76b3e69d"
+fast_files = [
+  "config.json=git-sha1:117ac2c1222152ef378abaad1f605e18c4a18ab0",
+  "model.safetensors=sha256:8bac5ac767ceb8deb511b272d32883f811768d56cb25e920cea94ba9b979ca14",
+  "special_tokens_map.json=git-sha1:ef5f0f7d7baf4947564eafcf79972d272cd80a15",
+  "tokenizer_config.json=git-sha1:80100348e3f2b8ab05b59f3352ea7631685083cd",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+official_repo = "airkingbd/dplm_150m"
+official_revision = "49b7125a5d28c6418fcc2f3c4fe799352ac1488b"
+official_files = [
+  "config.json=git-sha1:4910cb02f1840e9ac577026f601829604af58c74",
+  "pytorch_model.bin=sha256:ea4eaa99536b60ed76f945f71a1a5e604f08447ec3def5104a93ca6001a59961",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:dbcdd9fb2e742627ee310713615e0d7aeed0c34e",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+
+[[models]]
+id = "dplm_650m"
+family = "dplm"
+size_category = "large"
+generation_contract = "required"
+official_golden = { metadata = "tests/goldens/dplm_650m.json=sha256:bf58d0ce73aaac7e6fb1923ef3d9adad67122df2a3dd414c3229488ef9587a6d", tensors = "tests/goldens/dplm_650m.safetensors=sha256:073f0a6abea7e48f28c2d921ff8329a28e22627f01979277cb324908a01b3378" }
+fast_repo = "Synthyra/DPLM-650M"
+fast_revision = "05dc16d97c5c028aed924c9ed681cee4ab609760"
+fast_files = [
+  "config.json=git-sha1:3537150eb87b213a676d5840548625e220b60e8b",
+  "model.safetensors=sha256:e27a47b8ec1c078b3fccb36542210e20f0380c88828db2ca9acf3d8a25048bd8",
+  "special_tokens_map.json=git-sha1:ef5f0f7d7baf4947564eafcf79972d272cd80a15",
+  "tokenizer_config.json=git-sha1:80100348e3f2b8ab05b59f3352ea7631685083cd",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+official_repo = "airkingbd/dplm_650m"
+official_revision = "7a7e651baa667d094aba05e9dc1cf52a3332110a"
+official_files = [
+  "config.json=git-sha1:625574d625a4178ca6966e9545fee56026c0b634",
+  "pytorch_model.bin=sha256:db4e54343a89e7600f41c3aacbc593db1b0caee82ec28cab25ff2ae090eba39c",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:dbcdd9fb2e742627ee310713615e0d7aeed0c34e",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+
+[[models]]
+id = "dplm_3b"
+family = "dplm"
+size_category = "xlarge"
+generation_contract = "required"
+official_golden = { metadata = "tests/goldens/dplm_3b.json=sha256:a5b6df8b9c7b371976892ec1d6c45581a32ad3a6325c6c0a0b3267012848c8ed", tensors = "tests/goldens/dplm_3b.safetensors=sha256:75b0a0854fc391133920b0feaaeb8f69ab7568a88b3759627aca1556c4338c1e" }
+fast_repo = "Synthyra/DPLM-3B"
+fast_revision = "7d764dd3d70ecf1ac0e64693de64a0064aacac65"
+fast_files = [
+  "config.json=git-sha1:7f5baf9426be06760c86882948b0f4af2e681e22",
+  "model-00001-of-00003.safetensors=sha256:37b54855d087ef3e7d883464ae9d5ea3127ec15a16c6323d91ad16a6b98305c9",
+  "model-00002-of-00003.safetensors=sha256:042604fefb05ea8c360a48416ce7ba662a4f90b176b4baf646c5c1814c35e6e8",
+  "model-00003-of-00003.safetensors=sha256:b9ae04012665163c3fc9781dd04fcd69738ac20c07e615e98fc4483fd2c4de45",
+  "special_tokens_map.json=git-sha1:ef5f0f7d7baf4947564eafcf79972d272cd80a15",
+  "tokenizer_config.json=git-sha1:80100348e3f2b8ab05b59f3352ea7631685083cd",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+official_repo = "airkingbd/dplm_3b"
+official_revision = "53849d4a7fe944ae0b9cf2bbc0d2cc0054795b51"
+official_files = [
+  "config.json=git-sha1:f6206456e8c2f22ebe1d37fce3b5d50fd8073e68",
+  "pytorch_model-00001-of-00004.bin=sha256:0bcb86a115fe744ed686756db143f78851304e855e2f83cec58681c6080ced5f",
+  "pytorch_model-00002-of-00004.bin=sha256:daf3324f3be949e7dd1c3c84b28da7fec5151b1890cb0904e73427266856a06f",
+  "pytorch_model-00003-of-00004.bin=sha256:dbbeb7924a21059854f994931e23590b054aa000b10370a71c052c4aa36e9246",
+  "pytorch_model-00004-of-00004.bin=sha256:21c01740d091487db43446489d8a893dea1fcc6f2e1c1991ece13945f7ab4e07",
+  "special_tokens_map.json=git-sha1:ba0f9b53dbbf27934f7555e5d31e37bdea9317f1",
+  "tokenizer_config.json=git-sha1:dbcdd9fb2e742627ee310713615e0d7aeed0c34e",
+  "vocab.txt=git-sha1:6b946952cc35537226f07fd70957ee2f848880d2",
+]
+
+[[models]]
+id = "dplm2_150m"
+family = "dplm2"
+size_category = "small"
+generation_contract = "required"
+official_golden = { metadata = "tests/goldens/dplm2_150m.json=sha256:d269de779ea1503de72c77e7b2e6224afc9797bd945b40c571ff6faec782e4aa", tensors = "tests/goldens/dplm2_150m.safetensors=sha256:17fc26600938ba5364b8ecb96750786d33e9f92bcd4ea4df3e12a389340748eb" }
+artifact_source = "official"
+canonical_state_sha256 = "82e1751f59052b8de72b082517557db47947e8d9b4ac2f11278369e6c0cbf001"
+fast_repo = "Synthyra/DPLM2-150M"
+fast_revision = "182745b8dc5661f898481a4fa60a7af9d53385c4"
+fast_files = [
+  "config.json=git-sha1:07905a2e4327d27d073cd0390f140aec2976125a",
+  "model.safetensors=sha256:0a7751b3113027b1d9c966a5bda2d6ab831855de7aaa047b911731665a7c3cc6",
+  "special_tokens_map.json=git-sha1:e6378d20e897b8806734e65fd3ef9cf42a17631b",
+  "tokenizer_config.json=git-sha1:f2090783e3368b7323aa877e2b740e09f0862259",
+  "vocab.txt=git-sha1:9706a4277a5c39dc9b4ec7b283e8eb130ceaa7f2",
+]
+official_repo = "airkingbd/dplm2_150m"
+official_revision = "3451d984d06497f835ed49634bd68c9dfb54d730"
+official_files = [
+  "config.json=git-sha1:20f1e55c64fdc4d1d30f7b1df64b6167fa23dc7c",
+  "pytorch_model.bin=sha256:be7f5cf9e421f59fcc437e63ce1c7391099a314a4e9a4f10b8688785fa581238",
+  "special_tokens_map.json=git-sha1:eb760e9f49a55145bbe0c64922d4ec2d3de1692a",
+  "tokenizer_config.json=git-sha1:fc8c21760dcff173955afb106859e5f015d4f757",
+  "vocab.txt=git-sha1:e133a3abd4350ddc3fc62548e162c8df7e62cf37",
+]
+
+[[models]]
+id = "dplm2_650m"
+family = "dplm2"
+size_category = "large"
+generation_contract = "required"
+official_golden = { metadata = "tests/goldens/dplm2_650m.json=sha256:d9a7548f9af657a72d441ca70f27379863724fcce8ddd3da4f672104b7bfb772", tensors = "tests/goldens/dplm2_650m.safetensors=sha256:c4e0e467c252c3ac813363d2d4b17a5e3bd99e75fad315e76d97689b4655ddac" }
+artifact_source = "official"
+canonical_state_sha256 = "cba76b6602d2258de9fffff953b608d93cb8ef4a9e89b0bbd27e160c81e78bb4"
+fast_repo = "Synthyra/DPLM2-650M"
+fast_revision = "b9d8527a9473a54954fa2764f590b9ea1b435bb2"
+fast_files = [
+  "config.json=git-sha1:3e079579b214d48a09db57f2c60be6a1acea5baf",
+  "model.safetensors=sha256:92db08c7dbfd6c5e03fbfeaea3f36b09640ee794dcf5ea8d550527869a9f1d63",
+  "special_tokens_map.json=git-sha1:e6378d20e897b8806734e65fd3ef9cf42a17631b",
+  "tokenizer_config.json=git-sha1:f2090783e3368b7323aa877e2b740e09f0862259",
+  "vocab.txt=git-sha1:9706a4277a5c39dc9b4ec7b283e8eb130ceaa7f2",
+]
+official_repo = "airkingbd/dplm2_650m"
+official_revision = "0bc69b644976c6680ab7e26669854d1979e8876e"
+official_files = [
+  "config.json=git-sha1:4cce8d9dc212cdace0e20e89169790bcf199c158",
+  "pytorch_model.bin=sha256:8d6e08cc05e4858064a714013c74cc88c9caa2cc8b12c34605a3c24bcd877cfb",
+  "special_tokens_map.json=git-sha1:eb760e9f49a55145bbe0c64922d4ec2d3de1692a",
+  "tokenizer_config.json=git-sha1:fc8c21760dcff173955afb106859e5f015d4f757",
+  "vocab.txt=git-sha1:e133a3abd4350ddc3fc62548e162c8df7e62cf37",
+]
+
+[[models]]
+id = "dplm2_3b"
+family = "dplm2"
+size_category = "xlarge"
+# The pinned public sampler fails before generation because cls_token_id is None.
+# State, tokenizer, and inference parity remain required for this checkpoint.
+generation_contract = "official_unavailable"
+official_golden = { metadata = "tests/goldens/dplm2_3b.json=sha256:d6e0e02af53b13cb129192f06e264758aa21c9ebf4ee82411cf67037082d2329", tensors = "tests/goldens/dplm2_3b.safetensors=sha256:838b11824d08f83bcb0c0b3268e579f3a87dbfb965370cfe5c3f8793b96b1964" }
+notes = "The pinned official DPLM2-3B sampler fails before generation, so live generation equivalence cannot be established for this checkpoint. State, tokenizer, and inference parity remain required."
+artifact_source = "official"
+canonical_state_sha256 = "8c46ec09115dbe6cbfb91d94ab5e906369d57e27fe620a7741c6f8cb1b6ca890"
+fast_repo = "Synthyra/DPLM2-3B"
+fast_revision = "2a63babe8848abf5233d31bd55891dff8285fc50"
+fast_files = [
+  "config.json=git-sha1:5932b1d501fed28b84614e0d2c1ecc4e89f10d6e",
+  "model-00001-of-00003.safetensors=sha256:2ff393f6e8df1568ce075d50de69ff4e5e9d9886e5ec47e43d6c24df23459be3",
+  "model-00002-of-00003.safetensors=sha256:feb3cea852c2aa849cc30783a984a97f0d076990ade6606cda5e38bf2a5a9621",
+  "model-00003-of-00003.safetensors=sha256:9be363ddb98436af20901981ffbed2f1097377424987f6c1baad27d512b62e71",
+  "special_tokens_map.json=git-sha1:e6378d20e897b8806734e65fd3ef9cf42a17631b",
+  "tokenizer_config.json=git-sha1:f2090783e3368b7323aa877e2b740e09f0862259",
+  "vocab.txt=git-sha1:9706a4277a5c39dc9b4ec7b283e8eb130ceaa7f2",
+]
+official_repo = "airkingbd/dplm2_3b"
+official_revision = "9e77567926f98d1b997ea9131a8eeb035b9bf827"
+official_files = [
+  "config.json=git-sha1:22d51ce44cd6da8d819e0d00566987bb51d74753",
+  "pytorch_model-00001-of-00004.bin=sha256:d8c641eae6bf891581ec64d543169891b093e296f5679ac75c695bcf596b4211",
+  "pytorch_model-00002-of-00004.bin=sha256:6478ad86ec5fef3d1d26580493af2d8666009d3ff884f3f88548080c8bbf94b5",
+  "pytorch_model-00003-of-00004.bin=sha256:dde8f88dac4a6355488c2fb433ee12cd69f1169950566624fba43684d4d99dc6",
+  "pytorch_model-00004-of-00004.bin=sha256:17ec0145152bc10e4dd3b4c2edff337979f6b99ee7c7bfd6cf4e6dbd7262d079",
+  "special_tokens_map.json=git-sha1:eb760e9f49a55145bbe0c64922d4ec2d3de1692a",
+  "tokenizer_config.json=git-sha1:fc8c21760dcff173955afb106859e5f015d4f757",
+  "vocab.txt=git-sha1:e133a3abd4350ddc3fc62548e162c8df7e62cf37",
+]
+
+[[models]]
+id = "ankh_base"
+family = "ankh"
+size_category = "medium"
+generation_contract = "required"
+official_golden = { metadata = "tests/goldens/ankh_base.json=sha256:ebce8d7de821827ee995789c9b38d79252d3b2f76888130b0a8a7eedafaefe2b", tensors = "tests/goldens/ankh_base.safetensors=sha256:f0e78aa15d11749e0c64ff57f9e88c51cec6538a0adf8951f839df70cc708b65" }
+notes = "ANKH parity covers the official encoder and sequence-to-sequence heads. AutoModelForMaskedLM exposes the separately named FastPLMs synthesized masked-LM extension and is not an official ANKH head."
+artifact_source = "official"
+canonical_state_sha256 = "cdd8d30d88e5bf41f44e1eef4470d8e46607aba5f7c7c805b06c035b89c8c16f"
+fast_repo = "Synthyra/ANKH_base"
+fast_revision = "7ec329aae8e3e174bf22a1eb9e0e9fcc12b53092"
+fast_files = [
+  "config.json=git-sha1:7e1cbce6d08f9bb64eee4410899b1c6b4054f418",
+  "model.safetensors=sha256:b0d3473cac1bda90e39cde54f2abe86da1fc84f872c833ca3415672776dccb95",
+  "special_tokens_map.json=git-sha1:a2d8d626c31389a935e197fb94072e2414a6e7d1",
+  "tokenizer.json=git-sha1:0734d752d12d0f46ac96467fbceb1c4bfbeee0be",
+  "tokenizer_config.json=git-sha1:db0b80de72d3b16242b9eda74ed4663e39c65bcf",
+]
+official_repo = "ElnaggarLab/ankh-base"
+official_revision = "d99cb6b966530dfc2ae96bc69d9255c2a07308b0"
+official_files = [
+  "config.json=git-sha1:abd44a36b5469e9a7cb019e4059b5ac1392d8422",
+  "pytorch_model.bin=sha256:9b2a886374f0ff4a893f4e7a989deed76bb2458c8998bd5202ea8e97d92ddcc3",
+  "special_tokens_map.json=git-sha1:55b145827029ae9672e50d4bb368540daacce791",
+  "tokenizer.json=git-sha1:212c5ef08819fa2463c6289ba4ef7db30e715c0a",
+  "tokenizer_config.json=git-sha1:a8a872ae3441e7cc85ce19210dff1e4c5d2d7bd0",
+]
+
+[[models]]
+id = "ankh_large"
+family = "ankh"
+size_category = "large"
+generation_contract = "required"
+official_golden = { metadata = "tests/goldens/ankh_large.json=sha256:59492518b021de5cfaea87d672c9448c8558e99a3443ba2cc7ab544963196ecb", tensors = "tests/goldens/ankh_large.safetensors=sha256:3fb8d3ac27716d15a9ea92aeef6acf2b977bcc887d9b535000539e523673459b" }
+notes = "ANKH parity covers the official encoder and sequence-to-sequence heads. AutoModelForMaskedLM exposes the separately named FastPLMs synthesized masked-LM extension and is not an official ANKH head."
+artifact_source = "official"
+canonical_state_sha256 = "e498a2e9aea76ef784cbe3e596c6b3f5e9a40e209ad837f7e3207099e4d74483"
+fast_repo = "Synthyra/ANKH_large"
+fast_revision = "3be3df34140f49dc4e65bd1f247e3ce819e7fc59"
+fast_files = [
+  "config.json=git-sha1:272509deedb527e5c2c95b0c269194a44148fdcc",
+  "model.safetensors=sha256:e70b8f9755ac6bfe95d18359060ae9fe38fac63b12a89a886c83349d1adbaa53",
+  "special_tokens_map.json=git-sha1:a2d8d626c31389a935e197fb94072e2414a6e7d1",
+  "tokenizer.json=git-sha1:0734d752d12d0f46ac96467fbceb1c4bfbeee0be",
+  "tokenizer_config.json=git-sha1:2bcaff2567826f5f51188b00600d2c6e7bcea56e",
+]
+official_repo = "ElnaggarLab/ankh-large"
+official_revision = "74b371dbfa3ee0a05d32ae74df0c2e0b82d6b9a6"
+official_files = [
+  "config.json=git-sha1:1abf33e52ee3d6be67d780ec57d32ac2b27b5306",
+  "pytorch_model.bin=sha256:517b6e8b279dedcb477af240b35c46bd6eb3307723eb281e60d4b2c8a87b889b",
+  "special_tokens_map.json=git-sha1:55b145827029ae9672e50d4bb368540daacce791",
+  "tokenizer.json=git-sha1:212c5ef08819fa2463c6289ba4ef7db30e715c0a",
+  "tokenizer_config.json=git-sha1:d7fe02ba6f2b18d9ccfa19ac129c9fdc9ec24d09",
+]
+
+[[models]]
+id = "ankh2_large"
+family = "ankh"
+size_category = "large"
+generation_contract = "required"
+official_golden = { metadata = "tests/goldens/ankh2_large.json=sha256:e8df38994ca1a1e0c598ace34a0b257b264937e4fdbb01bc41544985116b02a4", tensors = "tests/goldens/ankh2_large.safetensors=sha256:25fe1569f55c635fab8fa49c1d62a889a35a2a738bad921f5764a85b58fd4b5d" }
+notes = "ANKH parity covers the official encoder and sequence-to-sequence heads. AutoModelForMaskedLM exposes the separately named FastPLMs synthesized masked-LM extension and is not an official ANKH head."
+artifact_source = "official"
+canonical_state_sha256 = "597c4fe2fa8711f11a25317905f1d62fa92905e55fdd5c0a79614cd9c9d2bca3"
+fast_repo = "Synthyra/ANKH2_large"
+fast_revision = "392de5ed52bbfd73b45f545e378aaebcff096d0e"
+fast_files = [
+  "config.json=git-sha1:66b6adc7215743a98a3229958bbd1c9c42b6108b",
+  "model.safetensors=sha256:be8e6242388d93b51cd9719a0e32cfc17a2e804786570c795ba332197eccb915",
+  "special_tokens_map.json=git-sha1:a2d8d626c31389a935e197fb94072e2414a6e7d1",
+  "tokenizer.json=git-sha1:0734d752d12d0f46ac96467fbceb1c4bfbeee0be",
+  "tokenizer_config.json=git-sha1:db0b80de72d3b16242b9eda74ed4663e39c65bcf",
+]
+official_repo = "ElnaggarLab/ankh2-ext2"
+official_revision = "aa9b9fa72288c47d9f618ce80c011e24b54e17a8"
+official_files = [
+  "config.json=git-sha1:9286bed4ecbc4f7113024919d16ec9719b0c0748",
+  "generation_config.json=git-sha1:91f792e452403d46e170e206f9e50be5ddef9b9a",
+  "pytorch_model.bin=sha256:2df583f28f111276ee22a7b76007f4297e9a69766d60bccd9c8d7169c06ac606",
+  "special_tokens_map.json=git-sha1:55b145827029ae9672e50d4bb368540daacce791",
+  "tokenizer.json=git-sha1:212c5ef08819fa2463c6289ba4ef7db30e715c0a",
+  "tokenizer_config.json=git-sha1:854e5db75dae8b1e9dd39c5bae80dae5508b3e25",
+]
+
+[[models]]
+id = "ankh3_large"
+family = "ankh"
+size_category = "large"
+generation_contract = "required"
+official_golden = { metadata = "tests/goldens/ankh3_large.json=sha256:2e5bb05b3baa5baa78f61fef7d2a2c669b0da5dbfaf6b50b12abd3e17253a961", tensors = "tests/goldens/ankh3_large.safetensors=sha256:e5c494ac418e0a2fe7bdad1376676d48960d58ec9e044d19bfffccb8c3288513" }
+notes = "ANKH parity covers the official encoder and sequence-to-sequence heads. AutoModelForMaskedLM exposes the separately named FastPLMs synthesized masked-LM extension and is not an official ANKH head."
+artifact_source = "official"
+canonical_state_sha256 = "60acb7ef86e85dc0c51fc1edf4c8e69a0480049723b6b2c95e6e9faa720c112a"
+fast_repo = "Synthyra/ANKH3_large"
+fast_revision = "53600f175f328f986f43e55ca8ceb14935d337a4"
+fast_files = [
+  "config.json=git-sha1:432b09625d44a2eeab679fddb7495d42b560b7f9",
+  "model.safetensors=sha256:9f50f58cf5b3a537a0a41aa918695c3a26d7985dd0b2266642d6f86324c9e7a1",
+  "special_tokens_map.json=git-sha1:1fc3a4d6d4282e5201cd7c30d5c0a6a8bfa04f82",
+  "tokenizer.json=git-sha1:3d14291df2d6db3a183c5c4fe133afb330cc44cf",
+  "tokenizer_config.json=git-sha1:2005fec00a7ae9a49e248a1ecefbbd81c56674d6",
+]
+official_repo = "ElnaggarLab/ankh3-large"
+official_revision = "2be091622e8a393f0ef21735070084123c874b6e"
+official_files = [
+  "config.json=git-sha1:f5278f77d158cdd8a173df888e3ed365e84a80a3",
+  "generation_config.json=git-sha1:5767cc0cacebfd06884eb27ae1c796d3ca829fd2",
+  "pytorch_model.bin=sha256:26321a345e07a25b21c6c41b651c4db91b420892e52c0dcbc55bd7a8f510f95b",
+  "special_tokens_map.json=git-sha1:d596919b7fa2a197edd441ec3ec4685ecacd2de4",
+  "spiece.model=sha256:f2b5e1bbd110b71ca9b2878e1fcd3265610076ecc97bd696e8a745c9bacc54e0",
+  "tokenizer.json=git-sha1:90f0c94b43c81496b3ca81e3ec1c092ef2dd7fca",
+  "tokenizer_config.json=git-sha1:0e699eebfa778698473b4faf1e66ef363b93fb21",
+]
+
+[[models]]
+id = "ankh3_xl"
+family = "ankh"
+size_category = "xlarge"
+generation_contract = "required"
+official_golden = { metadata = "tests/goldens/ankh3_xl.json=sha256:66bb12e033e4163be225d636108a479393228a4f5061015c8af114e766c3c486", tensors = "tests/goldens/ankh3_xl.safetensors=sha256:72d34567d0228cb6f1ee701c578ed4039fead4346e3f161a52e0e74df28dc8ae" }
+notes = "ANKH parity covers the official encoder and sequence-to-sequence heads. AutoModelForMaskedLM exposes the separately named FastPLMs synthesized masked-LM extension and is not an official ANKH head. The official PyTorch shard index is deliberately excluded: the builder verifies every declared source shard directly and writes a new canonical safetensors index."
+artifact_source = "official"
+canonical_state_sha256 = "dd2188e0d2ca65232135714eef6de394239734d843ddae4928c7398685d858e7"
+fast_repo = "Synthyra/ANKH3_xl"
+fast_revision = "3cbf2c22c4f7d67bf0bfcbdcd500f41723e91d29"
+fast_files = [
+  "config.json=git-sha1:23f6d78ddcb3a031b88f876eaaf04c2fafaea46f",
+  "model-00001-of-00003.safetensors=sha256:39bd8f75cf98a67cf04055399f9fc401198f6fc2896b112aba9fd9ec9df52ab9",
+  "model-00002-of-00003.safetensors=sha256:9ff73233b39d2c200abb78e66b320c014ec61431bd6e1af36fb188a3cfa24c34",
+  "model-00003-of-00003.safetensors=sha256:c13125c02dbcd7f07bd412e9e085f2bca6624d2f1f45fedc95fb777f53161cbe",
+  "special_tokens_map.json=git-sha1:1fc3a4d6d4282e5201cd7c30d5c0a6a8bfa04f82",
+  "tokenizer.json=git-sha1:3d14291df2d6db3a183c5c4fe133afb330cc44cf",
+  "tokenizer_config.json=git-sha1:2005fec00a7ae9a49e248a1ecefbbd81c56674d6",
+]
+official_repo = "ElnaggarLab/ankh3-xl"
+official_revision = "e00113df5c95ef71df7ea3f5a73d56bd00e473a4"
+official_files = [
+  "config.json=git-sha1:f8997040e8913df75fd2eebe71a2a8eb750ed0d0",
+  "generation_config.json=git-sha1:91f792e452403d46e170e206f9e50be5ddef9b9a",
+  "pytorch_model-00001-of-00003.bin=sha256:2c9793cbee16697cd4149debe07d3a27143e280f6e970fa46042aae820fea981",
+  "pytorch_model-00002-of-00003.bin=sha256:31c5a860e414513c829ae52affb0970d7cef2c0545df2d6e1338b6806ab7174b",
+  "pytorch_model-00003-of-00003.bin=sha256:055a853bdd3623db95a637935aa299427e837cd8ea69fc04708b0262508bec75",
+  "special_tokens_map.json=git-sha1:d596919b7fa2a197edd441ec3ec4685ecacd2de4",
+  "spiece.model=sha256:f2b5e1bbd110b71ca9b2878e1fcd3265610076ecc97bd696e8a745c9bacc54e0",
+  "tokenizer.json=git-sha1:90f0c94b43c81496b3ca81e3ec1c092ef2dd7fca",
+  "tokenizer_config.json=git-sha1:0e699eebfa778698473b4faf1e66ef363b93fb21",
+]
+
+[[models]]
+id = "boltz2"
+family = "boltz2"
+size_category = "structure"
+generation_contract = "not_applicable"
+notes = "Boltz2 is provisional in FastPLMs 1.0. Exact configuration, the declared inference-core state, feature preparation, and seeded execution remain tested, but native-environment BF16 end-to-end inference currently exceeds the fixed numerical-equivalence limits. FastPLMs therefore does not claim official inference equivalence for this checkpoint yet. Work on that numerical gap continues independently of the ESM++ and ESMFold2 release gates."
+fast_repo = "Synthyra/Boltz2"
+fast_revision = "3b148fc5efea109c065ec82ba8683d024de7134e"
+fast_files = [
+  "config.json=git-sha1:8682ccb12e177e73bc7a351ff7e3af484bfb6fac",
+  "model.safetensors=sha256:5c863fd200a1613a0e311071e2ad73ab350635e3fd336e6822cf45c52cb960e5",
+]
+official_repo = "boltz-community/boltz-2"
+official_revision = "6fdef46d763fee7fbb83ca5501ccceff43b85607"
+official_files = [
+  "boltz2_conf.ckpt=sha256:090e82ac8c92f5e943fa1b39e7410a44027bea7243c0bbb3caa67a77fc1428e1",
+  "mols.tar=sha256:39e076d96dbec6b4e86982bbda16f3a53a2a60c9bdc17828d88f6f9a0c7d1fd7",
+]
+
+[[models]]
+id = "esmfold"
+family = "esmfold"
+size_category = "structure"
+generation_contract = "not_applicable"
+official_golden = { metadata = "tests/goldens/esmfold.json=sha256:380b9a96168410717d1f698feaabb826b1606444cbdeec86c2ea06d9ffe8f186", tensors = "tests/goldens/esmfold.safetensors=sha256:873b1b325a43d8e0f35f355c8914a2a9fe611cc48763875e9e6a22e09ec9ebcb" }
+fast_repo = "Synthyra/FastESMFold"
+fast_revision = "b88c8cb50d19b2cf7ab4fee4b0a61f5e02da7823"
+fast_files = [
+  "config.json=git-sha1:18e0091dcbf6140bf68924d53c4c8917b9cd90b1",
+  "model-00001-of-00003.safetensors=sha256:36fab9e5c96d409b2a34a8b4f1273acac8c07f119c32c4fcfa7d47bbcd55b83c",
+  "model-00002-of-00003.safetensors=sha256:34954aaa05bc91635776ba6672946da5822626753d80db97b38c0538e9525102",
+  "model-00003-of-00003.safetensors=sha256:2f1178cda0e6cff3b1e158e1acc59c83e3f4fc46e246388a5127bc56b8d9c4f2",
+  "special_tokens_map.json=git-sha1:53cd95604a28eb7e23da763c8da23f5006ab2179",
+  "tokenizer_config.json=git-sha1:10213f69b51b4b38876a29271b8f908e853a5800",
+  "vocab.txt=git-sha1:eee0a1fc93c82568f78f086550fbd7c591cf423a",
+]
+official_repo = "facebook/esmfold_v1"
+official_revision = "75a3841ee059df2bf4d56688166c8fb459ddd97a"
+official_files = [
+  "config.json=git-sha1:1232d0aee4be551021d8e70e66ed2b062df917bf",
+  "pytorch_model.bin=sha256:2ee07356b125d1e3e57503c204111fd7323347fc4735d41d3caac57c2a78e116",
+  "special_tokens_map.json=git-sha1:121c8d54f8ea66cdf678f48b3cb37c05b4de5c0d",
+  "tokenizer_config.json=git-sha1:aad24fba9f1bad2d74ed79d414ddcd60e6b0f812",
+  "vocab.txt=git-sha1:9abfdf5472c0ed970648b683b86ab131256b3e42",
+]
+
+[[models.oracle_assets]]
+role = "weights"
+path = "models/esmfold_3B_v1.pt"
+url = "https://dl.fbaipublicfiles.com/fair-esm/models/esmfold_3B_v1.pt"
+sha256 = "e9a52579027e77d2d2e0a18218e755821f395730e86624cab9413dc117f5ca62"
+size = 2771653574
+
+[[models]]
+id = "esmfold2"
+family = "esmfold2"
+size_category = "structure"
+generation_contract = "not_applicable"
+msa_conditioning = true
+official_golden = { metadata = "tests/goldens/esmfold2.json=sha256:f6e0ed1ec400b9a0fcc817db51774be968dc454b7a32645a07c479e42423ab20", tensors = "tests/goldens/esmfold2.safetensors=sha256:e4d6be4344c528e26b13f79a9303549e3de7e582da195c0078db3ce957fad420" }
+fast_repo = "Synthyra/ESMFold2"
+fast_revision = "cd5a0927cec585a778d983b99a8db23d2e9b281e"
+fast_files = [
+  "config.json=git-sha1:67e81ff571f393f0b630cd5a22398bd84979c030",
+  "model.safetensors=sha256:138fd4350d6892b81ce6be7ff9bf5a93ae9d4d3751f46a27438a3f9f0dcefa0e",
+]
+official_repo = "biohub/ESMFold2"
+official_revision = "1ebf0e3481a5184eb6171d40615c79e384b48796"
+official_files = [
+  "config.json=git-sha1:0300c084b990b2bd600efd9f538aa5de27109fea",
+  "model.safetensors=sha256:138fd4350d6892b81ce6be7ff9bf5a93ae9d4d3751f46a27438a3f9f0dcefa0e",
+]
+
+[[models]]
+id = "esmfold2_fast"
+family = "esmfold2"
+size_category = "structure"
+generation_contract = "not_applicable"
+msa_conditioning = false
+official_golden = { metadata = "tests/goldens/esmfold2_fast.json=sha256:091b004c0b330217b59c12acd6da3d6edaf91e48d95f6d5f40fc20399cef9478", tensors = "tests/goldens/esmfold2_fast.safetensors=sha256:6e2e1cd07401538b4d9df994f82abe7a5b38a01e8d1ee26681e1216d44a81990" }
+fast_repo = "Synthyra/ESMFold2-Fast"
+fast_revision = "407875bfcaa42552bfcb25acd67ee1888b790170"
+fast_files = [
+  "config.json=git-sha1:62ccca15a416a5dcbd02cd6ce161f432c7b4de58",
+  "model.safetensors=sha256:60ca19f2898188beba92944365f7b909efd9c99212f5018af75cc47cd9a6184a",
+]
+official_repo = "biohub/ESMFold2-Fast"
+official_revision = "b28d8ace5e05e61e5bec1e6820cfd3e221819d12"
+official_files = [
+  "config.json=git-sha1:c0ca526090fa7f8342ee4666d56e7fe3a4b8cbb2",
+  "model.safetensors=sha256:60ca19f2898188beba92944365f7b909efd9c99212f5018af75cc47cd9a6184a",
+]
+
+[[models]]
+id = "esmfold2_experimental_cutoff2025"
+family = "esmfold2"
+size_category = "structure"
+generation_contract = "not_applicable"
+msa_conditioning = true
+official_golden = { metadata = "tests/goldens/esmfold2_experimental_cutoff2025.json=sha256:cfd0e35b2bc468a0dc4f614d3acfa2fce004f96e9ae2433256ed095b829d55cc", tensors = "tests/goldens/esmfold2_experimental_cutoff2025.safetensors=sha256:9347466bbe803b6f5dc82e3356ca6cbbf2c2edd8765f9fd273385bda255019f6" }
+fast_repo = "Synthyra/ESMFold2-Experimental-Cutoff2025"
+fast_revision = "632ff4a9e68f1de78ee956a613267bdcdb5b354d"
+fast_files = [
+  "config.json=git-sha1:41119745d38bc5503a0212ad923e75211dec565f",
+  "model.safetensors=sha256:01358c317428d38535e3db513cab177336fc0f7fab0d84002e64b7741d5181b3",
+]
+official_repo = "biohub/ESMFold2-Experimental-Cutoff2025"
+official_revision = "56f94f5c1069ecde17512c96928850518340d287"
+official_files = [
+  "config.json=git-sha1:79ed0dc0f867b8f09bfa004d6f77397c2ab9b38d",
+  "model.safetensors=sha256:01358c317428d38535e3db513cab177336fc0f7fab0d84002e64b7741d5181b3",
+]
+auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2_experimental.ESMFold2ExperimentalModel" }
+
+[[models]]
+id = "esmfold2_experimental_fast_cutoff2025"
+family = "esmfold2"
+size_category = "structure"
+generation_contract = "not_applicable"
+msa_conditioning = false
+official_golden = { metadata = "tests/goldens/esmfold2_experimental_fast_cutoff2025.json=sha256:1d0b2da4f1579243f37ae04bd4b834b747005cd8e8e7665e00d088123c43afd9", tensors = "tests/goldens/esmfold2_experimental_fast_cutoff2025.safetensors=sha256:516e216d05d7e6bee59e77126d3e595e2bb7821929433f00c259c5d5241964bb" }
+fast_repo = "Synthyra/ESMFold2-Experimental-Fast-Cutoff2025"
+fast_revision = "8f022c2514a6c32692aaca078a8391d6bc6c4bac"
+fast_files = [
+  "config.json=git-sha1:b9d39e941050179ca51faaed58cbbd77778c1143",
+  "model.safetensors=sha256:4e903b740ad6ad704ec60881bfd593e0d6c874a630ffa0f0838276e0b665088f",
+]
+official_repo = "biohub/ESMFold2-Experimental-Fast-Cutoff2025"
+official_revision = "74b88548bf19688b8727432db0d698cb2e1d8783"
+official_files = [
+  "config.json=git-sha1:0333d68ddb12ed2f066741dcb801142f466c0a2c",
+  "model.safetensors=sha256:4e903b740ad6ad704ec60881bfd593e0d6c874a630ffa0f0838276e0b665088f",
+]
+auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2_experimental.ESMFold2ExperimentalModel" }

fastplms/models/__init__.py

@@ -0,0 +1,10 @@
+"""Lazy model-family namespace for FastPLMs.
+
+Model classes are resolved through Transformers AutoClasses and the typed
+registry. Importing this package therefore does not load checkpoints, create
+tokenizers, compile kernels, or initialize an accelerator runtime.
+"""
+
+from __future__ import annotations
+
+__all__: tuple[str, ...] = ()

fastplms/models/_diffusion_generation.py

@@ -0,0 +1,510 @@
+"""Discrete diffusion generation shared by DPLM and DPLM2.
+
+The implementation keeps model-specific vocabulary rules at the public entry
+points and shares only the categorical sampling and confidence-based remasking
+mechanism.  It has no dependency on the pinned upstream checkout.
+"""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Iterable, Mapping
+from contextlib import contextmanager
+from typing import Any, Protocol
+
+import torch
+from tqdm.auto import tqdm
+
+
+class _MaskedLanguageModel(Protocol):
+    """Structural type used by the two generation entry points."""
+
+    config: Any
+
+    def eval(self) -> Any: ...
+
+    def modules(self) -> Iterable[torch.nn.Module]: ...
+
+    def __call__(self, **kwargs: Any) -> Any: ...
+
+
+_DPLM2_AA_BOUNDARY = 33
+_DPLM2_AA_BOS = 0
+_DPLM2_PAD = 1
+_DPLM2_AA_EOS = 2
+_DPLM2_AA_UNK = 3
+_DPLM2_AA_X = 24
+_DPLM2_AA_B = 25
+_DPLM2_AA_U = 26
+_DPLM2_AA_Z = 27
+_DPLM2_AA_O = 28
+_DPLM2_AA_MASK = 32
+_DPLM2_STRUCT_BOS = 33
+_DPLM2_STRUCT_EOS = 34
+_DPLM2_STRUCT_UNK = 35
+
+
+@contextmanager
+def _temporary_eval(model: _MaskedLanguageModel):
+    """Run one generation forward in eval mode and restore every module flag."""
+    training_states = tuple((module, module.training) for module in model.modules())
+    model.eval()
+    try:
+        yield
+    finally:
+        for module, training in training_states:
+            module.training = training
+
+
+def _resolve_max_iter(model: _MaskedLanguageModel, max_iter: int | None) -> int:
+    if max_iter is None:
+        max_iter = int(getattr(model.config, "num_diffusion_timesteps", 500))
+    if isinstance(max_iter, bool) or not isinstance(max_iter, int) or max_iter <= 0:
+        raise ValueError("max_iter must be a positive integer")
+    return max_iter
+
+
+def _validate_inputs(
+    input_tokens: torch.Tensor,
+    partial_masks: torch.Tensor | None,
+) -> torch.Tensor | None:
+    if input_tokens.ndim != 2 or input_tokens.dtype not in {
+        torch.int8,
+        torch.int16,
+        torch.int32,
+        torch.int64,
+        torch.uint8,
+    }:
+        raise ValueError("input_tokens must be an integer tensor with shape (b, l)")
+    if input_tokens.shape[-1] == 0:
+        raise ValueError("input_tokens must contain at least one token")
+    if partial_masks is None:
+        return None
+    if partial_masks.shape != input_tokens.shape or partial_masks.dtype != torch.bool:
+        raise ValueError("partial_masks must be boolean with the same shape as input_tokens")
+    if partial_masks.device != input_tokens.device:
+        raise ValueError("partial_masks and input_tokens must be on the same device")
+    return partial_masks
+
+
+def _validate_temperature(temperature: float | None, *, default: float = 1.0) -> float:
+    if temperature is None:
+        temperature = default
+    temperature = float(temperature)
+    if not math.isfinite(temperature) or temperature < 0:
+        raise ValueError("temperature must be finite and non-negative")
+    return temperature
+
+
+def _steps(max_iter: int, *, show_progress: bool) -> Iterable[int]:
+    return tqdm(range(max_iter), desc="Decoding", disable=not show_progress)
+
+
+def _categorical(
+    logits: torch.Tensor,
+    *,
+    temperature: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    if temperature == 0:
+        scores, tokens = logits.log_softmax(dim=-1).max(dim=-1)
+        return tokens, scores
+    distribution = torch.distributions.Categorical(logits=logits.div(temperature))
+    tokens = distribution.sample()
+    return tokens, distribution.log_prob(tokens)
+
+
+def _gumbel_argmax(
+    logits: torch.Tensor,
+    *,
+    noise_scale: float,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    uniform = torch.rand_like(logits)
+    noise = -torch.log(-torch.log(uniform + 1e-8) + 1e-8)
+    return _categorical(logits + noise_scale * noise, temperature=0.0)
+
+
+def _top_p(logits: torch.Tensor, probability: float = 0.95) -> torch.Tensor:
+    """Apply the nucleus filter used by the official DPLM samplers."""
+
+    original_shape = logits.shape
+    flattened = logits.reshape(-1, original_shape[-1])
+    sorted_logits, sorted_indices = flattened.sort(dim=-1, descending=True)
+    cumulative = sorted_logits.softmax(dim=-1).cumsum(dim=-1)
+    remove = cumulative > probability
+    remove[..., 1:] = remove[..., :-1].clone()
+    remove[..., 0] = False
+    sorted_logits.masked_fill_(remove, -math.inf)
+    return sorted_logits.gather(1, sorted_indices.argsort(dim=-1)).reshape(original_shape)
+
+
+def _lowest_confidence_mask(
+    scores: torch.Tensor,
+    eligible: torch.Tensor,
+    *,
+    rate: float,
+    stochastic_temperature: float | None = None,
+) -> torch.Tensor:
+    selection_scores = scores.masked_fill(~eligible, 1000.0)
+    if stochastic_temperature is not None:
+        uniform = torch.rand_like(selection_scores)
+        noise = -torch.log(-torch.log(uniform + 1e-8) + 1e-8)
+        selection_scores = selection_scores + stochastic_temperature * rate * noise
+    cutoff_index = (eligible.sum(dim=-1, keepdim=True).to(scores.dtype) * rate).long()
+    cutoff_index.clamp_(min=0, max=scores.shape[-1] - 1)
+    sorted_scores = selection_scores.sort(dim=-1).values
+    cutoff = sorted_scores.gather(dim=-1, index=cutoff_index)
+    return (selection_scores < cutoff) & eligible
+
+
+def _reparameterize(
+    output_tokens: torch.Tensor,
+    output_scores: torch.Tensor,
+    candidate_tokens: torch.Tensor,
+    candidate_scores: torch.Tensor,
+    active_mask: torch.Tensor,
+    eligible: torch.Tensor,
+    *,
+    mask_token_id: int,
+    rate: float,
+    stochastic_temperature: float | None = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    remask = _lowest_confidence_mask(
+        candidate_scores,
+        eligible,
+        rate=rate,
+        stochastic_temperature=stochastic_temperature,
+    )
+    output_tokens.masked_fill_(remask, mask_token_id)
+    output_scores.masked_fill_(remask, -math.inf)
+    accept = active_mask & eligible & ~remask
+    output_tokens.masked_scatter_(accept, candidate_tokens[accept])
+    output_scores.masked_scatter_(accept, candidate_scores[accept])
+    return remask, output_tokens, output_scores
+
+
+def _logits(output: object) -> torch.Tensor:
+    value = output.get("logits") if isinstance(output, Mapping) else getattr(output, "logits", None)
+    if not torch.is_tensor(value):
+        raise RuntimeError("The masked-language model did not return logits")
+    return value
+
+
+def _suppress_token_ids(logits: torch.Tensor, token_ids: Iterable[int]) -> None:
+    vocabulary_size = logits.shape[-1]
+    for token_id in token_ids:
+        if 0 <= token_id < vocabulary_size:
+            logits[..., token_id] = -math.inf
+
+
+def _dplm_special_id(
+    model: _MaskedLanguageModel,
+    tokenizer: object | None,
+    name: str,
+    default: int,
+) -> int:
+    value = getattr(model.config, name, None)
+    if value is None:
+        if tokenizer is None:
+            tokenizer = getattr(model, "tokenizer", None)
+        value = getattr(tokenizer, name, None)
+    return default if value is None else int(value)
+
+
+def _dplm_resample_repeats(
+    model: _MaskedLanguageModel,
+    candidate_tokens: torch.Tensor,
+    candidate_scores: torch.Tensor,
+    *,
+    invalid_token_ids: tuple[int, ...],
+    mask_token_id: int,
+    ratio: float,
+) -> None:
+    selected_rows: list[int] = []
+    resample_tokens: list[torch.Tensor] = []
+    resample_scores: list[torch.Tensor] = []
+    resample_masks: list[torch.Tensor] = []
+    for row_index, row in enumerate(candidate_tokens):
+        positions: dict[int, list[int]] = {}
+        for position, token in enumerate(row.tolist()):
+            positions.setdefault(int(token), []).append(position)
+        repeated = [indices for indices in positions.values() if len(indices) > row.numel() * ratio]
+        if not repeated:
+            continue
+        M = torch.zeros_like(row, dtype=torch.bool)
+        for indices in repeated:
+            M[indices] = True
+        selected_rows.append(row_index)
+        resample_masks.append(M)
+        resample_tokens.append(row.masked_fill(M, mask_token_id))
+        resample_scores.append(candidate_scores[row_index])
+
+    if not selected_rows:
+        return
+    X = torch.stack(resample_tokens)
+    S = torch.stack(resample_scores)
+    M = torch.stack(resample_masks)
+    with _temporary_eval(model), torch.no_grad():
+        logits = _logits(model(input_ids=X, return_dict=True))
+    if logits.dtype != S.dtype:
+        logits = logits.to(S.dtype)
+    _suppress_token_ids(logits, invalid_token_ids)
+    logits = _top_p(logits)
+    sampled_tokens, sampled_scores = _gumbel_argmax(logits, noise_scale=1.0)
+    X.masked_scatter_(M, sampled_tokens[M])
+    S.masked_scatter_(M, sampled_scores[M])
+    candidate_tokens[selected_rows] = X
+    candidate_scores[selected_rows] = S
+
+
+def generate_dplm(
+    model: _MaskedLanguageModel,
+    input_tokens: torch.Tensor,
+    *,
+    tokenizer: object | None = None,
+    max_iter: int | None = None,
+    temperature: float | None = None,
+    partial_masks: torch.Tensor | None = None,
+    sampling_strategy: str = "gumbel_argmax",
+    disable_resample: bool = False,
+    resample_ratio: float = 0.25,
+    show_progress: bool = False,
+) -> torch.Tensor:
+    """Generate DPLM sequences with the official iterative unmasking process.
+
+    ``input_tokens`` is X with shape (b, l). ``partial_masks=True`` marks fixed
+    positions. The return value is the generated token tensor X with shape
+    (b, l), matching the official DPLM public API.
+    """
+
+    partial_masks = _validate_inputs(input_tokens, partial_masks)
+    max_iter = _resolve_max_iter(model, max_iter)
+    # Upstream treats ``None`` as the falsey, zero-temperature branch for
+    # vanilla categorical sampling. Gumbel and argmax strategies ignore it.
+    temperature = _validate_temperature(temperature, default=0.0)
+    if sampling_strategy not in {"vanilla", "argmax", "gumbel_argmax"}:
+        raise ValueError(f"Unsupported DPLM sampling strategy: {sampling_strategy!r}")
+    if not 0 < float(resample_ratio) <= 1:
+        raise ValueError("resample_ratio must be in (0, 1]")
+
+    pad_id = _dplm_special_id(model, tokenizer, "pad_token_id", 1)
+    bos_id = _dplm_special_id(model, tokenizer, "bos_token_id", 0)
+    eos_id = _dplm_special_id(model, tokenizer, "eos_token_id", 2)
+    mask_id = _dplm_special_id(model, tokenizer, "mask_token_id", 32)
+    x_id = 24
+    X = input_tokens.clone()
+    mutable = X.ne(pad_id) & X.ne(bos_id) & X.ne(eos_id)
+    if partial_masks is not None:
+        mutable &= ~partial_masks
+    X.masked_fill_(mutable, mask_id)
+    S = torch.zeros_like(X, dtype=torch.float32)
+    active = mutable.clone()
+    invalid_ids = (mask_id, x_id, pad_id, bos_id, eos_id)
+    for step in _steps(max_iter, show_progress=show_progress):
+        with _temporary_eval(model), torch.no_grad():
+            logits = _logits(model(input_ids=X, return_dict=True))
+        if logits.dtype != S.dtype:
+            logits = logits.to(S.dtype)
+        _suppress_token_ids(logits, invalid_ids)
+        if sampling_strategy == "vanilla":
+            candidate_tokens, candidate_scores = _categorical(
+                logits,
+                temperature=temperature,
+            )
+        elif sampling_strategy == "argmax":
+            candidate_scores, candidate_tokens = logits.max(dim=-1)
+        else:
+            candidate_tokens, candidate_scores = _gumbel_argmax(logits, noise_scale=1.0)
+            if not disable_resample:
+                _dplm_resample_repeats(
+                    model,
+                    candidate_tokens,
+                    candidate_scores,
+                    invalid_token_ids=invalid_ids,
+                    mask_token_id=mask_id,
+                    ratio=float(resample_ratio),
+                )
+
+        eligible = X.ne(pad_id) & X.ne(bos_id) & X.ne(eos_id)
+        if partial_masks is not None:
+            eligible &= ~partial_masks
+        rate = 1.0 - (step + 1) / max_iter
+        active, X, S = _reparameterize(
+            X.clone(),
+            S.clone(),
+            candidate_tokens,
+            candidate_scores,
+            active,
+            eligible,
+            mask_token_id=mask_id,
+            rate=rate,
+        )
+    return X
+
+
+def _normalize_dplm2_special_ids(X: torch.Tensor, vocabulary_size: int) -> torch.Tensor:
+    normalized = X.clone()
+    replacements = {
+        vocabulary_size: _DPLM2_AA_EOS,
+        vocabulary_size + 1: _DPLM2_AA_UNK,
+        vocabulary_size + 2: _DPLM2_AA_BOS,
+        vocabulary_size + 3: _DPLM2_AA_MASK,
+    }
+    for generic_id, native_id in replacements.items():
+        normalized.masked_fill_(X.eq(generic_id), native_id)
+    return normalized
+
+
+def _dplm2_types(X: torch.Tensor) -> torch.Tensor:
+    valid = X.ne(_DPLM2_PAD)
+    types = ((X < _DPLM2_AA_BOUNDARY) & valid).to(torch.int64)
+    types.masked_fill_(~valid, 2)
+    return types
+
+
+def _dplm2_mutable(X: torch.Tensor, partial_masks: torch.Tensor | None) -> torch.Tensor:
+    mutable = (
+        X.ne(_DPLM2_PAD)
+        & X.ne(_DPLM2_AA_BOS)
+        & X.ne(_DPLM2_AA_EOS)
+        & X.ne(_DPLM2_STRUCT_BOS)
+        & X.ne(_DPLM2_STRUCT_EOS)
+    )
+    if partial_masks is not None:
+        mutable &= ~partial_masks
+    return mutable
+
+
+def _dplm2_unmasking_temperature(strategy: str) -> float | None:
+    if strategy == "deterministic":
+        return None
+    if strategy.startswith("stochastic"):
+        suffix = strategy.removeprefix("stochastic")
+        value = 1.0 if not suffix else float(suffix)
+        if not math.isfinite(value) or value < 0:
+            raise ValueError("The stochastic unmasking temperature must be non-negative")
+        return value
+    raise ValueError(f"Unsupported DPLM2 unmasking strategy: {strategy!r}")
+
+
+def _annealing_temperature(strategy: str, step: int, max_iter: int) -> float | None:
+    if not strategy.startswith("annealing"):
+        return None
+    try:
+        maximum, minimum = map(float, strategy.split("@", maxsplit=1)[1].split(":"))
+    except (IndexError, ValueError) as error:
+        raise ValueError("Annealing must use the form 'annealing@maximum:minimum'") from error
+    if not all(math.isfinite(value) and value >= 0 for value in (maximum, minimum)):
+        raise ValueError("Annealing temperatures must be finite and non-negative")
+    rate = 1.0 - step / max_iter
+    return minimum + (maximum - minimum) * rate
+
+
+def generate_dplm2(
+    model: _MaskedLanguageModel,
+    input_tokens: torch.Tensor,
+    *,
+    max_iter: int | None = None,
+    temperature: float = 1.0,
+    partial_masks: torch.Tensor | None = None,
+    unmasking_strategy: str = "stochastic1.0",
+    sampling_strategy: str = "annealing@2.0:0.1",
+    show_progress: bool = False,
+) -> dict[str, torch.Tensor]:
+    """Generate packed DPLM2 sequence and structure tracks.
+
+    ``input_tokens`` is X with shape (b, l). A packed co-generation input has
+    two equal-length modality tracks. ``partial_masks=True`` marks fixed
+    positions. The output mapping matches the official DPLM2 public API.
+    """
+
+    partial_masks = _validate_inputs(input_tokens, partial_masks)
+    max_iter = _resolve_max_iter(model, max_iter)
+    temperature = _validate_temperature(temperature)
+    unmasking_temperature = _dplm2_unmasking_temperature(unmasking_strategy)
+    if sampling_strategy.startswith("annealing"):
+        _annealing_temperature(sampling_strategy, 0, max_iter)
+    elif sampling_strategy not in {"argmax", "gumbel_argmax"}:
+        raise ValueError(f"Unsupported DPLM2 sampling strategy: {sampling_strategy!r}")
+    vocabulary_size = int(model.config.vocab_size)
+    if vocabulary_size <= _DPLM2_STRUCT_UNK + 1:
+        raise ValueError("DPLM2 generation requires the multimodal vocabulary")
+    struct_mask_id = vocabulary_size - 1
+
+    X = _normalize_dplm2_special_ids(input_tokens, vocabulary_size)
+    if X.numel() and (X.min() < 0 or X.max() >= vocabulary_size):
+        raise ValueError("input_tokens contains an ID outside the DPLM2 vocabulary")
+    mutable = _dplm2_mutable(X, partial_masks)
+    types = _dplm2_types(X)
+    X.masked_fill_(mutable & types.eq(1), _DPLM2_AA_MASK)
+    X.masked_fill_(mutable & types.eq(0), struct_mask_id)
+    S = torch.zeros_like(X, dtype=torch.float32)
+    active = mutable.clone()
+    invalid_ids = (
+        _DPLM2_AA_BOS,
+        _DPLM2_AA_EOS,
+        _DPLM2_AA_MASK,
+        _DPLM2_STRUCT_BOS,
+        _DPLM2_STRUCT_EOS,
+        struct_mask_id,
+        _DPLM2_PAD,
+        _DPLM2_AA_UNK,
+        _DPLM2_STRUCT_UNK,
+        _DPLM2_AA_X,
+        _DPLM2_AA_B,
+        _DPLM2_AA_U,
+        _DPLM2_AA_Z,
+        _DPLM2_AA_O,
+    )
+    for step in _steps(max_iter, show_progress=show_progress):
+        eligible = _dplm2_mutable(X, partial_masks)
+        types = _dplm2_types(X)
+        with _temporary_eval(model), torch.no_grad():
+            logits = _logits(model(input_ids=X, return_dict=True)).log_softmax(dim=-1)
+        if logits.dtype != S.dtype:
+            logits = logits.to(S.dtype)
+        aa_rows, aa_columns = torch.where(types.eq(1) & eligible)
+        struct_rows, struct_columns = torch.where(types.eq(0) & eligible)
+        logits[aa_rows, aa_columns, _DPLM2_AA_BOUNDARY:] = -math.inf
+        logits[struct_rows, struct_columns, :_DPLM2_AA_BOUNDARY] = -math.inf
+        _suppress_token_ids(logits, invalid_ids)
+        logits = _top_p(logits)
+
+        if sampling_strategy == "argmax":
+            candidate_scores, candidate_tokens = logits.max(dim=-1)
+        elif sampling_strategy == "gumbel_argmax":
+            candidate_tokens, candidate_scores = _gumbel_argmax(
+                logits,
+                noise_scale=temperature,
+            )
+            candidate_tokens.masked_scatter_(~eligible, X[~eligible])
+        else:
+            annealed = _annealing_temperature(sampling_strategy, step, max_iter)
+            sample_temperature = temperature if annealed is None else annealed
+            candidate_tokens, candidate_scores = _categorical(
+                logits,
+                temperature=sample_temperature,
+            )
+
+        rate = 1.0 - (step + 1) / max_iter
+        new_active = torch.zeros_like(active)
+        for modality, mask_id in ((1, _DPLM2_AA_MASK), (0, struct_mask_id)):
+            modality_positions = types.eq(modality) & eligible
+            if not bool(modality_positions.any()):
+                continue
+            modality_active, X, S = _reparameterize(
+                X,
+                S,
+                candidate_tokens,
+                candidate_scores,
+                active,
+                modality_positions,
+                mask_token_id=mask_id,
+                rate=rate,
+                stochastic_temperature=unmasking_temperature,
+            )
+            new_active |= modality_active
+        active = new_active
+    return {"output_tokens": X}
+
+
+__all__ = ["generate_dplm", "generate_dplm2"]

fastplms/models/_esm_rotary.py

@@ -0,0 +1,81 @@
+"""Stable ESM rotary embeddings independent of Transformers internals.
+
+Transformers 5 changed both the name and call contract of its private ESM
+rotary helper. FastPLMs checkpoints use the earlier two-tensor contract, so
+the small mathematical primitive lives here instead of importing a private
+Transformers implementation.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+
+def _rotate_half(tensor: torch.Tensor) -> torch.Tensor:
+    """Rotate the final dimension of X by 90 degrees in paired subspaces."""
+
+    first, second = tensor.chunk(2, dim=-1)
+    return torch.cat((-second, first), dim=-1)
+
+
+def apply_rotary_pos_emb(
+    tensor: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+) -> torch.Tensor:
+    """Apply cached rotary factors to X with shape ``(b, h, l, d)``."""
+
+    cos = cos[:, :, : tensor.shape[-2], :]
+    sin = sin[:, :, : tensor.shape[-2], :]
+    return tensor * cos + _rotate_half(tensor) * sin
+
+
+class RotaryEmbedding(nn.Module):
+    """Apply rotary position embeddings to query and key tensors."""
+
+    inv_freq: torch.Tensor
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        frequencies = 1.0 / (10_000 ** (torch.arange(0, dim, 2, dtype=torch.int64).float() / dim))
+        # Keep this persistent to preserve the historical checkpoint schema.
+        self.register_buffer("inv_freq", frequencies)
+        self._seq_len_cached: int | None = None
+        self._cos_cached: torch.Tensor | None = None
+        self._sin_cached: torch.Tensor | None = None
+
+    def _update_cos_sin_tables(
+        self,
+        tensor: torch.Tensor,
+        seq_dimension: int = 2,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        seq_len = tensor.shape[seq_dimension]
+        cache_stale = (
+            self._cos_cached is None
+            or self._sin_cached is None
+            or self._seq_len_cached != seq_len
+            or self._cos_cached.device != tensor.device
+        )
+        if cache_stale:
+            self._seq_len_cached = seq_len
+            positions = torch.arange(seq_len, device=tensor.device).type_as(self.inv_freq)
+            angles = torch.outer(positions, self.inv_freq)
+            angles = torch.cat((angles, angles), dim=-1).to(tensor.device)
+            self._cos_cached = angles.cos()[None, None, :, :]
+            self._sin_cached = angles.sin()[None, None, :, :]
+
+        assert self._cos_cached is not None
+        assert self._sin_cached is not None
+        return self._cos_cached, self._sin_cached
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        cos, sin = self._update_cos_sin_tables(key, seq_dimension=-2)
+        return (
+            apply_rotary_pos_emb(query, cos, sin).to(dtype=query.dtype),
+            apply_rotary_pos_emb(key, cos, sin).to(dtype=key.dtype),
+        )

fastplms/models/dplm2/__init__.py

@@ -0,0 +1,5 @@
+"""DPLM2 model and tokenizer implementation."""
+
+from .tokenization_dplm2 import DPLM2Tokenizer
+
+__all__ = ["DPLM2Tokenizer"]

fastplms/models/dplm2/modeling_dplm2.py

@@ -0,0 +1,1457 @@
+"""
+FastPLMs-compatible DPLM2 implementation.
+"""
+
+from __future__ import annotations
+
+import contextlib
+import warnings
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any, ClassVar
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from torch.nn import functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
+from transformers import AutoTokenizer
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPoolingAndCrossAttentions,
+    MaskedLMOutput,
+    ModelOutput,
+    SequenceClassifierOutput,
+    TokenClassifierOutput,
+)
+from transformers.models.esm.configuration_esm import EsmConfig
+from transformers.models.esm.modeling_esm import (
+    EsmAttention,
+    EsmClassificationHead,
+    EsmContactPredictionHead,
+    EsmEmbeddings,
+    EsmEncoder,
+    EsmIntermediate,
+    EsmLayer,
+    EsmLMHead,
+    EsmOutput,
+    EsmPooler,
+    EsmPreTrainedModel,
+    EsmSelfAttention,
+    EsmSelfOutput,
+)
+
+from fastplms.models._diffusion_generation import generate_dplm2
+from fastplms.models._esm_rotary import RotaryEmbedding, apply_rotary_pos_emb
+from fastplms.models.dplm2.tokenization_dplm2 import DPLM2Tokenizer
+
+try:
+    from fastplms.attention import (
+        AttentionBackend,
+        FastPLMsAttentionMixin,
+        get_attention_mask,
+        resolve_attention_backend,
+        resolve_attention_backend_for_call,
+    )
+    from fastplms.embeddings import EmbeddingMixin, select_hidden_state_embeddings
+    from fastplms.models.ttt import FastPLMTestTimeTrainingMixin
+except ModuleNotFoundError as error:
+    _COMPOSITE_REQUIRED_NAMES = (
+        "AttentionBackend",
+        "EmbeddingMixin",
+        "FastPLMsAttentionMixin",
+        "FastPLMTestTimeTrainingMixin",
+        "get_attention_mask",
+        "resolve_attention_backend",
+        "resolve_attention_backend_for_call",
+        "select_hidden_state_embeddings",
+    )
+    if error.name != "fastplms" or any(
+        name not in globals() for name in _COMPOSITE_REQUIRED_NAMES
+    ):
+        raise
+    # Legacy flat Hub composites define every shared symbol above this block.
+
+
+def _infer_modality_type(input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
+    input_mask = attention_mask.bool()
+    modality_type = ((input_ids < 33) & input_mask).int()
+    modality_type[~input_mask] = 2
+    return modality_type
+
+
+def _normalize_dplm2_input_ids(input_ids: torch.Tensor, vocab_size: int) -> torch.Tensor:
+    if input_ids.numel() == 0:
+        return input_ids
+
+    normalized_input_ids = input_ids.clone()
+    generic_to_aa_special_ids = {
+        vocab_size: 2,
+        vocab_size + 1: 3,
+        vocab_size + 2: 0,
+        vocab_size + 3: 32,
+    }
+    for generic_id, aa_id in generic_to_aa_special_ids.items():
+        normalized_input_ids[input_ids == generic_id] = aa_id
+
+    valid_token_mask = normalized_input_ids.ge(0)
+    if valid_token_mask.any():
+        max_token_id = int(normalized_input_ids[valid_token_mask].max().item())
+        if max_token_id >= vocab_size:
+            raise ValueError(
+                f"Found token id {max_token_id} outside the DPLM2 embedding table "
+                f"(vocab_size={vocab_size}). Tokenizer special tokens must be normalized "
+                "before embedding."
+            )
+    return normalized_input_ids
+
+
+def _validate_dplm2_model_inputs(
+    *,
+    input_ids: torch.Tensor | None,
+    inputs_embeds: torch.Tensor | None,
+    attention_mask: torch.Tensor | None,
+    type_ids: torch.Tensor | None,
+    hidden_size: int,
+) -> tuple[int, int]:
+    if (input_ids is None) == (inputs_embeds is None):
+        raise ValueError("Specify exactly one of input_ids or inputs_embeds.")
+
+    if input_ids is not None:
+        if input_ids.ndim != 2:
+            raise ValueError(
+                f"input_ids must have shape (batch, seq_len), got {tuple(input_ids.shape)}."
+            )
+        batch_size, seq_len = input_ids.shape
+    else:
+        if inputs_embeds is None:  # Defensive guard for static narrowing.
+            raise RuntimeError("inputs_embeds validation reached an invalid state.")
+        if inputs_embeds.ndim != 3:
+            raise ValueError(
+                "inputs_embeds must have shape (batch, seq_len, hidden_size), "
+                f"got {tuple(inputs_embeds.shape)}."
+            )
+        if inputs_embeds.shape[-1] != hidden_size:
+            raise ValueError(
+                f"inputs_embeds hidden size must be {hidden_size}, got {inputs_embeds.shape[-1]}."
+            )
+        batch_size, seq_len = inputs_embeds.shape[:2]
+
+    expected_shape = (batch_size, seq_len)
+    for name, value in (("attention_mask", attention_mask), ("type_ids", type_ids)):
+        if value is not None and tuple(value.shape) != expected_shape:
+            raise ValueError(f"{name} must have shape {expected_shape}, got {tuple(value.shape)}.")
+    return expected_shape
+
+
+def _has_packed_multimodal_layout(
+    type_ids: torch.Tensor | None,
+    aa_type: int,
+    struct_type: int,
+    pad_type: int,
+) -> bool:
+    if type_ids is None:
+        return False
+    if type_ids.ndim != 2:
+        raise ValueError(
+            f"Expected type_ids to have shape (batch, seq_len), got {tuple(type_ids.shape)}"
+        )
+    seq_len = type_ids.shape[-1]
+    if seq_len % 2 != 0:
+        return False
+
+    half_len = seq_len // 2
+    first_half = type_ids[:, :half_len]
+    second_half = type_ids[:, half_len:]
+
+    first_is_aa = ((first_half == aa_type) | (first_half == pad_type)).all(dim=-1)
+    first_is_struct = ((first_half == struct_type) | (first_half == pad_type)).all(dim=-1)
+    second_is_aa = ((second_half == aa_type) | (second_half == pad_type)).all(dim=-1)
+    second_is_struct = ((second_half == struct_type) | (second_half == pad_type)).all(dim=-1)
+    first_count = first_half.ne(pad_type).sum(dim=-1)
+    second_count = second_half.ne(pad_type).sum(dim=-1)
+    modalities_are_separate = (first_is_aa & second_is_struct) | (first_is_struct & second_is_aa)
+    packed_rows = modalities_are_separate & first_count.gt(0) & first_count.eq(second_count)
+    return bool(packed_rows.all())
+
+
+@dataclass
+class DPLM2MaskedLMOutput(MaskedLMOutput):
+    """Masked-LM output with DPLM2 extensions after the HF fields."""
+
+    s_max: tuple[list[torch.Tensor], ...] | None = None
+    last_hidden_state: torch.Tensor | None = None
+
+
+@dataclass
+class DPLM2ModelOutput(BaseModelOutputWithPoolingAndCrossAttentions):
+    """Base-model output with optional attention diagnostics."""
+
+    s_max: tuple[list[torch.Tensor], ...] | None = None
+
+
+@dataclass
+class DPLM2SequenceClassifierOutput(SequenceClassifierOutput):
+    """Sequence-classification output with optional attention diagnostics."""
+
+    s_max: tuple[list[torch.Tensor], ...] | None = None
+
+
+@dataclass
+class DPLM2TokenClassifierOutput(TokenClassifierOutput):
+    """Token-classification output with optional attention diagnostics."""
+
+    s_max: tuple[list[torch.Tensor], ...] | None = None
+
+
+@dataclass
+class DPLM2EncoderOutput(ModelOutput):
+    last_hidden_state: torch.Tensor | None = None
+    hidden_states: tuple[torch.Tensor, ...] | None = None
+    attentions: tuple[torch.Tensor, ...] | None = None
+    s_max: tuple[list[torch.Tensor], ...] | None = None
+
+
+class DPLM2Config(EsmConfig):
+    model_type = "dplm2"
+
+    def __init__(
+        self,
+        attn_backend: str | None = "sdpa",
+        add_pooling_layer: bool = False,
+        aa_type: int = 1,
+        struct_type: int = 0,
+        pad_type: int = 2,
+        **kwargs,
+    ):
+        if kwargs.get("is_decoder", False) or kwargs.get("add_cross_attention", False):
+            raise ValueError(
+                "DPLM2 is encoder-only; is_decoder and add_cross_attention must be false."
+            )
+
+        # Published DPLM2 checkpoint configs inherited ``use_cache=true`` from
+        # EsmConfig even though the FastPLMs encoder has never implemented a KV
+        # cache.  Keep those legacy artifacts loadable, but make the effective
+        # and newly serialized contract explicit and fail closed.
+        if kwargs.get("use_cache") is True:
+            warnings.warn(
+                "Legacy DPLM2 config requested use_cache=True, but DPLM2 is encoder-only "
+                "and does not implement KV caching; normalizing use_cache to False.",
+                UserWarning,
+                stacklevel=2,
+            )
+        kwargs["is_decoder"] = False
+        kwargs["add_cross_attention"] = False
+        kwargs["use_cache"] = False
+        super().__init__(**kwargs)
+        # DPLM2's published implementation and manifest expose SDPA only. An
+        # older checkpoint may omit this FastPLMs field (or serialize it as
+        # null), so normalize that legacy representation to the same explicit
+        # backend before Transformers chooses its own generic eager default.
+        self.attn_backend = "sdpa" if attn_backend is None else attn_backend
+        self.add_pooling_layer = add_pooling_layer
+        self.aa_type = aa_type
+        self.struct_type = struct_type
+        self.pad_type = pad_type
+        self.tie_word_embeddings = False
+
+
+_TOKENIZER_LOAD_CONTEXT_KEYS = (
+    "cache_dir",
+    "force_download",
+    "local_files_only",
+    "proxies",
+    "revision",
+    "subfolder",
+    "token",
+    "trust_remote_code",
+)
+
+
+class DPLM2PreTrainedModel(FastPLMsAttentionMixin, EsmPreTrainedModel):
+    config_class = DPLM2Config
+    # All advertised wrappers install the encoder at ``self.esm``.  Transformers
+    # uses this name both for ``base_model`` and checkpoint prefix reconciliation.
+    base_model_prefix = "esm"
+    supports_gradient_checkpointing = True
+    all_tied_weights_keys: ClassVar[dict[str, str]] = {}
+    _supports_flex_attn = False
+    _supports_flash_attn = False
+    _supports_flash_attn_2 = False
+    _supports_flash_attn_3 = False
+    _fastplms_attention_implementations = ("sdpa",)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        load_context = {key: kwargs[key] for key in _TOKENIZER_LOAD_CONTEXT_KEYS if key in kwargs}
+        if "token" not in load_context and "use_auth_token" in kwargs:
+            load_context["token"] = kwargs["use_auth_token"]
+        load_context["source"] = pretrained_model_name_or_path
+
+        loaded = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+        model = loaded[0] if isinstance(loaded, tuple) else loaded
+        model.__dict__["_fastplms_tokenizer_load_context"] = load_context
+        model.__dict__["_fastplms_tokenizer"] = None
+        return loaded
+
+    @property
+    def tokenizer(self):
+        tokenizer = self.__dict__.get("_fastplms_tokenizer")
+        if tokenizer is None:
+            load_context = dict(self.__dict__.get("_fastplms_tokenizer_load_context") or {})
+            source = load_context.pop("source", None)
+            if source is None:
+                source = str(getattr(self.config, "_name_or_path", "")).strip()
+            if not source:
+                raise RuntimeError(
+                    "DPLM2 tokenizer loading requires a model loaded with from_pretrained "
+                    "so checkpoint provenance is available."
+                )
+            tokenizer_kwargs = {
+                key: value
+                for key, value in load_context.items()
+                if key in _TOKENIZER_LOAD_CONTEXT_KEYS and value is not None
+            }
+            resolved_revision = getattr(self.config, "_commit_hash", None)
+            if resolved_revision:
+                tokenizer_kwargs["revision"] = resolved_revision
+            tokenizer = DPLM2Tokenizer.from_pretrained(source, **tokenizer_kwargs)
+            self.__dict__["_fastplms_tokenizer"] = tokenizer
+        return tokenizer
+
+    @tokenizer.setter
+    def tokenizer(self, value) -> None:
+        self.__dict__["_fastplms_tokenizer"] = value
+
+    def _tokenize_sequence_batch(
+        self,
+        sequences: Sequence[str],
+        *,
+        tokenizer: Any | None = None,
+        **kwargs: Any,
+    ) -> Any:
+        """Tokenize raw amino-acid sequences with official DPLM2 boundaries."""
+
+        resolved = tokenizer if tokenizer is not None else self.tokenizer
+        sequence_list = [sequences] if isinstance(sequences, str) else sequences
+        formatted = [
+            f"{resolved.aa_cls_token}{sequence}{resolved.aa_eos_token}"
+            for sequence in sequence_list
+        ]
+        return resolved(formatted, add_special_tokens=False, **kwargs)
+
+    @property
+    def attn_backend(self) -> str:
+        return self.config.attn_backend
+
+    @attn_backend.setter
+    def attn_backend(self, backend: str) -> None:
+        if backend not in self._fastplms_attention_implementations:
+            raise ValueError(
+                f"DPLM2 does not support {backend!r}; expected one of "
+                f"{self._fastplms_attention_implementations}."
+            )
+        self.config.attn_backend = backend
+        resolved = resolve_attention_backend(backend)
+        for module in self.modules():
+            if isinstance(module, ModifiedEsmEncoder):
+                module.attention_backend = resolved
+            elif isinstance(module, ModifiedEsmSelfAttention):
+                module.attn_backend = resolved
+
+
+class ModifiedRotaryEmbedding(RotaryEmbedding):
+    def __init__(self, dim: int, aa_type: int, struct_type: int, pad_type: int):
+        super().__init__(dim)
+        self.aa_type = aa_type
+        self.struct_type = struct_type
+        self.pad_type = pad_type
+
+    def _has_multimodal_tokens(self, type_ids: torch.Tensor | None) -> bool:
+        # The split rotary path only works when the sequence tensor is already packed
+        # as two equal-length, modality-specific halves. Either track may come first.
+        # Plain protein batches can still contain high-ID special tokens, so mere
+        # modality presence is not enough.
+        return _has_packed_multimodal_layout(
+            type_ids=type_ids,
+            aa_type=self.aa_type,
+            struct_type=self.struct_type,
+            pad_type=self.pad_type,
+        )
+
+    def align_frequency_buffer(
+        self,
+        *,
+        device: torch.device,
+        dtype: torch.dtype,
+    ) -> None:
+        """Match the official model-wide ``to(device, dtype)`` conversion.
+
+        Transformers' meta-device loader converts parameters to the requested
+        dtype but can leave this persistent rotary buffer in FP32. The pinned
+        official implementation moves the complete module, including
+        ``inv_freq``. Aligning the buffer before building rotary factors keeps
+        Q, K, and V in one dtype for every attention backend.
+        """
+
+        if self.inv_freq.device == device and self.inv_freq.dtype == dtype:
+            return
+        self.inv_freq = self.inv_freq.to(device=device, dtype=dtype)
+        self._seq_len_cached = None
+        self._cos_cached = None
+        self._sin_cached = None
+
+    def _update_cos_sin_tables(
+        self,
+        x: torch.Tensor,
+        type_ids: torch.Tensor | None,
+        seq_dimension: int = 2,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        seq_len = x.shape[seq_dimension]
+        if self._has_multimodal_tokens(type_ids):
+            seq_len = seq_len // 2
+
+        cache_is_stale = (
+            self._cos_cached is None
+            or self._sin_cached is None
+            or seq_len != self._seq_len_cached
+            or self._cos_cached.device != x.device
+            or self._cos_cached.dtype != self.inv_freq.dtype
+        )
+        if cache_is_stale:
+            self._seq_len_cached = seq_len
+            t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
+            freqs = torch.outer(t, self.inv_freq)
+            # Match the official DPLM2 operation order: rotary factors inherit
+            # the frequency-buffer dtype. This keeps them in FP32 under BF16
+            # autocast, while a model explicitly converted to BF16 still builds
+            # BF16 factors and remains usable without autocast.
+            emb = torch.cat((freqs, freqs), dim=-1).to(device=x.device)
+            self._cos_cached = emb.cos()[None, None, :, :]
+            self._sin_cached = emb.sin()[None, None, :, :]
+
+        return self._cos_cached, self._sin_cached
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        type_ids: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        self._cos_cached, self._sin_cached = self._update_cos_sin_tables(
+            k,
+            type_ids=type_ids,
+            seq_dimension=-2,
+        )
+
+        if self._has_multimodal_tokens(type_ids):
+            q_1, q_2 = q.chunk(2, dim=-2)
+            k_1, k_2 = k.chunk(2, dim=-2)
+            q_1 = apply_rotary_pos_emb(q_1, self._cos_cached, self._sin_cached)
+            q_2 = apply_rotary_pos_emb(q_2, self._cos_cached, self._sin_cached)
+            k_1 = apply_rotary_pos_emb(k_1, self._cos_cached, self._sin_cached)
+            k_2 = apply_rotary_pos_emb(k_2, self._cos_cached, self._sin_cached)
+            return torch.cat((q_1, q_2), dim=-2), torch.cat((k_1, k_2), dim=-2)
+
+        return (
+            apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached),
+            apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached),
+        )
+
+
+class ModifiedEsmSelfAttention(EsmSelfAttention):
+    def __init__(self, config, position_embedding_type=None):
+        super().__init__(config, position_embedding_type)
+        self.config = config
+        self.scale = self.attention_head_size**-0.5
+        self.dropout_prob = config.attention_probs_dropout_prob
+        self.attn_backend = resolve_attention_backend(config.attn_backend)
+        self.rotary_embeddings = ModifiedRotaryEmbedding(
+            dim=self.attention_head_size,
+            aa_type=config.aa_type,
+            struct_type=config.struct_type,
+            pad_type=config.pad_type,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask_4d: torch.Tensor | None = None,
+        output_attentions: bool = False,
+        output_s_max: bool = False,
+        type_ids: torch.Tensor | None = None,
+        effective_backend: AttentionBackend | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]:
+        batch_size, seq_length = hidden_states.shape[:-1]
+        hidden_shape = (batch_size, seq_length, -1, self.attention_head_size)
+        query_heads = self.query(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_heads = self.key(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_heads = self.value(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        query_heads = query_heads * self.scale
+
+        if self.position_embedding_type == "rotary":
+            self.rotary_embeddings.align_frequency_buffer(
+                device=query_heads.device,
+                dtype=self.query.weight.dtype,
+            )
+            query_heads, key_heads = self.rotary_embeddings(query_heads, key_heads, type_ids)
+
+        attn_output, attn_weights, s_max = self._attn(
+            query_heads,
+            key_heads,
+            value_heads,
+            attention_mask_4d=attention_mask_4d,
+            output_attentions=output_attentions,
+            output_s_max=output_s_max,
+            effective_backend=effective_backend,
+        )
+        return attn_output, attn_weights, s_max
+
+    def _attn(
+        self,
+        query_heads: torch.Tensor,
+        key_heads: torch.Tensor,
+        value_heads: torch.Tensor,
+        attention_mask_4d: torch.Tensor | None = None,
+        output_attentions: bool = False,
+        output_s_max: bool = False,
+        effective_backend: AttentionBackend | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]:
+        if effective_backend is None:
+            effective_backend = resolve_attention_backend_for_call(
+                self.attn_backend,
+                output_attentions=output_attentions,
+            )
+
+        if effective_backend == AttentionBackend.EAGER:
+            attn_output, attn_weights, s_max = self._manual_attn(
+                query_heads, key_heads, value_heads, attention_mask_4d, output_s_max
+            )
+            return attn_output, attn_weights if output_attentions else None, s_max
+
+        if output_attentions:
+            raise AssertionError(
+                "DPLM2 output_attentions=True must resolve to eager attention for this call."
+            )
+
+        if effective_backend != AttentionBackend.SDPA:
+            raise AssertionError(f"Unsupported resolved backend: {effective_backend}")
+        attn_output, attn_weights = self._sdpa_attn(
+            query_heads,
+            key_heads,
+            value_heads,
+            attention_mask_4d,
+        )
+
+        s_max = self._compute_s_max(query_heads, key_heads) if output_s_max else None
+        return attn_output, attn_weights, s_max
+
+    @torch.no_grad()
+    def _compute_s_max(
+        self, query_heads: torch.Tensor, key_heads: torch.Tensor
+    ) -> list[torch.Tensor]:
+        q_norm = torch.linalg.vector_norm(query_heads, dim=-1)
+        k_norm = torch.linalg.vector_norm(key_heads, dim=-1)
+        s_max_bound = (q_norm.max(dim=-1).values * k_norm.max(dim=-1).values).max(dim=0).values
+        return [s_max_bound[h] for h in range(self.num_attention_heads)]
+
+    def _manual_attn(
+        self,
+        query_heads: torch.Tensor,
+        key_heads: torch.Tensor,
+        value_heads: torch.Tensor,
+        attention_mask_4d: torch.Tensor | None = None,
+        output_s_max: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor] | None]:
+        attn_weights = torch.matmul(query_heads, key_heads.transpose(-1, -2))
+        if attention_mask_4d is not None:
+            attn_weights = attn_weights.masked_fill(attention_mask_4d.logical_not(), float("-inf"))
+        attn_weights = F.softmax(attn_weights, dim=-1)
+        if self.dropout_prob > 0 and self.training:
+            attn_weights = F.dropout(attn_weights, p=self.dropout_prob, training=self.training)
+        context_heads = torch.matmul(attn_weights, value_heads)
+        attn_output = rearrange(context_heads, "b h s d -> b s (h d)")
+        s_max = self._compute_s_max(query_heads, key_heads) if output_s_max else None
+        return attn_output, attn_weights, s_max
+
+    def _sdpa_attn(
+        self,
+        query_heads: torch.Tensor,
+        key_heads: torch.Tensor,
+        value_heads: torch.Tensor,
+        attention_mask_4d: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, None]:
+        # Pinned DPLM2 uses PyTorch's efficient SDPA kernel for its non-null
+        # padding mask. Newer PyTorch releases otherwise select cuDNN on H100,
+        # which exceeds the fixed deep-BF16 parity target. This is still the
+        # public SDPA operation and raises if its required CUDA kernel is absent.
+        kernel_context = (
+            sdpa_kernel(SDPBackend.EFFICIENT_ATTENTION)
+            if query_heads.is_cuda
+            else contextlib.nullcontext()
+        )
+        with kernel_context:
+            context_heads = F.scaled_dot_product_attention(
+                query_heads,
+                key_heads,
+                value_heads,
+                attn_mask=attention_mask_4d,
+                dropout_p=self.dropout_prob if self.training else 0.0,
+                scale=1.0,
+            )
+        return rearrange(context_heads, "b h s d -> b s (h d)"), None
+
+
+class ModifiedEsmAttention(EsmAttention):
+    def __init__(self, config):
+        nn.Module.__init__(self)
+        self.self = ModifiedEsmSelfAttention(config)
+        self.output = EsmSelfOutput(config)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask_4d: torch.Tensor | None = None,
+        output_attentions: bool = False,
+        output_s_max: bool = False,
+        type_ids: torch.Tensor | None = None,
+        effective_backend: AttentionBackend | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]:
+        hidden_states_ln = self.LayerNorm(hidden_states)
+        attn_output, attn_weights, s_max = self.self(
+            hidden_states_ln,
+            attention_mask_4d=attention_mask_4d,
+            output_attentions=output_attentions,
+            output_s_max=output_s_max,
+            type_ids=type_ids,
+            effective_backend=effective_backend,
+        )
+        attention_output = self.output(attn_output, hidden_states)
+        return attention_output, attn_weights, s_max
+
+
+class ModifiedEsmLayer(EsmLayer):
+    def __init__(self, config):
+        nn.Module.__init__(self)
+        self.chunk_size_feed_forward = config.chunk_size_feed_forward
+        self.seq_len_dim = 1
+        self.attention = ModifiedEsmAttention(config)
+        self.intermediate = EsmIntermediate(config)
+        self.output = EsmOutput(config)
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask_4d: torch.Tensor | None = None,
+        output_attentions: bool = False,
+        output_s_max: bool = False,
+        type_ids: torch.Tensor | None = None,
+        effective_backend: AttentionBackend | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]:
+        attention_output, attn_weights, s_max = self.attention(
+            hidden_states,
+            attention_mask_4d=attention_mask_4d,
+            output_attentions=output_attentions,
+            output_s_max=output_s_max,
+            type_ids=type_ids,
+            effective_backend=effective_backend,
+        )
+        layer_output = self.feed_forward_chunk(attention_output)
+        return layer_output, attn_weights, s_max
+
+
+class ModifiedEsmEncoder(EsmEncoder):
+    def __init__(self, config):
+        nn.Module.__init__(self)
+        self.config = config
+        self.attention_backend = resolve_attention_backend(config.attn_backend)
+        self.layer = nn.ModuleList(
+            [ModifiedEsmLayer(config) for _ in range(config.num_hidden_layers)]
+        )
+        self.emb_layer_norm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.gradient_checkpointing = False
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+        output_s_max: bool = False,
+        type_ids: torch.Tensor | None = None,
+    ) -> DPLM2EncoderOutput:
+        first_parameter = next(self.parameters(), None)
+        if (
+            not self.training
+            and first_parameter is not None
+            and first_parameter.dtype == torch.bfloat16
+        ):
+            raise RuntimeError(
+                "DPLM2 BF16 inference requires FP32-resident parameters under "
+                "CUDA BF16 autocast; static BF16 parameters do not meet the "
+                "declared parity contract."
+            )
+        all_hidden_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+        full_s_max = () if output_s_max else None
+
+        effective_backend = resolve_attention_backend_for_call(
+            self.attention_backend,
+            output_attentions=output_attentions,
+        )
+        _, attention_mask_4d, _ = get_attention_mask(
+            effective_backend=effective_backend,
+            batch_size=hidden_states.shape[0],
+            seq_len=hidden_states.shape[1],
+            device=hidden_states.device,
+            attention_mask=attention_mask,
+            dtype=hidden_states.dtype,
+            mask_semantics="padding",
+        )
+
+        for layer_module in self.layer:
+            if output_hidden_states:
+                all_hidden_states = (*all_hidden_states, hidden_states)
+
+            if self.gradient_checkpointing and self.training:
+                hidden_states, attn_weights, s_max = self._gradient_checkpointing_func(
+                    layer_module.__call__,
+                    hidden_states,
+                    attention_mask_4d,
+                    output_attentions,
+                    output_s_max,
+                    type_ids,
+                    effective_backend,
+                )
+            else:
+                hidden_states, attn_weights, s_max = layer_module(
+                    hidden_states,
+                    attention_mask_4d=attention_mask_4d,
+                    output_attentions=output_attentions,
+                    output_s_max=output_s_max,
+                    type_ids=type_ids,
+                    effective_backend=effective_backend,
+                )
+
+            if all_attentions is not None:
+                all_attentions = (*all_attentions, attn_weights)
+            if full_s_max is not None:
+                full_s_max = (*full_s_max, s_max)
+
+        if self.emb_layer_norm_after:
+            hidden_states = self.emb_layer_norm_after(hidden_states)
+
+        if output_hidden_states:
+            all_hidden_states = (*all_hidden_states, hidden_states)
+
+        return DPLM2EncoderOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+            attentions=all_attentions,
+            s_max=full_s_max,
+        )
+
+
+class FAST_DPLM2_ENCODER(DPLM2PreTrainedModel, EmbeddingMixin):
+    """Inner encoder class that holds the actual ESM-style weights (embeddings, encoder)
+    so that the weight keys are prefixed with 'esm.' in the outer DPLM2Model,
+    matching pretrained DPLM2 checkpoints."""
+
+    def __init__(self, config, **kwargs):
+        DPLM2PreTrainedModel.__init__(self, config, **kwargs)
+        self.config = config
+        self.embeddings = EsmEmbeddings(config)
+        self.encoder = ModifiedEsmEncoder(config)
+        self.contact_head = EsmContactPredictionHead(
+            in_features=config.num_hidden_layers * config.num_attention_heads,
+            bias=True,
+        )
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings.word_embeddings = value
+
+    def predict_contacts(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Predict residue contacts with the checkpoint's tied contact head."""
+        input_ids = _normalize_dplm2_input_ids(input_ids, self.config.vocab_size)
+        if attention_mask is None:
+            attention_mask = input_ids.ne(self.config.pad_token_id)
+        type_ids = self._get_modality_type(input_ids, attention_mask)
+        attentions = self(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            type_ids=type_ids,
+            output_attentions=True,
+        ).attentions
+        if attentions is None:
+            raise RuntimeError("DPLM2 did not return attention maps for contact prediction.")
+        # A is the layer/head attention tensor; M marks valid tokens.
+        attention_tensor = torch.stack(attentions, dim=1)
+        residue_mask = attention_mask.to(dtype=attention_tensor.dtype)
+        attention_tensor = (
+            attention_tensor
+            * residue_mask[:, None, None, :, None]
+            * residue_mask[:, None, None, None, :]
+        )
+        return self.contact_head(input_ids, attention_tensor)
+
+    def _get_modality_type(
+        self, input_ids: torch.Tensor, attention_mask: torch.Tensor
+    ) -> torch.Tensor:
+        return _infer_modality_type(input_ids, attention_mask)
+
+    def _embed(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        hidden_state_index: int = -1,
+        store_all_hidden_states: bool = False,
+    ) -> torch.Tensor:
+        input_ids = _normalize_dplm2_input_ids(input_ids, self.config.vocab_size)
+        if attention_mask is None:
+            attention_mask = input_ids.ne(self.config.pad_token_id)
+        type_ids = _infer_modality_type(input_ids, attention_mask)
+        token_embedding_output = self.embeddings(input_ids, attention_mask=attention_mask)
+        output_hidden_states = store_all_hidden_states or hidden_state_index != -1
+        encoder_outputs = self.encoder(
+            token_embedding_output,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=False,
+            type_ids=type_ids,
+        )
+        return select_hidden_state_embeddings(
+            encoder_outputs.last_hidden_state,
+            encoder_outputs.hidden_states,
+            hidden_state_index=hidden_state_index,
+            store_all_hidden_states=store_all_hidden_states,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        output_s_max: bool | None = False,
+        return_dict: bool | None = None,
+        type_ids: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, ...] | DPLM2EncoderOutput:
+        _validate_dplm2_model_inputs(
+            input_ids=input_ids,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            type_ids=type_ids,
+            hidden_size=self.config.hidden_size,
+        )
+        output_attentions = (
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is not None:
+            input_ids = _normalize_dplm2_input_ids(input_ids, self.config.vocab_size)
+        token_embedding_output = self.embeddings(
+            input_ids=input_ids,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+        )
+        encoder_outputs = self.encoder(
+            token_embedding_output,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+            output_s_max=output_s_max,
+            type_ids=type_ids,
+        )
+
+        result = DPLM2EncoderOutput(
+            last_hidden_state=encoder_outputs.last_hidden_state,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+            s_max=encoder_outputs.s_max,
+        )
+        if not return_dict:
+            return result.to_tuple()
+        return result
+
+
+class DPLM2Model(DPLM2PreTrainedModel, EmbeddingMixin):
+    config_class = DPLM2Config
+
+    def __init__(self, config, add_pooling_layer: bool | None = None):
+        DPLM2PreTrainedModel.__init__(self, config)
+        self.config = config
+        self.esm = FAST_DPLM2_ENCODER(config)
+        if add_pooling_layer is None:
+            add_pooling_layer = config.add_pooling_layer
+        config.add_pooling_layer = bool(add_pooling_layer)
+        self.pooler = EsmPooler(config) if add_pooling_layer else None
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.esm.embeddings.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.esm.embeddings.word_embeddings = value
+
+    def predict_contacts(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        return self.esm.predict_contacts(input_ids, attention_mask)
+
+    def _embed(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        hidden_state_index: int = -1,
+        store_all_hidden_states: bool = False,
+    ) -> torch.Tensor:
+        return self.esm._embed(
+            input_ids,
+            attention_mask,
+            hidden_state_index=hidden_state_index,
+            store_all_hidden_states=store_all_hidden_states,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        output_s_max: bool | None = False,
+        return_dict: bool | None = None,
+        type_ids: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, ...] | DPLM2ModelOutput:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        direct_dplm_esm = getattr(self.config, "dplm_type", None) == "dplm_esm"
+        _validate_dplm2_model_inputs(
+            input_ids=input_ids,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            type_ids=type_ids,
+            hidden_size=self.config.hidden_size,
+        )
+        if inputs_embeds is not None and type_ids is None and not direct_dplm_esm:
+            raise ValueError(
+                "type_ids is required for multimodal DPLM2 calls that use inputs_embeds."
+            )
+        if input_ids is not None:
+            normalized_input_ids = _normalize_dplm2_input_ids(input_ids, self.config.vocab_size)
+            if attention_mask is None:
+                attention_mask = normalized_input_ids.ne(self.config.pad_token_id)
+            if type_ids is None and not direct_dplm_esm:
+                type_ids = _infer_modality_type(normalized_input_ids, attention_mask)
+            input_ids = normalized_input_ids
+
+        outputs = self.esm(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            output_s_max=output_s_max,
+            return_dict=True,
+            type_ids=type_ids,
+        )
+        sequence_output = outputs.last_hidden_state
+        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
+
+        result = DPLM2ModelOutput(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            s_max=outputs.s_max,
+        )
+        return result if return_dict else result.to_tuple()
+
+
+class DPLM2ForMaskedLM(FastPLMTestTimeTrainingMixin, DPLM2PreTrainedModel, EmbeddingMixin):
+    config_class = DPLM2Config
+
+    def __init__(
+        self,
+        config,
+        dropout: float | None = None,
+        vocab_size: int | None = None,
+    ):
+        if dropout is not None:
+            config.hidden_dropout_prob = dropout
+        config.tie_word_embeddings = False
+        if vocab_size is not None:
+            config.vocab_size = vocab_size
+        DPLM2PreTrainedModel.__init__(self, config)
+        self.esm = FAST_DPLM2_ENCODER(config)
+        self.lm_head = EsmLMHead(config)
+        self.loss_fct = nn.CrossEntropyLoss()
+        self.post_init()
+        self.pad_id = config.pad_token_id
+        self.contact_head = None
+        self.init_ttt({"lora_target_replace_module": "ModifiedEsmAttention"})
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.esm.get_input_embeddings()
+
+    def set_input_embeddings(self, value: nn.Module) -> None:
+        self.esm.set_input_embeddings(value)
+
+    def get_output_embeddings(self):
+        return self.lm_head.decoder
+
+    def set_output_embeddings(self, new_embeddings):
+        old_bias = self.lm_head.bias
+        new_vocab_size = int(new_embeddings.out_features)
+        if old_bias.shape[0] != new_vocab_size:
+            resized_bias = old_bias.new_zeros(new_vocab_size)
+            copy_length = min(old_bias.shape[0], new_vocab_size)
+            with torch.no_grad():
+                resized_bias[:copy_length].copy_(old_bias[:copy_length])
+            self.lm_head.bias = nn.Parameter(resized_bias)
+        # EsmLMHead.forward adds this standalone bias after the decoder. HF's
+        # generic LM-head resizer may create a biased Linear, which would apply
+        # the bias twice and introduce an undeclared shared tensor on save.
+        new_embeddings.bias = None
+        self.lm_head.decoder = new_embeddings
+
+    def generate(
+        self,
+        input_tokens: torch.Tensor,
+        max_iter: int | None = None,
+        temperature: float = 1.0,
+        partial_masks: torch.Tensor | None = None,
+        unmasking_strategy: str = "stochastic1.0",
+        sampling_strategy: str = "annealing@2.0:0.1",
+        show_progress: bool = False,
+        **kwargs,
+    ) -> dict[str, torch.Tensor]:
+        """Generate packed sequence and structure tokens with DPLM2 diffusion.
+
+        ``input_tokens`` is X with shape (b, l). Positions marked ``True`` in
+        ``partial_masks`` remain fixed. The returned mapping contains
+        ``output_tokens``, matching the official DPLM2 public API.
+        """
+
+        if kwargs:
+            names = ", ".join(sorted(kwargs))
+            raise TypeError(f"Unexpected DPLM2 generation arguments: {names}")
+        return generate_dplm2(
+            self,
+            input_tokens,
+            max_iter=max_iter,
+            temperature=temperature,
+            partial_masks=partial_masks,
+            unmasking_strategy=unmasking_strategy,
+            sampling_strategy=sampling_strategy,
+            show_progress=show_progress,
+        )
+
+    def predict_contacts(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        """Return the official ESM contact head output from the encoder."""
+        return self.esm.predict_contacts(input_ids, attention_mask)
+
+    def _get_modality_type(
+        self, input_ids: torch.Tensor, attention_mask: torch.Tensor
+    ) -> torch.Tensor:
+        input_ids = _normalize_dplm2_input_ids(input_ids, self.config.vocab_size)
+        return _infer_modality_type(input_ids, attention_mask)
+
+    def _embed(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        hidden_state_index: int = -1,
+        store_all_hidden_states: bool = False,
+    ) -> torch.Tensor:
+        if attention_mask is None:
+            attention_mask = input_ids.ne(self.pad_id)
+        type_ids = self._get_modality_type(input_ids, attention_mask)
+        output_hidden_states = store_all_hidden_states or hidden_state_index != -1
+        outputs = self.esm(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            type_ids=type_ids,
+            output_attentions=False,
+            output_hidden_states=output_hidden_states,
+            return_dict=True,
+        )
+        return select_hidden_state_embeddings(
+            outputs.last_hidden_state,
+            outputs.hidden_states,
+            hidden_state_index=hidden_state_index,
+            store_all_hidden_states=store_all_hidden_states,
+        )
+
+    def _ttt_get_trainable_modules(self) -> list[nn.Module]:
+        return [self.esm]
+
+    def _ttt_tokenize(
+        self,
+        seq: str | list[str] | None = None,
+        input_ids: torch.Tensor | None = None,
+        **kwargs: Any,
+    ) -> torch.Tensor:
+        del kwargs
+        if input_ids is not None:
+            return input_ids
+        if seq is None:
+            raise ValueError("Pass either seq or input_ids for TTT.")
+        sequences = [seq] if isinstance(seq, str) else seq
+        tokenized = self._tokenize_sequence_batch(
+            sequences,
+            return_tensors="pt",
+            padding=True,
+        )
+        return tokenized["input_ids"]
+
+    def _ttt_mask_token(self) -> int:
+        return int(self.tokenizer._token_to_id[self.tokenizer.aa_mask_token])
+
+    def _ttt_replacement_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
+        tokenizer = self.tokenizer
+        special_ids = set(tokenizer.all_special_ids)
+        struct_boundary = int(tokenizer._token_to_id[tokenizer.struct_cls_token])
+        residue_ids = []
+        for residue in "ACDEFGHIKLMNPQRSTVWY":
+            token_id = tokenizer._token_to_id.get(residue)
+            if (
+                isinstance(token_id, int)
+                and 0 <= token_id < struct_boundary
+                and token_id not in special_ids
+                and token_id not in residue_ids
+            ):
+                residue_ids.append(token_id)
+        if not residue_ids:
+            raise RuntimeError("DPLM2 TTT amino-acid replacement set is empty.")
+        if len(residue_ids) != 20:
+            raise RuntimeError(
+                "DPLM2 TTT requires all 20 canonical amino-acid replacement tokens; "
+                f"resolved {len(residue_ids)}."
+            )
+        return torch.tensor(residue_ids, device=input_ids.device, dtype=input_ids.dtype)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        type_ids: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        output_s_max: bool | None = False,
+        return_dict: bool | None = None,
+    ) -> tuple[torch.Tensor] | DPLM2MaskedLMOutput:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        direct_dplm_esm = getattr(self.config, "dplm_type", None) == "dplm_esm"
+        _validate_dplm2_model_inputs(
+            input_ids=input_ids,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            type_ids=type_ids,
+            hidden_size=self.config.hidden_size,
+        )
+
+        if attention_mask is None:
+            if input_ids is None:
+                raise ValueError(
+                    "attention_mask is required when DPLM2 is called with inputs_embeds."
+                )
+            attention_mask = input_ids.ne(self.pad_id)
+
+        encoder_input_ids = input_ids
+        if input_ids is not None:
+            input_ids = _normalize_dplm2_input_ids(input_ids, self.config.vocab_size)
+            encoder_input_ids = input_ids
+        if type_ids is None and not direct_dplm_esm:
+            if input_ids is None:
+                raise ValueError(
+                    "type_ids is required for multimodal DPLM2 calls that use inputs_embeds."
+                )
+            type_ids = self._get_modality_type(input_ids, attention_mask)
+
+        if input_ids is not None and inputs_embeds is None and not direct_dplm_esm:
+            # The official multimodal wrapper applies the embedding block
+            # once before entering EsmForDPLM2. The inner ESM model then
+            # applies it a second time using these intermediate embeddings.
+            inputs_embeds = self.esm.embeddings(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+            )
+            encoder_input_ids = None
+
+        outputs = self.esm(
+            input_ids=encoder_input_ids,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            output_s_max=output_s_max,
+            return_dict=True,
+            type_ids=type_ids,
+        )
+
+        sequence_output = outputs.last_hidden_state
+        logits = self.lm_head(sequence_output)
+        loss = None
+        if labels is not None:
+            labels = _normalize_dplm2_input_ids(labels, self.config.vocab_size)
+            labels = labels.to(logits.device)
+            loss = self.loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
+
+        result = DPLM2MaskedLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            s_max=outputs.s_max,
+            last_hidden_state=sequence_output,
+        )
+        return result if return_dict else result.to_tuple()
+
+
+class DPLM2ForSequenceClassification(DPLM2PreTrainedModel, EmbeddingMixin):
+    config_class = DPLM2Config
+
+    def __init__(self, config):
+        DPLM2PreTrainedModel.__init__(self, config)
+        self.num_labels = config.num_labels
+        self.esm = FAST_DPLM2_ENCODER(config)
+        self.classifier = EsmClassificationHead(config)
+        self.mse = nn.MSELoss()
+        self.ce = nn.CrossEntropyLoss()
+        self.bce = nn.BCEWithLogitsLoss()
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.esm.get_input_embeddings()
+
+    def set_input_embeddings(self, value: nn.Module) -> None:
+        self.esm.set_input_embeddings(value)
+
+    def _embed(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        hidden_state_index: int = -1,
+        store_all_hidden_states: bool = False,
+    ) -> torch.Tensor:
+        return self.esm._embed(
+            input_ids,
+            attention_mask,
+            hidden_state_index=hidden_state_index,
+            store_all_hidden_states=store_all_hidden_states,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        type_ids: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        output_s_max: bool | None = False,
+        return_dict: bool | None = None,
+    ) -> tuple[torch.Tensor, ...] | DPLM2SequenceClassifierOutput:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        direct_dplm_esm = getattr(self.config, "dplm_type", None) == "dplm_esm"
+        _validate_dplm2_model_inputs(
+            input_ids=input_ids,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            type_ids=type_ids,
+            hidden_size=self.config.hidden_size,
+        )
+        if inputs_embeds is not None and type_ids is None and not direct_dplm_esm:
+            raise ValueError(
+                "type_ids is required for multimodal DPLM2 calls that use inputs_embeds."
+            )
+        if input_ids is not None:
+            input_ids = _normalize_dplm2_input_ids(input_ids, self.config.vocab_size)
+            if attention_mask is None:
+                attention_mask = input_ids.ne(self.config.pad_token_id)
+        if type_ids is None and input_ids is not None and not direct_dplm_esm:
+            type_ids = _infer_modality_type(input_ids, attention_mask)
+
+        outputs = self.esm(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            type_ids=type_ids,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            output_s_max=output_s_max,
+            return_dict=True,
+        )
+        sequence_output = outputs.last_hidden_state
+        logits = self.classifier(sequence_output)
+
+        loss = None
+        if labels is not None:
+            labels = labels.to(logits.device)
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (
+                    labels.dtype == torch.long or labels.dtype == torch.int
+                ):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                if self.num_labels == 1:
+                    loss = self.mse(logits.squeeze(), labels.squeeze())
+                else:
+                    loss = self.mse(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss = self.ce(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss = self.bce(logits, labels)
+
+        result = DPLM2SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            s_max=outputs.s_max,
+        )
+        return result if return_dict else result.to_tuple()
+
+
+class DPLM2ForTokenClassification(DPLM2PreTrainedModel, EmbeddingMixin):
+    config_class = DPLM2Config
+
+    def __init__(self, config):
+        DPLM2PreTrainedModel.__init__(self, config)
+        self.num_labels = config.num_labels
+        self.esm = FAST_DPLM2_ENCODER(config)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+        self.loss_fct = nn.CrossEntropyLoss()
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Module:
+        return self.esm.get_input_embeddings()
+
+    def set_input_embeddings(self, value: nn.Module) -> None:
+        self.esm.set_input_embeddings(value)
+
+    def _embed(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        hidden_state_index: int = -1,
+        store_all_hidden_states: bool = False,
+    ) -> torch.Tensor:
+        return self.esm._embed(
+            input_ids,
+            attention_mask,
+            hidden_state_index=hidden_state_index,
+            store_all_hidden_states=store_all_hidden_states,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        type_ids: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        output_s_max: bool | None = False,
+        return_dict: bool | None = None,
+    ) -> tuple[torch.Tensor, ...] | DPLM2TokenClassifierOutput:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        direct_dplm_esm = getattr(self.config, "dplm_type", None) == "dplm_esm"
+        _validate_dplm2_model_inputs(
+            input_ids=input_ids,
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_mask,
+            type_ids=type_ids,
+            hidden_size=self.config.hidden_size,
+        )
+        if inputs_embeds is not None and type_ids is None and not direct_dplm_esm:
+            raise ValueError(
+                "type_ids is required for multimodal DPLM2 calls that use inputs_embeds."
+            )
+        if input_ids is not None:
+            input_ids = _normalize_dplm2_input_ids(input_ids, self.config.vocab_size)
+            if attention_mask is None:
+                attention_mask = input_ids.ne(self.config.pad_token_id)
+        if type_ids is None and input_ids is not None and not direct_dplm_esm:
+            type_ids = _infer_modality_type(input_ids, attention_mask)
+
+        outputs = self.esm(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            type_ids=type_ids,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            output_s_max=output_s_max,
+            return_dict=True,
+        )
+        sequence_output = self.dropout(outputs.last_hidden_state)
+        logits = self.classifier(sequence_output)
+
+        loss = None
+        if labels is not None:
+            labels = labels.to(logits.device)
+            loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        result = DPLM2TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            s_max=outputs.s_max,
+        )
+        return result if return_dict else result.to_tuple()
+
+
+# Importing the DPLM2 model implementation makes its paired tokenizer visible
+# to AutoTokenizer. This is registration only; it performs no I/O or downloads.
+try:
+    AutoTokenizer.register(
+        DPLM2Config,
+        tokenizer_class=DPLM2Tokenizer,
+        exist_ok=True,
+    )
+except TypeError:
+    # Transformers 4.x used this name; 5.x prefers tokenizer_class.
+    AutoTokenizer.register(
+        DPLM2Config,
+        slow_tokenizer_class=DPLM2Tokenizer,
+        exist_ok=True,
+    )

fastplms/models/dplm2/tokenization_dplm2.py

@@ -0,0 +1,93 @@
+"""Independent DPLM2 amino-acid and structure tokenizer.
+
+DPLM2 stores two token tracks in one vocabulary. Amino-acid tokens use their
+own boundary, unknown, and mask tokens, while structure codes use a separate
+set. The generic ``cls_token`` and ``eos_token`` aliases are deliberately not
+defined because a caller must choose the modality-specific boundaries.
+"""
+
+from __future__ import annotations
+
+from typing import ClassVar
+
+from transformers import AddedToken, EsmTokenizer, PreTrainedTokenizer
+
+
+class DPLM2Tokenizer(EsmTokenizer):
+    """Tokenize the official DPLM2 amino-acid and structure vocabulary.
+
+    Input text is split against the complete pinned vocabulary. Amino-acid
+    sequences may be passed as contiguous characters and structure sequences
+    as whitespace-separated four-digit codes. Callers constructing a model
+    input add ``aa_*`` or ``struct_*`` boundaries explicitly and then use
+    ``add_special_tokens=False``. The output token IDs preserve the official
+    multimodal vocabulary exactly.
+    """
+
+    SPECIAL_TOKENS_ATTRIBUTES: ClassVar[list[str]] = [
+        "aa_cls_token",
+        "aa_eos_token",
+        "aa_unk_token",
+        "aa_mask_token",
+        "struct_cls_token",
+        "struct_eos_token",
+        "struct_unk_token",
+        "struct_mask_token",
+        "pad_token",
+    ]
+    # The official tokenizer exposes no generic sequence-boundary aliases.
+    # Keeping these attributes explicitly set to None preserves that public
+    # contract on Transformers v5, whose custom special-token lookup is strict.
+    bos_token: ClassVar[None] = None
+    cls_token: ClassVar[None] = None
+    eos_token: ClassVar[None] = None
+    mask_token: ClassVar[None] = None
+    sep_token: ClassVar[None] = None
+    unk_token: ClassVar[None] = None
+    bos_token_id: ClassVar[None] = None
+    cls_token_id: ClassVar[None] = None
+    eos_token_id: ClassVar[None] = None
+    mask_token_id: ClassVar[None] = None
+    sep_token_id: ClassVar[None] = None
+    unk_token_id: ClassVar[None] = None
+
+    def __init__(
+        self,
+        vocab_file: str,
+        aa_cls_token: str | AddedToken = "<cls_aa>",
+        aa_eos_token: str | AddedToken = "<eos_aa>",
+        aa_unk_token: str | AddedToken = "<unk_aa>",
+        aa_mask_token: str | AddedToken = "<mask_aa>",
+        struct_cls_token: str | AddedToken = "<cls_struct>",
+        struct_eos_token: str | AddedToken = "<eos_struct>",
+        struct_unk_token: str | AddedToken = "<unk_struct>",
+        struct_mask_token: str | AddedToken = "<mask_struct>",
+        pad_token: str | AddedToken = "<pad>",
+        **kwargs: object,
+    ) -> None:
+        with open(vocab_file, encoding="utf-8") as handle:
+            self.all_tokens = [line.strip() for line in handle.read().splitlines()]
+        self._id_to_token = dict(enumerate(self.all_tokens))
+        self._token_to_id = {token: token_id for token_id, token in self._id_to_token.items()}
+
+        # EsmTokenizer would install generic ESM boundary aliases. DPLM2 has
+        # modality-specific boundaries instead, so initialize the common
+        # tokenizer base with only the nine official special-token fields.
+        PreTrainedTokenizer.__init__(
+            self,
+            aa_cls_token=aa_cls_token,
+            aa_eos_token=aa_eos_token,
+            aa_unk_token=aa_unk_token,
+            aa_mask_token=aa_mask_token,
+            struct_cls_token=struct_cls_token,
+            struct_eos_token=struct_eos_token,
+            struct_unk_token=struct_unk_token,
+            struct_mask_token=struct_mask_token,
+            pad_token=pad_token,
+            **kwargs,
+        )
+        self.unique_no_split_tokens = self.all_tokens
+        self._update_trie(self.unique_no_split_tokens)
+
+
+__all__ = ["DPLM2Tokenizer"]

fastplms/models/ttt.py

@@ -0,0 +1,866 @@
+from __future__ import annotations
+
+import contextlib
+import math
+import numbers
+import typing as T
+from dataclasses import asdict, dataclass, fields
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+_STANDARD_AMINO_ACIDS = "ACDEFGHIKLMNPQRSTVWY"
+_TTT_SERIALIZATION_VERSION = 1
+
+
+@dataclass
+class TTTConfig:
+    lr: float = 4e-4
+    steps: int = 30
+    ags: int = 16
+    batch_size: int = 2
+    mask_ratio: float = 0.15
+    crop_size: int = 1024
+    bert_leave_prob: float = 0.1
+    bert_replace_prob: float = 0.1
+    optimizer: str = "sgd"
+    momentum: float = 0.0
+    weight_decay: float = 0.0
+    seed: int | None = 0
+    lora_rank: int = 8
+    lora_alpha: float = 32.0
+    lora_target_replace_module: str | None = None
+    lora_target_modules: tuple[str, ...] | None = None
+    initial_state_reset: bool = True
+    automatic_best_state_reset: bool = False
+    eval_each_step: bool = False
+    gradient_clip: bool = False
+    gradient_clip_max_norm: float = 1.0
+
+    def __post_init__(self) -> None:
+        self.verify()
+
+    @classmethod
+    def from_kwargs(cls, **kwargs: T.Any) -> TTTConfig:
+        valid_names = {field.name for field in fields(cls)}
+        unknown_names = set(kwargs) - valid_names
+        if unknown_names:
+            raise ValueError(f"Unknown TTTConfig fields: {sorted(unknown_names)}")
+        # JSON has no tuple type. Normalize the serialized representation while
+        # keeping the public constructor and runtime overrides type-strict.
+        if isinstance(kwargs.get("lora_target_modules"), list):
+            kwargs["lora_target_modules"] = tuple(kwargs["lora_target_modules"])
+        return cls(**kwargs)
+
+    def merged(self, overrides: T.Mapping[str, T.Any] | TTTConfig | None) -> TTTConfig:
+        if overrides is None:
+            return self
+        if isinstance(overrides, TTTConfig):
+            return overrides
+        values = {field.name: self.__dict__[field.name] for field in fields(self)}
+        for name, value in overrides.items():
+            if name not in values:
+                raise ValueError(f"Unknown TTTConfig field: {name}")
+            values[name] = value
+        return TTTConfig(**values)
+
+    def to_dict(self) -> dict[str, T.Any]:
+        return asdict(self)
+
+    def verify(self) -> None:
+        numeric_fields = {
+            "lr": self.lr,
+            "mask_ratio": self.mask_ratio,
+            "lora_alpha": self.lora_alpha,
+            "bert_leave_prob": self.bert_leave_prob,
+            "bert_replace_prob": self.bert_replace_prob,
+            "gradient_clip_max_norm": self.gradient_clip_max_norm,
+            "momentum": self.momentum,
+            "weight_decay": self.weight_decay,
+        }
+        for name, value in numeric_fields.items():
+            if isinstance(value, bool) or not isinstance(value, numbers.Real):
+                raise TypeError(f"TTT {name} must be a real number.")
+            if not math.isfinite(float(value)):
+                raise ValueError(f"TTT {name} must be finite.")
+
+        integer_fields = {
+            "steps": self.steps,
+            "ags": self.ags,
+            "batch_size": self.batch_size,
+            "crop_size": self.crop_size,
+            "lora_rank": self.lora_rank,
+        }
+        for name, value in integer_fields.items():
+            if isinstance(value, bool) or not isinstance(value, int):
+                raise TypeError(f"TTT {name} must be an integer.")
+
+        if self.seed is not None and (
+            isinstance(self.seed, bool) or not isinstance(self.seed, int)
+        ):
+            raise TypeError("TTT seed must be None or an integer.")
+
+        boolean_fields = {
+            "initial_state_reset": self.initial_state_reset,
+            "automatic_best_state_reset": self.automatic_best_state_reset,
+            "eval_each_step": self.eval_each_step,
+            "gradient_clip": self.gradient_clip,
+        }
+        for name, value in boolean_fields.items():
+            if type(value) is not bool:
+                raise TypeError(f"TTT {name} must be a boolean.")
+
+        if self.lr <= 0.0:
+            raise ValueError("TTT learning rate must be positive.")
+        if self.steps < 1:
+            raise ValueError("TTT steps must be >= 1.")
+        if self.ags < 1:
+            raise ValueError("TTT gradient accumulation steps must be >= 1.")
+        if self.batch_size < 1:
+            raise ValueError("TTT batch_size must be >= 1.")
+        if not 0.0 < self.mask_ratio <= 1.0:
+            raise ValueError("TTT mask_ratio must be in (0, 1].")
+        if self.crop_size < 1:
+            raise ValueError("TTT crop_size must be >= 1.")
+        if self.lora_rank < 1:
+            raise ValueError("TTT v1 is LoRA-only, so lora_rank must be >= 1.")
+        if self.lora_alpha <= 0.0:
+            raise ValueError("TTT lora_alpha must be positive.")
+        if not isinstance(self.optimizer, str):
+            raise TypeError("TTT optimizer must be a string.")
+        if self.optimizer not in {"adamw", "sgd"}:
+            raise ValueError("TTT optimizer must be 'adamw' or 'sgd'.")
+        if self.momentum < 0.0:
+            raise ValueError("TTT momentum must be non-negative.")
+        if self.weight_decay < 0.0:
+            raise ValueError("TTT weight_decay must be non-negative.")
+        if not 0.0 <= self.bert_leave_prob <= 1.0:
+            raise ValueError("bert_leave_prob must be in [0, 1].")
+        if not 0.0 <= self.bert_replace_prob <= 1.0:
+            raise ValueError("bert_replace_prob must be in [0, 1].")
+        if self.bert_leave_prob + self.bert_replace_prob > 1.0:
+            raise ValueError("bert_leave_prob + bert_replace_prob must be <= 1.")
+        if self.gradient_clip and self.gradient_clip_max_norm <= 0.0:
+            raise ValueError("gradient_clip_max_norm must be positive.")
+        if self.lora_target_replace_module is not None:
+            if not isinstance(self.lora_target_replace_module, str):
+                raise TypeError("lora_target_replace_module must be None or a string.")
+            if not self.lora_target_replace_module.strip():
+                raise ValueError("lora_target_replace_module must not be empty.")
+        if self.lora_target_modules is not None:
+            if not isinstance(self.lora_target_modules, tuple):
+                raise TypeError("lora_target_modules must be None or a tuple of strings.")
+            if not self.lora_target_modules:
+                raise ValueError("lora_target_modules must not be empty.")
+            if any(not isinstance(name, str) for name in self.lora_target_modules):
+                raise TypeError("lora_target_modules must contain only strings.")
+            if any(not name.strip() for name in self.lora_target_modules):
+                raise ValueError(
+                    "lora_target_modules must contain only non-empty strings."
+                )
+            if len(set(self.lora_target_modules)) != len(self.lora_target_modules):
+                raise ValueError("lora_target_modules must not contain duplicates.")
+
+
+class LoraInjectedLinear(nn.Module):
+    """ProteinTTT-compatible low-rank adapter.
+
+    ``alpha`` is the direct adapter-output multiplier used by the pinned
+    ProteinTTT ``inject_trainable_lora(..., scale=lora_alpha)`` contract.  It
+    is intentionally not divided by ``rank`` as it would be in the common
+    PEFT LoRA convention.
+    """
+
+    def __init__(
+        self,
+        linear: nn.Module,
+        rank: int,
+        alpha: float,
+        generator: torch.Generator | None = None,
+    ) -> None:
+        super().__init__()
+        weight = linear._parameters.get("weight")
+        if not isinstance(weight, torch.Tensor):
+            raise TypeError("LoRA targets must expose a tensor weight parameter.")
+        if weight.ndim != 2:
+            raise ValueError("LoRA can only wrap 2D linear weights.")
+        self.linear = linear
+        self.linear.requires_grad_(False)
+        self.rank = rank
+        # ProteinTTT names this setting ``lora_alpha`` but passes it directly
+        # to cloneofsimo/lora's ``scale`` argument.  Preserve that numerical
+        # contract for parity and for saved FastPLMs TTT configurations.
+        self.scale = alpha
+        in_features = weight.shape[1]
+        out_features = weight.shape[0]
+        # ``nn.Linear`` initializes from the process-global CPU generator. Preserve
+        # that state when TTT supplies its own generator so lazy adapter injection
+        # is reproducible without perturbing the caller's RNG stream.
+        with torch.random.fork_rng(devices=[], enabled=generator is not None):
+            self.lora_down = nn.Linear(in_features, rank, bias=False, dtype=torch.float32)
+            self.lora_up = nn.Linear(rank, out_features, bias=False, dtype=torch.float32)
+        nn.init.normal_(self.lora_down.weight, std=1.0 / rank, generator=generator)
+        nn.init.zeros_(self.lora_up.weight)
+        self.lora_down.to(device=weight.device)
+        self.lora_up.to(device=weight.device)
+        self.register_buffer(
+            "_ttt_initial_lora_down",
+            self.lora_down.weight.detach().clone(),
+            persistent=True,
+        )
+        self.register_buffer(
+            "_ttt_initial_lora_up",
+            self.lora_up.weight.detach().clone(),
+            persistent=True,
+        )
+
+    @property
+    def weight(self) -> torch.Tensor:
+        return self.linear._parameters["weight"]
+
+    @property
+    def bias(self) -> torch.Tensor | None:
+        return self.linear._parameters["bias"]
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        base = self.linear(x)
+        delta = self.lora_up(self.lora_down(x.to(dtype=torch.float32))) * self.scale
+        return base + delta.to(dtype=base.dtype)
+
+    def reset_lora_parameters(self) -> None:
+        with torch.no_grad():
+            self.lora_down.weight.copy_(self._ttt_initial_lora_down)
+            self.lora_up.weight.copy_(self._ttt_initial_lora_up)
+
+
+class FastPLMTestTimeTrainingMixin:
+    def init_ttt(self, ttt_config: TTTConfig | T.Mapping[str, T.Any] | None = None) -> None:
+        base_config = self.__dict__.get("_ttt_cfg")
+        if base_config is None:
+            base_config = TTTConfig()
+        if not isinstance(base_config, TTTConfig):
+            raise TypeError("Existing TTT configuration must be a TTTConfig instance.")
+        configured = base_config.merged(ttt_config)
+        serialized = getattr(getattr(self, "config", None), "fastplms_ttt", None)
+        serialized_initialized = False
+        if serialized is not None:
+            if not isinstance(serialized, T.Mapping):
+                raise ValueError("config.fastplms_ttt must be a mapping.")
+            version = serialized.get("version")
+            if version != _TTT_SERIALIZATION_VERSION:
+                raise ValueError(
+                    "Unsupported FastPLMs TTT serialization version "
+                    f"{version!r}; expected {_TTT_SERIALIZATION_VERSION}."
+                )
+            serialized_config = serialized.get("config")
+            if not isinstance(serialized_config, T.Mapping):
+                raise ValueError("Serialized FastPLMs TTT state is missing its config mapping.")
+            configured = TTTConfig.from_kwargs(**dict(serialized_config))
+            initialized_value = serialized.get("initialized", False)
+            if type(initialized_value) is not bool:
+                raise ValueError("Serialized FastPLMs TTT initialized flag must be a boolean.")
+            serialized_initialized = initialized_value
+
+        self._ttt_cfg = configured
+        self._ttt_cfg.verify()
+        self._ttt_initialized = False
+        if serialized_initialized:
+            self._ttt_inject_lora()
+            self._ttt_initialized = True
+
+    @property
+    def ttt_config(self) -> TTTConfig:
+        if "_ttt_cfg" not in self.__dict__:
+            self.init_ttt()
+        return self._ttt_cfg
+
+    def _ttt_get_trainable_modules(self) -> list[nn.Module]:
+        return [self]
+
+    def _ttt_get_frozen_modules(self) -> list[nn.Module]:
+        return []
+
+    def _ttt_tokenize(
+        self,
+        seq: str | list[str] | None = None,
+        input_ids: torch.Tensor | None = None,
+        **kwargs: T.Any,
+    ) -> torch.Tensor | dict[str, torch.Tensor]:
+        del kwargs
+        if input_ids is not None:
+            return input_ids
+        if seq is None:
+            raise ValueError("Pass either seq or input_ids for TTT.")
+        tokenized = self.tokenizer(seq, return_tensors="pt", padding=True)
+        return tokenized["input_ids"]
+
+    def _ttt_mask_token(self) -> int:
+        return int(self.tokenizer.mask_token_id)
+
+    def _ttt_padding_token(self) -> int:
+        return int(self.tokenizer.pad_token_id)
+
+    def _ttt_replacement_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
+        tokenizer = self.tokenizer
+        special_ids = set(tokenizer.all_special_ids)
+        vocab_size = int(self.config.vocab_size)
+        unknown_id = getattr(tokenizer, "unk_token_id", None)
+        if unknown_id is not None:
+            special_ids.add(int(unknown_id))
+
+        vocab: T.Mapping[str, T.Any] = {}
+        get_vocab = getattr(tokenizer, "get_vocab", None)
+        if callable(get_vocab):
+            vocab = get_vocab()
+        elif isinstance(getattr(tokenizer, "vocab", None), T.Mapping):
+            vocab = tokenizer.vocab
+        elif isinstance(getattr(tokenizer, "_token_to_id", None), T.Mapping):
+            vocab = tokenizer._token_to_id
+
+        ids: list[int] = []
+        convert = getattr(tokenizer, "convert_tokens_to_ids", None)
+        for amino_acid in _STANDARD_AMINO_ACIDS:
+            token_id = convert(amino_acid) if callable(convert) else vocab.get(amino_acid)
+            if (
+                isinstance(token_id, int)
+                and 0 <= token_id < vocab_size
+                and token_id not in special_ids
+                and token_id not in ids
+            ):
+                ids.append(token_id)
+        if not ids:
+            raise ValueError(
+                "TTT could not resolve any canonical amino-acid token IDs from the tokenizer; "
+                "refusing to sample arbitrary or reserved vocabulary entries."
+            )
+        return torch.tensor(ids, device=input_ids.device, dtype=input_ids.dtype)
+
+    def _ttt_predict_logits(
+        self,
+        batch: torch.Tensor | dict[str, torch.Tensor],
+        **kwargs: T.Any,
+    ) -> torch.Tensor:
+        del kwargs
+        if isinstance(batch, dict):
+            output = self(**batch)
+            return output.logits
+        attention_mask = batch.ne(self._ttt_padding_token())
+        output = self(input_ids=batch, attention_mask=attention_mask)
+        return output.logits
+
+    def _ttt_eval_step(
+        self,
+        step: int,
+        loss: float,
+        seq: str | list[str] | None = None,
+        input_ids: torch.Tensor | None = None,
+        **kwargs: T.Any,
+    ) -> tuple[dict[str, T.Any], float | None]:
+        del step, loss, seq, input_ids, kwargs
+        return {}, None
+
+    def _ttt_is_lora_target(
+        self,
+        name: str,
+        full_name: str,
+        module: nn.Module,
+        active: bool,
+        target_modules: tuple[str, ...] | None,
+    ) -> bool:
+        if not active:
+            return False
+        if isinstance(module, LoraInjectedLinear):
+            return False
+        if (
+            target_modules is not None
+            and name not in target_modules
+            and full_name not in target_modules
+        ):
+            return False
+        if isinstance(module, nn.Linear):
+            return True
+        if "weight" not in module._parameters:
+            return False
+        weight = module._parameters["weight"]
+        if weight is None or weight.ndim != 2:
+            return False
+        return "Linear" in module.__class__.__name__
+
+    def _ttt_inject_lora(self) -> int:
+        cfg = self.ttt_config
+        cfg.verify()
+        target_class = cfg.lora_target_replace_module
+        target_modules = cfg.lora_target_modules
+        wrapped = 0
+        generator = None
+        if cfg.seed is not None:
+            generator = torch.Generator(device="cpu")
+            generator.manual_seed(cfg.seed)
+
+        def inject(module: nn.Module, prefix: str, active: bool) -> None:
+            nonlocal wrapped
+            for name, child in list(module.named_children()):
+                full_name = f"{prefix}.{name}" if prefix else name
+                child_active = active
+                if target_class is not None:
+                    child_active = active or child.__class__.__name__ == target_class
+                if self._ttt_is_lora_target(name, full_name, child, child_active, target_modules):
+                    setattr(
+                        module,
+                        name,
+                        LoraInjectedLinear(
+                            child,
+                            rank=cfg.lora_rank,
+                            alpha=cfg.lora_alpha,
+                            generator=generator,
+                        ),
+                    )
+                    wrapped += 1
+                    continue
+                inject(child, full_name, child_active)
+
+        for trainable_module in self._ttt_get_trainable_modules():
+            inject(trainable_module, "", target_class is None)
+        if wrapped == 0:
+            raise ValueError("TTT LoRA injection did not find any target modules.")
+        return wrapped
+
+    def _ttt_lora_modules(self) -> list[LoraInjectedLinear]:
+        return [module for module in self.modules() if isinstance(module, LoraInjectedLinear)]
+
+    def _ttt_lora_parameters(self) -> list[nn.Parameter]:
+        params: list[nn.Parameter] = []
+        for module in self._ttt_lora_modules():
+            params.extend(module.lora_down.parameters())
+            params.extend(module.lora_up.parameters())
+        if not params:
+            raise RuntimeError("TTT has no LoRA parameters.")
+        return params
+
+    def _ttt_snapshot_lora_state(self) -> list[dict[str, torch.Tensor]]:
+        snapshot = []
+        for module in self._ttt_lora_modules():
+            snapshot.append(
+                {
+                    "lora_down.weight": module.lora_down.weight.detach().clone(),
+                    "lora_up.weight": module.lora_up.weight.detach().clone(),
+                }
+            )
+        if not snapshot:
+            raise RuntimeError("TTT has no LoRA state to snapshot.")
+        return snapshot
+
+    def _ttt_restore_lora_state(self, state: list[dict[str, torch.Tensor]]) -> None:
+        modules = self._ttt_lora_modules()
+        if len(modules) != len(state):
+            raise RuntimeError("TTT LoRA state/module count mismatch.")
+        with torch.no_grad():
+            for module, module_state in zip(modules, state, strict=True):
+                module.lora_down.weight.copy_(module_state["lora_down.weight"])
+                module.lora_up.weight.copy_(module_state["lora_up.weight"])
+
+    def _ttt_ensure_initialized(self) -> None:
+        if "_ttt_cfg" not in self.__dict__:
+            self.init_ttt()
+        if self._ttt_initialized:
+            return
+        self._ttt_inject_lora()
+        self._ttt_initialized = True
+
+    def ttt_reset(self) -> None:
+        self._ttt_ensure_initialized()
+        for module in self._ttt_lora_modules():
+            module.reset_lora_parameters()
+
+    def _ttt_serialized_contract(self) -> dict[str, T.Any]:
+        return {
+            "version": _TTT_SERIALIZATION_VERSION,
+            "initialized": bool(self._ttt_initialized),
+            "config": self.ttt_config.to_dict(),
+        }
+
+    def save_pretrained(self, save_directory: T.Any, *args: T.Any, **kwargs: T.Any) -> T.Any:
+        """Save initialized adapters, their reset baseline, and the TTT config.
+
+        Adapter injection changes the module tree, so the serialized config must
+        reconstruct that tree before Transformers loads the state dict. Models
+        whose own state-dict hooks omit their trainable TTT modules fail closed
+        instead of producing an artifact that cannot restore the adaptation.
+        """
+
+        if self._ttt_initialized:
+            state_keys = set(self.state_dict())
+            missing_adapter_keys = [
+                name
+                for name, _ in self.named_parameters()
+                if ".lora_" in name and name not in state_keys
+            ]
+            if missing_adapter_keys:
+                raise RuntimeError(
+                    "This model attaches TTT adapters to transient modules that its "
+                    "checkpoint excludes, so save_pretrained cannot persist the adapted "
+                    "state safely. Reset the model or use a model-specific adapter export."
+                )
+            self.config.fastplms_ttt = self._ttt_serialized_contract()
+        return super().save_pretrained(save_directory, *args, **kwargs)
+
+    def _ttt_make_optimizer(self) -> torch.optim.Optimizer:
+        cfg = self.ttt_config
+        params = self._ttt_lora_parameters()
+        if cfg.optimizer == "sgd":
+            return torch.optim.SGD(
+                params,
+                lr=cfg.lr,
+                momentum=cfg.momentum,
+                weight_decay=cfg.weight_decay,
+            )
+        return torch.optim.AdamW(params, lr=cfg.lr, weight_decay=cfg.weight_decay)
+
+    def _ttt_to_device(
+        self,
+        batch: torch.Tensor | dict[str, torch.Tensor],
+        device: torch.device,
+    ) -> torch.Tensor | dict[str, torch.Tensor]:
+        if isinstance(batch, dict):
+            return {name: tensor.to(device) for name, tensor in batch.items()}
+        return batch.to(device)
+
+    def _ttt_input_ids_from_batch(
+        self,
+        batch: torch.Tensor | dict[str, torch.Tensor],
+    ) -> torch.Tensor:
+        if isinstance(batch, dict):
+            return batch["input_ids"]
+        return batch
+
+    def _ttt_set_input_ids(
+        self,
+        batch: torch.Tensor | dict[str, torch.Tensor],
+        input_ids: torch.Tensor,
+    ) -> torch.Tensor | dict[str, torch.Tensor]:
+        if isinstance(batch, dict):
+            updated = dict(batch)
+            updated["input_ids"] = input_ids
+            return updated
+        return input_ids
+
+    def _ttt_non_special_mask(self, input_ids: torch.Tensor) -> torch.Tensor:
+        residue_ids = self._ttt_replacement_tokens(input_ids)
+        return torch.isin(input_ids, residue_ids)
+
+    def _ttt_validate_tokenized_batch(
+        self,
+        batch: torch.Tensor | dict[str, torch.Tensor],
+    ) -> None:
+        input_ids = self._ttt_input_ids_from_batch(batch)
+        if input_ids.ndim != 2 or input_ids.shape[0] == 0 or input_ids.shape[1] == 0:
+            raise ValueError(
+                "TTT input_ids must have non-empty shape (batch, sequence); got "
+                f"{tuple(input_ids.shape)}."
+            )
+
+        if str(getattr(self.config, "model_type", "")) == "dplm2":
+            tokenizer = self.tokenizer
+            token_to_id = getattr(tokenizer, "_token_to_id", {})
+            struct_cls_token = getattr(tokenizer, "struct_cls_token", None)
+            struct_boundary = token_to_id.get(struct_cls_token)
+            if struct_boundary is None:
+                raise ValueError(
+                    "DPLM2 TTT could not resolve the structure-token boundary safely."
+                )
+            pad_token = self._ttt_padding_token()
+            generic_aa_special_ids = torch.tensor(
+                [int(self.config.vocab_size) + offset for offset in range(4)],
+                device=input_ids.device,
+                dtype=input_ids.dtype,
+            )
+            is_structure = input_ids.ge(int(struct_boundary)) & input_ids.ne(pad_token)
+            is_structure &= ~torch.isin(input_ids, generic_aa_special_ids)
+            if bool(is_structure.any()):
+                raise ValueError(
+                    "DPLM2 TTT currently supports amino-acid-only inputs. Packed or "
+                    "structure-token inputs require a modality-specific corruption objective."
+                )
+
+            if isinstance(batch, dict) and "type_ids" in batch:
+                type_ids = batch["type_ids"]
+                attention_mask = batch.get("attention_mask", input_ids.ne(pad_token)).bool()
+                if bool(((type_ids == int(self.config.struct_type)) & attention_mask).any()):
+                    raise ValueError(
+                        "DPLM2 TTT currently supports amino-acid-only inputs; structure "
+                        "type_ids are not accepted."
+                    )
+
+        if not bool(self._ttt_non_special_mask(input_ids).any()):
+            raise ValueError(
+                "TTT input contains no trainable biological residue tokens after excluding "
+                "padding, boundary, mask, and reserved tokens."
+            )
+
+    def _ttt_sample_crop(
+        self,
+        batch: torch.Tensor | dict[str, torch.Tensor],
+        generator: torch.Generator,
+    ) -> torch.Tensor | dict[str, torch.Tensor]:
+        input_ids = self._ttt_input_ids_from_batch(batch)
+        cfg = self.ttt_config
+        if input_ids.shape[1] <= cfg.crop_size:
+            return batch
+        position_has_residue = self._ttt_non_special_mask(input_ids).any(dim=0).to(torch.int64)
+        prefix = F.pad(position_has_residue.cumsum(dim=0), (1, 0))
+        window_counts = prefix[cfg.crop_size :] - prefix[: -cfg.crop_size]
+        valid_starts = torch.where(window_counts > 0)[0]
+        if valid_starts.numel() == 0:
+            raise ValueError("TTT could not find a crop containing a biological residue token.")
+        selected = torch.randint(
+            valid_starts.numel(),
+            (1,),
+            generator=generator,
+            device=input_ids.device,
+        )
+        start = int(valid_starts[selected].item())
+        end = start + cfg.crop_size
+        if isinstance(batch, dict):
+            cropped = {}
+            for name, tensor in batch.items():
+                if tensor.ndim >= 2 and tensor.shape[1] == input_ids.shape[1]:
+                    cropped[name] = tensor[:, start:end]
+                else:
+                    cropped[name] = tensor
+            return cropped
+        return input_ids[:, start:end]
+
+    def _ttt_sample_batch(
+        self,
+        tokenized: torch.Tensor | dict[str, torch.Tensor],
+        generator: torch.Generator,
+    ) -> tuple[torch.Tensor | dict[str, torch.Tensor], torch.Tensor]:
+        cfg = self.ttt_config
+        batch = self._ttt_sample_crop(tokenized, generator)
+        input_ids = self._ttt_input_ids_from_batch(batch)
+        row_has_residue = self._ttt_non_special_mask(input_ids).any(dim=1)
+        eligible_rows = torch.where(row_has_residue)[0]
+        if eligible_rows.numel() == 0:
+            raise ValueError(
+                "TTT sampled batch contains no trainable biological residue tokens."
+            )
+        sampled_row_indices = torch.randint(
+            eligible_rows.numel(),
+            (cfg.batch_size,),
+            generator=generator,
+            device=input_ids.device,
+        )
+        rows = eligible_rows[sampled_row_indices]
+        if isinstance(batch, dict):
+            sampled: torch.Tensor | dict[str, torch.Tensor] = {}
+            for name, tensor in batch.items():
+                if tensor.ndim >= 1 and tensor.shape[0] == input_ids.shape[0]:
+                    sampled[name] = tensor.index_select(0, rows)
+                else:
+                    sampled[name] = tensor
+        else:
+            sampled = input_ids.index_select(0, rows)
+
+        sampled_ids = self._ttt_input_ids_from_batch(sampled)
+        labels = sampled_ids.clone()
+        non_special = self._ttt_non_special_mask(sampled_ids)
+        label_mask = torch.zeros_like(non_special)
+        for row_idx in range(sampled_ids.shape[0]):
+            candidate_positions = torch.where(non_special[row_idx])[0]
+            if candidate_positions.numel() == 0:
+                continue
+            num_mask = max(1, round(candidate_positions.numel() * cfg.mask_ratio))
+            order = torch.randperm(
+                candidate_positions.numel(),
+                generator=generator,
+                device=sampled_ids.device,
+            )
+            chosen = candidate_positions[order[:num_mask]]
+            label_mask[row_idx, chosen] = True
+        labels = labels.masked_fill(~label_mask, -100)
+
+        masked_ids = sampled_ids.clone()
+        chosen_positions = torch.where(label_mask)
+        if chosen_positions[0].numel() > 0:
+            random_values = torch.rand(
+                chosen_positions[0].shape,
+                generator=generator,
+                device=sampled_ids.device,
+            )
+            leave = random_values < cfg.bert_leave_prob
+            replace = (random_values >= cfg.bert_leave_prob) & (
+                random_values < cfg.bert_leave_prob + cfg.bert_replace_prob
+            )
+            mask = ~(leave | replace)
+            if mask.any():
+                masked_ids[
+                    chosen_positions[0][mask],
+                    chosen_positions[1][mask],
+                ] = self._ttt_mask_token()
+            if replace.any():
+                replacement_tokens = self._ttt_replacement_tokens(sampled_ids)
+                replacement_idx = torch.randint(
+                    replacement_tokens.shape[0],
+                    (int(replace.sum().item()),),
+                    generator=generator,
+                    device=sampled_ids.device,
+                )
+                masked_ids[
+                    chosen_positions[0][replace],
+                    chosen_positions[1][replace],
+                ] = replacement_tokens[replacement_idx]
+
+        return self._ttt_set_input_ids(sampled, masked_ids), labels
+
+    @contextlib.contextmanager
+    def _ttt_seed_scope(self, seed: int | None) -> T.Iterator[None]:
+        if seed is None:
+            yield
+            return
+        cuda_devices = sorted(
+            {
+                parameter.device.index
+                for parameter in self.parameters()
+                if parameter.device.type == "cuda" and parameter.device.index is not None
+            }
+        )
+        with torch.random.fork_rng(devices=cuda_devices):
+            torch.random.default_generator.manual_seed(seed)
+            for device_index in cuda_devices:
+                with torch.cuda.device(device_index):
+                    torch.cuda.manual_seed(seed)
+            yield
+
+    def ttt(
+        self,
+        seq: str | list[str] | None = None,
+        input_ids: torch.Tensor | None = None,
+        ttt_config: TTTConfig | T.Mapping[str, T.Any] | None = None,
+        **kwargs: T.Any,
+    ) -> dict[str, T.Any]:
+        if ttt_config is not None:
+            if "_ttt_initialized" in self.__dict__ and self._ttt_initialized:
+                next_cfg = self.ttt_config.merged(ttt_config)
+                current_cfg = self.ttt_config
+                if next_cfg.lora_rank != current_cfg.lora_rank:
+                    raise ValueError(
+                        "Changing lora_rank after TTT initialization is not supported."
+                    )
+                if next_cfg.lora_alpha != current_cfg.lora_alpha:
+                    raise ValueError(
+                        "Changing lora_alpha after TTT initialization is not supported."
+                    )
+                if (
+                    next_cfg.lora_target_replace_module
+                    != current_cfg.lora_target_replace_module
+                ):
+                    raise ValueError(
+                        "Changing LoRA target class after TTT initialization is not supported."
+                    )
+                if next_cfg.lora_target_modules != current_cfg.lora_target_modules:
+                    raise ValueError(
+                        "Changing LoRA target modules after TTT initialization is not supported."
+                    )
+                self._ttt_cfg = next_cfg
+            else:
+                # Family constructors preconfigure the attention class that may
+                # receive LoRA adapters. A first-call mapping changes only the
+                # requested fields; rebuilding from TTTConfig defaults here
+                # would erase that family target immediately before injection.
+                self._ttt_cfg = self.ttt_config.merged(ttt_config)
+                self._ttt_cfg.verify()
+
+        cfg = self.ttt_config
+        device = next(self.parameters()).device
+        tokenized = self._ttt_tokenize(seq=seq, input_ids=input_ids, **kwargs)
+        tokenized = self._ttt_to_device(tokenized, device)
+        self._ttt_validate_tokenized_batch(tokenized)
+        self._ttt_ensure_initialized()
+        if cfg.initial_state_reset:
+            self.ttt_reset()
+
+        generator_device = device if device.type == "cuda" else torch.device("cpu")
+        generator = torch.Generator(device=generator_device)
+        if cfg.seed is not None:
+            generator.manual_seed(cfg.seed)
+
+        module_modes = {module: module.training for module in self.modules()}
+        requires_grad = {param: param.requires_grad for param in self.parameters()}
+        losses: list[float] = []
+        step_metrics: list[dict[str, T.Any]] = []
+        best_state: list[dict[str, torch.Tensor]] | None = None
+        best_metric: float | None = None
+        best_step = 0
+
+        with self._ttt_seed_scope(cfg.seed):
+            try:
+                self.train()
+                for param in self.parameters():
+                    param.requires_grad_(False)
+                for param in self._ttt_lora_parameters():
+                    param.requires_grad_(True)
+
+                optimizer = self._ttt_make_optimizer()
+                optimizer.zero_grad(set_to_none=True)
+                total_micro_steps = cfg.steps * cfg.ags
+                for micro_step in range(total_micro_steps):
+                    batch, labels = self._ttt_sample_batch(tokenized, generator)
+                    if not bool(labels.ne(-100).any()):
+                        raise RuntimeError(
+                            "TTT produced an all-ignored label batch; refusing a NaN update."
+                        )
+                    logits = self._ttt_predict_logits(batch, **kwargs)
+                    labels = labels.to(device=logits.device)
+                    loss = F.cross_entropy(
+                        logits.reshape(-1, logits.shape[-1]),
+                        labels.reshape(-1),
+                        ignore_index=-100,
+                    )
+                    if not bool(torch.isfinite(loss)):
+                        raise FloatingPointError(
+                            f"TTT loss is non-finite at micro-step {micro_step + 1}."
+                        )
+                    (loss / cfg.ags).backward()
+                    if (micro_step + 1) % cfg.ags != 0:
+                        continue
+
+                    if cfg.gradient_clip:
+                        torch.nn.utils.clip_grad_norm_(
+                            self._ttt_lora_parameters(),
+                            cfg.gradient_clip_max_norm,
+                        )
+                    optimizer.step()
+                    optimizer.zero_grad(set_to_none=True)
+                    step = (micro_step + 1) // cfg.ags
+                    loss_value = float(loss.detach().item())
+                    losses.append(loss_value)
+                    if cfg.eval_each_step:
+                        metrics, metric = self._ttt_eval_step(
+                            step=step,
+                            loss=loss_value,
+                            seq=seq,
+                            input_ids=input_ids,
+                            **kwargs,
+                        )
+                        if len(metrics) > 0:
+                            step_metrics.append(metrics)
+                        if metric is not None and (best_metric is None or metric > best_metric):
+                            best_metric = metric
+                            best_step = step
+                            best_state = self._ttt_snapshot_lora_state()
+
+                if cfg.automatic_best_state_reset and best_state is not None:
+                    self._ttt_restore_lora_state(best_state)
+            finally:
+                for param, value in requires_grad.items():
+                    param.requires_grad_(value)
+                for module, training in module_modes.items():
+                    module.train(training)
+
+        return {
+            "losses": losses,
+            "step_metrics": step_metrics,
+            "best_step": best_step,
+            "best_metric": best_metric,
+        }

fastplms/registry.py

@@ -0,0 +1,1486 @@
+"""Typed access to the FastPLMs model and provenance manifest.
+
+The registry is intentionally independent of Torch and Transformers. Tooling can
+therefore inspect supported checkpoints, licenses, and reference sources without
+initializing a model runtime or downloading any files.
+"""
+
+from __future__ import annotations
+
+import re
+import tomllib
+from collections.abc import Iterator, Mapping
+from dataclasses import dataclass
+from functools import lru_cache
+from importlib import resources
+from pathlib import Path, PurePosixPath, PureWindowsPath
+from types import MappingProxyType
+from typing import Any, Literal, cast
+from urllib.parse import urlparse
+
+_HEX_RE = re.compile(r"^[0-9a-f]+$")
+_IDENTIFIER_RE = re.compile(r"^[a-z0-9][a-z0-9_-]*$")
+_HUB_LICENSE_NAME_RE = re.compile(r"[^a-z0-9.]+")
+_WINDOWS_INVALID_PATH_CHARACTERS = frozenset('<>:"|?*')
+_WINDOWS_RESERVED_PATH_NAMES = frozenset(
+    {"AUX", "CON", "NUL", "PRN"}
+    | {f"COM{index}" for index in range(1, 10)}
+    | {f"LPT{index}" for index in range(1, 10)}
+)
+_REPOSITORY_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9_.-]*/[A-Za-z0-9][A-Za-z0-9_.-]*$")
+_REFERENCE_CONTAINER_RE = re.compile(r"^reference-[a-z0-9]+(?:-[a-z0-9]+)*$")
+_REFERENCE_ADAPTER_RE = re.compile(
+    r"^tests\.parity\.support\.reference_adapters\.[a-z_][a-z0-9_]*$"
+)
+_DOCUMENTATION_FRAGMENT_RE = re.compile(r"^[a-z0-9]+(?:-[a-z0-9]+)*$")
+_ALLOWED_ATTENTION = frozenset(
+    {"eager", "sdpa", "flex_attention", "flash_attention_2", "flash_attention_3"}
+)
+_ALLOWED_DTYPES = frozenset({"float32", "bfloat16"})
+_ALLOWED_PRECISIONS = frozenset({"default", "auto", "fp32", "bf16", "fp8"})
+_ALLOWED_BF16_EXECUTIONS = frozenset({"static_parameters", "fp32_parameters_autocast"})
+HUB_LICENSE_IDENTIFIERS = frozenset({"mit", "apache-2.0", "cc-by-nc-sa-4.0", "other"})
+_ALLOWED_TOKENIZER_MODES = frozenset({"tokenizer", "sequence", "structure"})
+_ALLOWED_SIZE_CATEGORIES = frozenset({"small", "medium", "large", "xlarge", "structure"})
+RuntimeExtra = Literal["core", "structure"]
+TestTier = Literal["check", "compliance", "structure", "feature", "artifact", "benchmark"]
+VramTier = Literal["sequence", "large-sequence", "structure", "structure-6b"]
+GenerationContract = Literal["not_applicable", "required", "official_unavailable"]
+RuntimeAssetTrustKind = Literal["hash_pinned_pickle"]
+Bf16Execution = Literal["static_parameters", "fp32_parameters_autocast"]
+DtypeName = Literal["float32", "bfloat16"]
+_ALLOWED_EXTRAS = frozenset({"core", "structure"})
+_ALLOWED_TEST_TIERS = frozenset(
+    {"check", "compliance", "structure", "feature", "artifact", "benchmark"}
+)
+_ALLOWED_VRAM_TIERS = frozenset({"sequence", "large-sequence", "structure", "structure-6b"})
+_ALLOWED_GENERATION_CONTRACTS = frozenset({"not_applicable", "required", "official_unavailable"})
+_ALLOWED_RUNTIME_ASSET_TRUST_KINDS = frozenset({"hash_pinned_pickle"})
+_ALLOWED_RUNTIME_ASSET_OFFLINE_BEHAVIORS = frozenset({"requires_cached_verified_file"})
+_ALLOWED_AUTO_CLASSES = frozenset(
+    {
+        "AutoConfig",
+        "AutoModel",
+        "AutoModelForMaskedLM",
+        "AutoModelForProteinFolding",
+        "AutoModelForSequenceClassification",
+        "AutoModelForSeq2SeqLM",
+        "AutoModelForTokenClassification",
+    }
+)
+_WEIGHT_SUFFIXES = (".bin", ".ckpt", ".pt", ".pth", ".safetensors")
+_ALLOWED_ORACLE_ASSET_ROLES = frozenset({"weights", "contact_regression"})
+_FAIR_ESM_ASSET_HOST = "dl.fbaipublicfiles.com"
+_ROOT_FIELDS = frozenset(
+    {
+        "schema_version",
+        "legal_files",
+        "attention_kernels",
+        "upstreams",
+        "families",
+        "models",
+        "runtime_assets",
+    }
+)
+_UPSTREAM_FIELDS = frozenset(
+    {
+        "id",
+        "path",
+        "url",
+        "revision",
+        "license",
+        "license_files",
+        "license_digests",
+        "distribution_files",
+    }
+)
+_FAMILY_FIELDS = frozenset(
+    {
+        "architecture",
+        "upstreams",
+        "tokenizer_mode",
+        "public_input",
+        "extra",
+        "reference_container",
+        "reference_adapter",
+        "attention",
+        "dtypes",
+        "bf16_execution",
+        "precisions",
+        "experimental_precisions",
+        "vram_tier",
+        "checkpoint_license",
+        "hub_license",
+        "hub_license_name",
+        "hub_license_link",
+        "state_transform",
+        "conversion_provenance",
+        "representative",
+        "documentation",
+        "test_tiers",
+        "runtime_paths",
+        "requires_complete_weight_publication",
+        "weights_publication_allowed",
+        "auto_map",
+        "tokenizer_class",
+        "backbone_model",
+    }
+)
+_MODEL_FIELDS = frozenset(
+    {
+        "id",
+        "family",
+        "size_category",
+        "generation_contract",
+        "fast_repo",
+        "fast_revision",
+        "fast_files",
+        "fast_unresolved_files",
+        "official_repo",
+        "official_revision",
+        "official_files",
+        "official_unresolved_files",
+        "oracle_assets",
+        "official_golden",
+        "artifact_source",
+        "canonical_state_sha256",
+        "tokenizer_source",
+        "auto_map",
+        "notes",
+        "msa_conditioning",
+    }
+)
+_RUNTIME_ASSET_FIELDS = frozenset(
+    {
+        "id",
+        "repository",
+        "revision",
+        "path",
+        "sha256",
+        "size",
+        "consumer_family",
+        "trust_kind",
+        "license",
+        "offline_behavior",
+    }
+)
+
+
+class RegistryError(ValueError):
+    """Raised when the model manifest is incomplete or internally inconsistent."""
+
+
+def _portable_relative_path(value: str, context: str) -> PurePosixPath:
+    """Return one normalized cross-platform relative path or fail closed."""
+
+    posix = PurePosixPath(value)
+    windows = PureWindowsPath(value)
+    unsafe_windows_part = any(
+        part.rstrip(" .") != part
+        or part.split(".", maxsplit=1)[0].upper() in _WINDOWS_RESERVED_PATH_NAMES
+        or any(
+            ord(character) < 32 or character in _WINDOWS_INVALID_PATH_CHARACTERS
+            for character in part
+        )
+        for part in posix.parts
+    )
+    if (
+        not value
+        or not posix.parts
+        or posix == PurePosixPath(".")
+        or posix.is_absolute()
+        or windows.is_absolute()
+        or windows.drive
+        or "\\" in value
+        or "." in posix.parts
+        or ".." in posix.parts
+        or value != posix.as_posix()
+        or any(
+            part.lower() in {".git", ".cache", "__pycache__"}
+            for part in posix.parts
+        )
+        or unsafe_windows_part
+    ):
+        raise RegistryError(f"{context} is not portable: {value!r}")
+    return posix
+
+
+@dataclass(frozen=True, slots=True)
+class FileDigest:
+    """Expected content identity for one pinned file."""
+
+    path: str
+    algorithm: str
+    digest: str
+
+    @classmethod
+    def parse(cls, value: str) -> FileDigest:
+        try:
+            path, encoded_digest = value.split("=", maxsplit=1)
+            algorithm, digest = encoded_digest.split(":", maxsplit=1)
+        except ValueError as error:
+            raise RegistryError("File digests must use '<path>=<algorithm>:<digest>'.") from error
+
+        _portable_relative_path(path, "Checkpoint file path")
+
+        expected_length = {"git-sha1": 40, "sha256": 64}.get(algorithm)
+        if expected_length is None:
+            raise RegistryError(f"Unsupported file digest algorithm: {algorithm!r}")
+        if len(digest) != expected_length or _HEX_RE.fullmatch(digest) is None:
+            raise RegistryError(f"Invalid {algorithm} digest for {path!r}: {digest!r}")
+        return cls(path=path, algorithm=algorithm, digest=digest)
+
+    @property
+    def encoded(self) -> str:
+        return f"{self.algorithm}:{self.digest}"
+
+
+@dataclass(frozen=True, slots=True)
+class CheckpointSource:
+    """One immutable Hugging Face repository snapshot."""
+
+    repo_id: str
+    revision: str
+    files: tuple[FileDigest, ...]
+    unresolved_files: tuple[str, ...] = ()
+
+    @property
+    def file_map(self) -> Mapping[str, FileDigest]:
+        return MappingProxyType({item.path: item for item in self.files})
+
+
+@dataclass(frozen=True, slots=True)
+class OracleAsset:
+    """Hash-pinned external file required by a native parity oracle."""
+
+    role: str
+    path: str
+    url: str
+    sha256: str
+    size: int
+
+
+@dataclass(frozen=True, slots=True)
+class RuntimeAsset:
+    """Immutable runtime data with an explicit deserialization trust boundary."""
+
+    id: str
+    repository: str
+    revision: str
+    path: str
+    sha256: str
+    size: int
+    consumer_family: str
+    trust_kind: RuntimeAssetTrustKind
+    license_expression: str
+    offline_behavior: str
+
+
+@dataclass(frozen=True, slots=True)
+class OfficialGolden:
+    """Hash-pinned official output bundle required by the check tier."""
+
+    metadata: FileDigest
+    tensors: FileDigest
+
+
+@dataclass(frozen=True, slots=True)
+class UpstreamSource:
+    """Pinned official implementation used as a parity oracle."""
+
+    id: str
+    path: str
+    url: str
+    revision: str
+    license_expression: str
+    license_files: tuple[str, ...]
+    license_digests: tuple[FileDigest, ...] = ()
+    distribution_files: tuple[FileDigest, ...] = ()
+
+
+@dataclass(frozen=True, slots=True)
+class AttentionKernelSpec:
+    """Immutable Hugging Face kernel used by one attention backend."""
+
+    implementation: str
+    repository: str
+    revision: str
+    version: int
+    expected_variant: str
+    dtypes: tuple[DtypeName, ...]
+
+
+@dataclass(frozen=True, slots=True)
+class ModelFamily:
+    """Shared runtime and compliance contract for one architecture family."""
+
+    id: str
+    architecture: str
+    upstreams: tuple[str, ...]
+    tokenizer_mode: str
+    public_input: str
+    extra: RuntimeExtra
+    reference_container: str
+    reference_adapter: str
+    attention: tuple[str, ...]
+    dtypes: tuple[DtypeName, ...]
+    bf16_execution: Bf16Execution
+    precisions: tuple[str, ...]
+    vram_tier: VramTier
+    checkpoint_license: str
+    hub_license: str
+    state_transform: str
+    representative: str
+    documentation: str
+    test_tiers: tuple[TestTier, ...]
+    runtime_paths: tuple[str, ...]
+    auto_map_items: tuple[tuple[str, str], ...]
+    requires_complete_weight_publication: bool = False
+    weights_publication_allowed: bool = False
+    experimental_precisions: tuple[str, ...] = ()
+    tokenizer_class: str | None = None
+    hub_license_name: str | None = None
+    hub_license_link: str | None = None
+    conversion_provenance: str = ""
+    backbone_model: str | None = None
+
+    @property
+    def auto_map(self) -> Mapping[str, str]:
+        return MappingProxyType(dict(self.auto_map_items))
+
+    @property
+    def hub_license_metadata(self) -> Mapping[str, str]:
+        """Return valid Hugging Face model-card license fields."""
+
+        metadata = {"license": self.hub_license}
+        if self.hub_license_name is not None:
+            # Hugging Face validates custom license names as lowercase slugs,
+            # while the manifest retains the reader-facing display name used
+            # in generated prose.
+            metadata["license_name"] = _HUB_LICENSE_NAME_RE.sub(
+                "-",
+                self.hub_license_name.lower(),
+            ).strip("-.")
+        if self.hub_license_link is not None:
+            metadata["license_link"] = self.hub_license_link
+        return MappingProxyType(metadata)
+
+    @property
+    def stable_precisions(self) -> tuple[str, ...]:
+        """Return precision policies covered by the release contract."""
+
+        experimental = set(self.experimental_precisions)
+        return tuple(precision for precision in self.precisions if precision not in experimental)
+
+
+@dataclass(frozen=True, slots=True)
+class ModelSpec:
+    """Complete immutable source and runtime contract for one checkpoint."""
+
+    id: str
+    family: ModelFamily
+    fast: CheckpointSource
+    official: CheckpointSource
+    size_category: str
+    generation_contract: GenerationContract = "not_applicable"
+    oracle_assets: tuple[OracleAsset, ...] = ()
+    official_golden: OfficialGolden | None = None
+    artifact_source: str = "fast"
+    canonical_state_sha256: str | None = None
+    tokenizer_source_id: str | None = None
+    auto_map_items: tuple[tuple[str, str], ...] = ()
+    notes: str = ""
+    msa_conditioning: bool | None = None
+
+    @property
+    def is_deep_reference(self) -> bool:
+        return self.id == self.family.representative
+
+    @property
+    def auto_map(self) -> Mapping[str, str]:
+        if self.auto_map_items:
+            return MappingProxyType(dict(self.auto_map_items))
+        return self.family.auto_map
+
+    @property
+    def artifact_checkpoint(self) -> CheckpointSource:
+        """Return the checkpoint selected for local artifact construction."""
+
+        return self.fast if self.artifact_source == "fast" else self.official
+
+    @property
+    def oracle_asset_map(self) -> Mapping[str, OracleAsset]:
+        """Return native oracle assets keyed by their declared role."""
+
+        return MappingProxyType({asset.role: asset for asset in self.oracle_assets})
+
+
+class ModelRegistry(Mapping[str, ModelSpec]):
+    """Validated mapping of model IDs to typed model specifications."""
+
+    def __init__(
+        self,
+        *,
+        schema_version: int,
+        upstreams: Mapping[str, UpstreamSource],
+        families: Mapping[str, ModelFamily],
+        models: Mapping[str, ModelSpec],
+        runtime_assets: Mapping[str, RuntimeAsset] = MappingProxyType({}),
+        attention_kernels: Mapping[str, AttentionKernelSpec] = MappingProxyType({}),
+        legal_files: tuple[FileDigest, ...] = (),
+    ) -> None:
+        self.schema_version = schema_version
+        self.upstreams = MappingProxyType(dict(upstreams))
+        self.attention_kernels = MappingProxyType(dict(attention_kernels))
+        self.families = MappingProxyType(dict(families))
+        self._models = MappingProxyType(dict(models))
+        self.runtime_assets = MappingProxyType(dict(runtime_assets))
+        self.legal_files = legal_files
+
+    def __getitem__(self, key: str) -> ModelSpec:
+        return self._models[key]
+
+    def __iter__(self) -> Iterator[str]:
+        return iter(self._models)
+
+    def __len__(self) -> int:
+        return len(self._models)
+
+    def by_family(self, family_id: str) -> tuple[ModelSpec, ...]:
+        if family_id not in self.families:
+            raise KeyError(family_id)
+        return tuple(model for model in self._models.values() if model.family.id == family_id)
+
+    def supported_attention_dtypes(
+        self,
+        family_id: str,
+        implementation: str,
+    ) -> tuple[DtypeName, ...]:
+        """Return manifest-supported dtypes for one family/backend pair."""
+
+        family = self.families[family_id]
+        if implementation not in family.attention:
+            raise KeyError(
+                f"Family {family_id!r} does not advertise attention backend "
+                f"{implementation!r}."
+            )
+        kernel = self.attention_kernels.get(implementation)
+        if kernel is None:
+            return family.dtypes
+        return tuple(dtype for dtype in family.dtypes if dtype in kernel.dtypes)
+
+    def require_resolved(self, model_id: str | None = None) -> None:
+        """Fail release validation when required file identities remain unresolved."""
+
+        selected = self._models.values() if model_id is None else (self._models[model_id],)
+        unresolved: list[str] = []
+        for model in selected:
+            for label, checkpoint in (("fast", model.fast), ("official", model.official)):
+                for path in checkpoint.unresolved_files:
+                    unresolved.append(f"{model.id}.{label}:{path}")
+        if unresolved:
+            detail = ", ".join(unresolved)
+            raise RegistryError(f"Release provenance is unresolved: {detail}")
+
+
+def _reject_unknown_fields(
+    table: Mapping[str, Any],
+    allowed: frozenset[str],
+    context: str,
+) -> None:
+    unknown = sorted(set(table).difference(allowed))
+    if unknown:
+        raise RegistryError(f"{context} contains unknown fields: {unknown}.")
+
+
+def _require_str(table: Mapping[str, Any], key: str, context: str) -> str:
+    value = table.get(key)
+    if not isinstance(value, str) or not value.strip():
+        raise RegistryError(f"{context}.{key} must be a non-empty string.")
+    return value
+
+
+def _require_enum(
+    table: Mapping[str, Any],
+    key: str,
+    context: str,
+    allowed: frozenset[str],
+) -> str:
+    value = _require_str(table, key, context)
+    if value not in allowed:
+        raise RegistryError(
+            f"{context}.{key} must be one of {sorted(allowed)}; received {value!r}."
+        )
+    return value
+
+
+def _parse_reference_container(table: Mapping[str, Any], context: str) -> str:
+    value = _require_str(table, "reference_container", context)
+    if _REFERENCE_CONTAINER_RE.fullmatch(value) is None:
+        raise RegistryError(
+            f"{context}.reference_container must be a portable 'reference-<name>' target."
+        )
+    return value
+
+
+def _parse_reference_adapter(table: Mapping[str, Any], context: str) -> str:
+    value = _require_str(table, "reference_adapter", context)
+    if _REFERENCE_ADAPTER_RE.fullmatch(value) is None:
+        raise RegistryError(
+            f"{context}.reference_adapter must name one module under "
+            "tests.parity.support.reference_adapters."
+        )
+    return value
+
+
+def _parse_documentation_path(table: Mapping[str, Any], context: str) -> str:
+    value = _require_str(table, "documentation", context)
+    if value.count("#") > 1 or "\\" in value:
+        raise RegistryError(f"{context}.documentation must be a portable documentation path.")
+    raw_path, separator, fragment = value.partition("#")
+    path = PurePosixPath(raw_path)
+    if (
+        path.is_absolute()
+        or ".." in path.parts
+        or len(path.parts) < 2
+        or path.parts[0] != "docs"
+        or path.suffix != ".md"
+        or path.as_posix() != raw_path
+    ):
+        raise RegistryError(
+            f"{context}.documentation must reference a normalized Markdown file under docs/."
+        )
+    if separator and _DOCUMENTATION_FRAGMENT_RE.fullmatch(fragment) is None:
+        raise RegistryError(f"{context}.documentation has an invalid heading fragment.")
+    return value
+
+
+def _require_str_list(table: Mapping[str, Any], key: str, context: str) -> tuple[str, ...]:
+    value = table.get(key)
+    if not isinstance(value, list) or not value or any(not isinstance(item, str) for item in value):
+        raise RegistryError(f"{context}.{key} must be a non-empty string array.")
+    result = tuple(value)
+    if len(set(result)) != len(result):
+        raise RegistryError(f"{context}.{key} contains duplicate values.")
+    return result
+
+
+def _optional_str_list(table: Mapping[str, Any], key: str, context: str) -> tuple[str, ...]:
+    value = table.get(key, [])
+    if not isinstance(value, list) or any(not isinstance(item, str) for item in value):
+        raise RegistryError(f"{context}.{key} must be a string array.")
+    result = tuple(value)
+    if len(set(result)) != len(result):
+        raise RegistryError(f"{context}.{key} contains duplicate values.")
+    return result
+
+
+def _optional_str(table: Mapping[str, Any], key: str, context: str) -> str | None:
+    value = table.get(key)
+    if value is None:
+        return None
+    if (
+        not isinstance(value, str)
+        or not value.strip()
+        or value != value.strip()
+        or "\n" in value
+        or "\r" in value
+    ):
+        raise RegistryError(f"{context}.{key} must be a non-empty single-line string.")
+    return value
+
+
+def _parse_hub_license(
+    table: Mapping[str, Any],
+    *,
+    checkpoint_license: str,
+    context: str,
+) -> tuple[str, str | None, str | None]:
+    expected_fields = {"hub_license", "hub_license_name", "hub_license_link"}
+    unknown_fields = sorted(
+        key for key in table if key.startswith("hub_") and key not in expected_fields
+    )
+    if unknown_fields:
+        raise RegistryError(f"{context} contains unsupported Hub license fields: {unknown_fields}.")
+    identifier = _require_str(table, "hub_license", context)
+    if identifier not in HUB_LICENSE_IDENTIFIERS:
+        raise RegistryError(
+            f"{context}.hub_license must be a supported Hugging Face license identifier."
+        )
+    expected_identifier: str | None = None
+    for prefix, candidate in (
+        ("MIT", "mit"),
+        ("Apache-2.0", "apache-2.0"),
+        ("CC-BY-NC-SA-4.0", "cc-by-nc-sa-4.0"),
+        ("Profluent-E1-Agreement", "other"),
+        ("Unresolved", "other"),
+    ):
+        if checkpoint_license.startswith(prefix):
+            expected_identifier = candidate
+            break
+    if expected_identifier is None:
+        raise RegistryError(
+            f"{context}.checkpoint_license has no declared Hugging Face identifier mapping."
+        )
+    if identifier != expected_identifier:
+        raise RegistryError(
+            f"{context}.hub_license must be {expected_identifier!r} for "
+            f"checkpoint terms {checkpoint_license!r}."
+        )
+
+    name = _optional_str(table, "hub_license_name", context)
+    link = _optional_str(table, "hub_license_link", context)
+    if identifier != "other":
+        if name is not None or link is not None:
+            raise RegistryError(
+                f"{context} may define hub_license_name and hub_license_link only "
+                "when hub_license='other'."
+            )
+        return identifier, None, None
+    if name is None or link is None:
+        raise RegistryError(
+            f"{context} must define hub_license_name and hub_license_link when hub_license='other'."
+        )
+    parsed_link = urlparse(link)
+    if (
+        parsed_link.scheme != "https"
+        or not parsed_link.netloc
+        or not parsed_link.path
+        or parsed_link.username is not None
+        or parsed_link.password is not None
+    ):
+        raise RegistryError(f"{context}.hub_license_link must be an absolute HTTPS URL.")
+    return identifier, name, link
+
+
+def _require_digest_list(
+    table: Mapping[str, Any], key: str, context: str
+) -> tuple[FileDigest, ...]:
+    encoded = _require_str_list(table, key, context)
+    result = tuple(FileDigest.parse(value) for value in encoded)
+    paths = [item.path for item in result]
+    if len(paths) != len(set(paths)):
+        raise RegistryError(f"{context}.{key} contains duplicate paths.")
+    return result
+
+
+def _validate_revision(revision: str, context: str) -> None:
+    if len(revision) != 40 or _HEX_RE.fullmatch(revision) is None:
+        raise RegistryError(f"{context} must be an immutable 40-character commit revision.")
+
+
+def _parse_checkpoint(table: Mapping[str, Any], prefix: str, context: str) -> CheckpointSource:
+    repo_id = _require_str(table, f"{prefix}_repo", context)
+    if _REPOSITORY_ID_RE.fullmatch(repo_id) is None:
+        raise RegistryError(f"{context}.{prefix}_repo must be a Hugging Face repository ID.")
+    revision = _require_str(table, f"{prefix}_revision", context)
+    _validate_revision(revision, f"{context}.{prefix}_revision")
+    encoded_files = _require_str_list(table, f"{prefix}_files", context)
+    files = tuple(FileDigest.parse(value) for value in encoded_files)
+    paths = [item.path for item in files]
+    if len(paths) != len(set(paths)):
+        raise RegistryError(f"{context}.{prefix}_files contains duplicate paths.")
+    if not any(item.path.endswith(_WEIGHT_SUFFIXES) for item in files):
+        raise RegistryError(f"{context}.{prefix}_files does not identify a weight file.")
+    unresolved_files = _optional_str_list(table, f"{prefix}_unresolved_files", context)
+    for unresolved_path in unresolved_files:
+        _portable_relative_path(unresolved_path, "Unresolved checkpoint path")
+        if unresolved_path in paths:
+            raise RegistryError(
+                f"{context}.{prefix} marks {unresolved_path!r} both resolved and unresolved."
+            )
+    return CheckpointSource(
+        repo_id=repo_id,
+        revision=revision,
+        files=files,
+        unresolved_files=unresolved_files,
+    )
+
+
+def _parse_oracle_assets(table: Mapping[str, Any], context: str) -> tuple[OracleAsset, ...]:
+    raw = table.get("oracle_assets", [])
+    if not isinstance(raw, list):
+        raise RegistryError(f"{context}.oracle_assets must be an array of tables.")
+    result: list[OracleAsset] = []
+    for index, value in enumerate(raw):
+        asset_context = f"{context}.oracle_assets[{index}]"
+        if not isinstance(value, dict):
+            raise RegistryError(f"{asset_context} must be a table.")
+        expected_fields = {"role", "path", "url", "sha256", "size"}
+        if set(value) != expected_fields:
+            raise RegistryError(f"{asset_context} must contain exactly {sorted(expected_fields)}.")
+        role = _require_str(value, "role", asset_context)
+        if role not in _ALLOWED_ORACLE_ASSET_ROLES:
+            raise RegistryError(f"Unsupported oracle asset role: {role!r}.")
+        path = _require_str(value, "path", asset_context)
+        try:
+            normalized_path = _portable_relative_path(path, "Oracle asset path")
+        except RegistryError as error:
+            raise RegistryError(f"Invalid oracle asset path: {path!r}.") from error
+        if normalized_path.suffix != ".pt":
+            raise RegistryError(f"Invalid oracle asset path: {path!r}.")
+        url = _require_str(value, "url", asset_context)
+        parsed_url = urlparse(url)
+        if (
+            parsed_url.scheme != "https"
+            or parsed_url.hostname != _FAIR_ESM_ASSET_HOST
+            or parsed_url.path != f"/fair-esm/{path}"
+            or parsed_url.params
+            or parsed_url.query
+            or parsed_url.fragment
+        ):
+            raise RegistryError(f"Invalid fair-esm oracle asset URL: {url!r}.")
+        sha256 = _require_str(value, "sha256", asset_context)
+        if len(sha256) != 64 or _HEX_RE.fullmatch(sha256) is None:
+            raise RegistryError(f"Invalid oracle asset SHA-256 for {path!r}.")
+        size = value.get("size")
+        if isinstance(size, bool) or not isinstance(size, int) or size <= 0:
+            raise RegistryError(f"{asset_context}.size must be a positive byte count.")
+        result.append(
+            OracleAsset(
+                role=role,
+                path=path,
+                url=url,
+                sha256=sha256,
+                size=size,
+            )
+        )
+    roles = [asset.role for asset in result]
+    paths = [asset.path for asset in result]
+    urls = [asset.url for asset in result]
+    if (
+        len(roles) != len(set(roles))
+        or len(paths) != len(set(paths))
+        or len(urls) != len(set(urls))
+    ):
+        raise RegistryError(f"{context}.oracle_assets contains duplicate identities.")
+    return tuple(result)
+
+
+def _parse_official_golden(
+    table: Mapping[str, Any],
+    model_id: str,
+    context: str,
+) -> OfficialGolden | None:
+    raw = table.get("official_golden")
+    if raw is None:
+        return None
+    if not isinstance(raw, dict) or set(raw) != {"metadata", "tensors"}:
+        raise RegistryError(
+            f"{context}.official_golden must contain exactly 'metadata' and 'tensors'."
+        )
+    parsed: dict[str, FileDigest] = {}
+    for role in ("metadata", "tensors"):
+        value = raw[role]
+        if not isinstance(value, str):
+            raise RegistryError(f"{context}.official_golden.{role} must be a file digest.")
+        digest = FileDigest.parse(value)
+        if digest.algorithm != "sha256":
+            raise RegistryError(
+                f"{context}.official_golden.{role} must use an immutable SHA-256 digest."
+            )
+        expected = f"tests/goldens/{model_id}.{'json' if role == 'metadata' else 'safetensors'}"
+        if digest.path != expected:
+            raise RegistryError(f"{context}.official_golden.{role} must use path {expected!r}.")
+        parsed[role] = digest
+    return OfficialGolden(metadata=parsed["metadata"], tensors=parsed["tensors"])
+
+
+def _parse_attention_kernels(raw: object) -> dict[str, AttentionKernelSpec]:
+    if not isinstance(raw, list) or not raw:
+        raise RegistryError("The manifest must contain [[attention_kernels]] entries.")
+    result: dict[str, AttentionKernelSpec] = {}
+    expected_variants = {
+        "flash_attention_2": "flash_attn2",
+        "flash_attention_3": "flash_attn3",
+    }
+    for index, value in enumerate(raw):
+        context = f"attention_kernels[{index}]"
+        if not isinstance(value, dict):
+            raise RegistryError(f"{context} must be a table.")
+        expected_fields = frozenset(
+            {
+                "implementation",
+                "repository",
+                "revision",
+                "version",
+                "expected_variant",
+                "dtypes",
+            }
+        )
+        _reject_unknown_fields(value, expected_fields, context)
+        implementation = _require_str(value, "implementation", context)
+        if implementation not in expected_variants:
+            raise RegistryError(f"Unsupported attention kernel {implementation!r}.")
+        if implementation in result:
+            raise RegistryError(f"Duplicate attention kernel {implementation!r}.")
+        repository = _require_str(value, "repository", context)
+        if _REPOSITORY_ID_RE.fullmatch(repository) is None:
+            raise RegistryError(f"Invalid attention-kernel repository {repository!r}.")
+        revision = _require_str(value, "revision", context)
+        _validate_revision(revision, f"{context}.revision")
+        kernel_version = value.get("version")
+        if (
+            isinstance(kernel_version, bool)
+            or not isinstance(kernel_version, int)
+            or kernel_version <= 0
+        ):
+            raise RegistryError(f"{context}.version must be a positive integer.")
+        expected_variant = _require_str(value, "expected_variant", context)
+        if expected_variant != expected_variants[implementation]:
+            raise RegistryError(
+                f"{context}.expected_variant must be {expected_variants[implementation]!r}."
+            )
+        dtypes = _require_str_list(value, "dtypes", context)
+        if not set(dtypes).issubset(_ALLOWED_DTYPES):
+            raise RegistryError(f"{context}.dtypes contains unsupported dtypes.")
+        result[implementation] = AttentionKernelSpec(
+            implementation=implementation,
+            repository=repository,
+            revision=revision,
+            version=kernel_version,
+            expected_variant=expected_variant,
+            dtypes=cast(tuple[DtypeName, ...], dtypes),
+        )
+    if set(result) != set(expected_variants):
+        raise RegistryError("The manifest must pin both FlashAttention kernel versions.")
+    return result
+
+
+def _parse_upstreams(raw: object) -> dict[str, UpstreamSource]:
+    if not isinstance(raw, list) or not raw:
+        raise RegistryError("The manifest must contain at least one [[upstreams]] entry.")
+    result: dict[str, UpstreamSource] = {}
+    paths: set[str] = set()
+    for index, value in enumerate(raw):
+        context = f"upstreams[{index}]"
+        if not isinstance(value, dict):
+            raise RegistryError(f"{context} must be a table.")
+        _reject_unknown_fields(value, _UPSTREAM_FIELDS, context)
+        source_id = _require_str(value, "id", context)
+        if _IDENTIFIER_RE.fullmatch(source_id) is None:
+            raise RegistryError(f"Invalid upstream ID: {source_id!r}")
+        if source_id in result:
+            raise RegistryError(f"Duplicate upstream ID: {source_id!r}")
+        revision = _require_str(value, "revision", context)
+        _validate_revision(revision, f"{context}.revision")
+        path = _require_str(value, "path", context)
+        try:
+            normalized_path = _portable_relative_path(path, f"{context}.path")
+        except RegistryError as error:
+            raise RegistryError(
+                f"{context}.path must be a normalized directory directly under "
+                "'vendor/upstream/'."
+            ) from error
+        if (
+            normalized_path.parts[:2] != ("vendor", "upstream")
+            or len(normalized_path.parts) != 3
+        ):
+            raise RegistryError(
+                f"{context}.path must be a normalized directory directly under "
+                "'vendor/upstream/'."
+            )
+        if path in paths:
+            raise RegistryError(f"Duplicate upstream path: {path!r}")
+        paths.add(path)
+        url = _require_str(value, "url", context)
+        if not url.startswith("https://github.com/") or not url.endswith(".git"):
+            raise RegistryError(f"{context}.url must be an HTTPS GitHub clone URL.")
+        license_files = _require_str_list(value, "license_files", context)
+        license_digests = _require_digest_list(value, "license_digests", context)
+        if tuple(item.path for item in license_digests) != license_files:
+            raise RegistryError(
+                f"{context}.license_digests must cover license_files in the same order."
+            )
+        distribution_files = _require_digest_list(value, "distribution_files", context)
+        distribution_map = {item.path: item for item in distribution_files}
+        for canonical in license_digests:
+            distributed = distribution_map.get(canonical.path)
+            if distributed is None or distributed.encoded != canonical.encoded:
+                raise RegistryError(
+                    f"{context}.distribution_files must include an exact copy of "
+                    f"{canonical.path!r}."
+                )
+        if source_id == "e1":
+            required_e1 = {
+                "LICENSE",
+                "ATTRIBUTION",
+                "NOTICE",
+                "Apache-2.0.txt",
+                "BSD-3-Clause.txt",
+                "MODIFICATIONS.md",
+            }
+            missing_e1 = sorted(required_e1.difference(distribution_map))
+            if missing_e1:
+                raise RegistryError(f"{context} is missing E1 legal files: {missing_e1}")
+        result[source_id] = UpstreamSource(
+            id=source_id,
+            path=path,
+            url=url,
+            revision=revision,
+            license_expression=_require_str(value, "license", context),
+            license_files=license_files,
+            license_digests=license_digests,
+            distribution_files=distribution_files,
+        )
+    return result
+
+
+def _parse_families(
+    raw: object,
+    upstreams: Mapping[str, UpstreamSource],
+) -> dict[str, ModelFamily]:
+    if not isinstance(raw, dict) or not raw:
+        raise RegistryError("The manifest must contain [families.<id>] tables.")
+    result: dict[str, ModelFamily] = {}
+    for family_id, value in raw.items():
+        context = f"families.{family_id}"
+        if _IDENTIFIER_RE.fullmatch(family_id) is None or not isinstance(value, dict):
+            raise RegistryError(f"Invalid family table: {family_id!r}")
+        checkpoint_license = _require_str(value, "checkpoint_license", context)
+        hub_license, hub_license_name, hub_license_link = _parse_hub_license(
+            value,
+            checkpoint_license=checkpoint_license,
+            context=context,
+        )
+        _reject_unknown_fields(value, _FAMILY_FIELDS, context)
+        source_ids = _require_str_list(value, "upstreams", context)
+        unknown_sources = sorted(set(source_ids).difference(upstreams))
+        if unknown_sources:
+            raise RegistryError(f"{context} references unknown upstreams: {unknown_sources}")
+        tokenizer_mode = _require_str(value, "tokenizer_mode", context)
+        if tokenizer_mode not in _ALLOWED_TOKENIZER_MODES:
+            raise RegistryError(f"Unsupported tokenizer mode in {context}: {tokenizer_mode!r}")
+        public_input = _require_str(value, "public_input", context)
+        attention = _require_str_list(value, "attention", context)
+        if not set(attention).issubset(_ALLOWED_ATTENTION):
+            raise RegistryError(f"Unsupported attention implementation in {context}.")
+        dtypes = _require_str_list(value, "dtypes", context)
+        if not set(dtypes).issubset(_ALLOWED_DTYPES):
+            raise RegistryError(f"Unsupported dtype in {context}.")
+        bf16_execution = cast(
+            Bf16Execution,
+            _require_enum(
+                value,
+                "bf16_execution",
+                context,
+                _ALLOWED_BF16_EXECUTIONS,
+            ),
+        )
+        precisions = _require_str_list(value, "precisions", context)
+        if not set(precisions).issubset(_ALLOWED_PRECISIONS):
+            raise RegistryError(f"Unsupported precision policy in {context}.")
+        experimental_precisions = _optional_str_list(
+            value,
+            "experimental_precisions",
+            context,
+        )
+        unknown_experimental_precisions = sorted(
+            set(experimental_precisions).difference(precisions)
+        )
+        if unknown_experimental_precisions:
+            raise RegistryError(
+                f"{context}.experimental_precisions must be a subset of precisions; "
+                f"unknown values: {unknown_experimental_precisions}."
+            )
+        extra = cast(RuntimeExtra, _require_enum(value, "extra", context, _ALLOWED_EXTRAS))
+        vram_tier = cast(
+            VramTier,
+            _require_enum(value, "vram_tier", context, _ALLOWED_VRAM_TIERS),
+        )
+        test_tiers_raw = _require_str_list(value, "test_tiers", context)
+        unknown_test_tiers = sorted(set(test_tiers_raw).difference(_ALLOWED_TEST_TIERS))
+        if unknown_test_tiers:
+            raise RegistryError(
+                f"{context}.test_tiers contains unsupported tiers: {unknown_test_tiers}."
+            )
+        test_tiers = cast(tuple[TestTier, ...], test_tiers_raw)
+        reference_container = _parse_reference_container(value, context)
+        reference_adapter = _parse_reference_adapter(value, context)
+        documentation = _parse_documentation_path(value, context)
+        runtime_paths = _require_str_list(value, "runtime_paths", context)
+        if len(runtime_paths) != len(set(runtime_paths)):
+            raise RegistryError(f"{context}.runtime_paths must not contain duplicates.")
+        for runtime_path in runtime_paths:
+            try:
+                _portable_relative_path(runtime_path, f"{context}.runtime_paths entry")
+            except RegistryError as error:
+                raise RegistryError(
+                    f"Unsafe runtime path in {context}: {runtime_path!r}"
+                ) from error
+            if runtime_path.startswith("vendor/"):
+                raise RegistryError(f"Unsafe runtime path in {context}: {runtime_path!r}")
+        requires_complete_weight_publication = value.get(
+            "requires_complete_weight_publication",
+            False,
+        )
+        if not isinstance(requires_complete_weight_publication, bool):
+            raise RegistryError(
+                f"{context}.requires_complete_weight_publication must be a boolean."
+            )
+        if "weights_publication_allowed" not in value:
+            raise RegistryError(
+                f"{context}.weights_publication_allowed must be declared explicitly."
+            )
+        weights_publication_allowed = value["weights_publication_allowed"]
+        if not isinstance(weights_publication_allowed, bool):
+            raise RegistryError(f"{context}.weights_publication_allowed must be a boolean.")
+        raw_auto_map = value.get("auto_map")
+        if not isinstance(raw_auto_map, dict) or not raw_auto_map:
+            raise RegistryError(f"{context}.auto_map must be a non-empty table.")
+        auto_map: list[tuple[str, str]] = []
+        for auto_class, class_path in raw_auto_map.items():
+            if auto_class not in _ALLOWED_AUTO_CLASSES or not isinstance(class_path, str):
+                raise RegistryError(f"Invalid AutoClass mapping in {context}: {auto_class!r}")
+            if not class_path.startswith("fastplms.") or class_path.count(".") < 2:
+                raise RegistryError(f"Invalid Python class path in {context}: {class_path!r}")
+            auto_map.append((auto_class, class_path))
+        tokenizer_class = value.get("tokenizer_class")
+        if tokenizer_class is not None:
+            if tokenizer_mode != "tokenizer":
+                raise RegistryError(
+                    f"{context}.tokenizer_class requires tokenizer_mode='tokenizer'."
+                )
+            if (
+                not isinstance(tokenizer_class, str)
+                or not tokenizer_class.startswith("fastplms.")
+                or tokenizer_class.count(".") < 2
+            ):
+                raise RegistryError(
+                    f"Invalid tokenizer class path in {context}: {tokenizer_class!r}"
+                )
+        backbone_model = value.get("backbone_model")
+        if backbone_model is not None and (
+            not isinstance(backbone_model, str)
+            or _IDENTIFIER_RE.fullmatch(backbone_model) is None
+        ):
+            raise RegistryError(
+                f"{context}.backbone_model must be a valid manifest model ID."
+            )
+        state_transform = _require_str(value, "state_transform", context)
+        conversion_provenance = _require_str(value, "conversion_provenance", context)
+        required_sections = ("Input:", "Transformation:", "Output:", "Validation:", "Limitation:")
+        missing_sections = [
+            section for section in required_sections if section not in conversion_provenance
+        ]
+        if missing_sections or state_transform not in conversion_provenance:
+            raise RegistryError(
+                f"{context}.conversion_provenance must identify {state_transform!r} and "
+                f"contain mechanism-first sections; missing {missing_sections}."
+            )
+        result[family_id] = ModelFamily(
+            id=family_id,
+            architecture=_require_str(value, "architecture", context),
+            upstreams=source_ids,
+            tokenizer_mode=tokenizer_mode,
+            public_input=public_input,
+            extra=extra,
+            reference_container=reference_container,
+            reference_adapter=reference_adapter,
+            attention=attention,
+            dtypes=cast(tuple[DtypeName, ...], dtypes),
+            bf16_execution=bf16_execution,
+            precisions=precisions,
+            experimental_precisions=experimental_precisions,
+            vram_tier=vram_tier,
+            checkpoint_license=checkpoint_license,
+            hub_license=hub_license,
+            state_transform=state_transform,
+            representative=_require_str(value, "representative", context),
+            documentation=documentation,
+            test_tiers=test_tiers,
+            runtime_paths=runtime_paths,
+            auto_map_items=tuple(auto_map),
+            requires_complete_weight_publication=requires_complete_weight_publication,
+            weights_publication_allowed=weights_publication_allowed,
+            tokenizer_class=tokenizer_class,
+            hub_license_name=hub_license_name,
+            hub_license_link=hub_license_link,
+            conversion_provenance=conversion_provenance,
+            backbone_model=backbone_model,
+        )
+    return result
+
+
+def _parse_runtime_assets(
+    raw: object,
+    families: Mapping[str, ModelFamily],
+) -> dict[str, RuntimeAsset]:
+    if not isinstance(raw, list) or not raw:
+        raise RegistryError("The manifest must contain at least one [[runtime_assets]] entry.")
+    result: dict[str, RuntimeAsset] = {}
+    identities: set[tuple[str, str, str]] = set()
+    for index, value in enumerate(raw):
+        context = f"runtime_assets[{index}]"
+        if not isinstance(value, dict):
+            raise RegistryError(f"{context} must be a table.")
+        _reject_unknown_fields(value, _RUNTIME_ASSET_FIELDS, context)
+        asset_id = _require_str(value, "id", context)
+        if _IDENTIFIER_RE.fullmatch(asset_id) is None:
+            raise RegistryError(f"Invalid runtime asset ID: {asset_id!r}")
+        if asset_id in result:
+            raise RegistryError(f"Duplicate runtime asset ID: {asset_id!r}")
+        repository = _require_str(value, "repository", context)
+        if _REPOSITORY_ID_RE.fullmatch(repository) is None:
+            raise RegistryError(f"{context}.repository must be a Hugging Face repository ID.")
+        revision = _require_str(value, "revision", context)
+        _validate_revision(revision, f"{context}.revision")
+        path = _require_str(value, "path", context)
+        try:
+            normalized_path = _portable_relative_path(path, "Runtime asset path")
+        except RegistryError as error:
+            raise RegistryError(f"Runtime asset path is not portable: {path!r}") from error
+        sha256 = _require_str(value, "sha256", context)
+        if len(sha256) != 64 or _HEX_RE.fullmatch(sha256) is None:
+            raise RegistryError(f"Invalid runtime asset SHA-256 for {path!r}.")
+        size = value.get("size")
+        if isinstance(size, bool) or not isinstance(size, int) or size <= 0:
+            raise RegistryError(f"{context}.size must be a positive byte count.")
+        consumer_family = _require_str(value, "consumer_family", context)
+        if consumer_family not in families:
+            raise RegistryError(
+                f"{context}.consumer_family references unknown family {consumer_family!r}."
+            )
+        trust_kind = cast(
+            RuntimeAssetTrustKind,
+            _require_enum(
+                value,
+                "trust_kind",
+                context,
+                _ALLOWED_RUNTIME_ASSET_TRUST_KINDS,
+            ),
+        )
+        license_expression = _require_str(value, "license", context)
+        offline_behavior = _require_str(value, "offline_behavior", context)
+        if offline_behavior not in _ALLOWED_RUNTIME_ASSET_OFFLINE_BEHAVIORS:
+            raise RegistryError(
+                f"{context}.offline_behavior is unsupported: {offline_behavior!r}."
+            )
+        if trust_kind == "hash_pinned_pickle" and normalized_path.suffix != ".pkl":
+            raise RegistryError(
+                f"{context}.path must end in '.pkl' for trust_kind='hash_pinned_pickle'."
+            )
+        identity = (repository, revision, path)
+        if identity in identities:
+            raise RegistryError(f"Duplicate runtime asset identity: {identity!r}")
+        identities.add(identity)
+        result[asset_id] = RuntimeAsset(
+            id=asset_id,
+            repository=repository,
+            revision=revision,
+            path=path,
+            sha256=sha256,
+            size=size,
+            consumer_family=consumer_family,
+            trust_kind=trust_kind,
+            license_expression=license_expression,
+            offline_behavior=offline_behavior,
+        )
+    return result
+
+
+def _parse_models(
+    raw: object,
+    families: Mapping[str, ModelFamily],
+) -> dict[str, ModelSpec]:
+    if not isinstance(raw, list) or not raw:
+        raise RegistryError("The manifest must contain at least one [[models]] entry.")
+    result: dict[str, ModelSpec] = {}
+    fast_repositories: set[str] = set()
+    for index, value in enumerate(raw):
+        context = f"models[{index}]"
+        if not isinstance(value, dict):
+            raise RegistryError(f"{context} must be a table.")
+        _reject_unknown_fields(value, _MODEL_FIELDS, context)
+        model_id = _require_str(value, "id", context)
+        if _IDENTIFIER_RE.fullmatch(model_id) is None:
+            raise RegistryError(f"Invalid model ID: {model_id!r}")
+        if model_id in result:
+            raise RegistryError(f"Duplicate model ID: {model_id!r}")
+        family_id = _require_str(value, "family", context)
+        if family_id not in families:
+            raise RegistryError(f"{context} references unknown family {family_id!r}.")
+        fast = _parse_checkpoint(value, "fast", context)
+        official = _parse_checkpoint(value, "official", context)
+        if fast.repo_id in fast_repositories:
+            raise RegistryError(f"Duplicate FastPLMs repository ID: {fast.repo_id!r}")
+        fast_repositories.add(fast.repo_id)
+        family = families[family_id]
+        oracle_assets = _parse_oracle_assets(value, context)
+        official_golden = _parse_official_golden(value, model_id, context)
+        size_category = _require_str(value, "size_category", context)
+        if size_category not in _ALLOWED_SIZE_CATEGORIES:
+            raise RegistryError(f"Unsupported size category in {context}: {size_category!r}")
+        generation_contract = cast(
+            GenerationContract,
+            _require_enum(
+                value,
+                "generation_contract",
+                context,
+                _ALLOWED_GENERATION_CONTRACTS,
+            ),
+        )
+        if family.tokenizer_mode == "structure" and size_category != "structure":
+            raise RegistryError(
+                f"Structure checkpoint {model_id!r} must use size_category='structure'."
+            )
+        artifact_source = value.get("artifact_source", "fast")
+        if artifact_source not in {"fast", "official"}:
+            raise RegistryError(f"{context}.artifact_source must be 'fast' or 'official'.")
+        canonical_state_sha256 = value.get("canonical_state_sha256")
+        if artifact_source == "official":
+            if (
+                not isinstance(canonical_state_sha256, str)
+                or len(canonical_state_sha256) != 64
+                or _HEX_RE.fullmatch(canonical_state_sha256) is None
+            ):
+                raise RegistryError(
+                    f"{context}.canonical_state_sha256 must be a SHA-256 commitment "
+                    "for an official-source artifact."
+                )
+        elif canonical_state_sha256 is not None:
+            raise RegistryError(
+                f"{context}.canonical_state_sha256 is restricted to official-source artifacts."
+            )
+        if family.tokenizer_mode == "tokenizer" and not any(
+            "tokenizer" in item.path or "vocab" in item.path for item in fast.files
+        ):
+            raise RegistryError(f"{context} does not pin a tokenizer asset.")
+        tokenizer_source_id = value.get("tokenizer_source")
+        if tokenizer_source_id is not None and (
+            family.tokenizer_mode != "tokenizer"
+            or not isinstance(tokenizer_source_id, str)
+            or _IDENTIFIER_RE.fullmatch(tokenizer_source_id) is None
+        ):
+            raise RegistryError(f"{context}.tokenizer_source is invalid.")
+        notes = value.get("notes", "")
+        if not isinstance(notes, str):
+            raise RegistryError(f"{context}.notes must be a string.")
+        msa_conditioning = value.get("msa_conditioning")
+        if family_id == "esmfold2":
+            if not isinstance(msa_conditioning, bool):
+                raise RegistryError(
+                    f"{context}.msa_conditioning must be an explicit boolean for "
+                    "ESMFold2 checkpoints."
+                )
+        elif "msa_conditioning" in value:
+            raise RegistryError(
+                f"{context}.msa_conditioning is only valid for ESMFold2 checkpoints."
+            )
+        raw_auto_map = value.get("auto_map")
+        auto_map: list[tuple[str, str]] = []
+        if raw_auto_map is not None:
+            if not isinstance(raw_auto_map, dict) or not raw_auto_map:
+                raise RegistryError(f"{context}.auto_map must be a non-empty table.")
+            for auto_class, class_path in raw_auto_map.items():
+                if auto_class not in _ALLOWED_AUTO_CLASSES or not isinstance(class_path, str):
+                    raise RegistryError(f"Invalid AutoClass mapping in {context}: {auto_class!r}")
+                if not class_path.startswith("fastplms.") or class_path.count(".") < 2:
+                    raise RegistryError(f"Invalid Python class path in {context}: {class_path!r}")
+                auto_map.append((auto_class, class_path))
+        result[model_id] = ModelSpec(
+            id=model_id,
+            family=family,
+            fast=fast,
+            official=official,
+            size_category=size_category,
+            generation_contract=generation_contract,
+            oracle_assets=oracle_assets,
+            official_golden=official_golden,
+            artifact_source=artifact_source,
+            canonical_state_sha256=canonical_state_sha256,
+            tokenizer_source_id=tokenizer_source_id,
+            auto_map_items=tuple(auto_map),
+            notes=notes,
+            msa_conditioning=msa_conditioning,
+        )
+    return result
+
+
+def _validate_registry(
+    upstreams: Mapping[str, UpstreamSource],
+    attention_kernels: Mapping[str, AttentionKernelSpec],
+    families: Mapping[str, ModelFamily],
+    models: Mapping[str, ModelSpec],
+) -> None:
+    for spec in models.values():
+        if spec.tokenizer_source_id is None:
+            continue
+        source = models.get(spec.tokenizer_source_id)
+        if source is None:
+            raise RegistryError(
+                f"Model {spec.id!r} references unknown tokenizer source "
+                f"{spec.tokenizer_source_id!r}."
+            )
+        if not any(
+            PurePosixPath(item.path).name
+            in {
+                "added_tokens.json",
+                "merges.txt",
+                "sentencepiece.bpe.model",
+                "special_tokens_map.json",
+                "spiece.model",
+                "tokenizer.json",
+                "tokenizer_config.json",
+                "vocab.json",
+                "vocab.txt",
+            }
+            for item in source.official.files
+        ):
+            raise RegistryError(
+                f"Tokenizer source {source.id!r} has no official tokenizer assets."
+            )
+    expected_esmfold2 = {
+        "esmfold2": ("Synthyra/ESMFold2", "biohub/ESMFold2"),
+        "esmfold2_fast": ("Synthyra/ESMFold2-Fast", "biohub/ESMFold2-Fast"),
+        "esmfold2_experimental_cutoff2025": (
+            "Synthyra/ESMFold2-Experimental-Cutoff2025",
+            "biohub/ESMFold2-Experimental-Cutoff2025",
+        ),
+        "esmfold2_experimental_fast_cutoff2025": (
+            "Synthyra/ESMFold2-Experimental-Fast-Cutoff2025",
+            "biohub/ESMFold2-Experimental-Fast-Cutoff2025",
+        ),
+    }
+    actual_esmfold2 = {
+        model.id: (model.fast.repo_id, model.official.repo_id)
+        for model in models.values()
+        if model.family.id == "esmfold2"
+    }
+    if actual_esmfold2 != expected_esmfold2:
+        raise RegistryError(
+            "ESMFold2 support must contain exactly the four approved model IDs and "
+            "official/Synthyra repositories."
+        )
+
+    golden_paths: list[str] = []
+    for model in models.values():
+        if model.official_golden is not None:
+            golden_paths.extend(
+                (
+                    model.official_golden.metadata.path,
+                    model.official_golden.tensors.path,
+                )
+            )
+    if len(golden_paths) != len(set(golden_paths)):
+        raise RegistryError("Official golden paths must be unique across model declarations.")
+    unused_upstreams = sorted(
+        set(upstreams).difference(
+            source for family in families.values() for source in family.upstreams
+        )
+    )
+    if unused_upstreams:
+        raise RegistryError(
+            f"Upstream sources are not connected to a model family: {unused_upstreams}"
+        )
+    advertised_flash = {
+        implementation
+        for family in families.values()
+        for implementation in family.attention
+        if implementation.startswith("flash_attention_")
+    }
+    missing_kernels = sorted(advertised_flash.difference(attention_kernels))
+    if missing_kernels:
+        raise RegistryError(
+            f"Advertised FlashAttention backends lack kernel specs: {missing_kernels}."
+        )
+    for family in families.values():
+        for implementation in family.attention:
+            kernel = attention_kernels.get(implementation)
+            if kernel is not None and not set(family.dtypes).intersection(kernel.dtypes):
+                raise RegistryError(
+                    f"Family {family.id!r} and attention kernel {implementation!r} "
+                    "have no supported dtype in common."
+                )
+        family_models = [model for model in models.values() if model.family.id == family.id]
+        if not family_models:
+            raise RegistryError(f"Family {family.id!r} has no checkpoints.")
+        representative = models.get(family.representative)
+        if representative is None or representative.family.id != family.id:
+            raise RegistryError(
+                f"Family {family.id!r} has invalid representative {family.representative!r}."
+            )
+        if family.backbone_model is not None and family.backbone_model not in models:
+            raise RegistryError(
+                f"Family {family.id!r} references unknown backbone model "
+                f"{family.backbone_model!r}."
+            )
+
+
+def _load_manifest_bytes(raw_bytes: bytes) -> ModelRegistry:
+    try:
+        data = tomllib.loads(raw_bytes.decode("utf-8"))
+    except (UnicodeDecodeError, tomllib.TOMLDecodeError) as error:
+        raise RegistryError(f"Unable to parse model manifest: {error}") from error
+    _reject_unknown_fields(data, _ROOT_FIELDS, "manifest")
+    if data.get("schema_version") != 1:
+        raise RegistryError("Unsupported model manifest schema_version; expected 1.")
+    legal_files = _require_digest_list(data, "legal_files", "manifest")
+    required_legal_paths = {"LICENSE", "THIRD_PARTY_NOTICES.md"}
+    if {item.path for item in legal_files} != required_legal_paths:
+        raise RegistryError("manifest.legal_files must contain LICENSE and THIRD_PARTY_NOTICES.md.")
+    attention_kernels = _parse_attention_kernels(data.get("attention_kernels"))
+    upstreams = _parse_upstreams(data.get("upstreams"))
+    families = _parse_families(data.get("families"), upstreams)
+    runtime_assets = _parse_runtime_assets(data.get("runtime_assets"), families)
+    models = _parse_models(data.get("models"), families)
+    _validate_registry(upstreams, attention_kernels, families, models)
+    return ModelRegistry(
+        schema_version=1,
+        upstreams=upstreams,
+        attention_kernels=attention_kernels,
+        families=families,
+        models=models,
+        runtime_assets=runtime_assets,
+        legal_files=legal_files,
+    )
+
+
+def load_model_registry(path: str | Path | None = None) -> ModelRegistry:
+    """Load and validate a model manifest without importing model code."""
+
+    if path is None:
+        manifest = resources.files("fastplms").joinpath("models.toml")
+        return _load_manifest_bytes(manifest.read_bytes())
+    return _load_manifest_bytes(Path(path).read_bytes())
+
+
+@lru_cache(maxsize=1)
+def get_model_registry() -> ModelRegistry:
+    """Return the validated package registry, cached after its first read."""
+
+    return load_model_registry()
+
+
+def get_model_spec(model_id: str) -> ModelSpec:
+    """Return one model specification by its stable manifest ID."""
+
+    try:
+        return get_model_registry()[model_id]
+    except KeyError as error:
+        supported = ", ".join(get_model_registry())
+        raise KeyError(
+            f"Unknown FastPLMs model ID {model_id!r}. Supported IDs: {supported}"
+        ) from error
+
+
+__all__ = [
+    "HUB_LICENSE_IDENTIFIERS",
+    "CheckpointSource",
+    "FileDigest",
+    "GenerationContract",
+    "ModelFamily",
+    "ModelRegistry",
+    "ModelSpec",
+    "OracleAsset",
+    "RegistryError",
+    "RuntimeAsset",
+    "RuntimeAssetTrustKind",
+    "RuntimeExtra",
+    "TestTier",
+    "UpstreamSource",
+    "VramTier",
+    "get_model_registry",
+    "get_model_spec",
+    "load_model_registry",
+]

fastplms/runtime.py

@@ -0,0 +1,68 @@
+"""Explicit, reversible Torch runtime configuration.
+
+Importing FastPLMs does not change global Torch settings. Callers that want a
+runtime profile opt in with :func:`runtime_profile` and receive their previous
+settings back when the context exits.
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Literal
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
+
+
+MatmulPrecision = Literal["highest", "high", "medium"]
+
+
+@dataclass(frozen=True, slots=True)
+class RuntimeProfile:
+    """Requested Torch settings for a bounded inference or training block."""
+
+    float32_matmul_precision: MatmulPrecision = "highest"
+    allow_tf32: bool | None = None
+
+
+@contextmanager
+def runtime_profile(profile: RuntimeProfile | None = None) -> Iterator[None]:
+    """Apply a Torch runtime profile and restore the previous global settings.
+
+    The default profile requests the highest float32 matrix-multiplication
+    precision and leaves TF32 policy unchanged. Torch is imported only when the
+    context is entered.
+    """
+
+    import torch
+
+    selected = profile or RuntimeProfile()
+    previous_matmul_precision = torch.get_float32_matmul_precision()
+    matmul_backend = getattr(getattr(torch.backends, "cuda", None), "matmul", None)
+    cudnn_backend = getattr(torch.backends, "cudnn", None)
+    previous_matmul_tf32 = (
+        getattr(matmul_backend, "allow_tf32", None) if matmul_backend is not None else None
+    )
+    previous_cudnn_tf32 = (
+        getattr(cudnn_backend, "allow_tf32", None) if cudnn_backend is not None else None
+    )
+
+    torch.set_float32_matmul_precision(selected.float32_matmul_precision)
+    if selected.allow_tf32 is not None:
+        if matmul_backend is not None and hasattr(matmul_backend, "allow_tf32"):
+            matmul_backend.allow_tf32 = selected.allow_tf32
+        if cudnn_backend is not None and hasattr(cudnn_backend, "allow_tf32"):
+            cudnn_backend.allow_tf32 = selected.allow_tf32
+    try:
+        yield
+    finally:
+        torch.set_float32_matmul_precision(previous_matmul_precision)
+        if selected.allow_tf32 is not None:
+            if matmul_backend is not None and previous_matmul_tf32 is not None:
+                matmul_backend.allow_tf32 = previous_matmul_tf32
+            if cudnn_backend is not None and previous_cudnn_tf32 is not None:
+                cudnn_backend.allow_tf32 = previous_cudnn_tf32
+
+
+__all__ = ["MatmulPrecision", "RuntimeProfile", "runtime_profile"]

modeling_fastplms.py

@@ -0,0 +1,239 @@
+"""Generated bridge to the unchanged FastPLMs package sources."""
+
+import base64
+import hashlib
+import importlib
+import importlib.util
+import sys
+import tempfile
+from importlib.metadata import PackageNotFoundError, distribution
+from io import BytesIO
+from pathlib import Path
+from zipfile import ZIP_DEFLATED, ZipFile
+
+from .fastplms_bundle import RUNTIME_DATA, RUNTIME_HASH
+
+if RUNTIME_HASH != "3ca371517961576b5af3c51ca4f07e486d91a11c8da3b0d7916d0bbc8f76be72":
+    raise RuntimeError("FastPLMs runtime identity differs from the bridge.")
+
+_RUNTIME_TEMPORARIES = []
+
+def _archive_runtime_hashes(payload):
+    result = {}
+    with ZipFile(BytesIO(payload)) as archive:
+        for member in archive.infolist():
+            name = member.filename
+            parts = Path(name).parts
+            if (
+                member.is_dir()
+                or "\\" in name
+                or not parts
+                or parts[0] != "fastplms"
+                or len(parts) < 2
+                or any(part in {"", ".", ".."} for part in parts)
+                or Path(name).suffix in {".pyc", ".pyo"}
+                or member.flag_bits & 0x1
+                or member.compress_type != ZIP_DEFLATED
+                or member.external_attr >> 16 != 0o100644
+            ):
+                raise RuntimeError("Embedded FastPLMs archive has an unsafe path.")
+            relative = Path(*parts[1:]).as_posix()
+            if relative in result:
+                raise RuntimeError("Embedded FastPLMs archive repeats a path.")
+            result[relative] = hashlib.sha256(archive.read(member)).hexdigest()
+    return result
+
+def _ensure_runtime():
+    payload = base64.b85decode("".join(RUNTIME_DATA))
+    if hashlib.sha256(payload).hexdigest() != RUNTIME_HASH:
+        raise RuntimeError("Embedded FastPLMs runtime hash mismatch.")
+    expected = _archive_runtime_hashes(payload)
+    temporary = tempfile.TemporaryDirectory(prefix="fastplms-artifact-runtime-")
+    try:
+        runtime_root = Path(temporary.name)
+        with ZipFile(BytesIO(payload)) as archive:
+            for member in archive.infolist():
+                target = runtime_root.joinpath(*Path(member.filename).parts)
+                target.parent.mkdir(parents=True, exist_ok=True)
+                with target.open("xb") as handle:
+                    handle.write(archive.read(member))
+        package_root = runtime_root / "fastplms"
+        if _runtime_file_hashes(package_root) != expected:
+            raise RuntimeError(
+                "Private FastPLMs runtime differs from the embedded archive."
+            )
+    except BaseException:
+        temporary.cleanup()
+        raise
+    _RUNTIME_TEMPORARIES.append(temporary)
+    return package_root
+
+def _runtime_file_hashes(package_root):
+    result = {}
+    for path in sorted(package_root.rglob("*")):
+        relative = path.relative_to(package_root)
+        if path.is_symlink():
+            raise RuntimeError("Private FastPLMs runtime contains a symlink.")
+        if path.is_dir():
+            continue
+        if path.suffix in {".pyc", ".pyo"}:
+            raise RuntimeError("Private FastPLMs runtime contains bytecode.")
+        if not path.is_file():
+            raise RuntimeError("Private FastPLMs runtime contains a non-file entry.")
+        result[relative.as_posix()] = hashlib.sha256(path.read_bytes()).hexdigest()
+    return result
+
+def _installed_runtime_digest(installed_root, relative):
+    candidate = installed_root / relative
+    if candidate.is_file():
+        return hashlib.sha256(candidate.read_bytes()).hexdigest()
+    if relative != "kernels.lock":
+        return None
+    try:
+        installed_distribution = distribution("fastplms")
+    except PackageNotFoundError:
+        return None
+    for entry in installed_distribution.files or ():
+        normalized = str(entry).replace("\\", "/")
+        if normalized.endswith(".dist-info/kernels.lock"):
+            lock_path = Path(installed_distribution.locate_file(entry))
+            if lock_path.is_file():
+                return hashlib.sha256(lock_path.read_bytes()).hexdigest()
+    return None
+
+def _extend_loaded_package_paths(package_root):
+    for name, module in list(sys.modules.items()):
+        if name != "fastplms" and not name.startswith("fastplms."):
+            continue
+        paths = getattr(module, "__path__", None)
+        if paths is None:
+            continue
+        relative = name.split(".")[1:]
+        candidate = package_root.joinpath(*relative)
+        candidate_text = str(candidate)
+        if candidate.is_dir() and candidate_text not in paths:
+            paths.append(candidate_text)
+
+def _merge_runtime(installed, package_root):
+    incoming = _runtime_file_hashes(package_root)
+    known = dict(getattr(installed, "__fastplms_artifact_runtime_files__", {}))
+    installed_root_text = getattr(
+        installed, "__fastplms_artifact_installed_root__", None
+    )
+    if not known:
+        installed_file = getattr(installed, "__file__", None)
+        if installed_file is None:
+            raise RuntimeError(
+                "The loaded fastplms package has no source path and cannot be verified "
+                "against the embedded artifact runtime."
+            )
+        installed_root = Path(installed_file).resolve().parent
+        for relative, digest in incoming.items():
+            if _installed_runtime_digest(installed_root, relative) != digest:
+                raise RuntimeError(
+                    "The installed FastPLMs runtime differs from this artifact at "
+                    f"{relative!r}. Install the artifact's matching FastPLMs release "
+                    "or use a separate Python process."
+                )
+        installed_root_text = str(installed_root)
+        installed.__fastplms_artifact_installed_root__ = installed_root_text
+    conflicts = sorted(
+        relative
+        for relative, digest in incoming.items()
+        if relative in known and known[relative] != digest
+    )
+    if conflicts:
+        raise RuntimeError(
+            "FastPLMs artifacts contain incompatible runtime sources at "
+            + ", ".join(repr(path) for path in conflicts[:5])
+            + ". Load incompatible releases in separate Python processes."
+        )
+    if installed_root_text is not None:
+        installed_root = Path(installed_root_text)
+        for relative, digest in incoming.items():
+            if relative in known:
+                continue
+            if _installed_runtime_digest(installed_root, relative) != digest:
+                raise RuntimeError(
+                    "The installed FastPLMs runtime differs from this artifact at "
+                    f"{relative!r}. Install the artifact's matching FastPLMs release "
+                    "or use a separate Python process."
+                )
+    known.update(incoming)
+    installed.__fastplms_artifact_runtime_files__ = known
+    roots = list(getattr(installed, "__fastplms_artifact_runtime_roots__", ()))
+    if str(package_root) not in roots:
+        roots.append(str(package_root))
+    installed.__fastplms_artifact_runtime_roots__ = tuple(roots)
+    temporaries = list(
+        getattr(installed, "__fastplms_artifact_runtime_temporaries__", ())
+    )
+    for temporary in _RUNTIME_TEMPORARIES:
+        if temporary not in temporaries:
+            temporaries.append(temporary)
+    installed.__fastplms_artifact_runtime_temporaries__ = tuple(temporaries)
+    hashes = set(getattr(installed, "__fastplms_artifact_runtime_hashes__", ()))
+    hashes.add(RUNTIME_HASH)
+    installed.__fastplms_artifact_runtime_hashes__ = frozenset(hashes)
+    _extend_loaded_package_paths(package_root)
+    return installed
+
+def _import_without_bytecode(module_name):
+    previous = sys.dont_write_bytecode
+    sys.dont_write_bytecode = True
+    try:
+        return importlib.import_module(module_name)
+    finally:
+        sys.dont_write_bytecode = previous
+
+def _install_runtime():
+    installed = sys.modules.get("fastplms")
+    hashes = getattr(installed, "__fastplms_artifact_runtime_hashes__", ())
+    if RUNTIME_HASH in hashes:
+        return installed
+    package_root = _ensure_runtime()
+    if installed is not None:
+        return _merge_runtime(installed, package_root)
+    spec = importlib.util.spec_from_file_location(
+        "fastplms",
+        package_root / "__init__.py",
+        submodule_search_locations=[str(package_root)],
+    )
+    if spec is None or spec.loader is None:
+        raise ImportError("Unable to load the embedded FastPLMs runtime.")
+    package = importlib.util.module_from_spec(spec)
+    package.__fastplms_artifact_runtime_hash__ = RUNTIME_HASH
+    package.__fastplms_artifact_runtime_hashes__ = frozenset({RUNTIME_HASH})
+    package.__fastplms_artifact_runtime_files__ = _runtime_file_hashes(package_root)
+    package.__fastplms_artifact_runtime_roots__ = (str(package_root),)
+    package.__fastplms_artifact_runtime_temporaries__ = tuple(
+        _RUNTIME_TEMPORARIES
+    )
+    sys.modules["fastplms"] = package
+    previous = sys.dont_write_bytecode
+    sys.dont_write_bytecode = True
+    try:
+        try:
+            spec.loader.exec_module(package)
+        except BaseException:
+            sys.modules.pop("fastplms", None)
+            raise
+    finally:
+        sys.dont_write_bytecode = previous
+    return package
+
+_install_runtime()
+_module_225 = _import_without_bytecode("fastplms.models.dplm2.modeling_dplm2")
+DPLM2Config = _module_225.DPLM2Config
+DPLM2Config.__module__ = __name__
+DPLM2ForMaskedLM = _module_225.DPLM2ForMaskedLM
+DPLM2ForMaskedLM.__module__ = __name__
+DPLM2ForSequenceClassification = _module_225.DPLM2ForSequenceClassification
+DPLM2ForSequenceClassification.__module__ = __name__
+DPLM2ForTokenClassification = _module_225.DPLM2ForTokenClassification
+DPLM2ForTokenClassification.__module__ = __name__
+DPLM2Model = _module_225.DPLM2Model
+DPLM2Model.__module__ = __name__
+_module_236 = _import_without_bytecode("fastplms.models.dplm2.tokenization_dplm2")
+DPLM2Tokenizer = _module_236.DPLM2Tokenizer
+DPLM2Tokenizer.__module__ = __name__

runtime-attestation.json

@@ -0,0 +1,49 @@
+{
+  "files": {
+    "LICENSES/FastPLMs-Apache-2.0.txt": "sha256:2d2b50c7b1414bff1189a1db1f0cfb92e3e064b50f4c2b1019827b683e1b629a",
+    "LICENSES/dplm/LICENSE": "sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30",
+    "LICENSES/dplm/PROVENANCE.md": "sha256:a659f74be9073cf1ad2d2f7071531ca56959b421f111152cf4c41184ace5970e",
+    "README.md": "sha256:5f2a653497b6807eebe5d83681babe4629a5a4636545e9652e77aaf8b3595170",
+    "THIRD_PARTY_NOTICES.md": "sha256:25704b3c76404696cae52e7fca13088d329f70f412687340351259e86cd62baa",
+    "config.json": "sha256:50a24887d94862f1505c3f629794b69667d84c0c4aef7edeb43a8095179a6dfc",
+    "fastplms/__init__.py": "sha256:4fb3196022ca8ec699d59d09bdbc5f0184195552b773698ab9b061fe3cd7df12",
+    "fastplms/attention/__init__.py": "sha256:f60b9fecfb4bcb37a4e7c26dc2f752b9035f9cbad627b4a84213f3a92ec88f7d",
+    "fastplms/attention/_core.py": "sha256:8f7ec5b65bd8b6c6fa4951d50d1c0e499abf03ae00914794b51fc410201e3e33",
+    "fastplms/attention/_kernel_lock.py": "sha256:85d8521a2af5f94fad3948af3814db0c866c414ee4d43df797b9bd6f980e947b",
+    "fastplms/attention/interfaces.py": "sha256:1c6f06a8e411e0f9bf6d230522205c93ae46ea58864006fbb892aa05e5ca5749",
+    "fastplms/embeddings/__init__.py": "sha256:47ff8cdf682d44037dd9edab133e2e60675d60786bd5cf0bffd1998f31985555",
+    "fastplms/embeddings/pooling.py": "sha256:a140266ed6b1cc344c8507edc5c6c4f2dce464c3db70ba4b16c7ac2ba2fad96e",
+    "fastplms/embeddings/runner.py": "sha256:23ee4727a918d6d331f7a0f89b823d149f1a791f0c5586e3496d7b6eb2ce97e0",
+    "fastplms/embeddings/storage.py": "sha256:3fbe2bab75092e5a4cadf4d27e4752181d597469a65a55db085ceef808ed418e",
+    "fastplms/embeddings/types.py": "sha256:119718a20989d1ae5a60fabc0f5e98bdc172c5163b04db3d4554ac3956b30e52",
+    "fastplms/models.toml": "sha256:05a8399f084a5babb5f0916cee7e564c4030767f3ff0230c4f46e539209847d1",
+    "fastplms/models/__init__.py": "sha256:5e48c2cb3877aa6f42f3b5411d53b16bba2e32827bbde634f47f174c5cb36f86",
+    "fastplms/models/_diffusion_generation.py": "sha256:00a644adf604a1d7760b11f75be5a823bcbfbf894bc7a02b0640b96245b5ae5b",
+    "fastplms/models/_esm_rotary.py": "sha256:4dcee8af74adf2cfcabeeaadf4915f1da5fb6c0672d583bf3ee5478ac1290805",
+    "fastplms/models/dplm2/__init__.py": "sha256:1d823d50970e7dc9ad2e7e06b6532c70612ce533dca2bdc49c687d4707c41232",
+    "fastplms/models/dplm2/modeling_dplm2.py": "sha256:c3af6614cd85ac290c9c310b553c213755ffb05f2c6067cf3bddfaaf19e20d6b",
+    "fastplms/models/dplm2/tokenization_dplm2.py": "sha256:ce4ce471f751255dbd44e2ebbfaffe5875a99cea28dfc9d30582aa977cfa96d5",
+    "fastplms/models/ttt.py": "sha256:a0df4e98b02120d423e3c7ca9b866a8d0e3748b9076042a0102a838a11aed046",
+    "fastplms/registry.py": "sha256:afca271911b651a882345b74a58366494a1d784e4a48c8683a87f5508f4ba16e",
+    "fastplms/runtime.py": "sha256:110018646d6f248cedab140a030c3065e1b062b61f6aff659c231e538614bc01",
+    "fastplms_bundle.py": "sha256:e5a7998e7dcb788abcffc57eb3697fd1a445cbb9c00e8d7d33e0577d23172f3e",
+    "modeling_fastplms.py": "sha256:99e57753e146b9568fc5df47bf9c5d6a0bf5a3028a11c1acf3d4ebb1e4c1f1a2",
+    "special_tokens_map.json": "sha256:10d6448302e43bd18bc2b648a750ade73a1464d77340f44fc6360151f89bd593",
+    "tokenizer_config.json": "sha256:cdfc691425b5a32516704396ebc4bc1f604f2e18de354fbd608558f5b6097dc3",
+    "vocab.txt": "sha256:d6bd0a2b1a0bd0f8e10568cfd75443e62a4bcf794907e93f41a5f945baf3b900"
+  },
+  "model_id": "dplm2_150m",
+  "redistributable": true,
+  "release_tool_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "release_tool_sha256": "1459b5d7d13d9b07bd97b3eee764f2ce73623e15e32d07ddf6825c2a9509afb9",
+  "runtime_bundle_sha256": "3ca371517961576b5af3c51ca4f07e486d91a11c8da3b0d7916d0bbc8f76be72",
+  "runtime_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "schema_version": 2,
+  "scope": "runtime-only",
+  "source_tree_sha256": "5de51d1c5c03b023e3db0af2d6a8849140af8011164caf84722faa356428efe1",
+  "weights": {
+    "repo_id": "Synthyra/DPLM2-150M",
+    "revision": "182745b8dc5661f898481a4fa60a7af9d53385c4"
+  },
+  "weights_license_status": "resolved"
+}

special_tokens_map.json

@@ -1,13 +1,11 @@
 {
-  "cls_token": "<cls>",
-  "eos_token": "<eos>",
-  "mask_token": "<mask>",
-  "pad_token": {
-    "content": "<pad>",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  },
-  "unk_token": "<unk>"
+  "aa_cls_token": "<cls_aa>",
+  "aa_eos_token": "<eos_aa>",
+  "aa_mask_token": "<mask_aa>",
+  "aa_unk_token": "<unk_aa>",
+  "struct_cls_token": "<cls_struct>",
+  "struct_eos_token": "<eos_struct>",
+  "struct_mask_token": "<mask_struct>",
+  "struct_unk_token": "<unk_struct>",
+  "pad_token": "<pad>"
 }

tokenizer_config.json

@@ -69,31 +69,7 @@
       "special": true
     },
     "8229": {
-      "content": "<eos>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "8230": {
-      "content": "<unk>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "8231": {
-      "content": "<cls>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "8232": {
-      "content": "<mask>",
+      "content": "<mask_struct>",
       "lstrip": false,
       "normalized": false,
       "rstrip": false,
@@ -101,17 +77,18 @@
       "special": true
     }
   },
+  "auto_map": {
+    "AutoTokenizer": [
+      "modeling_fastplms.DPLM2Tokenizer",
+      null
+    ]
+  },
   "clean_up_tokenization_spaces": true,
-  "cls_token": "<cls>",
-  "eos_token": "<eos>",
-  "extra_special_tokens": {},
-  "mask_token": "<mask>",
   "model_max_length": 1000000000000000019884624838656,
   "pad_token": "<pad>",
   "struct_cls_token": "<cls_struct>",
   "struct_eos_token": "<eos_struct>",
   "struct_mask_token": "<mask_struct>",
   "struct_unk_token": "<unk_struct>",
-  "tokenizer_class": "EsmTokenizer",
-  "unk_token": "<unk>"
+  "tokenizer_class": "DPLM2Tokenizer"
 }

vocab.txt

@@ -8226,4 +8226,4 @@ O
 8189
 8190
 8191
-<mask_struct>
+<mask_struct>