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/fair-esm/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) Meta Platforms, Inc. and affiliates.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

LICENSES/fair-esm/PROVENANCE.md

@@ -0,0 +1,8 @@
+# Meta ESM provenance
+
+FastPLMs uses `facebookresearch/esm` revision
+`2b369911bb5b4b0dda914521b9475cad1656b2ac` as the official parity oracle for
+ESM2 and ESMFold. The repository is pinned at
+`vendor/upstream/fair-esm/` and is not a production dependency or runtime image
+component. The accompanying `LICENSE` is the verbatim MIT text from that
+revision.

LICENSES/openfold/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/openfold/MODIFICATIONS.md

@@ -0,0 +1,25 @@
+# OpenFold reference-build modification
+
+FastPLMs pins `aqlaboratory/openfold` revision
+`4b41059694619831a7db195b7e0988fc4ff3a307` as an ESMFold parity oracle.
+The pinned checkout under `vendor/upstream/openfold/` is not modified.
+
+The `reference-esmfold` image applies
+`docker/constraints/openfold-sm90.patch` to the copied `setup.py` before
+installing OpenFold. The patch replaces OpenFold's build-time fallback CUDA
+architecture list with `sm90` and changes the extension build flag from C++14
+to C++17. Docker BuildKit cannot observe the workstation GPU, so the original
+setup otherwise requests legacy architectures that CUDA 12.1 no longer
+accepts. PyTorch 2.2 requires C++17, and the H100 requires an `sm90` extension.
+
+This is a build-only packaging change. It does not alter OpenFold model classes,
+extension source, checkpoint data, or the public API used by the native oracle.
+The resulting reference image is specific to the declared H100 validation
+environment.
+
+OpenFold imports PyTorch Lightning and NVIDIA DLLogger from its package
+initializers even though ESMFold inference does not use their training or
+logging features. The native image therefore pins PyTorch Lightning `1.9.5`,
+TorchMetrics `0.11.4`, Lightning Utilities `0.15.2`, and NVIDIA DLLogger commit
+`0478734ff7be75adde8d160e04872664d1c62e5f`. These Apache-2.0 dependencies are
+native-reference-only and are not FastPLMs runtime dependencies.

LICENSES/openfold/PROVENANCE.md

@@ -0,0 +1,7 @@
+# OpenFold provenance
+
+FastPLMs uses `aqlaboratory/openfold` revision
+`4b41059694619831a7db195b7e0988fc4ff3a307` as an official ESMFold parity
+oracle. The repository is pinned at `vendor/upstream/openfold/` and is not a
+production dependency or runtime image component. The accompanying `LICENSE`
+is the verbatim Apache-2.0 text from that revision.

README.md

