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/biohub-esm/LICENSE.md

@@ -0,0 +1,9 @@
+**License (MIT)**
+
+Copyright 2026 Chan Zuckerberg Biohub, Inc.
+
+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/biohub-esm/THIRD_PARTY_NOTICE.md

@@ -0,0 +1,13 @@
+The code in this repository depends on the following third-party libraries:
+
+| Library | License | Link |
+|----------|----------|----------|
+| flash-attn | BSD | https://github.com/Dao-AILab/flash-attention/blob/main/LICENSE |
+| PyTorch | BSD | https://github.com/pytorch/pytorch/blob/main/LICENSE |
+| xformers | BSD | https://github.com/facebookresearch/xformers/blob/main/LICENSE |
+| jaxtyping | MIT | https://github.com/patrick-kidger/jaxtyping/blob/main/LICENSE |
+| einops | MIT | https://github.com/arogozhnikov/einops/blob/main/LICENSE |
+| omegaconf | BSD | https://github.com/omry/omegaconf/blob/master/LICENSE |
+| attrs | MIT | https://github.com/python-attrs/attrs/blob/main/LICENSE |
+| scipy | BSD-3-Clause | https://github.com/scipy/scipy/blob/main/LICENSE.txt<br>https://github.com/scipy/scipy/blob/main/LICENSES_bundled.txt |
+| lightning / torchmetrics | Apache 2.0 | https://github.com/Lightning-AI/torchmetrics/blob/master/LICENSE |

LICENSES/biohub-transformers/LICENSE

@@ -0,0 +1,203 @@
+Copyright 2018- The Hugging Face team. All rights reserved.
+
+                                 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/protein-ttt/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 Anton Bushuiev
+
+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/protein-ttt/PROVENANCE.md

@@ -0,0 +1,8 @@
+# ProteinTTT provenance
+
+FastPLMs uses `anton-bushuiev/ProteinTTT` revision
+`fde2817cd84b936167cc76ccabf31e5c0fe49962` as the official reference for the
+optional protein test-time training workflow. The repository is pinned at
+`vendor/upstream/protein-ttt/` and is not a production dependency or runtime
+image component. The accompanying `LICENSE` is the verbatim MIT text from that
+revision.

README.md

@@ -1,247 +1,268 @@
----
-library_name: transformers
-tags:
-  - biology
-  - protein-structure
-  - esmfold2
-  - multimodal-protein-model
----
-
-# FastPLMs ESMFold2
-
-FastPLMs ESMFold2 is a self-contained Hugging Face `AutoModel` wrapper for
-Biohub's ESMFold2, ESMFold2-Fast, and experimental ESMFold2 structure
-predictors. It vendors the released Biohub ESMFold2 model code, input builder,
-MSA helpers, and structure export utilities, while loading the PLM backbone
-through FastPLMs ESM++.
-
-## Load With AutoModel
-
-```python
-import torch
-from transformers import AutoModel
-
-model = AutoModel.from_pretrained(
-    "Synthyra/ESMFold2-Fast",
-    trust_remote_code=True,
-    dtype=torch.float32,
-).eval().cuda()
-```
-
-Use `Synthyra/ESMFold2` for the full model, `Synthyra/ESMFold2-Fast` for the
-faster release variant, and the `Synthyra/ESMFold2-Experimental*` checkpoints
-for differentiable binder design and experimental critic ensembles.
-The folding trunk runs in fp32; the 6B FastPLMs ESM++ backbone is loaded in
-bf16 by default via `esmc_precision="bf16"` and uses the flex attention backend
-by default inside ESMFold2.
-
-## Fold One Protein
-
-```python
-sequence = "MKTLLILAVVAAALA"
-
-result = model.fold_protein(
-    sequence,
-    num_loops=3,
-    num_sampling_steps=50,
-    num_diffusion_samples=1,
-    seed=0,
-)
-
-print(float(result.plddt.mean()))
-print(float(result.ptm))
-```
-
-## Experimental Test-Time Training
-
-TTT is disabled by default. Standard `fold_protein(...)`, `fold(...)`, raw tensor
-inference, and `state_dict()` keys are unchanged unless you explicitly pass
-`ttt=True` or call `fold_protein_ttt(...)`.
-
-The ESMFold2 TTT path is experimental and protein-only in v1. It trains local
-LoRA adapters only on `_esmc` with a masked language modeling objective. The
-folding trunk, confidence head, diffusion head, and structure input pipeline are
-frozen. TTT can improve difficult low-confidence folds, but it adds substantial
-test-time compute and can degrade already confident predictions.
-
-```python
-result = model.fold_protein(
-    "MSTNPKPQRKTKRNT",
-    num_loops=1,
-    num_sampling_steps=10,
-    num_diffusion_samples=1,
-    seed=0,
-    ttt=True,
-    ttt_config={
-        "steps": 1,
-        "ags": 1,
-        "batch_size": 1,
-        "lora_rank": 8,
-        "lora_alpha": 32.0,
-    },
-)
-
-print(result.ttt_metrics["losses"])
-print(result.ttt_metrics["step_plddts"])
-print(result.ttt_metrics["best_step"])
-```
-
-`load_esmc=True` is required for TTT because the ESM++ MLM head is loaded lazily
-from `config.esmc_id`. If that pretrained MLM head cannot be loaded, TTT raises
-an assertion instead of silently using a random head.
-
-## Save mmCIF or PDB
-
-```python
-model.save_as_cif(result, "prediction.cif")
-model.save_as_pdb(result, "prediction.pdb")
-
-cif_text = model.result_to_cif(result)
-pdb_text = model.result_to_pdb(result)
-```
-
-`result_to_cif` preserves the full `MolecularComplex`. `result_to_pdb` converts through Biohub's protein-only `ProteinComplex` representation, so use mmCIF for complexes with ligands or nucleic acids.
-
-## Fold Complexes
-
-```python
-types = model.input_types
-
-complex_input = types.StructurePredictionInput(
-    sequences=[
-        types.ProteinInput(id="A", sequence="MKTLLILAVVAAALA"),
-        types.DNAInput(id="B", sequence="GATAGC"),
-        types.LigandInput(id="L", ccd=["SAH"]),
-    ]
-)
-
-result = model.fold(
-    complex_input,
-    num_loops=3,
-    num_sampling_steps=50,
-    num_diffusion_samples=1,
-    seed=0,
-)
-
-model.save_as_cif(result, "complex_prediction.cif")
-```
-
-## Binder Design With FastPLMs ESMFold2
-
-FastPLMs includes a FastPLMs-only port of the Biohub ESMFold2 binder design
-tutorial at `cookbook/tutorials/binder_design_fastplms.py`. The workflow uses
-ESMFold2 experimental checkpoints for differentiable folding losses, ESM++ for
-sequence regularization, and ESMFold2 hero critics for final confidence scoring.
-
-![FastPLMs EGFR minibinder design](https://raw.githubusercontent.com/Synthyra/FastPLMs/main/docs/assets/egfr_fastplms_binder_design.png)
-
-The optimizer follows the official strategy:
-
-1. Optimize mutable `#` residues as continuous amino acid logits.
-2. Suppress cysteine design by masking cysteine logits and gradients.
-3. Backpropagate through ESMFold2 `res_type_soft` using intra-contact,
-   inter-contact, and globularity losses from the distogram.
-4. Add an ESM++ masked-LM pseudoperplexity regularizer on mutable binder
-   residues.
-5. Keep the late-trajectory sequence with the best iPTM.
-6. Fold the selected sequence with the final critic ensemble and write
-   `results.parquet`, `selection.parquet`, `trajectory.jsonl`,
-   `best_sequences.fasta`, and per-critic PDB/CIF/logit files.
-
-Run the verified EGFR 128 amino acid de novo minibinder example:
-
-```bash
-cd /home/ubuntu/FastPLMs
-
-sudo -n docker run --gpus all --rm \
-  -v /home/ubuntu/FastPLMs:/app \
-  -v /home/ubuntu/FastPLMs:/workspace \
-  -v /home/ubuntu/.cache/huggingface:/workspace/.cache/huggingface \
-  -w /workspace fastplms-esmfold2 \
-  python /app/cookbook/tutorials/binder_design_fastplms.py \
-    --backend local \
-    --target-name egfr \
-    --binder-sequence '################################################################################################################################' \
-    --not-antibody \
-    --steps 150 \
-    --batch-size 1 \
-    --seed 103 \
-    --output-dir /workspace/campaign_egfr_len128_b1_s150_seed103_consensus_cli
-```
-
-Verified result:
-
-| Metric | Value |
-| :--- | :--- |
-| Binder length | `128` |
-| Seed | `103` |
-| Steps | `150` |
-| Hero mean iPTM | `0.913870` |
-| Hero min iPTM | `0.904600` |
-| All four hero critics above 0.9 | `True` |
-
-Binder sequence:
-
-```text
-SAVKHLLEIVKYLEEAIEKALEVDPVFLVPPAAEELLIAAKVIKELAKENPELIEVYELLMKAVKGLKKLVRSNDKEILREVIRLLRKAAKVIREILKNNPDLDPELRKALEELAKVLEEIAEVLEQQ
-```
-
-See the full guide in [`docs/binder_design.md`](https://github.com/Synthyra/FastPLMs/blob/main/docs/binder_design.md)
-for Modal execution, official pI and selection scoring, per-critic metrics, and
-the tested cheaper step-count boundary.
-
-## Use MSAs
-
-```python
-types = model.input_types
-
-msa = types.MSA.from_a3m("query.a3m", max_sequences=128)
-input_with_msa = types.StructurePredictionInput(
-    sequences=[
-        types.ProteinInput(id="A", sequence=msa.query, msa=msa),
-    ]
-)
-
-result = model.fold(input_with_msa, num_sampling_steps=50, seed=0)
-```
-
-## Raw Tensor Inference
-
-```python
-features, chain_infos = model.prepare_structure_input(complex_input, seed=0)
-
-with torch.inference_mode():
-    output = model(
-        **features,
-        num_loops=3,
-        num_sampling_steps=50,
-        num_diffusion_samples=1,
-    )
-
-decoded = model.input_builder.decode(output, features, chain_infos)
-```
-
-Set `load_esmc=False` when loading if you want to provide precomputed `lm_hidden_states` manually or run folding-trunk tests without loading the 6B ESM++ backbone:
-
-```python
-model = AutoModel.from_pretrained(
-    "Synthyra/ESMFold2-Fast",
-    trust_remote_code=True,
-    load_esmc=False,
-).cuda().eval()
-```
-
-For FP8 LM inference, install `transformer_engine.pytorch` in a CUDA
-environment with FP8-capable hardware and load the shared FastPLMs ESM++
-backbone with:
-
-```python
-model = AutoModel.from_pretrained(
-    "Synthyra/ESMFold2-Fast",
-    trust_remote_code=True,
-    esmc_precision="fp8",
-).cuda().eval()
-```
-
-FP8 is inference-only for the ESMFold2 LM backbone. TTT remains a bf16/fp32
-path.
+---
+library_name: transformers
+license: "mit"
+tags:
+  - protein-language-model
+  - fastplms
+---
+
+<!-- Generated from src/fastplms/models.toml. Do not edit. -->
+
+# Synthyra/ESMFold2-Experimental-Cutoff2025
+
+This checkpoint packages the FastPLMs `ESMFold2` implementation.
+
+Accepted inputs are raw amino-acid sequences or typed molecular-complex
+specifications; low-level forward accepts prepared feature 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/ESMFold2-Experimental-Cutoff2025"
+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/ESMFold2-Experimental-Cutoff2025` 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.
+
+## Alignment-conditioning contract
+
+This is a full 48-block ESMFold2 checkpoint. It supports both
+single-sequence inference and optional MSA-conditioned inference. Typed
+multichain and multimolecule inputs may attach an MSA to each applicable
+protein chain.
+
+
+## Protein folding
+
+The single-protein helper returns typed structure and confidence outputs:
+
+```python
+result = model.fold_protein(
+    "MSTNPKPQRKTKRNT",
+    num_loops=1,
+    num_sampling_steps=200,
+    num_diffusion_samples=1,
+    seed=7,
+)
+pdb_text = model.result_to_pdb(result)
+cif_text = model.result_to_cif(result)
+print(result.ptm, result.plddt.mean().item())
+```
+
+No target structure is required. For complexes, construct the input from the
+types exposed by the loaded artifact:
+
+```python
+types = model.input_types
+complex_input = types.StructurePredictionInput(
+    sequences=[
+        types.ProteinInput(id="A", sequence="MSTNPKPQRKTKRNT"),
+        types.ProteinInput(id="B", sequence="MKTIIALSYIFCLVFA"),
+        types.DNAInput(id="C", sequence="ATGC"),
+        types.LigandInput(id="L", smiles="O"),
+    ]
+)
+complex_result = model.fold(
+    complex_input,
+    num_loops=1,
+    num_sampling_steps=200,
+    seed=7,
+)
+print(complex_result.ptm, complex_result.plddt.mean().item())
+```
+
+The typed interface also supports RNA, protein MSAs, modifications, covalent
+bonds, and distogram conditioning. The public schema recognizes
+`PocketConditioning`, but the pinned official runtime discards it and hard-codes
+a zero pocket feature. FastPLMs therefore rejects non-null pocket conditioning
+instead of silently ignoring it. Prepared `ref_pos` values are component
+reference geometries created during featurization, not target coordinates.
+Predicted coordinates and confidence scores are outputs and do not establish
+biochemical activity.
+
+## Learned representation and ESMC precision
+
+ESMFold2 combines the ordered 81 ESMC-6B states `H: (b, l, 81, 2560)` with the
+checkpoint's learned projection. Retrieve the resulting residue representation
+through the public embedding API:
+
+```python
+representations = model.embed_dataset(
+    ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"],
+    batch_size=2,
+    full_embeddings=True,
+)
+print(representations[0].tensor.shape)  # (sequence_length, 256)
+```
+
+`model.embed_dataset(..., full_embeddings=True)` returns one `(l, 256)` residue
+tensor per single-chain input. It rejects complexes, ligands, MSAs,
+chain-separated inputs, `cls`, and `parti` in the embedding path.
+
+Set `esmc_precision` to `auto`, `bf16`, `fp32`, or `fp8` when loading.
+`auto` always resolves to BF16. Explicit FP8 is experimental, inference-only,
+and strict:
+
+```python
+model.reload_esmc(precision="fp8", device="cuda:0")
+print(model.esmc_precision_status)
+```
+
+FP8 raises when the validated CUDA and Transformer Engine path is unavailable.
+Canonical BF16 weights are retained, and transient quantization state is never
+serialized.
+
+The ESMC backbone uses SDPA as the recommended highest-fidelity path. Flex
+Attention is supported and non-experimental but can be numerically divergent;
+ESMFold2 does not advertise FlashAttention for the folding interface.
+
+| Backend | Support | Measurement status |
+| --- | --- | --- |
+| `sdpa` | Recommended fidelity path | Pending complete validated 30-record frozen-head GH200/aarch64 set |
+| `eager` | Supported | Pending complete validated 30-record frozen-head GH200/aarch64 set |
+| `flex_attention` | Supported, numerically divergent | Pending complete validated 30-record frozen-head GH200/aarch64 set |
+
+No threshold, report from another checkpoint, or result from another
+accelerator is substituted for a measurement. A release set contains all
+30 model/backend/panel records from one exact GH200 device and aarch64 runtime:
+18 eager/SDPA/Flex measurements include relative L2, Q99.9, residue
+cosine, pooled cosine, top-1, and Jensen-Shannon distributions; 12
+FlashAttention 2/3 records explicitly attest locked-platform unavailability.
+
+## Locked oracle package compatibility exception
+
+The frozen oracle lock permits exactly one nonzero `pip check` diagnostic:
+`nvidia-cusparselt-cu13 0.8.1 is not supported on this platform`. It applies only to
+`nvidia-cusparselt-cu13==0.8.1` on
+`NVIDIA GH200 480GB` / `linux` /
+`aarch64`. The vendor filename tag is
+`py3-none-manylinux2014_aarch64`, while the wheel metadata declares
+`py3-none-manylinux2014_sbsa`. The exact wheel is
+`nvidia_cusparselt_cu13-0.8.1-py3-none-manylinux2014_aarch64.whl` with SHA-256 `4dca476c50bf4780d46cd0bfbd82e2bc10a08e4fef7950917ce8d7578d22a23f`.
+FastPLMs accepts this vendor metadata mismatch only after the lock, installed
+inventory, wheel bytes, metadata tag, and target identity all match. The wheel
+is not rewritten (`validated-vendor-metadata-exception-no-wheel-rewrite`). Any additional diagnostic or
+identity drift fails closed.
+
+
+Metrics must be tied to the exact ESMFold2 and ESMC revisions, dtype, current
+GH200/aarch64 device and container images, dependency lock, source attestations,
+and sequence panel. Pending cells are not performance or parity claims.
+
+## Hash-pinned CCD runtime asset
+
+Structure preparation requires `ccd.pkl` from
+`biohub/ESMFold2@1ebf0e3481a5184eb6171d40615c79e384b48796`. The manifest pins
+its 417,306,584-byte size and SHA-256
+`9ff44b1927c6b9198e38ffe0928706827a09a350c15530beeeabebfa88038fc5`
+under MIT terms. This is a trusted-deserialization boundary: FastPLMs only
+allows the exact manifest repository/revision snapshot link to resolve within
+that repository's contained blob directory; user-supplied asset and `cache_dir`
+symlinks are rejected. The loader creates a private temporary snapshot, verifies
+its size and SHA-256, and unpickles only that loader-owned snapshot, closing
+path-replacement and in-place source-write races. Offline execution requires the
+exact cache object and never downloads a replacement.
+
+## Binder-design research example
+
+The FastPLMs binder-design workflow uses the experimental Fast Cutoff2025
+checkpoint for differentiable inversion, both experimental Cutoff2025
+checkpoints as critics, and ESM++ as the sequence prior:
+
+![FastPLMs EGFR minibinder design](https://raw.githubusercontent.com/Synthyra/FastPLMs/main/docs/assets/egfr_fastplms_binder_design.png)
+
+```bash
+python examples/binder_design_fastplms.py \
+  --target-name pd-l1 \
+  --binder-name minibinder \
+  --batch-size 4 \
+  --steps 150 \
+  --output-dir artifacts/binder-design
+```
+
+The workflow ranks candidates by mean iPTM across the approved critics after
+the minibinder isoelectric-point filter. These are model-based prioritization
+signals, not experimental evidence of affinity or specificity. See the
+[complete workflow](https://github.com/Synthyra/FastPLMs/blob/main/docs/binder_design.md).
+
+## Runtime contract
+
+- Public input: Raw amino-acid sequences or typed molecular-complex specifications; low-level forward accepts prepared feature tensors
+- Advertised AutoClasses: `AutoConfig`, `AutoModel`
+- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`
+- Attention implementations: `eager`, `sdpa`, `flex_attention`
+- Precision policies: `auto`, `fp32`, `bf16`, `fp8` (experimental)
+- 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/ESMFold2-Experimental-Cutoff2025@632ff4a9e68f1de78ee956a613267bdcdb5b354d`
+- Runtime revision: `1b9ce023f1e06571cf3e6324be0610ffa53e0a4a`
+- Runtime source-tree SHA-256: `15e781c5f1cd2ba8486e22076df15ffab37d3c00a689bf25280d803f2d60ee74`
+- Runtime bundle SHA-256: `278bb01ff0e426ae5f707c7a93ee720a0e87dfade5afb658921528d784720232`
+- Generator/schema version and complete/runtime-only attestations: recorded in `provenance.json`
+- Official checkpoint: `biohub/ESMFold2-Experimental-Cutoff2025@56f94f5c1069ecde17512c96928850518340d287`
+- Artifact source: `fast`
+- State transform: `identity`
+- BF16 execution: `fp32_parameters_autocast`
+- Pinned upstreams: `biohub-esm`, `biohub-transformers`, `protein-ttt`
+- Reference container: `reference-esmfold2`
+- 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

@@ -3,8 +3,8 @@
     "ESMFold2ExperimentalModel"
   ],
   "auto_map": {
-    "AutoConfig": "configuration_esmfold2.ESMFold2Config",
-    "AutoModel": "modeling_esmfold2_experimental.ESMFold2ExperimentalModel"
+    "AutoConfig": "modeling_fastplms.ESMFold2Config",
+    "AutoModel": "modeling_fastplms.ESMFold2ExperimentalModel"
   },
   "confidence_head": {
     "distogram_bins": 128,
@@ -25,6 +25,16 @@
   "disable_msa_features": false,
   "dtype": "float32",
   "esmc_id": "biohub/ESMC-6B",
+  "fastplms_checkpoint_hash": "3a40758a12594cab337bbe7e664d305df722dc1eb1f0cc1ba0c840bf46d22834",
+  "fastplms_checkpoint_repo_id": "Synthyra/ESMFold2-Experimental-Cutoff2025",
+  "fastplms_checkpoint_revision": "632ff4a9e68f1de78ee956a613267bdcdb5b354d",
+  "fastplms_model_id": "esmfold2_experimental_cutoff2025",
+  "fastplms_release_tool_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "fastplms_release_tool_sha256": "1459b5d7d13d9b07bd97b3eee764f2ce73623e15e32d07ddf6825c2a9509afb9",
+  "fastplms_runtime_bundle_sha256": "278bb01ff0e426ae5f707c7a93ee720a0e87dfade5afb658921528d784720232",
+  "fastplms_runtime_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "fastplms_source_tree_sha256": "15e781c5f1cd2ba8486e22076df15ffab37d3c00a689bf25280d803f2d60ee74",
+  "fastplms_weights_revision": "632ff4a9e68f1de78ee956a613267bdcdb5b354d",
   "folding_trunk": {
     "dropout": 0.25,
     "n_heads": 8,
@@ -56,6 +66,7 @@
   },
   "lm_num_layers": 80,
   "model_type": "esmfold2",
+  "msa_conditioning": true,
   "msa_encoder": {
     "d_hidden": 32,
     "d_msa": 128,

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_plusplus/modeling_esm_plusplus.py

@@ -0,0 +1,1552 @@
+"""Hugging Face-compatible ESMC models implemented by FastPLMs."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from functools import partial
+from typing import ClassVar
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from tokenizers import Tokenizer
+from tokenizers.models import BPE
+from tokenizers.processors import TemplateProcessing
+from transformers import PretrainedConfig, PreTrainedModel, PreTrainedTokenizerFast
+from transformers.modeling_outputs import (
+    MaskedLMOutput,
+    ModelOutput,
+    SequenceClassifierOutput,
+    TokenClassifierOutput,
+)
+
+try:
+    from fastplms.attention import (
+        AttentionBackend,
+        BlockMask,
+        FastPLMsAttentionMixin,
+        _get_flex_attention_fn,
+        _get_flex_block_mask,
+        flex_attention,
+        get_attention_mask,
+        kernels_flash_attention_func,
+        resolve_attention_backend,
+        resolve_attention_backend_for_call,
+    )
+    from fastplms.embeddings import EmbeddingMixin, Pooler, select_hidden_state_embeddings
+    from fastplms.models.ttt import FastPLMTestTimeTrainingMixin
+except ModuleNotFoundError as error:
+    _COMPOSITE_REQUIRED_NAMES = (
+        "AttentionBackend",
+        "BlockMask",
+        "EmbeddingMixin",
+        "FastPLMsAttentionMixin",
+        "FastPLMTestTimeTrainingMixin",
+        "Pooler",
+        "_get_flex_attention_fn",
+        "_get_flex_block_mask",
+        "flex_attention",
+        "get_attention_mask",
+        "kernels_flash_attention_func",
+        "resolve_attention_backend",
+        "resolve_attention_backend_for_call",
+        "select_hidden_state_embeddings",
+    )
+    if error.name != "fastplms" or any(
+        name not in globals() for name in _COMPOSITE_REQUIRED_NAMES
+    ):
+        raise
+    # Legacy flat Hub composites define every shared symbol above this block.
+
+
+class ESMplusplusConfig(PretrainedConfig):
+    """Configuration class for ESM++ model.
+
+    Args:
+        vocab_size: Size of the vocabulary
+        hidden_size: Dimension of hidden layers
+        num_attention_heads: Number of attention heads
+        num_hidden_layers: Number of transformer layers
+        num_labels: Number of output labels for classification
+        problem_type: Type of problem - regression, single/multi label classification
+    """
+
+    model_type = "ESMplusplus"
+
+    def __init__(
+        self,
+        vocab_size: int = 64,
+        hidden_size: int = 960,
+        num_attention_heads: int = 15,
+        num_hidden_layers: int = 30,
+        num_labels: int | None = None,
+        problem_type: str | None = None,
+        dropout: float = 0.0,
+        initializer_range: float = 0.02,
+        classifier_dropout: float = 0.1,
+        classifier_pooling_types: list[str] | None = None,
+        attn_backend: str | None = None,
+        pad_token_id: int = 1,
+        mask_token_id: int = 32,
+        **kwargs,
+    ):
+        if num_labels is None:
+            configured_labels = kwargs.get("id2label")
+            num_labels = len(configured_labels) if configured_labels else 2
+        super().__init__(
+            pad_token_id=pad_token_id,
+            mask_token_id=mask_token_id,
+            num_labels=num_labels,
+            **kwargs,
+        )
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_attention_heads = num_attention_heads
+        self.num_hidden_layers = num_hidden_layers
+        self.problem_type = problem_type
+        self.dropout = dropout
+        self.initializer_range = initializer_range
+        self.classifier_dropout = classifier_dropout
+        self.classifier_pooling_types = (
+            list(classifier_pooling_types) if classifier_pooling_types is not None else None
+        )
+        self.tie_word_embeddings = False
+        self.attn_backend = attn_backend
+
+
+### Rotary Embeddings
+def rotate_half(x: torch.Tensor, interleaved: bool = False) -> torch.Tensor:
+    """Rotate the final axis of X by 90 degrees in each two-dimensional plane."""
+    if interleaved:
+        paired = x.unflatten(-1, (-1, 2))
+        return torch.stack((-paired[..., 1], paired[..., 0]), dim=-1).flatten(-2)
+
+    # torch.chunk assigns an odd remainder to the first half. Express the same
+    # public behavior explicitly while keeping the ESMC path branch-free.
+    midpoint = (x.shape[-1] + 1) // 2
+    return torch.cat((-x[..., midpoint:], x[..., :midpoint]), dim=-1)
+
+
+def apply_rotary_emb_torch(
+    x: torch.Tensor,
+    cos: torch.Tensor,
+    sin: torch.Tensor,
+    interleaved: bool = False,
+    _inplace: bool = False,
+) -> torch.Tensor:
+    """Apply cached rotary angles to X while preserving any unrotated features."""
+    del _inplace  # Kept in the signature for checkpoint remote-code compatibility.
+    rotary_width = 2 * cos.shape[-1]
+    if rotary_width > x.shape[-1]:
+        raise AssertionError("rotary width exceeds the attention head dimension")
+
+    token_count = x.shape[1]
+    cos_full = torch.cat((cos[:token_count], cos[:token_count]), dim=-1).unsqueeze(1)
+    sin_full = torch.cat((sin[:token_count], sin[:token_count]), dim=-1).unsqueeze(1)
+    x_rotary = x[..., :rotary_width]
+    y_rotary = x_rotary * cos_full + rotate_half(x_rotary, interleaved) * sin_full
+    if rotary_width == x.shape[-1]:
+        return y_rotary
+    return torch.cat((y_rotary, x[..., rotary_width:]), dim=-1)
+
+
+class RotaryEmbedding(torch.nn.Module):
+    """Rotary position embeddings.
+
+    Based on the paper "RoFormer: Enhanced Transformer with Rotary Position Embedding"
+
+    Args:
+        dim: Dimension of the embedding
+        base: Base for computing angular frequencies
+        interleaved: Whether to use interleaved rotations
+        scale_base: Base for scaling
+        scaling_factor: Factor for scaling positions
+        pos_idx_in_fp32: Whether to compute position indices in fp32
+        device: Computation device
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        base: float = 10000.0,
+        interleaved: bool = False,
+        scale_base: float | None = None,
+        scaling_factor: float = 1.0,
+        pos_idx_in_fp32: bool = True,
+        device: torch.device | None = None,
+    ):
+        super().__init__()
+        self.dim, self.base = dim, float(base)
+        self.interleaved, self.scale_base = interleaved, scale_base
+        self.scaling_factor, self.pos_idx_in_fp32 = scaling_factor, pos_idx_in_fp32
+        self.device = device
+        self._clear_cache()
+        self.reset_parameters()
+
+    def _clear_cache(self) -> None:
+        self._seq_len_cached = 0
+        self._cos_cached: torch.Tensor | None = None
+        self._sin_cached: torch.Tensor | None = None
+        self._cos_k_cached: torch.Tensor | None = None
+        self._sin_k_cached: torch.Tensor | None = None
+
+    def reset_parameters(self, device: torch.device | str | None = None):
+        """Rebuild the non-persistent frequency buffers on ``device``."""
+        if device is not None:
+            buffer_device = torch.device(device)
+        elif "inv_freq" in self._buffers and isinstance(self._buffers["inv_freq"], torch.Tensor):
+            buffer_device = self._buffers["inv_freq"].device
+        else:
+            buffer_device = self.device
+        inv_freq = self._compute_inv_freq(buffer_device)
+        self._clear_cache()
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        arange = torch.arange(0, self.dim, 2, device=buffer_device, dtype=torch.float32)
+        scale = (
+            (arange + 0.4 * self.dim) / (1.4 * self.dim) if self.scale_base is not None else None
+        )
+        self.register_buffer("scale", scale)
+
+    def _compute_inv_freq(self, device: torch.device | None = None) -> torch.Tensor:
+        """Compute inverse frequency bands on their execution device."""
+        return 1 / (
+            self.base
+            ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim)
+        )
+
+    def _apply(self, fn, recurse: bool = True):
+        """Move the module, then regenerate device-specific RoPE frequencies."""
+        if self.inv_freq.is_meta:
+            self.reset_parameters(device="cpu")
+        result = super()._apply(fn, recurse=recurse)
+        self.register_buffer(
+            "inv_freq",
+            self._compute_inv_freq(self.inv_freq.device),
+            persistent=False,
+        )
+        self._clear_cache()
+        return result
+
+    def _cache_is_current(
+        self,
+        token_count: int,
+        device: torch.device | None,
+        dtype: torch.dtype | None,
+    ) -> bool:
+        cached = self._cos_cached
+        return (
+            cached is not None
+            and self._seq_len_cached >= token_count
+            and cached.device == device
+            and cached.dtype == dtype
+            and not (self.training and cached.is_inference())
+        )
+
+    def _rotary_angles(
+        self,
+        token_count: int,
+        device: torch.device | None,
+    ) -> torch.Tensor:
+        position_dtype = torch.float32 if self.pos_idx_in_fp32 else self.inv_freq.dtype
+        positions = torch.arange(token_count, device=device, dtype=position_dtype)
+        positions.div_(self.scaling_factor)
+        frequencies = (
+            self.inv_freq.to(torch.float32)
+            if self.pos_idx_in_fp32 and self.inv_freq.dtype != torch.float32
+            else self.inv_freq
+        )
+        return torch.outer(positions, frequencies)
+
+    def _update_cos_sin_cache(
+        self, seqlen: int, device: torch.device | None = None, dtype: torch.dtype | None = None
+    ) -> None:
+        """Build angle tables when the requested cache identity has changed."""
+        if self._cache_is_current(seqlen, device, dtype):
+            return
+
+        self._seq_len_cached = seqlen
+        angles = self._rotary_angles(seqlen, device)
+        cos_angles = torch.cos(angles)
+        sin_angles = torch.sin(angles)
+        if self.scale is None:
+            self._cos_cached = cos_angles.to(dtype)
+            self._sin_cached = sin_angles.to(dtype)
+            return
+
+        centered_positions = (
+            torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - seqlen // 2
+        ) / self.scale_base
+        scale = self.scale ** centered_positions.unsqueeze(-1)
+        self._cos_cached = (cos_angles * scale).to(dtype)
+        self._sin_cached = (sin_angles * scale).to(dtype)
+        self._cos_k_cached = (cos_angles / scale).to(dtype)
+        self._sin_k_cached = (sin_angles / scale).to(dtype)
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        """Apply rotary embeddings to queries and keys.
+
+        Args:
+            q: Query tensor Q with shape (b, l, h, d).
+            k: Key tensor K with shape (b, l, h, d).
+
+        Returns:
+            Tuple of rotated query and key tensors
+        """
+        # The pinned Biohub Transformers oracle recomputes inverse frequencies
+        # on the execution device. CPU and CUDA differ by about one FP32 ULP in
+        # some bands, which is immaterial in BF16 but accumulates measurably in
+        # deep FP32 execution.
+        self._update_cos_sin_cache(q.shape[1], device=q.device, dtype=q.dtype)
+        if self._cos_cached is None or self._sin_cached is None:
+            raise RuntimeError(
+                "Rotary cache initialization did not produce cosine and sine values."
+            )
+        if self.scale is not None:
+            raise AssertionError("Scaled rotary embeddings are unsupported for ESMC.")
+
+        cos_angles = self._cos_cached
+        sin_angles = self._sin_cached
+        return (
+            apply_rotary_emb_torch(q, cos_angles, sin_angles, self.interleaved, True),
+            apply_rotary_emb_torch(k, cos_angles, sin_angles, self.interleaved, True),
+        )
+
+
+### Feedforward Network Components
+def swiglu_correction_fn(expansion_ratio: float, d_model: int) -> int:
+    """Compute corrected dimension for SwiGLU."""
+    return int(((expansion_ratio * d_model) + 255) // 256 * 256)
+
+
+class SwiGLU(nn.Module):
+    """SwiGLU activation function."""
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x1, x2 = x.chunk(2, dim=-1)
+        return F.silu(x1) * x2
+
+
+def swiglu_ln_ffn(d_model: int, expansion_ratio: float) -> nn.Sequential:
+    """Create SwiGLU feedforward network with layer normalization."""
+    return nn.Sequential(
+        nn.LayerNorm(d_model),
+        nn.Linear(d_model, swiglu_correction_fn(expansion_ratio, d_model) * 2, bias=False),
+        SwiGLU(),
+        nn.Linear(swiglu_correction_fn(expansion_ratio, d_model), d_model, bias=False),
+    )
+
+
+### Attention
+class MultiHeadAttention(nn.Module):
+    """Multi-head attention with rotary embeddings and configurable backend.
+
+    Args:
+        d_model: Model dimension
+        n_heads: Number of attention heads
+        attn_backend: One of "eager", "sdpa", or "flex_attention".
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        n_heads: int,
+        attn_backend: str = "sdpa",
+    ):
+        super().__init__()
+        self.d_model = d_model
+        self.n_heads = n_heads
+        self.d_head = self.d_model // self.n_heads
+        self.scale = 1.0 / math.sqrt(self.d_head)
+        self.attn_backend = resolve_attention_backend(attn_backend)
+        self.layernorm_qkv = nn.Sequential(
+            nn.LayerNorm(d_model), nn.Linear(d_model, d_model * 3, bias=False)
+        )
+        self.out_proj = nn.Linear(d_model, d_model, bias=False)
+        self.q_ln = nn.LayerNorm(d_model, bias=False)
+        self.k_ln = nn.LayerNorm(d_model, bias=False)
+        self.reshaper = partial(rearrange, pattern="b s (h d) -> b h s d", h=n_heads)
+        self.rotary = RotaryEmbedding(d_model // n_heads)
+
+    def _apply_rotary(self, q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        q = q.unflatten(-1, (self.n_heads, self.d_head))
+        k = k.unflatten(-1, (self.n_heads, self.d_head))
+        q, k = self.rotary(q, k)
+        q = q.flatten(-2, -1)
+        k = k.flatten(-2, -1)
+        return q, k
+
+    def forward(
+        self,
+        x: 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,
+        output_s_max: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]:
+        qkv = self.layernorm_qkv(x)
+        query_sequence, key_sequence, value_sequence = torch.chunk(qkv, 3, dim=-1)
+        query_sequence, key_sequence = (
+            self.q_ln(query_sequence).to(query_sequence.dtype),
+            self.k_ln(key_sequence).to(query_sequence.dtype),
+        )
+        query_sequence, key_sequence = self._apply_rotary(query_sequence, key_sequence)
+        query_heads, key_heads, value_heads = map(
+            self.reshaper, (query_sequence, key_sequence, value_sequence)
+        )
+
+        attn_output, attn_weights, s_max = 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,
+            output_s_max=output_s_max,
+        )
+
+        output = self.out_proj(attn_output)
+        return output, attn_weights, s_max
+
+    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,
+        output_s_max: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]:
+        if output_attentions:
+            return self._manual_attn(
+                query_heads, key_heads, value_heads, attention_mask_4d, output_s_max
+            )
+
+        if self.attn_backend == AttentionBackend.EAGER:
+            attn_output, _, s_max = self._manual_attn(
+                query_heads, key_heads, value_heads, attention_mask_4d, output_s_max
+            )
+            return attn_output, None, s_max
+        if self.attn_backend.is_flash:
+            attn_output, attn_weights = self._kernels_flash_attn(
+                query_heads, key_heads, value_heads, attention_mask_2d
+            )
+        elif self.attn_backend == AttentionBackend.FLEX:
+            attn_output, attn_weights = self._flex_attn(
+                query_heads,
+                key_heads,
+                value_heads,
+                flex_block_mask,
+                attention_mask_2d,
+            )
+        elif self.attn_backend == AttentionBackend.SDPA:
+            attn_output, attn_weights = self._sdpa_attn(
+                query_heads, key_heads, value_heads, attention_mask_4d
+            )
+        else:
+            raise AssertionError(f"Unsupported resolved backend: {self.attn_backend}")
+
+        s_max = self._compute_s_max(query_heads, key_heads) if output_s_max else None
+        return attn_output, attn_weights, s_max
+
+    @torch.no_grad()
+    def _compute_s_max(
+        self, query_heads: torch.Tensor, key_heads: torch.Tensor
+    ) -> list[torch.Tensor]:
+        q_norm = torch.linalg.vector_norm(query_heads, dim=-1)
+        k_norm = torch.linalg.vector_norm(key_heads, dim=-1)
+        s_max_bound = (q_norm.max(dim=-1).values * k_norm.max(dim=-1).values).max(
+            dim=0
+        ).values * self.scale
+        return [s_max_bound[h] for h in range(self.n_heads)]
+
+    def _manual_attn(
+        self,
+        query_heads: torch.Tensor,
+        key_heads: torch.Tensor,
+        value_heads: torch.Tensor,
+        attention_mask_4d: torch.Tensor | None = None,
+        output_s_max: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor] | None]:
+        attn_weights = torch.matmul(query_heads, key_heads.transpose(-2, -1)) * self.scale
+        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)
+        context_heads = torch.matmul(attn_weights, value_heads)
+        attn_output = rearrange(context_heads, "b h s d -> b s (h d)")
+        s_max = self._compute_s_max(query_heads, key_heads) if output_s_max else None
+        return attn_output, attn_weights, s_max
+
+    def _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()
+        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,
+            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=self.scale,
+            kernel_options={"PRESCALE_QK": True, "BLOCK_N": 32},
+        )
+        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,
+            scale=self.scale,
+        )
+        return rearrange(context_heads, "b h s d -> b s (h d)"), None
+
+
+### Regression Head
+def RegressionHead(d_model: int, output_dim: int, hidden_dim: int | None = None) -> nn.Module:
+    """Create a regression head with optional hidden dimension.
+
+    Args:
+        d_model: Input dimension
+        output_dim: Output dimension
+        hidden_dim: Optional hidden dimension (defaults to d_model)
+    """
+    hidden_dim = hidden_dim if hidden_dim is not None else d_model
+    return nn.Sequential(
+        nn.Linear(d_model, hidden_dim),
+        nn.GELU(),
+        nn.LayerNorm(hidden_dim),
+        nn.Linear(hidden_dim, output_dim),
+    )
+
+
+### Transformer Block
+class UnifiedTransformerBlock(nn.Module):
+    """Transformer block with attention and feedforward layers."""
+
+    def __init__(
+        self,
+        d_model: int,
+        n_heads: int,
+        residue_scaling_factor: float = 1,
+        expansion_ratio: float = 8 / 3,
+        dropout: float = 0.0,
+        attn_backend: str = "sdpa",
+    ):
+        super().__init__()
+        self.attn = MultiHeadAttention(d_model=d_model, n_heads=n_heads, attn_backend=attn_backend)
+        self.ffn = swiglu_ln_ffn(d_model, expansion_ratio)
+        self.scaling_factor = residue_scaling_factor
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(
+        self,
+        x: 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,
+        output_s_max: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]:
+        attn_output, attn_weights, s_max = self.attn(
+            x,
+            attention_mask_2d=attention_mask_2d,
+            attention_mask_4d=attention_mask_4d,
+            flex_block_mask=flex_block_mask,
+            output_attentions=output_attentions,
+            output_s_max=output_s_max,
+        )
+        x = x + self.dropout(attn_output) / self.scaling_factor
+        x = x + self.dropout(self.ffn(x)) / self.scaling_factor
+        return x, attn_weights, s_max
+
+
+### Model Outputs
+@dataclass
+class TransformerOutput(ModelOutput):
+    """Output type for transformer encoder."""
+
+    last_hidden_state: torch.Tensor | None = None
+    hidden_states: tuple[torch.Tensor] | None = None
+    attentions: tuple[torch.Tensor] | None = None
+    s_max: tuple[list[torch.Tensor], ...] | None = None
+
+
+@dataclass
+class ESMplusplusOutput(MaskedLMOutput):
+    """Masked-LM output with FastPLMs fields after the HF contract."""
+
+    s_max: tuple[list[torch.Tensor], ...] | None = None
+    last_hidden_state: torch.Tensor | None = None
+
+
+@dataclass
+class ESMplusplusSequenceClassifierOutput(SequenceClassifierOutput):
+    """Sequence-classification output with optional attention diagnostics."""
+
+    s_max: tuple[list[torch.Tensor], ...] | None = None
+
+
+@dataclass
+class ESMplusplusTokenClassifierOutput(TokenClassifierOutput):
+    """Token-classification output with optional attention diagnostics."""
+
+    s_max: tuple[list[torch.Tensor], ...] | None = None
+
+
+### Transformer Stack
+class TransformerStack(nn.Module):
+    """Stack of transformer blocks."""
+
+    def __init__(
+        self,
+        d_model: int,
+        n_heads: int,
+        n_layers: int,
+        dropout: float = 0.0,
+        attn_backend: str = "sdpa",
+    ):
+        super().__init__()
+        self.attention_backend = resolve_attention_backend(attn_backend)
+        self.blocks = nn.ModuleList(
+            [
+                UnifiedTransformerBlock(
+                    d_model,
+                    n_heads,
+                    residue_scaling_factor=math.sqrt(n_layers / 36),
+                    dropout=dropout,
+                    attn_backend=attn_backend,
+                )
+                for i in range(n_layers)
+            ]
+        )
+        self.norm = nn.LayerNorm(d_model, bias=False)
+        self.gradient_checkpointing = False
+
+    @property
+    def attn_backend(self) -> AttentionBackend:
+        return self.attention_backend
+
+    @attn_backend.setter
+    def attn_backend(self, backend: str) -> None:
+        resolved = resolve_attention_backend(backend)
+        self.attention_backend = resolved
+        for block in self.blocks:
+            block.attn.attn_backend = resolved
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        sequence_id: torch.Tensor | None = None,
+        output_hidden_states: bool | None = False,
+        output_attentions: bool | None = False,
+        output_s_max: bool | None = False,
+        esmfold2_hidden_states: bool = False,
+    ) -> TransformerOutput:
+        hidden_states = () if output_hidden_states else None
+        attentions = () if output_attentions else None
+        full_s_max = () if output_s_max else None
+        # Match the pinned Biohub Transformers contract: a supplied sequence_id
+        # is authoritative and must encode padding as -1.  attention_mask is
+        # ignored in that mode rather than intersected with the chain mask.
+        if sequence_id is None and attention_mask is not None:
+            expected_shape = (x.shape[0], x.shape[1])
+            if attention_mask.ndim != 2 or tuple(attention_mask.shape) != expected_shape:
+                raise ValueError(
+                    f"attention_mask must have shape {expected_shape}; "
+                    f"received {tuple(attention_mask.shape)}."
+                )
+            attention_mask = attention_mask.to(device=x.device, dtype=torch.bool)
+            if not bool(attention_mask.any(dim=1).all()):
+                raise ValueError("attention_mask must keep at least one valid key per batch row.")
+        effective_backend = resolve_attention_backend_for_call(
+            self.attention_backend,
+            output_attentions=bool(output_attentions),
+        )
+
+        if sequence_id is None and attention_mask is not None:
+            attention_mask_2d, attention_mask_4d, flex_block_mask = (
+                self._sequence_id_attention_masks(
+                    sequence_id=attention_mask.to(device=x.device, dtype=torch.bool),
+                    batch_size=x.shape[0],
+                    seq_len=x.shape[1],
+                    device=x.device,
+                    dtype=x.dtype,
+                    effective_backend=effective_backend,
+                )
+            )
+        elif sequence_id is None:
+            attention_mask_2d, attention_mask_4d, flex_block_mask = get_attention_mask(
+                effective_backend=effective_backend,
+                batch_size=x.shape[0],
+                seq_len=x.shape[1],
+                device=x.device,
+                attention_mask=attention_mask,
+                dtype=x.dtype,
+                mask_semantics="padding",
+            )
+        else:
+            attention_mask_2d, attention_mask_4d, flex_block_mask = (
+                self._sequence_id_attention_masks(
+                    sequence_id=sequence_id,
+                    batch_size=x.shape[0],
+                    seq_len=x.shape[1],
+                    device=x.device,
+                    dtype=x.dtype,
+                    effective_backend=effective_backend,
+                )
+            )
+
+        for block in self.blocks:
+            if output_hidden_states:
+                if hidden_states is None:
+                    raise RuntimeError(
+                        "Hidden-state collection was not initialized for an enabled request."
+                    )
+                # Biohub Transformers records the input to each block followed
+                # by the final normalized state. This gives n_layers + 1 states
+                # and, for ESMC-6B, the 81-state order consumed by ESMFold2.
+                hidden_states += (x,)
+            if self.gradient_checkpointing and self.training:
+                x, attn_weights, s_max = self._gradient_checkpointing_func(
+                    block.__call__,
+                    x=x,
+                    attention_mask_2d=attention_mask_2d,
+                    attention_mask_4d=attention_mask_4d,
+                    flex_block_mask=flex_block_mask,
+                    output_attentions=output_attentions,
+                    output_s_max=output_s_max,
+                )
+            else:
+                x, attn_weights, s_max = block(
+                    x=x,
+                    attention_mask_2d=attention_mask_2d,
+                    attention_mask_4d=attention_mask_4d,
+                    flex_block_mask=flex_block_mask,
+                    output_attentions=output_attentions,
+                    output_s_max=output_s_max,
+                )
+
+            if attentions is not None:
+                attentions += (attn_weights,)
+            if full_s_max is not None:
+                full_s_max += (s_max,)
+
+        last_hidden_state = self.norm(x)
+        if output_hidden_states:
+            hidden_states += (last_hidden_state,)
+
+        return TransformerOutput(
+            last_hidden_state=last_hidden_state,
+            hidden_states=hidden_states,
+            attentions=attentions,
+            s_max=full_s_max,
+        )
+
+    def _sequence_id_attention_masks(
+        self,
+        sequence_id: torch.Tensor,
+        batch_size: int,
+        seq_len: int,
+        device: torch.device,
+        dtype: torch.dtype | None = None,
+        effective_backend: AttentionBackend | None = None,
+    ) -> tuple[torch.Tensor | None, torch.Tensor | None, BlockMask | None]:
+        expected_shape = (batch_size, seq_len)
+        if sequence_id.ndim != 2 or tuple(sequence_id.shape) != expected_shape:
+            raise ValueError(
+                f"sequence_id must have shape {expected_shape}; "
+                f"received {tuple(sequence_id.shape)}."
+            )
+        if sequence_id.device != device:
+            sequence_id = sequence_id.to(device=device)
+        backend = (
+            self.attention_backend
+            if effective_backend is None
+            else resolve_attention_backend(effective_backend)
+        )
+        if sequence_id.dtype == torch.bool:
+            attention_mask_2d = sequence_id
+            # Biohub's boolean single-chain form groups biological positions
+            # together and padding positions together. Padding queries remain
+            # finite without allowing their states to enter residue attention.
+            attention_mask_4d = sequence_id[:, None, :, None] == sequence_id[:, None, None, :]
+        else:
+            attention_mask_2d = sequence_id != -1
+            attention_mask_4d = (sequence_id.unsqueeze(-1) == sequence_id.unsqueeze(-2)).unsqueeze(
+                1
+            )
+        if not bool(attention_mask_2d.any(dim=1).all()):
+            raise ValueError("attention_mask must keep at least one valid key per batch row.")
+
+        if backend.is_flash:
+            if sequence_id.dtype != torch.bool:
+                raise ValueError(
+                    "ESM++ FlashAttention only supports boolean sequence_id padding masks. "
+                    "Use eager, sdpa, or flex_attention for chain-aware integer sequence_id "
+                    "masks."
+                )
+            return attention_mask_2d, attention_mask_4d, None
+
+        if backend == AttentionBackend.FLEX:
+            if sequence_id.dtype == torch.bool:
+
+                def mask_mod(batch_idx, head_idx, q_idx, kv_idx):
+                    del head_idx
+                    return sequence_id[batch_idx, q_idx] == sequence_id[batch_idx, kv_idx]
+
+            else:
+
+                def mask_mod(batch_idx, head_idx, q_idx, kv_idx):
+                    del head_idx
+                    q_id = sequence_id[batch_idx, q_idx]
+                    kv_id = sequence_id[batch_idx, kv_idx]
+                    return q_id == kv_id
+
+            flex_block_mask = _get_flex_block_mask(
+                mask_pattern=sequence_id,
+                batch_size=batch_size,
+                query_length=seq_len,
+                key_value_length=seq_len,
+                device=device,
+                dtype=dtype,
+                mask_semantics=(
+                    "boolean_sequence_id"
+                    if sequence_id.dtype == torch.bool
+                    else "integer_sequence_id"
+                ),
+                mask_mod=mask_mod,
+            )
+            return attention_mask_2d, attention_mask_4d, flex_block_mask
+
+        return attention_mask_2d, attention_mask_4d, None
+
+
+class PreTrainedESMplusplusModel(FastPLMsAttentionMixin, PreTrainedModel):
+    """
+    init weights for ESM++ models
+    """
+
+    config_class = ESMplusplusConfig
+    base_model_prefix = "esm++"
+    supports_gradient_checkpointing = True
+    all_tied_weights_keys: ClassVar[dict[str, str]] = {}
+    _supports_flash_attn = True
+    _supports_flash_attn_2 = True
+    _supports_flash_attn_3 = True
+    _fastplms_attention_implementations = (
+        "eager",
+        "sdpa",
+        "flex_attention",
+        "flash_attention_2",
+        "flash_attention_3",
+    )
+
+    @property
+    def tokenizer(self) -> EsmSequenceTokenizer:
+        """Construct the sequence tokenizer only when a raw-sequence API needs it."""
+
+        tokenizer = self.__dict__.get("_fastplms_tokenizer")
+        if tokenizer is None:
+            tokenizer = EsmSequenceTokenizer()
+            self.__dict__["_fastplms_tokenizer"] = tokenizer
+        return tokenizer
+
+    @tokenizer.setter
+    def tokenizer(self, value: EsmSequenceTokenizer | None) -> None:
+        self.__dict__["_fastplms_tokenizer"] = value
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        # HF from_pretrained marks loaded parameters with `_is_hf_initialized`.
+        # Skip this module if any local parameter is already marked as loaded.
+        for parameter in module.parameters(recurse=False):
+            if parameter.__dict__.get("_is_hf_initialized"):
+                return
+
+        if isinstance(module, nn.Linear):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+        elif isinstance(module, nn.Embedding):
+            nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                with torch.no_grad():
+                    module.weight[module.padding_idx].zero_()
+        elif isinstance(module, nn.LayerNorm):
+            if module.bias is not None:
+                nn.init.zeros_(module.bias)
+            nn.init.ones_(module.weight)
+
+    @property
+    def attn_backend(self) -> str:
+        return self.config.attn_backend
+
+    @attn_backend.setter
+    def attn_backend(self, backend: str) -> None:
+        if backend not in self._fastplms_attention_implementations:
+            raise ValueError(
+                f"{type(self).__name__} does not support {backend!r}; expected one of "
+                f"{self._fastplms_attention_implementations}."
+            )
+        self.set_attn_implementation(backend)
+
+    def _reset_rotary_embeddings(self):
+        """Refresh non-persistent rotary buffers after checkpoint loading."""
+        for module in self.modules():
+            if isinstance(module, RotaryEmbedding):
+                module.reset_parameters()
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        output_loading_info = (
+            bool(kwargs["output_loading_info"]) if "output_loading_info" in kwargs else False
+        )
+        loaded = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+        if output_loading_info:
+            model, loading_info = loaded
+            model._reset_rotary_embeddings()
+            return model, loading_info
+        loaded._reset_rotary_embeddings()
+        return loaded
+
+
+### ESM++ Models
+class ESMplusplusModel(PreTrainedESMplusplusModel, EmbeddingMixin):
+    """
+    ESM++ transformer backbone.
+
+    Official ESM++ checkpoints contain the sequence head even when loaded through
+    ``AutoModel``.  Keep that module in the base class so the checkpoint has one
+    exact state-dict contract across ``AutoModel`` and ``AutoModelForMaskedLM``;
+    the base forward path intentionally does not compute or return logits.
+    """
+
+    config_class = ESMplusplusConfig
+
+    def __init__(self, config: ESMplusplusConfig, **kwargs):
+        PreTrainedESMplusplusModel.__init__(self, config, **kwargs)
+        self.config = config
+        self.vocab_size = config.vocab_size
+        self.embed = nn.Embedding(self.vocab_size, config.hidden_size)
+        self.transformer = TransformerStack(
+            d_model=config.hidden_size,
+            n_heads=config.num_attention_heads,
+            n_layers=config.num_hidden_layers,
+            dropout=config.dropout,
+            attn_backend=config.attn_backend,
+        )
+        self.sequence_head = RegressionHead(config.hidden_size, self.vocab_size)
+        self.init_weights()
+
+    def get_input_embeddings(self):
+        return self.embed
+
+    def set_input_embeddings(self, value):
+        self.embed = value
+
+    def get_output_embeddings(self):
+        return self.sequence_head[-1]
+
+    def set_output_embeddings(self, new_embeddings):
+        self.sequence_head[-1] = new_embeddings
+
+    def _embed(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        hidden_state_index: int = -1,
+        store_all_hidden_states: bool = False,
+    ) -> torch.Tensor:
+        if attention_mask is None:
+            attention_mask = input_ids.ne(self.config.pad_token_id)
+        x = self.embed(input_ids)
+        output_hidden_states = store_all_hidden_states or hidden_state_index != -1
+        output = self.transformer(
+            x=x,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=False,
+        )
+        return select_hidden_state_embeddings(
+            output.last_hidden_state,
+            output.hidden_states,
+            hidden_state_index=hidden_state_index,
+            store_all_hidden_states=store_all_hidden_states,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        sequence_id: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        output_s_max: bool | None = False,
+        esmfold2_hidden_states: bool = False,
+        return_dict: bool | None = None,
+    ) -> TransformerOutput | tuple[torch.Tensor, ...]:
+        """Run ESMC inference with the pinned Biohub mask precedence.
+
+        ``sequence_id`` is authoritative when supplied: non-negative integers
+        identify chains and ``-1`` identifies padding.  In that mode
+        ``attention_mask`` is ignored, matching the official implementation.
+        Without ``sequence_id``, ``attention_mask`` is the ordinary padding
+        mask and defaults to ``input_ids != pad_token_id``.
+        """
+        if input_ids is None and inputs_embeds is None:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        output_attentions = (
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if attention_mask is None and sequence_id is None and input_ids is not None:
+            attention_mask = input_ids.ne(self.config.pad_token_id)
+
+        x = self.embed(input_ids) if inputs_embeds is None else inputs_embeds
+
+        transformer_output = self.transformer(
+            x=x,
+            attention_mask=attention_mask,
+            sequence_id=sequence_id,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+            output_s_max=output_s_max,
+            esmfold2_hidden_states=esmfold2_hidden_states,
+        )
+        result = TransformerOutput(
+            last_hidden_state=transformer_output.last_hidden_state,
+            hidden_states=transformer_output.hidden_states,
+            attentions=transformer_output.attentions,
+            s_max=transformer_output.s_max,
+        )
+        return result if return_dict else result.to_tuple()
+
+
+class ESMplusplusForMaskedLM(
+    FastPLMTestTimeTrainingMixin, PreTrainedESMplusplusModel, EmbeddingMixin
+):
+    """
+    ESM++ model for masked language modeling.
+    Implements the base ESM++ architecture with a masked language modeling head.
+    """
+
+    config_class = ESMplusplusConfig
+
+    def __init__(self, config: ESMplusplusConfig, **kwargs):
+        PreTrainedESMplusplusModel.__init__(self, config, **kwargs)
+        self.config = config
+        self.vocab_size = config.vocab_size
+        self.embed = nn.Embedding(self.vocab_size, config.hidden_size)
+        self.transformer = TransformerStack(
+            d_model=config.hidden_size,
+            n_heads=config.num_attention_heads,
+            n_layers=config.num_hidden_layers,
+            dropout=config.dropout,
+            attn_backend=config.attn_backend,
+        )
+        self.sequence_head = RegressionHead(config.hidden_size, self.vocab_size)
+        self.ce_loss = nn.CrossEntropyLoss()
+        self.init_weights()
+        self.init_ttt({"lora_target_replace_module": "MultiHeadAttention"})
+
+    def get_input_embeddings(self):
+        return self.embed
+
+    def set_input_embeddings(self, value):
+        self.embed = value
+
+    def get_output_embeddings(self):
+        return self.sequence_head[-1]
+
+    def set_output_embeddings(self, new_embeddings):
+        self.sequence_head[-1] = new_embeddings
+
+    def _embed(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        hidden_state_index: int = -1,
+        store_all_hidden_states: bool = False,
+    ) -> torch.Tensor:
+        if attention_mask is None:
+            attention_mask = input_ids.ne(self.config.pad_token_id)
+        x = self.embed(input_ids)
+        output_hidden_states = store_all_hidden_states or hidden_state_index != -1
+        output = self.transformer(
+            x=x,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=False,
+        )
+        return select_hidden_state_embeddings(
+            output.last_hidden_state,
+            output.hidden_states,
+            hidden_state_index=hidden_state_index,
+            store_all_hidden_states=store_all_hidden_states,
+        )
+
+    def _ttt_get_trainable_modules(self) -> list[nn.Module]:
+        return [self.transformer]
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        sequence_id: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        output_s_max: bool | None = False,
+        esmfold2_hidden_states: bool = False,
+        return_dict: bool | None = None,
+        compute_logits: bool = True,
+    ) -> ESMplusplusOutput | tuple[torch.Tensor, ...]:
+        if input_ids is None and inputs_embeds is None:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        if labels is not None and not compute_logits:
+            raise ValueError("labels require compute_logits=True.")
+        output_attentions = (
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        if attention_mask is None and sequence_id is None and input_ids is not None:
+            attention_mask = input_ids.ne(self.config.pad_token_id)
+
+        x = self.embed(input_ids) if inputs_embeds is None else inputs_embeds
+
+        output = self.transformer(
+            x=x,
+            attention_mask=attention_mask,
+            sequence_id=sequence_id,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+            output_s_max=output_s_max,
+            esmfold2_hidden_states=esmfold2_hidden_states,
+        )
+
+        last_hidden_state = output.last_hidden_state
+        logits = self.sequence_head(last_hidden_state) if compute_logits else None
+        loss = None
+        if labels is not None:
+            if logits is None:
+                raise ValueError("labels require compute_logits=True.")
+            labels = labels.to(logits.device)
+            loss = self.ce_loss(logits.view(-1, self.vocab_size), labels.view(-1))
+
+        result = ESMplusplusOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=output.hidden_states,
+            attentions=output.attentions,
+            s_max=output.s_max,
+            last_hidden_state=last_hidden_state,
+        )
+        return result if return_dict else result.to_tuple()
+
+
+class ESMplusplusForSequenceClassification(ESMplusplusForMaskedLM, EmbeddingMixin):
+    """
+    ESM++ model for sequence classification.
+    Extends the base ESM++ model with a classification head.
+    """
+
+    def __init__(self, config: ESMplusplusConfig, **kwargs):
+        pooling_types = kwargs.pop("pooling_types", None)
+        if pooling_types is None:
+            pooling_types = config.classifier_pooling_types or ["mean", "var"]
+        elif not isinstance(pooling_types, list):
+            raise TypeError("pooling_types must be a non-empty list of strings.")
+        elif not pooling_types:
+            raise ValueError("pooling_types must contain at least one pooling operation.")
+        elif not all(isinstance(pooling_type, str) for pooling_type in pooling_types):
+            raise TypeError("pooling_types must be a non-empty list of strings.")
+        if "parti" in pooling_types:
+            raise ValueError(
+                "pooling_types cannot contain 'parti' for sequence classification "
+                "because the classifier does not expose layer attentions to its pooler."
+            )
+        config.classifier_pooling_types = list(pooling_types)
+
+        ESMplusplusForMaskedLM.__init__(self, config, **kwargs)
+        self.config = config
+        self.num_labels = config.num_labels
+        self.classifier = RegressionHead(
+            config.hidden_size * len(pooling_types),
+            config.num_labels,
+            config.hidden_size * 4,
+        )
+        # Large intermediate projections help with sequence classification tasks (*4)
+        self.mse = nn.MSELoss()
+        self.ce = nn.CrossEntropyLoss()
+        self.bce = nn.BCEWithLogitsLoss()
+        self.pooler = Pooler(pooling_types)
+        self.init_weights()
+
+    def _embed(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        hidden_state_index: int = -1,
+        store_all_hidden_states: bool = False,
+    ) -> torch.Tensor:
+        x = self.embed(input_ids)
+        output_hidden_states = store_all_hidden_states or hidden_state_index != -1
+        output = self.transformer(
+            x=x,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=False,
+        )
+        return select_hidden_state_embeddings(
+            output.last_hidden_state,
+            output.hidden_states,
+            hidden_state_index=hidden_state_index,
+            store_all_hidden_states=store_all_hidden_states,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        sequence_id: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        output_s_max: bool | None = False,
+        return_dict: bool | None = None,
+    ) -> ESMplusplusSequenceClassifierOutput | tuple[torch.Tensor, ...]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        pooling_mask = attention_mask
+        if pooling_mask is None:
+            if sequence_id is not None:
+                pooling_mask = (
+                    sequence_id if sequence_id.dtype == torch.bool else sequence_id.ne(-1)
+                )
+            elif input_ids is not None:
+                pooling_mask = input_ids.ne(self.config.pad_token_id)
+            else:
+                if inputs_embeds is None:
+                    raise ValueError("You have to specify either input_ids or inputs_embeds")
+                pooling_mask = torch.ones(
+                    inputs_embeds.shape[:2],
+                    dtype=torch.bool,
+                    device=inputs_embeds.device,
+                )
+
+        output = super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            sequence_id=sequence_id,
+            inputs_embeds=inputs_embeds,
+            labels=None,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            output_s_max=output_s_max,
+            return_dict=True,
+            compute_logits=False,
+        )
+
+        last_hidden_state = output.last_hidden_state
+        features = self.pooler(last_hidden_state, pooling_mask)
+        logits = self.classifier(features)
+
+        loss = None
+        if labels is not None:
+            labels = labels.to(logits.device)
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (
+                    labels.dtype == torch.long or labels.dtype == torch.int
+                ):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                if self.num_labels == 1:
+                    loss = self.mse(logits.flatten(), labels.flatten())
+                else:
+                    loss = self.mse(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss = self.ce(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss = self.bce(logits, labels)
+
+        result = ESMplusplusSequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=output.hidden_states,
+            attentions=output.attentions,
+            s_max=output.s_max,
+        )
+        return result if return_dict else result.to_tuple()
+
+
+class ESMplusplusForTokenClassification(ESMplusplusForMaskedLM, EmbeddingMixin):
+    """
+    ESM++ model for token classification.
+    Extends the base ESM++ model with a token classification head.
+    """
+
+    def __init__(self, config: ESMplusplusConfig, **kwargs):
+        ESMplusplusForMaskedLM.__init__(self, config, **kwargs)
+        self.config = config
+        self.num_labels = config.num_labels
+        self.classifier = RegressionHead(
+            config.hidden_size, config.num_labels, config.hidden_size * 4
+        )
+        # Large intermediate projections help with sequence classification tasks (*4)
+        self.loss_fct = nn.CrossEntropyLoss()
+        self.init_weights()
+
+    def _embed(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        hidden_state_index: int = -1,
+        store_all_hidden_states: bool = False,
+    ) -> torch.Tensor:
+        x = self.embed(input_ids)
+        output_hidden_states = store_all_hidden_states or hidden_state_index != -1
+        output = self.transformer(
+            x,
+            attention_mask,
+            output_hidden_states=output_hidden_states,
+            output_attentions=False,
+        )
+        return select_hidden_state_embeddings(
+            output.last_hidden_state,
+            output.hidden_states,
+            hidden_state_index=hidden_state_index,
+            store_all_hidden_states=store_all_hidden_states,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        sequence_id: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        output_s_max: bool | None = False,
+        return_dict: bool | None = None,
+    ) -> ESMplusplusTokenClassifierOutput | tuple[torch.Tensor, ...]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        output = super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            sequence_id=sequence_id,
+            inputs_embeds=inputs_embeds,
+            labels=None,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            output_s_max=output_s_max,
+            return_dict=True,
+            compute_logits=False,
+        )
+
+        last_hidden_state = output.last_hidden_state
+        logits = self.classifier(last_hidden_state)
+        loss = None
+        if labels is not None:
+            labels = labels.to(logits.device)
+            loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        result = ESMplusplusTokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=output.hidden_states,
+            attentions=output.attentions,
+            s_max=output.s_max,
+        )
+        return result if return_dict else result.to_tuple()
+
+
+### Tokenization
+SEQUENCE_VOCAB = [
+    "<cls>",
+    "<pad>",
+    "<eos>",
+    "<unk>",
+    "L",
+    "A",
+    "G",
+    "V",
+    "S",
+    "E",
+    "R",
+    "T",
+    "I",
+    "D",
+    "P",
+    "K",
+    "Q",
+    "N",
+    "F",
+    "Y",
+    "M",
+    "H",
+    "W",
+    "C",
+    "X",
+    "B",
+    "U",
+    "Z",
+    "O",
+    ".",
+    "-",
+    "|",
+    "<mask>",
+]
+
+
+def _build_sequence_tokenizer_backend(
+    *,
+    unk_token: str,
+    cls_token: str,
+    pad_token: str,
+    mask_token: str,
+    eos_token: str,
+    chain_break_token: str,
+) -> Tokenizer:
+    """Build the fixed ESMC character vocabulary and boundary-token policy."""
+    vocabulary = dict(zip(SEQUENCE_VOCAB, range(len(SEQUENCE_VOCAB)), strict=True))
+    backend = Tokenizer(BPE(vocabulary, merges=[], unk_token=unk_token))
+    backend.add_special_tokens([cls_token, pad_token, mask_token, eos_token, chain_break_token])
+    backend.post_processor = TemplateProcessing(
+        single="<cls> $A <eos>",
+        pair="<cls>:0 $A:0 <eos>:0 $B:1 <eos>:1",
+        special_tokens=[
+            ("<cls>", backend.token_to_id("<cls>")),
+            ("<eos>", backend.token_to_id("<eos>")),
+        ],
+    )
+    return backend
+
+
+class EsmSequenceTokenizer(PreTrainedTokenizerFast):
+    model_input_names: ClassVar[list[str]] = ["input_ids", "attention_mask"]
+
+    def __init__(
+        self,
+        unk_token="<unk>",
+        cls_token="<cls>",
+        pad_token="<pad>",
+        mask_token="<mask>",
+        eos_token="<eos>",
+        chain_break_token="|",
+        **kwargs,
+    ):
+        backend = _build_sequence_tokenizer_backend(
+            unk_token=unk_token,
+            cls_token=cls_token,
+            pad_token=pad_token,
+            mask_token=mask_token,
+            eos_token=eos_token,
+            chain_break_token=chain_break_token,
+        )
+        self.cb_token = chain_break_token
+        super().__init__(
+            tokenizer_object=backend,
+            unk_token=unk_token,
+            cls_token=cls_token,
+            pad_token=pad_token,
+            mask_token=mask_token,
+            eos_token=eos_token,
+            additional_special_tokens=[chain_break_token],
+            **kwargs,
+        )
+
+    # These are a footgun, we never use the `bos` token anywhere so we're just overriding it here.
+    @property
+    def bos_token(self):
+        return self.cls_token
+
+    @property
+    def bos_token_id(self):
+        return self.cls_token_id
+
+    @property
+    def chain_break_token(self):
+        return self.cb_token
+
+    @property
+    def chain_break_token_id(self):
+        return self.convert_tokens_to_ids(self.chain_break_token)
+
+    @property
+    def all_token_ids(self):
+        return list(range(self.vocab_size))
+
+    @property
+    def special_token_ids(self):
+        return self.all_special_ids

fastplms/models/esmfold2/__init__.py

@@ -0,0 +1,39 @@
+"""ESMFold2 public classes, imported lazily to keep optional extras isolated."""
+
+from __future__ import annotations
+
+from importlib import import_module
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from .configuration_esmfold2 import ESMFold2Config as ESMFold2Config
+    from .modeling_esmfold2 import ESMFold2Model as ESMFold2Model
+    from .modeling_esmfold2 import ESMFold2Output as ESMFold2Output
+    from .modeling_esmfold2_experimental import (
+        ESMFold2ExperimentalModel as ESMFold2ExperimentalModel,
+    )
+    from .reproducibility import seed_context as seed_context
+
+_EXPORT_MODULES = {
+    "ESMFold2Config": ".configuration_esmfold2",
+    "ESMFold2ExperimentalModel": ".modeling_esmfold2_experimental",
+    "ESMFold2Model": ".modeling_esmfold2",
+    "ESMFold2Output": ".modeling_esmfold2",
+    "seed_context": ".reproducibility",
+}
+
+
+def __getattr__(name: str) -> Any:
+    module_name = _EXPORT_MODULES.get(name)
+    if module_name is None:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    value = getattr(import_module(module_name, __name__), name)
+    globals()[name] = value
+    return value
+
+
+def __dir__() -> list[str]:
+    return sorted(set(globals()) | set(_EXPORT_MODULES))
+
+
+__all__ = list(_EXPORT_MODULES)

fastplms/models/esmfold2/attention.py

@@ -0,0 +1,48 @@
+"""Transformers-compatible attention selection for ESMFold2's ESMC backbone."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+
+from ...attention import FastPLMsAttentionMixin, get_attn_implementation
+
+
+class ESMFold2AttentionMixin(FastPLMsAttentionMixin):
+    """Route the outer Transformers attention API into the loaded ESMC model."""
+
+    _supports_attention_backend = True
+    _supports_sdpa = True
+    _supports_flex_attn = True
+    _supports_flash_attn_2 = False
+    _supports_flash_attn_3 = False
+    _fastplms_attention_implementations = (
+        "eager",
+        "sdpa",
+        "flex_attention",
+    )
+
+    def __init__(self, config, *args, **kwargs) -> None:
+        super().__init__(config, *args, **kwargs)
+        config.esmc_attn_backend = get_attn_implementation(config)
+
+    def set_attn_implementation(
+        self,
+        attn_implementation: str | Mapping[str, str],
+        allow_all_kernels: bool = False,
+    ) -> None:
+        """Set one canonical backend on ESMFold2 and its loaded ESMC model."""
+
+        if allow_all_kernels:
+            raise ValueError(
+                "ESMFold2 accepts only its declared built-in attention backends; "
+                "external attention kernels are not supported."
+            )
+        super().set_attn_implementation(attn_implementation)
+        resolved = get_attn_implementation(self.config)
+        self.config.esmc_attn_backend = resolved
+        esmc = getattr(self, "_esmc", None)
+        if esmc is not None:
+            esmc.set_attn_implementation(resolved)
+
+
+__all__ = ["ESMFold2AttentionMixin"]

fastplms/models/esmfold2/configuration_esmfold2.py

@@ -0,0 +1,306 @@
+# Copyright 2026 Biohub. All rights reserved.
+#
+# 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.
+
+"""Configuration schema for release and experimental ESMFold2 checkpoints."""
+
+from __future__ import annotations
+
+from dataclasses import asdict, dataclass, field
+from typing import Any, TypeVar, cast
+
+from transformers.configuration_utils import PretrainedConfig
+
+_ESMC_ATTENTION_IMPLEMENTATIONS = frozenset({"eager", "flex_attention", "sdpa"})
+_ESMC_PRECISIONS = frozenset({"auto", "bf16", "fp32", "fp8"})
+
+
+def _esmc_backbone_checkpoint_ids() -> tuple[str, str]:
+    """Return the manifest-pinned official and FastPLMs ESMC repositories."""
+
+    from fastplms.registry import RegistryError, get_model_registry
+
+    registry = get_model_registry()
+    family = registry.families["esmfold2"]
+    if family.backbone_model is None:
+        raise RegistryError("families.esmfold2 must declare backbone_model.")
+    backbone = registry[family.backbone_model]
+    return backbone.official.repo_id, backbone.fast.repo_id
+
+
+def normalize_esmc_id(esmc_id: str) -> str:
+    """Resolve an official ESMC identifier to its FastPLMs checkpoint mirror."""
+
+    official_repo, fast_repo = _esmc_backbone_checkpoint_ids()
+    return fast_repo if esmc_id == official_repo else esmc_id
+
+
+def normalize_esmc_attention_implementation(
+    implementation: str | dict[str, str] | None,
+) -> str | None:
+    """Validate the ESMC backend and translate the historical ``flex`` name."""
+
+    if isinstance(implementation, dict):
+        if tuple(implementation) != ("",):
+            raise ValueError(
+                "ESMFold2 has one ESMC attention backbone; use a string or {'': implementation}."
+            )
+        implementation = implementation[""]
+    canonical = "flex_attention" if implementation == "flex" else implementation
+    if canonical is not None and canonical not in _ESMC_ATTENTION_IMPLEMENTATIONS:
+        expected = sorted(_ESMC_ATTENTION_IMPLEMENTATIONS)
+        raise ValueError(
+            f"Unsupported ESMFold2 attention implementation {canonical!r}; "
+            f"expected one of {expected}."
+        )
+    return canonical
+
+
+NestedConfig = TypeVar("NestedConfig")
+
+
+def _nested_config(value: Any, config_type: type[NestedConfig]) -> NestedConfig:
+    if isinstance(value, config_type):
+        return value
+    return config_type(**value) if isinstance(value, dict) else config_type()
+
+
+def _coerce_nested_field(
+    value: NestedConfig | dict[str, Any], config_type: type[NestedConfig]
+) -> NestedConfig:
+    """Convert serialized nested dictionaries while retaining supplied objects."""
+
+    return config_type(**value) if isinstance(value, dict) else value
+
+
+@dataclass
+class AtomAttentionConfig:
+    """Sliding-window atom attention and three-dimensional RoPE settings."""
+
+    d_atom: int = field(default=128)
+    d_token: int = field(default=768)
+    n_blocks: int = field(default=3)
+    n_heads: int = field(default=4)
+    swa_window_size: int = field(default=128)
+    expansion_ratio: int = field(default=2)
+    spatial_rope_base_frequency: float = field(default=20.0)
+    n_spatial_rope_pairs_per_axis: int = field(default=2)
+    n_uid_rope_pairs: int = field(default=10)
+    uid_rope_base_frequency: float = field(default=10000.0)
+
+
+@dataclass
+class DiffusionModuleConfig:
+    """Dimensions and depth of the coordinate diffusion network."""
+
+    sigma_data: float = field(default=16.0)
+    c_atom: int = field(default=128)
+    c_token: int = field(default=768)
+    c_z: int = field(default=256)
+    c_s_inputs: int = field(default=451)
+    fourier_dim: int = field(default=256)
+    relpos_r_max: int = field(default=32)
+    relpos_s_max: int = field(default=2)
+    atom_num_blocks: int = field(default=3)
+    atom_num_heads: int = field(default=4)
+    token_num_blocks: int = field(default=12)
+    token_num_heads: int = field(default=16)
+    transition_multiplier: int = field(default=2)
+
+
+@dataclass
+class FoldingTrunkConfig:
+    """Iterative pair/single trunk dimensions."""
+
+    n_layers: int = field(default=24)
+    n_heads: int = field(default=8)
+    dropout: float = field(default=0.0)
+
+
+@dataclass
+class InputsEmbedderConfig:
+    """Input feature width and atom encoder settings."""
+
+    d_inputs: int = field(default=451)
+    atom_encoder: AtomAttentionConfig = field(default_factory=AtomAttentionConfig)
+
+    def __post_init__(self) -> None:
+        self.atom_encoder = _coerce_nested_field(self.atom_encoder, AtomAttentionConfig)
+
+
+@dataclass
+class DiffusionStructureHeadConfig:
+    """Training and inference schedules for coordinate denoising."""
+
+    diffusion_module: DiffusionModuleConfig = field(default_factory=DiffusionModuleConfig)
+    distogram_bins: int = field(default=128)
+    train_noise_log_mean: float = field(default=-1.2)
+    train_noise_log_std: float = field(default=1.5)
+    gamma_0: float = field(default=0.605)
+    gamma_min: float = field(default=1.107)
+    noise_scale: float = field(default=0.0)
+    step_scale: float = field(default=1.0)
+    inference_s_max: float = field(default=160.0)
+    inference_s_min: float = field(default=4e-4)
+    inference_p: float = field(default=8.0)
+    inference_num_steps: int = field(default=68)
+
+    def __post_init__(self) -> None:
+        self.diffusion_module = _coerce_nested_field(self.diffusion_module, DiffusionModuleConfig)
+
+
+@dataclass
+class ConfidenceHeadConfig:
+    """Confidence-bin definitions and the compact confidence trunk."""
+
+    enabled: bool = field(default=True)
+    num_plddt_bins: int = field(default=50)
+    num_pde_bins: int = field(default=64)
+    num_pae_bins: int = field(default=64)
+    min_dist: float = field(default=2.0)
+    max_dist: float = field(default=52.0)
+    distogram_bins: int = field(default=128)
+    folding_trunk: FoldingTrunkConfig = field(
+        default_factory=lambda: FoldingTrunkConfig(n_layers=4)
+    )
+
+    def __post_init__(self) -> None:
+        self.folding_trunk = _coerce_nested_field(self.folding_trunk, FoldingTrunkConfig)
+
+
+@dataclass
+class MSAEncoderConfig:
+    """Optional multiple-sequence-alignment encoder settings."""
+
+    enabled: bool = field(default=False)
+    d_msa: int = field(default=128)
+    d_hidden: int = field(default=32)
+    n_layers: int = field(default=4)
+    n_heads_msa: int = field(default=8)
+    msa_head_width: int = field(default=32)
+
+
+@dataclass
+class LMEncoderConfig:
+    """Release-model pair encoder derived from language-model states."""
+
+    enabled: bool = field(default=True)
+    n_layers: int = field(default=4)
+    lm_dropout: float = field(default=0.25)
+    per_loop_lm_dropout: bool = field(default=True)
+
+
+@dataclass
+class ParcaeConfig:
+    """Release-model diffusion-loop scheduler settings."""
+
+    enabled: bool = field(default=True)
+    poisson_mean: float = field(default=3.0)
+    min_steps: int = field(default=1)
+    max_steps: int | None = field(default=6)
+    coda_n_layers: int = field(default=2)
+
+
+_SCALAR_DEFAULTS: tuple[tuple[str, Any], ...] = (
+    ("d_single", 384),
+    ("d_pair", 256),
+    ("n_relative_residx_bins", 32),
+    ("n_relative_chain_bins", 2),
+    ("num_loops", 10),
+    ("num_diffusion_samples", 8),
+    ("disable_msa_features", False),
+    ("lm_dropout", 0.0),
+    ("force_lm_dropout_during_inference", False),
+    ("lm_mask_pct", 0.0),
+    ("lm_d_model", 2560),
+    ("lm_num_layers", 80),
+)
+_NESTED_CONFIGS = (
+    ("inputs", InputsEmbedderConfig),
+    ("folding_trunk", FoldingTrunkConfig),
+    ("structure_head", DiffusionStructureHeadConfig),
+    ("confidence_head", ConfidenceHeadConfig),
+    ("msa_encoder", MSAEncoderConfig),
+    ("parcae", ParcaeConfig),
+    ("lm_encoder", LMEncoderConfig),
+)
+
+
+class ESMFold2Config(PretrainedConfig):
+    """Serializable ESMFold2 architecture, runtime, and precision settings."""
+
+    model_type = "esmfold2"
+    has_no_defaults_at_init = True
+
+    def __init__(self, **kwargs: Any) -> None:
+        legacy_backend = normalize_esmc_attention_implementation(kwargs.get("esmc_attn_backend"))
+        requested_backend = normalize_esmc_attention_implementation(
+            kwargs.get("attn_implementation")
+        )
+        resolved_backend = requested_backend or legacy_backend
+        kwargs["attn_implementation"] = resolved_backend
+        super().__init__(**kwargs)
+
+        self.type = kwargs.get("type", "release")
+        if self.type not in {"experimental", "release"}:
+            raise ValueError(
+                f"ESMFold2Config.type must be 'release' or 'experimental', got {self.type!r}"
+            )
+
+        for name, default in _SCALAR_DEFAULTS:
+            setattr(self, name, kwargs.get(name, default))
+
+        _official_esmc_repo, default_esmc_repo = _esmc_backbone_checkpoint_ids()
+        self.esmc_id = normalize_esmc_id(kwargs.get("esmc_id", default_esmc_repo))
+        self.esmc_attn_backend = resolved_backend
+        self.esmc_precision = str(kwargs.get("esmc_precision", "auto"))
+        if self.esmc_precision not in _ESMC_PRECISIONS:
+            raise ValueError(
+                "esmc_precision must be 'auto', 'bf16', 'fp32', or 'fp8', "
+                f"got {self.esmc_precision!r}."
+            )
+
+        for name, config_type in _NESTED_CONFIGS:
+            setattr(self, name, _nested_config(kwargs.get(name), config_type))
+        if not isinstance(self.msa_encoder.enabled, bool):
+            raise TypeError("msa_encoder.enabled must be a boolean.")
+        declared_msa_conditioning = kwargs.get("msa_conditioning")
+        if "msa_conditioning" in kwargs and not isinstance(declared_msa_conditioning, bool):
+            raise TypeError("msa_conditioning must be a boolean when provided.")
+        self.msa_conditioning = (
+            self.msa_encoder.enabled
+            if "msa_conditioning" not in kwargs
+            else declared_msa_conditioning
+        )
+        if self.msa_conditioning != self.msa_encoder.enabled:
+            raise ValueError(
+                "msa_conditioning must match msa_encoder.enabled; received "
+                f"{self.msa_conditioning!r} and {self.msa_encoder.enabled!r}."
+            )
+        self.msa_encoder_overwrite = bool(kwargs.get("msa_encoder_overwrite", True))
+
+    def to_dict(self) -> dict[str, Any]:
+        output = cast(dict[str, Any], super().to_dict())
+        for name, _config_type in _NESTED_CONFIGS:
+            output[name] = asdict(getattr(self, name))
+        return output
+
+
+__all__ = [
+    "ESMFold2Config",
+    "LMEncoderConfig",
+    "MSAEncoderConfig",
+    "ParcaeConfig",
+    "normalize_esmc_attention_implementation",
+    "normalize_esmc_id",
+]

fastplms/models/esmfold2/embedding.py

@@ -0,0 +1,105 @@
+"""ESMFold2 integration for the shared FastPLMs embedding API."""
+
+from __future__ import annotations
+
+from typing import Any, ClassVar
+
+import torch
+from torch import Tensor
+
+from ...embeddings import EmbeddingBatch, EmbeddingResult, embed_dataset
+from .esmfold2_constants_esm3 import SEQUENCE_PAD_TOKEN, SEQUENCE_VOCAB
+
+_TOKEN_TO_ID = {token: index for index, token in enumerate(SEQUENCE_VOCAB)}
+_VALID_RESIDUES = frozenset(SEQUENCE_VOCAB[4:31]) - {".", "-", "|"}
+
+
+def _encode_single_chain(sequence: str) -> list[int]:
+    normalized = sequence.upper()
+    if not normalized:
+        raise ValueError("ESMFold2 dataset embedding requires at least one protein residue.")
+    invalid = sorted(set(normalized) - _VALID_RESIDUES)
+    if invalid:
+        raise ValueError(
+            "ESMFold2 dataset embedding accepts one ungapped protein chain; "
+            f"invalid symbols: {invalid}."
+        )
+    return [_TOKEN_TO_ID[residue] for residue in normalized]
+
+
+class ESMFold2EmbeddingMixin:
+    """Learned ESMC sequence summaries for ESMFold2 models."""
+
+    embedding_unsupported_pooling = frozenset({"cls", "parti"})
+    embedding_layer = "all_81_esmc_states"
+    embedding_projection = "esmfold2_learned_sequence_summary"
+    embedding_token_policy: ClassVar[dict[str, object]] = {
+        "unit": "residue",
+        "normalization": "uppercase",
+        "include": ["single-chain protein residues"],
+        "exclude": [
+            "BOS",
+            "EOS",
+            "padding",
+            "chain delimiters",
+            "non-protein tokens",
+        ],
+    }
+
+    def project_esmc_hidden_states(
+        self,
+        hidden_states: Tensor,
+        residue_mask: Tensor | None = None,
+    ) -> Tensor:
+        """Project H from ``(b, l, 81, 2560)`` to Z with shape ``(b, l, 256)``."""
+
+        if hidden_states.ndim != 4 or hidden_states.shape[-2] != 81:
+            raise ValueError(
+                "ESMFold2 projection requires the official ordered 81-state "
+                "ESMC tensor H with shape (b, l, 81, d_model)."
+            )
+        return self.language_model.project_sequence(hidden_states, residue_mask)
+
+    def _embedding_batch(self, sequences: list[str], **kwargs: Any) -> EmbeddingBatch:
+        if kwargs:
+            raise TypeError(f"Unexpected ESMFold2 embedding options: {', '.join(sorted(kwargs))}.")
+        if self._esmc is None:
+            raise RuntimeError("ESMFold2 embeddings require load_esmc=True.")
+        encoded = [_encode_single_chain(sequence) for sequence in sequences]
+        sequence_length = max(map(len, encoded))
+        b = len(encoded)
+        device = self.device
+        input_ids = torch.full(
+            (b, sequence_length),
+            SEQUENCE_PAD_TOKEN,
+            dtype=torch.long,
+            device=device,
+        )
+        residue_mask = torch.zeros((b, sequence_length), dtype=torch.bool, device=device)
+        for batch_index, token_ids in enumerate(encoded):
+            length = len(token_ids)
+            input_ids[batch_index, :length] = torch.tensor(
+                token_ids, dtype=torch.long, device=device
+            )
+            residue_mask[batch_index, :length] = True
+
+        residue_index = torch.arange(sequence_length, device=device).expand(b, -1)
+        asym_id = torch.zeros_like(input_ids)
+        mol_type = torch.zeros_like(input_ids)
+        hidden_states = self._compute_lm_hidden_states(
+            input_ids,
+            asym_id,
+            residue_index,
+            mol_type,
+            residue_mask,
+        )
+        projected = self.project_esmc_hidden_states(hidden_states, residue_mask)
+        return EmbeddingBatch(X=projected, residue_mask=residue_mask)
+
+    def embed_dataset(self, inputs: Any, **kwargs: Any) -> EmbeddingResult:
+        """Embed single-chain proteins using the learned 256-wide ESMFold2 summary."""
+
+        return embed_dataset(self, inputs, **kwargs)
+
+
+__all__ = ["ESMFold2EmbeddingMixin"]

fastplms/models/esmfold2/esmfold2_affine3d.py

@@ -0,0 +1,605 @@
+"""Differentiable rigid rotations and affine transforms for ESMFold2."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Self
+
+import torch
+from torch.nn import functional as F
+
+from .esmfold2_misc import fp32_autocast_context
+
+
+def _index_tuple(index: Any) -> tuple[Any, ...]:
+    if isinstance(index, int) or index is None:
+        return (index,)
+    return tuple(index)
+
+
+def _sqrt_subgradient(values: torch.Tensor) -> torch.Tensor:
+    """Square root with a zero subgradient for non-positive inputs."""
+
+    result = torch.zeros_like(values)
+    positive = values > 0
+    result[positive] = torch.sqrt(values[positive])
+    return result
+
+
+def _quat_invert(quaternion: torch.Tensor) -> torch.Tensor:
+    conjugate_sign = torch.tensor([1, -1, -1, -1], device=quaternion.device)
+    return quaternion * conjugate_sign
+
+
+def _quat_mult(left: torch.Tensor, right: torch.Tensor) -> torch.Tensor:
+    """Hamilton product for real-first quaternion tensors."""
+
+    aw, ax, ay, az = torch.unbind(left, -1)
+    bw, bx, by, bz = torch.unbind(right, -1)
+    return torch.stack(
+        (
+            aw * bw - ax * bx - ay * by - az * bz,
+            aw * bx + ax * bw + ay * bz - az * by,
+            aw * by - ax * bz + ay * bw + az * bx,
+            aw * bz + ax * by - ay * bx + az * bw,
+        ),
+        -1,
+    )
+
+
+def _quat_rotation(
+    quaternion: torch.Tensor,
+    points: torch.Tensor,
+) -> torch.Tensor:
+    """Rotate points using normalized real-first quaternions."""
+
+    aw, ax, ay, az = torch.unbind(quaternion, -1)
+    bx, by, bz = torch.unbind(points, -1)
+    product = torch.stack(
+        (
+            -ax * bx - ay * by - az * bz,
+            aw * bx + ay * bz - az * by,
+            aw * by - ax * bz + az * bx,
+            aw * bz + ax * by - ay * bx,
+        ),
+        -1,
+    )
+    return _quat_mult(product, _quat_invert(quaternion))[..., 1:]
+
+
+def _graham_schmidt(
+    x_axis: torch.Tensor,
+    xy_plane: torch.Tensor,
+    eps: float = 1e-12,
+) -> torch.Tensor:
+    """Construct a right-handed orthonormal frame from two directions."""
+
+    with fp32_autocast_context(x_axis.device.type):
+        e1 = xy_plane
+        denominator = torch.sqrt((x_axis**2).sum(dim=-1, keepdim=True) + eps)
+        x_axis = x_axis / denominator
+        projection = (x_axis * e1).sum(dim=-1, keepdim=True)
+        e1 = e1 - x_axis * projection
+        denominator = torch.sqrt((e1**2).sum(dim=-1, keepdim=True) + eps)
+        e1 = e1 / denominator
+        e2 = torch.cross(x_axis, e1, dim=-1)
+        return torch.stack([x_axis, e1, e2], dim=-1)
+
+
+class Rotation:
+    """Common interface for matrix-backed and quaternion-backed rotations."""
+
+    @classmethod
+    def identity(cls, shape: tuple[int, ...], **tensor_kwargs) -> Self: ...
+
+    @classmethod
+    def random(cls, shape: tuple[int, ...], **tensor_kwargs) -> Self: ...
+
+    def __getitem__(self, idx: Any) -> Self: ...
+
+    @property
+    def tensor(self) -> torch.Tensor: ...
+
+    @property
+    def shape(self) -> torch.Size: ...
+
+    def as_matrix(self) -> RotationMatrix: ...
+
+    def as_quat(self, normalize: bool = False) -> RotationQuat: ...
+
+    def compose(self, other: Self) -> Self: ...
+
+    def convert_compose(self, other: Self) -> Self: ...
+
+    def apply(self, points: torch.Tensor) -> torch.Tensor: ...
+
+    def invert(self) -> Self: ...
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.tensor.dtype
+
+    @property
+    def device(self) -> torch.device:
+        return self.tensor.device
+
+    @property
+    def requires_grad(self) -> bool:
+        return self.tensor.requires_grad
+
+    @classmethod
+    def _from_tensor(cls, tensor: torch.Tensor) -> Self:
+        return cls(tensor)  # type: ignore[call-arg]
+
+    def to(self, **kwargs) -> Self:
+        return self._from_tensor(self.tensor.to(**kwargs))
+
+    def detach(self, *args, **kwargs) -> Self:
+        return self._from_tensor(self.tensor.detach(**kwargs))
+
+    def tensor_apply(self, func) -> Self:
+        transformed = [func(component) for component in self.tensor.unbind(dim=-1)]
+        return self._from_tensor(torch.stack(transformed, dim=-1))
+
+
+class RotationQuat(Rotation):
+    """A rotation represented by a real-first quaternion."""
+
+    def __init__(self, quats: torch.Tensor, normalized: bool = False):
+        if not isinstance(quats, torch.Tensor):
+            raise TypeError("quats must be a Torch tensor.")
+        if quats.ndim == 0 or quats.shape[-1] != 4:
+            raise ValueError(
+                f"quats must have trailing dimension 4, got shape {tuple(quats.shape)}."
+            )
+        if not isinstance(normalized, bool):
+            raise TypeError("normalized must be a boolean.")
+        self._normalized = normalized
+        if normalized:
+            quats = F.normalize(quats.to(torch.float32), dim=-1)
+            self._quats = quats.where(quats[..., :1] >= 0, -quats)
+        else:
+            self._quats = quats.to(torch.float32)
+
+    @property
+    def tensor(self) -> torch.Tensor:
+        return self._quats
+
+    @property
+    def shape(self) -> torch.Size:
+        return self._quats.shape[:-1]
+
+    @classmethod
+    def identity(cls, shape, **tensor_kwargs) -> RotationQuat:
+        quaternions = torch.ones((*shape, 4), **tensor_kwargs)
+        selector = torch.tensor([1, 0, 0, 0], device=quaternions.device)
+        return cls(quaternions * selector)
+
+    @classmethod
+    def random(cls, shape, **tensor_kwargs) -> RotationQuat:
+        return cls(torch.randn((*shape, 4), **tensor_kwargs), normalized=True)
+
+    def __getitem__(self, idx: Any) -> RotationQuat:
+        indices = _index_tuple(idx)
+        return RotationQuat(self._quats[(*indices, slice(None))])
+
+    def normalized(self) -> RotationQuat:
+        if self._normalized:
+            return self
+        return RotationQuat(self._quats, normalized=True)
+
+    def as_quat(self, normalize: bool = False) -> RotationQuat:
+        return self
+
+    def as_matrix(self) -> RotationMatrix:
+        quaternion = self.normalized().tensor
+        r, i, j, k = torch.unbind(quaternion, -1)
+        scale = 2.0 / torch.linalg.norm(quaternion, dim=-1)
+        elements = torch.stack(
+            (
+                1 - scale * (j * j + k * k),
+                scale * (i * j - k * r),
+                scale * (i * k + j * r),
+                scale * (i * j + k * r),
+                1 - scale * (i * i + k * k),
+                scale * (j * k - i * r),
+                scale * (i * k - j * r),
+                scale * (j * k + i * r),
+                1 - scale * (i * i + j * j),
+            ),
+            -1,
+        )
+        return RotationMatrix(elements.reshape((*quaternion.shape[:-1], 3, 3)))
+
+    def compose(self, other: RotationQuat) -> RotationQuat:
+        with fp32_autocast_context(self.device.type):
+            return RotationQuat(_quat_mult(self._quats, other._quats))
+
+    def convert_compose(self, other: Rotation) -> RotationQuat:
+        return self.compose(other.as_quat())
+
+    def apply(self, points: torch.Tensor) -> torch.Tensor:
+        return _quat_rotation(self.normalized()._quats, points)
+
+    def invert(self) -> RotationQuat:
+        return RotationQuat(_quat_invert(self._quats))
+
+
+class RotationMatrix(Rotation):
+    """A rotation represented by a dense FP32 matrix."""
+
+    def __init__(self, rots: torch.Tensor):
+        if not isinstance(rots, torch.Tensor):
+            raise TypeError("rots must be a Torch tensor.")
+        if rots.ndim > 0 and rots.shape[-1] == 9:
+            rots = rots.unflatten(-1, (3, 3))
+        if rots.ndim < 2 or rots.shape[-2:] != (3, 3):
+            raise ValueError(
+                "rots must have trailing shape (3, 3) or flattened width 9, got "
+                f"shape {tuple(rots.shape)}."
+            )
+        self._rots = rots.to(torch.float32)
+
+    @property
+    def tensor(self) -> torch.Tensor:
+        return self._rots.flatten(-2)
+
+    @property
+    def shape(self) -> torch.Size:
+        return self._rots.shape[:-2]
+
+    @classmethod
+    def identity(cls, shape, **tensor_kwargs) -> RotationMatrix:
+        matrix = torch.eye(3, **tensor_kwargs)
+        matrix = matrix.view(*(1 for _ in shape), 3, 3)
+        return cls(matrix.expand(*shape, -1, -1))
+
+    @classmethod
+    def random(cls, shape, **tensor_kwargs) -> RotationMatrix:
+        return RotationQuat.random(shape, **tensor_kwargs).as_matrix()
+
+    @staticmethod
+    def from_graham_schmidt(
+        x_axis: torch.Tensor,
+        xy_plane: torch.Tensor,
+        eps: float = 1e-12,
+    ) -> RotationMatrix:
+        return RotationMatrix(_graham_schmidt(x_axis, xy_plane, eps))
+
+    def __getitem__(self, idx: Any) -> RotationMatrix:
+        indices = _index_tuple(idx)
+        return RotationMatrix(self._rots[(*indices, slice(None), slice(None))])
+
+    def as_matrix(self) -> RotationMatrix:
+        return self
+
+    def to_3x3(self) -> torch.Tensor:
+        return self._rots
+
+    def as_quat(self, normalize: bool = False) -> RotationQuat:
+        m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind(
+            self._rots.flatten(-2),
+            dim=-1,
+        )
+        q_abs = _sqrt_subgradient(
+            torch.stack(
+                (
+                    1.0 + m00 + m11 + m22,
+                    1.0 + m00 - m11 - m22,
+                    1.0 - m00 + m11 - m22,
+                    1.0 - m00 - m11 + m22,
+                ),
+                dim=-1,
+            )
+        )
+        products = torch.stack(
+            (
+                q_abs[..., 0] ** 2,
+                m21 - m12,
+                m02 - m20,
+                m10 - m01,
+                m21 - m12,
+                q_abs[..., 1] ** 2,
+                m10 + m01,
+                m02 + m20,
+                m02 - m20,
+                m10 + m01,
+                q_abs[..., 2] ** 2,
+                m12 + m21,
+                m10 - m01,
+                m20 + m02,
+                m21 + m12,
+                q_abs[..., 3] ** 2,
+            ),
+            dim=-1,
+        ).unflatten(-1, (4, 4))
+        floor = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device)
+        candidates = products / (2.0 * q_abs[..., None].max(floor))
+        best = torch.zeros_like(q_abs, dtype=torch.bool)
+        best.scatter_(-1, q_abs.argmax(dim=-1, keepdim=True), True)
+        quaternion = candidates[best, :].reshape(q_abs.shape)
+        return RotationQuat(quaternion)
+
+    def compose(self, other: RotationMatrix) -> RotationMatrix:
+        with fp32_autocast_context(self.device.type):
+            return RotationMatrix(self._rots @ other._rots)
+
+    def convert_compose(self, other: Rotation) -> RotationMatrix:
+        return self.compose(other.as_matrix())
+
+    def apply(self, points: torch.Tensor) -> torch.Tensor:
+        with fp32_autocast_context(self.device.type):
+            if self._rots.shape[-3] == 1:
+                return points @ self._rots.transpose(-1, -2).squeeze(-3)
+            return torch.einsum("...ij,...j", self._rots, points)
+
+    def invert(self) -> RotationMatrix:
+        return RotationMatrix(self._rots.transpose(-1, -2))
+
+
+@dataclass(frozen=True)
+class Affine3D:
+    """A rigid transform with translation and rotation components."""
+
+    trans: torch.Tensor
+    rot: Rotation
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.trans, torch.Tensor):
+            raise TypeError("trans must be a Torch tensor.")
+        if not isinstance(self.rot, Rotation):
+            raise TypeError("rot must implement the ESMFold2 Rotation interface.")
+        if self.trans.ndim == 0 or self.trans.shape[-1] != 3:
+            raise ValueError(
+                "trans must have trailing dimension 3, got "
+                f"shape {tuple(self.trans.shape)}."
+            )
+        if self.trans.shape[:-1] != self.rot.shape:
+            raise ValueError(
+                "translation and rotation batch shapes must match, got "
+                f"{tuple(self.trans.shape[:-1])} and {tuple(self.rot.shape)}."
+            )
+
+    @property
+    def shape(self) -> torch.Size:
+        return self.trans.shape[:-1]
+
+    @property
+    def dtype(self) -> torch.dtype:
+        return self.trans.dtype
+
+    @property
+    def device(self) -> torch.device:
+        return self.trans.device
+
+    @property
+    def requires_grad(self) -> bool:
+        return self.trans.requires_grad
+
+    @property
+    def tensor(self) -> torch.Tensor:
+        return torch.cat((self.rot.tensor, self.trans), dim=-1)
+
+    @staticmethod
+    def identity(
+        shape_or_affine: tuple[int, ...] | Affine3D,
+        rotation_type: type[Rotation] = RotationMatrix,
+        **tensor_kwargs,
+    ) -> Affine3D:
+        if isinstance(shape_or_affine, Affine3D):
+            kwargs = {
+                "dtype": shape_or_affine.dtype,
+                "device": shape_or_affine.device,
+            }
+            kwargs.update(tensor_kwargs)
+            shape = shape_or_affine.shape
+            rotation_type = type(shape_or_affine.rot)
+        else:
+            kwargs = tensor_kwargs
+            shape = shape_or_affine
+        return Affine3D(
+            torch.zeros((*shape, 3), **kwargs),
+            rotation_type.identity(shape, **kwargs),
+        )
+
+    @staticmethod
+    def random(
+        shape: tuple[int, ...],
+        std: float = 1,
+        rotation_type: type[Rotation] = RotationMatrix,
+        **tensor_kwargs,
+    ) -> Affine3D:
+        translation = torch.randn((*shape, 3), **tensor_kwargs).mul(std)
+        rotation = rotation_type.random(shape, **tensor_kwargs)
+        return Affine3D(trans=translation, rot=rotation)
+
+    @staticmethod
+    def from_tensor(tensor: torch.Tensor) -> Affine3D:
+        if not isinstance(tensor, torch.Tensor):
+            raise TypeError("tensor must be a Torch tensor.")
+        if tensor.ndim == 0:
+            raise ValueError("tensor must have at least one dimension.")
+        width = tensor.shape[-1]
+        if width == 4:
+            if tensor.ndim < 2 or tensor.shape[-2] not in (3, 4):
+                raise ValueError(
+                    "matrix-form affine tensors must have trailing shape (3, 4) or "
+                    f"(4, 4), got {tuple(tensor.shape)}."
+                )
+            translation = tensor[..., :3, 3]
+            rotation: Rotation = RotationMatrix(tensor[..., :3, :3])
+        elif width == 6:
+            translation = tensor[..., -3:]
+            rotation = RotationQuat(F.pad(tensor[..., :3], (1, 0), value=1))
+        elif width == 7:
+            translation = tensor[..., -3:]
+            rotation = RotationQuat(tensor[..., :4])
+        elif width == 12:
+            translation = tensor[..., -3:]
+            rotation = RotationMatrix(tensor[..., :-3].unflatten(-1, (3, 3)))
+        else:
+            raise RuntimeError(
+                f"Cannot detect rotation format from {tensor.shape[-1] - 3}-d flat vector"
+            )
+        return Affine3D(translation, rotation)
+
+    @staticmethod
+    def from_tensor_pair(
+        translation: torch.Tensor,
+        rotation: torch.Tensor,
+    ) -> Affine3D:
+        return Affine3D(translation, RotationMatrix(rotation))
+
+    @staticmethod
+    def from_graham_schmidt(
+        neg_x_axis: torch.Tensor,
+        origin: torch.Tensor,
+        xy_plane: torch.Tensor,
+        eps: float = 1e-10,
+    ) -> Affine3D:
+        x_axis = origin - neg_x_axis
+        plane_direction = xy_plane - origin
+        rotation = RotationMatrix.from_graham_schmidt(
+            x_axis,
+            plane_direction,
+            eps,
+        )
+        return Affine3D(trans=origin, rot=rotation)
+
+    @staticmethod
+    def cat(affines: list[Affine3D], dim: int = 0) -> Affine3D:
+        if not affines:
+            raise ValueError("affines must contain at least one transform.")
+        if any(not isinstance(affine, Affine3D) for affine in affines):
+            raise TypeError("affines must contain only Affine3D instances.")
+        if dim < 0:
+            dim = len(affines[0].shape) + dim
+        return Affine3D.from_tensor(torch.cat([affine.tensor for affine in affines], dim=dim))
+
+    def __getitem__(self, idx: Any) -> Affine3D:
+        indices = _index_tuple(idx)
+        translation = self.trans[(*indices, slice(None))]
+        return Affine3D(trans=translation, rot=self.rot[idx])
+
+    def to(self, **kwargs) -> Affine3D:
+        return Affine3D(self.trans.to(**kwargs), self.rot.to(**kwargs))
+
+    def detach(self, *args, **kwargs) -> Affine3D:
+        return Affine3D(
+            self.trans.detach(**kwargs),
+            self.rot.detach(**kwargs),
+        )
+
+    def tensor_apply(self, func) -> Affine3D:
+        components = [func(value) for value in self.tensor.unbind(dim=-1)]
+        return Affine3D.from_tensor(torch.stack(components, dim=-1))
+
+    def as_matrix(self) -> Affine3D:
+        return Affine3D(trans=self.trans, rot=self.rot.as_matrix())
+
+    def as_quat(self, normalize: bool = False) -> Affine3D:
+        return Affine3D(
+            trans=self.trans,
+            rot=self.rot.as_quat(normalize),
+        )
+
+    def compose(
+        self,
+        other: Affine3D,
+        autoconvert: bool = False,
+    ) -> Affine3D:
+        compose_rotation = self.rot.convert_compose if autoconvert else self.rot.compose
+        rotation = compose_rotation(other.rot)
+        translation = self.rot.apply(other.trans) + self.trans
+        return Affine3D(trans=translation, rot=rotation)
+
+    def compose_rotation(
+        self,
+        other: Rotation,
+        autoconvert: bool = False,
+    ) -> Affine3D:
+        compose = self.rot.convert_compose if autoconvert else self.rot.compose
+        return Affine3D(trans=self.trans, rot=compose(other))
+
+    def scale(self, value: torch.Tensor | float) -> Affine3D:
+        return Affine3D(self.trans * value, self.rot)
+
+    def mask(self, mask: torch.Tensor, with_zero: bool = False) -> Affine3D:
+        if with_zero:
+            masked = torch.zeros_like(self.tensor).where(
+                mask[..., None],
+                self.tensor,
+            )
+            return Affine3D.from_tensor(masked)
+        identity = self.identity(
+            self.shape,
+            rotation_type=type(self.rot),
+            device=self.device,
+            dtype=self.dtype,
+        ).tensor
+        return Affine3D.from_tensor(identity.where(mask[..., None], self.tensor))
+
+    def apply(self, points: torch.Tensor) -> torch.Tensor:
+        return self.rot.apply(points) + self.trans
+
+    def invert(self) -> Affine3D:
+        rotation = self.rot.invert()
+        return Affine3D(trans=-rotation.apply(self.trans), rot=rotation)
+
+
+def build_affine3d_from_coordinates(
+    coords: torch.Tensor,
+) -> tuple[Affine3D, torch.Tensor]:
+    """Build residue frames from X with shape (b, l, 3, 3)."""
+
+    if not isinstance(coords, torch.Tensor):
+        raise TypeError("coords must be a Torch tensor.")
+    if coords.ndim != 4 or coords.shape[-2:] != (3, 3):
+        raise ValueError(
+            "coords must have shape (batch, length, 3, 3), got "
+            f"{tuple(coords.shape)}."
+        )
+
+    maximum_distance = 1e6
+    coord_mask = torch.all(
+        torch.all(
+            torch.isfinite(coords) & (coords < maximum_distance),
+            dim=-1,
+        ),
+        dim=-1,
+    )
+
+    def backbone_affine(positions: torch.Tensor) -> Affine3D:
+        n, ca, c = positions.unbind(dim=-2)
+        return Affine3D.from_graham_schmidt(c, ca, n)
+
+    coords = coords.clone().float()
+    coords[~coord_mask] = 0
+    average = coords.masked_fill(~coord_mask[..., None, None], 0).sum(1) / (
+        coord_mask.sum(-1)[..., None, None] + 1e-8
+    )
+    average_affine = backbone_affine(average.float()).as_matrix()
+
+    b, length, _, _ = coords.shape
+    rotation = average_affine.rot.tensor[..., None, :].expand(b, length, 9)
+    translation = average_affine.trans[..., None, :].expand(b, length, 3)
+    identity = RotationMatrix.identity(
+        (b, length),
+        dtype=torch.float32,
+        device=coords.device,
+        requires_grad=False,
+    )
+    rotation = rotation.where(
+        coord_mask.any(-1)[..., None, None],
+        identity.tensor,
+    )
+    missing_frame = Affine3D(translation, RotationMatrix(rotation))
+
+    residue_frame = backbone_affine(coords.float())
+    residue_frame = Affine3D.from_tensor(
+        residue_frame.tensor.where(
+            coord_mask[..., None],
+            missing_frame.tensor,
+        )
+    )
+    return residue_frame, coord_mask

fastplms/models/esmfold2/esmfold2_aligner.py

@@ -0,0 +1,87 @@
+"""Rigid alignment for structure dataclasses."""
+
+from __future__ import annotations
+
+from dataclasses import Field, replace
+from typing import Any, ClassVar, Protocol, TypeVar
+
+import numpy as np
+import torch
+from torch import Tensor
+
+from .esmfold2_protein_structure import compute_affine_and_rmsd
+
+
+class Alignable(Protocol):
+    """Minimum structure interface accepted by :class:`Aligner`."""
+
+    __dataclass_fields__: ClassVar[dict[str, Field[Any]]]
+
+    @property
+    def atom37_positions(self) -> np.ndarray: ...
+
+    @property
+    def atom37_mask(self) -> np.ndarray: ...
+
+    def __len__(self) -> int: ...
+
+
+AlignableT = TypeVar("AlignableT", bound=Alignable)
+
+
+def _coordinate_batch(structure: Alignable) -> Tensor:
+    return torch.as_tensor(structure.atom37_positions, dtype=torch.double).unsqueeze(0)
+
+
+def _shared_atom_mask(mobile: Alignable, target: Alignable, backbone_only: bool) -> Tensor:
+    shared = np.asarray(mobile.atom37_mask, dtype=bool) & np.asarray(
+        target.atom37_mask,
+        dtype=bool,
+    )
+    if backbone_only:
+        shared = shared.copy()
+        shared[:, 3:] = False
+    return torch.from_numpy(shared).unsqueeze(0)
+
+
+class Aligner:
+    """Fit a mobile structure onto a target with masked Kabsch alignment."""
+
+    def __init__(
+        self,
+        mobile: Alignable,
+        target: Alignable,
+        only_use_backbone: bool = False,
+        use_reflection: bool = False,
+    ) -> None:
+        if len(mobile) != len(target):
+            raise AssertionError("mobile and target must contain the same residue count")
+
+        mobile_coordinates = _coordinate_batch(mobile)
+        target_coordinates = _coordinate_batch(target)
+        if use_reflection:
+            target_coordinates = -target_coordinates
+        atom_mask = _shared_atom_mask(mobile, target, only_use_backbone)
+        self._affine3D, rmsd = compute_affine_and_rmsd(
+            mobile_coordinates,
+            target_coordinates,
+            atom_exists_mask=atom_mask,
+        )
+        self._rmsd = rmsd.item()
+
+    @property
+    def rmsd(self) -> float:
+        return self._rmsd
+
+    def apply(self, mobile: AlignableT) -> AlignableT:
+        """Return a dataclass copy with all present atom coordinates aligned."""
+
+        present = np.asarray(mobile.atom37_mask, dtype=bool)
+        packed = torch.as_tensor(
+            mobile.atom37_positions[present],
+            dtype=torch.float32,
+        ).unsqueeze(0)
+        aligned = self._affine3D.apply(packed).squeeze(0).cpu().numpy()
+        atom37_positions = np.full_like(mobile.atom37_positions, np.nan)
+        atom37_positions[present] = aligned
+        return replace(mobile, atom37_positions=atom37_positions)

fastplms/models/esmfold2/esmfold2_atom_indexer.py

@@ -0,0 +1,30 @@
+"""Name-based views into an atom-axis property."""
+
+from __future__ import annotations
+
+from operator import attrgetter
+from typing import Any
+
+import numpy as np
+
+from .esmfold2_protein_structure import index_by_atom_name
+
+
+class AtomIndexer:
+    """Select named atoms from one property of a structure-like object.
+
+    The wrapper intentionally remains small because ``ProteinChain.atom37`` and
+    related public properties expose it directly.
+    """
+
+    __slots__ = ("_get_property", "dim", "property", "structure")
+
+    def __init__(self, structure: Any, property: str, dim: int):
+        self.structure = structure
+        self.property = property
+        self.dim = dim
+        self._get_property = attrgetter(property)
+
+    def __getitem__(self, atom_names: str | list[str]) -> np.ndarray:
+        values = self._get_property(self.structure)
+        return index_by_atom_name(values, atom_names, dim=self.dim)

fastplms/models/esmfold2/esmfold2_conformers.py

@@ -0,0 +1,402 @@
+"""Lazy access to Chemical Component Dictionary conformers.
+
+The feature pipeline depends on atom names, formal charges, bonds, leaving-atom
+flags, and one preferred reference conformer. Asset resolution is explicit at
+``load_ccd`` time; importing this module performs no download or file access.
+"""
+
+from __future__ import annotations
+
+import os
+import pickle
+import stat
+import tempfile
+from collections.abc import Iterator
+from contextlib import contextmanager
+from dataclasses import dataclass
+from hashlib import file_digest
+from pathlib import Path
+from typing import Any, BinaryIO
+
+import numpy as np
+from huggingface_hub import hf_hub_download
+from huggingface_hub.constants import HF_HUB_CACHE
+
+from fastplms.registry import RuntimeAsset, get_model_registry
+
+from .esmfold2_constants import RES_TYPE_TO_CCD
+
+_CCD_ENVIRONMENT_VARIABLE = "ESMCFOLD_CCD_PATH"
+_CCD_ASSET_ID = "esmfold2_ccd"
+
+
+@dataclass(frozen=True)
+class _ResolvedAsset:
+    path: Path
+    trusted_hub_cache_root: Path | None = None
+
+
+def _asset_contract() -> RuntimeAsset:
+    """Return the manifest-owned identity of the trusted CCD pickle."""
+
+    try:
+        asset = get_model_registry().runtime_assets[_CCD_ASSET_ID]
+    except KeyError as error:
+        raise RuntimeError(
+            f"The package manifest does not declare runtime asset {_CCD_ASSET_ID!r}."
+        ) from error
+    if asset.trust_kind != "hash_pinned_pickle":
+        raise RuntimeError(
+            f"Runtime asset {_CCD_ASSET_ID!r} must use the hash_pinned_pickle trust policy."
+        )
+    return asset
+
+
+@contextmanager
+def _open_verified_asset(
+    asset_path: Path,
+    contract: RuntimeAsset,
+    *,
+    trusted_hub_cache_root: Path | None = None,
+) -> Iterator[BinaryIO]:
+    """Yield a private snapshot containing exactly the verified pickle bytes."""
+
+    try:
+        path_state = asset_path.lstat()
+    except FileNotFoundError as error:
+        raise FileNotFoundError(f"CCD asset does not exist: {asset_path}") from error
+    opened_path = asset_path
+    if stat.S_ISLNK(path_state.st_mode):
+        if trusted_hub_cache_root is None:
+            raise ValueError(f"CCD asset must not be a symlink: {asset_path}")
+        opened_path = _resolve_trusted_hub_snapshot_link(
+            asset_path,
+            contract,
+            trusted_hub_cache_root,
+        )
+        path_state = opened_path.lstat()
+    if not stat.S_ISREG(path_state.st_mode):
+        raise ValueError(f"CCD asset must be a regular file: {asset_path}")
+
+    flags = os.O_RDONLY | getattr(os, "O_BINARY", 0) | getattr(os, "O_CLOEXEC", 0)
+    flags |= getattr(os, "O_NOFOLLOW", 0)
+    descriptor: int | None = None
+    try:
+        descriptor = os.open(opened_path, flags)
+        opened_state = os.fstat(descriptor)
+        if not stat.S_ISREG(opened_state.st_mode):
+            raise ValueError(f"CCD asset must be a regular file: {asset_path}")
+        if (path_state.st_dev, path_state.st_ino) != (
+            opened_state.st_dev,
+            opened_state.st_ino,
+        ):
+            raise ValueError(f"CCD asset changed while it was being opened: {asset_path}")
+
+        source = os.fdopen(descriptor, "rb")
+        descriptor = None
+        with source, tempfile.TemporaryFile(mode="w+b") as snapshot:
+            actual_size = opened_state.st_size
+            if actual_size != contract.size:
+                raise ValueError(
+                    "CCD asset size mismatch: "
+                    f"expected {contract.size} bytes, received {actual_size}."
+                )
+            # Copy into a loader-owned OS temporary file. Hashing and
+            # deserialization then consume the same immutable snapshot, so a
+            # path replacement or in-place source write cannot substitute
+            # unverified pickle bytes after validation.
+            remaining = contract.size
+            while remaining:
+                chunk = source.read(min(1024 * 1024, remaining))
+                if not chunk:
+                    break
+                snapshot.write(chunk)
+                remaining -= len(chunk)
+            copied_size = snapshot.tell()
+            extra_byte = source.read(1)
+            if remaining or extra_byte:
+                observed_size = copied_size if remaining else copied_size + len(extra_byte)
+                raise ValueError(
+                    "CCD asset size changed while it was being copied: "
+                    f"expected {contract.size} bytes, received at least {observed_size}."
+                )
+            snapshot.flush()
+            snapshot.seek(0)
+            actual_hash = file_digest(snapshot, "sha256").hexdigest()
+            if actual_hash != contract.sha256:
+                raise ValueError(
+                    "CCD asset SHA256 mismatch; refusing to cross the "
+                    "trusted-pickle boundary."
+                )
+            snapshot.seek(0)
+            yield snapshot
+    finally:
+        if descriptor is not None:
+            os.close(descriptor)
+
+
+def _resolve_trusted_hub_snapshot_link(
+    asset_path: Path,
+    contract: RuntimeAsset,
+    cache_root: Path,
+) -> Path:
+    """Resolve only the immutable Hub snapshot link declared by the manifest."""
+
+    root = cache_root.expanduser().resolve(strict=True)
+    if len(contract.revision) != 40 or any(
+        character not in "0123456789abcdef" for character in contract.revision.lower()
+    ):
+        raise ValueError("CCD Hub asset revision must be an immutable 40-character commit.")
+    relative_asset = Path(contract.path)
+    if relative_asset.is_absolute() or ".." in relative_asset.parts:
+        raise ValueError(f"CCD Hub asset path is unsafe: {contract.path!r}")
+    repository_cache = root / f"models--{contract.repository.replace('/', '--')}"
+    try:
+        repository_cache.resolve(strict=True).relative_to(root)
+    except (FileNotFoundError, ValueError) as error:
+        raise ValueError(
+            f"CCD Hub repository cache escapes the effective Hub cache root: {repository_cache}"
+        ) from error
+    snapshot_root = repository_cache / "snapshots" / contract.revision
+    expected_path = snapshot_root / relative_asset
+    lexical_path = Path(os.path.abspath(asset_path))
+    if lexical_path != Path(os.path.abspath(expected_path)):
+        raise ValueError(
+            "CCD Hub symlink is not the manifest-owned immutable snapshot path: "
+            f"{asset_path}"
+        )
+
+    try:
+        asset_path.parent.resolve(strict=True).relative_to(root)
+    except (FileNotFoundError, ValueError) as error:
+        raise ValueError(
+            f"CCD Hub snapshot path escapes the effective Hub cache root: {asset_path}"
+        ) from error
+
+    resolved = asset_path.resolve(strict=True)
+    blob_root = (repository_cache / "blobs").resolve(strict=True)
+    try:
+        blob_root.relative_to(root)
+        resolved.relative_to(blob_root)
+    except ValueError as error:
+        raise ValueError(
+            f"CCD Hub snapshot link escapes its repository blob cache: {asset_path}"
+        ) from error
+    if not resolved.is_file() or resolved.is_symlink():
+        raise ValueError(f"CCD Hub snapshot target must be a regular file: {resolved}")
+    return resolved
+
+
+class _ChemicalComponentStore:
+    def __init__(self) -> None:
+        self.molecules: dict[str, Any] | None = None
+        self.conformers: dict[str, dict[str, np.ndarray]] = {}
+        self.atoms: dict[str, list[tuple[str, str, int]]] = {}
+        self.bonds: dict[str, list[tuple[str, str]]] = {}
+        self.leaving_atoms: dict[str, set[str]] = {}
+        self.standard_positions: dict[tuple[int, str], np.ndarray | None] = {}
+        self.ligand_positions: dict[tuple[str, str], np.ndarray | None] = {}
+
+    def load(self, cache_dir: Path | str | None = None) -> dict[str, Any]:
+        if self.molecules is not None:
+            return self.molecules
+        contract = _asset_contract()
+        resolved = self._resolve_asset_location(cache_dir, contract)
+        asset = resolved.path
+        try:
+            # SECURITY: the private snapshot is both hash-validated and
+            # deserialized, closing path-replacement and in-place-write races.
+            with _open_verified_asset(
+                asset,
+                contract,
+                trusted_hub_cache_root=resolved.trusted_hub_cache_root,
+            ) as handle:
+                loaded = pickle.load(handle)
+        except FileNotFoundError:
+            raise
+        except Exception as error:
+            raise ValueError(f"Could not read the CCD asset at {asset}: {error}") from error
+        if loaded is not None and not isinstance(loaded, dict):
+            raise TypeError("The CCD asset must contain a component dictionary.")
+        self.molecules = loaded or {}
+        return self.molecules
+
+    @staticmethod
+    def _resolve_asset(cache_dir: Path | str | None) -> Path:
+        contract = _asset_contract()
+        return _ChemicalComponentStore._resolve_asset_location(cache_dir, contract).path
+
+    @staticmethod
+    def _resolve_asset_location(
+        cache_dir: Path | str | None,
+        contract: RuntimeAsset,
+    ) -> _ResolvedAsset:
+        configured = os.environ.get(_CCD_ENVIRONMENT_VARIABLE)
+        if configured:
+            asset = Path(configured).expanduser()
+        elif cache_dir is not None:
+            asset = Path(cache_dir).expanduser() / contract.path
+        else:
+            try:
+                asset = Path(
+                    hf_hub_download(
+                        repo_id=contract.repository,
+                        filename=contract.path,
+                        revision=contract.revision,
+                    )
+                )
+            except Exception as error:
+                raise FileNotFoundError(
+                    "Could not resolve the ESMFold2 CCD asset. Set "
+                    f"{_CCD_ENVIRONMENT_VARIABLE} or populate the Hugging Face cache."
+                ) from error
+            return _ResolvedAsset(
+                path=asset,
+                trusted_hub_cache_root=Path(HF_HUB_CACHE),
+            )
+        return _ResolvedAsset(path=asset)
+
+    def _component_with_conformer(self, component_id: str):
+        molecule = self.load().get(component_id)
+        if molecule is None or molecule.GetNumConformers() == 0:
+            return None, None
+
+        conformers = list(molecule.GetConformers())
+        priority = {"Computed": 0, "Ideal": 1}
+        selected_index = min(
+            range(len(conformers)),
+            key=lambda index: priority.get(conformers[index].GetPropsAsDict().get("name"), 2),
+        )
+
+        from rdkit import Chem
+
+        heavy_molecule = Chem.RemoveHs(molecule, sanitize=False)
+        if heavy_molecule.GetNumConformers() == 0:
+            return None, None
+        conformer_index = min(selected_index, heavy_molecule.GetNumConformers() - 1)
+        return heavy_molecule, heavy_molecule.GetConformer(conformer_index)
+
+    def conformer(self, component_id: str) -> dict[str, np.ndarray] | None:
+        if component_id not in self.conformers:
+            molecule, conformer = self._component_with_conformer(component_id)
+            positions: dict[str, np.ndarray] = {}
+            if molecule is not None and conformer is not None:
+                for atom in molecule.GetAtoms():
+                    atom_name = atom.GetPropsAsDict().get("name")
+                    if not isinstance(atom_name, str) or not atom_name:
+                        continue
+                    point = conformer.GetAtomPosition(atom.GetIdx())
+                    positions[atom_name] = np.asarray((point.x, point.y, point.z), dtype=np.float32)
+            self.conformers[component_id] = positions
+        result = self.conformers[component_id]
+        return result or None
+
+    def atom_records(self, component_id: str) -> list[tuple[str, str, int]] | None:
+        if component_id not in self.atoms:
+            molecule, _conformer = self._component_with_conformer(component_id)
+            records: list[tuple[str, str, int]] = []
+            if molecule is not None:
+                for atom in molecule.GetAtoms():
+                    atom_name = atom.GetPropsAsDict().get("name")
+                    if isinstance(atom_name, str) and atom_name:
+                        records.append((atom_name, atom.GetSymbol(), atom.GetFormalCharge()))
+            self.atoms[component_id] = records
+        result = self.atoms[component_id]
+        return result or None
+
+    def bond_records(self, component_id: str) -> list[tuple[str, str]] | None:
+        if component_id not in self.bonds:
+            molecule, _conformer = self._component_with_conformer(component_id)
+            records: list[tuple[str, str]] = []
+            if molecule is not None:
+                names = {
+                    atom.GetIdx(): atom.GetPropsAsDict().get("name") for atom in molecule.GetAtoms()
+                }
+                for bond in molecule.GetBonds():
+                    first = names.get(bond.GetBeginAtomIdx())
+                    second = names.get(bond.GetEndAtomIdx())
+                    if isinstance(first, str) and first and isinstance(second, str) and second:
+                        records.append((first, second))
+            self.bonds[component_id] = records
+        result = self.bonds[component_id]
+        return result or None
+
+    def component_leaving_atoms(self, component_id: str) -> set[str]:
+        if component_id not in self.leaving_atoms:
+            molecule = self.load().get(component_id)
+            names: set[str] = set()
+            if molecule is not None:
+                for atom in molecule.GetAtoms():
+                    if atom.HasProp("leaving_atom") and atom.GetProp("leaving_atom") == "1":
+                        name = atom.GetProp("name") if atom.HasProp("name") else ""
+                        if name:
+                            names.add(name)
+            self.leaving_atoms[component_id] = names
+        return self.leaving_atoms[component_id]
+
+
+_STORE = _ChemicalComponentStore()
+
+
+def load_ccd(cache_dir: Path | str | None = None) -> dict[str, Any]:
+    """Load and cache the CCD asset, resolving it only when called."""
+
+    return _STORE.load(cache_dir)
+
+
+def get_ccd_conformer(component_id: str) -> dict[str, np.ndarray] | None:
+    """Return the preferred heavy-atom conformer by atom name."""
+
+    return _STORE.conformer(component_id)
+
+
+def get_idealized_atom_pos(res_type: int, atom_name: str) -> np.ndarray | None:
+    """Return one standard-residue atom position from the preferred conformer."""
+
+    key = (res_type, atom_name)
+    if key not in _STORE.standard_positions:
+        component_id = RES_TYPE_TO_CCD.get(res_type)
+        conformer = _STORE.conformer(component_id) if component_id is not None else None
+        _STORE.standard_positions[key] = None if conformer is None else conformer.get(atom_name)
+    return _STORE.standard_positions[key]
+
+
+def get_ligand_idealized_atom_pos(residue_name: str, atom_name: str) -> np.ndarray | None:
+    """Return one ligand atom position from the preferred conformer."""
+
+    key = (residue_name, atom_name)
+    if key not in _STORE.ligand_positions:
+        conformer = _STORE.conformer(residue_name)
+        _STORE.ligand_positions[key] = None if conformer is None else conformer.get(atom_name)
+    return _STORE.ligand_positions[key]
+
+
+def get_ligand_ccd_atoms_with_charges(
+    component_id: str,
+) -> list[tuple[str, str, int]] | None:
+    """Return heavy-atom name, element, and formal-charge records."""
+
+    return _STORE.atom_records(component_id)
+
+
+def get_ligand_ccd_bonds(component_id: str) -> list[tuple[str, str]] | None:
+    """Return bonds as component atom-name pairs."""
+
+    return _STORE.bond_records(component_id)
+
+
+def get_ccd_leaving_atoms(component_id: str) -> set[str]:
+    """Return atoms removed when a CCD component is polymerized."""
+
+    return _STORE.component_leaving_atoms(component_id)
+
+
+__all__ = [
+    "get_ccd_conformer",
+    "get_ccd_leaving_atoms",
+    "get_idealized_atom_pos",
+    "get_ligand_ccd_atoms_with_charges",
+    "get_ligand_ccd_bonds",
+    "get_ligand_idealized_atom_pos",
+    "load_ccd",
+]

fastplms/models/esmfold2/esmfold2_constants.py

@@ -0,0 +1,156 @@
+"""Declarative molecular schema for ESMFold2 feature preparation.
+
+The package manifest owns the upstream revision and license provenance. This
+module expresses the corresponding checkpoint-facing integer schema as compact
+ordered records, then derives lookup tables from those records. The generated
+tables are validated at import without reading files, downloading assets, or
+mutating process state.
+"""
+
+from __future__ import annotations
+
+SCHEMA_PROVENANCE = {
+    "manifest_family": "esmfold2",
+    "contract": "biohub_esmfold2_input_v1",
+}
+
+
+def _words(value: str) -> list[str]:
+    return value.split()
+
+
+MOL_TYPE_PROTEIN = 0
+MOL_TYPE_DNA = 1
+MOL_TYPE_RNA = 2
+MOL_TYPE_NONPOLYMER = 3
+
+# The record order is part of the checkpoint input contract. Residue indices
+# start at two because zero and one are reserved by the model feature schema.
+_PROTEIN_SCHEMA = tuple(
+    tuple(record.split(":"))
+    for record in (
+        "ALA:A:N CA C O CB",
+        "ARG:R:N CA C O CB CG CD NE CZ NH1 NH2",
+        "ASN:N:N CA C O CB CG OD1 ND2",
+        "ASP:D:N CA C O CB CG OD1 OD2",
+        "CYS:C:N CA C O CB SG",
+        "GLN:Q:N CA C O CB CG CD OE1 NE2",
+        "GLU:E:N CA C O CB CG CD OE1 OE2",
+        "GLY:G:N CA C O",
+        "HIS:H:N CA C O CB CG ND1 CD2 CE1 NE2",
+        "ILE:I:N CA C O CB CG1 CG2 CD1",
+        "LEU:L:N CA C O CB CG CD1 CD2",
+        "LYS:K:N CA C O CB CG CD CE NZ",
+        "MET:M:N CA C O CB CG SD CE",
+        "PHE:F:N CA C O CB CG CD1 CD2 CE1 CE2 CZ",
+        "PRO:P:N CA C O CB CG CD",
+        "SER:S:N CA C O CB OG",
+        "THR:T:N CA C O CB OG1 CG2",
+        "TRP:W:N CA C O CB CG CD1 CD2 NE1 CE2 CE3 CZ2 CZ3 CH2",
+        "TYR:Y:N CA C O CB CG CD1 CD2 CE1 CE2 CZ OH",
+        "VAL:V:N CA C O CB CG1 CG2",
+    )
+)
+
+PROTEIN_RESIDUE_TO_RES_TYPE = {
+    residue: index for index, (residue, _letter, _atoms) in enumerate(_PROTEIN_SCHEMA, 2)
+}
+PROTEIN_RESIDUE_TO_RES_TYPE["MSE"] = PROTEIN_RESIDUE_TO_RES_TYPE["MET"]
+PROTEIN_UNK_RES_TYPE = 22
+
+RNA_RESIDUE_TO_RES_TYPE = dict(zip("AGCU", range(23, 27), strict=True))
+RNA_UNK_RES_TYPE = 27
+DNA_RESIDUE_TO_RES_TYPE = dict(zip(("DA", "DG", "DC", "DT"), range(28, 32), strict=True))
+DNA_UNK_RES_TYPE = 32
+GAP_RES_TYPE = DNA_UNK_RES_TYPE
+
+PROTEIN_3TO1 = {residue: letter for residue, letter, _atoms in _PROTEIN_SCHEMA}
+PROTEIN_3TO1["MSE"] = "M"
+PROTEIN_1TO3 = {letter: residue for residue, letter, _atoms in _PROTEIN_SCHEMA}
+PROTEIN_1TO3["X"] = "UNK"
+DNA_1TO3 = dict(zip("ATCG", ("DA", "DT", "DC", "DG"), strict=True))
+RNA_1TO3 = {letter: letter for letter in "AUCG"}
+
+_ESM_RESIDUE_ORDER = "LAGVSERTIDPKQNFYM HWC".replace(" ", "")
+ESM_PROTEIN_VOCAB = {residue: token_id for token_id, residue in enumerate(_ESM_RESIDUE_ORDER, 4)}
+ESM_PROTEIN_VOCAB["X"] = 3
+DNA_RNA_LIGAND_INPUT_ID = 24
+MSA_PAD_TOKEN_ID = 0
+MSA_GAP_TOKEN_ID = 1
+
+RES_TYPE_TO_CCD = {
+    **{
+        index: residue for residue, index in PROTEIN_RESIDUE_TO_RES_TYPE.items() if residue != "MSE"
+    },
+    22: "UNK",
+    **dict(zip(range(23, 28), ("A", "G", "C", "U", "N"), strict=True)),
+    **dict(zip(range(28, 33), ("DA", "DG", "DC", "DT", "DN"), strict=True)),
+}
+
+_CHARGE_SCHEMA = _words(
+    "LYS:NZ:1 ARG:NH2:1 HIS:ND1:1 PO4:O2:-1 PO4:O3:-1 PO4:O4:-1 "
+    "SO4:O3:-1 SO4:O4:-1 MG:MG:2 ZN:ZN:2 CA:CA:2 FE2:FE:2 MN:MN:2 "
+    "CO:CO:2 NCO:CO:3 CU:CU:2 NI:NI:2 K:K:1 NA:NA:1 CD:CD:2 CL:CL:-1 "
+    "ACT:OXT:-1 NAD:O2N:-1 NAD:N1N:1 NAP:O2N:-1 NAP:N1N:1 IMD:N3:1 "
+    "SAM:SD:1 FE:FE:3 A1BH3:N3:1"
+)
+CHARGED_ATOMS = {
+    (component, atom): int(charge)
+    for component, atom, charge in (record.split(":") for record in _CHARGE_SCHEMA)
+}
+
+_PERIODIC_SYMBOLS = _words(
+    "H HE LI BE B C N O F NE NA MG AL SI P S CL AR K CA SC TI V CR MN FE CO NI CU ZN "
+    "GA GE AS SE BR KR RB SR Y ZR NB MO TC RU RH PD AG CD IN SN SB TE I XE CS BA LA CE "
+    "PR ND PM SM EU GD TB DY HO ER TM YB LU HF TA W RE OS IR PT AU HG TL PB BI PO AT RN "
+    "FR RA AC TH PA U"
+)
+ELEMENT_TO_ATOMIC_NUM = {
+    symbol: atomic_number
+    for atomic_number, symbol in enumerate(_PERIODIC_SYMBOLS, 1)
+    if symbol != "HE"
+}
+ELEMENT_NUMBER_TO_SYMBOL = {
+    atomic_number: symbol for symbol, atomic_number in ELEMENT_TO_ATOMIC_NUM.items()
+}
+
+PROTEIN_HEAVY_ATOMS = {
+    residue: atom_string.split() for residue, _letter, atom_string in _PROTEIN_SCHEMA
+}
+PROTEIN_HEAVY_ATOMS["MSE"] = PROTEIN_HEAVY_ATOMS["MET"].copy()
+PROTEIN_HEAVY_ATOMS["UNK"] = _words("N CA C O")
+
+DNA_BACKBONE_ATOMS = _words("P OP1 OP2 O5' C5' C4' O4' C3' O3' C2' C1'")
+RNA_BACKBONE_ATOMS = _words("P OP1 OP2 O5' C5' C4' O4' C3' O3' C2' O2' C1'")
+_NUCLEOBASE_ATOMS = {
+    "A": _words("N9 C8 N7 C5 C6 N6 N1 C2 N3 C4"),
+    "G": _words("N9 C8 N7 C5 C6 O6 N1 C2 N2 N3 C4"),
+    "C": _words("N1 C2 O2 N3 C4 N4 C5 C6"),
+    "U": _words("N1 C2 O2 N3 C4 O4 C5 C6"),
+    "T": _words("N1 C2 O2 N3 C4 O4 C5 C7 C6"),
+}
+DNA_HEAVY_ATOMS = {
+    "DA": DNA_BACKBONE_ATOMS + _NUCLEOBASE_ATOMS["A"],
+    "DG": DNA_BACKBONE_ATOMS + _NUCLEOBASE_ATOMS["G"],
+    "DC": DNA_BACKBONE_ATOMS + _NUCLEOBASE_ATOMS["C"],
+    "DT": DNA_BACKBONE_ATOMS + _NUCLEOBASE_ATOMS["T"],
+}
+RNA_HEAVY_ATOMS = {residue: RNA_BACKBONE_ATOMS + _NUCLEOBASE_ATOMS[residue] for residue in "AGCU"}
+
+
+def _validate_schema() -> None:
+    if sorted(set(PROTEIN_RESIDUE_TO_RES_TYPE.values())) != list(range(2, 22)):
+        raise RuntimeError("Protein residue indices must cover the checkpoint interval 2..21.")
+    if len(_ESM_RESIDUE_ORDER) != 20 or len(set(_ESM_RESIDUE_ORDER)) != 20:
+        raise RuntimeError("The ESM residue vocabulary must contain 20 canonical residues.")
+    if RES_TYPE_TO_CCD[14] != "MET" or PROTEIN_RESIDUE_TO_RES_TYPE["MSE"] != 14:
+        raise RuntimeError("Selenomethionine must share the methionine residue index.")
+    if ELEMENT_TO_ATOMIC_NUM.get("U") != 92 or 2 in ELEMENT_NUMBER_TO_SYMBOL:
+        raise RuntimeError("The element schema must preserve the training-time atomic-number map.")
+    if set(DNA_HEAVY_ATOMS) != {"DA", "DG", "DC", "DT"}:
+        raise RuntimeError("The DNA atom schema is incomplete.")
+
+
+_validate_schema()
+
+__all__ = [name for name in globals() if name.isupper()]

fastplms/models/esmfold2/esmfold2_constants_esm3.py

@@ -0,0 +1,98 @@
+"""Token schemas needed by the ESMC encoder inside ESMFold2.
+
+The values implement the published Biohub ESM sequence-token contract pinned by
+``models.toml``. They are generated from ordered schemas so token position and
+special-token relationships are explicit and independently testable. This
+module performs no downloads and resolves no model assets at import time.
+"""
+
+from __future__ import annotations
+
+from types import MappingProxyType
+
+
+def _words(value: str) -> list[str]:
+    return value.split()
+
+
+SEQUENCE_VOCAB = _words(
+    "<cls> <pad> <eos> <unk> L A G V S E R T I D P K Q N F Y M H W C X B U Z O . - | <mask>"
+)
+
+_sequence_token_ids = MappingProxyType({token: index for index, token in enumerate(SEQUENCE_VOCAB)})
+SEQUENCE_BOS_TOKEN = _sequence_token_ids["<cls>"]
+SEQUENCE_PAD_TOKEN = _sequence_token_ids["<pad>"]
+SEQUENCE_EOS_TOKEN = _sequence_token_ids["<eos>"]
+SEQUENCE_CHAINBREAK_TOKEN = _sequence_token_ids["|"]
+SEQUENCE_MASK_TOKEN = _sequence_token_ids["<mask>"]
+SEQUENCE_STANDARD_AA_MIN_TOKEN = _sequence_token_ids["L"]
+SEQUENCE_STANDARD_AA_MAX_TOKEN = _sequence_token_ids["X"]
+
+VQVAE_CODEBOOK_SIZE = 4096
+VQVAE_SPECIAL_TOKENS = {
+    name: VQVAE_CODEBOOK_SIZE + offset
+    for offset, name in enumerate(("MASK", "EOS", "BOS", "PAD", "CHAINBREAK"))
+}
+VQVAE_DIRECTION_LOSS_BINS = 16
+VQVAE_PAE_BINS = 64
+VQVAE_MAX_PAE_BIN = 31.0
+VQVAE_PLDDT_BINS = 50
+
+STRUCTURE_MASK_TOKEN = VQVAE_SPECIAL_TOKENS["MASK"]
+STRUCTURE_EOS_TOKEN = VQVAE_SPECIAL_TOKENS["EOS"]
+STRUCTURE_BOS_TOKEN = VQVAE_SPECIAL_TOKENS["BOS"]
+STRUCTURE_PAD_TOKEN = VQVAE_SPECIAL_TOKENS["PAD"]
+STRUCTURE_CHAINBREAK_TOKEN = VQVAE_SPECIAL_TOKENS["CHAINBREAK"]
+STRUCTURE_UNDEFINED_TOKEN = 955
+
+SASA_PAD_TOKEN = 0
+SS8_PAD_TOKEN = 0
+INTERPRO_PAD_TOKEN = 0
+RESIDUE_PAD_TOKEN = 0
+
+CHAIN_BREAK_STR = "|"
+SEQUENCE_BOS_STR = "<cls>"
+SEQUENCE_EOS_STR = "<eos>"
+MASK_STR_SHORT = "_"
+SEQUENCE_MASK_STR = "<mask>"
+SASA_MASK_STR = "<unk>"
+SS8_MASK_STR = "<unk>"
+
+SSE_8CLASS_VOCAB = "GHITEBSC"
+SSE_3CLASS_VOCAB = "HEC"
+SSE_8CLASS_TO_3CLASS_MAP = dict(zip(SSE_8CLASS_VOCAB, "HHHCEECC", strict=True))
+
+SASA_DISCRETIZATION_BOUNDARIES = [
+    0.8,
+    4.0,
+    9.6,
+    16.4,
+    24.5,
+    32.9,
+    42.0,
+    51.5,
+    61.2,
+    70.9,
+    81.6,
+    93.3,
+    107.2,
+    125.4,
+    151.4,
+]
+MAX_RESIDUE_ANNOTATIONS = 16
+TFIDF_VECTOR_SIZE = 58_641
+FUNCTION_TOKENS_DEPTH = 8
+
+
+def _validate_schema() -> None:
+    if len(SEQUENCE_VOCAB) != len(set(SEQUENCE_VOCAB)):
+        raise RuntimeError("The ESM sequence vocabulary contains duplicate tokens.")
+    if SEQUENCE_STANDARD_AA_MAX_TOKEN - SEQUENCE_STANDARD_AA_MIN_TOKEN != 20:
+        raise RuntimeError("The canonical residue interval must contain 20 tokens.")
+    if tuple(VQVAE_SPECIAL_TOKENS.values()) != tuple(range(4096, 4101)):
+        raise RuntimeError("The structure special-token interval is not contiguous.")
+
+
+_validate_schema()
+
+__all__ = [name for name in globals() if name.isupper()]

fastplms/models/esmfold2/esmfold2_input_builder.py

@@ -0,0 +1,244 @@
+"""Typed, JSON-safe inputs for ESMFold2 feature preparation."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Any, TypeAlias
+
+import numpy as np
+
+from .esmfold2_msa import MSA
+
+MSAInput: TypeAlias = MSA | None
+
+
+@dataclass
+class Modification:
+    """A zero-indexed residue substitution using a CCD component."""
+
+    position: int
+    ccd: str
+    smiles: str | None = None
+
+
+@dataclass
+class ProteinInput:
+    id: str | list[str]
+    sequence: str
+    modifications: list[Modification] | None = None
+    msa: MSAInput = None
+
+
+@dataclass
+class RNAInput:
+    id: str | list[str]
+    sequence: str
+    modifications: list[Modification] | None = None
+
+
+@dataclass
+class DNAInput:
+    id: str | list[str]
+    sequence: str
+    modifications: list[Modification] | None = None
+
+
+@dataclass
+class LigandInput:
+    id: str | list[str]
+    smiles: str | None = None
+    ccd: list[str] | None = None
+
+
+@dataclass
+class DistogramConditioning:
+    chain_id: str
+    distogram: np.ndarray
+
+
+@dataclass
+class PocketConditioning:
+    binder_chain_id: str
+    contacts: list[tuple[str, int]]
+
+
+@dataclass
+class CovalentBond:
+    chain_id1: str
+    res_idx1: int
+    atom_idx1: int
+    chain_id2: str
+    res_idx2: int
+    atom_idx2: int
+
+
+SequenceInput: TypeAlias = ProteinInput | RNAInput | DNAInput | LigandInput
+
+
+@dataclass
+class StructurePredictionInput:
+    sequences: Sequence[SequenceInput]
+    pocket: PocketConditioning | None = None
+    distogram_conditioning: list[DistogramConditioning] | None = None
+    covalent_bonds: list[CovalentBond] | None = None
+
+
+_CHAIN_TYPE = {
+    ProteinInput: "protein",
+    RNAInput: "rna",
+    DNAInput: "dna",
+}
+
+
+def _serialize_modifications(
+    modifications: list[Modification] | None,
+) -> list[dict[str, Any]] | None:
+    if not modifications:
+        return None
+    return [{"position": item.position, "ccd": item.ccd} for item in modifications]
+
+
+def _serialize_chain(chain: SequenceInput) -> dict[str, Any]:
+    if isinstance(chain, LigandInput):
+        return {
+            "smiles": chain.smiles,
+            "id": chain.id,
+            "ccd": chain.ccd,
+            "type": "ligand",
+        }
+
+    chain_type = _CHAIN_TYPE.get(type(chain))
+    if chain_type is None:
+        raise ValueError(f"Unsupported sequence input type: {type(chain)}")
+    serialized: dict[str, Any] = {
+        "sequence": chain.sequence,
+        "id": chain.id,
+        "type": chain_type,
+    }
+    if modifications := _serialize_modifications(chain.modifications):
+        serialized["modifications"] = modifications
+    if isinstance(chain, ProteinInput):
+        if chain.msa is not None and not isinstance(chain.msa, MSA):
+            raise AttributeError(f"MSA must be None or MSA. Got {chain.msa} instead.")
+        serialized["msa"] = None if chain.msa is None else {"sequences": chain.msa.sequences}
+    return serialized
+
+
+def serialize_structure_prediction_input(
+    structure_input: StructurePredictionInput,
+) -> dict[str, Any]:
+    """Convert an input object to a JSON-safe mapping."""
+
+    serialized: dict[str, Any] = {
+        "sequences": [_serialize_chain(chain) for chain in structure_input.sequences]
+    }
+    if structure_input.covalent_bonds is not None:
+        serialized["covalent_bonds"] = [
+            vars(bond).copy() for bond in structure_input.covalent_bonds
+        ]
+    if structure_input.pocket is not None:
+        serialized["pocket"] = {
+            "binder_chain_id": structure_input.pocket.binder_chain_id,
+            "contacts": structure_input.pocket.contacts,
+        }
+    if structure_input.distogram_conditioning is not None:
+        serialized["distogram_conditioning"] = [
+            {"chain_id": item.chain_id, "distogram": item.distogram.tolist()}
+            for item in structure_input.distogram_conditioning
+        ]
+    return serialized
+
+
+def _deserialize_modifications(chain: dict[str, Any]) -> list[Modification] | None:
+    raw = chain.get("modifications")
+    if not raw:
+        return None
+    return [Modification(position=item["position"], ccd=item["ccd"]) for item in raw]
+
+
+def _deserialize_msa(chain: dict[str, Any]) -> MSAInput:
+    raw = chain.get("msa")
+    if raw is None:
+        return None
+    if not isinstance(raw, dict) or not isinstance(raw.get("sequences"), list):
+        raise ValueError(f"Unexpected MSA value: {raw!r}")
+    return MSA.from_sequences(raw["sequences"])
+
+
+def _deserialize_chain(chain: dict[str, Any]) -> SequenceInput:
+    chain_type = chain.get("type")
+    common = {"id": chain["id"]}
+    if chain_type == "protein":
+        return ProteinInput(
+            **common,
+            sequence=chain["sequence"],
+            modifications=_deserialize_modifications(chain),
+            msa=_deserialize_msa(chain),
+        )
+    if chain_type == "rna":
+        return RNAInput(
+            **common,
+            sequence=chain["sequence"],
+            modifications=_deserialize_modifications(chain),
+        )
+    if chain_type == "dna":
+        return DNAInput(
+            **common,
+            sequence=chain["sequence"],
+            modifications=_deserialize_modifications(chain),
+        )
+    if chain_type == "ligand":
+        return LigandInput(**common, smiles=chain.get("smiles"), ccd=chain.get("ccd"))
+    raise ValueError(f"Unsupported sequence type: {chain_type!r}")
+
+
+def deserialize_structure_prediction_input(data: dict[str, Any]) -> StructurePredictionInput:
+    """Reconstruct the typed input represented by a serialized mapping."""
+
+    pocket_data = data.get("pocket")
+    pocket = None
+    if pocket_data is not None:
+        pocket = PocketConditioning(
+            binder_chain_id=pocket_data["binder_chain_id"],
+            contacts=[tuple(contact) for contact in pocket_data["contacts"]],
+        )
+
+    distogram_data = data.get("distogram_conditioning")
+    distograms = None
+    if distogram_data is not None:
+        distograms = [
+            DistogramConditioning(
+                chain_id=item["chain_id"], distogram=np.asarray(item["distogram"])
+            )
+            for item in distogram_data
+        ]
+
+    bond_data = data.get("covalent_bonds")
+    bonds = None
+    if bond_data is not None:
+        bonds = [CovalentBond(**item) for item in bond_data]
+
+    return StructurePredictionInput(
+        sequences=[_deserialize_chain(chain) for chain in data["sequences"]],
+        pocket=pocket,
+        distogram_conditioning=distograms,
+        covalent_bonds=bonds,
+    )
+
+
+__all__ = [
+    "CovalentBond",
+    "DNAInput",
+    "DistogramConditioning",
+    "LigandInput",
+    "MSAInput",
+    "Modification",
+    "PocketConditioning",
+    "ProteinInput",
+    "RNAInput",
+    "SequenceInput",
+    "StructurePredictionInput",
+    "deserialize_structure_prediction_input",
+    "serialize_structure_prediction_input",
+]

fastplms/models/esmfold2/esmfold2_metrics.py

@@ -0,0 +1,235 @@
+"""Contact, lDDT, RMSD, and GDT-TS metrics for structure validation."""
+
+from __future__ import annotations
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torch.amp import autocast  # type: ignore
+
+from . import esmfold2_residue_constants as residue_constants
+from .esmfold2_misc import binpack, unbinpack
+from .esmfold2_protein_structure import (
+    compute_alignment_tensors,
+    compute_gdt_ts_no_alignment,
+    compute_rmsd_no_alignment,
+)
+
+
+def _distance_matrix(positions: Tensor, eps: float) -> Tensor:
+    displacement = positions[..., None, :] - positions[..., None, :, :]
+    return torch.sqrt(eps + torch.sum(displacement**2, dim=-1))
+
+
+def compute_lddt_from_dmat(
+    dmat_pred: Tensor,
+    dmat_true: Tensor,
+    pairwise_mask: Tensor,
+    cutoff: float | Tensor = 15.0,
+    eps: float = 1e-10,
+    per_residue: bool = True,
+) -> Tensor:
+    """Score distance matrices ``D_pred`` and ``D_true`` with shape (..., l, l)."""
+
+    sequence_length = dmat_true.size(-1)
+    identity = torch.eye(sequence_length, device=dmat_true.device)
+    scored_pairs = (dmat_true < cutoff) * pairwise_mask * (1.0 - identity)
+    absolute_error = torch.abs(dmat_true - dmat_pred)
+    score = (
+        (absolute_error < 0.5).type(absolute_error.dtype)
+        + (absolute_error < 1.0).type(absolute_error.dtype)
+        + (absolute_error < 2.0).type(absolute_error.dtype)
+        + (absolute_error < 4.0).type(absolute_error.dtype)
+    ) * 0.25
+    dimensions = (-1,) if per_residue else (-2, -1)
+    normalization = 1.0 / (eps + scored_pairs.sum(dim=dimensions))
+    return normalization * (eps + (scored_pairs * score).sum(dim=dimensions))
+
+
+def compute_lddt(
+    all_atom_pred_pos: Tensor,
+    all_atom_positions: Tensor,
+    all_atom_mask: Tensor,
+    pairwise_all_atom_mask: Tensor | None = None,
+    cutoff: float | Tensor = 15.0,
+    eps: float = 1e-10,
+    per_residue: bool = True,
+    sequence_id: Tensor | None = None,
+) -> Tensor:
+    """Compute lDDT from coordinate tensors and atom masks."""
+
+    expanded_mask = all_atom_mask[..., None]
+    true_distances = _distance_matrix(all_atom_positions, eps)
+    predicted_distances = _distance_matrix(all_atom_pred_pos, eps)
+    pair_mask = expanded_mask * expanded_mask.transpose(-2, -1)
+    if pairwise_all_atom_mask is not None:
+        pair_mask = pair_mask * pairwise_all_atom_mask
+    if sequence_id is not None:
+        same_sequence = sequence_id[..., None] == sequence_id[..., None, :]
+        pair_mask = pair_mask * same_sequence.type_as(pair_mask)
+    return compute_lddt_from_dmat(
+        predicted_distances,
+        true_distances,
+        pair_mask,
+        cutoff=cutoff,
+        eps=eps,
+        per_residue=per_residue,
+    )
+
+
+def compute_lddt_ca(
+    all_atom_pred_pos: Tensor,
+    all_atom_positions: Tensor,
+    all_atom_mask: Tensor,
+    cutoff: float = 15.0,
+    eps: float = 1e-10,
+    per_residue: bool = True,
+    sequence_id: Tensor | None = None,
+) -> Tensor:
+    """Compute lDDT using only C-alpha coordinates."""
+
+    ca_index = residue_constants.atom_order["CA"]
+    predicted_ca = (
+        all_atom_pred_pos if all_atom_pred_pos.dim() == 3 else all_atom_pred_pos[..., ca_index, :]
+    )
+    return compute_lddt(
+        predicted_ca,
+        all_atom_positions[..., ca_index, :],
+        all_atom_mask[..., ca_index],
+        cutoff=cutoff,
+        eps=eps,
+        per_residue=per_residue,
+        sequence_id=sequence_id,
+    )
+
+
+@torch.no_grad()
+@autocast("cuda", enabled=False)
+def compute_rmsd(
+    mobile: Tensor,
+    target: Tensor,
+    atom_exists_mask: Tensor | None = None,
+    sequence_id: Tensor | None = None,
+    reduction: str = "batch",
+) -> Tensor:
+    """Align ``X`` to ``Y`` and compute RMSD."""
+
+    centered_mobile, _, centered_target, _, rotation, counts = compute_alignment_tensors(
+        mobile,
+        target,
+        atom_exists_mask,
+        sequence_id,
+    )
+    rmsd = compute_rmsd_no_alignment(
+        torch.matmul(centered_mobile, rotation),
+        centered_target,
+        counts,
+        reduction=reduction,
+    )
+    if reduction == "per_residue" and sequence_id is not None:
+        return binpack(rmsd, sequence_id, pad_value=0)
+    return rmsd
+
+
+def compute_gdt_ts(
+    mobile: Tensor,
+    target: Tensor,
+    atom_exists_mask: Tensor | None = None,
+    sequence_id: Tensor | None = None,
+    reduction: str = "per_sample",
+) -> Tensor:
+    """Align ``X`` to ``Y`` and compute GDT-TS."""
+
+    if atom_exists_mask is None:
+        atom_exists_mask = torch.isfinite(target).all(dim=-1)
+    centered_mobile, _, centered_target, _, rotation, _ = compute_alignment_tensors(
+        mobile,
+        target,
+        atom_exists_mask,
+        sequence_id,
+    )
+    if sequence_id is not None:
+        atom_exists_mask = unbinpack(atom_exists_mask, sequence_id, pad_value=False)
+    return compute_gdt_ts_no_alignment(
+        torch.matmul(centered_mobile, rotation),
+        centered_target,
+        atom_exists_mask,
+        reduction,
+    )
+
+
+def _batched_contacts(predictions: Tensor, targets: Tensor) -> tuple[Tensor, Tensor]:
+    if predictions.dim() == 2:
+        predictions = predictions.unsqueeze(0)
+    if targets.dim() == 2:
+        targets = targets.unsqueeze(0)
+    if predictions.size() != targets.size():
+        raise ValueError(
+            f"Size mismatch. Received predictions of size {predictions.size()}, "
+            f"targets of size {targets.size()}"
+        )
+    return predictions, targets
+
+
+def _valid_contact_mask(
+    targets: Tensor,
+    src_lengths: Tensor,
+    minsep: int,
+    maxsep: int | None,
+) -> Tensor:
+    sequence_length = targets.shape[-1]
+    positions = torch.arange(sequence_length, device=targets.device)
+    separation = (positions.unsqueeze(0) - positions.unsqueeze(1)).unsqueeze(0)
+    valid = (separation >= minsep) & (targets >= 0)
+    if maxsep is not None:
+        valid &= separation < maxsep
+    within_length = positions.unsqueeze(0) < src_lengths.unsqueeze(1)
+    return valid & within_length.unsqueeze(1) & within_length.unsqueeze(2)
+
+
+def contact_precision(
+    predictions: Tensor,
+    targets: Tensor,
+    src_lengths: Tensor | None = None,
+    minsep: int = 6,
+    maxsep: int | None = None,
+    override_length: int | None = None,
+) -> dict[str, Tensor]:
+    """Compute P@L, P@L/5, and binned area for contact probabilities."""
+
+    predictions, targets = _batched_contacts(predictions, targets)
+    batch_size, sequence_length, _ = predictions.shape
+    if src_lengths is None:
+        src_lengths = torch.full(
+            (batch_size,),
+            sequence_length,
+            dtype=torch.long,
+            device=predictions.device,
+        )
+    valid = _valid_contact_mask(targets, src_lengths, minsep, maxsep)
+    masked_predictions = predictions.masked_fill(~valid, float("-inf"))
+    row_index, column_index = np.triu_indices(sequence_length, minsep)
+    upper_predictions = masked_predictions[:, row_index, column_index]
+    upper_targets = targets[:, row_index, column_index]
+
+    topk = sequence_length if override_length is None else max(sequence_length, override_length)
+    ranked_indices = upper_predictions.argsort(dim=-1, descending=True)[:, :topk]
+    batch_indices = torch.arange(batch_size, device=ranked_indices.device).unsqueeze(1)
+    ranked_targets = upper_targets[batch_indices, ranked_indices]
+    if ranked_targets.size(1) < topk:
+        ranked_targets = F.pad(ranked_targets, [0, topk - ranked_targets.size(1)])
+    cumulative_contacts = ranked_targets.type_as(predictions).cumsum(dim=-1)
+
+    gather_lengths = src_lengths.unsqueeze(1)
+    if override_length is not None:
+        gather_lengths = override_length * torch.ones_like(gather_lengths)
+    fractions = torch.arange(0.1, 1.1, 0.1, device=predictions.device).unsqueeze(0)
+    gather_indices = (fractions * gather_lengths).type(torch.long).sub(1).clamp_min(0)
+    cumulative_bins = cumulative_contacts.gather(1, gather_indices)
+    precisions = cumulative_bins / (gather_indices + 1).type_as(cumulative_bins)
+    return {
+        "AUC": precisions.mean(dim=-1),
+        "P@L": precisions[:, 9],
+        "P@L5": precisions[:, 1],
+    }

fastplms/models/esmfold2/esmfold2_misc.py

@@ -0,0 +1,400 @@
+"""Small tensor, sequence, and annotation utilities used by ESMFold2.
+
+The helpers in this module are deliberately free of model state. Importing the
+module therefore performs no device selection, compilation, or remote access.
+"""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from collections.abc import Generator, Iterable, Sequence
+from contextlib import AbstractContextManager, nullcontext
+from dataclasses import is_dataclass
+from io import BytesIO
+from typing import Any, Protocol, TypeVar, runtime_checkable
+from warnings import warn
+
+import numpy as np
+import torch
+import zstandard
+
+from .esmfold2_constants_esm3 import CHAIN_BREAK_STR
+from .esmfold2_utils_types import FunctionAnnotation
+
+MAX_SUPPORTED_DISTANCE = 1e6
+
+TSequence = TypeVar("TSequence", bound=Sequence)
+
+
+@runtime_checkable
+class Concatable(Protocol):
+    """Protocol for sequence-like records with a class-level concatenator."""
+
+    @classmethod
+    def concat(cls, objs: list[Concatable]) -> Concatable: ...
+
+
+def fp32_autocast_context(
+    device_type: str,
+) -> AbstractContextManager[Any]:  # type: ignore
+    """Return a context that keeps numerically sensitive work in FP32."""
+
+    if device_type == "mps":
+        return nullcontext()
+    if device_type == "cpu":
+        return torch.amp.autocast(device_type, enabled=False)  # type: ignore
+    if device_type == "cuda":
+        return torch.amp.autocast(device_type, dtype=torch.float32)  # type: ignore
+    raise ValueError(f"Unsupported device type: {device_type}")
+
+
+def maybe_tensor(value, convert_none_to_nan: bool = False) -> torch.Tensor | None:
+    """Convert an optional array-like value to a tensor."""
+
+    if value is None:
+        return None
+    if isinstance(value, torch.Tensor):
+        return value
+    if isinstance(value, list) and all(isinstance(element, torch.Tensor) for element in value):
+        return torch.stack(value)
+    if convert_none_to_nan:
+        value = np.asarray(value, dtype=np.float32)
+        value = np.where(value is None, np.nan, value)
+    return torch.tensor(value)
+
+
+def maybe_list(value, convert_nan_to_none: bool = False) -> list | None:
+    """Convert an optional tensor or NumPy array to nested Python lists."""
+
+    if value is None:
+        return None
+    if not convert_nan_to_none:
+        return value.tolist()
+    if isinstance(value, torch.Tensor):
+        nan_mask = torch.isnan(value).cpu().numpy()
+        array = value.cpu().numpy().astype(object)
+    elif isinstance(value, np.ndarray):
+        nan_mask = np.isnan(value)
+        array = value.astype(object)
+    else:
+        raise TypeError("maybe_list can only work with torch.tensor or np.ndarray.")
+    array[nan_mask] = None
+    return array.tolist()
+
+
+def replace_inf(data):
+    """Replace infinite array values by the ESM API sentinel value."""
+
+    if data is None:
+        return None
+    array = np.asarray(data, dtype=np.float32)
+    return np.where(np.isinf(array), 1000, array).tolist()
+
+
+def slice_python_object_as_numpy(
+    obj: TSequence,
+    idx: int | list[int] | slice | np.ndarray,
+) -> TSequence:
+    """Apply NumPy-style scalar, mask, or index-array slicing to Python data."""
+
+    normalized_idx: list[int] | slice | np.ndarray = (
+        [int(idx)] if np.isscalar(idx) else idx  # type: ignore[arg-type]
+    )
+
+    if isinstance(normalized_idx, np.ndarray) and normalized_idx.dtype == bool:
+        selected = [obj[position] for position in np.flatnonzero(normalized_idx)]
+    elif isinstance(normalized_idx, slice):
+        selected = obj[normalized_idx]
+    else:
+        selected = [obj[position] for position in normalized_idx]
+
+    if isinstance(obj, str) and isinstance(selected, list):
+        return "".join(selected)  # type: ignore[return-value]
+    return obj.__class__(selected)  # type: ignore[call-arg,return-value]
+
+
+def slice_any_object(
+    obj: TSequence,
+    idx: int | list[int] | slice | np.ndarray,
+) -> TSequence:
+    """Slice tensors, arrays, dataclasses, and ordinary Python sequences."""
+
+    if isinstance(obj, (np.ndarray, torch.Tensor)) or is_dataclass(obj):
+        return obj[idx]  # type: ignore[index,return-value]
+    return slice_python_object_as_numpy(obj, idx)
+
+
+def join_lists(
+    lists: Sequence[Sequence[Any]],
+    separator: Sequence[Any] | None = None,
+) -> list[Any]:
+    """Join lists, inserting all elements of ``separator`` between inputs."""
+
+    if len(lists) == 0:
+        return []
+    joined = list(lists[0])
+    for values in lists[1:]:
+        if separator:
+            joined.extend(separator)
+        joined.extend(values)
+    return joined
+
+
+def iterate_with_intermediate(
+    lists: Iterable,
+    intermediate,
+) -> Generator[Any, None, None]:
+    """Yield an intermediate value between consecutive input values."""
+
+    iterator = iter(lists)
+    yield next(iterator)
+    for value in iterator:
+        yield intermediate
+        yield value
+
+
+def concat_objects(objs: Sequence[Any], separator: Any | None = None):
+    """Concatenate one supported homogeneous collection."""
+
+    if not objs:
+        raise ValueError("objs must contain at least one value.")
+    first = objs[0]
+    if isinstance(first, Concatable):
+        return first.__class__.concat(objs)
+    if isinstance(first, str):
+        if not isinstance(separator, str):
+            raise TypeError("separator must be a string when joining strings.")
+        return separator.join(objs)
+    if isinstance(first, list):
+        return join_lists(objs, None if separator is None else [separator])
+    if isinstance(first, np.ndarray):
+        pieces = (
+            objs
+            if separator is None
+            else list(iterate_with_intermediate(objs, np.array([separator])))
+        )
+        return np.concatenate(pieces)
+    if isinstance(first, torch.Tensor):
+        pieces = (
+            objs
+            if separator is None
+            else list(iterate_with_intermediate(objs, torch.tensor([separator])))
+        )
+        return torch.cat(pieces)  # type: ignore[arg-type]
+    raise TypeError(type(first))
+
+
+def rbf(values, v_min, v_max, n_bins=16):
+    """Encode values against evenly spaced radial basis centers."""
+
+    centers = torch.linspace(
+        v_min,
+        v_max,
+        n_bins,
+        dtype=values.dtype,
+        device=values.device,
+    )
+    centers = centers.reshape((1,) * values.ndim + (-1,))
+    standardized = (values.unsqueeze(-1) - centers) / ((v_max - v_min) / n_bins)
+    return torch.exp(-(standardized**2))
+
+
+def batched_gather(data, inds, dim=0, no_batch_dims=0):
+    """Gather along one data dimension while retaining leading batch axes."""
+
+    batch_indices = []
+    index_rank = len(inds.shape)
+    for axis, size in enumerate(data.shape[:no_batch_dims]):
+        shape = (1,) * axis + (-1,) + (1,) * (index_rank - axis - 1)
+        batch_indices.append(torch.arange(size).view(*shape))
+    tail = [slice(None)] * (len(data.shape) - no_batch_dims)
+    tail[dim - no_batch_dims if dim >= 0 else dim] = inds
+    return data[tuple(batch_indices + tail)]
+
+
+def node_gather(s: torch.Tensor, edges: torch.Tensor) -> torch.Tensor:
+    """Gather node features for each row of an edge-index tensor."""
+
+    return batched_gather(
+        s.unsqueeze(-3),
+        edges,
+        -2,
+        no_batch_dims=len(s.shape) - 1,
+    )
+
+
+def knn_graph(
+    coords: torch.Tensor,
+    coord_mask: torch.Tensor,
+    padding_mask: torch.Tensor,
+    sequence_id: torch.Tensor,
+    *,
+    no_knn: int,
+):
+    """Build nearest-neighbor edges, using sequence distance for missing geometry."""
+
+    length = coords.shape[-2]
+    coords = coords.nan_to_num()
+    missing_pair = ~(coord_mask[..., None, :] & coord_mask[..., :, None])
+    excluded_pair = padding_mask[..., None, :] | padding_mask[..., :, None]
+    if sequence_id is not None:
+        excluded_pair |= sequence_id.unsqueeze(1) != sequence_id.unsqueeze(2)
+
+    distances = (coords.unsqueeze(-2) - coords.unsqueeze(-3)).norm(dim=-1)
+    residue_index = torch.arange(length, device=coords.device)
+    sequence_distance = (residue_index.unsqueeze(-1) - residue_index.unsqueeze(-2)).abs()
+    if not (distances[~missing_pair] < MAX_SUPPORTED_DISTANCE).all():
+        raise ValueError(
+            "Coordinate pairwise distances exceed max supported distance "
+            f"({MAX_SUPPORTED_DISTANCE}). "
+        )
+
+    rank_distance = sequence_distance.to(distances.dtype).mul(1e2).add(MAX_SUPPORTED_DISTANCE)
+    rank_distance = rank_distance.where(missing_pair, distances)
+    rank_distance = rank_distance.masked_fill(excluded_pair, torch.inf)
+    sorted_distance, sorted_edge = rank_distance.sort(dim=-1, descending=False)
+    width = min(no_knn, length)
+    return sorted_edge[..., :width], sorted_distance[..., :width].isfinite()
+
+
+def stack_variable_length_tensors(
+    sequences: Sequence[torch.Tensor],
+    constant_value: int | float = 0,
+    dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    """Pad arbitrary tensor dimensions to their maxima, then stack."""
+
+    output_shape = [
+        len(sequences),
+        *np.max([sequence.shape for sequence in sequences], axis=0).tolist(),
+    ]
+    output = torch.full(
+        output_shape,
+        constant_value,
+        dtype=sequences[0].dtype if dtype is None else dtype,
+        device=sequences[0].device,
+    )
+    for destination, source in zip(output, sequences, strict=True):
+        destination[tuple(slice(size) for size in source.shape)] = source
+    return output
+
+
+def binpack(
+    tensor: torch.Tensor,
+    sequence_id: torch.Tensor | None,
+    pad_value: int | float,
+):
+    """Scatter a sequence-major tensor into the packed layout described by IDs."""
+
+    if sequence_id is None:
+        return tensor
+    sequence_counts = sequence_id.max(dim=-1).values + 1
+    output = torch.full(
+        sequence_id.shape + tensor.shape[2:],
+        fill_value=pad_value,
+        dtype=tensor.dtype,
+        device=tensor.device,
+    )
+    source_index = 0
+    for batch_index, (batch_ids, count) in enumerate(
+        zip(sequence_id, sequence_counts, strict=True)
+    ):
+        for seqid in range(count):
+            selection = batch_ids == seqid
+            output[batch_index, selection] = tensor[source_index, : selection.sum()]
+            source_index += 1
+    return output
+
+
+def unbinpack(
+    tensor: torch.Tensor,
+    sequence_id: torch.Tensor | None,
+    pad_value: int | float,
+):
+    """Restore sequence-major rows from a packed tensor and its sequence IDs."""
+
+    if sequence_id is None:
+        return tensor
+    rows = []
+    sequence_counts = sequence_id.max(dim=-1).values + 1
+    for batch_index, (batch_ids, count) in enumerate(
+        zip(sequence_id, sequence_counts, strict=True)
+    ):
+        for seqid in range(count):
+            rows.append(tensor[batch_index, batch_ids == seqid])
+    return stack_variable_length_tensors(rows, pad_value)
+
+
+def merge_ranges(
+    ranges: list[range],
+    merge_gap_max: int | None = None,
+) -> list[range]:
+    """Merge overlapping or sufficiently close ranges in positional order."""
+
+    maximum_gap = 0 if merge_gap_max is None else merge_gap_max
+    if not isinstance(maximum_gap, int) or isinstance(maximum_gap, bool):
+        raise TypeError("merge_gap_max must be an integer or None.")
+    if maximum_gap < 0:
+        raise ValueError(f"merge_gap_max must be non-negative, got {maximum_gap}.")
+    merged: list[range] = []
+    for current in sorted(ranges, key=lambda item: item.start):
+        if not merged or merged[-1].stop + maximum_gap < current.start:
+            merged.append(current)
+            continue
+        previous = merged[-1]
+        merged[-1] = range(previous.start, max(previous.stop, current.stop))
+    return merged
+
+
+def merge_annotations(
+    annotations: list[FunctionAnnotation],
+    merge_gap_max: int | None = None,
+) -> list[FunctionAnnotation]:
+    """Merge overlapping annotations independently for each label."""
+
+    grouped: dict[str, list[range]] = defaultdict(list)
+    for annotation in annotations:
+        grouped[annotation.label].append(range(annotation.start, annotation.end + 1))
+    result = []
+    for label, spans in grouped.items():
+        result.extend(
+            FunctionAnnotation(label=label, start=span.start, end=span.stop - 1)
+            for span in merge_ranges(spans, merge_gap_max=merge_gap_max)
+        )
+    return result
+
+
+def get_chainbreak_boundaries_from_sequence(
+    sequence: Sequence[str],
+) -> np.ndarray:
+    """Return half-open chain intervals split by chain-break tokens."""
+
+    boundaries = [0]
+    final_index = len(sequence) - 1
+    for index, residue in enumerate(sequence):
+        if residue != CHAIN_BREAK_STR:
+            continue
+        if index == final_index:
+            raise ValueError(
+                "Encountered chain break token at end of sequence, this is unexpected."
+            )
+        if index == final_index - 1:
+            warn(
+                "Encountered chain break token at penultimate position, this is unexpected.",
+                stacklevel=2,
+            )
+        boundaries.extend((index, index + 1))
+    boundaries.append(len(sequence))
+    assert len(boundaries) % 2 == 0
+    return np.asarray(boundaries).reshape(-1, 2)
+
+
+def deserialize_tensors(data: bytes) -> Any:
+    """Decompress a tensor-only Torch payload onto CPU."""
+
+    decompressed = zstandard.ZstdDecompressor().decompress(data)
+    return torch.load(
+        BytesIO(decompressed),
+        map_location="cpu",
+        weights_only=True,
+    )

fastplms/models/esmfold2/esmfold2_mmcif_parsing.py

@@ -0,0 +1,469 @@
+"""Biotite-backed mmCIF parsing used by ESMFold2 structure records."""
+
+from __future__ import annotations
+
+import functools
+import io
+import os
+from contextlib import suppress
+from dataclasses import dataclass
+from datetime import datetime
+
+import biotite.structure as bs
+import biotite.structure.io.pdbx as pdbx
+import numpy as np
+from biotite.structure.io.pdbx import CIFColumn, CIFData, CIFFile
+
+from . import esmfold2_residue_constants as residue_constants
+
+PathOrBuffer = str | os.PathLike | io.StringIO
+
+PLDDT_B_FACTOR_SCALE = 100.0
+_MMCIF_COLUMN_DECIMALS = {
+    "Cartn_x": 3,
+    "Cartn_y": 3,
+    "Cartn_z": 3,
+    "B_iso_or_equiv": 2,
+}
+_NONPOLYMER_ENTITY_TYPES = frozenset({"NON-POLYMER", "WATER", "BRANCHED"})
+
+
+class NoProteinError(Exception):
+    """Raised internally when an mmCIF block contains no model-one atoms."""
+
+
+@dataclass
+class Residue:
+    residue_number: int | None = None
+    insertion_code: str = ""
+    hetflag: bool = False
+
+
+@dataclass
+class MmcifHeader:
+    release_date: datetime | None = None
+    resolution: float | None = None
+    structure_method: str = "UNKNOWN"
+
+
+def round_mmcif_columns(cif_file: CIFFile) -> None:
+    """Round coordinate and confidence columns in place for stable exports."""
+
+    if "atom_site" not in cif_file.block:
+        return
+    atom_site = cif_file.block["atom_site"]
+    for name, decimals in _MMCIF_COLUMN_DECIMALS.items():
+        if name not in atom_site:
+            continue
+        original = atom_site[name]
+        values = original.as_array(np.float64)
+        strings = np.asarray(
+            [f"{value:.{decimals}f}" for value in values],
+            dtype=np.str_,
+        )
+        atom_site[name] = CIFColumn(
+            data=CIFData(array=strings, dtype=np.str_),
+            mask=original.mask,
+        )
+
+
+def _clean_chain_list(value: str) -> list[str]:
+    return [chain.strip() for chain in value.split(",") if chain.strip()]
+
+
+def _empty_residue() -> Residue:
+    return Residue(residue_number=None, insertion_code="", hetflag=False)
+
+
+def _header_from_block(
+    block,
+    header: MmcifHeader | None = None,
+) -> MmcifHeader:
+    header = MmcifHeader() if header is None else header
+    try:
+        if "pdbx_database_status" in block:
+            category = block["pdbx_database_status"]
+            if "recvd_initial_deposition_date" in category:
+                value = category["recvd_initial_deposition_date"].as_item()
+                if value and value != "?":
+                    with suppress(ValueError):
+                        header.release_date = datetime.strptime(value, "%Y-%m-%d")
+        if "refine" in block:
+            category = block["refine"]
+            if "ls_d_res_high" in category:
+                value = category["ls_d_res_high"].as_item()
+                if value and value != "?":
+                    with suppress(ValueError):
+                        header.resolution = float(value)
+        if "exptl" in block:
+            category = block["exptl"]
+            if "method" in category:
+                value = category["method"].as_item()
+                if value and value != "?":
+                    header.structure_method = value.upper()
+    except Exception:
+        pass
+    return header
+
+
+def _entities_from_block(
+    block,
+    entities: dict[int, list[str]] | None = None,
+) -> dict[int, list[str]]:
+    entities = {} if entities is None else entities
+    if "entity" in block:
+        category = block["entity"]
+        ids = category["id"].as_array(str)
+        types = category["type"].as_array(str)
+        for entity_id, _ in zip(ids, types, strict=False):
+            entities[int(entity_id)] = []
+    if "entity_poly" in block:
+        category = block["entity_poly"]
+        ids = category["entity_id"].as_array(str)
+        chain_lists = category["pdbx_strand_id"].as_array(str)
+        for raw_id, raw_chains in zip(ids, chain_lists, strict=False):
+            entity_id = int(raw_id)
+            if entity_id in entities:
+                entities[entity_id] = _clean_chain_list(raw_chains)
+    if "struct_asym" in block:
+        category = block["struct_asym"]
+        asym_ids = category["id"].as_array(str)
+        entity_ids = category["entity_id"].as_array(str)
+        for asym_id, raw_id in zip(asym_ids, entity_ids, strict=False):
+            entity_id = int(raw_id)
+            if entity_id in entities and not entities[entity_id]:
+                entities[entity_id].append(asym_id)
+    return entities
+
+
+def _polymer_sequences(block) -> dict[str, str]:
+    sequences: dict[str, str] = {}
+    if "entity_poly" not in block:
+        return sequences
+    category = block["entity_poly"]
+    entity_ids = category["entity_id"].as_array(str)
+    raw_sequences = category["pdbx_seq_one_letter_code_can"].as_array(str)
+    chain_lists = category["pdbx_strand_id"].as_array(str)
+    for _, raw_sequence, raw_chains in zip(
+        entity_ids,
+        raw_sequences,
+        chain_lists,
+        strict=False,
+    ):
+        sequence = "".join(raw_sequence.split())
+        for chain_id in _clean_chain_list(raw_chains):
+            sequences[chain_id] = sequence
+    return sequences
+
+
+def _scheme_columns(category):
+    asym_ids = category["asym_id"].as_array(str)
+    insertion_codes = (
+        category["pdb_ins_code"].as_array(str)
+        if "pdb_ins_code" in category
+        else [""] * len(asym_ids)
+    )
+    hetflags = category["hetflag"].as_array(str) if "hetflag" in category else ["N"] * len(asym_ids)
+    author_chains = (
+        category["pdb_strand_id"].as_array(str) if "pdb_strand_id" in category else asym_ids
+    )
+    return (
+        asym_ids,
+        category["seq_id"].as_array(str),
+        category["auth_seq_num"].as_array(str),
+        insertion_codes,
+        hetflags,
+        author_chains,
+    )
+
+
+def _scheme_residue_map(category):
+    (
+        asym_ids,
+        sequence_positions,
+        author_numbers,
+        insertion_codes,
+        hetflags,
+        author_chains,
+    ) = _scheme_columns(category)
+    asym_to_author = {
+        asym_id: author_id for asym_id, author_id in zip(asym_ids, author_chains, strict=False)
+    }
+    per_chain: dict[str, dict[int, Residue]] = {}
+    for asym_id, raw_position, raw_number, raw_code, raw_hetflag in zip(
+        asym_ids,
+        sequence_positions,
+        author_numbers,
+        insertion_codes,
+        hetflags,
+        strict=False,
+    ):
+        residues = per_chain.setdefault(asym_id, {})
+        try:
+            position = int(raw_position) - 1
+            residue_number = int(raw_number) if raw_number != "?" else None
+        except ValueError:
+            continue
+        if residue_number is None:
+            insertion_code = ""
+        else:
+            insertion_code = "" if raw_code in (".", "?") else raw_code
+        residues[position] = Residue(
+            residue_number=residue_number,
+            insertion_code=insertion_code,
+            hetflag=raw_hetflag.upper() == "Y",
+        )
+    return per_chain, asym_to_author
+
+
+def _renumber_duplicate_residues(
+    per_chain: dict[str, dict[int, Residue]],
+) -> None:
+    for residues in per_chain.values():
+        positions_by_number: dict[int, list[int]] = {}
+        for position, residue in residues.items():
+            if residue.residue_number is not None:
+                positions_by_number.setdefault(residue.residue_number, []).append(position)
+        for number, positions in positions_by_number.items():
+            if len(positions) <= 1:
+                continue
+            positions.sort()
+            for offset, position in enumerate(positions):
+                previous = residues[position]
+                residues[position] = Residue(
+                    residue_number=number + offset,
+                    insertion_code=previous.insertion_code,
+                    hetflag=previous.hetflag,
+                )
+
+
+def _ordered_scheme_mapping(
+    per_chain: dict[str, dict[int, Residue]],
+    asym_to_author: dict[str, str],
+    chain_sequences: dict[str, str],
+) -> dict[str, dict[int, Residue]]:
+    result: dict[str, dict[int, Residue]] = {}
+    for asym_id, residues in per_chain.items():
+        author_chain = asym_to_author.get(asym_id, asym_id)
+        if author_chain in chain_sequences:
+            result[author_chain] = {
+                position: residues.get(position, _empty_residue())
+                for position in range(len(chain_sequences[author_chain]))
+            }
+        elif residues:
+            result[author_chain] = {
+                index: residues[position] for index, position in enumerate(sorted(residues))
+            }
+    return result
+
+
+def _complete_polymer_mappings(
+    mappings: dict[str, dict[int, Residue]],
+    chain_sequences: dict[str, str],
+) -> None:
+    for chain_id, sequence in chain_sequences.items():
+        mapping = mappings.setdefault(chain_id, {})
+        for position in range(len(sequence)):
+            if position not in mapping:
+                mapping[position] = _empty_residue()
+
+
+def _add_structure_fallbacks(
+    mappings: dict[str, dict[int, Residue]],
+    structure: bs.AtomArray,
+) -> None:
+    if not (
+        structure
+        and hasattr(structure, "chain_id")
+        and structure.chain_id is not None
+        and hasattr(structure.chain_id, "__iter__")
+    ):
+        return
+    for chain_id in set(structure.chain_id):
+        if chain_id in mappings:
+            continue
+        chain = structure[structure.chain_id == chain_id]
+        if not (
+            hasattr(chain, "res_id")
+            and chain.res_id is not None
+            and hasattr(chain.res_id, "__iter__")
+        ):
+            continue
+        residue_ids = sorted(set(chain.res_id))
+        mappings[chain_id] = {
+            index: Residue(
+                residue_number=residue_id,
+                insertion_code="",
+                hetflag=False,
+            )
+            for index, residue_id in enumerate(residue_ids)
+        }
+
+
+def _nonpolymer_entity_ids(block) -> set[str]:
+    result = set()
+    if "entity" not in block:
+        return result
+    category = block["entity"]
+    ids = category["id"].as_array(str)
+    types = category["type"].as_array(str)
+    for entity_id, entity_type in zip(ids, types, strict=False):
+        if entity_type.upper() in _NONPOLYMER_ENTITY_TYPES:
+            result.add(entity_id)
+    return result
+
+
+def _nonpolymer_component_map(block, entity_ids: set[str]) -> dict[str, str]:
+    result = {}
+    if "pdbx_entity_nonpoly" not in block:
+        return result
+    category = block["pdbx_entity_nonpoly"]
+    ids = category["entity_id"].as_array(str)
+    components = category["comp_id"].as_array(str)
+    for entity_id, component in zip(ids, components, strict=False):
+        if entity_id in entity_ids:
+            result[entity_id] = component
+    return result
+
+
+class MmcifWrapper:
+    """Parsed model-one structure, metadata, sequences, and residue mappings."""
+
+    def __init__(self, id: str | None = None):
+        self.id = id or ""
+        self.raw: pdbx.CIFFile | None = None
+        self.structure: bs.AtomArray
+        self.header = MmcifHeader()
+        self.entities: dict[int, list[str]] = {}
+        self.chain_to_seqres: dict[str, str] = {}
+        self.seqres_to_structure: dict[str, dict[int, Residue]] = {}
+
+    @classmethod
+    def read(cls, path: PathOrBuffer, id: str | None = None) -> MmcifWrapper:
+        wrapper = cls(id=id)
+        wrapper._load(path)
+        return wrapper
+
+    def _load(self, path: PathOrBuffer, fileid: str | None = None) -> None:
+        self.raw = pdbx.CIFFile.read(path)
+        self._parse_structure()
+        self._parse_header()
+        self._parse_entities()
+        self._parse_sequences()
+
+    def _parse_structure(self) -> None:
+        try:
+            structure = pdbx.get_structure(self.raw, model=1)
+            if structure is None or not isinstance(structure, bs.AtomArray):
+                raise NoProteinError("No structure found in mmCIF file")
+            if len(structure) == 0:
+                raise NoProteinError("Empty structure in mmCIF file")
+            self.structure = structure
+        except Exception as error:
+            raise ValueError(f"Failed to parse structure: {error}") from error
+
+    def _parse_header(self) -> None:
+        if self.raw:
+            self.header = _header_from_block(self.raw.block, self.header)
+
+    def _parse_entities(self) -> None:
+        if not self.raw:
+            return
+        try:
+            self.entities = _entities_from_block(self.raw.block, self.entities)
+        except Exception:
+            if (
+                self.structure
+                and hasattr(self.structure, "chain_id")
+                and self.structure.chain_id is not None
+                and hasattr(self.structure.chain_id, "__iter__")
+            ):
+                self.entities = {1: list(set(self.structure.chain_id))}
+
+    def _parse_sequences(self) -> None:
+        if not self.raw:
+            return
+        block = self.raw.block
+        self.chain_to_seqres.update(_polymer_sequences(block))
+        if "pdbx_poly_seq_scheme" in block:
+            per_chain, asym_to_author = _scheme_residue_map(block["pdbx_poly_seq_scheme"])
+            _renumber_duplicate_residues(per_chain)
+            self.seqres_to_structure.update(
+                _ordered_scheme_mapping(
+                    per_chain,
+                    asym_to_author,
+                    self.chain_to_seqres,
+                )
+            )
+        _complete_polymer_mappings(
+            self.seqres_to_structure,
+            self.chain_to_seqres,
+        )
+        _add_structure_fallbacks(self.seqres_to_structure, self.structure)
+
+    def _parse_nonpoly_from_mmcif(self) -> dict[tuple, bs.AtomArray]:
+        assert self.raw is not None
+        block = self.raw.block
+        entity_ids = _nonpolymer_entity_ids(block)
+        _nonpolymer_component_map(block, entity_ids)
+        groups: dict[tuple[str, str], list[int]] = {}
+        if "atom_site" in block:
+            category = block["atom_site"]
+            chain_ids = category["label_asym_id"].as_array(str)
+            atom_entity_ids = category["label_entity_id"].as_array(str)
+            component_ids = category["label_comp_id"].as_array(str)
+            for index, (chain_id, entity_id, component_id) in enumerate(
+                zip(chain_ids, atom_entity_ids, component_ids, strict=False)
+            ):
+                if entity_id in entity_ids:
+                    groups.setdefault((component_id, chain_id), []).append(index)
+
+        coordinates = {}
+        for component_id, chain_id in groups:
+            selection = (self.structure.chain_id == chain_id) & (
+                self.structure.res_name == component_id
+            )
+            if not selection.any():
+                continue
+            atoms = self.structure[selection]
+            if isinstance(atoms, (bs.AtomArray, bs.AtomArrayStack)) and len(atoms) > 0:
+                coordinates[(component_id, chain_id)] = atoms
+        return coordinates
+
+    def _parse_nonpoly_fallback(self) -> dict[tuple, bs.AtomArray]:
+        result = {}
+        if not (self.structure and hasattr(self.structure, "chain_id")):
+            return result
+        standard_residues = set(residue_constants.resnames[:-1])
+        standard_residues.update({"A", "C", "G", "T", "U"})
+        if self.structure.chain_id is None:
+            return result
+        for chain_id in set(self.structure.chain_id):
+            chain = self.structure[self.structure.chain_id == chain_id]
+            if not (
+                hasattr(chain, "res_name")
+                and chain.res_name is not None
+                and hasattr(chain.res_name, "__iter__")
+            ):
+                continue
+            for residue_name in set(chain.res_name):
+                if residue_name in standard_residues:
+                    continue
+                selection = (chain.chain_id == chain_id) & (chain.res_name == residue_name)
+                if selection.any() and isinstance(
+                    chain,
+                    (bs.AtomArray, bs.AtomArrayStack),
+                ):
+                    result[(residue_name, chain_id)] = chain[selection]
+        return result
+
+    @functools.cached_property
+    def non_polymer_coords(self) -> dict[tuple, bs.AtomArray]:
+        """Map each non-polymer component and chain to its atoms."""
+
+        if not self.structure or not self.raw:
+            return {}
+        try:
+            return self._parse_nonpoly_from_mmcif()
+        except Exception:
+            return self._parse_nonpoly_fallback()

fastplms/models/esmfold2/esmfold2_molecular_complex.py

@@ -0,0 +1,1016 @@
+"""Flat molecular-complex records used by the ESMFold2 public API.
+
+The folding model operates on tokens and a single atom table.  This module owns
+that representation, its protein-only bridge, mmCIF I/O, structure metrics, and
+the compact wire format.  It deliberately has no dependency on the upstream
+Biohub package; the pinned submodule is used only by differential tests.
+"""
+
+from __future__ import annotations
+
+import io
+import os
+import re
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from subprocess import check_output
+from tempfile import TemporaryDirectory
+from typing import TYPE_CHECKING, Any
+
+import biotite.structure as bs
+import biotite.structure.io.pdbx as pdbx
+import brotli
+import msgpack
+import numpy as np
+import torch
+from biotite.structure.io.pdbx import (
+    CIFCategory,
+    CIFColumn,
+    CIFData,
+    CIFFile,
+    set_structure,
+)
+
+from . import esmfold2_residue_constants as residue_constants
+from .esmfold2_metrics import compute_lddt, compute_rmsd
+from .esmfold2_mmcif_parsing import PLDDT_B_FACTOR_SCALE, round_mmcif_columns
+from .esmfold2_protein_complex import ProteinComplex, ProteinComplexMetadata
+
+
+@dataclass
+class MolecularComplexResult:
+    """One folded complex and the optional model outputs associated with it."""
+
+    complex: MolecularComplex
+    plddt: torch.Tensor | None = None
+    ptm: float | None = None
+    iptm: float | None = None
+    pae: torch.Tensor | None = None
+    distogram: torch.Tensor | None = None
+    pair_chains_iptm: torch.Tensor | None = None
+    output_embedding_sequence: torch.Tensor | None = None
+    output_embedding_pair_pooled: torch.Tensor | None = None
+    residue_index: torch.Tensor | None = None
+    entity_id: torch.Tensor | None = None
+    sae_features: np.ndarray | None = None  # X has shape (l, n_features).
+    ttt_metrics: dict[str, Any] | None = None
+
+
+@dataclass
+class MolecularComplexMetadata:
+    """Entity and chain labels carried with a molecular complex."""
+
+    entity_lookup: dict[int, str]
+    chain_lookup: dict[int, str]
+    assembly_composition: dict[str, list[str]] | None = None
+
+
+@dataclass
+class Molecule:
+    """The atom slice represented by one model token."""
+
+    token: str
+    token_idx: int
+    atom_positions: np.ndarray  # P has shape (n_atoms, 3).
+    atom_elements: np.ndarray  # E has shape (n_atoms,).
+    atom_names: np.ndarray | None = None  # N has shape (n_atoms,) when present.
+    atom_hetero: np.ndarray | None = None  # M has shape (n_atoms,) when present.
+    residue_type: int = 0
+    molecule_type: int = 0
+    confidence: float = 0.0
+
+
+_NUCLEOTIDE_NAMES = frozenset({"A", "T", "G", "C", "U", "DA", "DT", "DG", "DC"})
+_SERIALIZED_ARRAYS = frozenset(
+    {
+        "atom_positions",
+        "atom_elements",
+        "atom_names",
+        "atom_hetero",
+        "token_to_atoms",
+        "chain_id",
+        "entity_id",
+        "sym_id",
+        "plddt",
+    }
+)
+
+
+def _assert_table_lengths(complex_value: MolecularComplex) -> None:
+    """Check that token and atom annotations align with their tables."""
+    if not isinstance(complex_value.sequence, list) or any(
+        not isinstance(token, str) for token in complex_value.sequence
+    ):
+        raise TypeError("sequence must be a list of token strings.")
+    n_tokens = len(complex_value.sequence)
+    if not isinstance(complex_value.atom_positions, np.ndarray):
+        raise TypeError("atom_positions must be a NumPy array.")
+    if complex_value.atom_positions.ndim != 2 or complex_value.atom_positions.shape[1:] != (
+        3,
+    ):
+        raise ValueError(
+            "atom_positions must have shape (n_atoms, 3), got "
+            f"{complex_value.atom_positions.shape}."
+        )
+    if not np.issubdtype(complex_value.atom_positions.dtype, np.number):
+        raise TypeError("atom_positions must use a numeric dtype.")
+    n_atoms = len(complex_value.atom_positions)
+    if not isinstance(complex_value.atom_elements, np.ndarray):
+        raise TypeError("atom_elements must be a NumPy array.")
+    if complex_value.atom_elements.shape != (n_atoms,):
+        raise ValueError(
+            f"atom_elements shape {complex_value.atom_elements.shape} != {n_atoms} atoms"
+        )
+    token_tables = {
+        "token_to_atoms": complex_value.token_to_atoms,
+        "chain_id": complex_value.chain_id,
+        "plddt": complex_value.plddt,
+    }
+    if complex_value.entity_id is not None:
+        token_tables["entity_id"] = complex_value.entity_id
+    if complex_value.sym_id is not None:
+        token_tables["sym_id"] = complex_value.sym_id
+    for label, values in token_tables.items():
+        if not isinstance(values, np.ndarray):
+            raise TypeError(f"{label} must be a NumPy array, got {type(values).__name__}.")
+        if values.ndim == 0 or values.shape[0] != n_tokens:
+            raise ValueError(f"{label} shape {values.shape} != {n_tokens} tokens")
+    if complex_value.token_to_atoms.shape != (n_tokens, 2):
+        raise ValueError(
+            "token_to_atoms must have shape "
+            f"({n_tokens}, 2), got {complex_value.token_to_atoms.shape}."
+        )
+    if not np.issubdtype(complex_value.token_to_atoms.dtype, np.integer):
+        raise TypeError("token_to_atoms must use an integer dtype.")
+    if complex_value.chain_id.shape != (n_tokens,):
+        raise ValueError(f"chain_id must have shape ({n_tokens},).")
+    for label, values in (
+        ("chain_id", complex_value.chain_id),
+        ("entity_id", complex_value.entity_id),
+        ("sym_id", complex_value.sym_id),
+    ):
+        if values is not None and values.shape != (n_tokens,):
+            raise ValueError(f"{label} must have shape ({n_tokens},).")
+        if values is not None and not np.issubdtype(values.dtype, np.integer):
+            raise TypeError(f"{label} must use an integer dtype.")
+    if complex_value.plddt.shape != (n_tokens,):
+        raise ValueError(f"plddt must have shape ({n_tokens},).")
+    if not np.issubdtype(complex_value.plddt.dtype, np.number):
+        raise TypeError("plddt must use a numeric dtype.")
+    if n_tokens:
+        starts = complex_value.token_to_atoms[:, 0]
+        stops = complex_value.token_to_atoms[:, 1]
+        if np.any(starts < 0) or np.any(stops < starts) or np.any(stops > n_atoms):
+            raise ValueError("token_to_atoms contains an invalid or out-of-bounds atom span.")
+    for label, values in (
+        ("atom_names", complex_value.atom_names),
+        ("atom_hetero", complex_value.atom_hetero),
+    ):
+        if values is not None and not isinstance(values, np.ndarray):
+            raise TypeError(f"{label} must be a NumPy array, got {type(values).__name__}.")
+        if isinstance(values, np.ndarray) and values.shape != (n_atoms,):
+            raise ValueError(f"{label} shape {values.shape} != {n_atoms} atoms")
+
+
+def _flat_protein_atoms(
+    protein: ProteinComplex,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    """Flatten the populated atom37 entries of a protein complex."""
+    positions: list[np.ndarray] = []
+    elements: list[str] = []
+    names: list[str] = []
+    hetero: list[bool] = []
+    spans: list[tuple[int, int]] = []
+
+    for sequence_index, residue in enumerate(protein.sequence):
+        if residue == "|":
+            continue
+        start = len(positions)
+        mask = protein.atom37_mask[sequence_index]
+        residue_positions = protein.atom37_positions[sequence_index]
+        for atom_index in np.flatnonzero(mask):
+            atom_name = residue_constants.atom_types[int(atom_index)]
+            positions.append(residue_positions[atom_index])
+            elements.append(atom_name[0] if atom_name else "C")
+            names.append(atom_name)
+            hetero.append(False)
+        spans.append((start, len(positions)))
+
+    return (
+        np.asarray(positions, dtype=np.float32),
+        np.asarray(elements, dtype=object),
+        np.asarray(names, dtype=object),
+        np.asarray(hetero, dtype=bool),
+        np.asarray(spans, dtype=np.int32),
+    )
+
+
+def _protein_sequence_and_indices(
+    complex_value: MolecularComplex,
+) -> tuple[list[int], str, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
+    protein_indices = [
+        index
+        for index, token in enumerate(complex_value.sequence)
+        if token in residue_constants.restype_3to1
+    ]
+    if not protein_indices:
+        raise ValueError("No protein tokens found in MolecularComplex")
+
+    chain_ids = complex_value.chain_id[protein_indices]
+    entity_ids = (
+        chain_ids
+        if complex_value.entity_id is None
+        else complex_value.entity_id[protein_indices]
+    )
+    sym_ids = (
+        np.zeros_like(chain_ids)
+        if complex_value.sym_id is None
+        else complex_value.sym_id[protein_indices]
+    )
+    confidences = complex_value.plddt[protein_indices]
+    sequence: list[str] = []
+    previous_instance: Any = None
+    preserve_instances = complex_value.sym_id is not None
+    for index, chain_id, sym_id in zip(
+        protein_indices, chain_ids, sym_ids, strict=True
+    ):
+        instance = (int(chain_id), int(sym_id)) if preserve_instances else int(chain_id)
+        if previous_instance is not None and instance != previous_instance:
+            sequence.append("|")
+        sequence.append(residue_constants.restype_3to1[complex_value.sequence[index]])
+        previous_instance = instance
+    return protein_indices, "".join(sequence), chain_ids, entity_ids, sym_ids, confidences
+
+
+def _protein_entity_metadata_value(value: int | str) -> int | str:
+    """Restore the numeric entity labels used by ProteinComplex metadata."""
+    if isinstance(value, str):
+        try:
+            return int(value)
+        except ValueError:
+            pass
+    return value
+
+
+def _atom37_from_flat(
+    complex_value: MolecularComplex, protein_indices: list[int]
+) -> tuple[np.ndarray, np.ndarray]:
+    n_residues = len(protein_indices)
+    positions = np.full((n_residues, 37, 3), np.nan, dtype=np.float32)
+    mask = np.zeros((n_residues, 37), dtype=bool)
+    if complex_value.atom_names is None:
+        return positions, mask
+
+    for residue_index, token_index in enumerate(protein_indices):
+        start, stop = complex_value.token_to_atoms[token_index]
+        seen: set[str] = set()
+        for atom_name, atom_position in zip(
+            complex_value.atom_names[start:stop],
+            complex_value.atom_positions[start:stop],
+            strict=True,
+        ):
+            normalized = str(atom_name).upper().strip()
+            if normalized in seen:
+                continue
+            seen.add(normalized)
+            atom37_index = residue_constants.atom_order.get(normalized)
+            if atom37_index is not None:
+                positions[residue_index, atom37_index] = atom_position
+                mask[residue_index, atom37_index] = True
+    return positions, mask
+
+
+def _expand_protein_rows(
+    sequence: str,
+    protein_chain_ids: np.ndarray,
+    protein_entity_ids: np.ndarray,
+    protein_sym_ids: np.ndarray,
+    confidences: np.ndarray,
+    compact_positions: np.ndarray,
+    compact_mask: np.ndarray,
+) -> dict[str, np.ndarray]:
+    """Insert empty rows at chain separators in a protein representation."""
+    n_positions = len(sequence)
+    expanded = {
+        "chain_id": np.full(n_positions, -1, dtype=np.int64),
+        "entity_id": np.full(n_positions, -1, dtype=np.int64),
+        "sym_id": np.zeros(n_positions, dtype=np.int64),
+        "residue_index": np.zeros(n_positions, dtype=np.int64),
+        "insertion_code": np.asarray([""] * n_positions, dtype=object),
+        "confidence": np.zeros(n_positions, dtype=np.float32),
+        "atom37_positions": np.full((n_positions, 37, 3), np.nan, dtype=np.float32),
+        "atom37_mask": np.zeros((n_positions, 37), dtype=bool),
+    }
+    residue_number = 0
+    compact_index = 0
+    for sequence_index, residue in enumerate(sequence):
+        if residue == "|":
+            residue_number = 0
+            continue
+        chain_id = protein_chain_ids[compact_index]
+        residue_number += 1
+        expanded["chain_id"][sequence_index] = chain_id
+        expanded["entity_id"][sequence_index] = protein_entity_ids[compact_index]
+        expanded["sym_id"][sequence_index] = protein_sym_ids[compact_index]
+        expanded["residue_index"][sequence_index] = residue_number
+        expanded["confidence"][sequence_index] = confidences[compact_index]
+        expanded["atom37_positions"][sequence_index] = compact_positions[compact_index]
+        expanded["atom37_mask"][sequence_index] = compact_mask[compact_index]
+        compact_index += 1
+    return expanded
+
+
+def _read_cif(source: str) -> CIFFile:
+    if os.path.exists(source):
+        return pdbx.CIFFile.read(source)
+    return pdbx.CIFFile.read(io.StringIO(source))
+
+
+def _read_structure(cif_file: CIFFile) -> Any:
+    try:
+        return pdbx.get_structure(cif_file, model=1, extra_fields=["b_factor"])
+    except (KeyError, ValueError):
+        try:
+            return pdbx.get_structure(cif_file)
+        except Exception:
+            return pdbx.get_structure(cif_file, model=None)
+
+
+def _column_array(category: Any, name: str) -> np.ndarray:
+    column = category[name]
+    if hasattr(column, "as_array"):
+        return column.as_array(str)
+    return np.asarray(list(column), dtype=str)
+
+
+def _label_asym_ids(cif_file: CIFFile, n_structure_atoms: int) -> list[str] | None:
+    """Return label-asym identifiers after applying Biohub's atom filters."""
+    block = cif_file.block
+    if "atom_site" not in block or "label_asym_id" not in block["atom_site"]:
+        return None
+    atom_site = block["atom_site"]
+    labels = _column_array(atom_site, "label_asym_id")
+    keep = np.ones(len(labels), dtype=bool)
+    if "pdbx_PDB_model_num" in atom_site:
+        keep &= _column_array(atom_site, "pdbx_PDB_model_num") == "1"
+    if "label_alt_id" in atom_site:
+        keep &= np.isin(_column_array(atom_site, "label_alt_id"), [".", "?", "", "A"])
+    filtered = labels[keep]
+    return filtered.tolist() if len(filtered) == n_structure_atoms else None
+
+
+def _entity_metadata(cif_file: CIFFile) -> dict[Any, Any]:
+    result: dict[Any, Any] = {}
+    try:
+        category = cif_file.block["entity"]
+        if "id" not in category or "type" not in category:
+            return result
+        for entity_id, entity_type in zip(category["id"], category["type"], strict=False):
+            result[entity_id] = entity_type
+    except Exception:
+        return {}
+    return result
+
+
+def _group_structure_atoms(
+    structure: Any, labels: list[str] | None
+) -> dict[str, dict[tuple[int, str], dict[str, Any]]]:
+    grouped: dict[str, dict[tuple[int, str], dict[str, Any]]] = {}
+    for atom_index, atom in enumerate(structure):
+        chain = labels[atom_index] if labels is not None else atom.chain_id
+        residues = grouped.setdefault(chain, {})
+        key = (atom.res_id, atom.res_name)
+        record = residues.setdefault(
+            key,
+            {"atoms": [], "res_name": atom.res_name, "is_hetero": atom.hetero},
+        )
+        record["atoms"].append(atom)
+    return grouped
+
+
+def _flatten_structure_groups(
+    grouped: dict[str, dict[tuple[int, str], dict[str, Any]]],
+) -> tuple[
+    list[str],
+    list[np.ndarray],
+    list[str],
+    list[str],
+    list[bool],
+    list[tuple[int, int]],
+    list[float],
+    list[int],
+    dict[str, int],
+]:
+    tokens: list[str] = []
+    positions: list[np.ndarray] = []
+    elements: list[str] = []
+    names: list[str] = []
+    hetero: list[bool] = []
+    spans: list[tuple[int, int]] = []
+    confidences: list[float] = []
+    token_chains: list[int] = []
+    chain_numbers = {chain: index for index, chain in enumerate(sorted(grouped))}
+
+    for chain in sorted(grouped):
+        for residue_key in sorted(grouped[chain]):
+            record = grouped[chain][residue_key]
+            if record["res_name"] == "HOH":
+                continue
+            atoms = record["atoms"]
+            tokens.append(record["res_name"])
+            token_chains.append(chain_numbers[chain])
+            start = len(positions)
+            positions.extend(atom.coord for atom in atoms)
+            elements.extend(atom.element for atom in atoms)
+            names.extend(atom.atom_name for atom in atoms)
+            hetero.extend(atom.hetero for atom in atoms)
+            spans.append((start, len(positions)))
+            b_factor = getattr(atoms[0], "b_factor", 50.0) if atoms else 50.0
+            confidences.append(min(b_factor / PLDDT_B_FACTOR_SCALE, 1.0))
+    return (
+        tokens,
+        positions,
+        elements,
+        names,
+        hetero,
+        spans,
+        confidences,
+        token_chains,
+        chain_numbers,
+    )
+
+
+def _chain_entity_maps(
+    complex_value: MolecularComplex,
+) -> tuple[dict[str, list[str]], dict[str, int], dict[int, tuple[str, ...]]]:
+    chains: dict[str, list[str]] = {}
+    for token_index, numeric_chain in enumerate(complex_value.chain_id):
+        numeric = int(numeric_chain)
+        label = complex_value.metadata.chain_lookup.get(numeric, chr(65 + numeric))
+        chains.setdefault(label, []).append(complex_value.sequence[token_index])
+
+    sequence_entities: dict[tuple[str, ...], int] = {}
+    chain_entities: dict[str, int] = {}
+    entity_sequences: dict[int, tuple[str, ...]] = {}
+    for label, sequence in chains.items():
+        key = tuple(sequence)
+        entity_id = sequence_entities.get(key)
+        if entity_id is None:
+            entity_id = len(sequence_entities) + 1
+            sequence_entities[key] = entity_id
+            entity_sequences[entity_id] = key
+        chain_entities[label] = entity_id
+    return chains, chain_entities, entity_sequences
+
+
+def _cif_column(values: list[str]) -> CIFColumn:
+    return CIFColumn(data=CIFData(array=np.asarray(values), dtype=np.str_))
+
+
+def _add_entity_categories(
+    cif_file: CIFFile,
+    complex_value: MolecularComplex,
+    entity_sequences: dict[int, tuple[str, ...]],
+) -> None:
+    ids: list[str] = []
+    types: list[str] = []
+    descriptions: list[str] = []
+    for entity_id in sorted(entity_sequences):
+        sequence = entity_sequences[entity_id]
+        protein = any(token in residue_constants.restype_3to1 for token in sequence)
+        nucleic = any(token in _NUCLEOTIDE_NAMES for token in sequence)
+        ids.append(str(entity_id))
+        types.append("polymer" if protein or nucleic else "non-polymer")
+        if protein:
+            descriptions.append(f"Polymer entity {entity_id} (protein)")
+        elif nucleic:
+            descriptions.append(f"Polymer entity {entity_id} (nucleic acid)")
+        else:
+            descriptions.append(f"Non-polymer entity {entity_id}")
+
+    if ids:
+        cif_file.block["entity"] = CIFCategory(
+            name="entity",
+            columns={
+                "id": _cif_column(ids),
+                "type": _cif_column(types),
+                "pdbx_description": _cif_column(descriptions),
+            },
+        )
+
+    _, chain_entities, _ = _chain_entity_maps(complex_value)
+    if chain_entities:
+        labels = sorted(chain_entities)
+        cif_file.block["struct_asym"] = CIFCategory(
+            name="struct_asym",
+            columns={
+                "id": _cif_column(labels),
+                "entity_id": _cif_column([str(chain_entities[label]) for label in labels]),
+            },
+        )
+
+    entity_chains: dict[int, list[str]] = {}
+    for chain, entity_id in chain_entities.items():
+        entity_chains.setdefault(entity_id, []).append(chain)
+    polymer_rows: list[tuple[str, str, str, str]] = []
+    residue_rows: list[tuple[str, str, str, str]] = []
+    for entity_id in sorted(entity_sequences):
+        sequence = entity_sequences[entity_id]
+        protein = any(token in residue_constants.restype_3to1 for token in sequence)
+        nucleic = any(token in _NUCLEOTIDE_NAMES for token in sequence)
+        if not (protein or nucleic):
+            continue
+        if protein:
+            polymer_type = "polypeptide(L)"
+            canonical = "".join(
+                residue_constants.restype_3to1.get(token, "(X)") for token in sequence
+            )
+        else:
+            polymer_type = (
+                "polyribonucleotide"
+                if "U" in sequence
+                else (
+                    "polydeoxyribonucleotide"
+                    if any(token in {"DA", "DT", "DG", "DC"} for token in sequence)
+                    else "polyribonucleotide"
+                )
+            )
+            nucleotide_letters = {"DA": "A", "DT": "T", "DG": "G", "DC": "C"}
+            canonical = "".join(nucleotide_letters.get(token, token) for token in sequence)
+        strand_ids = ",".join(sorted(entity_chains.get(entity_id, []))) or "?"
+        polymer_rows.append((str(entity_id), polymer_type, strand_ids, canonical))
+        residue_rows.extend(
+            (str(entity_id), str(number), token, "n")
+            for number, token in enumerate(sequence, start=1)
+        )
+
+    if polymer_rows:
+        columns = list(zip(*polymer_rows, strict=True))
+        cif_file.block["entity_poly"] = CIFCategory(
+            name="entity_poly",
+            columns={
+                "entity_id": _cif_column(list(columns[0])),
+                "type": _cif_column(list(columns[1])),
+                "pdbx_strand_id": _cif_column(list(columns[2])),
+                "pdbx_seq_one_letter_code_can": _cif_column(list(columns[3])),
+            },
+        )
+    if residue_rows:
+        columns = list(zip(*residue_rows, strict=True))
+        cif_file.block["entity_poly_seq"] = CIFCategory(
+            name="entity_poly_seq",
+            columns={
+                "entity_id": _cif_column(list(columns[0])),
+                "num": _cif_column(list(columns[1])),
+                "mon_id": _cif_column(list(columns[2])),
+                "hetero": _cif_column(list(columns[3])),
+            },
+        )
+
+
+def _fallback_atom_names(token: str, count: int) -> list[str]:
+    if token in residue_constants.restype_3to1:
+        names = list(residue_constants.residue_atoms.get(token, ["N", "CA", "C", "O"]))[:count]
+        names.extend(f"X{index + 1}" for index in range(len(names), count))
+        return names
+    return [f"C{index + 1}" for index in range(count)]
+
+
+def _as_atom_array(complex_value: MolecularComplex, chain_entities: dict[str, int]) -> bs.AtomArray:
+    n_atoms = len(complex_value.atom_positions)
+    atom_array = bs.AtomArray(length=n_atoms)
+    atom_array.coord = complex_value.atom_positions
+    residue_ids = np.zeros(n_atoms, dtype=np.int32)
+    chain_labels = np.empty(n_atoms, dtype=object)
+    residue_names = np.empty(n_atoms, dtype=object)
+    hetero = np.zeros(n_atoms, dtype=bool)
+    b_factors = np.zeros(n_atoms, dtype=np.float32)
+    atom_names = np.empty(n_atoms, dtype=object)
+    entity_ids = np.zeros(n_atoms, dtype=np.int32)
+    next_residue: dict[Any, int] = {}
+
+    for token_index, (start, stop) in enumerate(complex_value.token_to_atoms):
+        token = complex_value.sequence[token_index]
+        numeric_chain = complex_value.chain_id[token_index]
+        numeric = int(numeric_chain)
+        chain = complex_value.metadata.chain_lookup.get(numeric, chr(65 + numeric))
+        residue_id = next_residue.get(numeric_chain, 0) + 1
+        next_residue[numeric_chain] = residue_id
+        count = int(stop - start)
+        names = (
+            list(complex_value.atom_names[start:stop])
+            if complex_value.atom_names is not None
+            else _fallback_atom_names(token, count)
+        )
+        residue_ids[start:stop] = residue_id
+        chain_labels[start:stop] = chain
+        residue_names[start:stop] = token
+        hetero[start:stop] = (
+            complex_value.atom_hetero[start:stop]
+            if complex_value.atom_hetero is not None
+            else token not in residue_constants.restype_3to1
+        )
+        b_factors[start:stop] = complex_value.plddt[token_index] * PLDDT_B_FACTOR_SCALE
+        atom_names[start:stop] = names
+        entity_ids[start:stop] = chain_entities.get(chain, 1)
+
+    atom_array.res_id = residue_ids
+    atom_array.chain_id = np.asarray(chain_labels, dtype="U16")
+    atom_array.res_name = np.asarray(residue_names, dtype="U8")
+    atom_array.hetero = hetero
+    atom_array.atom_name = np.asarray(atom_names, dtype="U4")
+    atom_array.add_annotation("b_factor", dtype=float)
+    atom_array.b_factor = b_factors
+    atom_array.add_annotation("occupancy", dtype=float)
+    atom_array.occupancy = np.ones(n_atoms, dtype=np.float32)
+    atom_array.add_annotation("entity_id", dtype=int)
+    atom_array.entity_id = entity_ids
+    if complex_value.atom_elements is not None and len(complex_value.atom_elements) == n_atoms:
+        atom_array.element = np.asarray(complex_value.atom_elements, dtype="U4")
+    else:
+        atom_array.element = bs.infer_elements(atom_array)
+    return atom_array
+
+
+def _repair_label_entity_ids(cif_file: CIFFile, chain_entities: dict[str, int]) -> None:
+    if "atom_site" not in cif_file.block:
+        return
+    atom_site = cif_file.block["atom_site"]
+    if "label_asym_id" not in atom_site or "label_entity_id" not in atom_site:
+        return
+    labels = _column_array(atom_site, "label_asym_id").tolist()
+    if labels:
+        atom_site["label_entity_id"] = _cif_column(
+            [str(chain_entities.get(label, 1)) for label in labels]
+        )
+
+
+def _centroid_tensors(
+    mobile: MolecularComplex,
+    target: MolecularComplex,
+    *,
+    retain_missing: bool,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    if len(mobile) != len(target):
+        raise ValueError(
+            f"Complexes must have the same number of tokens: {len(mobile)} vs {len(target)}"
+        )
+    mobile_centers: list[np.ndarray] = []
+    target_centers: list[np.ndarray] = []
+    valid: list[bool] = []
+    for token_index in range(len(mobile)):
+        mobile_start, mobile_stop = mobile.token_to_atoms[token_index]
+        target_start, target_stop = target.token_to_atoms[token_index]
+        mobile_atoms = mobile.atom_positions[mobile_start:mobile_stop]
+        target_atoms = target.atom_positions[target_start:target_stop]
+        present = len(mobile_atoms) > 0 and len(target_atoms) > 0
+        if not present and not retain_missing:
+            continue
+        if present:
+            mobile_centers.append(mobile_atoms.mean(axis=0))
+            target_centers.append(target_atoms.mean(axis=0))
+        else:
+            mobile_centers.append(np.full(3, np.nan))
+            target_centers.append(np.full(3, np.nan))
+        valid.append(present)
+    if not any(valid):
+        metric = "LDDT" if retain_missing else "RMSD"
+        raise ValueError(f"No valid atoms found for {metric} computation")
+    return (
+        torch.from_numpy(np.stack(mobile_centers)).unsqueeze(0),
+        torch.from_numpy(np.stack(target_centers)).unsqueeze(0),
+        torch.as_tensor(valid, dtype=torch.bool).unsqueeze(0),
+    )
+
+
+@dataclass(frozen=True)
+class MolecularComplex:
+    """A token sequence backed by one contiguous atom table.
+
+    P stores atom coordinates with shape (n_atoms, 3).  Token span ``i`` is
+    ``P[token_to_atoms[i, 0]:token_to_atoms[i, 1]]``. ``chain_id`` identifies
+    the author chain, while optional ``entity_id`` and ``sym_id`` distinguish
+    biological entities and repeated chain instances.
+    """
+
+    id: str
+    sequence: list[str]
+    atom_positions: np.ndarray  # P has shape (n_atoms, 3).
+    atom_elements: np.ndarray  # E has shape (n_atoms,).
+    token_to_atoms: np.ndarray  # I has shape (n_tokens, 2).
+    chain_id: np.ndarray  # C has shape (n_tokens,).
+    plddt: np.ndarray  # S has shape (n_tokens,).
+    metadata: MolecularComplexMetadata
+    atom_names: np.ndarray | None = None  # N has shape (n_atoms,) when present.
+    atom_hetero: np.ndarray | None = None  # M has shape (n_atoms,) when present.
+    # These token-aligned IDs are optional for compatibility with older blobs.
+    # ProteinComplex adapters populate them so homomers and repeated author-chain
+    # labels survive a MolecularComplex round trip.
+    entity_id: np.ndarray | None = None
+    sym_id: np.ndarray | None = None
+
+    def __post_init__(self) -> None:
+        _assert_table_lengths(self)
+
+    def __len__(self) -> int:
+        return len(self.sequence)
+
+    def __getitem__(self, idx: int) -> Molecule:
+        if idx < 0 or idx >= len(self):
+            raise IndexError(f"Token index {idx} out of range for {len(self)} tokens")
+        start, stop = self.token_to_atoms[idx]
+        return Molecule(
+            token=self.sequence[idx],
+            token_idx=idx,
+            atom_positions=self.atom_positions[start:stop],
+            atom_elements=self.atom_elements[start:stop],
+            atom_names=None if self.atom_names is None else self.atom_names[start:stop],
+            atom_hetero=(None if self.atom_hetero is None else self.atom_hetero[start:stop]),
+            residue_type=0,
+            molecule_type=0,
+            confidence=self.plddt[idx],
+        )
+
+    @property
+    def atom_coordinates(self) -> np.ndarray:
+        """Return P, the flat atom-coordinate table with shape (n_atoms, 3)."""
+        return self.atom_positions
+
+    @classmethod
+    def from_protein_complex(cls, pc: ProteinComplex) -> MolecularComplex:
+        positions, elements, names, hetero, spans = _flat_protein_atoms(pc)
+        residue_positions = [index for index, value in enumerate(pc.sequence) if value != "|"]
+        metadata = MolecularComplexMetadata(
+            entity_lookup={key: str(value) for key, value in pc.metadata.entity_lookup.items()},
+            chain_lookup=dict(pc.metadata.chain_lookup),
+            assembly_composition=pc.metadata.assembly_composition,
+        )
+        return cls(
+            id=pc.id,
+            sequence=[
+                residue_constants.restype_1to3.get(pc.sequence[index], "UNK")
+                for index in residue_positions
+            ],
+            atom_positions=positions,
+            atom_elements=elements,
+            token_to_atoms=spans,
+            chain_id=np.asarray(pc.chain_id[residue_positions], dtype=np.int64),
+            plddt=np.asarray(pc.confidence[residue_positions], dtype=np.float32),
+            metadata=metadata,
+            atom_names=names,
+            atom_hetero=hetero,
+            entity_id=np.asarray(pc.entity_id[residue_positions], dtype=np.int64),
+            sym_id=np.asarray(pc.sym_id[residue_positions], dtype=np.int64),
+        )
+
+    def to_protein_complex(self) -> ProteinComplex:
+        (
+            protein_indices,
+            sequence,
+            chain_ids,
+            entity_ids,
+            sym_ids,
+            confidences,
+        ) = _protein_sequence_and_indices(self)
+        compact_positions, compact_mask = _atom37_from_flat(self, protein_indices)
+        arrays = _expand_protein_rows(
+            sequence,
+            chain_ids,
+            entity_ids,
+            sym_ids,
+            confidences,
+            compact_positions,
+            compact_mask,
+        )
+        unique_chains = np.unique(chain_ids)
+        unique_entities = np.unique(entity_ids)
+        metadata = ProteinComplexMetadata(
+            entity_lookup={
+                int(entity): _protein_entity_metadata_value(
+                    self.metadata.entity_lookup.get(int(entity), int(entity))
+                )
+                for entity in unique_entities
+            },
+            chain_lookup={
+                int(chain): self.metadata.chain_lookup.get(int(chain), chr(65 + int(chain)))
+                for chain in unique_chains
+            },
+            assembly_composition=self.metadata.assembly_composition,
+        )
+        return ProteinComplex(
+            id=self.id,
+            sequence=sequence,
+            entity_id=arrays["entity_id"],
+            chain_id=arrays["chain_id"],
+            sym_id=arrays["sym_id"],
+            residue_index=arrays["residue_index"],
+            insertion_code=arrays["insertion_code"],
+            atom37_positions=arrays["atom37_positions"],
+            atom37_mask=arrays["atom37_mask"],
+            confidence=arrays["confidence"],
+            metadata=metadata,
+        )
+
+    @classmethod
+    def from_mmcif(cls, inp: str, id: str | None = None) -> MolecularComplex:
+        cif_file = _read_cif(inp)
+        structure = _read_structure(cif_file)
+        if TYPE_CHECKING:
+            structure: Any = structure
+        labels = _label_asym_ids(cif_file, len(structure))
+        grouped = _group_structure_atoms(structure, labels)
+        (
+            tokens,
+            positions,
+            elements,
+            names,
+            hetero,
+            spans,
+            confidences,
+            token_chains,
+            chain_numbers,
+        ) = _flatten_structure_groups(grouped)
+        n_tokens = len(tokens)
+        if positions:
+            position_array = np.asarray(positions, dtype=np.float32)
+            element_array = np.asarray(elements, dtype=object)
+            name_array = np.asarray(names, dtype=object)
+            hetero_array = np.asarray(hetero, dtype=bool)
+            span_array = np.asarray(spans, dtype=np.int32)
+            chain_array = np.asarray(token_chains, dtype=np.int64)
+        else:
+            position_array = np.zeros((0, 3), dtype=np.float32)
+            element_array = np.zeros(0, dtype=object)
+            name_array = np.zeros(0, dtype=object)
+            hetero_array = np.zeros(0, dtype=bool)
+            span_array = np.zeros((n_tokens, 2), dtype=np.int32)
+            chain_array = (
+                np.asarray(token_chains, dtype=np.int64)
+                if token_chains
+                else np.zeros(n_tokens, dtype=np.int64)
+            )
+        complex_id = id or (Path(inp).stem if os.path.exists(inp) else "complex_from_string")
+        return cls(
+            id=complex_id,
+            sequence=tokens,
+            atom_positions=position_array,
+            atom_elements=element_array,
+            token_to_atoms=span_array,
+            chain_id=chain_array,
+            plddt=np.asarray(confidences, dtype=np.float32),
+            metadata=MolecularComplexMetadata(
+                entity_lookup=_entity_metadata(cif_file),
+                chain_lookup={number: chain for chain, number in chain_numbers.items()},
+                assembly_composition=None,
+            ),
+            atom_names=name_array,
+            atom_hetero=hetero_array,
+        )
+
+    def _get_entity_mapping(
+        self,
+    ) -> tuple[dict[str, list[str]], dict[str, int], dict[int, tuple[str, ...]]]:
+        return _chain_entity_maps(self)
+
+    def _add_entity_information(
+        self, cif_file: CIFFile, entity_sequences: dict[int, tuple[str, ...]]
+    ) -> None:
+        _add_entity_categories(cif_file, self, entity_sequences)
+
+    def to_mmcif(self) -> str:
+        _, chain_entities, entity_sequences = _chain_entity_maps(self)
+        atom_array = _as_atom_array(self, chain_entities)
+        cif_file = CIFFile()
+        set_structure(cif_file, atom_array, data_block=self.id)
+        _repair_label_entity_ids(cif_file, chain_entities)
+        _add_entity_categories(cif_file, self, entity_sequences)
+        round_mmcif_columns(cif_file)
+        output = io.StringIO()
+        cif_file.write(output)
+        return output.getvalue()
+
+    def dockq(self, native: MolecularComplex) -> Any:
+        try:
+            mobile = self.to_protein_complex().normalize_chain_ids_for_pdb()
+            target = native.to_protein_complex().normalize_chain_ids_for_pdb()
+        except ValueError as error:
+            raise ValueError(
+                f"Cannot convert MolecularComplex to ProteinComplex for DockQ: {error}"
+            ) from None
+        try:
+            return mobile.dockq(target)
+        except Exception:
+            return self._compute_dockq_manual(native)
+
+    def _compute_dockq_manual(self, native: MolecularComplex) -> Any:
+        try:
+            mobile = self.to_protein_complex().normalize_chain_ids_for_pdb()
+            target = native.to_protein_complex().normalize_chain_ids_for_pdb()
+        except ValueError as error:
+            raise ValueError(
+                f"Cannot convert MolecularComplex to ProteinComplex for DockQ: {error}"
+            ) from None
+        with TemporaryDirectory() as directory:
+            mobile_path = Path(directory) / "self.pdb"
+            target_path = Path(directory) / "native.pdb"
+            mobile.to_pdb(mobile_path)
+            target.to_pdb(target_path)
+            try:
+                raw_output = check_output(["DockQ", str(mobile_path), str(target_path)])
+                output = raw_output.decode()
+                score: float | None = None
+                for line in output.split("\n"):
+                    if "Total DockQ" in line:
+                        match = re.search(r"Total DockQ.*: ([\d.]+)", line)
+                        if match:
+                            score = float(match.group(1))
+                            break
+                if score is None:
+                    for line in output.split("\n"):
+                        if line.startswith("DockQ") and ":" in line:
+                            try:
+                                score = float(line.split(":")[1].strip())
+                                break
+                            except (ValueError, IndexError):
+                                continue
+                if score is None:
+                    raise ValueError("Could not parse DockQ score from output")
+                return {"total_dockq": score, "raw_output": output, "aligned": self}
+            except FileNotFoundError:
+                raise RuntimeError(
+                    "DockQ is not installed. Please install DockQ to use this method."
+                ) from None
+            except Exception as error:
+                raise RuntimeError(f"DockQ computation failed: {error}") from error
+
+    def rmsd(self, target: MolecularComplex, **kwargs: Any) -> float:
+        mobile, reference, mask = _centroid_tensors(self, target, retain_missing=False)
+        value = compute_rmsd(
+            mobile=mobile,
+            target=reference,
+            atom_exists_mask=mask,
+            reduction="batch",
+            **kwargs,
+        )
+        return float(value)
+
+    def lddt_ca(self, target: MolecularComplex, **kwargs: Any) -> float:
+        mobile, reference, mask = _centroid_tensors(self, target, retain_missing=True)
+        value = compute_lddt(
+            all_atom_pred_pos=mobile,
+            all_atom_positions=reference,
+            all_atom_mask=mask,
+            per_residue=False,
+            **kwargs,
+        )
+        return float(value)
+
+    def state_dict(self) -> dict[str, Any]:
+        state = dict(vars(self))
+        for optional_identity in ("entity_id", "sym_id"):
+            if state[optional_identity] is None:
+                state.pop(optional_identity)
+        for key, value in tuple(state.items()):
+            if isinstance(value, MolecularComplexMetadata):
+                state[key] = asdict(value)
+            elif isinstance(value, np.ndarray):
+                if value.dtype == np.int64:
+                    value = value.astype(np.int32)
+                elif value.dtype in (np.dtype(np.float64), np.dtype(np.float32)):
+                    value = value.astype(np.float16)
+                state[key] = value.tolist()
+        return state
+
+    def to_blob(self) -> bytes:
+        return brotli.compress(msgpack.dumps(self.state_dict()), quality=5)
+
+    @classmethod
+    def from_state_dict(cls, dct: dict[str, Any]) -> MolecularComplex:
+        dct = dict(dct)
+        for key, value in tuple(dct.items()):
+            if isinstance(value, list) and key in _SERIALIZED_ARRAYS:
+                dct[key] = np.asarray(value)
+        for key, value in tuple(dct.items()):
+            if not isinstance(value, np.ndarray):
+                continue
+            if key in {"atom_positions", "plddt"}:
+                dct[key] = value.astype(np.float32)
+            elif key == "token_to_atoms":
+                dct[key] = value.astype(np.int32)
+            elif key in {"chain_id", "entity_id", "sym_id"}:
+                dct[key] = value.astype(np.int64)
+        dct["metadata"] = MolecularComplexMetadata(**dct["metadata"])
+        if "chain_id" not in dct:
+            dct["chain_id"] = np.zeros(len(dct["sequence"]), dtype=np.int64)
+        return cls(**dct)
+
+    @classmethod
+    def from_blob(cls, input: Path | str | io.BytesIO | bytes) -> MolecularComplex:
+        if isinstance(input, (Path, str)):
+            payload = Path(input).read_bytes()
+        elif isinstance(input, io.BytesIO):
+            payload = input.getvalue()
+        else:
+            payload = input
+        state = msgpack.loads(brotli.decompress(payload), strict_map_key=False)
+        return cls.from_state_dict(state)

fastplms/models/esmfold2/esmfold2_msa.py

@@ -0,0 +1,577 @@
+"""Multiple-sequence-alignment value objects and lossless encodings."""
+
+from __future__ import annotations
+
+import dataclasses
+import string
+from collections.abc import Sequence
+from dataclasses import dataclass
+from functools import cached_property
+from itertools import islice
+from typing import Any
+
+import numpy as np
+from Bio import SeqIO
+from scipy.spatial.distance import cdist
+
+from .esmfold2_misc import slice_any_object
+from .esmfold2_msa_filter_sequences import greedy_select_indices, hhfilter
+from .esmfold2_parsing import FastaEntry, read_sequences, write_sequences
+from .esmfold2_sequential_dataclass import SequentialDataclass
+from .esmfold2_system import PathOrBuffer
+
+_A3M_INSERTION_DELETE_TABLE = str.maketrans(
+    dict.fromkeys(string.ascii_lowercase + ".")
+)
+_SERIALIZATION_VERSION = 1
+_UINT32_BYTES = 4
+
+
+def is_a3m_insertion(character: str) -> bool:
+    """Return whether a character is an A3M insertion marker."""
+
+    return character == "." or character.islower()
+
+
+def remove_insertions_from_sequence(sequence: str) -> str:
+    """Remove lowercase residues and dot insertion markers from an A3M row."""
+
+    return sequence.translate(_A3M_INSERTION_DELETE_TABLE)
+
+
+def a3m_deletion_counts(sequence: str) -> np.ndarray:
+    """Count insertions preceding each A3M match column."""
+
+    codes = np.frombuffer(sequence.encode("ascii"), dtype=np.uint8)
+    lowercase = (codes >= ord("a")) & (codes <= ord("z"))
+    insertion_mask = lowercase | (codes == ord("."))
+    prefix_counts = np.concatenate(([0], np.cumsum(insertion_mask)))
+    match_positions = np.flatnonzero(~insertion_mask)
+    return np.diff(prefix_counts[match_positions], prepend=0)
+
+
+def _parse_full_payload(data: bytes) -> tuple[np.ndarray, list[str]]:
+    version = int.from_bytes(data[:1], "little")
+    if version != _SERIALIZATION_VERSION:
+        raise ValueError(f"Unsupported version: {version}")
+    seqlen = int.from_bytes(data[1:5], "little")
+    depth = int.from_bytes(data[5:9], "little")
+    body = data[9:]
+    split = seqlen * depth
+    array = np.frombuffer(body[:split], dtype="|S1").reshape(depth, seqlen)
+    headers = [header for header in body[split:].decode().split("\n") if header]
+    if not headers and depth > 0:
+        headers = [""] * depth
+    return array, headers
+
+
+def _parse_sequence_payload(data: bytes) -> np.ndarray:
+    seqlen = int.from_bytes(data[:_UINT32_BYTES], "little")
+    return np.frombuffer(data[_UINT32_BYTES:], dtype="|S1").reshape(-1, seqlen)
+
+
+def _full_payload(array: np.ndarray, headers: Sequence[str]) -> bytes:
+    depth, seqlen = array.shape
+    prefix = b"".join(
+        (
+            _SERIALIZATION_VERSION.to_bytes(1, "little"),
+            seqlen.to_bytes(_UINT32_BYTES, "little"),
+            depth.to_bytes(_UINT32_BYTES, "little"),
+        )
+    )
+    return prefix + array.tobytes() + "\n".join(headers).encode()
+
+
+def _sequence_payload(array: np.ndarray) -> bytes:
+    return array.shape[1].to_bytes(_UINT32_BYTES, "little") + array.tobytes()
+
+
+def _random_row_indices(depth: int, count: int) -> np.ndarray:
+    sampled = np.random.choice(depth - 1, count - 1, replace=False) + 1
+    return np.sort(np.append(0, sampled))
+
+
+@dataclass(frozen=True)
+class FastMSA(SequentialDataclass):
+    """An MSA stored as a two-dimensional NumPy byte array."""
+
+    array: np.ndarray
+    headers: list[str] | None = None
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.array, np.ndarray):
+            raise TypeError("FastMSA array must be a NumPy array.")
+        if self.array.ndim != 2 or self.array.shape[0] == 0 or self.array.shape[1] == 0:
+            raise ValueError(
+                f"FastMSA array must have non-empty shape (depth, length), got {self.array.shape}."
+            )
+        if self.headers is not None and len(self.headers) != self.depth:
+            raise ValueError("Number of headers must match depth.")
+
+    @property
+    def depth(self) -> int:
+        return self.array.shape[0]
+
+    @property
+    def seqlen(self) -> int:
+        return self.array.shape[1]
+
+    def __len__(self) -> int:
+        return self.seqlen
+
+    @classmethod
+    def from_bytes(cls, data: bytes) -> FastMSA:
+        array, headers = _parse_full_payload(data)
+        return cls(array, headers)
+
+    @classmethod
+    def from_sequence_bytes(cls, data: bytes) -> FastMSA:
+        return cls(_parse_sequence_payload(data))
+
+    def __getitem__(
+        self,
+        indices: int | list[int] | slice | np.ndarray,
+    ) -> FastMSA:
+        column_indices = [indices] if isinstance(indices, int) else indices
+        return dataclasses.replace(self, array=self.array[:, column_indices])
+
+    def select_sequences(
+        self,
+        indices: Sequence[int] | np.ndarray,
+    ) -> FastMSA:
+        headers = None
+        if self.headers is not None:
+            headers = [self.headers[index] for index in indices]
+        return dataclasses.replace(
+            self,
+            array=self.array[indices],
+            headers=headers,
+        )
+
+    def select_random_sequences(self, num_seqs: int) -> FastMSA:
+        if num_seqs >= self.depth:
+            return self
+        return self.select_sequences(_random_row_indices(self.depth, num_seqs))
+
+    def pad_to_depth(self, depth: int) -> FastMSA:
+        if depth < self.depth:
+            raise ValueError(f"Cannot pad to depth {depth} when depth is {self.depth}")
+        if depth == self.depth:
+            return self
+        row_count = depth - self.depth
+        pad_value = ord("-") if self.array.dtype == np.uint8 else b"-"
+        array = np.pad(
+            self.array,
+            ((0, row_count), (0, 0)),
+            constant_values=pad_value,
+        )
+        headers = None if self.headers is None else self.headers + [""] * row_count
+        return dataclasses.replace(self, array=array, headers=headers)
+
+    @classmethod
+    def concat(
+        cls,
+        msas: Sequence[FastMSA],
+        join_token: str | None = None,
+        allow_depth_mismatch: bool = False,
+    ) -> FastMSA:
+        if not msas:
+            raise ValueError("Cannot concatenate an empty list of MSAs")
+        if join_token not in (None, ""):
+            raise NotImplementedError("join_token is not supported for FastMSA")
+        depths = [msa.depth for msa in msas]
+        if len(set(depths)) != 1:
+            if not allow_depth_mismatch:
+                raise ValueError("Depth mismatch in concatenating MSAs")
+            maximum_depth = max(depths)
+            msas = [msa.pad_to_depth(maximum_depth) for msa in msas]
+        header_columns = (
+            msa.headers if msa.headers is not None else [""] * msa.depth for msa in msas
+        )
+        headers = [
+            "|".join(str(header) for header in row) for row in zip(*header_columns, strict=False)
+        ]
+        return cls(
+            np.concatenate([msa.array for msa in msas], axis=1),
+            headers,
+        )
+
+    @classmethod
+    def stack(
+        cls,
+        msas: Sequence[FastMSA],
+        remove_query_from_later_msas: bool = True,
+    ) -> FastMSA:
+        if not msas:
+            raise ValueError("Cannot stack an empty list of MSAs")
+        arrays: list[np.ndarray] = []
+        headers: list[str] | None = [] if any(msa.headers is not None for msa in msas) else None
+        for index, msa in enumerate(msas):
+            start = 1 if index > 0 and remove_query_from_later_msas else 0
+            arrays.append(msa.array[start:])
+            if headers is not None:
+                source_headers = msa.headers or [""] * msa.depth
+                headers.extend(source_headers[start:])
+        return cls(np.concatenate(arrays, axis=0), headers)
+
+    def to_msa(self) -> MSA:
+        headers = self.headers
+        if headers is None:
+            headers = [f"seq{index}" for index in range(self.depth)]
+        entries = [
+            FastaEntry(header, b"".join(row).decode())
+            for header, row in zip(headers, self.array, strict=False)
+        ]
+        return MSA(entries)
+
+
+@dataclass(frozen=True)
+class MSA(SequentialDataclass):
+    """An ordered set of aligned protein sequences and optional A3M metadata."""
+
+    entries: list[FastaEntry]
+    deletions: np.ndarray | None = dataclasses.field(default=None, compare=False)
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.entries, list):
+            raise TypeError("MSA entries must be a list of FastaEntry rows.")
+        if not self.entries:
+            raise ValueError("MSA requires at least one aligned sequence.")
+        if any(not isinstance(entry, FastaEntry) for entry in self.entries):
+            raise TypeError("Every MSA entry must be a FastaEntry.")
+        expected_length = len(self.entries[0].sequence)
+        if expected_length == 0:
+            raise ValueError("MSA sequences must be non-empty.")
+        for row, entry in enumerate(self.entries[1:], start=1):
+            if len(entry.sequence) != expected_length:
+                raise ValueError(
+                    "MSA row length mismatch: "
+                    f"row 0 has {expected_length} columns, row {row} has "
+                    f"{len(entry.sequence)}."
+                )
+        deletions = self.deletions
+        if deletions is not None and not isinstance(deletions, np.ndarray):
+            raise TypeError("MSA deletions must be a NumPy array when provided.")
+        if isinstance(deletions, np.ndarray) and deletions.shape != (
+            len(self.entries),
+            expected_length,
+        ):
+            raise ValueError(
+                "MSA deletion matrix must have shape "
+                f"({len(self.entries)}, {expected_length}), got {deletions.shape}."
+            )
+
+    @cached_property
+    def sequences(self) -> list[str]:
+        return [entry.sequence for entry in self.entries]
+
+    @cached_property
+    def headers(self) -> list[str]:
+        return [entry.header for entry in self.entries]
+
+    @property
+    def depth(self) -> int:
+        return len(self.entries)
+
+    @property
+    def seqlen(self) -> int:
+        return len(self.entries[0].sequence)
+
+    @property
+    def query(self) -> str:
+        return self.entries[0].sequence
+
+    @cached_property
+    def array(self) -> np.ndarray:
+        return np.array([list(sequence) for sequence in self.sequences], dtype="|S1")
+
+    @cached_property
+    def seqid(self) -> np.ndarray:
+        byte_array = self.array.view(np.uint8)
+        return (1 - cdist(byte_array[0][None], byte_array, "hamming"))[0]
+
+    def __len__(self) -> int:
+        return self.seqlen
+
+    def __repr__(self) -> str:
+        return f"MSA({self.entries[0].header}: Depth={self.depth}, Length={self.seqlen})"
+
+    @classmethod
+    def from_a3m(
+        cls,
+        path: PathOrBuffer,
+        remove_insertions: bool = True,
+        max_sequences: int | None = None,
+    ) -> MSA:
+        entries = []
+        deletion_rows = []
+        for header, raw_sequence in islice(read_sequences(path), max_sequences):
+            if remove_insertions:
+                deletion_rows.append(a3m_deletion_counts(raw_sequence))
+            sequence = (
+                remove_insertions_from_sequence(raw_sequence) if remove_insertions else raw_sequence
+            )
+            if entries:
+                expected_length = len(entries[0].sequence)
+                if len(sequence) != expected_length:
+                    raise ValueError(
+                        "Sequence length mismatch. "
+                        f"Expected: {expected_length}, Received: {len(sequence)}"
+                    )
+            entries.append(FastaEntry(header, sequence))
+        deletions = None
+        if remove_insertions and deletion_rows:
+            deletions = np.stack(deletion_rows).astype(np.float32)
+        return cls(entries, deletions=deletions)
+
+    @classmethod
+    def from_stockholm(
+        cls,
+        path: PathOrBuffer,
+        remove_insertions: bool = True,
+        max_sequences: int | None = None,
+    ) -> MSA:
+        entries = []
+        for record in islice(SeqIO.parse(path, "stockholm"), max_sequences):
+            sequence = str(record.seq)
+            if entries:
+                expected_length = len(entries[0].sequence)
+                if len(sequence) != expected_length:
+                    raise ValueError(
+                        "Sequence length mismatch. "
+                        f"Expected: {expected_length}, Received: {len(sequence)}"
+                    )
+            entries.append(FastaEntry(f"{record.id} {record.description}", sequence))
+        msa = cls(entries)
+        if remove_insertions:
+            msa = msa.select_positions(
+                [index for index, residue in enumerate(msa.query) if residue != "-"]
+            )
+        return msa
+
+    @classmethod
+    def from_sequences(
+        cls,
+        sequences: list[str],
+        remove_insertions: bool = False,
+    ) -> MSA:
+        transform = (
+            remove_insertions_from_sequence if remove_insertions else lambda sequence: sequence
+        )
+        return cls([FastaEntry("", transform(sequence)) for sequence in sequences])
+
+    @classmethod
+    def from_bytes(cls, data: bytes) -> MSA:
+        array, headers = _parse_full_payload(data)
+        return cls(
+            [
+                FastaEntry(header, b"".join(row).decode())
+                for header, row in zip(headers, array, strict=False)
+            ]
+        )
+
+    @classmethod
+    def from_sequence_bytes(cls, data: bytes) -> MSA:
+        array = _parse_sequence_payload(data)
+        return cls([FastaEntry("", b"".join(row).decode()) for row in array])
+
+    @classmethod
+    def from_state_dict(cls, dct: dict[str, Any]) -> MSA:
+        deletions = dct.get("deletions")
+        return cls(
+            [FastaEntry("", sequence) for sequence in dct["sequences"]],
+            deletions=(None if deletions is None else np.asarray(deletions, dtype=np.float32)),
+        )
+
+    def to_a3m(self, path: PathOrBuffer) -> None:
+        write_sequences(self.entries, path)
+
+    def to_fast_msa(self) -> FastMSA:
+        return FastMSA(self.array, self.headers)
+
+    def to_bytes(self) -> bytes:
+        return _full_payload(self.array, self.headers)
+
+    def to_sequence_bytes(self) -> bytes:
+        """Serialize aligned sequences without their headers."""
+
+        return _sequence_payload(self.array)
+
+    def state_dict(self, json_serializable: bool = False) -> dict[str, Any]:
+        result: dict[str, Any] = {"sequences": self.sequences}
+        if self.deletions is not None:
+            result["deletions"] = self.deletions.tolist() if json_serializable else self.deletions
+        return result
+
+    def _aligned_deletions(self) -> np.ndarray | None:
+        if self.deletions is None:
+            return None
+        if self.deletions.shape != (self.depth, self.seqlen):
+            return None
+        return self.deletions
+
+    def _select_deletion_columns(self, indices) -> np.ndarray | None:
+        if self.deletions is None or self.deletions.shape[1] != self.seqlen:
+            return None
+        return self.deletions[:, indices]
+
+    def select_sequences(
+        self,
+        indices: Sequence[int] | np.ndarray,
+    ) -> MSA:
+        deletions = None if self.deletions is None else self.deletions[np.asarray(indices)]
+        return dataclasses.replace(
+            self,
+            entries=[self.entries[index] for index in indices],
+            deletions=deletions,
+        )
+
+    def select_positions(
+        self,
+        indices: Sequence[int] | np.ndarray,
+    ) -> MSA:
+        entries = [
+            FastaEntry(
+                entry.header,
+                "".join(entry.sequence[index] for index in indices),
+            )
+            for entry in self.entries
+        ]
+        return dataclasses.replace(
+            self,
+            entries=entries,
+            deletions=self._select_deletion_columns(indices),
+        )
+
+    def __getitem__(
+        self,
+        indices: int | list[int] | slice | np.ndarray,
+    ) -> MSA:
+        column_indices = [indices] if isinstance(indices, int) else indices
+        entries = [
+            FastaEntry(
+                entry.header,
+                slice_any_object(entry.sequence, column_indices),
+            )
+            for entry in self.entries
+        ]
+        return dataclasses.replace(
+            self,
+            entries=entries,
+            deletions=self._select_deletion_columns(column_indices),
+        )
+
+    def greedy_select(self, num_seqs: int, mode: str = "max") -> MSA:
+        if mode not in ("max", "min"):
+            raise ValueError(f"Unsupported MSA selection mode: {mode!r}.")
+        if self.depth <= num_seqs:
+            return self
+        return self.select_sequences(greedy_select_indices(self.array, num_seqs, mode))
+
+    def hhfilter(
+        self,
+        seqid: int = 90,
+        diff: int = 0,
+        cov: int = 0,
+        qid: int = 0,
+        qsc: float = -20.0,
+        binary: str = "hhfilter",
+    ) -> MSA:
+        indices = hhfilter(
+            self.sequences,
+            seqid=seqid,
+            diff=diff,
+            cov=cov,
+            qid=qid,
+            qsc=qsc,
+            binary=binary,
+        )
+        return self.select_sequences(indices)
+
+    def select_random_sequences(self, num_seqs: int) -> MSA:
+        if num_seqs >= self.depth:
+            return self
+        return self.select_sequences(_random_row_indices(self.depth, num_seqs))
+
+    def select_diverse_sequences(self, num_seqs: int) -> MSA:
+        if num_seqs >= self.depth:
+            return self
+        filtered = self.hhfilter(diff=num_seqs)
+        if num_seqs < filtered.depth:
+            filtered = filtered.select_random_sequences(num_seqs)
+        return filtered
+
+    def pad_to_depth(self, depth: int) -> MSA:
+        if depth < self.depth:
+            raise ValueError(f"Cannot pad to depth {depth} when depth is {self.depth}")
+        if depth == self.depth:
+            return self
+        count = depth - self.depth
+        extra = [FastaEntry("", "-" * self.seqlen) for _ in range(count)]
+        deletions = self._aligned_deletions()
+        if deletions is not None:
+            zero_rows = np.zeros((count, self.seqlen), dtype=deletions.dtype)
+            deletions = np.concatenate((deletions, zero_rows), axis=0)
+        return dataclasses.replace(
+            self,
+            entries=self.entries + extra,
+            deletions=deletions,
+        )
+
+    @classmethod
+    def stack(
+        cls,
+        msas: Sequence[MSA],
+        remove_query_from_later_msas: bool = True,
+    ) -> MSA:
+        entries = []
+        deletion_arrays = []
+        for index, msa in enumerate(msas):
+            start = 1 if index > 0 and remove_query_from_later_msas else 0
+            entries.extend(msa.entries[start:])
+            aligned = msa._aligned_deletions()
+            if aligned is not None:
+                deletion_arrays.append(aligned[start:])
+        deletions = None
+        if (
+            len(deletion_arrays) == len(msas)
+            and len({array.shape[1] for array in deletion_arrays}) == 1
+        ):
+            deletions = np.concatenate(deletion_arrays, axis=0)
+        return cls(entries=entries, deletions=deletions)
+
+    @classmethod
+    def concat(
+        cls,
+        msas: Sequence[MSA],
+        join_token: str | None = "|",
+        allow_depth_mismatch: bool = False,
+    ) -> MSA:
+        if not msas:
+            raise ValueError("Cannot concatenate an empty list of MSAs")
+        depths = [msa.depth for msa in msas]
+        if len(set(depths)) != 1:
+            if not allow_depth_mismatch:
+                raise ValueError("Depth mismatch in concatenating MSAs")
+            maximum_depth = max(depths)
+            msas = [msa.pad_to_depth(maximum_depth) for msa in msas]
+        headers = [
+            "|".join(str(header) for header in row)
+            for row in zip(*(msa.headers for msa in msas), strict=False)
+        ]
+        separator = "" if join_token is None else join_token
+        sequences = [
+            separator.join(row) for row in zip(*(msa.sequences for msa in msas), strict=False)
+        ]
+        deletions = None
+        if separator == "":
+            arrays = [msa._aligned_deletions() for msa in msas]
+            if all(array is not None for array in arrays):
+                deletions = np.concatenate(arrays, axis=1)  # type: ignore[arg-type]
+        return cls(
+            [
+                FastaEntry(header, sequence)
+                for header, sequence in zip(headers, sequences, strict=False)
+            ],
+            deletions=deletions,
+        )

fastplms/models/esmfold2/esmfold2_msa_filter_sequences.py

@@ -0,0 +1,111 @@
+"""Sequence selection for multiple-sequence alignments."""
+
+from __future__ import annotations
+
+import os
+import tempfile
+from pathlib import Path
+
+import numpy as np
+
+from .esmfold2_system import run_subprocess_with_errorcheck
+
+
+def _byte_matrix(array: np.ndarray) -> np.ndarray:
+    """Return a two-dimensional byte view used for Hamming comparisons."""
+
+    matrix = np.asarray(array).view(np.uint8)
+    return matrix.reshape(matrix.shape[0], -1)
+
+
+def _hamming_to_all(query: np.ndarray, sequences: np.ndarray) -> np.ndarray:
+    return np.not_equal(sequences, query).mean(axis=1, dtype=np.float64)
+
+
+def greedy_select_indices(array: np.ndarray, num_seqs: int, mode: str = "max") -> list[int]:
+    """Select MSA rows by greedy mean Hamming distance from the query row.
+
+    Row zero is always retained. At each step the selector chooses the remaining
+    row with greatest distance for ``mode="max"`` or least distance for
+    ``mode="min"``. Returned indices follow source order.
+    """
+
+    if not isinstance(array, np.ndarray):
+        raise TypeError("array must be a NumPy array")
+    if array.ndim != 2 or array.shape[0] == 0 or array.shape[1] == 0:
+        raise ValueError(
+            f"array must have non-empty shape (depth, length), got {array.shape}"
+        )
+    if isinstance(num_seqs, bool) or not isinstance(num_seqs, int):
+        raise TypeError("num_seqs must be an integer")
+    if num_seqs <= 0:
+        raise ValueError("num_seqs must be greater than zero")
+    if not isinstance(mode, str):
+        raise TypeError("mode must be a string")
+    if mode not in {"max", "min"}:
+        raise ValueError(f"unsupported selection mode: {mode}")
+    depth = array.shape[0]
+    if depth <= num_seqs:
+        return list(range(depth))
+
+    sequences = _byte_matrix(array)
+    selected = [0]
+    available = np.ones(depth, dtype=bool)
+    available[0] = False
+    distance_sum = _hamming_to_all(sequences[0], sequences)
+    choose = np.argmax if mode == "max" else np.argmin
+
+    while len(selected) < num_seqs:
+        candidates = np.flatnonzero(available)
+        candidate_scores = distance_sum[candidates] / len(selected)
+        next_index = int(candidates[int(choose(candidate_scores))])
+        selected.append(next_index)
+        available[next_index] = False
+        distance_sum += _hamming_to_all(sequences[next_index], sequences)
+    return sorted(selected)
+
+
+def _temporary_root() -> str | None:
+    shared_memory = Path("/dev/shm")
+    return os.fspath(shared_memory) if shared_memory.is_dir() else None
+
+
+def hhfilter(
+    sequences: list[str],
+    seqid: int = 90,
+    diff: int = 0,
+    cov: int = 0,
+    qid: int = 0,
+    qsc: float = -20.0,
+    binary: str = "hhfilter",
+) -> list[int]:
+    """Run HH-suite filtering and return source indices from its FASTA headers."""
+
+    with tempfile.TemporaryDirectory(dir=_temporary_root()) as directory:
+        work = Path(directory)
+        source_path = work / "input.fasta"
+        result_path = work / "output.fasta"
+        records = (f">{index}\n{sequence}" for index, sequence in enumerate(sequences))
+        source_path.write_text("\n".join(records), encoding="utf-8")
+        command = [
+            binary,
+            "-i",
+            os.fspath(source_path),
+            "-M",
+            "a3m",
+            "-o",
+            os.fspath(result_path),
+            "-id",
+            str(seqid),
+            "-diff",
+            str(diff),
+            "-cov",
+            str(cov),
+            "-qid",
+            str(qid),
+            "-qsc",
+            str(qsc),
+        ]
+        run_subprocess_with_errorcheck(command, capture_output=True)
+        headers = result_path.read_text(encoding="utf-8").splitlines()
+        return [int(line[1:].strip()) for line in headers if line.startswith(">")]

fastplms/models/esmfold2/esmfold2_normalize_coordinates.py

@@ -0,0 +1,67 @@
+"""Rigid-frame normalization for atom37 coordinates."""
+
+from __future__ import annotations
+
+from typing import TypeVar
+
+import numpy as np
+import torch
+from torch import Tensor
+
+from . import esmfold2_residue_constants as residue_constants
+from .esmfold2_affine3d import Affine3D
+
+ArrayOrTensor = TypeVar("ArrayOrTensor", np.ndarray, Tensor)
+
+
+def atom3_to_backbone_frames(bb_positions: Tensor) -> Affine3D:
+    """Construct a frame from N, C-alpha, and C positions in ``X``."""
+
+    n_position, ca_position, c_position = bb_positions.unbind(dim=-2)
+    return Affine3D.from_graham_schmidt(c_position, ca_position, n_position)
+
+
+def index_by_atom_name(
+    atom37: ArrayOrTensor,
+    atom_names: str | list[str],
+    dim: int = -2,
+) -> ArrayOrTensor:
+    """Select one or more named atoms along an atom37 axis."""
+
+    single_atom = isinstance(atom_names, str)
+    names = [atom_names] if single_atom else atom_names
+    indices = [residue_constants.atom_order[name] for name in names]
+    axis = dim % atom37.ndim
+    if isinstance(atom37, Tensor):
+        index = torch.tensor(indices, dtype=torch.long, device=atom37.device)
+        selected = torch.index_select(atom37, axis, index)
+    else:
+        selected = np.take(atom37, indices, axis=axis)
+    return selected.squeeze(axis) if single_atom else selected  # type: ignore[return-value]
+
+
+def get_protein_normalization_frame(coords: Tensor) -> Affine3D:
+    """Build one frame from backbone coordinates ``X`` with shape (l, 37, 3)."""
+
+    backbone = index_by_atom_name(coords, ["N", "CA", "C"], dim=-2)
+    residue_is_valid = torch.isfinite(backbone).all(dim=-1).all(dim=-1)
+    weights = residue_is_valid[..., None, None]
+    coordinate_sum = backbone.masked_fill(~weights, 0).sum(dim=-3)
+    count = residue_is_valid.sum(dim=-1)[..., None, None]
+    mean_backbone = coordinate_sum / (count + 1e-8)
+    return atom3_to_backbone_frames(mean_backbone.float())
+
+
+def apply_frame_to_coords(coords: Tensor, frame: Affine3D) -> Tensor:
+    """Express atom coordinates ``X`` in the inverse of ``frame``."""
+
+    transformed = frame[..., None, None].invert().apply(coords)
+    frame_is_valid = frame.trans.norm(dim=-1) > 0
+    normalized = torch.where(frame_is_valid[..., None, None, None], transformed, coords)
+    return normalized.masked_fill(torch.isinf(coords), torch.inf)
+
+
+def normalize_coordinates(coords: Tensor) -> Tensor:
+    """Normalize ``X`` with shape (..., l, 37, 3) to its backbone frame."""
+
+    return apply_frame_to_coords(coords, get_protein_normalization_frame(coords))

fastplms/models/esmfold2/esmfold2_output.py

@@ -0,0 +1,201 @@
+"""Convert ESMFold2 coordinate tensors into molecular-complex records."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from dataclasses import dataclass, field
+from itertools import groupby
+from typing import Any
+
+import numpy as np
+import torch
+
+from .esmfold2_constants import ELEMENT_NUMBER_TO_SYMBOL, MOL_TYPE_NONPOLYMER
+from .esmfold2_molecular_complex import MolecularComplex, MolecularComplexMetadata
+
+
+def get_element_symbol(atomic_number: int) -> str:
+    """Map a training-time atomic number to a chemical symbol."""
+
+    return ELEMENT_NUMBER_TO_SYMBOL.get(atomic_number, "X")
+
+
+def _decode_atom_name(encoded_name: Any) -> str:
+    values = encoded_name.tolist() if hasattr(encoded_name, "tolist") else encoded_name
+    return "".join(chr(int(value) + 32) for value in values if int(value)).strip()
+
+
+@dataclass
+class _ComplexRecords:
+    sequence: list[str] = field(default_factory=list)
+    chain_ids: list[int] = field(default_factory=list)
+    token_to_atoms: list[list[int]] = field(default_factory=list)
+    confidence: list[float] = field(default_factory=list)
+    positions: list[list[float]] = field(default_factory=list)
+    elements: list[str] = field(default_factory=list)
+    atom_names: list[str] = field(default_factory=list)
+    atom_hetero: list[bool] = field(default_factory=list)
+    chain_lookup: dict[int, str] = field(default_factory=dict)
+    entity_lookup: dict[int, str] = field(default_factory=dict)
+
+    def add_token(
+        self,
+        *,
+        residue_name: str,
+        asym_id: int,
+        plddt: float,
+        atoms: Iterable[tuple[list[float], str, str]],
+        hetero: bool,
+    ) -> None:
+        atom_start = len(self.positions)
+        for position, element, atom_name in atoms:
+            self.positions.append(position)
+            self.elements.append(element)
+            self.atom_names.append(atom_name)
+            self.atom_hetero.append(hetero)
+        self.sequence.append(residue_name)
+        self.chain_ids.append(asym_id)
+        self.confidence.append(plddt)
+        self.token_to_atoms.append([atom_start, len(self.positions)])
+
+    def build(self, complex_id: str) -> MolecularComplex:
+        return MolecularComplex(
+            id=complex_id,
+            sequence=self.sequence,
+            atom_positions=np.asarray(self.positions, dtype=np.float32).reshape(-1, 3),
+            atom_elements=np.asarray(self.elements, dtype=object),
+            token_to_atoms=np.asarray(self.token_to_atoms, dtype=np.int32).reshape(-1, 2),
+            chain_id=np.asarray(self.chain_ids, dtype=np.int64),
+            plddt=np.asarray(self.confidence, dtype=np.float32),
+            atom_names=np.asarray(self.atom_names, dtype=object),
+            atom_hetero=np.asarray(self.atom_hetero, dtype=bool),
+            metadata=MolecularComplexMetadata(
+                entity_lookup=self.entity_lookup,
+                chain_lookup=self.chain_lookup,
+                assembly_composition=None,
+            ),
+        )
+
+
+def build_molecular_complex_from_features(
+    coords: torch.Tensor,
+    plddt: torch.Tensor,
+    atom_mask: torch.Tensor,
+    ref_element: torch.Tensor,
+    ref_atom_name_chars: torch.Tensor,
+    chain_infos: list[Any],
+    complex_id: str,
+) -> MolecularComplex:
+    """Decode model features into one complex without intermediate structure files.
+
+    Protein, DNA, and RNA tokens are grouped by residue index. Ligand atom
+    tokens are collapsed into one non-polymer residue per chain.
+    """
+
+    M = atom_mask.bool().cpu().numpy()
+    X = coords.float().cpu().numpy()
+    atom_names = ref_atom_name_chars.cpu().numpy()
+    elements = ref_element.cpu().numpy()
+    confidence = plddt.float().cpu().numpy()
+    records = _ComplexRecords()
+
+    def decode_atoms(tokens: Iterable[Any]):
+        for token in tokens:
+            for atom_index in range(token.atom_start, token.atom_start + token.atom_count):
+                if M[atom_index]:
+                    yield (
+                        X[atom_index].tolist(),
+                        get_element_symbol(int(elements[atom_index])),
+                        _decode_atom_name(atom_names[atom_index]),
+                    )
+
+    for chain in chain_infos:
+        is_nonpolymer = chain.mol_type == MOL_TYPE_NONPOLYMER
+        records.chain_lookup[chain.asym_id] = chain.chain_id
+        records.entity_lookup[chain.entity_id] = "non-polymer" if is_nonpolymer else "polymer"
+
+        if is_nonpolymer:
+            mean_confidence = (
+                float(np.mean([confidence[token.token_index] for token in chain.tokens]))
+                if chain.tokens
+                else 0.0
+            )
+            records.add_token(
+                residue_name=chain.tokens[0].residue_name if chain.tokens else "LIG",
+                asym_id=chain.asym_id,
+                plddt=mean_confidence,
+                atoms=decode_atoms(chain.tokens),
+                hetero=True,
+            )
+            continue
+
+        residue_groups = groupby(chain.tokens, key=lambda token: token.residue_index)
+        for _residue_index, group in residue_groups:
+            residue_tokens = list(group)
+            records.add_token(
+                residue_name=residue_tokens[0].residue_name,
+                asym_id=chain.asym_id,
+                plddt=float(np.mean([confidence[token.token_index] for token in residue_tokens])),
+                atoms=decode_atoms(residue_tokens),
+                hetero=False,
+            )
+
+    return records.build(complex_id)
+
+
+def build_molecular_complex(
+    structure: Any,
+    coords: torch.Tensor,
+    plddt: torch.Tensor,
+    complex_id: str,
+) -> MolecularComplex:
+    """Decode coordinates using the atom and residue arrays of a prepared structure."""
+
+    records = _ComplexRecords()
+    coordinate_index = 0
+    confidence_index = 0
+
+    for chain in structure.chains:
+        asym_id = int(chain["asym_id"])
+        mol_type = int(chain["mol_type"])
+        is_nonpolymer = mol_type == MOL_TYPE_NONPOLYMER
+        records.chain_lookup[asym_id] = str(chain["name"])
+        records.entity_lookup[int(chain["entity_id"])] = (
+            "non-polymer" if is_nonpolymer else "polymer"
+        )
+
+        residue_start = int(chain["res_idx"])
+        residue_stop = residue_start + int(chain["res_num"])
+        for residue in structure.residues[residue_start:residue_stop]:
+            atom_start = int(residue["atom_idx"])
+            atom_stop = atom_start + int(residue["atom_num"])
+            decoded_atoms: list[tuple[list[float], str, str]] = []
+            for atom in structure.atoms[atom_start:atom_stop]:
+                if not atom["is_present"]:
+                    continue
+                decoded_atoms.append(
+                    (
+                        coords[coordinate_index].tolist(),
+                        get_element_symbol(int(atom["element"].item())),
+                        _decode_atom_name(atom["name"]),
+                    )
+                )
+                coordinate_index += 1
+
+            records.add_token(
+                residue_name=str(residue["name"]),
+                asym_id=asym_id,
+                plddt=float(plddt[confidence_index].item()),
+                atoms=decoded_atoms,
+                hetero=is_nonpolymer,
+            )
+            confidence_index += 1
+
+    return records.build(complex_id)
+
+
+__all__ = [
+    "build_molecular_complex",
+    "build_molecular_complex_from_features",
+    "get_element_symbol",
+]

fastplms/models/esmfold2/esmfold2_paired_msa.py

@@ -0,0 +1,282 @@
+"""Construct taxonomy-paired MSA features for multichain folding."""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+
+import numpy as np
+
+from .esmfold2_constants import (
+    MSA_GAP_TOKEN_ID,
+    PROTEIN_3TO1,
+    PROTEIN_RESIDUE_TO_RES_TYPE,
+    PROTEIN_UNK_RES_TYPE,
+)
+from .esmfold2_msa import MSA
+
+_TAXONOMY_PATTERN = re.compile(r"key=(-?\d+)")
+
+
+def protein_letter_to_res_type() -> dict[str, int]:
+    """Return the one-letter residue vocabulary used by the MSA encoder."""
+
+    vocabulary = {
+        one_letter: PROTEIN_RESIDUE_TO_RES_TYPE[three_letter]
+        for three_letter, one_letter in PROTEIN_3TO1.items()
+        if three_letter in PROTEIN_RESIDUE_TO_RES_TYPE
+    }
+    vocabulary.update({"-": MSA_GAP_TOKEN_ID, "X": PROTEIN_UNK_RES_TYPE})
+    return vocabulary
+
+
+def _taxonomy_from_header(header: str) -> int:
+    match = _TAXONOMY_PATTERN.search(header) if header else None
+    return int(match.group(1)) if match is not None else -1
+
+
+def _emitted_length(sequence: str) -> int:
+    return sum(character != "." and not character.islower() for character in sequence)
+
+
+def _decode_a3m_row(
+    sequence: str,
+    sequence_length: int,
+    vocabulary: dict[str, int],
+) -> tuple[np.ndarray, np.ndarray]:
+    residues = np.full(sequence_length, MSA_GAP_TOKEN_ID, dtype=np.int64)
+    deletions = np.zeros(sequence_length, dtype=np.float32)
+    column = 0
+    insertion_count = 0
+    for character in sequence:
+        if character == "." or character.islower():
+            insertion_count += 1
+            continue
+        if column == sequence_length:
+            break
+        residues[column] = (
+            MSA_GAP_TOKEN_ID
+            if character == "-"
+            else vocabulary.get(character.upper(), PROTEIN_UNK_RES_TYPE)
+        )
+        if insertion_count:
+            deletions[column] = float(insertion_count)
+            insertion_count = 0
+        column += 1
+    return residues, deletions
+
+
+def msa_to_res_type_and_deletions(
+    msa: MSA,
+    letter_to_res_type: dict[str, int],
+) -> tuple[np.ndarray, np.ndarray]:
+    """Decode an A3M alignment into arrays ``X`` and ``D`` with shape (m, l)."""
+
+    sequence_length = _emitted_length(msa.entries[0].sequence)
+    residue_rows: list[np.ndarray] = []
+    deletion_rows: list[np.ndarray] = []
+    for entry in msa.entries:
+        residues, deletions = _decode_a3m_row(
+            entry.sequence,
+            sequence_length,
+            letter_to_res_type,
+        )
+        residue_rows.append(residues)
+        deletion_rows.append(deletions)
+    return np.stack(residue_rows), np.stack(deletion_rows)
+
+
+@dataclass(frozen=True)
+class _ChainAlignment:
+    residues: np.ndarray
+    deletions: np.ndarray
+    taxonomies: list[int]
+
+
+def _chain_alignment(
+    msa: MSA | None,
+    query_res_types: np.ndarray,
+    vocabulary: dict[str, int],
+) -> _ChainAlignment:
+    if msa is None or msa.depth == 0:
+        return _ChainAlignment(
+            residues=query_res_types[None, :],
+            deletions=np.zeros((1, query_res_types.shape[0]), dtype=np.float32),
+            taxonomies=[-1],
+        )
+    residues, deletions = msa_to_res_type_and_deletions(msa, vocabulary)
+    taxonomies = [_taxonomy_from_header(entry.header) for entry in msa.entries]
+    return _ChainAlignment(residues, deletions, taxonomies)
+
+
+def _taxonomy_groups(
+    chain_ids: list[int],
+    alignments: dict[int, _ChainAlignment],
+) -> dict[int, list[tuple[int, int]]]:
+    groups: dict[int, list[tuple[int, int]]] = {}
+    for chain_id in chain_ids:
+        for row, taxonomy in enumerate(alignments[chain_id].taxonomies):
+            if row and taxonomy != -1:
+                groups.setdefault(taxonomy, []).append((chain_id, row))
+    return {taxonomy: rows for taxonomy, rows in groups.items() if len(rows) > 1}
+
+
+def _available_rows(
+    chain_ids: list[int],
+    alignments: dict[int, _ChainAlignment],
+    groups: dict[int, list[tuple[int, int]]],
+) -> dict[int, list[int]]:
+    used = {row for group in groups.values() for row in group}
+    return {
+        chain_id: [
+            row
+            for row in range(1, len(alignments[chain_id].taxonomies))
+            if (chain_id, row) not in used
+        ]
+        for chain_id in chain_ids
+    }
+
+
+def _append_taxonomy_rows(
+    rows: list[dict[int, int]],
+    paired_flags: list[dict[int, int]],
+    chain_ids: list[int],
+    groups: dict[int, list[tuple[int, int]]],
+    available: dict[int, list[int]],
+    max_pairs: int,
+) -> None:
+    ordered_groups = sorted(
+        groups.values(),
+        key=lambda group: len({chain_id for chain_id, _row in group}),
+        reverse=True,
+    )
+    for group in ordered_groups:
+        rows_by_chain: dict[int, list[int]] = {}
+        for chain_id, row in group:
+            rows_by_chain.setdefault(chain_id, []).append(row)
+        for occurrence in range(max(map(len, rows_by_chain.values()))):
+            selected: dict[int, int] = {}
+            flags: dict[int, int] = {}
+            for chain_id, candidates in rows_by_chain.items():
+                selected[chain_id] = candidates[occurrence % len(candidates)]
+                flags[chain_id] = 1
+            for chain_id in chain_ids:
+                if chain_id not in selected:
+                    flags[chain_id] = 0
+                    selected[chain_id] = available[chain_id].pop(0) if available[chain_id] else -1
+            rows.append(selected)
+            paired_flags.append(flags)
+            if len(rows) >= max_pairs:
+                break
+        if len(rows) >= max_pairs:
+            break
+
+
+def _append_unpaired_rows(
+    rows: list[dict[int, int]],
+    paired_flags: list[dict[int, int]],
+    chain_ids: list[int],
+    available: dict[int, list[int]],
+    max_total: int,
+) -> None:
+    max_remaining = max((len(indices) for indices in available.values()), default=0)
+    for _ in range(min(max_total - len(rows), max_remaining)):
+        rows.append(
+            {
+                chain_id: available[chain_id].pop(0) if available[chain_id] else -1
+                for chain_id in chain_ids
+            }
+        )
+        paired_flags.append({chain_id: 0 for chain_id in chain_ids})
+        if len(rows) >= max_total:
+            break
+
+
+def _pairing_plan(
+    chain_ids: list[int],
+    alignments: dict[int, _ChainAlignment],
+    max_pairs: int,
+    max_total: int,
+    max_seqs: int,
+) -> tuple[list[dict[int, int]], list[dict[int, int]]]:
+    groups = _taxonomy_groups(chain_ids, alignments)
+    available = _available_rows(chain_ids, alignments, groups)
+    rows = [{chain_id: 0 for chain_id in chain_ids}]
+    flags = [{chain_id: 1 for chain_id in chain_ids}]
+    _append_taxonomy_rows(rows, flags, chain_ids, groups, available, max_pairs)
+    _append_unpaired_rows(rows, flags, chain_ids, available, max_total)
+    return rows[:max_seqs], flags[:max_seqs]
+
+
+def _project_alignment_rows(
+    chain_ids: list[int],
+    alignments: dict[int, _ChainAlignment],
+    rows: list[dict[int, int]],
+    flags: list[dict[int, int]],
+    token_asym_ids: np.ndarray,
+    token_res_ids: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    m, t = len(rows), len(token_asym_ids)
+    residues = np.full((m, t), MSA_GAP_TOKEN_ID, dtype=np.int64)
+    deletions = np.zeros((m, t), dtype=np.float32)
+    paired_mask = np.zeros((m, t), dtype=np.float32)
+    for chain_id in chain_ids:
+        alignment = alignments[chain_id]
+        selected_rows = np.asarray([row[chain_id] for row in rows], dtype=np.int64)
+        chain_flags = np.asarray([row[chain_id] for row in flags], dtype=np.float32)
+        token_mask = token_asym_ids == chain_id
+        if not token_mask.any():
+            continue
+        columns = np.minimum(token_res_ids[token_mask], alignment.residues.shape[1] - 1)
+        valid_rows = selected_rows >= 0
+        if valid_rows.any():
+            output_rows = np.flatnonzero(valid_rows)
+            output_columns = np.flatnonzero(token_mask)
+            residues[np.ix_(output_rows, output_columns)] = alignment.residues[
+                selected_rows[valid_rows]
+            ][:, columns]
+            deletions[np.ix_(output_rows, output_columns)] = alignment.deletions[
+                selected_rows[valid_rows]
+            ][:, columns]
+        paired_mask[:, token_mask] = chain_flags[:, None]
+    return residues, deletions, paired_mask
+
+
+def construct_paired_msa(
+    chain_msas: dict[int, MSA | None],
+    chain_query_res_types: dict[int, np.ndarray],
+    token_asym_ids: np.ndarray,
+    token_res_ids: np.ndarray,
+    letter_to_res_type: dict[str, int] | None = None,
+    *,
+    max_pairs: int = 8192,
+    max_total: int = 16384,
+    max_seqs: int = 16384,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+    """Return residue, deletion, and pairing arrays with shape (m, t)."""
+
+    vocabulary = protein_letter_to_res_type() if letter_to_res_type is None else letter_to_res_type
+    chain_ids = sorted(chain_msas)
+    alignments = {
+        chain_id: _chain_alignment(
+            chain_msas[chain_id],
+            chain_query_res_types[chain_id],
+            vocabulary,
+        )
+        for chain_id in chain_ids
+    }
+    rows, flags = _pairing_plan(
+        chain_ids,
+        alignments,
+        max_pairs,
+        max_total,
+        max_seqs,
+    )
+    return _project_alignment_rows(
+        chain_ids,
+        alignments,
+        rows,
+        flags,
+        token_asym_ids,
+        token_res_ids,
+    )

fastplms/models/esmfold2/esmfold2_parsing.py

@@ -0,0 +1,126 @@
+"""FASTA parsing and writing with explicit stream ownership."""
+
+from __future__ import annotations
+
+import gzip
+import io
+from collections.abc import Generator, Iterable
+from contextlib import nullcontext
+from pathlib import Path
+from typing import NamedTuple, TextIO
+
+from .esmfold2_utils_types import PathOrBuffer
+
+
+class FastaEntry(NamedTuple):
+    """One FASTA record in source order."""
+
+    header: str
+    sequence: str
+
+
+def parse_fasta(text: str) -> Generator[FastaEntry, None, None]:
+    """Yield records from FASTA text without normalizing sequence symbols."""
+
+    header: str | None = None
+    sequence_lines: list[str] = []
+    found_record = False
+
+    for line in text.splitlines():
+        if not line or line.startswith("#"):
+            continue
+        if line.startswith(">"):
+            if header is not None:
+                found_record = True
+                yield FastaEntry(header, "".join(sequence_lines))
+            header = line[1:].strip()
+            sequence_lines.clear()
+        elif header is not None:
+            sequence_lines.append(line)
+
+    if header is not None:
+        found_record = True
+        yield FastaEntry(header, "".join(sequence_lines))
+    if not found_record:
+        raise ValueError("Found no sequences in input")
+
+
+def _open_reader(source: PathOrBuffer):
+    if isinstance(source, io.TextIOBase):
+        return nullcontext(source)
+    path = Path(source)
+    if path.suffix.lower() == ".gz":
+        return gzip.open(path, mode="rt", encoding="utf-8")
+    return path.open(mode="r", encoding="utf-8")
+
+
+def read_sequences(source: PathOrBuffer) -> Generator[FastaEntry, None, None]:
+    """Read FASTA records while leaving caller-owned streams open."""
+
+    with _open_reader(source) as handle:
+        yield from parse_fasta(handle.read())
+
+
+def read_first_sequence(source: PathOrBuffer) -> FastaEntry:
+    """Return the first FASTA record from a path or text stream."""
+
+    return next(read_sequences(source))
+
+
+def count_fasta_sequences(path: str | Path) -> int:
+    """Count FASTA headers without parsing sequence bodies."""
+
+    source = Path(path)
+    if not source.exists():
+        return 0
+    with source.open(encoding="utf-8") as handle:
+        return sum(line.startswith(">") for line in handle)
+
+
+def append_fasta_sequence(header: str, sequence: str, path: str | Path) -> None:
+    """Append one record, inserting a separator if the file lacks a final newline."""
+
+    destination = Path(path)
+    destination.parent.mkdir(parents=True, exist_ok=True)
+    needs_separator = (
+        destination.exists()
+        and destination.stat().st_size > 0
+        and destination.read_bytes()[-1:] != b"\n"
+    )
+    with destination.open(mode="a", encoding="utf-8") as handle:
+        if needs_separator:
+            handle.write("\n")
+        handle.write(f">{header}\n{sequence}\n")
+
+
+def _open_writer(destination: PathOrBuffer):
+    if isinstance(destination, io.TextIOBase):
+        return nullcontext(destination)
+    path = Path(destination)
+    path.parent.mkdir(parents=True, exist_ok=True)
+    return path.open(mode="w", encoding="utf-8")
+
+
+def write_sequences(sequences: Iterable[tuple[str, str]], destination: PathOrBuffer) -> None:
+    """Write records with one blank-line-free separator between entries."""
+
+    with _open_writer(destination) as handle:
+        _write_records(handle, sequences)
+
+
+def _write_records(handle: TextIO, sequences: Iterable[tuple[str, str]]) -> None:
+    for index, (header, sequence) in enumerate(sequences):
+        if index:
+            handle.write("\n")
+        handle.write(f">{header}\n{sequence}")
+
+
+__all__ = [
+    "FastaEntry",
+    "append_fasta_sequence",
+    "count_fasta_sequences",
+    "parse_fasta",
+    "read_first_sequence",
+    "read_sequences",
+    "write_sequences",
+]

fastplms/models/esmfold2/esmfold2_predicted_aligned_error.py

@@ -0,0 +1,127 @@
+"""Predicted-aligned-error scores and training loss."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+from .esmfold2_affine3d import Affine3D
+
+_CPU_DEVICE = torch.device("cpu")
+
+
+def _compute_pae_masks(mask: Tensor) -> Tensor:
+    residue_mask = mask.bool()
+    return residue_mask.unsqueeze(-1) & residue_mask.unsqueeze(-2)
+
+
+def _pae_bins(
+    max_bin: float = 31,
+    num_bins: int = 64,
+    device: torch.device = _CPU_DEVICE,
+) -> Tensor:
+    """Return the representative distance for each PAE probability bin."""
+
+    boundaries = torch.linspace(0, max_bin, steps=num_bins - 1, device=device)
+    width = max_bin / (num_bins - 2)
+    centers = boundaries + width / 2
+    overflow_center = centers[-1:] + width
+    return torch.cat((centers, overflow_center))
+
+
+def _masked_probabilities(logits: Tensor, pair_mask: Tensor) -> Tensor:
+    masked_logits = logits.masked_fill(
+        ~pair_mask.unsqueeze(-1),
+        torch.finfo(logits.dtype).min,
+    )
+    return masked_logits.softmax(dim=-1)
+
+
+def masked_mean(
+    mask: Tensor,
+    value: Tensor,
+    dim: int | tuple[int, ...] | None = None,
+    eps: float = 1e-10,
+) -> Tensor:
+    """Average values over true entries of a broadcast-compatible mask."""
+
+    weights = mask.expand_as(value)
+    weighted_sum = torch.sum(weights * value, dim=dim)
+    weight_sum = torch.sum(weights, dim=dim)
+    return weighted_sum / (weight_sum + eps)
+
+
+def compute_predicted_aligned_error(
+    logits: Tensor,
+    aa_mask: Tensor,
+    sequence_id: Tensor | None = None,
+    max_bin: float = 31,
+) -> Tensor:
+    """Convert PAE logits ``X`` with shape (..., l, l, n) to distances."""
+
+    del sequence_id
+    pair_mask = _compute_pae_masks(aa_mask)
+    probabilities = _masked_probabilities(logits, pair_mask)
+    centers = _pae_bins(max_bin, logits.shape[-1], logits.device)
+    return torch.sum(probabilities * centers, dim=-1)
+
+
+@torch.no_grad()
+def compute_tm(logits: Tensor, aa_mask: Tensor, max_bin: float = 31.0) -> Tensor:
+    """Estimate TM score from pairwise PAE logits."""
+
+    pair_mask = _compute_pae_masks(aa_mask)
+    sequence_lengths = aa_mask.sum(dim=-1, keepdim=True)
+    centers = _pae_bins(max_bin, logits.shape[-1], logits.device)
+    distance_scale = 1.24 * (sequence_lengths.clamp_min(19) - 15) ** (1 / 3) - 1.8
+    tm_weights = 1.0 / (1 + (centers / distance_scale.unsqueeze(-1)) ** 2)
+    probabilities = _masked_probabilities(logits, pair_mask)
+    score_per_pair = torch.sum(probabilities * tm_weights.unsqueeze(-2), dim=-1)
+    score_per_anchor = masked_mean(pair_mask, score_per_pair, dim=-1)
+    return score_per_anchor.max(dim=-1).values
+
+
+def _local_coordinates(frames: Affine3D) -> Tensor:
+    origins = frames.trans[..., None, :, :]
+    return frames.invert()[..., None].apply(origins)
+
+
+def tm_loss(
+    logits: Tensor,
+    pred_affine: Tensor,
+    targ_affine: Tensor,
+    targ_mask: Tensor,
+    tm_mask: Tensor | None = None,
+    sequence_id: Tensor | None = None,
+    max_bin: float = 31,
+) -> Tensor:
+    """Cross-entropy loss for discretized aligned-position errors."""
+
+    del sequence_id
+    predicted_frames = Affine3D.from_tensor(pred_affine)
+    target_frames = Affine3D.from_tensor(targ_affine)
+    with torch.no_grad():
+        squared_error = (
+            (_local_coordinates(predicted_frames) - _local_coordinates(target_frames))
+            .square()
+            .sum(dim=-1)
+        )
+        boundaries = torch.linspace(
+            0,
+            max_bin,
+            logits.shape[-1] - 1,
+            device=logits.device,
+        ).square()
+        target_bins = (squared_error[..., None] > boundaries).sum(dim=-1).long()
+
+    cross_entropy = F.cross_entropy(
+        logits.movedim(3, 1),
+        target_bins,
+        reduction="none",
+    )
+    pair_mask = _compute_pae_masks(targ_mask)
+    loss_per_sample = masked_mean(pair_mask, cross_entropy, dim=(-1, -2))
+    if tm_mask is None:
+        return loss_per_sample.mean()
+    return masked_mean(tm_mask, loss_per_sample)

fastplms/models/esmfold2/esmfold2_prepare_input.py

@@ -0,0 +1,1130 @@
+"""Translate typed sequence inputs into the tensors consumed by ESMFold2.
+
+The conversion has four explicit stages: entity and chain assignment, residue
+tokenization, structural feature construction, and atom-table padding.  Keeping
+those stages separate makes the biological indexing rules testable without
+loading model weights.
+"""
+
+from __future__ import annotations
+
+import math
+import warnings
+from collections import defaultdict
+from contextlib import suppress
+from dataclasses import dataclass, field
+from itertools import combinations
+from typing import Any
+
+import numpy as np
+import torch
+
+from .esmfold2_conformers import (
+    get_ccd_leaving_atoms,
+    get_idealized_atom_pos,
+    get_ligand_ccd_atoms_with_charges,
+    get_ligand_ccd_bonds,
+    get_ligand_idealized_atom_pos,
+)
+from .esmfold2_constants import (
+    CHARGED_ATOMS,
+    DNA_1TO3,
+    DNA_BACKBONE_ATOMS,
+    DNA_HEAVY_ATOMS,
+    DNA_RESIDUE_TO_RES_TYPE,
+    DNA_RNA_LIGAND_INPUT_ID,
+    DNA_UNK_RES_TYPE,
+    ELEMENT_TO_ATOMIC_NUM,
+    ESM_PROTEIN_VOCAB,
+    MOL_TYPE_DNA,
+    MOL_TYPE_NONPOLYMER,
+    MOL_TYPE_PROTEIN,
+    MOL_TYPE_RNA,
+    MSA_GAP_TOKEN_ID,
+    PROTEIN_1TO3,
+    PROTEIN_3TO1,
+    PROTEIN_HEAVY_ATOMS,
+    PROTEIN_RESIDUE_TO_RES_TYPE,
+    PROTEIN_UNK_RES_TYPE,
+    RNA_1TO3,
+    RNA_BACKBONE_ATOMS,
+    RNA_HEAVY_ATOMS,
+    RNA_RESIDUE_TO_RES_TYPE,
+    RNA_UNK_RES_TYPE,
+)
+from .esmfold2_types import (
+    MSA,
+    DNAInput,
+    LigandInput,
+    Modification,
+    ProteinInput,
+    RNAInput,
+    StructurePredictionInput,
+)
+
+_ZERO_POS = np.zeros(3, dtype=np.float32)
+_ENCODE_ATOM_NAME_CACHE: dict[str, list[int]] = {}
+_ELEMENT_ATOMIC_NUM_CACHE: dict[str, int] = {}
+_TWO_LETTER_ELEMENTS = frozenset({"FE", "ZN", "MG", "MN", "CO", "NI", "CU", "SE", "BR"})
+
+
+@dataclass
+class AtomInfo:
+    """One row in the unpadded atom table."""
+
+    name: str
+    element: str
+    charge: int
+    ref_pos: np.ndarray  # R has shape (3,).
+    pos: np.ndarray  # X has shape (3,).
+    token_index: int = -1
+    atom_index: int = -1
+    space_uid: int = -1
+    is_valid: bool = True
+
+
+@dataclass
+class TokenInfo:
+    """Biological and atom-span annotations for one model token."""
+
+    token_index: int
+    residue_index: int
+    residue_name: str
+    mol_type: int
+    res_type: int
+    input_id: int
+    asym_id: int
+    sym_id: int
+    entity_id: int
+    atom_start: int
+    atom_count: int
+
+
+@dataclass
+class ChainInfo:
+    """One input chain after entity and symmetry assignment."""
+
+    chain_id: str
+    asym_id: int
+    entity_id: int
+    sym_id: int
+    mol_type: int
+    tokens: list[TokenInfo] = field(default_factory=list)
+    ligand_bonds: list[tuple[str, str]] = field(default_factory=list)
+
+
+@dataclass
+class _TokenizationState:
+    """Mutable cursor shared by residue tokenizers."""
+
+    token_index: int
+    atom_index: int
+    space_uid: int
+    tokens: list[TokenInfo] = field(default_factory=list)
+    atoms: list[AtomInfo] = field(default_factory=list)
+
+    def _append_atom(
+        self,
+        name: str,
+        element: str,
+        charge: int,
+        ref_pos: np.ndarray | None,
+    ) -> None:
+        self.atoms.append(
+            AtomInfo(
+                name=name,
+                element=element,
+                charge=charge,
+                ref_pos=(ref_pos.copy() if ref_pos is not None else _ZERO_POS.copy()),
+                pos=_ZERO_POS.copy(),
+                token_index=self.token_index,
+                atom_index=self.atom_index,
+                space_uid=self.space_uid,
+            )
+        )
+        self.atom_index += 1
+
+    def _append_token(
+        self,
+        *,
+        residue_index: int,
+        residue_name: str,
+        mol_type: int,
+        res_type: int,
+        input_id: int,
+        asym_id: int,
+        sym_id: int,
+        entity_id: int,
+        atom_start: int,
+        atom_count: int,
+    ) -> None:
+        self.tokens.append(
+            TokenInfo(
+                token_index=self.token_index,
+                residue_index=residue_index,
+                residue_name=residue_name,
+                mol_type=mol_type,
+                res_type=res_type,
+                input_id=input_id,
+                asym_id=asym_id,
+                sym_id=sym_id,
+                entity_id=entity_id,
+                atom_start=atom_start,
+                atom_count=atom_count,
+            )
+        )
+        self.token_index += 1
+
+    def add_residue_token(
+        self,
+        atom_specs: list[tuple[str, str, int, np.ndarray | None]],
+        **token_fields: Any,
+    ) -> None:
+        start = self.atom_index
+        for atom_spec in atom_specs:
+            self._append_atom(*atom_spec)
+        self._append_token(
+            atom_start=start,
+            atom_count=len(atom_specs),
+            **token_fields,
+        )
+        self.space_uid += 1
+
+    def add_atom_tokens(
+        self,
+        atom_specs: list[tuple[str, str, int, np.ndarray | None]],
+        **token_fields: Any,
+    ) -> None:
+        for atom_spec in atom_specs:
+            start = self.atom_index
+            self._append_atom(*atom_spec)
+            self._append_token(atom_start=start, atom_count=1, **token_fields)
+        self.space_uid += 1
+
+
+def encode_atom_name(name: str) -> list[int]:
+    """Encode a four-character atom name with the model's ASCII offset."""
+    cached = _ENCODE_ATOM_NAME_CACHE.get(name)
+    if cached is None:
+        cached = [0 if char == " " else ord(char) - 32 for char in name.ljust(4)[:4]]
+        _ENCODE_ATOM_NAME_CACHE[name] = cached
+    return cached
+
+
+def get_element_atomic_num(element: str) -> int:
+    """Map an element symbol to the model's atomic-number vocabulary."""
+    cached = _ELEMENT_ATOMIC_NUM_CACHE.get(element)
+    if cached is None:
+        cached = ELEMENT_TO_ATOMIC_NUM.get(element.upper(), 0)
+        _ELEMENT_ATOMIC_NUM_CACHE[element] = cached
+    return cached
+
+
+def _infer_element(atom_name: str) -> str:
+    normalized = atom_name.strip()
+    if not normalized:
+        return "C"
+    if normalized[0].isdigit():
+        return normalized[1] if len(normalized) > 1 else "H"
+    if len(normalized) == 2 and normalized in _TWO_LETTER_ELEMENTS:
+        return normalized
+    return normalized[0]
+
+
+def _compute_res_type(name: str, mol_type: int) -> int:
+    if mol_type == MOL_TYPE_PROTEIN:
+        return PROTEIN_RESIDUE_TO_RES_TYPE.get(name, PROTEIN_UNK_RES_TYPE)
+    if mol_type == MOL_TYPE_DNA:
+        return DNA_RESIDUE_TO_RES_TYPE.get(
+            name, RNA_RESIDUE_TO_RES_TYPE.get(name, DNA_UNK_RES_TYPE)
+        )
+    if mol_type == MOL_TYPE_RNA:
+        return RNA_RESIDUE_TO_RES_TYPE.get(
+            name, DNA_RESIDUE_TO_RES_TYPE.get(name, RNA_UNK_RES_TYPE)
+        )
+    return PROTEIN_UNK_RES_TYPE
+
+
+def _compute_esm_input_id(name: str, mol_type: int) -> int:
+    if mol_type != MOL_TYPE_PROTEIN:
+        return DNA_RNA_LIGAND_INPUT_ID
+    letter = PROTEIN_3TO1.get(name)
+    return (
+        DNA_RNA_LIGAND_INPUT_ID
+        if letter is None
+        else ESM_PROTEIN_VOCAB.get(letter, ESM_PROTEIN_VOCAB["X"])
+    )
+
+
+def _apply_modifications(residues: list[str], modifications: list[Modification] | None) -> set[int]:
+    changed: set[int] = set()
+    for modification in modifications or ():
+        residues[modification.position] = modification.ccd
+        changed.add(modification.position)
+    return changed
+
+
+def _ideal_atom_specs(
+    residue_name: str,
+    residue_type: int,
+    atom_names: list[str],
+    *,
+    charges: bool = True,
+) -> list[tuple[str, str, int, np.ndarray | None]]:
+    return [
+        (
+            atom_name,
+            _infer_element(atom_name),
+            CHARGED_ATOMS.get((residue_name, atom_name), 0) if charges else 0,
+            get_idealized_atom_pos(residue_type, atom_name),
+        )
+        for atom_name in atom_names
+    ]
+
+
+def _ccd_atom_specs(
+    residue_name: str,
+    atoms: list[tuple[str, str, int]],
+    excluded: set[str],
+    *,
+    force_zero: bool = False,
+) -> list[tuple[str, str, int, np.ndarray | None]]:
+    return [
+        (
+            atom_name,
+            element,
+            charge,
+            None if force_zero else get_ligand_idealized_atom_pos(residue_name, atom_name),
+        )
+        for atom_name, element, charge in atoms
+        if atom_name not in excluded
+    ]
+
+
+def tokenize_protein(
+    sequence: str,
+    modifications: list[Modification] | None,
+    entity_id: int,
+    asym_id: int,
+    sym_id: int,
+    token_offset: int,
+    atom_offset: int,
+    space_uid_offset: int,
+) -> tuple[list[TokenInfo], list[AtomInfo]]:
+    """Tokenize protein residues, atom-tokenizing modified CCD components."""
+    residues = [PROTEIN_1TO3.get(letter, "UNK") for letter in sequence]
+    modified = _apply_modifications(residues, modifications)
+    state = _TokenizationState(token_offset, atom_offset, space_uid_offset)
+
+    for residue_index, residue_name in enumerate(residues):
+        canonical_name = "MET" if residue_name == "MSE" else residue_name
+        common_fields = {
+            "residue_index": residue_index,
+            "mol_type": MOL_TYPE_PROTEIN,
+            "asym_id": asym_id,
+            "sym_id": sym_id,
+            "entity_id": entity_id,
+        }
+        if residue_index not in modified and canonical_name in PROTEIN_HEAVY_ATOMS:
+            residue_type = _compute_res_type(canonical_name, MOL_TYPE_PROTEIN)
+            state.add_residue_token(
+                _ideal_atom_specs(
+                    canonical_name,
+                    residue_type,
+                    PROTEIN_HEAVY_ATOMS[canonical_name],
+                ),
+                residue_name=canonical_name,
+                res_type=residue_type,
+                input_id=_compute_esm_input_id(canonical_name, MOL_TYPE_PROTEIN),
+                **common_fields,
+            )
+            continue
+
+        ccd_atoms = get_ligand_ccd_atoms_with_charges(residue_name)
+        if ccd_atoms is None:
+            ccd_atoms = [
+                (_infer_element(name), _infer_element(name), 0) for name in ("N", "CA", "C", "O")
+            ]
+        excluded = (
+            set() if residue_index == len(residues) - 1 else get_ccd_leaving_atoms(residue_name)
+        )
+        retained = [atom for atom in ccd_atoms if atom[0] not in excluded]
+        state.add_atom_tokens(
+            _ccd_atom_specs(
+                residue_name,
+                retained,
+                set(),
+                force_zero=len(retained) == 1,
+            ),
+            residue_name=residue_name,
+            res_type=PROTEIN_UNK_RES_TYPE,
+            input_id=DNA_RNA_LIGAND_INPUT_ID,
+            **common_fields,
+        )
+    return state.tokens, state.atoms
+
+
+def tokenize_nucleotide(
+    sequence: str,
+    modifications: list[Modification] | None,
+    mol_type: int,
+    entity_id: int,
+    asym_id: int,
+    sym_id: int,
+    token_offset: int,
+    atom_offset: int,
+    space_uid_offset: int,
+) -> tuple[list[TokenInfo], list[AtomInfo]]:
+    """Tokenize DNA or RNA, retaining backbone atoms for unknown bases."""
+    dna = mol_type == MOL_TYPE_DNA
+    letter_map = DNA_1TO3 if dna else RNA_1TO3
+    heavy_atoms = DNA_HEAVY_ATOMS if dna else RNA_HEAVY_ATOMS
+    backbone_atoms = DNA_BACKBONE_ATOMS if dna else RNA_BACKBONE_ATOMS
+    unknown_type = DNA_UNK_RES_TYPE if dna else RNA_UNK_RES_TYPE
+    residues = [letter_map.get(letter, "UNK") for letter in sequence]
+    modified = _apply_modifications(residues, modifications)
+    state = _TokenizationState(token_offset, atom_offset, space_uid_offset)
+
+    for residue_index, residue_name in enumerate(residues):
+        common_fields = {
+            "residue_index": residue_index,
+            "residue_name": residue_name,
+            "mol_type": mol_type,
+            "asym_id": asym_id,
+            "sym_id": sym_id,
+            "entity_id": entity_id,
+            "input_id": DNA_RNA_LIGAND_INPUT_ID,
+        }
+        if residue_index not in modified and residue_name in heavy_atoms:
+            residue_type = _compute_res_type(residue_name, mol_type)
+            state.add_residue_token(
+                _ideal_atom_specs(residue_name, residue_type, heavy_atoms[residue_name]),
+                res_type=residue_type,
+                **common_fields,
+            )
+            continue
+        if residue_index not in modified and residue_name == "UNK":
+            state.add_residue_token(
+                [(atom_name, _infer_element(atom_name), 0, None) for atom_name in backbone_atoms],
+                res_type=unknown_type,
+                **common_fields,
+            )
+            continue
+
+        ccd_atoms = get_ligand_ccd_atoms_with_charges(residue_name)
+        if ccd_atoms is None:
+            ccd_atoms = [(_infer_element(name), _infer_element(name), 0) for name in backbone_atoms]
+        excluded = (
+            set() if residue_index == len(residues) - 1 else get_ccd_leaving_atoms(residue_name)
+        )
+        state.add_atom_tokens(
+            _ccd_atom_specs(residue_name, ccd_atoms, excluded),
+            res_type=PROTEIN_UNK_RES_TYPE,
+            **common_fields,
+        )
+    return state.tokens, state.atoms
+
+
+def tokenize_ligand_ccd(
+    ccd_codes: list[str],
+    entity_id: int,
+    asym_id: int,
+    sym_id: int,
+    token_offset: int,
+    atom_offset: int,
+    space_uid_offset: int,
+    has_covalent_bond: bool,
+) -> tuple[list[TokenInfo], list[AtomInfo]]:
+    """Tokenize CCD ligands with one model token per retained atom."""
+    state = _TokenizationState(token_offset, atom_offset, space_uid_offset)
+    for residue_index, code in enumerate(ccd_codes):
+        ccd_atoms = get_ligand_ccd_atoms_with_charges(code)
+        if ccd_atoms is None:
+            raise ValueError(f"CCD component {code} not found")
+        excluded = get_ccd_leaving_atoms(code) if has_covalent_bond else set()
+        state.add_atom_tokens(
+            _ccd_atom_specs(code, ccd_atoms, excluded),
+            residue_index=residue_index,
+            residue_name=code,
+            mol_type=MOL_TYPE_NONPOLYMER,
+            res_type=PROTEIN_UNK_RES_TYPE,
+            input_id=DNA_RNA_LIGAND_INPUT_ID,
+            asym_id=asym_id,
+            sym_id=sym_id,
+            entity_id=entity_id,
+        )
+    return state.tokens, state.atoms
+
+
+def tokenize_ligand_smiles(
+    smiles: str,
+    entity_id: int,
+    asym_id: int,
+    sym_id: int,
+    token_offset: int,
+    atom_offset: int,
+    space_uid_offset: int,
+    seed: int | None = None,
+) -> tuple[list[TokenInfo], list[AtomInfo], list[tuple[str, str]]]:
+    """Generate a conformer and tokenize each heavy atom of a SMILES ligand."""
+    from rdkit import Chem
+    from rdkit.Chem import AllChem
+
+    molecule = Chem.MolFromSmiles(smiles)
+    if molecule is None:
+        raise ValueError(f"Failed to parse SMILES: {smiles}")
+    molecule = Chem.AddHs(molecule)
+    canonical_order = AllChem.CanonicalRankAtoms(molecule)  # type: ignore[attr-defined]
+    for atom, canonical_index in zip(molecule.GetAtoms(), canonical_order, strict=True):
+        name = atom.GetSymbol().upper() + str(canonical_index + 1)
+        if len(name) > 4:
+            raise ValueError(f"SMILES {smiles} has atom name longer than 4 chars: {name}")
+        atom.SetProp("name", name)
+
+    options = AllChem.ETKDGv3()  # type: ignore[attr-defined]
+    options.clearConfs = False
+    if seed is not None:
+        options.randomSeed = seed
+    conformer_id = AllChem.EmbedMolecule(molecule, options)  # type: ignore[attr-defined]
+    if conformer_id == -1:
+        options.useRandomCoords = True
+        conformer_id = AllChem.EmbedMolecule(molecule, options)  # type: ignore[attr-defined]
+    if conformer_id != -1:
+        with suppress(RuntimeError, ValueError):
+            AllChem.UFFOptimizeMolecule(  # type: ignore[attr-defined]
+                molecule, confId=conformer_id, maxIters=1000
+            )
+
+    molecule = Chem.RemoveHs(molecule)
+    if molecule.GetNumConformers() == 0:
+        raise ValueError(f"Failed to generate conformer for SMILES: {smiles}")
+    conformer = molecule.GetConformer(0)
+    atom_specs: list[tuple[str, str, int, np.ndarray | None]] = []
+    for atom in molecule.GetAtoms():
+        position = conformer.GetAtomPosition(atom.GetIdx())
+        atom_specs.append(
+            (
+                atom.GetProp("name"),
+                atom.GetSymbol(),
+                atom.GetFormalCharge(),
+                np.asarray([position.x, position.y, position.z], dtype=np.float32),
+            )
+        )
+    state = _TokenizationState(token_offset, atom_offset, space_uid_offset)
+    state.add_atom_tokens(
+        atom_specs,
+        residue_index=0,
+        residue_name="LIG",
+        mol_type=MOL_TYPE_NONPOLYMER,
+        res_type=PROTEIN_UNK_RES_TYPE,
+        input_id=DNA_RNA_LIGAND_INPUT_ID,
+        asym_id=asym_id,
+        sym_id=sym_id,
+        entity_id=entity_id,
+    )
+    bonds = [
+        (
+            bond.GetBeginAtom().GetProp("name"),
+            bond.GetEndAtom().GetProp("name"),
+        )
+        for bond in molecule.GetBonds()
+    ]
+    return state.tokens, state.atoms, bonds
+
+
+def _get_sequence_key(item: Any) -> str:
+    if isinstance(item, ProteinInput):
+        return f"PROTEIN:{item.sequence}"
+    if isinstance(item, DNAInput):
+        return f"DNA:{item.sequence}"
+    if isinstance(item, RNAInput):
+        return f"RNA:{item.sequence}"
+    if isinstance(item, LigandInput):
+        return f"LIGAND_CCD:{','.join(item.ccd)}" if item.ccd else f"LIGAND_SMILES:{item.smiles}"
+    raise ValueError(f"Unknown input type: {type(item)}")
+
+
+def _tokenize_chain(
+    item: Any,
+    chain_id: str,
+    *,
+    entity_id: int,
+    asym_id: int,
+    sym_id: int,
+    token_offset: int,
+    atom_offset: int,
+    space_uid_offset: int,
+    covalent_chains: set[str],
+    seed: int | None,
+) -> tuple[list[TokenInfo], list[AtomInfo], list[tuple[str, str]]]:
+    common = {
+        "entity_id": entity_id,
+        "asym_id": asym_id,
+        "sym_id": sym_id,
+        "token_offset": token_offset,
+        "atom_offset": atom_offset,
+        "space_uid_offset": space_uid_offset,
+    }
+    if isinstance(item, ProteinInput):
+        if item.msa is None:
+            warnings.warn(
+                f"No MSA provided for {item.id}, using single sequence mode",
+                stacklevel=2,
+            )
+        tokens, atoms = tokenize_protein(item.sequence, item.modifications, **common)
+        return tokens, atoms, []
+    if isinstance(item, (DNAInput, RNAInput)):
+        mol_type = MOL_TYPE_DNA if isinstance(item, DNAInput) else MOL_TYPE_RNA
+        tokens, atoms = tokenize_nucleotide(
+            item.sequence, item.modifications, mol_type=mol_type, **common
+        )
+        return tokens, atoms, []
+    if not isinstance(item, LigandInput):
+        raise ValueError(f"Unknown input type: {type(item)}")
+    if item.ccd is not None:
+        if item.smiles is not None:
+            warnings.warn("Both ccd and smiles provided, using ccd", stacklevel=2)
+        tokens, atoms = tokenize_ligand_ccd(
+            item.ccd,
+            has_covalent_bond=chain_id in covalent_chains,
+            **common,
+        )
+        return tokens, atoms, []
+    if item.smiles is not None:
+        return tokenize_ligand_smiles(item.smiles, seed=seed, **common)
+    raise ValueError("LigandInput must have either ccd or smiles")
+
+
+def build_chains_from_input(
+    input: StructurePredictionInput, seed: int | None = None
+) -> tuple[list[ChainInfo], list[TokenInfo], list[AtomInfo]]:
+    """Assign entities and symmetry copies, then tokenize every input chain."""
+    chains: list[ChainInfo] = []
+    tokens: list[TokenInfo] = []
+    atoms: list[AtomInfo] = []
+    entity_for_sequence: dict[str, int] = {}
+    next_symmetry: dict[int, int] = {}
+    covalent_chains = {
+        chain_id
+        for bond in input.covalent_bonds or ()
+        for chain_id in (bond.chain_id1, bond.chain_id2)
+    }
+    space_uid_offset = 0
+
+    for item in input.sequences:
+        key = _get_sequence_key(item)
+        entity_id = entity_for_sequence.setdefault(key, len(entity_for_sequence))
+        chain_ids = [item.id] if isinstance(item.id, str) else item.id
+        for chain_id in chain_ids:
+            sym_id = next_symmetry.get(entity_id, 0)
+            next_symmetry[entity_id] = sym_id + 1
+            asym_id = len(chains)
+            new_tokens, new_atoms, ligand_bonds = _tokenize_chain(
+                item,
+                chain_id,
+                entity_id=entity_id,
+                asym_id=asym_id,
+                sym_id=sym_id,
+                token_offset=len(tokens),
+                atom_offset=len(atoms),
+                space_uid_offset=space_uid_offset,
+                covalent_chains=covalent_chains,
+                seed=seed,
+            )
+            chains.append(
+                ChainInfo(
+                    chain_id=chain_id,
+                    asym_id=asym_id,
+                    entity_id=entity_id,
+                    sym_id=sym_id,
+                    mol_type=(new_tokens[0].mol_type if new_tokens else MOL_TYPE_PROTEIN),
+                    tokens=new_tokens,
+                    ligand_bonds=ligand_bonds,
+                )
+            )
+            tokens.extend(new_tokens)
+            atoms.extend(new_atoms)
+            space_uid_offset += len({atom.space_uid for atom in new_atoms})
+    return chains, tokens, atoms
+
+
+def _atom_indices_by_name(atoms: list[AtomInfo]) -> dict[int, dict[str, int]]:
+    result: dict[int, dict[str, int]] = defaultdict(dict)
+    for atom in atoms:
+        if atom.is_valid:
+            result[atom.token_index][atom.name] = atom.atom_index
+    return result
+
+
+def _ligand_frames(
+    tokens: list[TokenInfo],
+    atoms: list[AtomInfo],
+    atom_indices: dict[int, dict[str, int]],
+) -> dict[int, tuple[int, int, int]]:
+    atom_for_token: dict[int, int] = {}
+    tokens_by_residue: dict[tuple[int, int], list[int]] = defaultdict(list)
+    for token in tokens:
+        if token.mol_type != MOL_TYPE_NONPOLYMER:
+            continue
+        named_atoms = atom_indices.get(token.token_index)
+        if named_atoms:
+            atom_for_token[token.token_index] = next(iter(named_atoms.values()))
+        tokens_by_residue[(token.asym_id, token.residue_index)].append(token.token_index)
+
+    frames: dict[int, tuple[int, int, int]] = {}
+    for residue_tokens in tokens_by_residue.values():
+        residue_atoms = [
+            atom_for_token[token] for token in residue_tokens if token in atom_for_token
+        ]
+        if len(residue_atoms) < 3:
+            for token in residue_tokens:
+                if token in atom_for_token:
+                    atom_index = atom_for_token[token]
+                    frames[token] = (atom_index, atom_index, atom_index)
+            continue
+        R = np.asarray([atoms[index].ref_pos for index in residue_atoms])
+        distances = np.sqrt(((R[:, None] - R[None]) ** 2).sum(-1))
+        nearest = np.argsort(distances, axis=1)
+        local = np.column_stack((nearest[:, 1], nearest[:, 0], nearest[:, 2]))
+        local_index = {atom_index: index for index, atom_index in enumerate(residue_atoms)}
+        for token in residue_tokens:
+            atom_index = atom_for_token.get(token)
+            if atom_index is None:
+                continue
+            selected = local[local_index[atom_index]]
+            frames[token] = tuple(residue_atoms[int(index)] for index in selected)
+    return frames
+
+
+def _frame_for_token(
+    token: TokenInfo,
+    named_atoms: dict[str, int],
+    ligand_frames: dict[int, tuple[int, int, int]],
+) -> tuple[int, int, int]:
+    fallback = next(iter(named_atoms.values()), 0)
+    if token.mol_type == MOL_TYPE_PROTEIN:
+        return (
+            (fallback, fallback, fallback)
+            if token.res_type == PROTEIN_UNK_RES_TYPE
+            else (
+                named_atoms.get("N", 0),
+                named_atoms.get("CA", 0),
+                named_atoms.get("C", 0),
+            )
+        )
+    if token.mol_type in (MOL_TYPE_DNA, MOL_TYPE_RNA):
+        return (
+            (fallback, fallback, fallback)
+            if token.res_type == PROTEIN_UNK_RES_TYPE
+            else (
+                named_atoms.get("C1'", 0),
+                named_atoms.get("C3'", 0),
+                named_atoms.get("C4'", 0),
+            )
+        )
+    if token.mol_type == MOL_TYPE_NONPOLYMER:
+        return ligand_frames.get(token.token_index, (fallback, fallback, fallback))
+    return fallback, fallback, fallback
+
+
+def _resolved_frames(
+    frames: np.ndarray, tokens: list[TokenInfo], atoms: list[AtomInfo]
+) -> np.ndarray:
+    if not tokens:
+        return np.zeros(0, dtype=bool)
+    X = (
+        np.asarray([atom.pos for atom in atoms], dtype=np.float32)
+        if atoms
+        else np.zeros((0, 3), dtype=np.float32)
+    )
+    valid_atoms = (
+        np.asarray([atom.is_valid for atom in atoms], dtype=bool)
+        if atoms
+        else np.zeros(0, dtype=bool)
+    )
+    resolved_atoms = valid_atoms & np.any(X != 0, axis=1)
+    origin = X[frames[:, 1]]
+    left = X[frames[:, 0]] - origin
+    right = X[frames[:, 2]] - origin
+    left_norm = np.linalg.norm(left, axis=1)
+    right_norm = np.linalg.norm(right, axis=1)
+    valid_norms = (left_norm >= 1e-6) & (right_norm >= 1e-6)
+    cosine = np.zeros(len(tokens), dtype=np.float32)
+    if np.any(valid_norms):
+        cosine[valid_norms] = np.sum(left[valid_norms] * right[valid_norms], axis=1) / (
+            left_norm[valid_norms] * right_norm[valid_norms]
+        )
+    angle = np.degrees(np.arccos(np.abs(np.clip(cosine, -1, 1))))
+    all_resolved = resolved_atoms[frames].all(axis=1)
+    repeated = (frames[:, 0] == frames[:, 1]) & (frames[:, 1] == frames[:, 2])
+    return all_resolved & ~repeated & valid_norms & (angle >= 25)
+
+
+def compute_frame_indices(
+    tokens: list[TokenInfo], atoms: list[AtomInfo]
+) -> tuple[np.ndarray, np.ndarray]:
+    """Return frame atom indices F with shape (l, 3) and validity M with shape (l,)."""
+    named_atoms = _atom_indices_by_name(atoms)
+    ligand_frames = _ligand_frames(tokens, atoms, named_atoms)
+    frames = np.asarray(
+        [
+            _frame_for_token(token, named_atoms.get(token.token_index, {}), ligand_frames)
+            for token in tokens
+        ],
+        dtype=np.int64,
+    )
+    return frames, _resolved_frames(frames, tokens, atoms)
+
+
+def _atom_tokenized_residues(
+    tokens: list[TokenInfo], atoms: list[AtomInfo]
+) -> dict[tuple[int, int], list[tuple[str, int]]]:
+    grouped: dict[tuple[int, int], list[tuple[str, int]]] = defaultdict(list)
+    for atom in atoms:
+        if not atom.is_valid or atom.token_index >= len(tokens):
+            continue
+        token = tokens[atom.token_index]
+        if token.mol_type == MOL_TYPE_NONPOLYMER or token.res_type == PROTEIN_UNK_RES_TYPE:
+            grouped[(token.asym_id, token.residue_index)].append((atom.name, atom.token_index))
+    return grouped
+
+
+def _backbone_token(
+    residue_tokens: list[TokenInfo], atom_name: str, atoms: list[AtomInfo]
+) -> int | None:
+    if len(residue_tokens) == 1 and residue_tokens[0].res_type != PROTEIN_UNK_RES_TYPE:
+        return residue_tokens[0].token_index
+    for token in residue_tokens:
+        for atom_index in range(token.atom_start, token.atom_start + token.atom_count):
+            if atom_index < len(atoms) and atoms[atom_index].name == atom_name:
+                return token.token_index
+    return residue_tokens[0].token_index if residue_tokens else None
+
+
+def compute_token_bonds(
+    tokens: list[TokenInfo],
+    atoms: list[AtomInfo],
+    input: StructurePredictionInput,
+    chains: list[ChainInfo],
+) -> torch.Tensor:
+    """Build the symmetric token-bond matrix M with shape (l, l, 1)."""
+    edges: set[tuple[int, int]] = set()
+
+    def connect(left: int | None, right: int | None) -> None:
+        if left is not None and right is not None and left != right:
+            edges.add((min(left, right), max(left, right)))
+
+    explicit_bonds = {
+        (chain.asym_id, 0): chain.ligand_bonds for chain in chains if chain.ligand_bonds
+    }
+    for residue_key, atom_list in _atom_tokenized_residues(tokens, atoms).items():
+        if not atom_list:
+            continue
+        residue_name = tokens[atom_list[0][1]].residue_name
+        token_for_name = {name: token_index for name, token_index in atom_list}
+        bonds = explicit_bonds.get(residue_key)
+        if bonds is None:
+            bonds = get_ligand_ccd_bonds(residue_name)
+        if bonds:
+            for left_name, right_name in bonds:
+                if left_name in token_for_name and right_name in token_for_name:
+                    connect(token_for_name[left_name], token_for_name[right_name])
+        else:
+            for left, right in combinations([token_index for _, token_index in atom_list], 2):
+                connect(left, right)
+
+    if input.covalent_bonds:
+        chain_for_id = {chain.chain_id: chain for chain in chains}
+        residue_atoms: dict[tuple[int, int], list[AtomInfo]] = defaultdict(list)
+        for atom in atoms:
+            if atom.is_valid and atom.token_index < len(tokens):
+                token = tokens[atom.token_index]
+                residue_atoms[(token.asym_id, token.residue_index)].append(atom)
+        for bond in input.covalent_bonds:
+            left_chain = chain_for_id.get(bond.chain_id1)
+            right_chain = chain_for_id.get(bond.chain_id2)
+            if left_chain is None or right_chain is None:
+                continue
+            left_atoms = residue_atoms.get((left_chain.asym_id, bond.res_idx1), [])
+            right_atoms = residue_atoms.get((right_chain.asym_id, bond.res_idx2), [])
+            if bond.atom_idx1 < len(left_atoms) and bond.atom_idx2 < len(right_atoms):
+                connect(
+                    left_atoms[bond.atom_idx1].token_index,
+                    right_atoms[bond.atom_idx2].token_index,
+                )
+
+    protein_residues: dict[tuple[int, int], list[TokenInfo]] = defaultdict(list)
+    for token in tokens:
+        if token.mol_type == MOL_TYPE_PROTEIN:
+            protein_residues[(token.asym_id, token.residue_index)].append(token)
+    for (asym_id, residue_index), residue_tokens in protein_residues.items():
+        if not any(token.res_type == PROTEIN_UNK_RES_TYPE for token in residue_tokens):
+            continue
+        previous = protein_residues.get((asym_id, residue_index - 1))
+        following = protein_residues.get((asym_id, residue_index + 1))
+        if previous:
+            connect(
+                _backbone_token(previous, "C", atoms),
+                _backbone_token(residue_tokens, "N", atoms),
+            )
+        if following:
+            connect(
+                _backbone_token(residue_tokens, "C", atoms),
+                _backbone_token(following, "N", atoms),
+            )
+
+    matrix = torch.zeros(len(tokens), len(tokens), 1, dtype=torch.float32)
+    for left, right in edges:
+        matrix[left, right, 0] = 1.0
+        matrix[right, left, 0] = 1.0
+    return matrix
+
+
+def compute_representative_atoms(tokens: list[TokenInfo], atoms: list[AtomInfo]) -> torch.Tensor:
+    """Choose one distogram atom per token and return indices I with shape (l,)."""
+    named_atoms = _atom_indices_by_name(atoms)
+    representatives = torch.zeros(len(tokens), dtype=torch.int64)
+    for token in tokens:
+        names = named_atoms.get(token.token_index, {})
+        fallback = next(iter(names.values()), 0)
+        if token.mol_type == MOL_TYPE_PROTEIN:
+            representative = names.get("CB", names.get("CA", fallback))
+        elif token.mol_type in (MOL_TYPE_DNA, MOL_TYPE_RNA):
+            if token.res_type in (27, 32):
+                representative = names.get("C1'", fallback)
+            elif token.res_type in (23, 24, 28, 29):
+                representative = names.get("C4", names.get("C1'", fallback))
+            else:
+                representative = names.get("C2", names.get("C1'", fallback))
+        else:
+            representative = fallback
+        representatives[token.token_index] = representative
+    return representatives
+
+
+def _msa_assignments(
+    input: StructurePredictionInput, chains: list[ChainInfo]
+) -> dict[int, MSA | None]:
+    chain_msas: dict[int, MSA | None] = {}
+    chain_index = 0
+    for item in input.sequences:
+        chain_ids = [item.id] if isinstance(item.id, str) else list(item.id)
+        for _ in chain_ids:
+            chain = chains[chain_index]
+            if isinstance(item, ProteinInput):
+                chain_msas[chain.asym_id] = (
+                    MSA.from_sequences([item.sequence]) if item.msa is None else item.msa
+                )
+            else:
+                chain_msas[chain.asym_id] = None
+            chain_index += 1
+    return chain_msas
+
+
+def compute_msa_features(
+    input: StructurePredictionInput,
+    chains: list[ChainInfo],
+    tokens: list[TokenInfo],
+    max_seqs: int = 16384,
+) -> dict[str, torch.Tensor]:
+    """Pair per-chain MSAs and return row features with shape (m, l)."""
+    from .esmfold2_paired_msa import (
+        construct_paired_msa,
+        protein_letter_to_res_type,
+    )
+
+    chain_msas = _msa_assignments(input, chains)
+    query_types = {
+        chain.asym_id: np.asarray(
+            [token.res_type for token in tokens if token.asym_id == chain.asym_id],
+            dtype=np.int64,
+        )
+        for chain in chains
+    }
+    msa_residues, deletion_counts, _ = construct_paired_msa(
+        chain_msas,
+        query_types,
+        np.asarray([token.asym_id for token in tokens], dtype=np.int64),
+        np.asarray([token.residue_index for token in tokens], dtype=np.int64),
+        letter_to_res_type=protein_letter_to_res_type(),
+        max_seqs=max_seqs,
+    )
+    for token in tokens:
+        if chain_msas.get(token.asym_id) is None:
+            msa_residues[:, token.token_index] = MSA_GAP_TOKEN_ID
+            msa_residues[0, token.token_index] = token.res_type
+    if msa_residues.shape[0] == 0:
+        msa_residues = np.full((1, len(tokens)), MSA_GAP_TOKEN_ID, dtype=np.int64)
+        deletion_counts = np.zeros((1, len(tokens)), dtype=np.float32)
+
+    msa = torch.from_numpy(msa_residues)
+    deletion_count = torch.from_numpy(deletion_counts)
+    deletion_value = (np.pi / 2) * torch.arctan(deletion_count / 3)
+    return {
+        "msa": msa,
+        "deletion_value": deletion_value,
+        "has_deletion": deletion_count > 0,
+        "deletion_mean": deletion_value.mean(dim=0),
+        "msa_attention_mask": torch.ones_like(msa, dtype=torch.bool),
+    }
+
+
+def compute_distogram_conditioning(
+    input: StructurePredictionInput,
+    chains: list[ChainInfo],
+    tokens: list[TokenInfo],
+    disto_center: torch.Tensor,
+    min_dist: float = 2.0,
+    max_dist: float = 22.0,
+    num_bins: int = 64,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Bin user distances into D and return D plus its Boolean mask M."""
+    del disto_center
+    n_tokens = len(tokens)
+    bins = torch.zeros((n_tokens, n_tokens), dtype=torch.long)
+    mask = torch.zeros((n_tokens, n_tokens), dtype=torch.bool)
+    if not input.distogram_conditioning:
+        return bins, mask
+    asym_for_chain = {chain.chain_id: chain.asym_id for chain in chains}
+    tokens_for_asym: dict[int, list[int]] = defaultdict(list)
+    for token in tokens:
+        tokens_for_asym[token.asym_id].append(token.token_index)
+    boundaries = torch.linspace(min_dist, max_dist, num_bins + 1)
+
+    for conditioning in input.distogram_conditioning:
+        asym_id = asym_for_chain.get(conditioning.chain_id)
+        if asym_id is None:
+            continue
+        indices = tokens_for_asym[asym_id]
+        distances = torch.as_tensor(conditioning.distogram, dtype=torch.float32)
+        expected_shape = (len(indices), len(indices))
+        if distances.shape != expected_shape:
+            raise ValueError(
+                f"Distogram shape {distances.shape} doesn't match chain length {len(indices)}"
+            )
+        selected = torch.bucketize(distances, boundaries[:-1]).sub(1).clamp(0, num_bins - 1)
+        token_indices_tensor = torch.as_tensor(indices, dtype=torch.long)
+        bins[token_indices_tensor[:, None], token_indices_tensor[None, :]] = selected
+        mask[token_indices_tensor[:, None], token_indices_tensor[None, :]] = True
+    return bins, mask
+
+
+def _padded_atoms(atoms: list[AtomInfo]) -> list[AtomInfo]:
+    target = math.ceil(len(atoms) / 32) * 32 if atoms else 32
+    padding = [
+        AtomInfo(
+            name="",
+            element="",
+            charge=0,
+            ref_pos=_ZERO_POS.copy(),
+            pos=_ZERO_POS.copy(),
+            token_index=0,
+            atom_index=index,
+            space_uid=0,
+            is_valid=False,
+        )
+        for index in range(len(atoms), target)
+    ]
+    return [*atoms, *padding]
+
+
+def _token_tensors(tokens: list[TokenInfo]) -> dict[str, torch.Tensor]:
+    fields = {
+        "token_index": "token_index",
+        "residue_index": "residue_index",
+        "asym_id": "asym_id",
+        "sym_id": "sym_id",
+        "entity_id": "entity_id",
+        "mol_type": "mol_type",
+        "res_type": "res_type",
+        "input_ids": "input_id",
+    }
+    return {
+        output_name: torch.from_numpy(
+            np.asarray([getattr(token, attribute) for token in tokens], dtype=np.int64)
+        )
+        for output_name, attribute in fields.items()
+    }
+
+
+def _atom_tensors(atoms: list[AtomInfo]) -> dict[str, torch.Tensor]:
+    n_atoms = len(atoms)
+    ref_pos = np.zeros((n_atoms, 3), dtype=np.float32)
+    ref_element = np.zeros(n_atoms, dtype=np.int64)
+    ref_charge = np.zeros(n_atoms, dtype=np.int8)
+    ref_name = np.zeros((n_atoms, 4), dtype=np.int64)
+    ref_space = np.zeros(n_atoms, dtype=np.int64)
+    atom_mask = np.zeros(n_atoms, dtype=np.bool_)
+    atom_to_token = np.zeros(n_atoms, dtype=np.int64)
+    positions = np.zeros((n_atoms, 3), dtype=np.float64)
+    valid = np.zeros(n_atoms, dtype=np.bool_)
+    for index, atom in enumerate(atoms):
+        if atom.ref_pos is not None:
+            ref_pos[index] = atom.ref_pos
+        ref_charge[index] = atom.charge
+        ref_space[index] = atom.space_uid if atom.space_uid >= 0 else atom.token_index
+        atom_mask[index] = atom.is_valid
+        valid[index] = atom.is_valid
+        positions[index] = atom.pos
+        if atom.is_valid:
+            ref_element[index] = get_element_atomic_num(atom.element)
+            ref_name[index] = encode_atom_name(atom.name)
+            atom_to_token[index] = atom.token_index
+
+    resolved = valid & np.any(positions != 0, axis=1)
+    X = torch.from_numpy(positions)
+    resolved_mask = torch.from_numpy(resolved)
+    valid_mask = torch.from_numpy(valid)
+    if resolved_mask.any():
+        X = X - X[resolved_mask].mean(dim=0, keepdim=True)
+        X[~valid_mask] = 0.0
+    return {
+        "ref_pos": torch.from_numpy(ref_pos),
+        "ref_element": torch.from_numpy(ref_element),
+        "ref_charge": torch.from_numpy(ref_charge),
+        "ref_atom_name_chars": torch.from_numpy(ref_name),
+        "ref_space_uid": torch.from_numpy(ref_space),
+        "gt_coords": X.float().unsqueeze(0),
+        "atom_attention_mask": torch.from_numpy(atom_mask),
+        "atom_to_token": torch.from_numpy(atom_to_token),
+        "is_resolved": torch.tensor(resolved, dtype=torch.bool),
+    }
+
+
+def build_feature_tensors(
+    chains: list[ChainInfo],
+    tokens: list[TokenInfo],
+    atoms: list[AtomInfo],
+    input: StructurePredictionInput,
+) -> dict[str, torch.Tensor]:
+    """Assemble the complete unbatched ESMFold2 feature dictionary."""
+    token_features = _token_tensors(tokens)
+    atom_features = _atom_tensors(_padded_atoms(atoms))
+    frames, _ = compute_frame_indices(tokens, atoms)
+    msa_features = compute_msa_features(input, chains, tokens)
+    distogram, distogram_mask = compute_distogram_conditioning(
+        input,
+        chains,
+        tokens,
+        torch.zeros(len(tokens), 3, dtype=torch.float32),
+    )
+    return {
+        **token_features,
+        "token_bonds": compute_token_bonds(tokens, atoms, input, chains),
+        "token_attention_mask": torch.ones(len(tokens), dtype=torch.bool),
+        "pocket_feature": torch.zeros(len(tokens), dtype=torch.long),
+        **atom_features,
+        "distogram_atom_idx": compute_representative_atoms(tokens, atoms),
+        "frames_idx": torch.from_numpy(frames).to(torch.int64),
+        "disto_cond": distogram,
+        "disto_cond_mask": distogram_mask,
+        **msa_features,
+    }
+
+
+def prepare_esmfold2_input(
+    input: StructurePredictionInput, seed: int | None = None
+) -> tuple[dict[str, torch.Tensor], list[ChainInfo]]:
+    """Convert one typed request to model features and output-chain metadata."""
+    chains, tokens, atoms = build_chains_from_input(input, seed)
+    return build_feature_tensors(chains, tokens, atoms, input), chains

fastplms/models/esmfold2/esmfold2_processor.py

@@ -0,0 +1,332 @@
+"""Input preparation and output decoding for ESMFold2 inference."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import torch
+from torch import Tensor
+
+from .esmfold2_conformers import load_ccd
+from .esmfold2_molecular_complex import MolecularComplexResult
+from .esmfold2_output import build_molecular_complex_from_features
+from .esmfold2_prepare_input import ChainInfo, prepare_esmfold2_input
+from .esmfold2_types import MSA, Modification, ProteinInput, StructurePredictionInput
+from .modeling_esmfold2_common import MSA_CONDITIONING_INPUT_NAMES
+from .reproducibility import seed_context
+
+# Backward-compatible private alias for the pinned parity helpers. New callers
+# should import ``seed_context`` from the public ``fastplms.models.esmfold2``
+# package instead of reaching into implementation modules.
+_seed_context = seed_context
+
+
+@dataclass(frozen=True)
+class _SplitProteinState:
+    ids: dict[str, list[str]]
+    modifications: dict[str, list[Modification]]
+    msas: dict[str, MSA | None]
+
+
+@dataclass(frozen=True)
+class _PendingProtein:
+    source: ProteinInput
+    sequence: str
+    state: _SplitProteinState
+
+
+def _chain_starts(chains: list[str]) -> list[int]:
+    starts: list[int] = []
+    position = 0
+    for chain in chains:
+        starts.append(position)
+        position += len(chain) + 1
+    return starts
+
+
+def _split_modifications(
+    item: ProteinInput,
+    chains: list[str],
+    starts: list[int],
+) -> dict[str, list[Modification]]:
+    grouped: dict[str, list[Modification]] = {}
+    if item.modifications is None:
+        return grouped
+    for chain, start in zip(chains, starts, strict=True):
+        end = start + len(chain)
+        adjusted = [
+            Modification(position=modification.position - start, ccd=modification.ccd)
+            for modification in item.modifications
+            if start <= modification.position < end
+        ]
+        grouped.setdefault(chain, []).extend(adjusted)
+    return grouped
+
+
+def _split_msas(
+    item: ProteinInput,
+    chains: list[str],
+    starts: list[int],
+) -> dict[str, MSA | None]:
+    grouped: dict[str, MSA | None] = {}
+    if item.msa is None:
+        return grouped
+    for chain, start in zip(chains, starts, strict=True):
+        if chain not in grouped:
+            grouped[chain] = item.msa.select_positions(np.arange(start, start + len(chain)))
+    return grouped
+
+
+def _split_protein(item: ProteinInput) -> tuple[list[_PendingProtein], _SplitProteinState]:
+    chains = ":".join(item.sequence.split("|")).split(":")
+    starts = _chain_starts(chains)
+    base_id = item.id[0] if isinstance(item.id, list) else item.id
+    ids: dict[str, list[str]] = {}
+    for index, chain in enumerate(chains):
+        chain_ids = ids.setdefault(chain, [])
+        chain_ids.append(f"{base_id}_{index}")
+    state = _SplitProteinState(
+        ids=ids,
+        modifications=_split_modifications(item, chains, starts),
+        msas=_split_msas(item, chains, starts),
+    )
+    pending = [
+        _PendingProtein(item, chain, state)
+        for chain, chain_ids in ids.items()
+        if chain_ids
+    ]
+    return pending, state
+
+
+def _resolve_pending(pending: _PendingProtein) -> ProteinInput:
+    item = pending.source
+    sequence = pending.sequence
+    state = pending.state
+    return ProteinInput(
+        id=state.ids[sequence],
+        sequence=sequence,
+        msa=state.msas.get(sequence) if item.msa else None,
+        modifications=(state.modifications.get(sequence) if item.modifications else None),
+    )
+
+
+def clean_esmfold2_input(input: StructurePredictionInput) -> StructurePredictionInput:
+    """Expand chain delimiters and group repeated protein sequences by entity."""
+
+    if input.pocket is not None:
+        raise NotImplementedError(
+            "ESMFold2 pocket conditioning is present in the upstream input schema but "
+            "the published ESMFold2 feature pipeline drops it. FastPLMs refuses this "
+            "input instead of silently emitting an all-zero pocket feature."
+        )
+
+    cleaned: list[Any] = []
+    for item in input.sequences:
+        if not isinstance(item, ProteinInput):
+            cleaned.append(item)
+            continue
+        sequence = ":".join(item.sequence.split("|"))
+        if ":" not in sequence:
+            cleaned.append(item)
+            continue
+        if input.covalent_bonds is not None:
+            raise ValueError(
+                "Covalent bonds are not supported when using chainbreaks. "
+                "Chains must be separated into multiple ProteinInput objects."
+            )
+        pending, _state = _split_protein(item)
+        cleaned.extend(pending)
+
+    resolved = [
+        _resolve_pending(item) if isinstance(item, _PendingProtein) else item
+        for item in cleaned
+    ]
+    return StructurePredictionInput(
+        sequences=resolved,
+        pocket=input.pocket,
+        distogram_conditioning=input.distogram_conditioning,
+        covalent_bonds=input.covalent_bonds,
+    )
+
+
+def _batch_features(
+    features: dict[str, Any],
+    device: torch.device | str | None,
+) -> dict[str, Any]:
+    return {
+        name: (value[None].to(device) if device is not None else value[None])
+        if isinstance(value, Tensor)
+        else value
+        for name, value in features.items()
+    }
+
+
+def _sampler_overrides(
+    noise_scale: float | None,
+    step_scale: float | None,
+    max_inference_sigma: int | None,
+) -> dict[str, Any]:
+    values = {
+        "noise_scale": noise_scale,
+        "step_scale": step_scale,
+        "max_inference_sigma": max_inference_sigma,
+    }
+    return {name: value for name, value in values.items() if value is not None}
+
+
+class ESMFold2InputBuilder:
+    """Prepare public input objects, run folding, and decode model tensors."""
+
+    def __init__(self, ccd_cache: Path | None = None) -> None:
+        load_ccd(ccd_cache)
+
+    def prepare_input(
+        self,
+        input: StructurePredictionInput,
+        seed: int | None = None,
+        device: torch.device | str | None = None,
+    ) -> tuple[dict[str, Any], list[ChainInfo]]:
+        cleaned = clean_esmfold2_input(input)
+        with seed_context(seed):
+            features, chain_infos = prepare_esmfold2_input(cleaned, seed=seed)
+            return _batch_features(features, device), chain_infos
+
+    def prepare_model_input(
+        self,
+        model: Any,
+        input: StructurePredictionInput,
+        seed: int | None = None,
+        device: torch.device | str | None = None,
+    ) -> tuple[dict[str, Any], list[ChainInfo]]:
+        """Prepare features while enforcing the checkpoint's MSA contract."""
+
+        msa_conditioning = getattr(model.config, "msa_conditioning", None)
+        if not isinstance(msa_conditioning, bool):
+            raise RuntimeError("The ESMFold2 config has no Boolean msa_conditioning contract.")
+        if not msa_conditioning:
+            explicit_msa_ids = [
+                item.id
+                for item in input.sequences
+                if isinstance(item, ProteinInput) and item.msa is not None
+            ]
+            if explicit_msa_ids:
+                raise ValueError(
+                    "This ESMFold2 checkpoint was trained without MSA conditioning and "
+                    f"rejects explicit MSAs for protein inputs {explicit_msa_ids!r}."
+                )
+        features, chain_infos = self.prepare_input(input, seed=seed, device=device)
+        if not msa_conditioning:
+            for name in MSA_CONDITIONING_INPUT_NAMES:
+                features.pop(name, None)
+        return features, chain_infos
+
+    def __call__(
+        self,
+        input: StructurePredictionInput,
+        seed: int | None = None,
+        device: torch.device | str | None = None,
+    ) -> tuple[dict[str, Any], list[ChainInfo]]:
+        return self.prepare_input(input, seed=seed, device=device)
+
+    def _decode_sample(
+        self,
+        output: Mapping[str, Tensor],
+        features: dict[str, Tensor],
+        chain_infos: list[ChainInfo],
+        sample: int,
+        complex_id: str,
+    ) -> MolecularComplexResult:
+        plddt = output["plddt"][sample]
+        molecular_complex = build_molecular_complex_from_features(
+            coords=output["sample_atom_coords"][sample],
+            plddt=plddt,
+            atom_mask=features["atom_attention_mask"][0],
+            ref_element=features["ref_element"][0],
+            ref_atom_name_chars=features["ref_atom_name_chars"][0],
+            chain_infos=chain_infos,
+            complex_id=complex_id,
+        )
+
+        def sample_tensor(name: str) -> Tensor | None:
+            value = output.get(name)
+            return None if value is None else value[sample].detach().cpu()
+
+        def shared_tensor(name: str) -> Tensor | None:
+            value = output.get(name)
+            return None if value is None else value[0].detach().cpu()
+
+        ptm = output.get("ptm")
+        iptm = output.get("iptm")
+        return MolecularComplexResult(
+            complex=molecular_complex,
+            plddt=plddt.detach().cpu(),
+            ptm=float(ptm[sample].item()) if ptm is not None else None,
+            iptm=float(iptm[sample].item()) if iptm is not None else None,
+            pae=sample_tensor("pae"),
+            distogram=shared_tensor("distogram_logits"),
+            pair_chains_iptm=sample_tensor("pair_chains_iptm"),
+            residue_index=shared_tensor("residue_index"),
+            entity_id=shared_tensor("entity_id"),
+        )
+
+    def decode(
+        self,
+        output: Mapping[str, Tensor],
+        features: dict[str, Tensor],
+        chain_infos: list[ChainInfo],
+        *,
+        num_diffusion_samples: int = 1,
+        complex_id: str = "pred",
+    ) -> MolecularComplexResult | list[MolecularComplexResult]:
+        results = [
+            self._decode_sample(output, features, chain_infos, sample, complex_id)
+            for sample in range(output["sample_atom_coords"].shape[0])
+        ]
+        return results[0] if num_diffusion_samples == 1 and len(results) == 1 else results
+
+    def fold(
+        self,
+        model: Any,
+        input: StructurePredictionInput,
+        *,
+        num_loops: int = 3,
+        num_sampling_steps: int = 200,
+        num_diffusion_samples: int = 1,
+        seed: int | None = None,
+        noise_scale: float | None = None,
+        step_scale: float | None = None,
+        max_inference_sigma: int | None = None,
+        early_exit: bool = False,
+        complex_id: str = "pred",
+    ) -> MolecularComplexResult | list[MolecularComplexResult]:
+        features, chain_infos = self.prepare_model_input(
+            model,
+            input,
+            seed=seed,
+            device=model.device,
+        )
+        overrides = _sampler_overrides(noise_scale, step_scale, max_inference_sigma)
+        with torch.no_grad(), seed_context(seed):
+            output = model(
+                **features,
+                num_loops=num_loops,
+                num_sampling_steps=num_sampling_steps,
+                num_diffusion_samples=num_diffusion_samples,
+                early_exit=early_exit,
+                return_dict=True,
+                **overrides,
+            )
+        return self.decode(
+            output,
+            features,
+            chain_infos,
+            num_diffusion_samples=num_diffusion_samples,
+            complex_id=complex_id,
+        )
+
+
+__all__ = ["ESMFold2InputBuilder", "clean_esmfold2_input", "seed_context"]

fastplms/models/esmfold2/esmfold2_protein_chain.py

@@ -0,0 +1,1450 @@
+"""Protein-chain data, geometry, and serialization for ESMFold2."""
+
+from __future__ import annotations
+
+import io
+import warnings
+from collections.abc import Mapping, Sequence
+from dataclasses import asdict, dataclass, replace
+from functools import cached_property
+from pathlib import Path
+from typing import Any
+
+import biotite.structure as bs
+import brotli
+import msgpack
+import msgpack_numpy
+import numpy as np
+import torch
+from biotite.database import rcsb
+from biotite.structure.io.pdb import PDBFile
+from biotite.structure.io.pdbx import CIFCategory, CIFColumn, CIFData, CIFFile
+from biotite.structure.io.pdbx import set_structure as set_structure_pdbx
+from scipy.spatial import ConvexHull, KDTree
+from scipy.spatial.distance import cdist, pdist, squareform
+
+from . import esmfold2_residue_constants as residue_constants
+from .esmfold2_affine3d import Affine3D
+from .esmfold2_aligner import Aligner
+from .esmfold2_atom_indexer import AtomIndexer
+from .esmfold2_metrics import compute_gdt_ts, compute_lddt_ca
+from .esmfold2_misc import slice_python_object_as_numpy
+from .esmfold2_mmcif_parsing import (
+    PLDDT_B_FACTOR_SCALE,
+    MmcifWrapper,
+    Residue,
+    round_mmcif_columns,
+)
+from .esmfold2_normalize_coordinates import (
+    apply_frame_to_coords,
+    get_protein_normalization_frame,
+)
+from .esmfold2_protein_structure import index_by_atom_name
+from .esmfold2_utils_types import PathOrBuffer
+
+CHAIN_ID_CONST = "A"
+
+
+def _str_key_to_int_key(values: dict, ignore_keys: list[str] | None = None) -> dict:
+    """Restore integer dictionary keys after JSON-compatible serialization."""
+    ignored = frozenset(ignore_keys or ())
+    restored = {}
+    for key, value in values.items():
+        if isinstance(value, dict) and key not in ignored:
+            value = _str_key_to_int_key(value, ignore_keys=ignore_keys)
+        restored_key = int(key) if isinstance(key, str) and key.isdigit() else key
+        restored[restored_key] = value
+    return restored
+
+
+def _num_non_null_residues(seqres_to_structure_chain: Mapping[int, Residue]) -> int:
+    return sum(residue.residue_number is not None for residue in seqres_to_structure_chain.values())
+
+
+def infer_cb(
+    C,
+    N,
+    Ca,
+    bond_length: float = 1.522,
+    bond_angle: float = 1.927,
+    dihedral: float = -2.143,
+):
+    """Infer C-beta coordinates from C, N, and C-alpha coordinates."""
+
+    def normalize(X: np.ndarray) -> np.ndarray:
+        return X / np.sqrt(np.square(X).sum(-1, keepdims=True) + 1e-8)
+
+    with np.errstate(invalid="ignore"):
+        n_to_ca = N - Ca
+        n_to_c = N - C
+    axis = normalize(n_to_ca)
+    normal = normalize(np.cross(n_to_c, axis))
+    basis = (axis, np.cross(normal, axis), normal)
+    offsets = (
+        bond_length * np.cos(bond_angle),
+        bond_length * np.sin(bond_angle) * np.cos(dihedral),
+        -bond_length * np.sin(bond_angle) * np.sin(dihedral),
+    )
+    return Ca + sum(vector * offset for vector, offset in zip(basis, offsets, strict=False))
+
+
+def chain_to_ndarray(
+    atom_array: bs.AtomArray, mmcif: MmcifWrapper, chain_id: str, is_predicted=False
+):
+    if not isinstance(atom_array, bs.AtomArray):
+        raise TypeError("atom_array must be a biotite AtomArray.")
+    if not isinstance(mmcif, MmcifWrapper):
+        raise TypeError("mmcif must be an MmcifWrapper.")
+    if not isinstance(chain_id, str) or not chain_id:
+        raise ValueError("chain_id must be a non-empty string.")
+    if chain_id not in mmcif.chain_to_seqres or chain_id not in mmcif.seqres_to_structure:
+        raise ValueError(f"mmCIF data does not contain sequence mappings for chain {chain_id!r}.")
+    entity_id = None
+    for entity, chains in mmcif.entities.items():
+        if chain_id in chains:
+            entity_id = entity
+    num_res = len(mmcif.chain_to_seqres[chain_id])
+    sequence = mmcif.chain_to_seqres[chain_id]
+
+    atom_positions = np.full([num_res, residue_constants.atom_type_num, 3], np.nan)
+    atom_mask = np.full([num_res, residue_constants.atom_type_num], False, dtype=bool)
+    residue_index = np.full([num_res], -1, dtype=np.int64)
+    insertion_code = np.full([num_res], "", dtype="<U4")
+
+    confidence = np.ones([num_res], dtype=np.float32)
+
+    chain = atom_array[atom_array.chain_id == chain_id]
+    if not isinstance(chain, bs.AtomArray):
+        raise RuntimeError("Biotite selection did not return an AtomArray.")
+    for res_index in range(num_res):
+        res_at_position = mmcif.seqres_to_structure[chain_id][res_index]
+
+        if res_at_position.residue_number is None:
+            continue
+
+        residue_index[res_index] = res_at_position.residue_number
+        insertion_code[res_index] = res_at_position.insertion_code
+        res = chain[
+            (chain.res_id == res_at_position.residue_number)
+            & (chain.ins_code == res_at_position.insertion_code)
+            & (chain.hetero == res_at_position.hetflag)
+        ]
+        if not isinstance(res, bs.AtomArray):
+            raise RuntimeError("Biotite residue selection did not return an AtomArray.")
+
+        # Atom level features
+        for atom in res:
+            atom_name = atom.atom_name
+            if atom_name == "SE" and atom.res_name == "MSE":
+                # Put the coords of the selenium atom in the sulphur column
+                atom_name = "SD"
+
+            if atom_name in residue_constants.atom_order:
+                atom_positions[res_index, residue_constants.atom_order[atom_name]] = atom.coord
+                atom_mask[res_index, residue_constants.atom_order[atom_name]] = True
+                if is_predicted and atom_name == "CA":
+                    confidence[res_index] = atom.b_factor / PLDDT_B_FACTOR_SCALE
+
+    if not sequence or not all(sequence):
+        raise ValueError("Some residue name was not specified correctly.")
+    return (
+        sequence,
+        atom_positions,
+        atom_mask,
+        residue_index,
+        insertion_code,
+        confidence,
+        entity_id,
+    )
+
+
+@dataclass(frozen=True)
+class ProteinChain:
+    """Dataclass with atom37 representation of a single protein chain."""
+
+    id: str
+    sequence: str
+    chain_id: str  # author chain id - mutable
+    entity_id: int | None
+    residue_index: np.ndarray
+    insertion_code: np.ndarray
+    atom37_positions: np.ndarray
+    atom37_mask: np.ndarray
+    confidence: np.ndarray
+    mmcif: MmcifWrapper | None = None
+    atom37_confidence: np.ndarray | None = None  # P has shape (l, 37).
+
+    # Construction and parsing
+    @classmethod
+    def chain_iterable_from_mmcif(
+        cls,
+        path: PathOrBuffer | MmcifWrapper,
+        id: str | None = None,
+        is_predicted: bool = False,
+        keep_source: bool = False,
+    ):
+        """Yield every protein chain represented in an mmCIF structure."""
+        mmcif = path if isinstance(path, MmcifWrapper) else MmcifWrapper.read(path, id)
+        for chain in bs.chain_iter(mmcif.structure):
+            chain = chain[bs.filter_amino_acids(chain) & ~chain.hetero]
+            if len(chain) == 0:
+                continue
+            chain_id = chain.chain_id[0]
+            entity_id = None
+            for entity, chains in mmcif.entities.items():
+                if chain_id in chains:
+                    entity_id = entity
+            if entity_id is None:
+                raise ValueError(
+                    f"Failed to resolve entity identity for mmCIF chain {chain_id!r}."
+                )
+            (
+                sequence,
+                atom_positions,
+                atom_mask,
+                residue_index,
+                insertion_code,
+                confidence,
+                _,
+            ) = chain_to_ndarray(chain, mmcif, chain_id, is_predicted)
+            if not all(sequence):
+                raise ValueError("Some residue name was not specified correctly.")
+
+            yield cls(
+                id=mmcif.id,
+                sequence=sequence,
+                chain_id=chain_id,
+                entity_id=entity_id,
+                atom37_positions=atom_positions,
+                atom37_mask=atom_mask,
+                residue_index=residue_index,
+                insertion_code=insertion_code,
+                confidence=confidence,
+                mmcif=mmcif if keep_source else None,
+            )
+
+    @classmethod
+    def from_mmcif(
+        cls,
+        path: PathOrBuffer | MmcifWrapper,
+        chain_id: str | None = None,
+        entity_id: int | None = None,
+        id: str | None = None,
+        is_predicted: bool = False,
+        keep_source: bool = False,
+    ):
+        """Return a ProteinChain object from an mmcif file.
+
+        Args:
+            path: Uncompressed mmCIF path, buffer, or parsed wrapper.
+            id: Optional structure identifier when parsing a path or buffer.
+            is_predicted (bool): If True, reads b factor as the confidence readout. Default: False.
+            chain_id (str, optional): Select a chain corresponding to (author) chain id.
+            entity_id (int, optional): Select a chain corresponding to a particular entity.
+
+        If neither `chain_id` nor `entity_id` is specified, defaults to the first entity.
+        """
+        mmcif = path if isinstance(path, MmcifWrapper) else MmcifWrapper.read(path, id)
+
+        if chain_id is not None and entity_id is not None:
+            raise ValueError("Pass at most one of chain_id or entity_id.")
+
+        # If neither chain_id nor entity_id is specified, default to the first entity
+        if chain_id is None and entity_id is None:
+            if not mmcif.entities:
+                raise ValueError("Structure contains no entities")
+            entity_id = min(mmcif.entities.keys())  # Pick the first entity by ID
+
+        if entity_id is not None:
+            if entity_id not in mmcif.entities:
+                raise ValueError(
+                    f"Structure does not contain entity `{entity_id}`. "
+                    f"Valid entities: {mmcif.entities.keys()}"
+                )
+            chains = mmcif.entities[entity_id]
+
+            # Prefer the chain with the most resolved residues; ties preserve source order.
+            chain_id = max(
+                chains,
+                key=lambda chain: _num_non_null_residues(mmcif.seqres_to_structure[chain]),
+            )
+        else:
+            if chain_id is None:
+                raise RuntimeError("Failed to resolve an mmCIF chain selection.")
+            for entity, chains in mmcif.entities.items():
+                if chain_id in chains:
+                    entity_id = entity
+        if entity_id is None:
+            warnings.warn(
+                "Failed to detect entity_id from mmcif file, it may be malformed.",
+                stacklevel=2,
+            )
+
+        atom_array = mmcif.structure
+        (
+            sequence,
+            atom_positions,
+            atom_mask,
+            residue_index,
+            insertion_code,
+            confidence,
+            _,
+        ) = chain_to_ndarray(atom_array, mmcif, chain_id, is_predicted)
+        if not all(sequence):
+            raise ValueError("Some residue name was not specified correctly.")
+
+        return cls(
+            id=mmcif.id,
+            sequence=sequence,
+            chain_id=chain_id,
+            entity_id=entity_id,
+            atom37_positions=atom_positions,
+            atom37_mask=atom_mask.astype(bool),
+            residue_index=residue_index,
+            insertion_code=insertion_code,
+            confidence=confidence,
+            mmcif=mmcif if keep_source else None,
+        )
+
+    @classmethod
+    def from_atom37(
+        cls,
+        atom37_positions: np.ndarray | torch.Tensor,
+        *,
+        id: str | None = None,
+        sequence: str | None = None,
+        chain_id: str | None = None,
+        entity_id: int | None = None,
+        residue_index: np.ndarray | torch.Tensor | None = None,
+        insertion_code: np.ndarray | None = None,
+        confidence: np.ndarray | torch.Tensor | None = None,
+    ):
+        if isinstance(atom37_positions, torch.Tensor):
+            atom37_positions = atom37_positions.cpu().numpy()
+            if atom37_positions.ndim == 4:
+                if atom37_positions.shape[0] != 1:
+                    raise ValueError(
+                        "Cannot handle batched inputs, atom37_positions has shape "
+                        f"{atom37_positions.shape}"
+                    )
+                atom37_positions = atom37_positions[0]
+
+        if not isinstance(atom37_positions, np.ndarray):
+            raise TypeError("atom37_positions must be a NumPy array or Torch tensor.")
+        if atom37_positions.ndim != 3 or atom37_positions.shape[1:] != (37, 3):
+            raise ValueError(
+                "atom37_positions must have shape (length, 37, 3), got "
+                f"{atom37_positions.shape}."
+            )
+        seqlen = atom37_positions.shape[0]
+
+        atom_mask = np.isfinite(atom37_positions).all(-1)
+
+        if id is None:
+            id = ""
+
+        if sequence is None:
+            sequence = "A" * seqlen
+
+        if chain_id is None:
+            chain_id = "A"
+
+        if residue_index is None:
+            residue_index = np.arange(1, seqlen + 1)
+        elif isinstance(residue_index, torch.Tensor):
+            residue_index = residue_index.cpu().numpy()
+            if residue_index.ndim == 2:
+                if residue_index.shape[0] != 1:
+                    raise ValueError(
+                        "Cannot handle batched inputs, residue_index has shape "
+                        f"{residue_index.shape}"
+                    )
+                residue_index = residue_index[0]
+        if not isinstance(residue_index, np.ndarray):
+            raise TypeError("residue_index must be a NumPy array or Torch tensor.")
+
+        if insertion_code is None:
+            insertion_code = np.array(["" for _ in range(seqlen)])
+
+        if confidence is None:
+            confidence = np.ones(seqlen, dtype=np.float32)
+        elif isinstance(confidence, torch.Tensor):
+            confidence = confidence.cpu().numpy()
+            if confidence.ndim == 2:
+                if confidence.shape[0] != 1:
+                    raise ValueError(
+                        f"Cannot handle batched inputs, confidence has shape {confidence.shape}"
+                    )
+                confidence = confidence[0]
+        if not isinstance(confidence, np.ndarray):
+            raise TypeError("confidence must be a NumPy array or Torch tensor.")
+
+        return cls(
+            id=id,
+            sequence=sequence,  # type: ignore
+            chain_id=chain_id,
+            entity_id=entity_id,
+            atom37_positions=atom37_positions,
+            atom37_mask=atom_mask.astype(bool),
+            residue_index=residue_index,
+            insertion_code=insertion_code,
+            confidence=confidence,
+        )
+
+    @classmethod
+    def from_backbone_atom_coordinates(
+        cls, backbone_atom_coordinates: np.ndarray | torch.Tensor, **kwargs
+    ):
+        """Create a ProteinChain from a set of backbone atom coordinates.
+
+        This function simply expands the seqlen x 3 x 3 array of backbone atom
+        coordinates to a seqlen x 37 x 3 array of all atom coordinates, with the padded
+        positions set to infinity. This allows us to use from_atom37 to create the
+        appropriate ProteinChain object with the appropriate atom37_mask.
+
+        This function passes all kwargs to from_atom37.
+        """
+        if isinstance(backbone_atom_coordinates, torch.Tensor):
+            backbone_atom_coordinates = backbone_atom_coordinates.cpu().numpy()
+            if backbone_atom_coordinates.ndim == 4:
+                if backbone_atom_coordinates.shape[0] != 1:
+                    raise ValueError(
+                        f"Cannot handle batched inputs, backbone_atom_coordinates has "
+                        f"shape {backbone_atom_coordinates.shape}"
+                    )
+                backbone_atom_coordinates = backbone_atom_coordinates[0]
+
+        if not isinstance(backbone_atom_coordinates, np.ndarray):
+            raise TypeError(
+                "backbone_atom_coordinates must be a NumPy array or Torch tensor."
+            )
+        if backbone_atom_coordinates.ndim != 3 or backbone_atom_coordinates.shape[-2:] != (
+            3,
+            3,
+        ):
+            raise ValueError(
+                "backbone_atom_coordinates must have shape (length, 3, 3), got "
+                f"{backbone_atom_coordinates.shape}."
+            )
+
+        atom37_positions = np.full(
+            (backbone_atom_coordinates.shape[0], 37, 3),
+            np.inf,
+            dtype=backbone_atom_coordinates.dtype,
+        )
+        atom37_positions[:, :3, :] = backbone_atom_coordinates
+
+        return cls.from_atom37(atom37_positions=atom37_positions, **kwargs)
+
+    @classmethod
+    def from_pdb(
+        cls,
+        path: PathOrBuffer,
+        chain_id: str = "detect",
+        id: str | None = None,
+        is_predicted: bool = False,
+    ) -> ProteinChain:
+        """Return a ProteinChain object from an pdb file. NOTE: prefer mmcif for rcsb PDB files.
+        This function is mostly to interface with old PDB files and predicted structures -
+        it will not fill out the entity id correctly
+
+        Args:
+            path: PDB path or text buffer.
+            id: Optional structure identifier.
+            is_predicted (bool): If True, reads b factor as the confidence readout. Default: False.
+            chain_id: Author chain identifier. ``"detect"`` selects the first chain.
+        """
+
+        if id is not None:
+            file_id = id
+        else:
+            match path:
+                case Path() | str():
+                    file_id = Path(path).with_suffix("").name
+                case _:
+                    file_id = "null"
+
+        atom_array = PDBFile.read(path).get_structure(model=1, extra_fields=["b_factor"])
+        if len(atom_array) == 0:
+            raise ValueError("PDB contains no atoms.")
+        if chain_id == "detect":
+            chain_id = atom_array.chain_id[0]
+        atom_array = atom_array[
+            bs.filter_amino_acids(atom_array)
+            & ~atom_array.hetero
+            & (atom_array.chain_id == chain_id)
+        ]
+        if len(atom_array) == 0:
+            raise ValueError(f"PDB contains no amino-acid atoms for chain {chain_id!r}.")
+
+        entity_id = 1  # Not supplied in PDBfiles
+
+        sequence = "".join(
+            residue_constants.restype_3to1.get(monomer[0].res_name, "X")
+            for monomer in bs.residue_iter(atom_array)
+        )
+        num_res = len(sequence)
+
+        atom_positions = np.full(
+            [num_res, residue_constants.atom_type_num, 3], np.nan, dtype=np.float32
+        )
+        atom_mask = np.full([num_res, residue_constants.atom_type_num], False, dtype=bool)
+        residue_index = np.full([num_res], -1, dtype=np.int64)
+        insertion_code = np.full([num_res], "", dtype="<U4")
+
+        confidence = np.ones([num_res], dtype=np.float32)
+
+        for i, res in enumerate(bs.residue_iter(atom_array)):
+            res_index = res[0].res_id
+            residue_index[i] = res_index
+            insertion_code[i] = res[0].ins_code
+
+            # Atom level features
+            for atom in res:
+                atom_name = atom.atom_name
+                if atom_name == "SE" and atom.res_name == "MSE":
+                    # Put the coords of the selenium atom in the sulphur column
+                    atom_name = "SD"
+
+                if atom_name in residue_constants.atom_order:
+                    atom_positions[i, residue_constants.atom_order[atom_name]] = atom.coord
+                    atom_mask[i, residue_constants.atom_order[atom_name]] = True
+                    if is_predicted and atom_name == "CA":
+                        confidence[i] = atom.b_factor / PLDDT_B_FACTOR_SCALE
+
+        if not sequence or not all(sequence):
+            raise ValueError("Some residue name was not specified correctly.")
+
+        return cls(
+            id=file_id,
+            sequence=sequence,
+            chain_id=chain_id,
+            entity_id=entity_id,
+            atom37_positions=atom_positions,
+            atom37_mask=atom_mask.astype(bool),
+            residue_index=residue_index,
+            insertion_code=insertion_code,
+            confidence=confidence,
+            mmcif=None,
+        )
+
+    @classmethod
+    def from_mds(cls, data: dict[str, Any]) -> ProteinChain:
+        return cls(
+            id=data["id"],
+            chain_id=data["chain_id"],
+            entity_id=data["entity_id"],
+            sequence=data["sequence"],
+            residue_index=data["residue_index"],
+            insertion_code=np.asarray(data["insertion_code"]),
+            atom37_positions=data["atom37_positions"],
+            atom37_mask=data["atom37_mask"].astype(bool),
+            confidence=data["confidence"],
+            mmcif=None,
+        )
+
+    @classmethod
+    def from_rcsb(
+        cls,
+        pdb_id: str,
+        chain_id: str | None = None,
+        entity_id: int | None = None,
+        keep_source: bool = False,
+    ) -> ProteinChain:
+        f: io.StringIO = rcsb.fetch(pdb_id, "cif")  # type: ignore
+        return cls.from_mmcif(
+            f,
+            id=pdb_id,
+            chain_id=chain_id,
+            entity_id=entity_id,
+            keep_source=keep_source,
+            is_predicted=False,
+        )
+
+    @classmethod
+    def from_atomarray(
+        cls, atom_array: bs.AtomArray, id: str | None = None, is_predicted: bool = False
+    ) -> ProteinChain:
+        """A simple converter from bs.AtomArray -> ProteinChain.
+        Uses PDB file format as intermediate."""
+        atom_array = atom_array.copy()
+        atom_array.box = None  # remove surrounding box, from_pdb won't handle this
+        pdb_file = PDBFile()  # pyright: ignore
+        pdb_file.set_structure(atom_array)
+
+        buf = io.StringIO()
+        pdb_file.write(buf)
+        buf.seek(0)
+        return cls.from_pdb(buf, id=id, is_predicted=is_predicted)
+
+    # Object invariants and atom views
+    def __post_init__(self):
+        if not isinstance(self.id, str):
+            raise TypeError("id must be a string.")
+        if not isinstance(self.sequence, str):
+            raise TypeError("sequence must be a string.")
+        if not isinstance(self.chain_id, str) or not self.chain_id:
+            raise ValueError("chain_id must be a non-empty string.")
+        if self.entity_id is not None and (
+            not isinstance(self.entity_id, int) or isinstance(self.entity_id, bool)
+        ):
+            raise TypeError("entity_id must be an integer or None.")
+        sequence_length = len(self.sequence)
+        aligned = {
+            "atom37_positions": self.atom37_positions,
+            "atom37_mask": self.atom37_mask,
+            "residue_index": self.residue_index,
+            "insertion_code": self.insertion_code,
+            "confidence": self.confidence,
+        }
+        for name, values in aligned.items():
+            if not isinstance(values, np.ndarray):
+                raise TypeError(f"{name} must be a NumPy array, got {type(values).__name__}.")
+            if values.ndim == 0 or values.shape[0] != sequence_length:
+                raise ValueError(
+                    f"{name} shape {values.shape} does not align with "
+                    f"sequence length {sequence_length}."
+                )
+        if self.atom37_positions.shape != (sequence_length, 37, 3):
+            raise ValueError(
+                "atom37_positions must have shape "
+                f"({sequence_length}, 37, 3), got {self.atom37_positions.shape}."
+            )
+        if self.atom37_mask.shape != (sequence_length, 37):
+            raise ValueError(
+                "atom37_mask must have shape "
+                f"({sequence_length}, 37), got {self.atom37_mask.shape}."
+            )
+        if self.atom37_mask.dtype != bool:
+            raise TypeError(f"atom37_mask must have Boolean dtype, got {self.atom37_mask.dtype}.")
+        if not np.issubdtype(self.atom37_positions.dtype, np.number):
+            raise TypeError("atom37_positions must use a numeric dtype.")
+        if not np.issubdtype(self.residue_index.dtype, np.integer):
+            raise TypeError("residue_index must use an integer dtype.")
+        if self.insertion_code.dtype.kind not in {"U", "S", "O"}:
+            raise TypeError("insertion_code must use a string-compatible dtype.")
+        if any(not isinstance(value, str) for value in self.insertion_code.tolist()):
+            raise TypeError("insertion_code must contain only strings.")
+        for name, values in (
+            ("residue_index", self.residue_index),
+            ("insertion_code", self.insertion_code),
+            ("confidence", self.confidence),
+        ):
+            if values.shape != (sequence_length,):
+                raise ValueError(
+                    f"{name} must have shape ({sequence_length},), got {values.shape}."
+                )
+        if not np.issubdtype(self.confidence.dtype, np.number):
+            raise TypeError("confidence must use a numeric dtype.")
+        atom37_confidence = self.atom37_confidence
+        if atom37_confidence is not None and not isinstance(atom37_confidence, np.ndarray):
+            raise TypeError("atom37_confidence must be a NumPy array when provided.")
+        if (
+            isinstance(atom37_confidence, np.ndarray)
+            and atom37_confidence.shape != self.atom37_mask.shape
+        ):
+            raise ValueError(
+                "atom37_confidence shape must match atom37_mask: "
+                f"{atom37_confidence.shape} != {self.atom37_mask.shape}."
+            )
+        if isinstance(atom37_confidence, np.ndarray) and not np.issubdtype(
+            atom37_confidence.dtype, np.number
+        ):
+            raise TypeError("atom37_confidence must use a numeric dtype.")
+
+    @cached_property
+    def atoms(self) -> AtomIndexer:
+        return AtomIndexer(self, property="atom37_positions", dim=-2)
+
+    @cached_property
+    def atom_mask(self) -> AtomIndexer:
+        return AtomIndexer(self, property="atom37_mask", dim=-1)
+
+    @cached_property
+    def atom_array(self) -> bs.AtomArray:
+        atoms = []
+        for res_idx_i, (
+            res_name,
+            res_idx,
+            ins_code,
+            positions,
+            mask,
+            conf,
+        ) in enumerate(
+            zip(
+                self.sequence,
+                self.residue_index,
+                self.insertion_code,
+                self.atom37_positions,
+                self.atom37_mask.astype(bool),
+                self.confidence,
+                strict=False,
+            )
+        ):
+            for i, pos in zip(np.where(mask)[0], positions[mask], strict=False):
+                b_factor = (
+                    self.atom37_confidence[res_idx_i, i]
+                    if self.atom37_confidence is not None
+                    else conf
+                )
+                atom = bs.Atom(
+                    coord=pos,
+                    chain_id="A" if self.chain_id is None else self.chain_id,
+                    res_id=res_idx,
+                    ins_code=ins_code,
+                    res_name=residue_constants.restype_1to3.get(res_name, "UNK"),
+                    hetero=False,
+                    atom_name=residue_constants.atom_types[i],
+                    element=residue_constants.atom_types[i][0],
+                    b_factor=float(b_factor) * PLDDT_B_FACTOR_SCALE,
+                    occupancy=1.0,
+                )
+                atoms.append(atom)
+        return bs.array(atoms)
+
+    # Coordinate transformations and dataset adapters
+    def get_normalization_frame(self) -> Affine3D:
+        """Given a set of coordinates, compute a single frame.
+        The frame is built from the mean N, C-alpha, and C coordinates with
+        Gram-Schmidt orthogonalization. Its origin is the mean C-alpha position.
+
+        Returns:
+            Affine3D: [] tensor of Affine3D frame
+        """
+        coords = torch.from_numpy(self.atom37_positions)
+        frame = get_protein_normalization_frame(coords)
+
+        return frame
+
+    def apply_frame(self, frame: Affine3D) -> ProteinChain:
+        """Given a frame, apply the frame to the protein's coordinates.
+
+        Args:
+            frame (Affine3D): [] tensor of Affine3D frame
+
+        Returns:
+            ProteinChain: Transformed protein chain
+        """
+        coords = torch.from_numpy(self.atom37_positions).to(frame.trans.dtype)
+        coords = apply_frame_to_coords(coords, frame)
+        atom37_positions = coords.numpy()
+        return replace(self, atom37_positions=atom37_positions)
+
+    def normalize_coordinates(self) -> ProteinChain:
+        """Normalize the coordinates of the protein chain."""
+        return self.apply_frame(self.get_normalization_frame())
+
+    def infer_oxygen(self) -> ProteinChain:
+        """Oxygen position is fixed given N, CA, C atoms. Infer it if not provided."""
+        O_missing_indices = np.argwhere(~np.isfinite(self.atoms["O"]).all(axis=1)).squeeze()
+
+        O_vector = torch.tensor([0.6240, -1.0613, 0.0103], dtype=torch.float32)
+        N, CA, C = torch.from_numpy(self.atoms[["N", "CA", "C"]]).float().unbind(dim=1)
+        N = torch.roll(N, -3)
+        N[..., -1, :] = torch.nan
+
+        # Get the frame defined by the CA-C-N atom
+        frames = Affine3D.from_graham_schmidt(CA, C, N)
+        oxygen_coordinates = frames.apply(O_vector)
+        atom37_positions = self.atom37_positions.copy()
+        atom37_mask = self.atom37_mask.copy()
+
+        atom37_positions[O_missing_indices, residue_constants.atom_order["O"]] = oxygen_coordinates[
+            O_missing_indices
+        ].numpy()
+        atom37_mask[O_missing_indices, residue_constants.atom_order["O"]] = ~np.isnan(
+            atom37_positions[O_missing_indices, residue_constants.atom_order["O"]]
+        ).any(-1)
+        new_chain = replace(self, atom37_positions=atom37_positions, atom37_mask=atom37_mask)
+        return new_chain
+
+    @cached_property
+    def inferred_cbeta(self) -> np.ndarray:
+        """Infer cbeta positions based on N, C, CA."""
+        N, CA, C = np.moveaxis(self.atoms[["N", "CA", "C"]], 1, 0)
+        # See usage in trDesign codebase.
+        # https://github.com/gjoni/trDesign/blob/f2d5930b472e77bfacc2f437b3966e7a708a8d37/02-GD/utils.py#L140
+        CB = infer_cb(C, N, CA, 1.522, 1.927, -2.143)
+        return CB
+
+    def infer_cbeta(self, infer_cbeta_for_glycine: bool = False) -> ProteinChain:
+        """Return a new chain with inferred CB atoms at all residues except GLY.
+
+        Args:
+            infer_cbeta_for_glycine (bool): If True, infers a beta carbon for glycine
+                residues, even though that residue doesn't have one.  Default off.
+
+                NOTE(rverkuil): The reason for having this switch in the first place
+                is that sometimes we want a (inferred) CB coordinate for every residue,
+                for example for making a pairwise distance matrix, or doing an RMSD
+                calculation between two designs for a given structural template, w/
+                CB atoms.
+        """
+        atom37_positions = self.atom37_positions.copy()
+        atom37_mask = self.atom37_mask.copy()
+
+        inferred_cbeta_positions = self.inferred_cbeta
+        if not infer_cbeta_for_glycine:
+            inferred_cbeta_positions[np.array(list(self.sequence)) == "G", :] = np.nan
+
+        atom37_positions[:, residue_constants.atom_order["CB"]] = inferred_cbeta_positions
+        atom37_mask[:, residue_constants.atom_order["CB"]] = ~np.isnan(
+            atom37_positions[:, residue_constants.atom_order["CB"]]
+        ).any(-1)
+        new_chain = replace(self, atom37_positions=atom37_positions, atom37_mask=atom37_mask)
+        return new_chain
+
+    @cached_property
+    def pdist_CA(self) -> np.ndarray:
+        CA = self.atoms["CA"]
+        pdist_CA = squareform(pdist(CA))
+        return pdist_CA
+
+    @cached_property
+    def pdist_CB(self) -> np.ndarray:
+        pdist_CB = squareform(pdist(self.inferred_cbeta))
+        return pdist_CB
+
+    @classmethod
+    def as_complex(cls, chains: Sequence[ProteinChain]):
+        raise RuntimeError(
+            ".as_complex() has been deprecated in favor of .concat(). "
+            ".concat() will eventually be deprecated in favor of ProteinComplex..."
+        )
+
+    @classmethod
+    def concat(cls, chains: Sequence[ProteinChain], use_chainbreak: bool = True):
+        if not chains:
+            raise ValueError("chains must contain at least one ProteinChain.")
+        if any(not isinstance(chain, ProteinChain) for chain in chains):
+            raise TypeError("chains must contain only ProteinChain instances.")
+        sep_tokens = {
+            "residue_index": np.array([-1]),
+            "insertion_code": np.array([""]),
+            "atom37_positions": np.full([1, 37, 3], np.inf),
+            "atom37_mask": np.zeros([1, 37], dtype=bool),
+            "confidence": np.array([0]),
+        }
+
+        def join_arrays(arrays: Sequence[np.ndarray], sep: np.ndarray):
+            if use_chainbreak:
+                full_array = []
+                for array in arrays:
+                    full_array.append(array)
+                    full_array.append(sep)
+                full_array = full_array[:-1]
+                return np.concatenate(full_array, 0)
+            else:
+                return np.concatenate(arrays, 0)
+
+        array_args: dict[str, np.ndarray] = {
+            name: join_arrays([getattr(chain, name) for chain in chains], sep)
+            for name, sep in sep_tokens.items()
+        }
+
+        chain_break = residue_constants.CHAIN_BREAK_TOKEN if use_chainbreak else ""
+        return cls(
+            id=chains[0].id,
+            sequence=chain_break.join(chain.sequence for chain in chains),
+            chain_id="A",
+            entity_id=None,
+            mmcif=None,
+            **array_args,
+        )
+
+    def find_nonpolymer_contacts(self):
+        if self.mmcif is None:
+            raise ValueError(
+                "find_nonpolymer_contacts requires a chain loaded with keep_source=True."
+            )
+        nonpolymer_and_chain_id_to_array = self.mmcif.non_polymer_coords
+
+        results = []
+        for (
+            nonpolymer,
+            _,
+        ), nonpolymer_array in nonpolymer_and_chain_id_to_array.items():
+            if nonpolymer_array.coord is None:
+                raise ValueError(
+                    f"Non-polymer {nonpolymer.comp_id!r} has no coordinate table."
+                )
+            chain_coords = self.atom37_positions[self.atom37_mask]
+            distance = cdist(nonpolymer_array.coord, chain_coords)
+
+            is_contact = distance < 5
+            if not is_contact.any():
+                continue
+            contacting_atoms = np.where(is_contact.any(0))[0]
+            chain_index = np.where(self.atom37_mask)[0]
+            contacting_residues = np.unique(chain_index[contacting_atoms])
+
+            result = {
+                "ligand": nonpolymer.name,
+                "ligand_id": nonpolymer.comp_id,
+                "contacting_residues": contacting_residues.tolist(),
+            }
+            results.append(result)
+        return results
+
+    def select_residue_indices(
+        self, indices: list[int | str], ignore_x_mismatch: bool = False
+    ) -> ProteinChain:
+        numeric_indices = [idx if isinstance(idx, int) else int(idx[1:]) for idx in indices]
+        mask = np.isin(self.residue_index, numeric_indices)
+        new = self[mask]
+        mismatches = []
+        for aa, idx in zip(new.sequence, indices, strict=False):
+            if isinstance(idx, int):
+                continue
+            if aa == "X" and ignore_x_mismatch:
+                continue
+            if aa != idx[0]:
+                mismatches.append((aa, idx))
+        if mismatches:
+            mismatch_str = "; ".join(
+                f"Position {idx[1:]}, Expected: {idx[0]}, Received: {aa}" for aa, idx in mismatches
+            )
+            raise RuntimeError(mismatch_str)
+
+        return new
+
+    def to_structure_encoder_inputs(
+        self,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Convert protein chain to structure encoder inputs.
+
+        Returns:
+            tuple: (coordinates, plddt, residue_index) where:
+                - coordinates: X with shape (1, l, 37, 3), containing atom positions
+                - plddt: P with shape (1, l), containing confidence scores
+                - residue_index: R with shape (1, l), containing residue indices
+        """
+        # Convert to tensors and add batch dimension
+        coordinates = (
+            torch.from_numpy(self.atom37_positions).float().unsqueeze(0)
+        )  # X has shape (1, l, 37, 3).
+        plddt = torch.from_numpy(self.confidence).float().unsqueeze(0)  # P: (1, l)
+        residue_index = (
+            torch.from_numpy(self.residue_index).long().unsqueeze(0)
+        )  # R has shape (1, l).
+
+        return coordinates, plddt, residue_index
+
+    # Sequence access, interchange, and compact storage
+    def __getitem__(self, idx: int | list[int] | slice | np.ndarray | torch.Tensor):
+        if isinstance(idx, int):
+            idx = [idx]
+        if isinstance(idx, torch.Tensor):
+            idx = idx.cpu().numpy()
+
+        sequence = slice_python_object_as_numpy(self.sequence, idx)
+        return replace(
+            self,
+            sequence=sequence,
+            residue_index=self.residue_index[..., idx],
+            insertion_code=self.insertion_code[..., idx],
+            atom37_positions=self.atom37_positions[..., idx, :, :],
+            atom37_mask=self.atom37_mask[..., idx, :],
+            confidence=self.confidence[..., idx],
+            atom37_confidence=self.atom37_confidence[..., idx, :]
+            if self.atom37_confidence is not None
+            else None,
+        )
+
+    def __len__(self):
+        return len(self.sequence)
+
+    def cbeta_contacts(self, distance_threshold: float = 8.0) -> np.ndarray:
+        distance = self.pdist_CB
+        contacts = (distance < distance_threshold).astype(np.int64)
+        contacts[np.isnan(distance)] = -1
+        np.fill_diagonal(contacts, -1)
+        return contacts
+
+    def to_pdb(self, path: PathOrBuffer, include_insertions: bool = True):
+        """Dssp works better w/o insertions."""
+        f = PDBFile()
+        if not include_insertions:
+            f.set_structure(self.atom_array_no_insertions)
+        else:
+            f.set_structure(self.atom_array)
+        f.write(path)
+
+    def to_pdb_string(self, include_insertions: bool = True) -> str:
+        buf = io.StringIO()
+        self.to_pdb(buf, include_insertions=include_insertions)
+        buf.seek(0)
+        return buf.read()
+
+    def to_mmcif(self, path: PathOrBuffer):
+        f = CIFFile()
+        set_structure_pdbx(f, self.atom_array, data_block=self.id)
+
+        # incantations molstar needs to render pLDDT / confidence onto
+        # the structure with "alphafold-view"
+        f.block["ma_qa_metric"] = CIFCategory(
+            name="ma_qa_metric",
+            columns={
+                "id": CIFColumn(data=CIFData(array=np.array([1, 2]), dtype=np.int64)),
+                "mode": CIFColumn(data=CIFData(array=np.array(["global", "local"]), dtype=np.str_)),
+                "name": CIFColumn(data=CIFData(array=np.array(["pLDDT", "pLDDT"]), dtype=np.str_)),
+            },
+        )
+
+        # table is a duplicate of data already in the atom array, but
+        # needed by molstar to render pLDDT / confidence
+        resid_pldd_table = {
+            # hard coded to as we currently only support single chain structures
+            "label_asym_id": CIFColumn(
+                data=CIFData(array=[CHAIN_ID_CONST] * len(self.residue_index), dtype=np.str_)
+            ),
+            "label_comp_id": CIFColumn(
+                data=CIFData(
+                    array=[residue_constants.restype_1to3.get(c, "UNK") for c in self.sequence],
+                    dtype=np.str_,
+                )
+            ),
+            "label_seq_id": CIFColumn(data=CIFData(array=self.residue_index, dtype=np.int64)),
+            "ordinal_id": CIFColumn(data=CIFData(array=self.residue_index, dtype=np.int64)),
+            # hard coded to show these are all local plDDT values
+            "metric_id": CIFColumn(
+                data=CIFData(array=["2"] * len(self.residue_index), dtype=np.str_)
+            ),
+            "metric_value": CIFColumn(
+                data=CIFData(
+                    array=self.confidence * PLDDT_B_FACTOR_SCALE,
+                    dtype=np.float32,
+                )
+            ),
+            # hard coded to show there are the initial version, there are no revisions
+            "model_id": CIFColumn(
+                data=CIFData(array=["1"] * len(self.residue_index), dtype=np.str_)
+            ),
+        }
+        f.block["ma_qa_metric_local"] = CIFCategory(
+            name="ma_qa_metric_local", columns=resid_pldd_table
+        )
+        round_mmcif_columns(f)
+        f.write(path)
+
+    def to_mmcif_string(self) -> str:
+        buf = io.StringIO()
+        self.to_mmcif(buf)
+        buf.seek(0)
+        return buf.read()
+
+    def state_dict(self, backbone_only=False, json_serializable=False):
+        """This state dict is optimized for storage, so it turns things to fp16 whenever
+        possible. Note that we also only support int32 residue indices, I'm hoping we don't
+        need more than 2**32 residues..."""
+        dct = {k: v for k, v in asdict(self).items() if k not in ["mmcif"]}
+        if backbone_only:
+            dct["atom37_mask"][:, 3:] = False
+        dct["atom37_positions"] = dct["atom37_positions"][dct["atom37_mask"]]
+        if dct.get("atom37_confidence") is not None:
+            dct["atom37_confidence"] = dct["atom37_confidence"][dct["atom37_mask"]]
+        else:
+            dct.pop("atom37_confidence", None)
+
+        for k, v in dct.items():
+            if isinstance(v, np.ndarray):
+                match v.dtype:
+                    case np.int64:
+                        dct[k] = v.astype(np.int32)
+                    case np.float64 | np.float32:
+                        dct[k] = v.astype(np.float16)
+                    case _:
+                        pass
+                if json_serializable:
+                    dct[k] = v.tolist()
+        return dct
+
+    def to_blob(self, backbone_only=False) -> bytes:
+        payload = msgpack.dumps(self.state_dict(backbone_only), default=msgpack_numpy.encode)
+        return brotli.compress(payload, quality=5)
+
+    @classmethod
+    def from_open_source(cls, pc: ProteinChain):
+        return cls(**vars(pc))
+
+    @classmethod
+    def from_state_dict(cls, dct):
+        # Note: assembly_composition is *supposed* to have string keys.
+        dct = _str_key_to_int_key(dct, ignore_keys=["assembly_composition"])
+
+        for k, v in dct.items():
+            if isinstance(v, list):
+                dct[k] = np.array(v)
+
+        atom37 = np.full((*dct["atom37_mask"].shape, 3), np.nan)
+        atom37[dct["atom37_mask"]] = dct["atom37_positions"]
+        dct["atom37_positions"] = atom37
+        if "atom37_confidence" in dct:
+            atom37_conf = np.full(dct["atom37_mask"].shape, np.nan, dtype=np.float32)
+            atom37_conf[dct["atom37_mask"]] = dct["atom37_confidence"]
+            dct["atom37_confidence"] = atom37_conf
+        dct = {
+            k: (
+                v.astype(np.float32)
+                if k in ["atom37_positions", "confidence", "atom37_confidence"]
+                else v
+            )
+            for k, v in dct.items()
+            if not (k == "atom37_confidence" and v is None)
+        }
+        return cls(**dct, mmcif=None)
+
+    @classmethod
+    def from_blob(cls, input: Path | str | io.BytesIO | bytes):
+        """NOTE(@zlin): blob + sparse coding + brotli + fp16 reduces memory
+        of chains from 52G/1M chains to 20G/1M chains, I think this is a good first
+        shot at compressing and dumping chains to disk. I'm sure there's better ways."""
+        match input:
+            case Path() | str():
+                bytes = Path(input).read_bytes()
+            case io.BytesIO():
+                bytes = input.getvalue()
+            case _:
+                bytes = input
+        state = msgpack.loads(brotli.decompress(bytes), object_hook=msgpack_numpy.decode)
+        return cls.from_state_dict(state)
+
+    # Surface and structural comparison metrics
+    def sasa(self, by_residue: bool = True):
+        arr = self.atom_array_no_insertions
+        if len(arr) == 0:
+            raise ValueError("SASA requires at least one resolved atom.")
+        sasa_per_atom = bs.sasa(arr)  # type: ignore
+        if by_residue:
+            # Sum per-atom SASA into residue "bins", with np.bincount.
+            if arr.res_id is None:
+                raise RuntimeError("Biotite AtomArray is missing residue identifiers.")
+            # Residue IDs are one-indexed, so discard the unused zero bin.
+            # NOTE(aderry): We compute only for residues with coordinates, return NaN otherwise.
+            num_trailing_residues = len(self) - arr.res_id.max()
+            sasa_per_residue = np.concatenate(
+                [
+                    np.bincount(arr.res_id, weights=sasa_per_atom)[1:],
+                    np.zeros(num_trailing_residues),
+                ]
+            )
+            sasa_per_residue[~self.atom37_mask.any(-1)] = np.nan
+            if len(sasa_per_residue) != len(self):
+                raise RuntimeError("Residue SASA output does not align with the protein chain.")
+            return sasa_per_residue
+        return sasa_per_atom
+
+    def sap_score(self, aggregation: str = "atom") -> np.ndarray:
+        """Compute per-atom spatial aggregation propensity (SAP).
+
+        Residue aggregation averages resolved atoms and omits unresolved residues.
+        Protein aggregation sums positive atom scores, following Lauer et al. 2011.
+        """
+        sap_radius = 5.0
+        arr = self.atom_array_no_insertions
+        if len(arr) == 0:
+            raise ValueError("SAP requires at least one resolved atom.")
+
+        for name in ("res_id", "res_name", "atom_name", "coord"):
+            if getattr(arr, name) is None:
+                raise RuntimeError(f"Biotite AtomArray is missing required {name!r} data.")
+
+        # compute SASA and residue-specific properties
+        sasa_per_atom = self.sasa(by_residue=False)
+        resid_to_resname = dict(zip(arr.res_id, arr.res_name, strict=False))
+
+        max_side_chain_asa = np.full(len(self), np.nan)
+        res_hydrophobicity = np.full(len(self), np.nan)
+        resolved_res_mask = self.atom37_mask.any(-1)
+        num_trailing_residues = len(self) - arr.res_id.max()
+
+        max_side_chain_asa[resolved_res_mask] = np.array(
+            [residue_constants.side_chain_asa[resid_to_resname[i]] for i in np.unique(arr.res_id)]
+        )
+        res_hydrophobicity[resolved_res_mask] = np.array(
+            [residue_constants.hydrophobicity[resid_to_resname[i]] for i in np.unique(arr.res_id)]
+        )
+
+        # compute SAP score
+        is_side_chain = ~bs.filter_peptide_backbone(arr)
+        sasa_per_atom[is_side_chain] = 0
+        kdtree = KDTree(arr.coord)
+        neighbors = kdtree.query_ball_tree(kdtree, sap_radius, p=2.0)
+        sap_by_atom = np.zeros_like(sasa_per_atom)
+        for i, nn_list in enumerate(neighbors):
+            saa_nn = np.zeros_like(sasa_per_atom)
+            saa_nn[nn_list] = sasa_per_atom[nn_list]
+            sasa_within_r = np.concatenate(
+                [
+                    np.bincount(arr.res_id, weights=saa_nn)[1:],
+                    np.zeros(num_trailing_residues),
+                ]
+            )
+            sap = np.nansum((sasa_within_r / max_side_chain_asa) * res_hydrophobicity)
+            sap_by_atom[i] = sap
+
+        match aggregation:
+            case "atom":
+                return sap_by_atom
+            case "residue":
+                sap_by_residue = np.concatenate(
+                    [
+                        np.bincount(arr.res_id, weights=sap_by_atom)[1:],
+                        np.zeros(num_trailing_residues),
+                    ]
+                ) / (
+                    np.concatenate([np.bincount(arr.res_id)[1:], np.zeros(num_trailing_residues)])
+                    + 1e-8
+                )
+                sap_by_residue[~resolved_res_mask] = np.nan
+                if len(sap_by_residue) != len(self):
+                    raise RuntimeError("Residue SAP output does not align with the protein chain.")
+                return sap_by_residue
+            case "protein":
+                return sum(sap_by_atom[sap_by_atom > 0])  # pyright: ignore[reportReturnType]
+            case _:
+                raise ValueError(
+                    f"Invalid aggregation method: {aggregation}. Must be one of "
+                    "'atom', 'residue', or 'protein'"
+                )
+
+    def globularity(self) -> float:
+        # Computes globularity using total volumes divided by MVEE.
+        # We make the simplifying approximation that atoms never overlap.
+        # The globularity is only computed where structure exists.
+        # Besides the approximation above, this is inspired by:
+
+        # https://www.mdpi.com/2073-4352/11/12/1539
+        # The non-overlapping-atom approximation can produce globularity above one.
+        mask = self.atom37_mask.any(-1)
+        points = self.atom37_positions[self.atom37_mask]
+        sequence = [aa for aa, m in zip(self.sequence, mask, strict=False) if m]  # type: ignore
+        A, _ = self._mvee(points, tol=1e-3)
+        mvee_volume = (4 * np.pi) / (3 * np.sqrt(np.linalg.det(A)))
+        volume = sum(residue_constants.amino_acid_volumes[x] for x in sequence)
+        ratio = volume / mvee_volume
+
+        # The paper compares the ellipsoidal profile with scalar t, a measurement
+        # of elongation. We want a single number, so we multiply by 1/(2t), so
+        # that value is normalized between 0-1
+        eigenvalues = np.linalg.eigvals(A)
+        R = 1 / np.sqrt(eigenvalues)
+        # ellipsoid radii length triangle inequality coefficient
+        t = max(R[0] / (R[1] + R[2]), R[1] / (R[0] + R[2]), R[2] / (R[0] + R[1]))
+        elongation_metric = 1 / max(t, 1)
+        return ratio * elongation_metric
+
+    @staticmethod
+    def _mvee(P: np.ndarray, tol, max_iter=10000):
+        # Finds minimum volume enclosing ellipsoid of a set of points.
+        # Returns A, c where the ellipse is defined as:
+        #    (x-c).T @ A @ (x-c) = 1
+        hull = ConvexHull(P)
+        P = P[hull.vertices]
+        P = P.T
+
+        # Data points
+        d, n = P.shape
+        Q = np.zeros((d + 1, n))
+        Q[:d, :] = P[:d, :n]
+        Q[d, :] = np.ones((1, n))
+
+        # Initializations
+        count = 1
+        err = 1.0
+        u = np.full((n, 1), 1 / n)  # First iteration.
+
+        # Khachiyan Algorithm
+        for _ in range(max_iter):
+            X = Q.dot(np.diag(u.squeeze())) @ Q.T
+            M = np.diag(Q.T @ np.linalg.inv(X) @ Q)
+            maximum, j = np.max(M), np.argmax(M)
+            step_size = (maximum - d - 1) / ((d + 1) * (maximum - 1))
+            new_u = (1 - step_size) * u
+            new_u[j] += step_size
+            count += 1
+            err = np.linalg.norm(new_u - u)
+            u = new_u
+            if err < tol:
+                break
+        else:
+            raise ValueError("MVEE did not converge")
+
+        d = P.shape[0]  # Fixed: use P.shape[0] instead of P.shape
+        U = np.diag(u.squeeze())
+
+        # The A matrix for the ellipse
+        A = (1 / d) * np.linalg.inv(P @ U @ P.T - (P @ u) @ (P @ u).T)
+
+        # Center of the ellipse
+        c = P @ u
+
+        return A, c
+
+    def radius_of_gyration(self):
+        arr = self.atom_array_no_insertions
+        return bs.gyration_radius(arr)
+
+    def align(
+        self,
+        target: ProteinChain,
+        mobile_inds: list[int] | np.ndarray | None = None,
+        target_inds: list[int] | np.ndarray | None = None,
+        only_use_backbone: bool = False,
+    ):
+        """
+        Aligns the current protein to the provided target.
+
+        Args:
+            target (ProteinChain): The target protein to align to.
+            mobile_inds: Mobile atom indices, not residue indices.
+            target_inds: Target atom indices, not residue indices.
+            only_use_backbone (bool, optional): If True, only align the backbone atoms.
+        """
+        aligner = Aligner(
+            self if mobile_inds is None else self[mobile_inds],
+            target if target_inds is None else target[target_inds],
+            only_use_backbone,
+        )
+
+        return aligner.apply(self)
+
+    def rmsd(
+        self,
+        target: ProteinChain,
+        also_check_reflection: bool = False,
+        mobile_inds: list[int] | np.ndarray | None = None,
+        target_inds: list[int] | np.ndarray | None = None,
+        only_compute_backbone_rmsd: bool = False,
+    ):
+        """
+        Compute the RMSD between this protein chain and another.
+
+        Args:
+            target (ProteinChain): The target (other) protein chain to compare to.
+            also_check_reflection: Compare the reflected mobile coordinates too.
+            mobile_inds: Mobile atom indices, not residue indices.
+            target_inds: Target atom indices, not residue indices.
+            only_compute_backbone_rmsd: Restrict the score to backbone atoms.
+        """
+        if isinstance(target, bs.AtomArray):
+            raise ValueError(
+                "Support for bs.AtomArray removed, use ProteinChain.from_atomarry for ProteinChain."
+            )
+        aligner = Aligner(
+            self if mobile_inds is None else self[mobile_inds],
+            target if target_inds is None else target[target_inds],
+            only_compute_backbone_rmsd,
+        )
+        avg_rmsd = aligner.rmsd
+
+        if not also_check_reflection:
+            return avg_rmsd
+
+        aligner = Aligner(
+            self if mobile_inds is None else self[mobile_inds],
+            target if target_inds is None else target[target_inds],
+            only_compute_backbone_rmsd,
+            use_reflection=True,
+        )
+        avg_rmsd_neg = aligner.rmsd
+
+        return min(avg_rmsd, avg_rmsd_neg)
+
+    def lddt_ca(
+        self,
+        native: ProteinChain,
+        mobile_inds: list[int] | np.ndarray | None = None,
+        target_inds: list[int] | np.ndarray | None = None,
+        **kwargs,
+    ) -> float | np.ndarray:
+        """Compute the LDDT between this protein chain and another. NOTE: LDDT IS NOT SYMMETRIC.
+        The call should always be prediction.lddt_ca(native).
+
+        Arguments:
+            native (ProteinChain): The ground truth protein chain
+            mobile_inds: Mobile atom indices, not residue indices.
+            target_inds: Target atom indices, not residue indices.
+
+        Returns:
+            float | np.ndarray: The LDDT score between the two protein chains, either
+                a single float or per-residue LDDT scores if `per_residue` is True.
+        """
+        lddt = compute_lddt_ca(
+            torch.tensor(self.atom37_positions[mobile_inds]).unsqueeze(0),
+            torch.tensor(native.atom37_positions[target_inds]).unsqueeze(0),
+            torch.tensor(native.atom37_mask[mobile_inds]).unsqueeze(0),
+            **kwargs,
+        )
+        return float(lddt) if lddt.numel() == 1 else lddt.numpy().flatten()
+
+    def gdt_ts(
+        self,
+        target: ProteinChain,
+        mobile_inds: list[int] | np.ndarray | None = None,
+        target_inds: list[int] | np.ndarray | None = None,
+        **kwargs,
+    ) -> float | np.ndarray:
+        """Compute the GDT_TS between this protein chain and another.
+
+        Arguments:
+            target (ProteinChain): The other protein chain to compare to.
+            mobile_inds: Mobile atom indices, not residue indices.
+            target_inds: Target atom indices, not residue indices.
+
+        Returns:
+            float: The GDT_TS score between the two protein chains.
+        """
+        gdt_ts = compute_gdt_ts(
+            mobile=torch.tensor(
+                index_by_atom_name(self.atom37_positions[mobile_inds], "CA"),
+                dtype=torch.float32,
+            ).unsqueeze(0),
+            target=torch.tensor(
+                index_by_atom_name(target.atom37_positions[target_inds], "CA"),
+                dtype=torch.float32,
+            ).unsqueeze(0),
+            atom_exists_mask=torch.tensor(
+                index_by_atom_name(self.atom37_mask[mobile_inds], "CA", dim=-1)
+                & index_by_atom_name(target.atom37_mask[target_inds], "CA", dim=-1)
+            ).unsqueeze(0),
+            **kwargs,
+        )
+        return float(gdt_ts) if gdt_ts.numel() == 1 else gdt_ts.numpy().flatten()
+
+    @cached_property
+    def residue_index_no_insertions(self) -> np.ndarray:
+        return self.residue_index + np.cumsum(self.insertion_code != "")
+
+    @cached_property
+    def atom_array_no_insertions(self) -> bs.AtomArray:
+        atoms = []
+        for res_idx, (res_name, positions, mask, conf) in enumerate(
+            zip(
+                self.sequence,
+                self.atom37_positions,
+                self.atom37_mask.astype(bool),
+                self.confidence,
+                strict=False,
+            )
+        ):
+            for i, pos in zip(np.where(mask)[0], positions[mask], strict=False):
+                b_factor = (
+                    self.atom37_confidence[res_idx, i]
+                    if self.atom37_confidence is not None
+                    else conf
+                )
+                atom = bs.Atom(
+                    coord=pos,
+                    # hard coded to as we currently only support single chain structures
+                    chain_id=CHAIN_ID_CONST,
+                    res_id=res_idx + 1,
+                    res_name=residue_constants.restype_1to3.get(res_name, "UNK"),
+                    hetero=False,
+                    atom_name=residue_constants.atom_types[i],
+                    element=residue_constants.atom_types[i][0],
+                    b_factor=float(b_factor) * PLDDT_B_FACTOR_SCALE,
+                    occupancy=1.0,
+                )
+                atoms.append(atom)
+        return bs.array(atoms)

fastplms/models/esmfold2/esmfold2_protein_complex.py

@@ -0,0 +1,1240 @@
+"""Protein-complex data, assembly expansion, and geometry for ESMFold2."""
+
+from __future__ import annotations
+
+import io
+import itertools
+import random
+import re
+import warnings
+from collections.abc import Iterable, Sequence
+from dataclasses import asdict, dataclass, replace
+from functools import cached_property
+from pathlib import Path
+from subprocess import check_output
+from tempfile import TemporaryDirectory
+from typing import Any
+
+import biotite.structure as bs
+import brotli
+import msgpack
+import msgpack_numpy
+import numpy as np
+import torch
+from biotite.database import rcsb
+from biotite.file import InvalidFileError
+from biotite.structure.io.pdb import PDBFile
+from biotite.structure.io.pdbx import CIFCategory, CIFColumn, CIFData, CIFFile
+from biotite.structure.io.pdbx import set_structure as set_structure_pdbx
+from biotite.structure.io.pdbx.convert import _get_transformations, get_structure
+from biotite.structure.util import matrix_rotate
+from scipy.spatial import KDTree
+
+from . import esmfold2_residue_constants as residue_constants
+from .esmfold2_affine3d import Affine3D
+from .esmfold2_aligner import Aligner
+from .esmfold2_atom_indexer import AtomIndexer
+from .esmfold2_metrics import compute_gdt_ts, compute_lddt_ca
+from .esmfold2_misc import slice_python_object_as_numpy
+from .esmfold2_mmcif_parsing import (
+    MmcifWrapper,
+    NoProteinError,
+    round_mmcif_columns,
+)
+from .esmfold2_protein_chain import (
+    ProteinChain,
+    _str_key_to_int_key,
+    chain_to_ndarray,
+    index_by_atom_name,
+    infer_cb,
+)
+from .esmfold2_utils_types import PathOrBuffer
+
+SINGLE_LETTER_CHAIN_IDS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
+
+
+def _parse_operation_expression(expression: str) -> list[tuple[str, ...]]:
+    """Expand an mmCIF operation expression in application order."""
+
+    def expand_group(group: str) -> list[str]:
+        operation_ids: list[str] = []
+        for term in group.split(","):
+            if "-" not in term:
+                operation_ids.append(term)
+                continue
+            first, last = (int(value) for value in term.split("-"))
+            operation_ids.extend(str(value) for value in range(first, last + 1))
+        return operation_ids
+
+    groups = [group for group in expression.replace(")", "").split("(") if group]
+    groups.reverse()
+    return list(itertools.product(*(expand_group(group) for group in groups)))
+
+
+def _apply_transformations_fast(chains, transformation_dict, operations):
+    """Return transformed copies of each affected protein chain."""
+    transformed_chains = []
+    for chain in chains:
+        for operation in operations:
+            coordinates = chain.atom37_positions.copy()
+            for op_step in operation:
+                transform = transformation_dict[op_step]
+                coordinates = matrix_rotate(coordinates, transform.rotation)
+                coordinates += transform.target_translation
+            transformed_chains.append(replace(chain, atom37_positions=coordinates))
+    return transformed_chains
+
+
+@dataclass
+class ProteinComplexMetadata:
+    entity_lookup: dict[int, int | str]
+    chain_lookup: dict[int, str]
+    mmcif: MmcifWrapper | None = None
+    # This is a dictionary that maps assembly ids to the list of unique chains
+    # in that assembly. Allows for usage of `switch_assembly`.
+    assembly_composition: dict[str, list[str]] | None = None
+
+
+@dataclass
+class DockQSingleScore:
+    native_chains: tuple[str, str]
+    DockQ: float
+    interface_rms: float
+    ligand_rms: float
+    fnat: float
+    fnonnat: float
+    clashes: float
+    F1: float
+    DockQ_F1: float
+
+
+@dataclass
+class DockQResult:
+    total_dockq: float
+    native_interfaces: int
+    chain_mapping: dict[str, str]
+    interfaces: dict[tuple[str, str], DockQSingleScore]
+    # zip(aligned.chain_iter(), native.chain_iter()) gives you the pairing
+    # aligned.rmsd(native) should give you a low rmsd irrespective of shuffling
+    aligned: ProteinComplex
+    aligned_rmsd: float
+
+
+@dataclass(frozen=True)
+class ProteinComplex:
+    """Dataclass with atom37 representation of an entire protein complex."""
+
+    id: str
+    sequence: str
+    entity_id: np.ndarray  # entities map to unique sequences
+    chain_id: np.ndarray  # multiple chains might share an entity id
+    sym_id: np.ndarray  # complexes might be copies of the same chain
+    residue_index: np.ndarray
+    insertion_code: np.ndarray
+    atom37_positions: np.ndarray
+    atom37_mask: np.ndarray
+    confidence: np.ndarray
+    # This metadata is parsed from the MMCIF file. For synthetic data, we do a best effort.
+    metadata: ProteinComplexMetadata
+    atom37_confidence: np.ndarray | None = None  # P has shape (l, 37).
+
+    # Coordinate completion, concatenation, and comparison
+    def infer_oxygen(self) -> ProteinComplex:
+        """Oxygen position is fixed given N, CA, C atoms. Infer it if not provided."""
+        O_missing_indices = np.argwhere(~np.isfinite(self.atoms["O"]).all(axis=1)).squeeze()
+
+        O_vector = torch.tensor([0.6240, -1.0613, 0.0103], dtype=torch.float32)
+        N, CA, C = torch.from_numpy(self.atoms[["N", "CA", "C"]]).float().unbind(dim=1)
+        N = torch.roll(N, -3)
+        N[..., -1, :] = torch.nan
+
+        # Get the frame defined by the CA-C-N atom
+        frames = Affine3D.from_graham_schmidt(CA, C, N)
+        oxygen_coordinates = frames.apply(O_vector)
+        atom37_positions = self.atom37_positions.copy()
+        atom37_mask = self.atom37_mask.copy()
+
+        atom37_positions[O_missing_indices, residue_constants.atom_order["O"]] = oxygen_coordinates[
+            O_missing_indices
+        ].numpy()
+        atom37_mask[O_missing_indices, residue_constants.atom_order["O"]] = ~np.isnan(
+            atom37_positions[O_missing_indices, residue_constants.atom_order["O"]]
+        ).any(-1)
+        new_chain = replace(self, atom37_positions=atom37_positions, atom37_mask=atom37_mask)
+        return new_chain
+
+    def infer_cbeta(self, infer_cbeta_for_glycine: bool = False) -> ProteinComplex:
+        """Return a new chain with inferred CB atoms at all residues except GLY.
+
+        Args:
+            infer_cbeta_for_glycine (bool): If True, infers a beta carbon for glycine
+                residues, even though that residue doesn't have one.  Default off.
+
+                NOTE(rverkuil): The reason for having this switch in the first place
+                is that sometimes we want a (inferred) CB coordinate for every residue,
+                for example for making a pairwise distance matrix, or doing an RMSD
+                calculation between two designs for a given structural template, w/
+                CB atoms.
+        """
+        atom37_positions = self.atom37_positions.copy()
+        atom37_mask = self.atom37_mask.copy()
+
+        N, CA, C = np.moveaxis(self.atoms[["N", "CA", "C"]], 1, 0)
+        # See usage in trDesign codebase.
+        # https://github.com/gjoni/trDesign/blob/f2d5930b472e77bfacc2f437b3966e7a708a8d37/02-GD/utils.py#L140
+        inferred_cbeta_positions = infer_cb(C, N, CA, 1.522, 1.927, -2.143)
+        if not infer_cbeta_for_glycine:
+            inferred_cbeta_positions[np.array(list(self.sequence)) == "G", :] = np.nan
+
+        atom37_positions[:, residue_constants.atom_order["CB"]] = inferred_cbeta_positions
+        atom37_mask[:, residue_constants.atom_order["CB"]] = ~np.isnan(
+            atom37_positions[:, residue_constants.atom_order["CB"]]
+        ).any(-1)
+        new_chain = replace(self, atom37_positions=atom37_positions, atom37_mask=atom37_mask)
+        return new_chain
+
+    @classmethod
+    def from_open_source(cls, pc: ProteinComplex):
+        # TODO(@zeming): deprecated, should delete
+        return pc
+
+    @classmethod
+    def concat(cls, objs: list[ProteinComplex]) -> ProteinComplex:
+        pdb_ids = [obj.id for obj in objs]
+        if len(set(pdb_ids)) > 1:
+            raise RuntimeError(
+                "Concatention of protein complexes across different PDB ids is unsupported"
+            )
+        return ProteinComplex.from_chains(
+            list(itertools.chain.from_iterable(obj.chain_iter() for obj in objs))
+        )
+
+    def _sanity_check_complexes_are_comparable(self, other: ProteinComplex):
+        if len(self) != len(other):
+            raise ValueError("Protein complexes must have the same length")
+        if len(list(self.chain_iter())) != len(list(other.chain_iter())):
+            raise ValueError("Protein complexes must have the same number of chains")
+
+    def rmsd(
+        self,
+        target: ProteinComplex,
+        also_check_reflection: bool = False,
+        mobile_inds: list[int] | np.ndarray | None = None,
+        target_inds: list[int] | np.ndarray | None = None,
+        only_compute_backbone_rmsd: bool = False,
+        compute_chain_assignment: bool = True,
+    ):
+        """
+        Compute the RMSD between this protein chain and another.
+
+        Args:
+            target (ProteinComplex): The target (other) protein complex to compare to.
+            also_check_reflection: Compare the reflected mobile coordinates too.
+            mobile_inds: Mobile atom indices, not residue indices.
+            target_inds: Target atom indices, not residue indices.
+            only_compute_backbone_rmsd: Restrict the score to backbone atoms.
+        """
+        aligned = self.dockq(target).aligned if compute_chain_assignment else self
+
+        aligner = Aligner(
+            aligned if mobile_inds is None else aligned[mobile_inds],
+            target if target_inds is None else target[target_inds],
+            only_compute_backbone_rmsd,
+        )
+        avg_rmsd = aligner.rmsd
+
+        if not also_check_reflection:
+            return avg_rmsd
+
+        aligner = Aligner(
+            aligned if mobile_inds is None else aligned[mobile_inds],
+            target if target_inds is None else target[target_inds],
+            only_compute_backbone_rmsd,
+            use_reflection=True,
+        )
+        avg_rmsd_neg = aligner.rmsd
+
+        return min(avg_rmsd, avg_rmsd_neg)
+
+    def lddt_ca(
+        self,
+        target: ProteinComplex,
+        mobile_inds: list[int] | np.ndarray | None = None,
+        target_inds: list[int] | np.ndarray | None = None,
+        compute_chain_assignment: bool = True,
+        **kwargs,
+    ) -> float | np.ndarray:
+        """Compute the LDDT between this protein complex and another.
+
+        Arguments:
+            target (ProteinComplex): The other protein complex to compare to.
+            mobile_inds: Mobile atom indices, not residue indices.
+            target_inds: Target atom indices, not residue indices.
+
+        Returns:
+            float | np.ndarray: The LDDT score between the two protein chains, either
+                a single float or per-residue LDDT scores if `per_residue` is True.
+        """
+        aligned = self.dockq(target).aligned if compute_chain_assignment else self
+        lddt = compute_lddt_ca(
+            torch.tensor(aligned.atom37_positions[mobile_inds]).unsqueeze(0),
+            torch.tensor(target.atom37_positions[target_inds]).unsqueeze(0),
+            torch.tensor(aligned.atom37_mask[mobile_inds]).unsqueeze(0),
+            **kwargs,
+        )
+        return float(lddt) if lddt.numel() == 1 else lddt.numpy().flatten()
+
+    def gdt_ts(
+        self,
+        target: ProteinComplex,
+        mobile_inds: list[int] | np.ndarray | None = None,
+        target_inds: list[int] | np.ndarray | None = None,
+        compute_chain_assignment: bool = True,
+        **kwargs,
+    ) -> float | np.ndarray:
+        """Compute the GDT_TS between this protein complex and another.
+
+        Arguments:
+            target (ProteinComplex): The other protein complex to compare to.
+            mobile_inds: Mobile atom indices, not residue indices.
+            target_inds: Target atom indices, not residue indices.
+
+        Returns:
+            float: The GDT_TS score between the two protein chains.
+        """
+        aligned = self.dockq(target).aligned if compute_chain_assignment else self
+        gdt_ts = compute_gdt_ts(
+            mobile=torch.tensor(
+                index_by_atom_name(aligned.atom37_positions[mobile_inds], "CA"),
+                dtype=torch.float32,
+            ).unsqueeze(0),
+            target=torch.tensor(
+                index_by_atom_name(target.atom37_positions[target_inds], "CA"),
+                dtype=torch.float32,
+            ).unsqueeze(0),
+            atom_exists_mask=torch.tensor(
+                index_by_atom_name(aligned.atom37_mask[mobile_inds], "CA", dim=-1)
+                & index_by_atom_name(target.atom37_mask[target_inds], "CA", dim=-1)
+            ).unsqueeze(0),
+            **kwargs,
+        )
+        return float(gdt_ts) if gdt_ts.numel() == 1 else gdt_ts.numpy().flatten()
+
+    def dockq(self, native: ProteinComplex):
+        # This function uses dockqv2 to compute the DockQ score. Because it does a mapping
+        # over all possible chains, it's quite slow. Be careful not to use this in an inference loop
+        # or something that requires fast scoring. It defaults to 8 CPUs.
+        #
+        # TODO(@zeming): Because we haven't properly implemented protein complexes for mmcif,
+        # if your protein has multi-letter or repeated chain IDs, this will fail. Please call
+        # Normalize chain IDs before DockQ when IDs repeat or use multiple letters.
+
+        try:
+            pass
+        except BaseException:
+            raise RuntimeError("DockQ is not installed. Please update your environment.") from None
+        self._sanity_check_complexes_are_comparable(native)
+
+        def sanity_check_chain_ids(pc: ProteinComplex):
+            ids = []
+            for i, chain in enumerate(pc.chain_iter()):
+                if i > len(SINGLE_LETTER_CHAIN_IDS):
+                    raise ValueError("Too many chains to write to PDB file")
+                if len(chain.chain_id) > 1:
+                    raise ValueError("We only supports single letter chain IDs for DockQ")
+                ids.append(chain.chain_id)
+            if len(set(ids)) != len(ids):
+                raise ValueError(f"Duplicate chain IDs in protein complex: {ids}")
+            return ids
+
+        sanity_check_chain_ids(self)
+        sanity_check_chain_ids(native)
+
+        with TemporaryDirectory() as tdir:
+            dir = Path(tdir)
+            self.to_pdb(dir / "self.pdb")
+            native.to_pdb(dir / "native.pdb")
+
+            output = check_output(["DockQ", dir / "self.pdb", dir / "native.pdb"])
+        lines = output.decode().split("\n")
+
+        # Remove the header comments
+        start_index = next(i for i, line in enumerate(lines) if line.startswith("Model"))
+        lines = lines[start_index:]
+
+        result = {}
+        interfaces = []
+        current_interface: dict = {}
+
+        for line in lines:
+            line = line.strip()
+            if not line:
+                continue
+
+            if line.startswith(("Model  :", "Native :")):
+                pass  # Tmp pdb file location, it's useless...
+            elif line.startswith("Total DockQ"):
+                total_dockq_match = re.search(
+                    r"Total DockQ over (\d+) native interfaces: ([\d.]+) with "
+                    r"(.*) model:native mapping",
+                    line,
+                )
+                if total_dockq_match:
+                    result["value"] = float(total_dockq_match.group(2))
+                    result["native interfaces"] = int(total_dockq_match.group(1))
+                    native_chains, self_chains = total_dockq_match.group(3).split(":")
+                    result["mapping"] = dict(zip(native_chains, self_chains, strict=False))
+                else:
+                    raise RuntimeError(
+                        "Failed to parse DockQ output, maybe your DockQ version is wrong?"
+                    )
+            elif line.startswith("Native chains:"):
+                if current_interface:
+                    interfaces.append(current_interface)
+                current_interface = {"Native chains": line.split(":")[1].strip().split(", ")}
+            elif line.startswith("Model chains:"):
+                current_interface["Model chains"] = line.split(":")[1].strip().split(", ")
+            elif ":" in line:
+                key, value = line.split(":", 1)
+                current_interface[key.strip()] = float(value.strip())
+
+        if current_interface:
+            interfaces.append(current_interface)
+
+        def parse_dict(d: dict[str, Any]) -> DockQSingleScore:
+            return DockQSingleScore(
+                native_chains=tuple(d["Native chains"]),  # type: ignore
+                DockQ=float(d["DockQ"]),
+                interface_rms=float(d["irms"]),
+                ligand_rms=float(d["Lrms"]),  # Note the capitalization difference
+                fnat=float(d["fnat"]),
+                fnonnat=float(d["fnonnat"]),
+                clashes=float(d["clashes"]),
+                F1=float(d["F1"]),
+                DockQ_F1=float(d["DockQ_F1"]),
+            )
+
+        inv_mapping = {v: k for k, v in result["mapping"].items()}
+
+        self_chain_map = {c.chain_id: c for c in self.chain_iter()}
+        realigned = []
+        for chain in native.chain_iter():
+            realigned.append(self_chain_map[inv_mapping[chain.chain_id]])
+
+        realigned = ProteinComplex.from_chains(realigned)
+        aligner = Aligner(realigned, native)
+        realigned = aligner.apply(realigned)
+
+        result = DockQResult(
+            total_dockq=result["value"],
+            native_interfaces=result["native interfaces"],
+            chain_mapping=result["mapping"],
+            interfaces={
+                (i["Model chains"][0], i["Model chains"][1]): parse_dict(i) for i in interfaces
+            },
+            aligned=realigned,
+            aligned_rmsd=aligner.rmsd,
+        )
+
+        return result
+
+    # Object invariants, slicing, and chain views
+    def __post_init__(self):
+        if not isinstance(self.sequence, str):
+            raise TypeError("sequence must be a string.")
+        sequence_length = len(self.sequence)
+        aligned = {
+            "atom37_positions": self.atom37_positions,
+            "atom37_mask": self.atom37_mask,
+            "residue_index": self.residue_index,
+            "insertion_code": self.insertion_code,
+            "confidence": self.confidence,
+            "entity_id": self.entity_id,
+            "chain_id": self.chain_id,
+            "sym_id": self.sym_id,
+        }
+        for name, values in aligned.items():
+            if not isinstance(values, np.ndarray):
+                raise TypeError(f"{name} must be a NumPy array, got {type(values).__name__}.")
+            if values.ndim == 0 or values.shape[0] != sequence_length:
+                raise ValueError(
+                    f"{name} shape {values.shape} does not align with "
+                    f"sequence length {sequence_length}."
+                )
+        if self.atom37_positions.shape != (sequence_length, 37, 3):
+            raise ValueError(
+                "atom37_positions must have shape "
+                f"({sequence_length}, 37, 3), got {self.atom37_positions.shape}."
+            )
+        if self.atom37_mask.shape != (sequence_length, 37):
+            raise ValueError(
+                "atom37_mask must have shape "
+                f"({sequence_length}, 37), got {self.atom37_mask.shape}."
+            )
+        if self.atom37_mask.dtype != bool:
+            raise TypeError(f"atom37_mask must have Boolean dtype, got {self.atom37_mask.dtype}.")
+        if not np.issubdtype(self.atom37_positions.dtype, np.number):
+            raise TypeError("atom37_positions must use a numeric dtype.")
+        for name, values in (
+            ("residue_index", self.residue_index),
+            ("insertion_code", self.insertion_code),
+            ("confidence", self.confidence),
+            ("entity_id", self.entity_id),
+            ("chain_id", self.chain_id),
+            ("sym_id", self.sym_id),
+        ):
+            if values.shape != (sequence_length,):
+                raise ValueError(
+                    f"{name} must have shape ({sequence_length},), got {values.shape}."
+                )
+        if not np.issubdtype(self.confidence.dtype, np.number):
+            raise TypeError("confidence must use a numeric dtype.")
+        atom37_confidence = self.atom37_confidence
+        if atom37_confidence is not None and not isinstance(atom37_confidence, np.ndarray):
+            raise TypeError("atom37_confidence must be a NumPy array when provided.")
+        if (
+            isinstance(atom37_confidence, np.ndarray)
+            and atom37_confidence.shape != self.atom37_mask.shape
+        ):
+            raise ValueError(
+                "atom37_confidence shape must match atom37_mask: "
+                f"{atom37_confidence.shape} != {self.atom37_mask.shape}."
+            )
+
+    def __getitem__(self, idx: int | list[int] | slice | np.ndarray):
+        """This function slices protein complexes without consideration of chain breaks
+        NOTE: When slicing with a boolean mask, it's possible that the output array won't
+        be the expected length. This is because we do our best to preserve chainbreak tokens.
+        """
+
+        if isinstance(idx, int):
+            idx = [idx]
+        if isinstance(idx, list):
+            raise ValueError("ProteinComplex doesn't supports indexing with lists of indices")
+
+        if isinstance(idx, np.ndarray):
+            is_chainbreak = np.asarray([s == "|" for s in self.sequence])
+            idx = idx.astype(bool) | is_chainbreak
+
+        complex = self._unsafe_slice(idx)
+        if len(complex) == 0:
+            return complex
+
+        # detect runs of chainbreaks by searching for instances of '||' in complex.sequence
+        chainbreak_runs = np.asarray(
+            [complex.sequence[i : i + 2] == "||" for i in range(len(complex.sequence) - 1)]
+            + [complex.sequence[-1] == "|"]
+        )
+        # We should remove as many chainbreaks as possible from the start of the sequence
+        for i in range(len(chainbreak_runs)):
+            if complex.sequence[i] == "|":
+                chainbreak_runs[i] = True
+            else:
+                break
+        complex = complex._unsafe_slice(~chainbreak_runs)
+        return complex
+
+    def _unsafe_slice(self, idx: int | list[int] | slice | np.ndarray):
+        sequence = slice_python_object_as_numpy(self.sequence, idx)
+        return replace(
+            self,
+            sequence=sequence,
+            entity_id=self.entity_id[..., idx],
+            chain_id=self.chain_id[..., idx],
+            sym_id=self.sym_id[..., idx],
+            residue_index=self.residue_index[..., idx],
+            insertion_code=self.insertion_code[..., idx],
+            atom37_positions=self.atom37_positions[..., idx, :, :],
+            atom37_mask=self.atom37_mask[..., idx, :],
+            confidence=self.confidence[..., idx],
+            atom37_confidence=self.atom37_confidence[..., idx, :]
+            if self.atom37_confidence is not None
+            else None,
+        )
+
+    def __len__(self):
+        return len(self.sequence)
+
+    @property
+    def num_chains(self):
+        return len(self.chain_boundaries)
+
+    @cached_property
+    def atoms(self) -> AtomIndexer:
+        return AtomIndexer(self, property="atom37_positions", dim=-2)
+
+    @cached_property
+    def atom_mask(self) -> AtomIndexer:
+        return AtomIndexer(self, property="atom37_mask", dim=-1)
+
+    @cached_property
+    def chain_lengths(self) -> np.ndarray:
+        return np.diff(self.chain_boundaries, axis=1).flatten()
+
+    @cached_property
+    def chain_boundaries(self) -> list[tuple[int, int]]:
+        cb = [-1]
+        for i, s in enumerate(self.sequence):
+            if s == "|":
+                cb.append(i)
+        cb.append(len(self))
+        return [(cb[i] + 1, cb[i + 1]) for i in range(len(cb) - 1)]
+
+    def get_chain_by_index(self, index: int) -> ProteinChain:
+        try:
+            start, end = self.chain_boundaries[index]
+            return self[start:end].as_chain()
+        except IndexError:
+            raise IndexError(f"Chain index {index} out of bounds") from None
+
+    def get_chain_by_id(
+        self, chain_id: str, sample_chain_if_duplicate: bool = True
+    ) -> ProteinChain:
+        valid_indices = [
+            index
+            for index, id_of_index in self.metadata.chain_lookup.items()
+            if id_of_index == chain_id
+        ]
+        if not valid_indices:
+            raise KeyError(f"Chain ID {chain_id} not found")
+        if sample_chain_if_duplicate:
+            index_to_return = random.choice(valid_indices)
+            return self.get_chain_by_index(index_to_return)
+        else:
+            if len(valid_indices) > 1:
+                raise ValueError(f"Multiple chains with chain ID {chain_id} found")
+            return self.get_chain_by_index(valid_indices[0])
+
+    def chain_iter(self) -> Iterable[ProteinChain]:
+        for start, end in self.chain_boundaries:
+            c = self[start:end]
+            yield c.as_chain()
+
+    def as_chain(self, force_conversion: bool = False) -> ProteinChain:
+        """Convert the ProteinComplex to a ProteinChain.
+
+        Args:
+            force_conversion: Flatten multiple chains to access chain-only utilities.
+
+        """
+        if not force_conversion:
+            if len(np.unique(self.chain_id)) != 1:
+                raise ValueError(
+                    f"Protein complex {self.id!r} has multiple chains; "
+                    "pass force_conversion=True to flatten it."
+                )
+            if len(np.unique(self.entity_id)) != 1:
+                raise ValueError(
+                    f"Protein complex {self.id!r} has multiple entities; "
+                    "pass force_conversion=True to flatten it."
+                )
+            if self.chain_id[0] not in self.metadata.chain_lookup:
+                warnings.warn(
+                    "Chain ID not found in metadata, using 'A' as default",
+                    stacklevel=2,
+                )
+            if self.entity_id[0] not in self.metadata.entity_lookup:
+                warnings.warn(
+                    "Entity ID not found in metadata, using None as default",
+                    stacklevel=2,
+                )
+            chain_id = self.metadata.chain_lookup.get(self.chain_id[0], "A")
+            entity_id = self.metadata.entity_lookup.get(self.entity_id[0], None)
+        else:
+            chain_id = "A"
+            entity_id = None
+
+        return ProteinChain(
+            id=self.id,
+            sequence=self.sequence,
+            chain_id=chain_id,
+            entity_id=entity_id,
+            atom37_positions=self.atom37_positions,
+            atom37_mask=self.atom37_mask,
+            residue_index=self.residue_index,
+            insertion_code=self.insertion_code,
+            confidence=self.confidence,
+            mmcif=self.metadata.mmcif,
+            atom37_confidence=self.atom37_confidence,
+        )
+
+    # Contact topology and mmCIF export
+    @cached_property
+    def per_chain_kd_trees(self):
+        # Iterate over chains, build KDTree for each chain
+        kdtrees = []
+
+        CA = self.atoms["CA"]
+
+        for start, end in self.chain_boundaries:
+            chain_CA = CA[start:end]
+            chain_CA = chain_CA[np.isfinite(chain_CA).all(axis=-1)]
+            kdtrees.append(KDTree(chain_CA))
+
+        return kdtrees
+
+    def chain_adjacency(self, cutoff: float = 8.0) -> np.ndarray:
+        # Compute adjacency matrix for protein complex
+        num_chains = self.num_chains
+        adjacency = np.zeros((num_chains, num_chains), dtype=bool)
+        for (i, kdtree), (j, kdtree2) in itertools.combinations(
+            enumerate(self.per_chain_kd_trees), 2
+        ):
+            adj = kdtree.query_ball_tree(kdtree2, cutoff)
+            any_is_adjacent = any(len(a) > 0 for a in adj)
+            adjacency[i, j] = any_is_adjacent
+            adjacency[j, i] = any_is_adjacent
+        return adjacency
+
+    def chain_adjacency_by_index(self, index: int, cutoff: float = 8.0) -> np.ndarray:
+        num_chains = len(self.chain_boundaries)
+        adjacency = np.zeros(num_chains, dtype=bool)
+        for i, kdtree in enumerate(self.per_chain_kd_trees):
+            if i == index:
+                continue
+            adj = kdtree.query_ball_tree(self.per_chain_kd_trees[index], cutoff)
+            adjacency[i] = any(len(a) > 0 for a in adj)
+        return adjacency
+
+    def add_prefix_to_chain_ids(self, prefix: str) -> ProteinComplex:
+        """Rename all chains in the complex with a given prefix.
+
+        Args:
+            prefix (str): The prefix to use for the new chain IDs. Each chain will be
+                named as "{prefix}_{chain_id}".
+
+        Returns:
+            ProteinComplex: A new protein complex with renamed chains.
+        """
+        new_chains = []
+        for chain in self.chain_iter():
+            # Create new chain with updated chain_id
+            new_chain = replace(chain, chain_id=f"{prefix}_{chain.chain_id}")
+            new_chains.append(new_chain)
+        return ProteinComplex.from_chains(new_chains)
+
+    def sasa(self, by_residue: bool = True):
+        chain = self.as_chain(force_conversion=True)
+        return chain.sasa(by_residue=by_residue)
+
+    def to_mmcif_string(self) -> str:
+        """Convert the ProteinComplex to mmCIF format.
+
+        Returns:
+            str: The mmCIF content as a string.
+        """
+        # Convert the ProteinComplex to a biotite AtomArray
+        # Collect all atoms from all chains
+        all_atoms = []
+        for chain in self.chain_iter():
+            chain_atom_array = chain.atom_array
+            # Convert AtomArray to list of atoms and add to collection
+            all_atoms.extend(chain_atom_array)
+
+        # Create combined AtomArray from all atoms
+        if not all_atoms:
+            raise ValueError("No atoms found in protein complex")
+
+        atom_array = bs.array(all_atoms)
+
+        # Create CIF file
+        f = CIFFile()
+        set_structure_pdbx(f, atom_array, data_block=self.id)
+
+        # Add entity information for proper mmCIF structure
+        self._add_entity_information(f)
+        round_mmcif_columns(f)
+
+        # Write to string
+        output = io.StringIO()
+        f.write(output)
+        return output.getvalue()
+
+    def _add_entity_information(self, cif_file: CIFFile) -> None:
+        """Add entity, entity_poly, and struct_asym sections to CIF file."""
+
+        # Group chains by sequence to create unique entities
+        entity_map = {}  # sequence -> entity_id
+        chain_to_entity = {}  # chain_id -> entity_id
+        entity_sequences = {}  # entity_id -> sequence
+        entity_id_counter = 1
+
+        for chain in self.chain_iter():
+            sequence = chain.sequence
+            if sequence not in entity_map:
+                entity_map[sequence] = entity_id_counter
+                entity_sequences[entity_id_counter] = sequence
+                entity_id_counter += 1
+            chain_to_entity[chain.chain_id] = entity_map[sequence]
+
+        # Create _entity section
+        entity_ids = []
+        entity_types = []
+        entity_descriptions = []
+
+        for entity_id in sorted(entity_sequences.keys()):
+            entity_ids.append(str(entity_id))
+            entity_types.append("polymer")
+            entity_descriptions.append(f"Protein chain (entity {entity_id})")
+
+        cif_file.block["entity"] = CIFCategory(
+            name="entity",
+            columns={
+                "id": CIFColumn(data=CIFData(array=np.array(entity_ids), dtype=np.str_)),
+                "type": CIFColumn(data=CIFData(array=np.array(entity_types), dtype=np.str_)),
+                "pdbx_description": CIFColumn(
+                    data=CIFData(array=np.array(entity_descriptions), dtype=np.str_)
+                ),
+            },
+        )
+
+        # Create _entity_poly section
+        poly_entity_ids = []
+        poly_types = []
+        poly_nstd_linkages = []
+        poly_sequences = []
+
+        for entity_id in sorted(entity_sequences.keys()):
+            poly_entity_ids.append(str(entity_id))
+            poly_types.append("polypeptide(L)")
+            poly_nstd_linkages.append("no")
+            poly_sequences.append(entity_sequences[entity_id])
+
+        cif_file.block["entity_poly"] = CIFCategory(
+            name="entity_poly",
+            columns={
+                "entity_id": CIFColumn(
+                    data=CIFData(array=np.array(poly_entity_ids), dtype=np.str_)
+                ),
+                "type": CIFColumn(data=CIFData(array=np.array(poly_types), dtype=np.str_)),
+                "nstd_linkage": CIFColumn(
+                    data=CIFData(array=np.array(poly_nstd_linkages), dtype=np.str_)
+                ),
+                "pdbx_seq_one_letter_code": CIFColumn(
+                    data=CIFData(array=np.array(poly_sequences), dtype=np.str_)
+                ),
+            },
+        )
+
+        # Create _struct_asym section
+        asym_ids = []
+        asym_entity_ids = []
+        asym_details = []
+
+        for chain in self.chain_iter():
+            asym_ids.append(chain.chain_id)
+            asym_entity_ids.append(str(chain_to_entity[chain.chain_id]))
+            asym_details.append("")
+
+        cif_file.block["struct_asym"] = CIFCategory(
+            name="struct_asym",
+            columns={
+                "id": CIFColumn(data=CIFData(array=np.array(asym_ids), dtype=np.str_)),
+                "entity_id": CIFColumn(
+                    data=CIFData(array=np.array(asym_entity_ids), dtype=np.str_)
+                ),
+                "details": CIFColumn(data=CIFData(array=np.array(asym_details), dtype=np.str_)),
+            },
+        )
+
+    # Construction, PDB interchange, and compact storage
+    @classmethod
+    def from_pdb(
+        cls, path: PathOrBuffer, id: str | None = None, is_predicted: bool = False
+    ) -> ProteinComplex:
+        atom_array = PDBFile.read(path).get_structure(model=1, extra_fields=["b_factor"])
+
+        chains = []
+        for chain in bs.chain_iter(atom_array):
+            chain = chain[~chain.hetero]
+            if len(chain) == 0:
+                continue
+            chains.append(ProteinChain.from_atomarray(chain, id, is_predicted))
+        return ProteinComplex.from_chains(chains)
+
+    def to_pdb(self, path: PathOrBuffer, include_insertions: bool = True):
+        atom_array = None
+        for chain in self.chain_iter():
+            carr = chain.atom_array if include_insertions else chain.atom_array_no_insertions
+            atom_array = carr if atom_array is None else atom_array + carr
+        f = PDBFile()
+        f.set_structure(atom_array)
+        f.write(path)
+
+    def to_pdb_string(self, include_insertions: bool = True) -> str:
+        buf = io.StringIO()
+        self.to_pdb(buf, include_insertions=include_insertions)
+        buf.seek(0)
+        return buf.read()
+
+    def normalize_chain_ids_for_pdb(self):
+        # Since PDB files have 1-letter chain IDs and don't support the idea of a symmetric index,
+        # we can normalize it instead which might be necessary for DockQ and to_pdb.
+        ids = SINGLE_LETTER_CHAIN_IDS
+        chains = []
+        for i, chain in enumerate(self.chain_iter()):
+            chain = replace(chain, chain_id=ids[i])
+            if i > len(ids):
+                raise RuntimeError("Too many chains to write to PDB file")
+            chains.append(chain)
+
+        return ProteinComplex.from_chains(chains)
+
+    def find_assembly_ids_with_chain(self, id: str) -> list[str]:
+        good_chains = []
+        if (comp := self.metadata.assembly_composition) is not None:
+            for assembly_id, chain_ids in comp.items():
+                if id in chain_ids:
+                    good_chains.append(assembly_id)
+        else:
+            raise ValueError(
+                "Cannot switch assemblies on this ProteinComplex; construct it from "
+                "mmCIF to retain assembly metadata"
+            )
+        return good_chains
+
+    def switch_assembly(self, id: str):
+        if self.metadata.mmcif is None:
+            raise ValueError(
+                "Cannot switch assemblies without retained mmCIF source metadata."
+            )
+        return get_assembly_fast(self.metadata.mmcif, assembly_id=id)
+
+    def state_dict(self, backbone_only=False, json_serializable=False):
+        """This state dict is optimized for storage, so it turns things to fp16 whenever
+        possible. Note that we also only support int32 residue indices, I'm hoping we don't
+        need more than 2**32 residues..."""
+        dct = {k: v for k, v in vars(self).items()}
+        if backbone_only:
+            # Frozen dataclasses do not make their NumPy members immutable.  Work on a
+            # private mask so requesting a compact backbone payload cannot clear the
+            # caller's side-chain atoms in-place.
+            atom37_mask = dct["atom37_mask"].copy()
+            atom37_mask[:, 3:] = False
+            dct["atom37_mask"] = atom37_mask
+        dct["atom37_positions"] = dct["atom37_positions"][dct["atom37_mask"]]
+        if dct.get("atom37_confidence") is not None:
+            dct["atom37_confidence"] = dct["atom37_confidence"][dct["atom37_mask"]]
+        else:
+            dct.pop("atom37_confidence", None)
+        for k, v in dct.items():
+            if isinstance(v, np.ndarray):
+                match v.dtype:
+                    case np.int64:
+                        dct[k] = v.astype(np.int32)
+                    case np.float64 | np.float32:
+                        dct[k] = v.astype(np.float16)
+                    case _:
+                        pass
+                if json_serializable:
+                    dct[k] = v.tolist()
+            elif isinstance(v, ProteinComplexMetadata):
+                dct[k] = asdict(v)
+        dct["metadata"]["mmcif"] = None
+        # These can be populated with non-serializable objects and are not needed for reconstruction
+        dct.pop("atoms", None)
+        dct.pop("atom_mask", None)
+        dct.pop("per_chain_kd_trees", None)
+        return dct
+
+    def to_blob(self, backbone_only=False) -> bytes:
+        payload = msgpack.dumps(self.state_dict(backbone_only), default=msgpack_numpy.encode)
+        return brotli.compress(payload, quality=5)
+
+    @classmethod
+    def from_state_dict(cls, dct):
+        # Note: assembly_composition is *supposed* to have string keys.
+        dct = _str_key_to_int_key(dct, ignore_keys=["assembly_composition"])
+
+        for k, v in dct.items():
+            if isinstance(v, list):
+                dct[k] = np.array(v)
+
+        atom37 = np.full((*dct["atom37_mask"].shape, 3), np.nan)
+        atom37[dct["atom37_mask"]] = dct["atom37_positions"]
+        dct["atom37_positions"] = atom37
+        if "atom37_confidence" in dct:
+            atom37_conf = np.full(dct["atom37_mask"].shape, np.nan, dtype=np.float32)
+            atom37_conf[dct["atom37_mask"]] = dct["atom37_confidence"]
+            dct["atom37_confidence"] = atom37_conf
+        dct = {
+            k: (
+                v.astype(np.float32)
+                if k in ["atom37_positions", "confidence", "atom37_confidence"]
+                else v
+            )
+            for k, v in dct.items()
+        }
+        if "chain_boundaries" in dct:
+            del dct["chain_boundaries"]
+        if "chain_boundaries" in dct["metadata"]:
+            del dct["metadata"]["chain_boundaries"]
+        dct["metadata"] = ProteinComplexMetadata(**dct["metadata"])
+        return cls(**dct)
+
+    @classmethod
+    def from_blob(cls, input: Path | str | io.BytesIO | bytes):
+        """NOTE(@zlin): blob + sparse coding + brotli + fp16 reduces memory
+        of chains from 52G/1M chains to 20G/1M chains, I think this is a good first
+        shot at compressing and dumping chains to disk. I'm sure there's better ways."""
+        match input:
+            case Path() | str():
+                bytes = Path(input).read_bytes()
+            case io.BytesIO():
+                bytes = input.getvalue()
+            case _:
+                bytes = input
+        state = msgpack.loads(
+            brotli.decompress(bytes),
+            object_hook=msgpack_numpy.decode,
+            strict_map_key=False,
+        )
+        return cls.from_state_dict(state)
+
+    @classmethod
+    def from_rcsb(cls, pdb_id: str, keep_source: bool = False) -> ProteinComplex:
+        f: io.StringIO = rcsb.fetch(pdb_id, "cif")  # type: ignore
+        return cls.from_mmcif(f, id=pdb_id, keep_source=keep_source, is_predicted=False)
+
+    @classmethod
+    def from_mmcif(
+        cls,
+        path: PathOrBuffer,
+        id: str | None = None,
+        assembly_id: str | None = None,
+        is_predicted: bool = False,
+        keep_source: bool = False,
+    ):
+        """Return a ProteinComplex object from an mmcif file.
+        TODO(@zeming): there's actually multiple complexes per file, but for ease of implementation,
+        we only consider the first defined complex!
+
+        Args:
+            path: Uncompressed mmCIF path or text buffer.
+            id: Optional structure identifier.
+            is_predicted (bool): If True, reads b factor as the confidence readout. Default: False.
+            chain_id (str, optional): Select a chain corresponding to (author) chain id.
+        """
+        mmcif = MmcifWrapper.read(path, id)
+        return get_assembly_fast(mmcif, assembly_id=assembly_id)
+
+    @classmethod
+    def from_chains(
+        cls,
+        chains: Sequence[ProteinChain],
+        mmcif: MmcifWrapper | None = None,
+        all_assembly_metadata_dictionary: dict[str, list[str]] | None = None,
+    ):
+        if not chains:
+            raise ValueError("Cannot create a ProteinComplex from an empty list of chains")
+
+        # TODO(roshan): Make a proper protein complex class
+        def join_arrays(arrays: Sequence[np.ndarray], sep: np.ndarray):
+            full_array = []
+            for array in arrays:
+                full_array.append(array)
+                full_array.append(sep)
+            full_array = full_array[:-1]
+            return np.concatenate(full_array, 0)
+
+        sep_tokens = {
+            "residue_index": np.array([-1]),
+            "insertion_code": np.array([""]),
+            "atom37_positions": np.full([1, 37, 3], np.nan),
+            "atom37_mask": np.zeros([1, 37], dtype=bool),
+            "confidence": np.array([0]),
+        }
+
+        any_has_atom37_conf = any(c.atom37_confidence is not None for c in chains)
+        if any_has_atom37_conf:
+            sep_tokens["atom37_confidence"] = np.full([1, 37], np.nan, dtype=np.float32)
+
+        def _get_chain_attr(chain: ProteinChain, name: str) -> np.ndarray:
+            val = getattr(chain, name)
+            if val is None and name == "atom37_confidence":
+                return np.full([len(chain), 37], np.nan, dtype=np.float32)
+            return val
+
+        array_args: dict[str, np.ndarray] = {
+            name: join_arrays([_get_chain_attr(chain, name) for chain in chains], sep)
+            for name, sep in sep_tokens.items()
+        }
+
+        multimer_arrays = []
+        chain2num_max = -1
+        chain2num = {}
+        ent2num_max = -1
+        ent2num = {}
+        total_index = 0
+        for i, c in enumerate(chains):
+            num_res = c.residue_index.shape[0]
+            if c.chain_id not in chain2num:
+                chain2num[c.chain_id] = (chain2num_max := chain2num_max + 1)
+            chain_id_array = np.full([num_res], chain2num[c.chain_id], dtype=np.int64)
+
+            if c.entity_id is None:
+                entity_num = (ent2num_max := ent2num_max + 1)
+            else:
+                if c.entity_id not in ent2num:
+                    ent2num[c.entity_id] = (ent2num_max := ent2num_max + 1)
+                entity_num = ent2num[c.entity_id]
+            entity_id_array = np.full([num_res], entity_num, dtype=np.int64)
+
+            sym_id_array = np.full([num_res], i, dtype=np.int64)
+
+            multimer_arrays.append(
+                {
+                    "chain_id": chain_id_array,
+                    "entity_id": entity_id_array,
+                    "sym_id": sym_id_array,
+                }
+            )
+
+            total_index += num_res + 1
+
+        sep = np.array([-1])
+        update = {
+            name: join_arrays([dct[name] for dct in multimer_arrays], sep=sep)
+            for name in ["chain_id", "entity_id", "sym_id"]
+        }
+        array_args.update(update)
+
+        metadata = ProteinComplexMetadata(
+            mmcif=mmcif,
+            chain_lookup={v: k for k, v in chain2num.items()},
+            entity_lookup={v: k for k, v in ent2num.items()},
+            assembly_composition=all_assembly_metadata_dictionary,
+        )
+
+        return cls(
+            id=chains[0].id,
+            sequence=residue_constants.CHAIN_BREAK_TOKEN.join(chain.sequence for chain in chains),
+            metadata=metadata,
+            **array_args,
+        )
+
+
+# Biological-assembly expansion
+def get_assembly_fast(
+    mmcif: MmcifWrapper,
+    assembly_id=None,
+    model=None,
+    data_block=None,
+    altloc="first",
+    use_author_fields=True,
+):
+    pdbx_file = mmcif.raw
+    if pdbx_file is None:
+        raise InvalidFileError("No mmCIF data loaded")
+    assembly_gen_category = pdbx_file.block["pdbx_struct_assembly_gen"]
+    if assembly_gen_category is None:
+        raise InvalidFileError("File has no 'pdbx_struct_assembly_gen' category")
+
+    struct_oper_category = pdbx_file.block["pdbx_struct_oper_list"]
+    if struct_oper_category is None:
+        raise InvalidFileError("File has no 'pdbx_struct_oper_list' category")
+
+    if assembly_id is None:
+        assembly_id = assembly_gen_category["assembly_id"].data.array[0]
+    elif assembly_id not in assembly_gen_category["assembly_id"].data.array:
+        raise KeyError(f"File has no Assembly ID '{assembly_id}'")
+
+    ### Calculate all possible transformations
+    transformations = _get_transformations(struct_oper_category)
+
+    ### Get structure according to additional parameters
+    structure = get_structure(
+        pdbx_file, model, data_block, altloc, ["label_asym_id"], use_author_fields
+    )[0]  # type: ignore
+    # TODO(@zeming) This line will remove all non-protein structural elements,
+    # we should remove this when we want to parse these too.
+    structure: bs.AtomArray = structure[
+        bs.filter_amino_acids(structure) & ~structure.hetero  # type: ignore
+    ]
+    if len(structure) == 0:
+        raise NoProteinError
+    unique_asym_ids = np.unique(structure.label_asym_id)  # type: ignore
+    asym2chain = {}
+    asym2auth = {}
+    for asym_id in unique_asym_ids:
+        sub_structure: bs.AtomArray = structure[structure.label_asym_id == asym_id]  # type: ignore
+        chain_id: str = sub_structure[0].chain_id  # type: ignore
+        (
+            sequence,
+            atom_positions,
+            atom_mask,
+            residue_index,
+            insertion_code,
+            confidence,
+            entity_id,
+        ) = chain_to_ndarray(sub_structure, mmcif, chain_id, False)
+
+        asym2chain[asym_id] = ProteinChain(
+            id=mmcif.id or "unknown",
+            sequence=sequence,
+            chain_id=chain_id,
+            entity_id=entity_id,
+            atom37_positions=atom_positions,
+            atom37_mask=atom_mask,
+            residue_index=residue_index,
+            insertion_code=insertion_code,
+            confidence=confidence,
+            mmcif=None,
+        )
+        asym2auth[asym_id] = chain_id
+
+    ### Get transformations and apply them to the affected asym IDs
+    assembly = []
+    assembly_id_dict: dict[str, list[str]] = {}
+
+    # Process the target assembly ID
+    for aid, op_expr, asym_id_expr in zip(
+        assembly_gen_category["assembly_id"].data.array,
+        assembly_gen_category["oper_expression"].data.array,
+        assembly_gen_category["asym_id_list"].data.array,
+        strict=False,
+    ):
+        if aid == assembly_id:
+            # Parse operations and asym IDs for this specific entry
+            operations = _parse_operation_expression(op_expr)
+            asym_ids = asym_id_expr.split(",")
+
+            # Filter affected asym IDs to only protein chains, preserving order
+            sub_structures = [asym2chain[asym_id] for asym_id in asym_ids if asym_id in asym2chain]
+
+            # Apply transformations
+            sub_assembly = _apply_transformations_fast(sub_structures, transformations, operations)
+            assembly.extend(sub_assembly)
+
+            # Build assembly_id_dict for this entry
+            assembly_id_dict[aid] = assembly_id_dict.get(aid, []) + [
+                asym2auth[id_] for id_ in asym_ids if id_ in asym2auth
+            ]
+
+    if len(assembly) == 0:
+        raise NoProteinError
+    return ProteinComplex.from_chains(assembly, mmcif, assembly_id_dict)
+
+
+def protein_chain_to_protein_complex(chain: ProteinChain) -> ProteinComplex:
+    if "|" not in chain.sequence:
+        return ProteinComplex.from_chains([chain])
+    chain_breaks = np.array(list(chain.sequence)) == "|"
+    chain_break_inds = np.where(chain_breaks)[0]
+    chain_break_inds = np.concatenate([[0], chain_break_inds, [len(chain)]])
+    chain_break_inds = np.array(list(itertools.pairwise(chain_break_inds)))
+    complex_chains = []
+    for start, end in chain_break_inds:
+        if start != 0:
+            start += 1
+        complex_chains.append(chain[start:end])
+    complex_chains = [
+        ProteinChain.from_atom37(
+            chain.atom37_positions,
+            sequence=chain.sequence,
+            chain_id=SINGLE_LETTER_CHAIN_IDS[i],
+            entity_id=i,
+        )
+        for i, chain in enumerate(complex_chains)
+    ]
+    return ProteinComplex.from_chains(complex_chains)

fastplms/models/esmfold2/esmfold2_protein_structure.py

@@ -0,0 +1,251 @@
+"""Atom selection, rigid alignment, RMSD, and GDT-TS primitives."""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from typing import TypeVar
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+from torch.amp import autocast  # type: ignore
+
+from .esmfold2_affine3d import Affine3D
+from .esmfold2_misc import unbinpack
+from .esmfold2_normalize_coordinates import index_by_atom_name
+
+ArrayOrTensor = TypeVar("ArrayOrTensor", np.ndarray, Tensor)
+
+
+def _coordinate_operations(
+    coordinates: ArrayOrTensor,
+) -> tuple[Callable[[ArrayOrTensor], ArrayOrTensor], Callable[..., ArrayOrTensor]]:
+    if isinstance(coordinates, np.ndarray):
+
+        def normalize(X: ArrayOrTensor) -> ArrayOrTensor:
+            return X / np.linalg.norm(X, axis=-1, keepdims=True)
+
+        return normalize, np.cross
+    return F.normalize, torch.cross  # type: ignore[return-value]
+
+
+def infer_cbeta_from_atom37(
+    atom37: ArrayOrTensor,
+    bond_length: float = 1.522,
+    bond_angle: float = 1.927,
+    dihedral: float = -2.143,
+) -> ArrayOrTensor:
+    """Infer C-beta coordinates from backbone tensor ``X``.
+
+    The scalar keyword arguments encode the bond length, bond angle, and
+    dihedral in radians used by the checkpoint's training geometry.
+    """
+
+    n_position = index_by_atom_name(atom37, "N", dim=-2)
+    ca_position = index_by_atom_name(atom37, "CA", dim=-2)
+    c_position = index_by_atom_name(atom37, "C", dim=-2)
+    normalize, cross = _coordinate_operations(atom37)
+    with np.errstate(invalid="ignore"):
+        n_to_ca = n_position - ca_position
+        n_to_c = n_position - c_position
+    unit_n_to_ca = normalize(n_to_ca)
+    normal = normalize(cross(n_to_c, unit_n_to_ca))
+    basis = [unit_n_to_ca, cross(normal, unit_n_to_ca), normal]
+    coefficients = [
+        bond_length * np.cos(bond_angle),
+        bond_length * np.sin(bond_angle) * np.cos(dihedral),
+        -bond_length * np.sin(bond_angle) * np.sin(dihedral),
+    ]
+    offset = sum(
+        vector * coefficient for vector, coefficient in zip(basis, coefficients, strict=True)
+    )
+    return ca_position + offset
+
+
+def _unpack_alignment_inputs(
+    mobile: Tensor,
+    target: Tensor,
+    atom_mask: Tensor | None,
+    sequence_id: Tensor | None,
+) -> tuple[Tensor, Tensor, Tensor | None]:
+    if sequence_id is None:
+        return mobile, target, atom_mask
+    unpacked_mobile = unbinpack(mobile, sequence_id, pad_value=torch.nan)
+    unpacked_target = unbinpack(target, sequence_id, pad_value=torch.nan)
+    if atom_mask is None:
+        unpacked_mask = torch.isfinite(unpacked_target).all(dim=-1)
+    else:
+        unpacked_mask = unbinpack(atom_mask, sequence_id, pad_value=0)
+    return unpacked_mobile, unpacked_target, unpacked_mask
+
+
+def _flatten_atom_axes(
+    mobile: Tensor,
+    target: Tensor,
+    atom_mask: Tensor | None,
+) -> tuple[Tensor, Tensor, Tensor | None]:
+    b = mobile.shape[0]
+    flat_mobile = mobile.view(b, -1, 3) if mobile.dim() == 4 else mobile
+    flat_target = target.view(b, -1, 3) if target.dim() == 4 else target
+    flat_mask = atom_mask
+    if flat_mask is not None and flat_mask.dim() == 3:
+        flat_mask = flat_mask.view(b, -1)
+    return flat_mobile, flat_target, flat_mask
+
+
+def _masked_coordinates(
+    mobile: Tensor,
+    target: Tensor,
+    atom_mask: Tensor | None,
+) -> tuple[Tensor, Tensor, Tensor]:
+    if atom_mask is None:
+        atom_mask = torch.ones(
+            mobile.shape[:2],
+            dtype=torch.bool,
+            device=mobile.device,
+        )
+        return mobile, target, atom_mask
+    expanded_mask = atom_mask.unsqueeze(-1)
+    return (
+        mobile.masked_fill(~expanded_mask, 0),
+        target.masked_fill(~expanded_mask, 0),
+        atom_mask,
+    )
+
+
+@torch.no_grad()
+@autocast("cuda", enabled=False)
+def compute_alignment_tensors(
+    mobile: Tensor,
+    target: Tensor,
+    atom_exists_mask: Tensor | None = None,
+    sequence_id: Tensor | None = None,
+) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]:
+    """Center and align coordinate tensors ``X`` and ``Y``.
+
+    Inputs have shape (b, n, 3), or (b, l, n_atoms, 3). The returned rotation
+    tensor ``R`` has shape (b, 3, 3), and atom counts have shape (b, 1).
+    """
+
+    mobile, target, atom_exists_mask = _unpack_alignment_inputs(
+        mobile,
+        target,
+        atom_exists_mask,
+        sequence_id,
+    )
+    if mobile.shape != target.shape:
+        raise AssertionError("Batch structure shapes do not match!")
+    mobile, target, atom_exists_mask = _flatten_atom_axes(
+        mobile,
+        target,
+        atom_exists_mask,
+    )
+    mobile, target, atom_exists_mask = _masked_coordinates(
+        mobile,
+        target,
+        atom_exists_mask,
+    )
+
+    num_valid_atoms = atom_exists_mask.sum(dim=-1, keepdim=True)
+    centroid_mobile = mobile.sum(dim=-2, keepdim=True) / num_valid_atoms.unsqueeze(-1)
+    centroid_target = target.sum(dim=-2, keepdim=True) / num_valid_atoms.unsqueeze(-1)
+    centroid_mobile[num_valid_atoms == 0] = 0
+    centroid_target[num_valid_atoms == 0] = 0
+
+    expanded_mask = atom_exists_mask.unsqueeze(-1)
+    centered_mobile = (mobile - centroid_mobile).masked_fill(~expanded_mask, 0)
+    centered_target = (target - centroid_target).masked_fill(~expanded_mask, 0)
+    covariance = torch.matmul(centered_mobile.transpose(1, 2), centered_target)
+    left_vectors, _, right_vectors = torch.svd(covariance)
+    rotation = torch.matmul(left_vectors, right_vectors.transpose(1, 2))
+    return (
+        centered_mobile,
+        centroid_mobile,
+        centered_target,
+        centroid_target,
+        rotation,
+        num_valid_atoms,
+    )
+
+
+def _validate_reduction(reduction: str, allowed: tuple[str, ...]) -> None:
+    if reduction not in allowed:
+        raise ValueError("Unrecognized reduction: '{reduction}'")
+
+
+@torch.no_grad()
+@autocast("cuda", enabled=False)
+def compute_rmsd_no_alignment(
+    aligned: Tensor,
+    target: Tensor,
+    num_valid_atoms: Tensor,
+    reduction: str = "batch",
+) -> Tensor:
+    """Measure RMSD after alignment using a declared reduction."""
+
+    _validate_reduction(reduction, ("per_residue", "per_sample", "batch"))
+    difference = aligned - target
+    if reduction == "per_residue":
+        mean_squared_error = difference.square().view(difference.size(0), -1, 9).mean(-1)
+    else:
+        mean_squared_error = difference.square().sum(dim=(1, 2)) / num_valid_atoms.squeeze(-1)
+    rmsd = torch.sqrt(mean_squared_error)
+    if reduction in {"per_residue", "per_sample"}:
+        return rmsd
+    valid_samples = num_valid_atoms.squeeze(-1) > 0
+    return rmsd.masked_fill(~valid_samples, 0).sum() / (valid_samples.sum() + 1e-8)
+
+
+@torch.no_grad()
+@autocast("cuda", enabled=False)
+def compute_affine_and_rmsd(
+    mobile: Tensor,
+    target: Tensor,
+    atom_exists_mask: Tensor | None = None,
+    sequence_id: Tensor | None = None,
+) -> tuple[Affine3D, Tensor]:
+    """Fit ``X`` onto ``Y`` and return the rigid transform and batch RMSD."""
+
+    (
+        centered_mobile,
+        centroid_mobile,
+        centered_target,
+        centroid_target,
+        rotation,
+        num_valid_atoms,
+    ) = compute_alignment_tensors(mobile, target, atom_exists_mask, sequence_id)
+    translation = torch.matmul(-centroid_mobile, rotation) + centroid_target
+    affine = Affine3D.from_tensor_pair(
+        translation,
+        rotation.unsqueeze(dim=-3).transpose(-2, -1),
+    )
+    rotated_mobile = torch.matmul(centered_mobile, rotation)
+    rmsd = compute_rmsd_no_alignment(
+        rotated_mobile,
+        centered_target,
+        num_valid_atoms,
+        reduction="batch",
+    )
+    return affine, rmsd
+
+
+def compute_gdt_ts_no_alignment(
+    aligned: Tensor,
+    target: Tensor,
+    atom_exists_mask: Tensor,
+    reduction: str = "batch",
+) -> Tensor:
+    """Compute GDT-TS for already aligned coordinate tensors."""
+
+    _validate_reduction(reduction, ("per_sample", "batch"))
+    if atom_exists_mask is None:
+        atom_exists_mask = torch.isfinite(target).all(dim=-1)
+    deviation = torch.linalg.vector_norm(aligned - target, dim=-1)
+    counts = atom_exists_mask.sum(dim=-1)
+    score_1 = ((deviation < 1) * atom_exists_mask).sum(dim=-1) / counts
+    score_2 = ((deviation < 2) * atom_exists_mask).sum(dim=-1) / counts
+    score_4 = ((deviation < 4) * atom_exists_mask).sum(dim=-1) / counts
+    score_8 = ((deviation < 8) * atom_exists_mask).sum(dim=-1) / counts
+    score = (score_1 + score_2 + score_4 + score_8) * 0.25
+    return score.mean() if reduction == "batch" else score