@@ -1,234 +1,136 @@
----
-library_name: transformers
-tags:
-  - protein
-  - structure-prediction
-  - esmfold
-  - test-time-training
----
-
-# NOTE
-The GitHub with the implementation and requirements.txt can be found [here](https://github.com/Synthyra/FastPLMs.git)
-
-# FastESMFold
-
-FastESMFold is a self-contained, HuggingFace-compatible reimplementation of ESMFold with optional experimental **Test-Time Training (TTT)** and multi-backend attention (SDPA, Flash, Flex).
-
-No dependency on `fair-esm`, `proteinttt`, or `openfold`. Just `transformers`, `torch`, and `einops`.
-
-## Why Test-Time Training?
-
-Protein language models like ESM2 are trained on millions of sequences, but at inference time they process each new protein in a single forward pass with no adaptation. This is a missed opportunity: the input sequence itself contains structural signal that the model could learn from.
-
-**Test-Time Training (TTT)** adapts the model to each individual protein before predicting its structure. The idea is simple: before folding, we briefly train the ESM2 backbone on the input sequence using masked language modeling (the same objective it was pretrained with). This forces the model to "study" the specific sequence, strengthening its internal representation of that protein's structural features.
-
-TTT is disabled by default. Standard `fold_protein(...)`, `infer(...)`, and
-`state_dict()` behavior are unchanged unless you explicitly pass `ttt=True` or
-call `fold_protein_ttt(...)`.
-
-The adaptation uses **LoRA** (Low-Rank Adaptation) for efficiency: only small
-adapter weights are trained (~4.4M parameters out of 3.5B), and the base model
-is restored after each prediction. This takes 20-45 seconds per sequence on an
-A10G GPU. It can improve structure prediction quality on difficult targets
-where standard ESMFold produces low-confidence predictions, but it is
-experimental and can degrade predictions that already have high confidence.
-
-**When is TTT most useful?**
-- Sequences with low baseline pLDDT (< 0.5): TTT can improve pLDDT by 10-30+ points
-- Novel proteins with limited homology in training data
-- Disordered or multi-domain proteins where ESMFold struggles
-
-**When is TTT unnecessary?**
-- Sequences that already fold well (baseline pLDDT > 0.7): TTT rarely helps and may slightly degrade predictions
-- High-throughput screening where speed matters more than accuracy
-
-## Key Features
-
-- **Standard ESMFold**: Full ESMFold v1 structure prediction, loadable via `AutoModel`
-- **Optional experimental TTT**: Enable test-time training for difficult sequences with explicit `ttt=True`
-- **Best structure selection**: When TTT is enabled, folds after each step and returns the structure with the highest pLDDT
-- **FastESM2 attention**: SDPA/Flash/Flex backends for the 3B ESM2 backbone
-- **Self-contained LoRA**: lora_diffusion-compatible implementation (no peft dependency)
-- **3.5B parameters**: Full ESMFold v1 architecture (ESM2-3B backbone + folding trunk)
-
-## Use with transformers
-
-### Standard structure prediction (no TTT)
-
-```python
-import torch
-from transformers import AutoModel
-
-model = AutoModel.from_pretrained(
-    "Synthyra/FastESMFold",
-    trust_remote_code=True,
-    dtype=torch.float32,
-).cuda().eval()
-
-# Standard fold (no TTT)
-with torch.no_grad():
-    output = model.infer("MKTLLILAVVAAALA...")
-pdb_strings = model.output_to_pdb(output)
-plddt = output["plddt"].mean().item()
-print(f"pLDDT: {plddt:.3f}")
-```
-
-### Structure prediction with experimental TTT
-
-TTT adapts the ESM2 backbone to a specific input sequence via masked language
-modeling before folding. It can improve pLDDT on difficult sequences, but it is
-experimental, adds test-time compute, and should not be assumed to improve every
-sequence.
-
-```python
-# Configure TTT
-model._ttt_cfg.steps = 10      # 10 optimizer steps (default)
-model._ttt_cfg.lora_rank = 8   # LoRA rank (default)
-model._ttt_cfg.lora_alpha = 32 # LoRA scale (default)
-
-# ttt=True runs TTT, folds after each step, returns best structure
-result = model.fold_protein("MKTLLILAVVAAALA...", ttt=True)
-# Equivalent:
-# result = model.fold_protein_ttt("MKTLLILAVVAAALA...")
-print(f"pLDDT: {result['plddt']:.3f}")
-print(f"Best step: {result['best_step']} (0=baseline, 1-10=TTT steps)")
-print(f"Step pLDDTs: {[f'{p:.2f}' for p in result['step_plddts']]}")
-
-# Save PDB
-with open("structure.pdb", "w") as f:
-    f.write(result["pdb_string"])
-```
-
-### Return values
-
-`fold_protein(sequence)` returns a dict. Without `ttt=True`, `step_plddts`
-contains only the baseline pLDDT and `best_step` is `0`.
-
-| Key | Type | Description |
-|-----|------|-------------|
-| `plddt` | float | Mean pLDDT for the selected structure |
-| `ptm` | float | Predicted TM-score for the selected structure |
-| `pdb_string` | str | PDB format structure |
-| `step_plddts` | list[float] | Baseline pLDDT, plus per-step pLDDT when TTT is enabled |
-| `best_step` | int | Which step produced the selected structure (0=baseline) |
-
-### TTT default behavior
-
-TTT is disabled by default. Use FastESMFold as a standard ESMFold by calling
-`fold_protein(...)` or `infer(...)` without `ttt=True`:
-
-```python
-# Baseline fold, no TTT
-result = model.fold_protein("MKTLLILAVVAAALA...")
-print(result["best_step"])  # 0
-
-# Raw ESMFold output
-with torch.no_grad():
-    output = model.infer("MKTLLILAVVAAALA...")
-pdb_strings = model.output_to_pdb(output)
-```
-
-## Experimental TTT Benchmark
-
-This benchmark is provided as an example of where TTT can help. It is not a
-guarantee of improvement on every sequence.
-
-Tested on 10 difficult sequences on A10G GPU:
-
-| Metric | Value |
-|--------|-------|
-| Mean baseline pLDDT | 0.549 |
-| Mean best TTT pLDDT | 0.637 |
-| Mean improvement | +0.088 |
-| Sequences improved >5pt | 5/10 |
-| Time per sequence | ~20-45s |
-| GPU memory peak | 18.3 GB |
-
-On the hardest sequence (baseline pLDDT 0.38), TTT improves to 0.72 (+34 points).
-
-## Attention backends
-
-The ESM2 backbone supports multiple attention backends via `config.attn_backend`:
-
-| Backend | Key | Notes |
-| :--- | :--- | :--- |
-| PyTorch SDPA | `"sdpa"` | Default. Exact numerics, stable on all hardware. |
-| Flash Attention | `"kernels_flash"` | Fastest. Requires `pip install kernels`. |
-| Flex Attention | `"flex"` | Skips padding tokens via block mask. First use compiles a Triton kernel. |
-| Auto | `"auto"` | Picks best available: `kernels_flash` > `flex` > `sdpa`. |
-
-```python
-from transformers import AutoConfig, AutoModel
-
-config = AutoConfig.from_pretrained("Synthyra/FastESMFold", trust_remote_code=True)
-config.attn_backend = "kernels_flash"
-model = AutoModel.from_pretrained("Synthyra/FastESMFold", config=config, trust_remote_code=True)
-```
-
-## TTT Configuration
-
-TTT parameters are set via `config.ttt_config` (a dict) or by modifying `model._ttt_cfg` after loading:
-
-| Parameter | Default | Description |
-|-----------|---------|-------------|
-| `lr` | 4e-4 | Learning rate for SGD optimizer |
-| `steps` | 10 | Number of optimizer steps when TTT is explicitly enabled |
-| `ags` | 4 | Gradient accumulation steps per optimizer step |
-| `batch_size` | 4 | Batch size for masked language model training |
-| `mask_ratio` | 0.15 | Fraction of tokens to mask |
-| `lora_rank` | 8 | LoRA rank (0 for full backbone fine-tuning) |
-| `lora_alpha` | 32.0 | LoRA scaling factor (applied as `scale=alpha`, matching lora_diffusion) |
-| `seed` | 0 | Random seed for reproducible LoRA initialization and masking |
-| `lora_target_class` | `"EsmSelfAttention"` | Which module class to inject LoRA into |
-
-## How TTT Works
-
-1. **Baseline fold** (step 0): Standard ESMFold prediction
-2. **LoRA injection**: Rank-8 LoRA adapters on all `nn.Linear` layers inside ESM2 attention modules
-3. **Masked LM training**: 10 optimizer steps (each with 4 gradient accumulation sub-steps) of BERT-style masked language modeling on the input sequence
-4. **Per-step folding**: After each optimizer step, fold the sequence and record pLDDT
-5. **Best selection**: Return the structure with the highest pLDDT
-6. **Reset**: Restore LoRA weights to initial state for the next sequence
-
-## Citations
-
-```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
-@misc{bushuiev2026proteinneed,
-  title={One protein is all you need},
-  author={Anton Bushuiev and Roman Bushuiev and Olga Pimenova and Nikola Zadorozhny and Raman Samusevich and Elisabet Manaskova and Rachel Seongeun Kim and Hannes St\"ark and Jiri Sedlar and Martin Steinegger and Tom\'a\v{s} Pluskal and Josef Sivic},
-  year={2026},
-  eprint={2411.02109},
-  archivePrefix={arXiv},
-  primaryClass={cs.LG},
-  url={https://arxiv.org/abs/2411.02109}
-}
-```
-
-```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: "mit"
+tags:
+  - protein-language-model
+  - fastplms
+---
+
+<!-- Generated from src/fastplms/models.toml. Do not edit. -->
+
+# Synthyra/FastESMFold
+
+This checkpoint packages the FastPLMs `ESMFold` implementation.
+
+Accepted inputs are raw amino-acid sequences through folding helpers, or
+prepared residue tensors.
+Supported Transformers entry points are `AutoConfig`, `AutoModel`.
+
+## Install and platform requirements
+
+Install FastPLMs from the exact source revision paired with this model card:
+
+```bash
+python -m pip install \
+  "fastplms[structure] @ git+https://github.com/Synthyra/FastPLMs.git@1b9ce023f1e06571cf3e6324be0610ffa53e0a4a"
+```
+
+Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. Structure inference requires the `structure` extra and a CUDA device for the published execution contract. The current validated release target is the exact NVIDIA GH200 on Linux aarch64; Linux x86-64, CPU-only, Windows, and macOS structure runs are not current release evidence. 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/FastESMFold"
+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/FastESMFold` path, then pass `local_files_only=True`.
+
+Leave attention unspecified for the Transformers default. Supported explicit
+choices are `eager`, `sdpa`, `flex_attention`.
+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.
+
+## Protein structure prediction
+
+ESMFold accepts a raw sequence and returns structure tensors and confidence:
+
+```python
+import torch
+
+model = model.cuda().eval()
+with torch.inference_mode():
+    output = model.infer(
+        "MKTLLILAVVAAALA",
+        num_recycles=4,
+    )
+
+print(output["mean_plddt"])
+
+summary = model.fold_protein(
+    "MKTLLILAVVAAALA",
+    return_pdb_string=True,
+)
+with open("prediction.pdb", "w", encoding="utf-8") as handle:
+    handle.write(summary["pdb_string"])
+print(summary["plddt"], summary["ptm"])
+```
+
+FastPLMs does not expose ProteinTTT for ESMFold. The pinned folding checkpoint
+does not contain a trained masked-language-model head for that objective, so
+`ttt()` and TTT folding requests raise explicitly.
+
+## Runtime contract
+
+- Public input: Raw amino-acid sequences through folding helpers, or prepared residue tensors
+- Advertised AutoClasses: `AutoConfig`, `AutoModel`
+- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`
+- Attention implementations: `eager`, `sdpa`, `flex_attention`
+- Precision policies: `default`
+- BF16 execution: `fp32_parameters_autocast`
+- Generation contract: `not_applicable`
+- Optional dependency group: `structure`
+- Weight publication allowed: `true`
+- Weight license status: `resolved`
+- Redistributable: `true`
+- Complete weight publication required: `false`
+
+## Provenance
+
+- FastPLMs weights: `Synthyra/FastESMFold@b88c8cb50d19b2cf7ab4fee4b0a61f5e02da7823`
+- Runtime revision: `1b9ce023f1e06571cf3e6324be0610ffa53e0a4a`
+- Runtime source-tree SHA-256: `30b1fc728d3d1fe637bd841b4cbd1b951713e6d19bb46fb5a6aeeed84c5da9e6`
+- Runtime bundle SHA-256: `862caebb47b831591c549e4cd86d9a4c3f1eea550c044a8e617eb985d6209ab0`
+- Generator/schema version and complete/runtime-only attestations: recorded in `provenance.json`
+- Official checkpoint: `facebook/esmfold_v1@75a3841ee059df2bf4d56688166c8fb459ddd97a`
+- Artifact source: `fast`
+- State transform: `esmfold_meta_to_fastplms_v1`
+- BF16 execution: `fp32_parameters_autocast`
+- Pinned upstreams: `fair-esm`, `openfold`
+- Reference container: `reference-esmfold`
+- Release tiers: `check`, `compliance`, `structure`, `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: MIT. The Hub model-card identifier is
+`mit`. 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

@@ -5,8 +5,8 @@
   "attention_probs_dropout_prob": 0.0,
   "attn_backend": "sdpa",
   "auto_map": {
-    "AutoConfig": "modeling_fast_esmfold.FastEsmFoldConfig",
-    "AutoModel": "modeling_fast_esmfold.FastEsmForProteinFolding"
+    "AutoConfig": "modeling_fastplms.FastEsmFoldConfig",
+    "AutoModel": "modeling_fastplms.FastEsmForProteinFolding"
   },
   "classifier_dropout": null,
   "dtype": "float32",
@@ -52,6 +52,16 @@
     },
     "use_esm_attn_map": false
   },
+  "fastplms_checkpoint_hash": "50adf597cd90d0d7db0e3684f2e73c62297fc556cf1dad459c923977ad8108ba",
+  "fastplms_checkpoint_repo_id": "Synthyra/FastESMFold",
+  "fastplms_checkpoint_revision": "b88c8cb50d19b2cf7ab4fee4b0a61f5e02da7823",
+  "fastplms_model_id": "esmfold",
+  "fastplms_release_tool_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "fastplms_release_tool_sha256": "1459b5d7d13d9b07bd97b3eee764f2ce73623e15e32d07ddf6825c2a9509afb9",
+  "fastplms_runtime_bundle_sha256": "862caebb47b831591c549e4cd86d9a4c3f1eea550c044a8e617eb985d6209ab0",
+  "fastplms_runtime_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "fastplms_source_tree_sha256": "30b1fc728d3d1fe637bd841b4cbd1b951713e6d19bb46fb5a6aeeed84c5da9e6",
+  "fastplms_weights_revision": "b88c8cb50d19b2cf7ab4fee4b0a61f5e02da7823",
   "hidden_act": "gelu",
   "hidden_dropout_prob": 0.0,
   "hidden_size": 2560,

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/_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/esmfold/modeling_fast_esmfold.py

@@ -0,0 +1,894 @@
+"""FastESMFold with FastESM2 attention.
+
+Usage:
+    from transformers import AutoModel
+    model = AutoModel.from_pretrained("Synthyra/FastESMFold", trust_remote_code=True).cuda()
+
+    # Basic folding, no TTT
+    result = model.fold_protein("MKTLLILAVVA...")
+    print(result["plddt"], result["pdb_string"][:100])
+
+The runtime uses public Transformers folding components. It does not import or
+depend on the pinned fair-esm or OpenFold parity repositories.
+"""
+
+from __future__ import annotations
+
+from contextvars import ContextVar
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+import torch.nn as nn
+from einops import rearrange
+from torch.nn import functional as F
+from transformers.modeling_outputs import ModelOutput
+from transformers.models.esm.configuration_esm import EsmConfig
+from transformers.models.esm.modeling_esm import (
+    EsmEmbeddings,
+    EsmIntermediate,
+    EsmOutput,
+    EsmSelfOutput,
+)
+from transformers.models.esm.modeling_esmfold import (
+    EsmForProteinFolding,
+    collate_dense_tensors,
+)
+from transformers.models.esm.openfold_utils import residue_constants
+
+from fastplms.models._esm_rotary import RotaryEmbedding
+
+# Hub composite artifacts define these shared names earlier in the assembled file.
+try:
+    from fastplms.attention import (
+        AttentionBackend,
+        BlockMask,
+        FastPLMsAttentionMixin,
+        _get_flex_attention_fn,
+        flex_attention,
+        get_attention_mask,
+        kernels_flash_attention_func,
+        resolve_attention_backend,
+        resolve_attention_backend_for_call,
+    )
+except ModuleNotFoundError as error:
+    _COMPOSITE_REQUIRED_NAMES = (
+        "AttentionBackend",
+        "BlockMask",
+        "FastPLMsAttentionMixin",
+        "_get_flex_attention_fn",
+        "flex_attention",
+        "get_attention_mask",
+        "kernels_flash_attention_func",
+        "resolve_attention_backend",
+        "resolve_attention_backend_for_call",
+    )
+    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.
+
+
+# =============================================================================
+# Output Dataclass
+# =============================================================================
+
+
+@dataclass
+class FastEsmEncoderOutput(ModelOutput):
+    last_hidden_state: torch.Tensor | None = None
+    hidden_states: tuple[torch.Tensor, ...] | None = None
+    attentions: tuple[torch.Tensor, ...] | None = None
+
+
+@dataclass
+class FastEsmForProteinFoldingOutput(ModelOutput):
+    """Folding output with a standard Transformers AutoModel prefix."""
+
+    last_hidden_state: torch.Tensor | None = None
+    hidden_states: tuple[torch.Tensor, ...] | None = None
+    attentions: tuple[torch.Tensor, ...] | None = None
+    frames: torch.Tensor | None = None
+    sidechain_frames: torch.Tensor | None = None
+    unnormalized_angles: torch.Tensor | None = None
+    angles: torch.Tensor | None = None
+    positions: torch.Tensor | None = None
+    states: torch.Tensor | None = None
+    s_s: torch.Tensor | None = None
+    s_z: torch.Tensor | None = None
+    distogram_logits: torch.Tensor | None = None
+    lm_logits: torch.Tensor | None = None
+    aatype: torch.Tensor | None = None
+    atom14_atom_exists: torch.Tensor | None = None
+    residx_atom14_to_atom37: torch.Tensor | None = None
+    residx_atom37_to_atom14: torch.Tensor | None = None
+    atom37_atom_exists: torch.Tensor | None = None
+    residue_index: torch.Tensor | None = None
+    lddt_head: torch.Tensor | None = None
+    plddt: torch.Tensor | None = None
+    ptm_logits: torch.Tensor | None = None
+    ptm: torch.Tensor | None = None
+    aligned_confidence_probs: torch.Tensor | None = None
+    predicted_aligned_error: torch.Tensor | None = None
+    max_predicted_aligned_error: torch.Tensor | None = None
+    mlm_targets: torch.Tensor | None = None
+
+
+# ``EsmForProteinFolding.forward`` calls
+# ``compute_language_model_representations`` without forwarding output controls.
+# Context variables bridge that private call boundary without mutating the model
+# instance, so concurrent calls can independently request attention tensors.
+_ESMFOLD_OUTPUT_ATTENTIONS: ContextVar[bool] = ContextVar(
+    "fastplms_esmfold_output_attentions",
+    default=False,
+)
+_ESMFOLD_CAPTURED_ATTENTIONS: ContextVar[
+    tuple[torch.Tensor, ...] | None
+] = ContextVar(
+    "fastplms_esmfold_captured_attentions",
+    default=None,
+)
+
+
+def _align_internal_esm_attentions(
+    attentions: tuple[torch.Tensor, ...],
+    residue_mask: torch.Tensor,
+) -> tuple[torch.Tensor, ...]:
+    """Remove internal BOS/EOS positions while preserving public padding slots."""
+
+    batch_size, sequence_length = residue_mask.shape
+    public_positions = torch.arange(sequence_length, device=residue_mask.device)
+    valid_lengths = residue_mask.to(dtype=torch.int64).sum(dim=-1, keepdim=True)
+    # Internal layout is BOS, compact biological residues, EOS, then padding.
+    # Public padding positions therefore advance by two rather than one.
+    internal_positions = public_positions.unsqueeze(0) + 1
+    internal_positions = internal_positions + (
+        public_positions.unsqueeze(0) >= valid_lengths
+    ).to(dtype=torch.int64)
+
+    aligned: list[torch.Tensor] = []
+    for attention in attentions:
+        if attention.shape[0] != batch_size or attention.shape[-2:] != (
+            sequence_length + 2,
+            sequence_length + 2,
+        ):
+            raise RuntimeError(
+                "FastESM attention shape does not match the folding input: "
+                f"got {tuple(attention.shape)} for residue mask "
+                f"{tuple(residue_mask.shape)}."
+            )
+        query_index = internal_positions[:, None, :, None].expand(
+            batch_size,
+            attention.shape[1],
+            sequence_length,
+            attention.shape[-1],
+        )
+        query_aligned = torch.gather(attention, dim=2, index=query_index)
+        key_index = internal_positions[:, None, None, :].expand(
+            batch_size,
+            attention.shape[1],
+            sequence_length,
+            sequence_length,
+        )
+        aligned.append(torch.gather(query_aligned, dim=3, index=key_index))
+    return tuple(aligned)
+
+
+# =============================================================================
+# FastESM2 Attention Layers (multi-backend: SDPA, Flash, Flex)
+# =============================================================================
+
+
+class EsmSelfAttention(nn.Module):
+    def __init__(self, config, position_embedding_type: str | None = None):
+        super().__init__()
+        if config.hidden_size % config.num_attention_heads != 0:
+            raise ValueError(
+                f"The hidden size ({config.hidden_size}) is not a multiple of the number "
+                f"of attention heads ({config.num_attention_heads})."
+            )
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+        self.query = nn.Linear(config.hidden_size, self.all_head_size)
+        self.key = nn.Linear(config.hidden_size, self.all_head_size)
+        self.value = nn.Linear(config.hidden_size, self.all_head_size)
+        self.scale = self.attention_head_size**-0.5
+
+        self.dropout_prob = config.attention_probs_dropout_prob
+        self.config = config
+        self.attn_backend = resolve_attention_backend(config.attn_backend)
+        self.position_embedding_type = position_embedding_type or config.position_embedding_type
+        self.rotary_embeddings = None
+        if self.position_embedding_type == "rotary":
+            self.rotary_embeddings = RotaryEmbedding(dim=self.attention_head_size)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask_2d: torch.Tensor | None = None,
+        attention_mask_4d: torch.Tensor | None = None,
+        flex_block_mask: BlockMask | None = None,
+        output_attentions: bool = False,
+    ) -> tuple[torch.Tensor, 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":
+            query_heads, key_heads = self.rotary_embeddings(query_heads, key_heads)
+
+        attn_output, attn_weights = self._attn(
+            query_heads,
+            key_heads,
+            value_heads,
+            attention_mask_2d=attention_mask_2d,
+            attention_mask_4d=attention_mask_4d,
+            flex_block_mask=flex_block_mask,
+            output_attentions=output_attentions,
+        )
+        return attn_output, attn_weights
+
+    def _attn(
+        self,
+        query_heads: torch.Tensor,
+        key_heads: torch.Tensor,
+        value_heads: torch.Tensor,
+        attention_mask_2d: torch.Tensor | None = None,
+        attention_mask_4d: torch.Tensor | None = None,
+        flex_block_mask: BlockMask | None = None,
+        output_attentions: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        if output_attentions:
+            return self._manual_attn(query_heads, key_heads, value_heads, attention_mask_4d)
+
+        if (
+            self.training
+            and self.dropout_prob > 0
+            and self.attn_backend == AttentionBackend.FLEX_ATTENTION
+        ):
+            raise RuntimeError(
+                "ESMFold flex_attention is inference-only when attention dropout is "
+                "nonzero. Use eager or SDPA for this training configuration."
+            )
+
+        if self.attn_backend == AttentionBackend.EAGER:
+            attn_output, _ = self._manual_attn(
+                query_heads, key_heads, value_heads, attention_mask_4d
+            )
+            return attn_output, None
+        if self.attn_backend.is_flash:
+            return self._kernels_flash_attn(query_heads, key_heads, value_heads, attention_mask_2d)
+        elif self.attn_backend == AttentionBackend.FLEX:
+            return self._flex_attn(
+                query_heads,
+                key_heads,
+                value_heads,
+                flex_block_mask,
+                attention_mask_2d,
+            )
+        elif self.attn_backend == AttentionBackend.SDPA:
+            return self._sdpa_attn(query_heads, key_heads, value_heads, attention_mask_4d)
+        else:
+            raise AssertionError(f"Unsupported resolved backend: {self.attn_backend}")
+
+    def _manual_attn(
+        self,
+        query_heads: torch.Tensor,
+        key_heads: torch.Tensor,
+        value_heads: torch.Tensor,
+        attention_mask_4d: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        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)")
+        return attn_output, attn_weights
+
+    def _kernels_flash_attn(
+        self,
+        query_heads: torch.Tensor,
+        key_heads: torch.Tensor,
+        value_heads: torch.Tensor,
+        attention_mask_2d: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, None]:
+        query_tokens = query_heads.transpose(1, 2).contiguous()
+        key_tokens = key_heads.transpose(1, 2).contiguous()
+        value_tokens = value_heads.transpose(1, 2).contiguous()
+        # Q is pre-scaled by self.scale in forward() -- pass softmax_scale=1.0
+        # to prevent the kernel from applying its default 1/sqrt(head_dim).
+        attn_output = kernels_flash_attention_func(
+            query_states=query_tokens,
+            key_states=key_tokens,
+            value_states=value_tokens,
+            attention_mask_2d=attention_mask_2d,
+            causal=False,
+            softmax_scale=1.0,
+            implementation=self.attn_backend.value,
+        )
+        return rearrange(attn_output, "b s h d -> b s (h d)"), None
+
+    def _flex_attn(
+        self,
+        query_heads: torch.Tensor,
+        key_heads: torch.Tensor,
+        value_heads: torch.Tensor,
+        flex_block_mask: BlockMask | None = None,
+        attention_mask_2d: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, None]:
+        if flex_attention is None:
+            raise RuntimeError("Flex attention is not available in this environment.")
+        fn = _get_flex_attention_fn(
+            device=query_heads.device,
+            dtype=query_heads.dtype,
+            shape=tuple(query_heads.shape),
+            mask_semantics="padding",
+        )
+        context_heads = fn(
+            query_heads, key_heads, value_heads, block_mask=flex_block_mask, scale=1.0
+        )
+        return rearrange(context_heads, "b h s d -> b s (h d)"), None
+
+    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]:
+        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 EsmAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.self = EsmSelfAttention(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_2d: torch.Tensor | None = None,
+        attention_mask_4d: torch.Tensor | None = None,
+        flex_block_mask: BlockMask | None = None,
+        output_attentions: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        hidden_states_ln = self.LayerNorm(hidden_states)
+        attn_output, attn_weights = self.self(
+            hidden_states_ln,
+            attention_mask_2d=attention_mask_2d,
+            attention_mask_4d=attention_mask_4d,
+            flex_block_mask=flex_block_mask,
+            output_attentions=output_attentions,
+        )
+        attention_output = self.output(attn_output, hidden_states)
+        return attention_output, attn_weights
+
+
+class EsmLayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.attention = EsmAttention(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_2d: torch.Tensor | None = None,
+        attention_mask_4d: torch.Tensor | None = None,
+        flex_block_mask: BlockMask | None = None,
+        output_attentions: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None]:
+        attention_output, attn_weights = self.attention(
+            hidden_states,
+            attention_mask_2d=attention_mask_2d,
+            attention_mask_4d=attention_mask_4d,
+            flex_block_mask=flex_block_mask,
+            output_attentions=output_attentions,
+        )
+        layer_output = self._feed_forward(attention_output)
+        return layer_output, attn_weights
+
+    def _feed_forward(self, attention_output: torch.Tensor) -> torch.Tensor:
+        attention_output_ln = self.LayerNorm(attention_output)
+        intermediate_output = self.intermediate(attention_output_ln)
+        return self.output(intermediate_output, attention_output)
+
+
+class FastEsmEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.attention_backend = resolve_attention_backend(config.attn_backend)
+        self.layer = nn.ModuleList([EsmLayer(config) for _ in range(config.num_hidden_layers)])
+        self.emb_layer_norm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+    ) -> FastEsmEncoderOutput:
+        all_hidden_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        effective_backend = resolve_attention_backend_for_call(
+            self.attention_backend,
+            output_attentions=output_attentions,
+        )
+        attention_mask_2d, attention_mask_4d, flex_block_mask = 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)
+
+            hidden_states, attn_weights = layer_module(
+                hidden_states,
+                attention_mask_2d=attention_mask_2d,
+                attention_mask_4d=attention_mask_4d,
+                flex_block_mask=flex_block_mask,
+                output_attentions=output_attentions,
+            )
+
+            if all_attentions is not None:
+                all_attentions = (*all_attentions, attn_weights)
+
+        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 FastEsmEncoderOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+            attentions=all_attentions,
+        )
+
+
+# =============================================================================
+# FastESM Backbone (replaces EsmModel inside ESMFold)
+# =============================================================================
+
+
+class FastEsmBackbone(nn.Module):
+    """FastESM2 backbone with multi-backend attention. Drop-in replacement for
+    transformers.EsmModel inside EsmForProteinFolding.
+
+    Folding uses hidden states only. The standalone ESM2 contact regressor and
+    masked-LM head are therefore omitted from this structure-only backbone.
+    """
+
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.embeddings = EsmEmbeddings(config)
+        self.encoder = FastEsmEncoder(config)
+
+    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,
+        return_dict: bool | None = None,
+    ) -> FastEsmEncoderOutput | tuple[Any, ...]:
+        output_attentions = (
+            self.config.output_attentions
+            if output_attentions is None
+            else output_attentions
+        )
+        output_hidden_states = (
+            self.config.output_hidden_states
+            if output_hidden_states is None
+            else output_hidden_states
+        )
+        return_dict = self.config.use_return_dict if return_dict is None else return_dict
+
+        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 = FastEsmEncoderOutput(
+            last_hidden_state=encoder_outputs.last_hidden_state,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )
+        return output if return_dict else output.to_tuple()
+
+
+# =============================================================================
+# FastEsmFoldConfig
+# =============================================================================
+
+
+class FastEsmFoldConfig(EsmConfig):
+    model_type = "fast_esmfold"
+
+    def __init__(self, attn_backend: str | None = None, **kwargs: Any) -> None:
+        # Earlier mirrors serialized an untrained ESMFold-specific TTT policy.
+        # It is intentionally ignored because the official checkpoint has no
+        # trained masked-language-model head.
+        kwargs.pop("ttt_config", None)
+        super().__init__(**kwargs)
+        self.attn_backend = attn_backend
+
+
+# =============================================================================
+# FastEsmForProteinFolding
+# =============================================================================
+
+
+class FastEsmForProteinFolding(FastPLMsAttentionMixin, EsmForProteinFolding):
+    """ESMFold with FastESM2 attention backends.
+
+    Inherits all folding logic (trunk, structure module, output_to_pdb, infer)
+    from transformers.EsmForProteinFolding. Replaces the ESM2 backbone with
+    FastESM2 for selectable attention implementations.
+
+    Key API:
+        result = model.fold_protein("MKTL...")
+        # result = {"plddt": float, "ptm": float, "pdb_string": str}
+    """
+
+    config_class = FastEsmFoldConfig
+    _supports_flash_attn_2 = False
+    _supports_flash_attn_3 = False
+    _fastplms_attention_implementations = ("eager", "sdpa", "flex_attention")
+
+    def __init__(self, config: FastEsmFoldConfig):
+        super().__init__(config)
+
+        # Replace the standard ESM2 backbone with the multi-backend FastESM2
+        # implementation while retaining the canonical checkpoint key schema.
+        self.esm = FastEsmBackbone(config)
+        self.esm.requires_grad_(False)
+        if config.esmfold_config.fp16_esm:
+            self.esm.half()
+
+    def compute_language_model_representations(
+        self,
+        esmaa: torch.Tensor,
+    ) -> torch.Tensor:
+        """Run the internal ESM stem with a structured output unconditionally.
+
+        The outer folding model still honors ``config.return_dict``. This
+        internal call must remain structured because the folding stem selects
+        hidden states by name before constructing the public output.
+        """
+
+        device = next(self.parameters()).device
+        batch_size, sequence_length = esmaa.shape
+        output_attentions = _ESMFOLD_OUTPUT_ATTENTIONS.get()
+        if self.config.esmfold_config.bypass_lm:
+            if output_attentions:
+                _ESMFOLD_CAPTURED_ATTENTIONS.set(())
+            return torch.zeros(
+                batch_size,
+                sequence_length,
+                self.esm_s_combine.size(0),
+                self.esm_feats,
+                device=device,
+            )
+
+        bos = esmaa.new_full((batch_size, 1), self.esm_dict_cls_idx)
+        eos = esmaa.new_full((batch_size, 1), self.esm_dict_padding_idx)
+        residue_mask = esmaa != self.esm_dict_padding_idx
+        with_special_tokens = torch.cat([bos, esmaa, eos], dim=1)
+        with_special_tokens[
+            range(batch_size),
+            (with_special_tokens != self.esm_dict_padding_idx).sum(1),
+        ] = self.esm_dict_eos_idx
+        esm_output = self.esm(
+            with_special_tokens,
+            attention_mask=with_special_tokens != self.esm_dict_padding_idx,
+            output_attentions=output_attentions,
+            output_hidden_states=True,
+            return_dict=True,
+        )
+        if not isinstance(esm_output, FastEsmEncoderOutput):
+            raise TypeError("FastESM internal backbone did not return FastEsmEncoderOutput.")
+        if esm_output.hidden_states is None:
+            raise RuntimeError("FastESM internal backbone omitted requested hidden states.")
+        if output_attentions:
+            if esm_output.attentions is None:
+                raise RuntimeError("FastESM internal backbone omitted requested attentions.")
+            _ESMFOLD_CAPTURED_ATTENTIONS.set(
+                _align_internal_esm_attentions(esm_output.attentions, residue_mask)
+            )
+        return torch.stack(esm_output.hidden_states, dim=2)[:, 1:-1]
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        position_ids: torch.Tensor | None = None,
+        masking_pattern: torch.Tensor | None = None,
+        num_recycles: int | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+    ) -> FastEsmForProteinFoldingOutput | tuple[Any, ...]:
+        """Run folding with Meta ESMFold's 0-to-100 pLDDT convention."""
+
+        config = getattr(self, "config", None)
+        resolved_attentions = (
+            bool(getattr(config, "output_attentions", False))
+            if output_attentions is None
+            else output_attentions
+        )
+        resolved_hidden_states = (
+            bool(getattr(config, "output_hidden_states", False))
+            if output_hidden_states is None
+            else output_hidden_states
+        )
+        resolved_return_dict = (
+            bool(getattr(config, "use_return_dict", True))
+            if return_dict is None
+            else return_dict
+        )
+
+        request_token = _ESMFOLD_OUTPUT_ATTENTIONS.set(bool(resolved_attentions))
+        capture_token = _ESMFOLD_CAPTURED_ATTENTIONS.set(None)
+        captured_attentions: tuple[torch.Tensor, ...] | None = None
+        try:
+            output = super().forward(
+                input_ids,
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                masking_pattern=masking_pattern,
+                num_recycles=num_recycles,
+                output_hidden_states=resolved_hidden_states,
+            )
+            captured_attentions = _ESMFOLD_CAPTURED_ATTENTIONS.get()
+        finally:
+            _ESMFOLD_CAPTURED_ATTENTIONS.reset(capture_token)
+            _ESMFOLD_OUTPUT_ATTENTIONS.reset(request_token)
+        # Transformers 5.13 returns categorical lDDT probabilities on [0, 1],
+        # while Meta ESMFold's public forward output reports pLDDT on [0, 100].
+        output["plddt"] = output["plddt"] * 100
+        payload = dict(output)
+        payload.pop("last_hidden_state", None)
+        payload.pop("hidden_states", None)
+        parent_attentions = payload.pop("attentions", None)
+        if captured_attentions is None:
+            captured_attentions = parent_attentions
+        if resolved_attentions and captured_attentions is None:
+            # A bypassed or injected folding stem has no attention layers, but
+            # still honors the output contract without fabricating tensors.
+            captured_attentions = ()
+        sequence_state = payload.get("s_s")
+        structured = FastEsmForProteinFoldingOutput(
+            last_hidden_state=sequence_state,
+            hidden_states=(sequence_state,)
+            if resolved_hidden_states and sequence_state is not None
+            else None,
+            attentions=captured_attentions if resolved_attentions else None,
+            **payload,
+        )
+        return structured if resolved_return_dict else structured.to_tuple()
+
+    @torch.no_grad()
+    def infer(
+        self,
+        sequences: str | list[str],
+        residx: torch.Tensor | list[torch.Tensor] | None = None,
+        masking_pattern: torch.Tensor | None = None,
+        num_recycles: int | None = None,
+        residue_index_offset: int | None = 512,
+        chain_linker: str | None = "G" * 25,
+    ):
+        """Fold raw sequences through Meta ESMFold's public input contract.
+
+        Transformers v5 narrows ``infer`` even though ``forward`` retains the
+        required controls. This adapter restores recycle selection, explicit
+        residue indices, masking, and colon-delimited multimer preparation.
+        """
+
+        sequence_batch = [sequences] if isinstance(sequences, str) else sequences
+        linker = "" if chain_linker is None else chain_linker
+        index_offset = 0 if residue_index_offset is None else residue_index_offset
+        unknown_index = residue_constants.restype_order_with_x["X"]
+        aatype_batch: list[torch.Tensor] = []
+        residx_batch: list[torch.Tensor] = []
+        linker_mask_batch: list[torch.Tensor] = []
+        chain_index_batch: list[torch.Tensor] = []
+
+        for sequence in sequence_batch:
+            chains = sequence.split(":")
+            joined_sequence = linker.join(chains)
+            encoded = torch.tensor(
+                [
+                    residue_constants.restype_order_with_x.get(residue, unknown_index)
+                    for residue in joined_sequence
+                ],
+                dtype=torch.int64,
+            )
+            sequence_residx = torch.arange(len(encoded), dtype=torch.int64)
+            cursor = 0
+            for chain_number, chain in enumerate(chains):
+                segment_length = len(chain) + len(linker)
+                sequence_residx[cursor : cursor + segment_length] += chain_number * index_offset
+                cursor += segment_length
+
+            linker_mask = torch.ones_like(encoded, dtype=torch.float32)
+            chain_indices: list[int] = []
+            cursor = 0
+            for chain_number, chain in enumerate(chains):
+                if chain_number > 0:
+                    chain_indices.extend([chain_number - 1] * len(linker))
+                chain_indices.extend([chain_number] * len(chain))
+                cursor += len(chain)
+                linker_mask[cursor : cursor + len(linker)] = 0
+                cursor += len(linker)
+
+            aatype_batch.append(encoded)
+            residx_batch.append(sequence_residx)
+            linker_mask_batch.append(linker_mask)
+            chain_index_batch.append(torch.tensor(chain_indices, dtype=torch.int64))
+
+        aatype = collate_dense_tensors(aatype_batch)
+        attention_mask = collate_dense_tensors(
+            [aatype.new_ones(len(encoded)) for encoded in aatype_batch]
+        )
+        prepared_residx = collate_dense_tensors(residx_batch)
+        linker_mask = collate_dense_tensors(linker_mask_batch)
+        chain_index = collate_dense_tensors(chain_index_batch, pad_v=-1)
+        if residx is None:
+            residx = prepared_residx
+        elif not isinstance(residx, torch.Tensor):
+            residx = collate_dense_tensors(residx)
+
+        device = next(self.parameters()).device
+        aatype = aatype.to(device)
+        attention_mask = attention_mask.to(device)
+        residx = residx.to(device)
+        linker_mask = linker_mask.to(device)
+        output = self.forward(
+            aatype,
+            attention_mask,
+            position_ids=residx,
+            masking_pattern=masking_pattern,
+            num_recycles=num_recycles,
+        )
+        output["atom37_atom_exists"] = output["atom37_atom_exists"] * linker_mask.unsqueeze(2)
+        output["mean_plddt"] = (output["plddt"] * output["atom37_atom_exists"]).sum(
+            dim=(1, 2)
+        ) / output["atom37_atom_exists"].sum(dim=(1, 2))
+        output["chain_index"] = chain_index
+        return output
+
+    @staticmethod
+    def _ttt_unavailable() -> None:
+        raise RuntimeError(
+            "ESMFold TTT is unavailable: the pinned Meta ESMFold checkpoint does "
+            "not contain a trained masked-language-model head. FastPLMs does not "
+            "construct or serialize a random replacement head."
+        )
+
+    def ttt(self, seq: str, **kwargs: Any) -> None:
+        """Reject ESMFold-specific TTT because no faithful MLM objective exists."""
+
+        del seq, kwargs
+        self._ttt_unavailable()
+
+    def ttt_reset(self) -> None:
+        """Reject reset because ESMFold does not expose a faithful TTT path."""
+
+        self._ttt_unavailable()
+
+    def _fold_single(self, sequence: str, return_pdb_string: bool = True) -> dict[str, Any]:
+        """Fold a sequence once and return pLDDT, ptm, and optionally PDB string."""
+        with torch.no_grad():
+            output = self.infer(sequence)
+        if "mean_plddt" in output:
+            # ``infer`` masks multimer linker atoms before computing this mean.
+            # Reusing it prevents synthetic linker confidence from affecting the
+            # public fold_protein summary.
+            mean_plddt = float(output["mean_plddt"].reshape(-1)[0].item())
+        else:
+            plddt = output["plddt"]
+            # P has shape (b, l, 37), with confidence for each atom37 position.
+            # Use CA atom (index 1) only, matching PDB B-factor output.
+            if plddt.dim() == 3:
+                mean_plddt = float(plddt[:, :, 1].mean().item())
+            elif plddt.dim() == 2:
+                mean_plddt = float(plddt[:, 1].mean().item())
+            else:
+                mean_plddt = float(plddt.mean().item())
+        result = {
+            "plddt": mean_plddt,
+            "ptm": float(output["ptm"].item()) if "ptm" in output else None,
+        }
+        if return_pdb_string:
+            pdb_strings = self.output_to_pdb(output)
+            result["pdb_string"] = pdb_strings[0] if isinstance(pdb_strings, list) else pdb_strings
+        return result
+
+    def fold_protein(
+        self,
+        sequence: str,
+        return_pdb_string: bool = True,
+        ttt: bool = False,
+    ) -> dict[str, Any]:
+        """Fold a protein sequence.
+
+        Passing ``ttt=True`` fails explicitly because the official Meta
+        checkpoint contains no trained masked-language-model head.
+
+        Args:
+            sequence: Protein sequence (single-letter amino acid codes)
+            return_pdb_string: If True, include PDB string in output
+            ttt: Reserved rejection flag for unsupported ESMFold TTT
+
+        Returns:
+            Dict with keys:
+                - plddt: float, mean pLDDT
+                - ptm: float, predicted TM-score
+                - pdb_string: str (if return_pdb_string=True), PDB from best step
+                - step_plddts: list[float], baseline pLDDT when TTT is disabled
+                - best_step: int, 0 when TTT is disabled
+        """
+        if ttt:
+            return self.fold_protein_ttt(
+                sequence=sequence,
+                return_pdb_string=return_pdb_string,
+            )
+        result = self._fold_single(sequence, return_pdb_string=return_pdb_string)
+        return {
+            "plddt": result["plddt"],
+            "ptm": result["ptm"],
+            "pdb_string": result.get("pdb_string"),
+            "step_plddts": [result["plddt"]],
+            "best_step": 0,
+        }
+
+    def fold_protein_ttt(
+        self,
+        sequence: str,
+        return_pdb_string: bool = True,
+    ) -> None:
+        """Reject ESMFold TTT because the official checkpoint has no MLM head."""
+
+        del sequence, return_pdb_string
+        self._ttt_unavailable()

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,230 @@
+"""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 != "862caebb47b831591c549e4cd86d9a4c3f1eea550c044a8e617eb985d6209ab0":
+    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.esmfold.modeling_fast_esmfold")
+FastEsmFoldConfig = _module_225.FastEsmFoldConfig
+FastEsmFoldConfig.__module__ = __name__
+FastEsmForProteinFolding = _module_225.FastEsmForProteinFolding
+FastEsmForProteinFolding.__module__ = __name__

runtime-attestation.json

@@ -0,0 +1,49 @@
+{
+  "files": {
+    "LICENSES/FastPLMs-Apache-2.0.txt": "sha256:2d2b50c7b1414bff1189a1db1f0cfb92e3e064b50f4c2b1019827b683e1b629a",
+    "LICENSES/fair-esm/LICENSE": "sha256:da6d3703ed11cbe42bd212c725957c98da23cbff1998c05fa4b3d976d1a58e93",
+    "LICENSES/fair-esm/PROVENANCE.md": "sha256:950adb94daf15e646ddf226dacfe2a8e77801aa0793e439a9a3490a48eb666e7",
+    "LICENSES/openfold/LICENSE": "sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30",
+    "LICENSES/openfold/MODIFICATIONS.md": "sha256:fd6f0aa1086a0c996cf967b326d18e965660cda0ad5c7f36a3474a8490720da3",
+    "LICENSES/openfold/PROVENANCE.md": "sha256:48c903db43a217a3126afaefbac60b7ddac7efda2dfcc0cbff0bffc7d6c30081",
+    "README.md": "sha256:3cd1ee142900b7b7c3676877bc57d7b7e235e7808ff1492a129f200426d0c460",
+    "THIRD_PARTY_NOTICES.md": "sha256:25704b3c76404696cae52e7fca13088d329f70f412687340351259e86cd62baa",
+    "config.json": "sha256:fe19346a53c8d8a4777a8835334d3d013979e64e27b1d81f6dfd22ad89214325",
+    "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/_esm_rotary.py": "sha256:4dcee8af74adf2cfcabeeaadf4915f1da5fb6c0672d583bf3ee5478ac1290805",
+    "fastplms/models/esmfold/__init__.py": "sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
+    "fastplms/models/esmfold/modeling_fast_esmfold.py": "sha256:8c1d353f3203d8e6bfea6637463cb38da9e18b13ad6f60c672f213f5af4c2194",
+    "fastplms/registry.py": "sha256:afca271911b651a882345b74a58366494a1d784e4a48c8683a87f5508f4ba16e",
+    "fastplms/runtime.py": "sha256:110018646d6f248cedab140a030c3065e1b062b61f6aff659c231e538614bc01",
+    "fastplms_bundle.py": "sha256:a620d776d5d02beeab5aad1ac6d0418ba25a017cbd1510ca933a671944073e43",
+    "modeling_fastplms.py": "sha256:8942fdf5ce62c2e19cdc9a9bb53567086491c16780c58adc56d0b7ffb66733e5",
+    "special_tokens_map.json": "sha256:5f89dfed5b1bcb31b03a93b608b8c7fc5ac8465f2f4b839f6f6f4aa87d683407",
+    "tokenizer_config.json": "sha256:01218ac34c07219089c4c5de01c3196cc44bd0bb9cc08ae3777c1a6712ce2be6",
+    "vocab.txt": "sha256:927f51c652e9367fc6a8403c1993df3633df51d79df602ce1c207bf68ad4d867"
+  },
+  "model_id": "esmfold",
+  "redistributable": true,
+  "release_tool_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "release_tool_sha256": "1459b5d7d13d9b07bd97b3eee764f2ce73623e15e32d07ddf6825c2a9509afb9",
+  "runtime_bundle_sha256": "862caebb47b831591c549e4cd86d9a4c3f1eea550c044a8e617eb985d6209ab0",
+  "runtime_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "schema_version": 2,
+  "scope": "runtime-only",
+  "source_tree_sha256": "30b1fc728d3d1fe637bd841b4cbd1b951713e6d19bb46fb5a6aeeed84c5da9e6",
+  "weights": {
+    "repo_id": "Synthyra/FastESMFold",
+    "revision": "b88c8cb50d19b2cf7ab4fee4b0a61f5e02da7823"
+  },
+  "weights_license_status": "resolved"
+}