Update FastPLMs runtime files

Atomic FastPLMs complete publication. Checkpoint weights, tokenizer assets, runtime sources, model card, legal texts, and scoped attestations are updated together. Deletes are restricted to obsolete current-registry-pinned paths.

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

@@ -0,0 +1,353 @@
+Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International
+
+Creative Commons Corporation ("Creative Commons") is not a law firm and
+does not provide legal services or legal advice. Distribution of
+Creative Commons public licenses does not create a lawyer-client or
+other relationship. Creative Commons makes its licenses and related
+information available on an "as-is" basis. Creative Commons gives no
+warranties regarding its licenses, any material licensed under their
+terms and conditions, or any related information. Creative Commons
+disclaims all liability for damages resulting from their use to the
+fullest extent possible.
+
+Using Creative Commons Public Licenses
+
+Creative Commons public licenses provide a standard set of terms and
+conditions that creators and other rights holders may use to share
+original works of authorship and other material subject to copyright and
+certain other rights specified in the public license below. The
+following considerations are for informational purposes only, are not
+exhaustive, and do not form part of our licenses.
+
+Considerations for licensors: Our public licenses are intended for use
+by those authorized to give the public permission to use material in
+ways otherwise restricted by copyright and certain other rights. Our
+licenses are irrevocable. Licensors should read and understand the terms
+and conditions of the license they choose before applying it. Licensors
+should also secure all rights necessary before applying our licenses so
+that the public can reuse the material as expected. Licensors should
+clearly mark any material not subject to the license. This includes
+other CC-licensed material, or material used under an exception or
+limitation to copyright. More considerations for licensors :
+wiki.creativecommons.org/Considerations\_for\_licensors
+
+Considerations for the public: By using one of our public licenses, a
+licensor grants the public permission to use the licensed material under
+specified terms and conditions. If the licensor's permission is not
+necessary for any reason–for example, because of any applicable
+exception or limitation to copyright–then that use is not regulated by
+the license. Our licenses grant only permissions under copyright and
+certain other rights that a licensor has authority to grant. Use of the
+licensed material may still be restricted for other reasons, including
+because others have copyright or other rights in the material. A
+licensor may make special requests, such as asking that all changes be
+marked or described. Although not required by our licenses, you are
+encouraged to respect those requests where reasonable. More
+considerations for the public :
+wiki.creativecommons.org/Considerations\_for\_licensees
+
+Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International
+Public License
+
+By exercising the Licensed Rights (defined below), You accept and agree
+to be bound by the terms and conditions of this Creative Commons
+Attribution-NonCommercial-ShareAlike 4.0 International Public License
+("Public License"). To the extent this Public License may be interpreted
+as a contract, You are granted the Licensed Rights in consideration of
+Your acceptance of these terms and conditions, and the Licensor grants
+You such rights in consideration of benefits the Licensor receives from
+making the Licensed Material available under these terms and conditions.
+
+Section 1 – Definitions.
+
+-   a. Adapted Material means material subject to Copyright and Similar
+    Rights that is derived from or based upon the Licensed Material and
+    in which the Licensed Material is translated, altered, arranged,
+    transformed, or otherwise modified in a manner requiring permission
+    under the Copyright and Similar Rights held by the Licensor. For
+    purposes of this Public License, where the Licensed Material is a
+    musical work, performance, or sound recording, Adapted Material is
+    always produced where the Licensed Material is synched in timed
+    relation with a moving image.
+-   b. Adapter's License means the license You apply to Your Copyright
+    and Similar Rights in Your contributions to Adapted Material in
+    accordance with the terms and conditions of this Public License.
+-   c. BY-NC-SA Compatible License means a license listed at
+    creativecommons.org/compatiblelicenses, approved by Creative Commons
+    as essentially the equivalent of this Public License.
+-   d. Copyright and Similar Rights means copyright and/or similar
+    rights closely related to copyright including, without limitation,
+    performance, broadcast, sound recording, and Sui Generis Database
+    Rights, without regard to how the rights are labeled or categorized.
+    For purposes of this Public License, the rights specified in Section
+    2(b)(1)-(2) are not Copyright and Similar Rights.
+-   e. Effective Technological Measures means those measures that, in
+    the absence of proper authority, may not be circumvented under laws
+    fulfilling obligations under Article 11 of the WIPO Copyright Treaty
+    adopted on December 20, 1996, and/or similar international
+    agreements.
+-   f. Exceptions and Limitations means fair use, fair dealing, and/or
+    any other exception or limitation to Copyright and Similar Rights
+    that applies to Your use of the Licensed Material.
+-   g. License Elements means the license attributes listed in the name
+    of a Creative Commons Public License. The License Elements of this
+    Public License are Attribution, NonCommercial, and ShareAlike.
+-   h. Licensed Material means the artistic or literary work, database,
+    or other material to which the Licensor applied this Public License.
+-   i. Licensed Rights means the rights granted to You subject to the
+    terms and conditions of this Public License, which are limited to
+    all Copyright and Similar Rights that apply to Your use of the
+    Licensed Material and that the Licensor has authority to license.
+-   j. Licensor means the individual(s) or entity(ies) granting rights
+    under this Public License.
+-   k. NonCommercial means not primarily intended for or directed
+    towards commercial advantage or monetary compensation. For purposes
+    of this Public License, the exchange of the Licensed Material for
+    other material subject to Copyright and Similar Rights by digital
+    file-sharing or similar means is NonCommercial provided there is no
+    payment of monetary compensation in connection with the exchange.
+-   l. Share means to provide material to the public by any means or
+    process that requires permission under the Licensed Rights, such as
+    reproduction, public display, public performance, distribution,
+    dissemination, communication, or importation, and to make material
+    available to the public including in ways that members of the public
+    may access the material from a place and at a time individually
+    chosen by them.
+-   m. Sui Generis Database Rights means rights other than copyright
+    resulting from Directive 96/9/EC of the European Parliament and of
+    the Council of 11 March 1996 on the legal protection of databases,
+    as amended and/or succeeded, as well as other essentially equivalent
+    rights anywhere in the world.
+-   n. You means the individual or entity exercising the Licensed Rights
+    under this Public License. Your has a corresponding meaning.
+
+Section 2 – Scope.
+
+-   a. License grant.
+    -   1. Subject to the terms and conditions of this Public License,
+        the Licensor hereby grants You a worldwide, royalty-free,
+        non-sublicensable, non-exclusive, irrevocable license to
+        exercise the Licensed Rights in the Licensed Material to:
+        -   A. reproduce and Share the Licensed Material, in whole or in
+            part, for NonCommercial purposes only; and
+        -   B. produce, reproduce, and Share Adapted Material for
+            NonCommercial purposes only.
+    -   2. Exceptions and Limitations. For the avoidance of doubt, where
+        Exceptions and Limitations apply to Your use, this Public
+        License does not apply, and You do not need to comply with its
+        terms and conditions.
+    -   3. Term. The term of this Public License is specified in Section
+        6(a).
+    -   4. Media and formats; technical modifications allowed. The
+        Licensor authorizes You to exercise the Licensed Rights in all
+        media and formats whether now known or hereafter created, and to
+        make technical modifications necessary to do so. The Licensor
+        waives and/or agrees not to assert any right or authority to
+        forbid You from making technical modifications necessary to
+        exercise the Licensed Rights, including technical modifications
+        necessary to circumvent Effective Technological Measures. For
+        purposes of this Public License, simply making modifications
+        authorized by this Section 2(a)(4) never produces Adapted
+        Material.
+    -   5. Downstream recipients.
+        -   A. Offer from the Licensor – Licensed Material. Every
+            recipient of the Licensed Material automatically receives an
+            offer from the Licensor to exercise the Licensed Rights
+            under the terms and conditions of this Public License.
+        -   B. Additional offer from the Licensor – Adapted Material.
+            Every recipient of Adapted Material from You automatically
+            receives an offer from the Licensor to exercise the Licensed
+            Rights in the Adapted Material under the conditions of the
+            Adapter's License You apply.
+        -   C. No downstream restrictions. You may not offer or impose
+            any additional or different terms or conditions on, or apply
+            any Effective Technological Measures to, the Licensed
+            Material if doing so restricts exercise of the Licensed
+            Rights by any recipient of the Licensed Material.
+    -   6. No endorsement. Nothing in this Public License constitutes or
+        may be construed as permission to assert or imply that You are,
+        or that Your use of the Licensed Material is, connected with, or
+        sponsored, endorsed, or granted official status by, the Licensor
+        or others designated to receive attribution as provided in
+        Section 3(a)(1)(A)(i).
+-   b. Other rights.
+    -   1. Moral rights, such as the right of integrity, are not
+        licensed under this Public License, nor are publicity, privacy,
+        and/or other similar personality rights; however, to the extent
+        possible, the Licensor waives and/or agrees not to assert any
+        such rights held by the Licensor to the limited extent necessary
+        to allow You to exercise the Licensed Rights, but not otherwise.
+    -   2. Patent and trademark rights are not licensed under this
+        Public License.
+    -   3. To the extent possible, the Licensor waives any right to
+        collect royalties from You for the exercise of the Licensed
+        Rights, whether directly or through a collecting society under
+        any voluntary or waivable statutory or compulsory licensing
+        scheme. In all other cases the Licensor expressly reserves any
+        right to collect such royalties, including when the Licensed
+        Material is used other than for NonCommercial purposes.
+
+Section 3 – License Conditions.
+
+Your exercise of the Licensed Rights is expressly made subject to the
+following conditions.
+
+-   a. Attribution.
+    -   1. If You Share the Licensed Material (including in modified
+        form), You must:
+        -   A. retain the following if it is supplied by the Licensor
+            with the Licensed Material:
+            -   i. identification of the creator(s) of the Licensed
+                Material and any others designated to receive
+                attribution, in any reasonable manner requested by the
+                Licensor (including by pseudonym if designated);
+            -   ii. a copyright notice;
+            -   iii. a notice that refers to this Public License;
+            -   iv. a notice that refers to the disclaimer of
+                warranties;
+            -   v. a URI or hyperlink to the Licensed Material to the
+                extent reasonably practicable;
+
+        -   B. indicate if You modified the Licensed Material and retain
+            an indication of any previous modifications; and
+        -   C. indicate the Licensed Material is licensed under this
+            Public License, and include the text of, or the URI or
+            hyperlink to, this Public License.
+    -   2. You may satisfy the conditions in Section 3(a)(1) in any
+        reasonable manner based on the medium, means, and context in
+        which You Share the Licensed Material. For example, it may be
+        reasonable to satisfy the conditions by providing a URI or
+        hyperlink to a resource that includes the required information.
+    -   3. If requested by the Licensor, You must remove any of the
+        information required by Section 3(a)(1)(A) to the extent
+        reasonably practicable.
+-   b. ShareAlike.In addition to the conditions in Section 3(a), if You
+    Share Adapted Material You produce, the following conditions also
+    apply.
+    -   1. The Adapter's License You apply must be a Creative Commons
+        license with the same License Elements, this version or later,
+        or a BY-NC-SA Compatible License.
+    -   2. You must include the text of, or the URI or hyperlink to, the
+        Adapter's License You apply. You may satisfy this condition in
+        any reasonable manner based on the medium, means, and context in
+        which You Share Adapted Material.
+    -   3. You may not offer or impose any additional or different terms
+        or conditions on, or apply any Effective Technological Measures
+        to, Adapted Material that restrict exercise of the rights
+        granted under the Adapter's License You apply.
+
+Section 4 – Sui Generis Database Rights.
+
+Where the Licensed Rights include Sui Generis Database Rights that apply
+to Your use of the Licensed Material:
+
+-   a. for the avoidance of doubt, Section 2(a)(1) grants You the right
+    to extract, reuse, reproduce, and Share all or a substantial portion
+    of the contents of the database for NonCommercial purposes only;
+-   b. if You include all or a substantial portion of the database
+    contents in a database in which You have Sui Generis Database
+    Rights, then the database in which You have Sui Generis Database
+    Rights (but not its individual contents) is Adapted Material,
+    including for purposes of Section 3(b); and
+-   c. You must comply with the conditions in Section 3(a) if You Share
+    all or a substantial portion of the contents of the database.  
+    For the avoidance of doubt, this Section 4 supplements and does not
+    replace Your obligations under this Public License where the
+    Licensed Rights include other Copyright and Similar Rights.
+
+Section 5 – Disclaimer of Warranties and Limitation of Liability.
+
+-   a. Unless otherwise separately undertaken by the Licensor, to the
+    extent possible, the Licensor offers the Licensed Material as-is and
+    as-available, and makes no representations or warranties of any kind
+    concerning the Licensed Material, whether express, implied,
+    statutory, or other. This includes, without limitation, warranties
+    of title, merchantability, fitness for a particular purpose,
+    non-infringement, absence of latent or other defects, accuracy, or
+    the presence or absence of errors, whether or not known or
+    discoverable. Where disclaimers of warranties are not allowed in
+    full or in part, this disclaimer may not apply to You.
+-   b. To the extent possible, in no event will the Licensor be liable
+    to You on any legal theory (including, without limitation,
+    negligence) or otherwise for any direct, special, indirect,
+    incidental, consequential, punitive, exemplary, or other losses,
+    costs, expenses, or damages arising out of this Public License or
+    use of the Licensed Material, even if the Licensor has been advised
+    of the possibility of such losses, costs, expenses, or damages.
+    Where a limitation of liability is not allowed in full or in part,
+    this limitation may not apply to You.
+-   c. The disclaimer of warranties and limitation of liability provided
+    above shall be interpreted in a manner that, to the extent possible,
+    most closely approximates an absolute disclaimer and waiver of all
+    liability.
+
+Section 6 – Term and Termination.
+
+-   a. This Public License applies for the term of the Copyright and
+    Similar Rights licensed here. However, if You fail to comply with
+    this Public License, then Your rights under this Public License
+    terminate automatically.
+-   b. Where Your right to use the Licensed Material has terminated
+    under Section 6(a), it reinstates:
+
+    -   1. automatically as of the date the violation is cured, provided
+        it is cured within 30 days of Your discovery of the violation;
+        or
+    -   2. upon express reinstatement by the Licensor.
+
+    For the avoidance of doubt, this Section 6(b) does not affect any
+    right the Licensor may have to seek remedies for Your violations of
+    this Public License.
+
+-   c. For the avoidance of doubt, the Licensor may also offer the
+    Licensed Material under separate terms or conditions or stop
+    distributing the Licensed Material at any time; however, doing so
+    will not terminate this Public License.
+-   d. Sections 1, 5, 6, 7, and 8 survive termination of this Public
+    License.
+
+Section 7 – Other Terms and Conditions.
+
+-   a. The Licensor shall not be bound by any additional or different
+    terms or conditions communicated by You unless expressly agreed.
+-   b. Any arrangements, understandings, or agreements regarding the
+    Licensed Material not stated herein are separate from and
+    independent of the terms and conditions of this Public License.
+
+Section 8 – Interpretation.
+
+-   a. For the avoidance of doubt, this Public License does not, and
+    shall not be interpreted to, reduce, limit, restrict, or impose
+    conditions on any use of the Licensed Material that could lawfully
+    be made without permission under this Public License.
+-   b. To the extent possible, if any provision of this Public License
+    is deemed unenforceable, it shall be automatically reformed to the
+    minimum extent necessary to make it enforceable. If the provision
+    cannot be reformed, it shall be severed from this Public License
+    without affecting the enforceability of the remaining terms and
+    conditions.
+-   c. No term or condition of this Public License will be waived and no
+    failure to comply consented to unless expressly agreed to by the
+    Licensor.
+-   d. Nothing in this Public License constitutes or may be interpreted
+    as a limitation upon, or waiver of, any privileges and immunities
+    that apply to the Licensor or You, including from the legal
+    processes of any jurisdiction or authority.
+
+Creative Commons is not a party to its public licenses. Notwithstanding,
+Creative Commons may elect to apply one of its public licenses to
+material it publishes and in those instances will be considered the
+"Licensor." The text of the Creative Commons public licenses is
+dedicated to the public domain under the CC0 Public Domain Dedication.
+Except for the limited purpose of indicating that material is shared
+under a Creative Commons public license or as otherwise permitted by the
+Creative Commons policies published at creativecommons.org/policies,
+Creative Commons does not authorize the use of the trademark "Creative
+Commons" or any other trademark or logo of Creative Commons without its
+prior written consent including, without limitation, in connection with
+any unauthorized modifications to any of its public licenses or any
+other arrangements, understandings, or agreements concerning use of
+licensed material. For the avoidance of doubt, this paragraph does not
+form part of the public licenses.
+
+Creative Commons may be contacted at creativecommons.org.

README.md

@@ -1,153 +1,244 @@
----
-library_name: transformers
-tags:
-- protein language model
-- biology
----
-
-# FastANKH
-
-Fast, optimized implementations of ANKH protein language models (T5-based) with multi-backend attention support.
-
-**Requires PyTorch 2.11+** for Flash Attention 4 (FA4) backend support via flex attention.
-
-## Models
-
-| Model | Params | Layers | Hidden | Heads | Activation | Source |
-|-------|--------|--------|--------|-------|------------|--------|
-| ANKH_base | 453.3M | 48 | 768 | 12 | gelu_new | ElnaggarLab/ankh-base |
-| ANKH_large | 1.15B | 48 | 1536 | 16 | gelu_new | ElnaggarLab/ankh-large |
-| ANKH2_large | 1.15B | 24 | 1536 | 16 | silu | ElnaggarLab/ankh2-ext2 |
-| ANKH3_large | 1.15B | 48 | 1536 | 16 | silu | ElnaggarLab/ankh3-large |
-| ANKH3_xl | 3.49B | 48 | 2560 | 32 | silu | ElnaggarLab/ankh3-xl |
-
-## Usage
-
-```python
-from transformers import AutoModel, AutoTokenizer
-
-model = AutoModel.from_pretrained("Synthyra/ANKH_base", trust_remote_code=True)
-tokenizer = AutoTokenizer.from_pretrained("Synthyra/ANKH_base")
-
-# Set attention backend before loading for best performance
-# Options: "sdpa" (default, exact), "flex" (FA4 on Hopper/Blackwell)
-model.config.attn_backend = "flex"
-
-sequences = ["MKTLLILAVL", "ACDEFGHIKLMNPQRSTVWY"]
-inputs = tokenizer(sequences, return_tensors="pt", padding=True)
-outputs = model(**inputs)
-
-# Per-residue embeddings
-embeddings = outputs.last_hidden_state  # (batch, seq_len, hidden_dim)
-```
-
-## Experimental Test-Time Training
-
-TTT is disabled by default. Normal ANKH inference, embeddings, and
-`state_dict()` keys are unchanged unless you explicitly call `model.ttt(...)`.
-The current implementation is experimental and trains only local LoRA adapters
-with masked language modeling on the test protein. ANKH's FastPLMs MaskedLM
-head is encoder-only and not pretrained for standard MLM, so treat TTT results
-with extra caution.
-
-```python
-from transformers import AutoModelForMaskedLM
-
-mlm = AutoModelForMaskedLM.from_pretrained(
-    "Synthyra/ANKH_base",
-    trust_remote_code=True,
-).cuda().eval()
-
-metrics = mlm.ttt(
-    seq="MSTNPKPQRKTKRNT",
-    ttt_config={"steps": 3, "ags": 1, "batch_size": 1},
-)
-mlm.ttt_reset()
-print(metrics["losses"])
-```
-
-## Batch Embedding
-
-```python
-model = AutoModel.from_pretrained("Synthyra/ANKH_base", trust_remote_code=True).to("cuda")
-embeddings = model.embed_dataset(
-    sequences=["MKTLLILAVL", "ACDEFGHIKLMNPQRSTVWY"],
-    batch_size=8,
-    max_len=512,
-    full_embeddings=True,
-)
-```
-
-## Attention Backends
-
-| Backend | Key | Notes |
-|---------|-----|-------|
-| SDPA | `"sdpa"` | Default. Exact attention with position bias as additive mask. |
-| Flex | `"flex"` | Uses FA4 on Hopper/Blackwell GPUs (PyTorch 2.11+). Position bias computed via `score_mod`. Triton fallback on older hardware. |
-| Flash | `"kernels_flash"` | Not supported for ANKH (no arbitrary bias support). Falls back to flex/sdpa. |
-
-## Architecture
-
-ANKH models are T5 encoder-only architectures:
-- **No absolute position embeddings**: Uses T5-style relative position bias (log-bucketed, bidirectional)
-- **RMS LayerNorm**: No mean subtraction, no bias term
-- **Gated FFN**: `activation(wi_0(x)) * wi_1(x) -> wo(x)` with gelu_new (v1) or silu (v2/v3)
-- **Pre-layer normalization**: Norm before attention and FFN, residual after
-- **No bias in projections**: All q/k/v/o and FFN linear layers are bias=False
-
-The relative position bias is computed once (materialized as a full tensor) and shared across all encoder layers. For the flex backend, the bias is passed as a `score_mod` closure for optimal performance.
-
-## Notes
-
-- The `FastAnkhForMaskedLM` variant includes an LM head initialized from the shared embedding weights. The original ANKH models were trained with T5's span corruption objective using an encoder-decoder architecture. This encoder-only MaskedLM head is **not pre-trained for standard MLM** and requires additional fine-tuning.
-- `model.tokenizer` is loaded from the checkpoint hub id so each ANKH model uses its matching tokenizer.
-- ANKH3 models use a vocabulary of 256 tokens (vs 144 for v1/v2) and were trained with dual objectives ([NLU] for embeddings, [S2S] for generation).
-
-## Citations
-
-```bibtex
-@article{elnaggar2023ankh,
-  title={Ankh: Optimized Protein Language Model Unlocks General-Purpose Modelling},
-  author={Elnaggar, Ahmed and Essam, Hazem and Salah-Eldin, Wafaa and Moustafa, Walid and Elkerdawy, Mohamed and Rochereau, Charlotte and Rost, Burkhard},
-  journal={arXiv preprint arXiv:2301.06568},
-  year={2023}
-}
-```
-
-```bibtex
-@article{alsamkary2025ankh3,
-  title={Ankh3: Multi-Task Pretraining with Sequence Denoising and Completion Enhances Protein Representations},
-  author={Alsamkary, Hazem and Elshaffei, Mohamed and Elkerdawy, Mohamed and Elnaggar, Ahmed},
-  journal={arXiv preprint arXiv:2505.20052},
-  year={2025}
-}
-```
-
-```bibtex
-@misc{FastPLMs,
-  author={Hallee, Logan and Bichara, David and Gleghorn, Jason P.},
-  title={FastPLMs: Fast, efficient, protein language model inference from Huggingface AutoModel.},
-  year={2024},
-  url={https://huggingface.co/Synthyra/ESMplusplus_small},
-  DOI={10.57967/hf/3726},
-  publisher={Hugging Face}
-}
-```
-
-```bibtex
-@article{dong2024flexattention,
-  title={Flex Attention: A Programming Model for Generating Optimized Attention Kernels},
-  author={Dong, Juechu and Feng, Boyuan and Guessous, Driss and Liang, Yanbo and He, Horace},
-  journal={arXiv preprint arXiv:2412.05496},
-  year={2024}
-}
-```
-
-```bibtex
-@inproceedings{paszke2019pytorch,
-  title={PyTorch: An Imperative Style, High-Performance Deep Learning Library},
-  author={Paszke, Adam and Gross, Sam and Massa, Francisco and Lerer, Adam and Bradbury, James and Chanan, Gregory and Killeen, Trevor and Lin, Zeming and Gimelshein, Natalia and Antiga, Luca and Desmaison, Alban and K{\"o}pf, Andreas and Yang, Edward and DeVito, Zach and Raison, Martin and Tejani, Alykhan and Chilamkurthy, Sasank and Steiner, Benoit and Fang, Lu and Bai, Junjie and Chintala, Soumith},
-  booktitle={Advances in Neural Information Processing Systems 32},
-  year={2019}
-}
-```
+---
+library_name: transformers
+license: "cc-by-nc-sa-4.0"
+tags:
+  - protein-language-model
+  - fastplms
+---
+
+<!-- Generated from src/fastplms/models.toml. Do not edit. -->
+
+# Synthyra/ANKH2_large
+
+This checkpoint packages the FastPLMs `ANKH` implementation.
+
+Accepted inputs are amino-acid sequences tokenized for encoder or sequence-to-
+sequence use.
+Supported Transformers entry points are `AutoConfig`, `AutoModel`,
+`AutoModelForMaskedLM`, `AutoModelForSeq2SeqLM`,
+`AutoModelForSequenceClassification`, `AutoModelForTokenClassification`.
+
+## Install and platform requirements
+
+Install FastPLMs from the exact source revision paired with this model card:
+
+```bash
+python -m pip install \
+  "fastplms @ git+https://github.com/Synthyra/FastPLMs.git@1b9ce023f1e06571cf3e6324be0610ffa53e0a4a"
+```
+
+Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The declared CPU gate covers tiny offline contracts; published checkpoint throughput and parity require the documented device tier. The Hub quick start below requires network
+access on first download. For an air-gapped run, first build the manifest-pinned
+local artifact and use the offline form shown in the example.
+
+## Quick start
+
+```python
+from transformers import AutoModel
+
+model_id = "Synthyra/ANKH2_large"
+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/ANKH2_large` path, then pass `local_files_only=True`.
+
+Leave attention unspecified for the Transformers default. Supported explicit
+choices are `eager`, `sdpa`.
+Pass the selected name through `attn_implementation`.
+When an optimized backend cannot return full attention tensors,
+`output_attentions=True` emits one explicit runtime warning and uses a correctly
+masked eager implementation for that call only. The warning identifies the
+configured backend, effective backend, and reason. Configuration and later
+calls are unchanged.
+For BF16 execution, this family uses parameters loaded directly in BF16.
+
+## Tokenization and forward inference
+
+The live `Synthyra/ANKH2_large` revision `392de5ed52bbfd73b45f545e378aaebcff096d0e` is legacy
+encoder-only. The Hub quick start above is therefore an encoder-only
+`AutoModel` path, not evidence that decoder or language-model-head weights are
+already published.
+
+Use the tokenizer owned by the loaded model so tokenizer files, revision,
+offline/cache policy, and ANKH's residue-aware pre-tokenizer stay aligned.
+Pass raw protein strings without inserted residue spaces:
+
+```python
+import torch
+
+tokenizer = model.tokenizer
+batch = tokenizer(
+    ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"],
+    padding=True,
+    return_tensors="pt",
+)
+
+with torch.inference_mode():
+    output = model(**batch)
+
+print(output.last_hidden_state.shape)
+```
+
+## Dataset embeddings
+
+The current live Hub revision is legacy encoder-only. It supports encoder
+dataset embeddings, which default to the encoder's final hidden state. Layer
+indices use the selected stack's native hidden-state order:
+
+```python
+encoder_result = model.embed_dataset(
+    ["MSTNPKPQRKTKRNT"],
+    hidden_state_source="encoder",
+    hidden_state_index=-1,
+    full_embeddings=True,
+)
+all_encoder_layers = model.embed_dataset(
+    ["MSTNPKPQRKTKRNT"],
+    hidden_state_source="encoder",
+    store_all_hidden_states=True,
+    full_embeddings=True,
+)
+```
+
+Decoder representations require `AutoModelForSeq2SeqLM` and exactly one
+explicit, aligned `decoder_inputs` sequence or `decoder_input_ids` tensor. ANKH
+does not infer a shifted source sequence because official tasks use prompts,
+sentinel tokens, or generated tokens that depend on the task. Protein inputs
+remain raw residue strings and sentinel prompts remain tight, as in
+`M<extra_id_0>`. Until the atomic Hub replacement is published, load the
+validated complete local artifact and fail closed on its registry-bound
+attestation:
+
+```python
+from pathlib import Path
+from transformers import AutoModelForSeq2SeqLM
+from fastplms.registry import get_model_registry
+from tools.artifacts.build import validate_artifact
+
+artifact = Path("dist/hub/ANKH2_large").resolve()
+registry = get_model_registry()
+validate_artifact(artifact, spec=registry["ankh2_large"], registry=registry)
+seq2seq = AutoModelForSeq2SeqLM.from_pretrained(
+    artifact,
+    trust_remote_code=True,
+    local_files_only=True,
+).eval()
+decoder_result = seq2seq.embed_dataset(
+    ["MSTNPKPQRKTKRNT"],
+    hidden_state_source="decoder",
+    hidden_state_index=-1,
+    decoder_inputs=["M<extra_id_0>"],
+    full_embeddings=True,
+)
+```
+
+`decoder_attention_mask` is accepted only with `decoder_input_ids`. Decoder
+pooling excludes start, EOS, padding, sentinel, and other tokenizer-special
+positions. Persisted results record the selected stack and layer, decoder input
+and mask fingerprints, input-position alignment, and biological-mask policy.
+
+## Encoder and sequence-to-sequence use
+
+The current manifest Hub revision `392de5ed52bbfd73b45f545e378aaebcff096d0e` for
+`Synthyra/ANKH2_large` is legacy encoder-only. It supports the `AutoModel`
+encoder path, but it is not the full FastPLMs 1.0 sequence-to-sequence
+artifact. Do not load `AutoModelForSeq2SeqLM` from that live revision.
+
+The full encoder-decoder replacement is still pending atomic publication. Use
+sequence-to-sequence behavior only from a locally built artifact whose complete
+weight, runtime, provenance, and registry validation has passed. The following
+snippet fails closed if that artifact is missing or invalid:
+
+```python
+import torch
+from pathlib import Path
+from transformers import AutoModelForSeq2SeqLM
+from fastplms.registry import get_model_registry
+from tools.artifacts.build import validate_artifact
+
+artifact = Path("dist/hub/ANKH2_large").resolve()
+registry = get_model_registry()
+validate_artifact(artifact, spec=registry["ankh2_large"], registry=registry)
+seq2seq = AutoModelForSeq2SeqLM.from_pretrained(
+    artifact,
+    trust_remote_code=True,
+    local_files_only=True,
+).eval()
+tokenizer = seq2seq.tokenizer
+batch = tokenizer("MSTNPKPQRKTKRNT", return_tensors="pt")
+
+with torch.inference_mode():
+    generated_ids = seq2seq.generate(**batch, max_new_tokens=16)
+print(tokenizer.batch_decode(generated_ids, skip_special_tokens=True))
+```
+
+ANKH artifacts retain CC BY-NC-SA 4.0 terms. The notes below distinguish the
+official heads from FastPLMs extensions. Once validated and published, the 1.0
+replacement will increase the default repository size while preserving
+encoder-output parity. Runtime code,
+configuration, tokenizer, card, provenance, and every weight shard must be
+published atomically. Files-only publication is forbidden for this migration.
+
+## Notes and limitations
+
+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.
+
+## Runtime contract
+
+- Public input: Amino-acid sequences tokenized for encoder or sequence-to-sequence use
+- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSeq2SeqLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification`
+- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `FastPLMs extension`, `AutoModelForSeq2SeqLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head`
+- Attention implementations: `eager`, `sdpa`
+- Precision policies: `default`
+- BF16 execution: `static_parameters`
+- Generation contract: `required`
+- Optional dependency group: `core`
+- Weight publication allowed: `true`
+- Weight license status: `resolved`
+- Redistributable: `true`
+- Complete weight publication required: `true`
+
+## Provenance
+
+- FastPLMs weights: `Synthyra/ANKH2_large@392de5ed52bbfd73b45f545e378aaebcff096d0e`
+- Runtime revision: `1b9ce023f1e06571cf3e6324be0610ffa53e0a4a`
+- Runtime source-tree SHA-256: `5eb598fd991b5447e3bb6a18c92bd2300be1d94a1e169504725bc56127c1d6bc`
+- Runtime bundle SHA-256: `17b8f83a33e63d941e3edfb8db2d8381286046b572f94194e576ca394d997597`
+- Generator/schema version and complete/runtime-only attestations: recorded in `provenance.json`
+- Canonical transformed state SHA-256: `597c4fe2fa8711f11a25317905f1d62fa92905e55fdd5c0a79614cd9c9d2bca3`
+- Conversion equality attestation: recorded in `provenance.json`
+- Official checkpoint: `ElnaggarLab/ankh2-ext2@aa9b9fa72288c47d9f618ce80c011e24b54e17a8`
+- Artifact source: `official`
+- State transform: `ankh_t5_to_fastplms_v1`
+- BF16 execution: `static_parameters`
+- Pinned upstreams: `ankh`
+- Reference container: `reference-ankh`
+- Release tiers: `check`, `compliance`, `feature`, `artifact`, `benchmark`
+- Unresolved required file identities: `0`
+
+The local artifact records exact file identities, conversion provenance, source
+revisions, and legal texts in `provenance.json`. A nonzero unresolved count is a
+release blocker.
+
+## Validation boundary
+
+For tiers declared by the manifest, the release contract compares applicable
+semantic configuration, tokenizer behavior, state keys, shapes, dtypes,
+values, aliases, and representative inference with the pinned official
+implementation. This metadata does not by itself claim that a particular build
+passed, that one backend is faster, or that an output has biological or
+therapeutic validity.
+
+## License
+
+Checkpoint terms: CC-BY-NC-SA-4.0. The Hub model-card identifier is
+`cc-by-nc-sa-4.0`. Applicable source licenses, notices, attribution,
+and conversion records are distributed with the local artifact. Review them
+before use.

THIRD_PARTY_NOTICES.md

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

artifact-manifest.json

@@ -0,0 +1,35 @@
+{
+  "LICENSES/FastPLMs-Apache-2.0.txt": "sha256:2d2b50c7b1414bff1189a1db1f0cfb92e3e064b50f4c2b1019827b683e1b629a",
+  "LICENSES/ankh/LICENSE.md": "sha256:cd041d7f9f52936e8824ac3f754e9c67410763205fc8a7020ba74fc8b6edc088",
+  "README.md": "sha256:d34607e87e0df11a9d8cea1e8dd541d1fe79761f90cfc2d95c9cb5135e0ccbbc",
+  "THIRD_PARTY_NOTICES.md": "sha256:25704b3c76404696cae52e7fca13088d329f70f412687340351259e86cd62baa",
+  "config.json": "sha256:fe9df500b8e73df29eff467f41a92d9bf1489ba7ac8c151f237f700f41b1b7a1",
+  "fastplms/__init__.py": "sha256:4fb3196022ca8ec699d59d09bdbc5f0184195552b773698ab9b061fe3cd7df12",
+  "fastplms/attention/__init__.py": "sha256:f60b9fecfb4bcb37a4e7c26dc2f752b9035f9cbad627b4a84213f3a92ec88f7d",
+  "fastplms/attention/_core.py": "sha256:8f7ec5b65bd8b6c6fa4951d50d1c0e499abf03ae00914794b51fc410201e3e33",
+  "fastplms/attention/_kernel_lock.py": "sha256:85d8521a2af5f94fad3948af3814db0c866c414ee4d43df797b9bd6f980e947b",
+  "fastplms/attention/interfaces.py": "sha256:1c6f06a8e411e0f9bf6d230522205c93ae46ea58864006fbb892aa05e5ca5749",
+  "fastplms/embeddings/__init__.py": "sha256:47ff8cdf682d44037dd9edab133e2e60675d60786bd5cf0bffd1998f31985555",
+  "fastplms/embeddings/pooling.py": "sha256:a140266ed6b1cc344c8507edc5c6c4f2dce464c3db70ba4b16c7ac2ba2fad96e",
+  "fastplms/embeddings/runner.py": "sha256:23ee4727a918d6d331f7a0f89b823d149f1a791f0c5586e3496d7b6eb2ce97e0",
+  "fastplms/embeddings/storage.py": "sha256:3fbe2bab75092e5a4cadf4d27e4752181d597469a65a55db085ceef808ed418e",
+  "fastplms/embeddings/types.py": "sha256:119718a20989d1ae5a60fabc0f5e98bdc172c5163b04db3d4554ac3956b30e52",
+  "fastplms/models.toml": "sha256:05a8399f084a5babb5f0916cee7e564c4030767f3ff0230c4f46e539209847d1",
+  "fastplms/models/__init__.py": "sha256:5e48c2cb3877aa6f42f3b5411d53b16bba2e32827bbde634f47f174c5cb36f86",
+  "fastplms/models/ankh/__init__.py": "sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
+  "fastplms/models/ankh/modeling_ankh.py": "sha256:b114c73e4aaedc567e3351b425ddddacf35661f09b3414cf59b6db0b712c82c3",
+  "fastplms/models/ttt.py": "sha256:a0df4e98b02120d423e3c7ca9b866a8d0e3748b9076042a0102a838a11aed046",
+  "fastplms/registry.py": "sha256:afca271911b651a882345b74a58366494a1d784e4a48c8683a87f5508f4ba16e",
+  "fastplms/runtime.py": "sha256:110018646d6f248cedab140a030c3065e1b062b61f6aff659c231e538614bc01",
+  "fastplms_bundle.py": "sha256:ebf70cf52fb0a07025edf06bd50d2e945d9ed09fe12e6d3ae3503001ec21c9f9",
+  "generation_config.json": "sha256:53e38d7cba6a16e598cd75a7acae25573674dc0ecdbf290ed19b46e0b4dd8c75",
+  "model-00001-of-00002.safetensors": "sha256:7c0c297f60bcf81c732cdfeae6e99e140272807eb52afd70356fc6fdfa94e5a8",
+  "model-00002-of-00002.safetensors": "sha256:f3d425d3e8741ccbdd925446559a9bf317c2c91e328f2eee44924423b56e3a3d",
+  "model.safetensors.index.json": "sha256:eb9f1a3b03bd5fc5dd912aa4244181ed09e19b99a57c2e1d95bdf952e510cea3",
+  "modeling_fastplms.py": "sha256:df7f6d96802eff127bc8270cde5fe00c8f5e77921c1cc893021fb9f75a3e4c77",
+  "provenance.json": "sha256:47a519cbd68eefcd5baac2bd4e2331be145e861d3b549e0ac41b83500f2c9bbe",
+  "runtime-attestation.json": "sha256:92db6caed50673b1bf451b43847d01e356e6e87ba6fc17e7ddd315faded2cb97",
+  "special_tokens_map.json": "sha256:c8995d2f8037fe3a8cfdef30475365e1c314c417880b75abaf8296e3c05d42d6",
+  "tokenizer.json": "sha256:b4533f607d9fd665f2d9d94b0cf71870a6fc2fc2ae7cbd516d0e43a9efb406fd",
+  "tokenizer_config.json": "sha256:dba1f9315ded007fc2e5bcb2ed4bfb6fff81c9becdc97b73a5713ba5ddc01afc"
+}

config.json

@@ -1,26 +1,49 @@
 {
-  "attn_backend": "sdpa",
+  "_name_or_path": "agemagician/protx-large-1gspan-partreconst-15mlmp-encl48-decl24-ramd128-ranb64-dmodel1536-silu-seq768-tpuv4",
+  "architectures": [
+    "T5ForConditionalGeneration"
+  ],
   "auto_map": {
-    "AutoConfig": "modeling_ankh.FastAnkhConfig",
-    "AutoModel": "modeling_ankh.FastAnkhModel",
-    "AutoModelForMaskedLM": "modeling_ankh.FastAnkhForMaskedLM",
-    "AutoModelForSequenceClassification": "modeling_ankh.FastAnkhForSequenceClassification",
-    "AutoModelForTokenClassification": "modeling_ankh.FastAnkhForTokenClassification"
+    "AutoConfig": "modeling_fastplms.FastAnkhConfig",
+    "AutoModel": "modeling_fastplms.FastAnkhModel",
+    "AutoModelForMaskedLM": "modeling_fastplms.FastAnkhForMaskedLMExtension",
+    "AutoModelForSeq2SeqLM": "modeling_fastplms.FastAnkhForConditionalGeneration",
+    "AutoModelForSequenceClassification": "modeling_fastplms.FastAnkhForSequenceClassification",
+    "AutoModelForTokenClassification": "modeling_fastplms.FastAnkhForTokenClassification"
   },
   "d_ff": 3840,
   "d_kv": 64,
   "d_model": 1536,
+  "decoder_start_token_id": 0,
   "dense_act_fn": "silu",
+  "dropout_rate": 0.0,
   "eos_token_id": 1,
+  "fastplms_checkpoint_hash": "ca31eec9942e836d0b557a94d3b50e642af3738627a348d6526a6968b8646f8c",
+  "fastplms_checkpoint_repo_id": "ElnaggarLab/ankh2-ext2",
+  "fastplms_checkpoint_revision": "aa9b9fa72288c47d9f618ce80c011e24b54e17a8",
+  "fastplms_model_id": "ankh2_large",
+  "fastplms_release_tool_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "fastplms_release_tool_sha256": "1459b5d7d13d9b07bd97b3eee764f2ce73623e15e32d07ddf6825c2a9509afb9",
+  "fastplms_runtime_bundle_sha256": "17b8f83a33e63d941e3edfb8db2d8381286046b572f94194e576ca394d997597",
+  "fastplms_runtime_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "fastplms_source_tree_sha256": "5eb598fd991b5447e3bb6a18c92bd2300be1d94a1e169504725bc56127c1d6bc",
+  "fastplms_weights_revision": "aa9b9fa72288c47d9f618ce80c011e24b54e17a8",
+  "feed_forward_proj": "gated-silu",
   "initializer_factor": 1.0,
+  "is_encoder_decoder": true,
+  "is_gated_act": true,
   "layer_norm_epsilon": 1e-06,
-  "model_type": "fast_ankh",
+  "model_type": "t5",
+  "num_decoder_layers": 24,
   "num_heads": 16,
   "num_layers": 48,
+  "output_past": true,
   "pad_token_id": 0,
   "relative_attention_max_distance": 128,
   "relative_attention_num_buckets": 64,
   "tie_word_embeddings": false,
-  "transformers_version": "4.57.6",
+  "torch_dtype": "float32",
+  "transformers_version": "4.26.1",
+  "use_cache": true,
   "vocab_size": 144
 }

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/ankh/modeling_ankh.py

@@ -0,0 +1,1627 @@
+from __future__ import annotations
+
+import math
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from numbers import Real
+from typing import Any, ClassVar
+
+import torch
+import torch.nn as nn
+from tokenizers import pre_tokenizers
+from torch.nn import functional as F
+from transformers import (
+    AutoTokenizer,
+    PretrainedConfig,
+    PreTrainedModel,
+    T5ForConditionalGeneration,
+)
+from transformers.modeling_outputs import (
+    MaskedLMOutput,
+    ModelOutput,
+    SequenceClassifierOutput,
+    TokenClassifierOutput,
+)
+
+try:
+    from fastplms.attention import (
+        AttentionBackend,
+        FastPLMsAttentionMixin,
+        bool_to_additive_mask,
+        get_attention_mask,
+        resolve_attention_backend,
+        resolve_attention_backend_for_call,
+        set_config_attn_implementation,
+    )
+    from fastplms.embeddings import (
+        EmbeddingBatch,
+        EmbeddingMixin,
+        select_hidden_state_embeddings,
+    )
+    from fastplms.models.ttt import FastPLMTestTimeTrainingMixin
+except ModuleNotFoundError as error:
+    _COMPOSITE_REQUIRED_NAMES = (
+        "AttentionBackend",
+        "EmbeddingBatch",
+        "EmbeddingMixin",
+        "FastPLMsAttentionMixin",
+        "FastPLMTestTimeTrainingMixin",
+        "bool_to_additive_mask",
+        "get_attention_mask",
+        "resolve_attention_backend",
+        "resolve_attention_backend_for_call",
+        "select_hidden_state_embeddings",
+        "set_config_attn_implementation",
+    )
+    if error.name != "fastplms" or any(
+        name not in globals() for name in _COMPOSITE_REQUIRED_NAMES
+    ):
+        raise
+    # Legacy flat Hub composites define every shared symbol above this block.
+
+
+# ---------------------------------------------------------------------------
+# Output dataclasses
+# ---------------------------------------------------------------------------
+
+
+@dataclass
+class AnkhEncoderOutput(ModelOutput):
+    last_hidden_state: torch.Tensor | None = None
+    hidden_states: tuple[torch.Tensor, ...] | None = None
+    attentions: tuple[torch.Tensor, ...] | None = None
+
+
+@dataclass
+class AnkhMaskedLMOutput(ModelOutput):
+    loss: torch.Tensor | None = None
+    logits: torch.Tensor | None = None
+    last_hidden_state: torch.Tensor | None = None
+    hidden_states: tuple[torch.Tensor, ...] | None = None
+    attentions: tuple[torch.Tensor, ...] | None = None
+
+
+# ---------------------------------------------------------------------------
+# Config
+# ---------------------------------------------------------------------------
+
+
+class FastAnkhConfig(PretrainedConfig):
+    model_type = "fast_ankh"
+    attribute_map: ClassVar[dict[str, str]] = {
+        "head_dim": "d_kv",
+        "hidden_size": "d_model",
+        "num_attention_heads": "num_heads",
+        "num_hidden_layers": "num_layers",
+    }
+
+    def __init__(
+        self,
+        vocab_size: int = 144,
+        d_model: int = 768,
+        d_kv: int = 64,
+        d_ff: int = 3072,
+        num_heads: int = 12,
+        num_layers: int = 48,
+        num_decoder_layers: int | None = None,
+        relative_attention_num_buckets: int = 64,
+        relative_attention_max_distance: int = 128,
+        dense_act_fn: str = "gelu_new",
+        feed_forward_proj: str | None = None,
+        dropout_rate: float = 0.0,
+        layer_norm_epsilon: float = 1e-6,
+        initializer_factor: float = 1.0,
+        pad_token_id: int = 0,
+        eos_token_id: int = 1,
+        decoder_start_token_id: int | None = None,
+        use_cache: bool = True,
+        tie_word_embeddings: bool = True,
+        attn_backend: str | None = None,
+        **kwargs,
+    ):
+        if feed_forward_proj is None:
+            feed_forward_proj = (
+                "gated-gelu" if dense_act_fn == "gelu_new" else f"gated-{dense_act_fn}"
+            )
+        if decoder_start_token_id is None:
+            decoder_start_token_id = pad_token_id
+        if isinstance(dropout_rate, bool) or not isinstance(dropout_rate, Real):
+            raise TypeError("dropout_rate must be a real number in [0, 1).")
+        dropout_rate = float(dropout_rate)
+        if not 0.0 <= dropout_rate < 1.0:
+            raise ValueError("dropout_rate must be in [0, 1).")
+        serialized_encoder_decoder = kwargs.pop("is_encoder_decoder", True)
+        if serialized_encoder_decoder is not True:
+            raise ValueError(
+                "FastAnkhConfig requires is_encoder_decoder=true to match the official "
+                "T5 configuration."
+            )
+        super().__init__(
+            pad_token_id=pad_token_id,
+            eos_token_id=eos_token_id,
+            decoder_start_token_id=decoder_start_token_id,
+            is_encoder_decoder=True,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        self.d_kv = d_kv
+        self.d_ff = d_ff
+        self.num_heads = num_heads
+        self.num_layers = num_layers
+        self.num_decoder_layers = num_layers if num_decoder_layers is None else num_decoder_layers
+        self.relative_attention_num_buckets = relative_attention_num_buckets
+        self.relative_attention_max_distance = relative_attention_max_distance
+        self.dense_act_fn = dense_act_fn
+        self.feed_forward_proj = feed_forward_proj
+        self.is_gated_act = feed_forward_proj.startswith("gated-")
+        self.dropout_rate = dropout_rate
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.initializer_factor = initializer_factor
+        self.use_cache = use_cache
+        self.scale_decoder_outputs = tie_word_embeddings
+        self.tie_word_embeddings = tie_word_embeddings
+        self.attn_backend = attn_backend
+
+    def to_dict(self) -> dict[str, Any]:
+        output = super().to_dict()
+        return output
+
+
+_TOKENIZER_LOAD_CONTEXT_KEYS = (
+    "cache_dir",
+    "force_download",
+    "local_files_only",
+    "proxies",
+    "subfolder",
+    "token",
+    "trust_remote_code",
+)
+
+
+def configure_ankh_tokenizer(tokenizer: Any) -> Any:
+    """Apply ANKH's residue-aware pre-tokenizer to a tokenizer instance.
+
+    The tokenizer files published by the official checkpoints use a leading
+    metaspace convention intended for natural-language text.  Protein inputs
+    are already residue-delimited, so retaining that convention emits a
+    leading ``<unk>`` token.  FastPLMs configures the fast tokenizer to split
+    raw residue strings and tight sentinel prompts without manufacturing a
+    whitespace token.
+    """
+
+    backend = getattr(tokenizer, "backend_tokenizer", None)
+    if backend is None:
+        if getattr(tokenizer, "is_fast", None) is False:
+            raise TypeError(
+                "ANKH requires a fast tokenizer so its residue-aware pre-tokenizer "
+                "can be configured."
+            )
+        # Lightweight tokenizer doubles used by offline CPU contracts need not
+        # expose a Rust tokenizer backend.
+        return tokenizer
+    backend.pre_tokenizer = pre_tokenizers.Metaspace(
+        replacement="\u2581",
+        prepend_scheme="never",
+        split=True,
+    )
+    return tokenizer
+
+
+def normalize_ankh_sequence(sequence: str) -> str:
+    """Return one ANKH protein sequence in canonical raw-residue form."""
+
+    if not isinstance(sequence, str):
+        raise TypeError("ANKH protein sequences must be strings.")
+    normalized = "".join(sequence.split())
+    if not normalized:
+        raise ValueError("ANKH protein sequences must not be empty or whitespace-only.")
+    return normalized
+
+
+def normalize_ankh_decoder_prompt(prompt: str) -> str:
+    """Return a decoder prompt with residues and sentinels directly adjacent."""
+
+    if not isinstance(prompt, str):
+        raise TypeError("ANKH decoder prompts must be strings.")
+    normalized = "".join(prompt.split())
+    if not normalized:
+        raise ValueError("ANKH decoder prompts must not be empty or whitespace-only.")
+    return normalized
+
+
+def _normalize_ankh_text_batch(
+    values: str | Sequence[str],
+    *,
+    field: str,
+) -> str | list[str]:
+    normalizer = normalize_ankh_sequence if field == "sequence" else normalize_ankh_decoder_prompt
+    if isinstance(values, str):
+        return normalizer(values)
+    if isinstance(values, bytes) or not isinstance(values, Sequence):
+        raise TypeError(f"ANKH {field} inputs must be a string or a sequence of strings.")
+    return [normalizer(value) for value in values]
+
+
+def tokenize_ankh_sequences(
+    tokenizer: Any,
+    sequences: str | Sequence[str],
+    **tokenizer_kwargs: Any,
+) -> Any:
+    """Tokenize raw ANKH protein sequences with one model-wide contract."""
+
+    configured = configure_ankh_tokenizer(tokenizer)
+    normalized = _normalize_ankh_text_batch(sequences, field="sequence")
+    return configured(normalized, **tokenizer_kwargs)
+
+
+def tokenize_ankh_decoder_prompts(
+    tokenizer: Any,
+    prompts: str | Sequence[str],
+    **tokenizer_kwargs: Any,
+) -> Any:
+    """Tokenize explicit ANKH decoder prompts without whitespace ``<unk>`` tokens."""
+
+    configured = configure_ankh_tokenizer(tokenizer)
+    normalized = _normalize_ankh_text_batch(prompts, field="decoder prompt")
+    return configured(normalized, **tokenizer_kwargs)
+
+
+def _load_ankh_tokenizer(
+    config: FastAnkhConfig,
+    load_context: Mapping[str, Any] | None = None,
+):
+    """Load the tokenizer from the same immutable checkpoint as the model."""
+    name_or_path = str(getattr(config, "_name_or_path", "")).strip()
+    if not name_or_path:
+        raise RuntimeError(
+            "ANKH tokenizer loading requires a model loaded with from_pretrained "
+            "so checkpoint provenance is available."
+        )
+    tokenizer_kwargs = {
+        key: value
+        for key, value in dict(load_context or {}).items()
+        if key in _TOKENIZER_LOAD_CONTEXT_KEYS and value is not None
+    }
+    # The resolved commit is authoritative. In particular, do not reload the
+    # tokenizer from a moving branch when Transformers resolved model weights to
+    # an immutable Hub commit.
+    revision = getattr(config, "_commit_hash", None)
+    if revision:
+        tokenizer_kwargs["revision"] = revision
+    tokenizer = AutoTokenizer.from_pretrained(name_or_path, **tokenizer_kwargs)
+    return configure_ankh_tokenizer(tokenizer)
+
+
+class _AnkhTokenizerLoadMixin:
+    """Keep tokenizer loading scoped to the model instance and weight request."""
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        load_context = {key: kwargs[key] for key in _TOKENIZER_LOAD_CONTEXT_KEYS if key in kwargs}
+        loaded = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+        model = loaded[0] if isinstance(loaded, tuple) else loaded
+        model.__dict__["_fastplms_tokenizer_load_context"] = load_context
+        model.__dict__["_fastplms_tokenizer"] = None
+        return loaded
+
+    @property
+    def tokenizer(self):
+        tokenizer = self.__dict__.get("_fastplms_tokenizer")
+        if tokenizer is None:
+            tokenizer = _load_ankh_tokenizer(
+                self.config,
+                self.__dict__.get("_fastplms_tokenizer_load_context"),
+            )
+            self.__dict__["_fastplms_tokenizer"] = tokenizer
+        return tokenizer
+
+    @tokenizer.setter
+    def tokenizer(self, value) -> None:
+        self.__dict__["_fastplms_tokenizer"] = configure_ankh_tokenizer(value)
+
+    def _tokenize_sequence_batch(
+        self,
+        sequences: Sequence[str],
+        *,
+        tokenizer: Any | None = None,
+        **tokenizer_kwargs: Any,
+    ) -> Any:
+        resolved_tokenizer = tokenizer if tokenizer is not None else self.tokenizer
+        return tokenize_ankh_sequences(
+            resolved_tokenizer,
+            sequences,
+            **tokenizer_kwargs,
+        )
+
+    def embed_dataset(self, inputs: Any, **kwargs: Any) -> Any:
+        explicit_tokenizer = kwargs.get("tokenizer")
+        if explicit_tokenizer is not None:
+            kwargs["tokenizer"] = configure_ankh_tokenizer(explicit_tokenizer)
+        decoder_inputs = kwargs.get("decoder_inputs")
+        if (
+            decoder_inputs is not None
+            and not isinstance(decoder_inputs, (str, bytes))
+            and isinstance(decoder_inputs, Sequence)
+        ):
+            kwargs["decoder_inputs"] = [
+                normalize_ankh_decoder_prompt(value) for value in decoder_inputs
+            ]
+        return EmbeddingMixin.embed_dataset(self, inputs, **kwargs)
+
+
+def _validate_hidden_state_source(hidden_state_source: str) -> str:
+    if hidden_state_source not in {"encoder", "decoder"}:
+        raise ValueError(
+            "hidden_state_source must be either 'encoder' or 'decoder'; "
+            f"received {hidden_state_source!r}."
+        )
+    return hidden_state_source
+
+
+def _require_encoder_embedding_source(
+    hidden_state_source: str,
+    *,
+    decoder_inputs: Sequence[str] | None = None,
+    decoder_input_ids: torch.Tensor | None = None,
+    decoder_attention_mask: torch.Tensor | None = None,
+) -> None:
+    source = _validate_hidden_state_source(hidden_state_source)
+    if source == "decoder":
+        raise ValueError(
+            "Decoder hidden states require FastAnkhForConditionalGeneration loaded "
+            "through AutoModelForSeq2SeqLM; the encoder-only ANKH view does not "
+            "allocate a decoder."
+        )
+    decoder_values = (decoder_inputs, decoder_input_ids, decoder_attention_mask)
+    if any(value is not None for value in decoder_values):
+        raise ValueError(
+            "decoder_inputs, decoder_input_ids, and decoder_attention_mask are only "
+            "valid when hidden_state_source='decoder'."
+        )
+
+
+def _biological_token_mask(
+    input_ids: torch.Tensor,
+    attention_mask: torch.Tensor,
+    tokenizer: Any,
+) -> torch.Tensor:
+    mask = attention_mask.to(device=input_ids.device, dtype=torch.bool)
+    special_ids = tuple(int(value) for value in getattr(tokenizer, "all_special_ids", ()))
+    if special_ids:
+        mask = mask & ~torch.isin(
+            input_ids,
+            torch.tensor(special_ids, device=input_ids.device, dtype=input_ids.dtype),
+        )
+    return mask
+
+
+# ---------------------------------------------------------------------------
+# Submodules
+# ---------------------------------------------------------------------------
+
+
+class AnkhRMSNorm(nn.Module):
+    """T5-style RMS layer norm: scales without mean subtraction or bias."""
+
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
+        return self.weight * hidden_states.to(self.weight.dtype)
+
+
+def _gelu_new(x: torch.Tensor) -> torch.Tensor:
+    return (
+        0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))
+    )
+
+
+class AnkhGatedFFN(nn.Module):
+    """T5-style gated feed-forward: activation(wi_0(x)) * wi_1(x) -> wo."""
+
+    def __init__(self, config: FastAnkhConfig):
+        super().__init__()
+        self.wi_0 = nn.Linear(config.d_model, config.d_ff, bias=False)
+        self.wi_1 = nn.Linear(config.d_model, config.d_ff, bias=False)
+        self.wo = nn.Linear(config.d_ff, config.d_model, bias=False)
+        self.act = F.silu if config.dense_act_fn == "silu" else _gelu_new
+        self.dropout = nn.Dropout(config.dropout_rate)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.act(self.wi_0(hidden_states)) * self.wi_1(hidden_states)
+        return self.wo(self.dropout(hidden_states))
+
+
+# ---------------------------------------------------------------------------
+# Attention
+# ---------------------------------------------------------------------------
+
+
+class AnkhSelfAttention(nn.Module):
+    """T5-style self-attention with relative position bias and multi-backend dispatch.
+
+    Only layer 0 has ``has_relative_attention_bias=True`` and owns the
+    ``nn.Embedding`` that produces the position bias.  All other layers
+    receive the precomputed bias through the forward call.
+    """
+
+    def __init__(self, config: FastAnkhConfig, has_relative_attention_bias: bool = False):
+        super().__init__()
+        self.num_heads = config.num_heads
+        self.d_kv = config.d_kv
+        self.inner_dim = self.num_heads * self.d_kv
+        self.has_relative_attention_bias = has_relative_attention_bias
+        self.relative_attention_num_buckets = config.relative_attention_num_buckets
+        self.relative_attention_max_distance = config.relative_attention_max_distance
+
+        self.q = nn.Linear(config.d_model, self.inner_dim, bias=False)
+        self.k = nn.Linear(config.d_model, self.inner_dim, bias=False)
+        self.v = nn.Linear(config.d_model, self.inner_dim, bias=False)
+        self.o = nn.Linear(self.inner_dim, config.d_model, bias=False)
+        # T5/ANKH attention is unscaled: scores = Q K^T (no 1/sqrt(d_kv)).
+        # The learned relative position bias absorbs any temperature.
+        self.scale = 1.0
+        self.dropout_prob = float(config.dropout_rate)
+
+        if self.has_relative_attention_bias:
+            self.relative_attention_bias = nn.Embedding(
+                config.relative_attention_num_buckets, config.num_heads
+            )
+
+        self.attn_backend: AttentionBackend = AttentionBackend.SDPA  # set by encoder
+
+    # ---- T5 relative position bucketing ----
+
+    @staticmethod
+    def _relative_position_bucket(
+        relative_position: torch.Tensor,
+        num_buckets: int = 32,
+        max_distance: int = 128,
+    ) -> torch.Tensor:
+        """Bidirectional log-bucketed relative position mapping (T5 style)."""
+        # Bidirectional: half buckets for negative, half for positive
+        num_buckets //= 2
+        relative_buckets = (relative_position > 0).to(torch.long) * num_buckets
+        relative_position = torch.abs(relative_position)
+
+        max_exact = num_buckets // 2
+        is_small = relative_position < max_exact
+
+        relative_position_if_large = max_exact + (
+            torch.log(relative_position.float() / max_exact)
+            / math.log(max_distance / max_exact)
+            * (num_buckets - max_exact)
+        ).to(torch.long)
+        relative_position_if_large = torch.clamp(relative_position_if_large, max=num_buckets - 1)
+
+        relative_buckets += torch.where(is_small, relative_position, relative_position_if_large)
+        return relative_buckets
+
+    def compute_bias(
+        self, query_length: int, key_length: int, device: torch.device
+    ) -> torch.Tensor:
+        """Compute the position-bias tensor A with shape (1, h, q, k)."""
+        context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None]
+        memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :]
+        relative_position = memory_position - context_position
+        buckets = self._relative_position_bucket(
+            relative_position,
+            num_buckets=self.relative_attention_num_buckets,
+            max_distance=self.relative_attention_max_distance,
+        )
+        values = self.relative_attention_bias(buckets)  # A has shape (q, k, h).
+        return values.permute(2, 0, 1).unsqueeze(0)  # A has shape (1, h, q, k).
+
+    # ---- Forward ----
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask_4d: torch.Tensor | None = None,
+        position_bias: torch.Tensor | None = None,
+        output_attentions: bool = False,
+        effective_backend: AttentionBackend | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+        """Returns (attn_output, attn_weights_or_none, position_bias)."""
+        batch_size, seq_length = hidden_states.shape[:2]
+        hidden_shape = (batch_size, seq_length, self.num_heads, self.d_kv)
+
+        query_heads = self.q(hidden_states).view(hidden_shape).transpose(1, 2)
+        key_heads = self.k(hidden_states).view(hidden_shape).transpose(1, 2)
+        value_heads = self.v(hidden_states).view(hidden_shape).transpose(1, 2)
+
+        # The first layer computes the bias once; later layers reuse it.
+        if position_bias is None and self.has_relative_attention_bias:
+            position_bias = self.compute_bias(seq_length, seq_length, hidden_states.device)
+            # Fold padding mask into position bias so layers don't need separate mask.
+            if attention_mask_4d is not None:
+                position_bias = position_bias + bool_to_additive_mask(
+                    attention_mask_4d, position_bias.dtype
+                )
+
+        if effective_backend is None:
+            effective_backend = resolve_attention_backend_for_call(
+                self.attn_backend,
+                output_attentions=output_attentions,
+            )
+        if output_attentions:
+            attn_output, attn_weights = self._manual_attn(
+                query_heads, key_heads, value_heads, position_bias
+            )
+            return self.o(attn_output), attn_weights, position_bias
+
+        if effective_backend == AttentionBackend.EAGER:
+            attn_output, _ = self._manual_attn(query_heads, key_heads, value_heads, position_bias)
+        elif effective_backend == AttentionBackend.SDPA:
+            attn_output = self._sdpa_attn(query_heads, key_heads, value_heads, position_bias)
+        else:
+            raise AssertionError(f"Unsupported backend for ANKH: {effective_backend}")
+
+        return self.o(attn_output), None, position_bias
+
+    def _sdpa_attn(
+        self,
+        query_heads: torch.Tensor,
+        key_heads: torch.Tensor,
+        value_heads: torch.Tensor,
+        position_bias: torch.Tensor | None,
+    ) -> torch.Tensor:
+        # A is the additive position bias with shape (1, h, q, k), including padding.
+        # Never mutate torch.backends.cuda process-global reduction policy from
+        # a model forward. Concurrent model requests must not change each
+        # other's numerical behavior or restore a stale process setting.
+        context_heads = F.scaled_dot_product_attention(
+            query_heads,
+            key_heads,
+            value_heads,
+            attn_mask=position_bias,
+            dropout_p=self.dropout_prob if self.training else 0.0,
+            scale=self.scale,
+        )
+        return (
+            context_heads.transpose(1, 2)
+            .contiguous()
+            .view(query_heads.shape[0], -1, self.inner_dim)
+        )
+
+    def _manual_attn(
+        self,
+        query_heads: torch.Tensor,
+        key_heads: torch.Tensor,
+        value_heads: torch.Tensor,
+        position_bias: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        attn_weights = torch.matmul(query_heads, key_heads.transpose(-1, -2)) * self.scale
+        if position_bias is not None:
+            attn_weights = attn_weights + position_bias
+        attn_weights = F.softmax(attn_weights.float(), dim=-1).type_as(attn_weights)
+        if self.dropout_prob > 0 and self.training:
+            attn_weights = F.dropout(
+                attn_weights,
+                p=self.dropout_prob,
+                training=self.training,
+            )
+        context_heads = torch.matmul(attn_weights, value_heads)
+        attn_output = (
+            context_heads.transpose(1, 2)
+            .contiguous()
+            .view(query_heads.shape[0], -1, self.inner_dim)
+        )
+        return attn_output, attn_weights
+
+
+# ---------------------------------------------------------------------------
+# Encoder block & stack (T5-compatible key naming)
+# ---------------------------------------------------------------------------
+
+
+class AnkhSelfAttentionLayer(nn.Module):
+    """Wraps AnkhSelfAttention + layer_norm to match T5Block.layer[0] key naming."""
+
+    def __init__(self, config: FastAnkhConfig, has_relative_attention_bias: bool = False):
+        super().__init__()
+        self.SelfAttention = AnkhSelfAttention(config, has_relative_attention_bias)
+        self.layer_norm = AnkhRMSNorm(config.d_model, eps=config.layer_norm_epsilon)
+        self.dropout = nn.Dropout(config.dropout_rate)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask_4d: torch.Tensor | None = None,
+        position_bias: torch.Tensor | None = None,
+        output_attentions: bool = False,
+        effective_backend: AttentionBackend | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+        normed = self.layer_norm(hidden_states)
+        attn_output, attn_weights, position_bias = self.SelfAttention(
+            normed,
+            attention_mask_4d=attention_mask_4d,
+            position_bias=position_bias,
+            output_attentions=output_attentions,
+            effective_backend=effective_backend,
+        )
+        hidden_states = hidden_states + self.dropout(attn_output)
+        return hidden_states, attn_weights, position_bias
+
+
+class AnkhFFLayer(nn.Module):
+    """Wraps AnkhGatedFFN + layer_norm to match T5Block.layer[1] key naming."""
+
+    def __init__(self, config: FastAnkhConfig):
+        super().__init__()
+        self.DenseReluDense = AnkhGatedFFN(config)
+        self.layer_norm = AnkhRMSNorm(config.d_model, eps=config.layer_norm_epsilon)
+        self.dropout = nn.Dropout(config.dropout_rate)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        normed = self.layer_norm(hidden_states)
+        hidden_states = hidden_states + self.dropout(self.DenseReluDense(normed))
+        return hidden_states
+
+
+class AnkhBlock(nn.Module):
+    """Single transformer block with T5-compatible .layer ModuleList naming."""
+
+    def __init__(self, config: FastAnkhConfig, has_relative_attention_bias: bool = False):
+        super().__init__()
+        self.layer = nn.ModuleList(
+            [
+                AnkhSelfAttentionLayer(config, has_relative_attention_bias),
+                AnkhFFLayer(config),
+            ]
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask_4d: torch.Tensor | None = None,
+        position_bias: torch.Tensor | None = None,
+        output_attentions: bool = False,
+        effective_backend: AttentionBackend | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+        hidden_states, attn_weights, position_bias = self.layer[0](
+            hidden_states,
+            attention_mask_4d=attention_mask_4d,
+            position_bias=position_bias,
+            output_attentions=output_attentions,
+            effective_backend=effective_backend,
+        )
+        hidden_states = self.layer[1](hidden_states)
+        return hidden_states, attn_weights, position_bias
+
+
+# ---------------------------------------------------------------------------
+# PreTrainedModel base
+# ---------------------------------------------------------------------------
+
+
+class AnkhPreTrainedModel(
+    _AnkhTokenizerLoadMixin,
+    FastPLMsAttentionMixin,
+    PreTrainedModel,
+):
+    config_class = FastAnkhConfig
+    base_model_prefix = "encoder"
+    supports_gradient_checkpointing = True
+    _no_split_modules: ClassVar[list[str]] = ["AnkhBlock"]
+    _supports_flash_attn_2 = False
+    _supports_flash_attn_3 = False
+    _supports_flex_attn = False
+    _fastplms_attention_implementations = ("eager", "sdpa")
+    embedding_unsupported_pooling = ("cls",)
+
+    def __init__(self, config: FastAnkhConfig, *args, **kwargs) -> None:
+        super().__init__(config, *args, **kwargs)
+        self.__dict__["_fastplms_tokenizer"] = None
+        self.__dict__["_fastplms_tokenizer_load_context"] = {}
+
+    @torch.no_grad()
+    def _init_weights(self, module: nn.Module) -> None:
+        factor = self.config.initializer_factor
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=factor * (self.config.d_model**-0.5))
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=factor * 1.0)
+        elif isinstance(module, AnkhRMSNorm):
+            module.weight.data.fill_(1.0)
+
+    def post_init(self) -> None:
+        super().post_init()
+
+    def get_output_embeddings(self):
+        return None
+
+    def _embedding_metadata(self, **context: Any) -> Mapping[str, Any]:
+        source = _validate_hidden_state_source(context.get("hidden_state_source", "encoder"))
+        if source != "encoder":
+            raise ValueError(
+                "Decoder hidden states require FastAnkhForConditionalGeneration loaded "
+                "through AutoModelForSeq2SeqLM."
+            )
+        return {
+            "architecture": "ANKH-T5",
+            "hidden_state_stack": "encoder",
+            "layer_order": "embedding-plus-transformer-blocks",
+        }
+
+    @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.config.attn_backend = backend
+        resolved = resolve_attention_backend(backend)
+        for module in self.modules():
+            if isinstance(module, FAST_ANKH_ENCODER):
+                module.attention_backend = resolved
+            elif isinstance(module, AnkhSelfAttention):
+                module.attn_backend = resolved
+
+
+# ---------------------------------------------------------------------------
+# FAST_ANKH_ENCODER (mirrors T5Stack key naming)
+# ---------------------------------------------------------------------------
+
+
+class FAST_ANKH_ENCODER(AnkhPreTrainedModel, EmbeddingMixin):
+    """Inner encoder that mirrors T5Stack attribute naming for weight compliance.
+
+    State dict keys: embed_tokens.*, block.{i}.layer.0.SelfAttention.*,
+    block.{i}.layer.1.DenseReluDense.*, final_layer_norm.*.
+    """
+
+    def __init__(self, config: FastAnkhConfig, **kwargs):
+        AnkhPreTrainedModel.__init__(self, config, **kwargs)
+        self.config = config
+
+        resolved = resolve_attention_backend(config.attn_backend)
+        if resolved.is_flash:
+            raise ValueError(
+                "ANKH does not support FlashAttention because it requires relative position bias."
+            )
+        self.attention_backend = resolved
+
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.d_model)
+        self.block = nn.ModuleList(
+            [
+                AnkhBlock(config, has_relative_attention_bias=(i == 0))
+                for i in range(config.num_layers)
+            ]
+        )
+        for blk in self.block:
+            blk.layer[0].SelfAttention.attn_backend = self.attention_backend
+
+        self.final_layer_norm = AnkhRMSNorm(config.d_model, eps=config.layer_norm_epsilon)
+        self.dropout = nn.Dropout(config.dropout_rate)
+        self.gradient_checkpointing = False
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.embed_tokens = value
+
+    def _embed(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        hidden_state_index: int = -1,
+        store_all_hidden_states: bool = False,
+        hidden_state_source: str = "encoder",
+        decoder_inputs: Sequence[str] | None = None,
+        decoder_input_ids: torch.Tensor | None = None,
+        decoder_attention_mask: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        _require_encoder_embedding_source(
+            hidden_state_source,
+            decoder_inputs=decoder_inputs,
+            decoder_input_ids=decoder_input_ids,
+            decoder_attention_mask=decoder_attention_mask,
+        )
+        hidden_states = self.embed_tokens(input_ids)
+        output_hidden_states = store_all_hidden_states or hidden_state_index != -1
+        encoder_output = self._run_encoder(
+            hidden_states,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states,
+        )
+        return select_hidden_state_embeddings(
+            encoder_output.last_hidden_state,
+            encoder_output.hidden_states,
+            hidden_state_index=hidden_state_index,
+            store_all_hidden_states=store_all_hidden_states,
+        )
+
+    def _run_encoder(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        output_hidden_states: bool = False,
+        output_attentions: bool = False,
+    ) -> AnkhEncoderOutput:
+        # T5Stack applies this module both to the input embeddings and after
+        # final normalization. Keeping those as separate calls preserves the
+        # official training-time stochastic path without affecting eval mode.
+        hidden_states = self.dropout(hidden_states)
+        all_hidden_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+
+        batch_size, seq_len = hidden_states.shape[:2]
+        effective_backend = resolve_attention_backend_for_call(
+            self.attention_backend,
+            output_attentions=output_attentions,
+        )
+        _, attention_mask_4d, _ = get_attention_mask(
+            effective_backend=effective_backend,
+            batch_size=batch_size,
+            seq_len=seq_len,
+            device=hidden_states.device,
+            attention_mask=attention_mask,
+            dtype=hidden_states.dtype,
+            mask_semantics="padding",
+        )
+
+        position_bias = None
+
+        for layer_module in self.block:
+            if output_hidden_states:
+                all_hidden_states = (*all_hidden_states, hidden_states)
+
+            if self.gradient_checkpointing and self.training:
+                hidden_states, attn_weights, position_bias = self._gradient_checkpointing_func(
+                    layer_module.__call__,
+                    hidden_states,
+                    attention_mask_4d,
+                    position_bias,
+                    output_attentions,
+                    effective_backend,
+                )
+            else:
+                hidden_states, attn_weights, position_bias = layer_module(
+                    hidden_states,
+                    attention_mask_4d=attention_mask_4d,
+                    position_bias=position_bias,
+                    output_attentions=output_attentions,
+                    effective_backend=effective_backend,
+                )
+
+            if all_attentions is not None:
+                all_attentions = (*all_attentions, attn_weights)
+
+        hidden_states = self.dropout(self.final_layer_norm(hidden_states))
+
+        if output_hidden_states:
+            all_hidden_states = (*all_hidden_states, hidden_states)
+
+        return AnkhEncoderOutput(
+            last_hidden_state=hidden_states,
+            hidden_states=all_hidden_states,
+            attentions=all_attentions,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        output_attentions: bool | None = None,
+        return_dict: bool | None = None,
+    ) -> AnkhEncoderOutput | tuple[torch.Tensor, ...]:
+        output_attentions = (
+            output_attentions if output_attentions is not None else self.config.output_attentions
+        )
+        output_hidden_states = (
+            output_hidden_states
+            if output_hidden_states is not None
+            else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            hidden_states = self.embed_tokens(input_ids)
+        elif inputs_embeds is not None:
+            hidden_states = inputs_embeds
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        outputs = self._run_encoder(
+            hidden_states,
+            attention_mask=attention_mask,
+            output_hidden_states=output_hidden_states or False,
+            output_attentions=output_attentions or False,
+        )
+        return outputs if return_dict else outputs.to_tuple()
+
+
+# ---------------------------------------------------------------------------
+# Model classes
+# ---------------------------------------------------------------------------
+
+
+class FastAnkhModel(AnkhPreTrainedModel, EmbeddingMixin):
+    """ANKH encoder model for embedding extraction."""
+
+    _tied_weights_keys: ClassVar[dict[str, str]] = {"encoder.embed_tokens.weight": "shared.weight"}
+    # The published ANKH checkpoint is the complete official T5 state. AutoModel
+    # intentionally exposes only its encoder view without allocating a decoder.
+    _keys_to_ignore_on_load_unexpected: ClassVar[list[str]] = [
+        r"^decoder\.",
+        r"^lm_head\.",
+    ]
+
+    def __init__(self, config: FastAnkhConfig, **kwargs):
+        AnkhPreTrainedModel.__init__(self, config, **kwargs)
+        self.config = config
+        self.shared = nn.Embedding(config.vocab_size, config.d_model)
+        self.encoder = FAST_ANKH_ENCODER(config)
+        self.encoder.embed_tokens = self.shared
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.encoder.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.shared = value
+        self.encoder.embed_tokens = value
+
+    def _embed(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        hidden_state_index: int = -1,
+        store_all_hidden_states: bool = False,
+        **embedding_kwargs,
+    ) -> torch.Tensor:
+        return self.encoder._embed(
+            input_ids,
+            attention_mask,
+            hidden_state_index=hidden_state_index,
+            store_all_hidden_states=store_all_hidden_states,
+            **embedding_kwargs,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        output_attentions: bool | None = None,
+        return_dict: bool | None = None,
+    ) -> AnkhEncoderOutput | tuple[torch.Tensor, ...]:
+        return self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+            return_dict=return_dict,
+        )
+
+
+class FastAnkhForMaskedLMExtension(
+    FastPLMTestTimeTrainingMixin, AnkhPreTrainedModel, EmbeddingMixin
+):
+    """ANKH encoder with LM head for masked language modeling.
+
+    NOTE: The LM head is initialized from the shared embedding weights but is NOT
+    tied. The original ANKH models were trained with T5's span corruption objective
+    using an encoder-decoder architecture. This encoder-only MaskedLM variant is
+    not pre-trained for standard MLM and requires additional fine-tuning.
+    """
+
+    _tied_weights_keys: ClassVar[dict[str, str]] = {"encoder.embed_tokens.weight": "shared.weight"}
+    _keys_to_ignore_on_load_unexpected: ClassVar[list[str]] = [r"^decoder\."]
+
+    def __init__(self, config: FastAnkhConfig, **kwargs):
+        # The historical Synthyra extension stores an independent output head.
+        config.tie_word_embeddings = False
+        AnkhPreTrainedModel.__init__(self, config, **kwargs)
+        self.config = config
+        self.shared = nn.Embedding(config.vocab_size, config.d_model)
+        self.encoder = FAST_ANKH_ENCODER(config)
+        self.encoder.embed_tokens = self.shared
+        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
+        self.loss_fct = nn.CrossEntropyLoss()
+        self.post_init()
+        self.init_ttt({"lora_target_replace_module": "AnkhSelfAttention"})
+
+    def get_input_embeddings(self):
+        return self.encoder.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.shared = value
+        self.encoder.embed_tokens = value
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings):
+        self.lm_head = 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,
+        **embedding_kwargs,
+    ) -> torch.Tensor:
+        return self.encoder._embed(
+            input_ids,
+            attention_mask,
+            hidden_state_index=hidden_state_index,
+            store_all_hidden_states=store_all_hidden_states,
+            **embedding_kwargs,
+        )
+
+    def _ttt_get_trainable_modules(self) -> list[nn.Module]:
+        return [self.encoder]
+
+    def _ttt_tokenize(
+        self,
+        seq: str | list[str] | None = None,
+        input_ids: torch.Tensor | None = None,
+        **kwargs,
+    ) -> torch.Tensor:
+        del kwargs
+        if input_ids is not None:
+            return input_ids
+        if seq is None:
+            raise ValueError("Pass either seq or input_ids for ANKH TTT.")
+        sequences = [seq] if isinstance(seq, str) else seq
+        tokenized = tokenize_ankh_sequences(
+            self.tokenizer,
+            sequences,
+            return_tensors="pt",
+            padding=True,
+        )
+        return tokenized["input_ids"]
+
+    def _ttt_replacement_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
+        amino_acids = "ACDEFGHIKLMNPQRSTVWY"
+        ids = [self.tokenizer.convert_tokens_to_ids(aa) for aa in amino_acids]
+        return torch.tensor(ids, device=input_ids.device, dtype=input_ids.dtype)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        output_attentions: bool | None = None,
+        return_dict: bool | None = None,
+    ) -> MaskedLMOutput | tuple[torch.Tensor, ...]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        outputs = self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+            return_dict=True,
+        )
+        sequence_output = outputs.last_hidden_state
+        logits = self.lm_head(sequence_output)
+
+        loss = None
+        if labels is not None:
+            labels = labels.to(logits.device)
+            loss = self.loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
+
+        if not return_dict:
+            output = (logits, *outputs.to_tuple()[1:])
+            return (loss, *output) if loss is not None else output
+
+        return MaskedLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class FastAnkhForConditionalGeneration(
+    _AnkhTokenizerLoadMixin,
+    T5ForConditionalGeneration,
+    EmbeddingMixin,
+):
+    """Official ANKH sequence-to-sequence architecture with exact T5 state keys.
+
+    ANKH generation checkpoints are ordinary T5 conditional-generation models.
+    This class intentionally delegates their decoder, cross-attention, language
+    model head, caching, generation, and tied-weight behavior to Transformers.
+    The optimized encoder-only implementation remains available through
+    :class:`FastAnkhModel`.
+    """
+
+    config_class = FastAnkhConfig
+    embedding_unsupported_pooling = ("cls",)
+    _fastplms_attention_implementations = ("eager",)
+
+    def __init__(self, config: FastAnkhConfig, **kwargs):
+        requested_backend = getattr(config, "_attn_implementation", None) or config.attn_backend
+        if requested_backend not in (None, "eager"):
+            raise ValueError(
+                "ANKH sequence-to-sequence checkpoints support only eager attention; "
+                f"received {requested_backend!r}. Use FastAnkhModel for optimized "
+                "encoder embeddings."
+            )
+        set_config_attn_implementation(config, "eager")
+        super().__init__(config, **kwargs)
+        self.__dict__["_fastplms_tokenizer"] = None
+        self.__dict__["_fastplms_tokenizer_load_context"] = {}
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        decoder_input_ids: torch.Tensor | None = None,
+        decoder_attention_mask: torch.Tensor | None = None,
+        encoder_outputs: Any | None = None,
+        past_key_values: Any | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        decoder_inputs_embeds: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        return_dict: bool | None = None,
+    ) -> Any:
+        """Run official T5 seq2seq behavior with a fail-closed public signature.
+
+        Transformers' T5 forward currently accepts and silently ignores arbitrary
+        keyword arguments. FastPLMs keeps the supported T5 arguments explicit so a
+        misspelled generation, cache, or conditioning argument cannot appear to
+        have taken effect.
+        """
+
+        return super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            decoder_input_ids=decoder_input_ids,
+            decoder_attention_mask=decoder_attention_mask,
+            encoder_outputs=encoder_outputs,
+            past_key_values=past_key_values,
+            inputs_embeds=inputs_embeds,
+            decoder_inputs_embeds=decoder_inputs_embeds,
+            labels=labels,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+    def _prepare_decoder_embedding_inputs(
+        self,
+        *,
+        batch_size: int,
+        decoder_inputs: Sequence[str] | None,
+        decoder_input_ids: torch.Tensor | None,
+        decoder_attention_mask: torch.Tensor | None,
+        tokenizer: Any | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, Any]:
+        if (decoder_inputs is None) == (decoder_input_ids is None):
+            raise ValueError(
+                "hidden_state_source='decoder' requires exactly one of "
+                "decoder_inputs or decoder_input_ids. Decoder inputs are task-specific "
+                "and FastPLMs will not synthesize shifted encoder tokens."
+            )
+        resolved_tokenizer = configure_ankh_tokenizer(
+            tokenizer if tokenizer is not None else self.tokenizer
+        )
+        if decoder_inputs is not None:
+            if decoder_attention_mask is not None:
+                raise ValueError(
+                    "decoder_attention_mask may only accompany decoder_input_ids; "
+                    "decoder_inputs are tokenized with their own attention mask."
+                )
+            values = [decoder_inputs] if isinstance(decoder_inputs, str) else list(decoder_inputs)
+            if len(values) != batch_size:
+                raise ValueError(
+                    "decoder_inputs must align one-to-one with encoder inputs; "
+                    f"expected {batch_size}, received {len(values)}."
+                )
+            encoded = tokenize_ankh_decoder_prompts(
+                resolved_tokenizer,
+                values,
+                return_tensors="pt",
+                padding=True,
+                truncation=False,
+            )
+            decoder_input_ids = encoded["input_ids"]
+            decoder_attention_mask = encoded.get("attention_mask")
+        if decoder_input_ids is None:
+            raise RuntimeError(
+                "Decoder input resolution completed without decoder_input_ids."
+            )
+        if decoder_input_ids.ndim != 2 or decoder_input_ids.shape[0] != batch_size:
+            raise ValueError(
+                "decoder_input_ids must have shape (batch, decoder_sequence_length); "
+                f"expected batch {batch_size}, received {tuple(decoder_input_ids.shape)}."
+            )
+        if decoder_attention_mask is None:
+            pad_token_id = self.config.pad_token_id
+            decoder_attention_mask = (
+                torch.ones_like(decoder_input_ids, dtype=torch.bool)
+                if pad_token_id is None
+                else decoder_input_ids.ne(pad_token_id)
+            )
+            decoder_start_token_id = self.config.decoder_start_token_id
+            if (
+                decoder_input_ids.shape[1] > 0
+                and decoder_start_token_id is not None
+                and decoder_start_token_id == pad_token_id
+            ):
+                decoder_attention_mask[:, 0] |= decoder_input_ids[:, 0].eq(decoder_start_token_id)
+        if tuple(decoder_attention_mask.shape) != tuple(decoder_input_ids.shape):
+            raise ValueError(
+                "decoder_attention_mask must have the same shape as decoder_input_ids; "
+                f"received {tuple(decoder_attention_mask.shape)} and "
+                f"{tuple(decoder_input_ids.shape)}."
+            )
+        device = self.shared.weight.device
+        return (
+            decoder_input_ids.to(device=device),
+            decoder_attention_mask.to(device=device),
+            resolved_tokenizer,
+        )
+
+    def _extract_embedding_stack(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None,
+        *,
+        hidden_state_source: str,
+        hidden_state_index: int,
+        store_all_hidden_states: bool,
+        decoder_inputs: Sequence[str] | None = None,
+        decoder_input_ids: torch.Tensor | None = None,
+        decoder_attention_mask: torch.Tensor | None = None,
+        tokenizer: Any | None = None,
+        output_attentions: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, tuple[torch.Tensor, ...] | None]:
+        source = _validate_hidden_state_source(hidden_state_source)
+        device = self.shared.weight.device
+        input_ids = input_ids.to(device=device)
+        if attention_mask is None:
+            attention_mask = input_ids.ne(self.config.pad_token_id)
+        attention_mask = attention_mask.to(device=device)
+        need_hidden_states = store_all_hidden_states or hidden_state_index != -1
+        encoder_outputs = self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_hidden_states=need_hidden_states,
+            output_attentions=output_attentions and source == "encoder",
+            return_dict=True,
+        )
+        if source == "encoder":
+            if any(
+                value is not None
+                for value in (decoder_inputs, decoder_input_ids, decoder_attention_mask)
+            ):
+                raise ValueError(
+                    "Decoder inputs are only valid when hidden_state_source='decoder'."
+                )
+            X = select_hidden_state_embeddings(
+                encoder_outputs.last_hidden_state,
+                encoder_outputs.hidden_states,
+                hidden_state_index=hidden_state_index,
+                store_all_hidden_states=store_all_hidden_states,
+            )
+            return X, input_ids, attention_mask, encoder_outputs.attentions
+
+        decoder_input_ids, decoder_attention_mask, _ = self._prepare_decoder_embedding_inputs(
+            batch_size=input_ids.shape[0],
+            decoder_inputs=decoder_inputs,
+            decoder_input_ids=decoder_input_ids,
+            decoder_attention_mask=decoder_attention_mask,
+            tokenizer=tokenizer,
+        )
+        decoder_outputs = self.decoder(
+            input_ids=decoder_input_ids,
+            attention_mask=decoder_attention_mask,
+            encoder_hidden_states=encoder_outputs.last_hidden_state,
+            encoder_attention_mask=attention_mask,
+            use_cache=False,
+            output_hidden_states=need_hidden_states,
+            output_attentions=output_attentions,
+            return_dict=True,
+        )
+        X = select_hidden_state_embeddings(
+            decoder_outputs.last_hidden_state,
+            decoder_outputs.hidden_states,
+            hidden_state_index=hidden_state_index,
+            store_all_hidden_states=store_all_hidden_states,
+        )
+        return X, decoder_input_ids, decoder_attention_mask, decoder_outputs.attentions
+
+    def _embed(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        hidden_state_index: int = -1,
+        store_all_hidden_states: bool = False,
+        hidden_state_source: str = "encoder",
+        decoder_inputs: Sequence[str] | None = None,
+        decoder_input_ids: torch.Tensor | None = None,
+        decoder_attention_mask: torch.Tensor | None = None,
+        tokenizer: Any | None = None,
+    ) -> torch.Tensor:
+        X, _, _, _ = self._extract_embedding_stack(
+            input_ids,
+            attention_mask,
+            hidden_state_source=hidden_state_source,
+            hidden_state_index=hidden_state_index,
+            store_all_hidden_states=store_all_hidden_states,
+            decoder_inputs=decoder_inputs,
+            decoder_input_ids=decoder_input_ids,
+            decoder_attention_mask=decoder_attention_mask,
+            tokenizer=tokenizer,
+        )
+        return X
+
+    def _embedding_batch(
+        self,
+        sequences: Sequence[str],
+        *,
+        tokenizer: Any | None = None,
+        max_length: int | None = None,
+        truncate: bool = True,
+        need_attentions: bool = False,
+        hidden_state_source: str = "encoder",
+        hidden_state_index: int = -1,
+        store_all_hidden_states: bool = False,
+        decoder_inputs: Sequence[str] | None = None,
+        decoder_input_ids: torch.Tensor | None = None,
+        decoder_attention_mask: torch.Tensor | None = None,
+    ) -> EmbeddingBatch:
+        del max_length, truncate  # The shared runner already crops biological residues.
+        resolved_tokenizer = configure_ankh_tokenizer(
+            tokenizer if tokenizer is not None else self.tokenizer
+        )
+        encoded = tokenize_ankh_sequences(
+            resolved_tokenizer,
+            list(sequences),
+            return_tensors="pt",
+            padding=True,
+            truncation=False,
+        )
+        device = self.shared.weight.device
+        input_ids = encoded["input_ids"].to(device=device)
+        attention_mask = encoded.get("attention_mask", input_ids.new_ones(input_ids.shape)).to(
+            device=device
+        )
+        X, selected_ids, selected_attention_mask, attentions = self._extract_embedding_stack(
+            input_ids,
+            attention_mask,
+            hidden_state_source=hidden_state_source,
+            hidden_state_index=hidden_state_index,
+            store_all_hidden_states=store_all_hidden_states,
+            decoder_inputs=decoder_inputs,
+            decoder_input_ids=decoder_input_ids,
+            decoder_attention_mask=decoder_attention_mask,
+            tokenizer=resolved_tokenizer,
+            output_attentions=need_attentions,
+        )
+        residue_mask = _biological_token_mask(
+            selected_ids,
+            selected_attention_mask,
+            resolved_tokenizer,
+        )
+        return EmbeddingBatch(X=X, residue_mask=residue_mask, attentions=attentions)
+
+    def _embedding_metadata(self, **context: Any) -> Mapping[str, Any]:
+        source = _validate_hidden_state_source(context.get("hidden_state_source", "encoder"))
+        return {
+            "architecture": "ANKH-T5",
+            "hidden_state_stack": source,
+            "layer_order": "embedding-plus-transformer-blocks",
+            "decoder_inputs": "explicit-task-inputs-required" if source == "decoder" else None,
+            "decoder_residue_mask": (
+                "attention-mask-minus-tokenizer-specials" if source == "decoder" else None
+            ),
+        }
+
+
+class FastAnkhForSequenceClassification(AnkhPreTrainedModel, EmbeddingMixin):
+    _tied_weights_keys: ClassVar[dict[str, str]] = {"encoder.embed_tokens.weight": "shared.weight"}
+    _keys_to_ignore_on_load_unexpected: ClassVar[list[str]] = [
+        r"^decoder\.",
+        r"^lm_head\.",
+    ]
+
+    def __init__(self, config: FastAnkhConfig, **kwargs):
+        AnkhPreTrainedModel.__init__(self, config, **kwargs)
+        self.num_labels = config.num_labels
+        self.config = config
+        self.shared = nn.Embedding(config.vocab_size, config.d_model)
+        self.encoder = FAST_ANKH_ENCODER(config)
+        self.encoder.embed_tokens = self.shared
+        self.classifier = nn.Linear(config.d_model, config.num_labels)
+        self.mse = nn.MSELoss()
+        self.ce = nn.CrossEntropyLoss()
+        self.bce = nn.BCEWithLogitsLoss()
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.encoder.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.shared = value
+        self.encoder.embed_tokens = value
+
+    def _embed(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        hidden_state_index: int = -1,
+        store_all_hidden_states: bool = False,
+        **embedding_kwargs,
+    ) -> torch.Tensor:
+        return self.encoder._embed(
+            input_ids,
+            attention_mask,
+            hidden_state_index=hidden_state_index,
+            store_all_hidden_states=store_all_hidden_states,
+            **embedding_kwargs,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        output_attentions: bool | None = None,
+        return_dict: bool | None = None,
+    ) -> SequenceClassifierOutput | tuple[torch.Tensor, ...]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        outputs = self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+            return_dict=True,
+        )
+        # Pool: mean over non-padding tokens
+        sequence_output = outputs.last_hidden_state
+        if attention_mask is not None:
+            mask = attention_mask.unsqueeze(-1).to(sequence_output.dtype)
+            pooled = (sequence_output * mask).sum(dim=1) / mask.sum(dim=1).clamp(min=1)
+        else:
+            pooled = sequence_output.mean(dim=1)
+        logits = self.classifier(pooled)
+
+        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":
+                loss = (
+                    self.mse(logits.squeeze(), labels.squeeze())
+                    if self.num_labels == 1
+                    else 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)
+
+        if not return_dict:
+            output = (logits, *outputs.to_tuple()[1:])
+            return (loss, *output) if loss is not None else output
+
+        return SequenceClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+class FastAnkhForTokenClassification(AnkhPreTrainedModel, EmbeddingMixin):
+    _tied_weights_keys: ClassVar[dict[str, str]] = {"encoder.embed_tokens.weight": "shared.weight"}
+    _keys_to_ignore_on_load_unexpected: ClassVar[list[str]] = [
+        r"^decoder\.",
+        r"^lm_head\.",
+    ]
+
+    def __init__(self, config: FastAnkhConfig, **kwargs):
+        AnkhPreTrainedModel.__init__(self, config, **kwargs)
+        self.num_labels = config.num_labels
+        self.shared = nn.Embedding(config.vocab_size, config.d_model)
+        self.encoder = FAST_ANKH_ENCODER(config)
+        self.encoder.embed_tokens = self.shared
+        self.classifier = nn.Linear(config.d_model, config.num_labels)
+        self.loss_fct = nn.CrossEntropyLoss()
+        self.post_init()
+
+    def get_input_embeddings(self):
+        return self.encoder.embed_tokens
+
+    def set_input_embeddings(self, value):
+        self.shared = value
+        self.encoder.embed_tokens = value
+
+    def _embed(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        hidden_state_index: int = -1,
+        store_all_hidden_states: bool = False,
+        **embedding_kwargs,
+    ) -> torch.Tensor:
+        return self.encoder._embed(
+            input_ids,
+            attention_mask,
+            hidden_state_index=hidden_state_index,
+            store_all_hidden_states=store_all_hidden_states,
+            **embedding_kwargs,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        labels: torch.Tensor | None = None,
+        output_hidden_states: bool | None = None,
+        output_attentions: bool | None = None,
+        return_dict: bool | None = None,
+    ) -> TokenClassifierOutput | tuple[torch.Tensor, ...]:
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        outputs = self.encoder(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            inputs_embeds=inputs_embeds,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+            return_dict=True,
+        )
+        sequence_output = outputs.last_hidden_state
+        logits = self.classifier(sequence_output)
+
+        loss = None
+        if labels is not None:
+            labels = labels.to(logits.device)
+            loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        if not return_dict:
+            output = (logits, *outputs.to_tuple()[1:])
+            return (loss, *output) if loss is not None else output
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

fastplms/models/ttt.py

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

fastplms/registry.py

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

fastplms/runtime.py

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

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "decoder_start_token_id": 0,
+  "eos_token_id": 1,
+  "pad_token_id": 0,
+  "transformers_version": "4.26.1"
+}

model.safetensors → model-00001-of-00002.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:be8e6242388d93b51cd9719a0e32cfc17a2e804786570c795ba332197eccb915
-size 4608653848
+oid sha256:7c0c297f60bcf81c732cdfeae6e99e140272807eb52afd70356fc6fdfa94e5a8
+size 5332161168

model-00002-of-00002.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f3d425d3e8741ccbdd925446559a9bf317c2c91e328f2eee44924423b56e3a3d
+size 2184524400

model.safetensors.index.json

@@ -0,0 +1,783 @@
+{
+  "metadata": {
+    "total_size": 7516590080
+  },
+  "weight_map": {
+    "decoder.block.0.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.0.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.0.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.0.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.0.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.0.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.0.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.0.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.0.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.0.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.0.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.0.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.0.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.0.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.1.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.1.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.1.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.1.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.1.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.1.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.1.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.1.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.1.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.1.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.1.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.1.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.1.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.1.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.10.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.10.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.10.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.10.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.10.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.10.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.10.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.10.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.10.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.10.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.10.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.10.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.10.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.10.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.11.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.11.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.11.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.11.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.11.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.11.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.11.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.11.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.11.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.11.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.11.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.11.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.11.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.11.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.12.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.12.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.12.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.12.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.12.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.12.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.12.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.12.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.12.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.12.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.12.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.12.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.12.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.12.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.13.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.13.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.13.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.13.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.13.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.13.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.13.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.13.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.13.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.13.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.13.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.13.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.13.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.13.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.14.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.14.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.14.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.14.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.14.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.14.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.14.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.14.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.14.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.14.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.14.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.14.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.14.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.14.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.15.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.15.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.15.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.15.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.15.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.15.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.15.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.15.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.15.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.15.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.15.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.15.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.15.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.15.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.16.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.16.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.16.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.16.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.16.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.16.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.16.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.16.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.16.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.16.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.16.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.16.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.16.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.16.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.17.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.17.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.17.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.17.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.17.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.17.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.17.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.17.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.17.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.17.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.17.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.17.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.17.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.17.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.18.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.18.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.18.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.18.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.18.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.18.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.18.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.18.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.18.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.18.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.18.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.18.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.18.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.18.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.19.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.19.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.19.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.19.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.19.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.19.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.19.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.19.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.19.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.19.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.19.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.19.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.19.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.19.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.2.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.2.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.2.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.2.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.2.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.2.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.2.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.2.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.2.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.2.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.2.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.2.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.2.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.2.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.20.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.20.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.20.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.20.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.20.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.20.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.20.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.20.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.20.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.20.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.20.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.20.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.20.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.20.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.21.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.21.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.21.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.21.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.21.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.21.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.21.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.21.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.21.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.21.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.21.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.21.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.21.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.21.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.22.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.22.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.22.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.22.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.22.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.22.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.22.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.22.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.22.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.22.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.22.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.22.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.22.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.22.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.23.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.23.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.23.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.23.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.23.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.23.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.23.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.23.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.23.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.23.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.23.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.23.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.23.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.23.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.3.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.3.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.3.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.3.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.3.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.3.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.3.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.3.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.3.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.3.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.3.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.3.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.3.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.3.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.4.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.4.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.4.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.4.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.4.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.4.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.4.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.4.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.4.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.4.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.4.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.4.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.4.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.4.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.5.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.5.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.5.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.5.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.5.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.5.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.5.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.5.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.5.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.5.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.5.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.5.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.5.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.5.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.6.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.6.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.6.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.6.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.6.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.6.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.6.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.6.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.6.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.6.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.6.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.6.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.6.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.6.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.7.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.7.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.7.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.7.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.7.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.7.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.7.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.7.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.7.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.7.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.7.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.7.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.7.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.7.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.8.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.8.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.8.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.8.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.8.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.8.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.8.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.8.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.8.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.8.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.8.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.8.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.8.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.8.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.9.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.9.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.9.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.9.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.9.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.9.layer.1.EncDecAttention.k.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.9.layer.1.EncDecAttention.o.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.9.layer.1.EncDecAttention.q.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.9.layer.1.EncDecAttention.v.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.9.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.9.layer.2.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.9.layer.2.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.9.layer.2.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "decoder.block.9.layer.2.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "decoder.embed_tokens.weight": "model-00001-of-00002.safetensors",
+    "decoder.final_layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.0.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.0.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.0.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.0.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.0.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.0.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.0.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.0.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.0.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.1.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.1.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.1.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.1.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.1.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.1.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.1.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.1.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.1.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.10.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.10.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.10.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.10.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.10.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.10.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.10.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.10.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.10.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.11.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.11.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.11.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.11.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.11.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.11.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.11.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.11.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.11.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.12.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.12.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.12.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.12.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.12.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.12.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.12.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.12.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.12.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.13.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.13.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.13.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.13.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.13.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.13.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.13.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.13.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.13.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.14.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.14.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.14.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.14.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.14.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.14.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.14.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.14.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.14.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.15.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.15.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.15.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.15.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.15.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.15.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.15.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.15.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.15.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.16.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.16.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.16.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.16.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.16.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.16.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.16.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.16.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.16.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.17.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.17.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.17.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.17.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.17.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.17.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.17.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.17.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.17.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.18.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.18.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.18.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.18.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.18.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.18.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.18.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.18.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.18.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.19.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.19.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.19.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.19.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.19.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.19.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.19.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.19.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.19.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.2.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.2.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.2.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.2.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.2.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.2.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.2.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.2.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.2.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.20.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.20.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.20.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.20.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.20.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.20.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.20.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.20.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.20.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.21.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.21.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.21.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.21.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.21.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.21.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.21.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.21.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.21.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.22.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.22.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.22.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.22.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.22.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.22.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.22.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.22.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.22.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.23.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.23.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.23.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.23.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.23.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.23.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.23.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.23.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.23.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.24.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.24.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.24.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.24.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.24.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.24.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.24.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.24.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.24.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.25.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.25.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.25.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.25.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.25.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.25.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.25.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.25.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.25.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.26.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.26.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.26.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.26.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.26.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.26.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.26.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.26.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.26.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.27.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.27.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.27.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.27.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.27.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.27.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.27.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.27.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.27.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.28.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.28.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.28.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.28.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.28.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.28.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.28.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.28.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.28.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.29.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.29.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.29.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.29.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.29.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.29.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.29.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.29.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.29.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.3.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.3.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.3.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.3.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.3.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.3.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.3.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.3.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.3.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.30.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.30.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.30.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.30.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.30.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.30.layer.1.DenseReluDense.wi_0.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.30.layer.1.DenseReluDense.wi_1.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.30.layer.1.DenseReluDense.wo.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.30.layer.1.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.31.layer.0.SelfAttention.k.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.31.layer.0.SelfAttention.o.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.31.layer.0.SelfAttention.q.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.31.layer.0.SelfAttention.v.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.31.layer.0.layer_norm.weight": "model-00001-of-00002.safetensors",
+    "encoder.block.31.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.31.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.31.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.31.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.32.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.32.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.32.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.32.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.32.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.32.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.32.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.32.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.32.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.33.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.33.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.33.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.33.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.33.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.33.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.33.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.33.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.33.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.34.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.34.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.34.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.34.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.34.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.34.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.34.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.34.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.34.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.35.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.35.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.35.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.35.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.35.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.35.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.35.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.35.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.35.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.36.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.36.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.36.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.36.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.36.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.36.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.36.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.36.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.36.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.37.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.37.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.37.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.37.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.37.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.37.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.37.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.37.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.37.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.38.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.38.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.38.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.38.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.38.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.38.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.38.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.38.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.38.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.39.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.39.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.39.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.39.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.39.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.39.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.39.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.39.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.39.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.4.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.4.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.4.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.4.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.4.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.4.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.4.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.4.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.4.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.40.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.40.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.40.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.40.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.40.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.40.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.40.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.40.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.40.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.41.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.41.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.41.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.41.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.41.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.41.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.41.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.41.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.41.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.42.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.42.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.42.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.42.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.42.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.42.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.42.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.42.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.42.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.43.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.43.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.43.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.43.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.43.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.43.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.43.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.43.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.43.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.44.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.44.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.44.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.44.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.44.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.44.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.44.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.44.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.44.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.45.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.45.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.45.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.45.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.45.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.45.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.45.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.45.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.45.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.46.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.46.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.46.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.46.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.46.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.46.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.46.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.46.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.46.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.47.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.47.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.47.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.47.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.47.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.47.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.47.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.47.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.47.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.5.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.5.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.5.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.5.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.5.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.5.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.5.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.5.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.5.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.6.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.6.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.6.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.6.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.6.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.6.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.6.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.6.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.6.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.7.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.7.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.7.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.7.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.7.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.7.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.7.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.7.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.7.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.8.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.8.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.8.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.8.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.8.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.8.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.8.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.8.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.8.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.9.layer.0.SelfAttention.k.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.9.layer.0.SelfAttention.o.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.9.layer.0.SelfAttention.q.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.9.layer.0.SelfAttention.v.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.9.layer.0.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.9.layer.1.DenseReluDense.wi_0.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.9.layer.1.DenseReluDense.wi_1.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.9.layer.1.DenseReluDense.wo.weight": "model-00002-of-00002.safetensors",
+    "encoder.block.9.layer.1.layer_norm.weight": "model-00002-of-00002.safetensors",
+    "encoder.embed_tokens.weight": "model-00002-of-00002.safetensors",
+    "encoder.final_layer_norm.weight": "model-00002-of-00002.safetensors",
+    "lm_head.weight": "model-00002-of-00002.safetensors",
+    "shared.weight": "model-00002-of-00002.safetensors"
+  }
+}

modeling_fastplms.py

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

provenance.json

@@ -0,0 +1,165 @@
+{
+  "architecture": "ANKH",
+  "artifact_checkpoint": {
+    "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"
+    },
+    "repo_id": "ElnaggarLab/ankh2-ext2",
+    "revision": "aa9b9fa72288c47d9f618ce80c011e24b54e17a8",
+    "unresolved_files": []
+  },
+  "artifact_source": "official",
+  "attestations": {
+    "complete_artifact": {
+      "redistributable": true,
+      "release_tool_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+      "release_tool_sha256": "1459b5d7d13d9b07bd97b3eee764f2ce73623e15e32d07ddf6825c2a9509afb9",
+      "runtime_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+      "scope": "weights+runtime",
+      "weights_license_status": "resolved",
+      "weights_revision": "aa9b9fa72288c47d9f618ce80c011e24b54e17a8"
+    },
+    "runtime_update": {
+      "path": "runtime-attestation.json",
+      "redistributable": true,
+      "release_tool_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+      "release_tool_sha256": "1459b5d7d13d9b07bd97b3eee764f2ce73623e15e32d07ddf6825c2a9509afb9",
+      "scope": "runtime-only",
+      "weights_license_status": "resolved",
+      "weights_repo_id": "Synthyra/ANKH2_large",
+      "weights_revision": "392de5ed52bbfd73b45f545e378aaebcff096d0e"
+    }
+  },
+  "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"
+  },
+  "bf16_execution": "static_parameters",
+  "canonical_weights": {
+    "format": "safetensors",
+    "index": "model.safetensors.index.json",
+    "index_digest": "sha256:eb9f1a3b03bd5fc5dd912aa4244181ed09e19b99a57c2e1d95bdf952e510cea3",
+    "max_shard_bytes": 5368709120,
+    "shards": {
+      "model-00001-of-00002.safetensors": "sha256:7c0c297f60bcf81c732cdfeae6e99e140272807eb52afd70356fc6fdfa94e5a8",
+      "model-00002-of-00002.safetensors": "sha256:f3d425d3e8741ccbdd925446559a9bf317c2c91e328f2eee44924423b56e3a3d"
+    },
+    "source_schema": "official",
+    "state_digest": {
+      "algorithm": "sha256",
+      "schema_version": 1,
+      "sha256": "597c4fe2fa8711f11a25317905f1d62fa92905e55fdd5c0a79614cd9c9d2bca3"
+    },
+    "state_transform": "ankh_t5_to_fastplms_v1",
+    "tensor_count": 776,
+    "total_size": 7516590080
+  },
+  "checkpoint_license": "CC-BY-NC-SA-4.0",
+  "conversion": {
+    "id": "ankh_t5_to_fastplms_v1",
+    "record": "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."
+  },
+  "conversion_equality_attestation": {
+    "attestation_sha256": "4b5acf56ebc107b0177b873116d4a498c7da56bcf41cecce250310ad302fa075",
+    "canonical_state": {
+      "algorithm": "sha256",
+      "schema_version": 1,
+      "sha256": "597c4fe2fa8711f11a25317905f1d62fa92905e55fdd5c0a79614cd9c9d2bca3"
+    },
+    "conversion_record_sha256": "34015dc5c07465d93b54349e6ac3b68179101bb231d7bac41af308287d69dc7d",
+    "model_id": "ankh2_large",
+    "schema_version": 1,
+    "source_checkpoint": {
+      "identity_sha256": "ca31eec9942e836d0b557a94d3b50e642af3738627a348d6526a6968b8646f8c",
+      "repo_id": "ElnaggarLab/ankh2-ext2",
+      "revision": "aa9b9fa72288c47d9f618ce80c011e24b54e17a8"
+    },
+    "state_transform": "ankh_t5_to_fastplms_v1"
+  },
+  "fast_checkpoint": {
+    "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"
+    },
+    "repo_id": "Synthyra/ANKH2_large",
+    "revision": "392de5ed52bbfd73b45f545e378aaebcff096d0e",
+    "unresolved_files": []
+  },
+  "fastplms_version": "1.0.0",
+  "generator": {
+    "name": "tools.artifacts.build",
+    "version": 3
+  },
+  "hub_license_metadata": {
+    "license": "cc-by-nc-sa-4.0"
+  },
+  "legal_files": {
+    "LICENSE": "sha256:2d2b50c7b1414bff1189a1db1f0cfb92e3e064b50f4c2b1019827b683e1b629a",
+    "THIRD_PARTY_NOTICES.md": "sha256:25704b3c76404696cae52e7fca13088d329f70f412687340351259e86cd62baa"
+  },
+  "model_id": "ankh2_large",
+  "official_checkpoint": {
+    "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"
+    },
+    "repo_id": "ElnaggarLab/ankh2-ext2",
+    "revision": "aa9b9fa72288c47d9f618ce80c011e24b54e17a8",
+    "unresolved_files": []
+  },
+  "oracle_assets": [],
+  "redistributable": true,
+  "release_tool_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "release_tool_sha256": "1459b5d7d13d9b07bd97b3eee764f2ce73623e15e32d07ddf6825c2a9509afb9",
+  "runtime_assets": [],
+  "runtime_bundle_sha256": "17b8f83a33e63d941e3edfb8db2d8381286046b572f94194e576ca394d997597",
+  "runtime_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "schema_version": 4,
+  "source_tree_sha256": "5eb598fd991b5447e3bb6a18c92bd2300be1d94a1e169504725bc56127c1d6bc",
+  "state_transform": "ankh_t5_to_fastplms_v1",
+  "tokenizer_auto_map": null,
+  "tokenizer_checkpoint": {
+    "files": {
+      "special_tokens_map.json": "git-sha1:55b145827029ae9672e50d4bb368540daacce791",
+      "tokenizer.json": "git-sha1:212c5ef08819fa2463c6289ba4ef7db30e715c0a",
+      "tokenizer_config.json": "git-sha1:854e5db75dae8b1e9dd39c5bae80dae5508b3e25"
+    },
+    "repo_id": "ElnaggarLab/ankh2-ext2",
+    "revision": "aa9b9fa72288c47d9f618ce80c011e24b54e17a8",
+    "unresolved_files": []
+  },
+  "tokenizer_class": null,
+  "upstreams": [
+    {
+      "canonical_license_files": {
+        "LICENSE.md": "sha256:cd041d7f9f52936e8824ac3f754e9c67410763205fc8a7020ba74fc8b6edc088"
+      },
+      "distribution_files": {
+        "LICENSE.md": "sha256:cd041d7f9f52936e8824ac3f754e9c67410763205fc8a7020ba74fc8b6edc088"
+      },
+      "id": "ankh",
+      "license": "CC-BY-NC-SA-4.0",
+      "path": "vendor/upstream/ankh",
+      "revision": "02b4e25ce5389b9e771c9df6e546c62af1216f8e",
+      "url": "https://github.com/agemagician/Ankh.git"
+    }
+  ],
+  "weights_license_status": "resolved",
+  "weights_revision": "aa9b9fa72288c47d9f618ce80c011e24b54e17a8"
+}

runtime-attestation.json

@@ -0,0 +1,45 @@
+{
+  "files": {
+    "LICENSES/FastPLMs-Apache-2.0.txt": "sha256:2d2b50c7b1414bff1189a1db1f0cfb92e3e064b50f4c2b1019827b683e1b629a",
+    "LICENSES/ankh/LICENSE.md": "sha256:cd041d7f9f52936e8824ac3f754e9c67410763205fc8a7020ba74fc8b6edc088",
+    "README.md": "sha256:d34607e87e0df11a9d8cea1e8dd541d1fe79761f90cfc2d95c9cb5135e0ccbbc",
+    "THIRD_PARTY_NOTICES.md": "sha256:25704b3c76404696cae52e7fca13088d329f70f412687340351259e86cd62baa",
+    "config.json": "sha256:fe9df500b8e73df29eff467f41a92d9bf1489ba7ac8c151f237f700f41b1b7a1",
+    "fastplms/__init__.py": "sha256:4fb3196022ca8ec699d59d09bdbc5f0184195552b773698ab9b061fe3cd7df12",
+    "fastplms/attention/__init__.py": "sha256:f60b9fecfb4bcb37a4e7c26dc2f752b9035f9cbad627b4a84213f3a92ec88f7d",
+    "fastplms/attention/_core.py": "sha256:8f7ec5b65bd8b6c6fa4951d50d1c0e499abf03ae00914794b51fc410201e3e33",
+    "fastplms/attention/_kernel_lock.py": "sha256:85d8521a2af5f94fad3948af3814db0c866c414ee4d43df797b9bd6f980e947b",
+    "fastplms/attention/interfaces.py": "sha256:1c6f06a8e411e0f9bf6d230522205c93ae46ea58864006fbb892aa05e5ca5749",
+    "fastplms/embeddings/__init__.py": "sha256:47ff8cdf682d44037dd9edab133e2e60675d60786bd5cf0bffd1998f31985555",
+    "fastplms/embeddings/pooling.py": "sha256:a140266ed6b1cc344c8507edc5c6c4f2dce464c3db70ba4b16c7ac2ba2fad96e",
+    "fastplms/embeddings/runner.py": "sha256:23ee4727a918d6d331f7a0f89b823d149f1a791f0c5586e3496d7b6eb2ce97e0",
+    "fastplms/embeddings/storage.py": "sha256:3fbe2bab75092e5a4cadf4d27e4752181d597469a65a55db085ceef808ed418e",
+    "fastplms/embeddings/types.py": "sha256:119718a20989d1ae5a60fabc0f5e98bdc172c5163b04db3d4554ac3956b30e52",
+    "fastplms/models.toml": "sha256:05a8399f084a5babb5f0916cee7e564c4030767f3ff0230c4f46e539209847d1",
+    "fastplms/models/__init__.py": "sha256:5e48c2cb3877aa6f42f3b5411d53b16bba2e32827bbde634f47f174c5cb36f86",
+    "fastplms/models/ankh/__init__.py": "sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
+    "fastplms/models/ankh/modeling_ankh.py": "sha256:b114c73e4aaedc567e3351b425ddddacf35661f09b3414cf59b6db0b712c82c3",
+    "fastplms/models/ttt.py": "sha256:a0df4e98b02120d423e3c7ca9b866a8d0e3748b9076042a0102a838a11aed046",
+    "fastplms/registry.py": "sha256:afca271911b651a882345b74a58366494a1d784e4a48c8683a87f5508f4ba16e",
+    "fastplms/runtime.py": "sha256:110018646d6f248cedab140a030c3065e1b062b61f6aff659c231e538614bc01",
+    "fastplms_bundle.py": "sha256:ebf70cf52fb0a07025edf06bd50d2e945d9ed09fe12e6d3ae3503001ec21c9f9",
+    "modeling_fastplms.py": "sha256:df7f6d96802eff127bc8270cde5fe00c8f5e77921c1cc893021fb9f75a3e4c77",
+    "special_tokens_map.json": "sha256:c8995d2f8037fe3a8cfdef30475365e1c314c417880b75abaf8296e3c05d42d6",
+    "tokenizer.json": "sha256:b4533f607d9fd665f2d9d94b0cf71870a6fc2fc2ae7cbd516d0e43a9efb406fd",
+    "tokenizer_config.json": "sha256:dba1f9315ded007fc2e5bcb2ed4bfb6fff81c9becdc97b73a5713ba5ddc01afc"
+  },
+  "model_id": "ankh2_large",
+  "redistributable": true,
+  "release_tool_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "release_tool_sha256": "1459b5d7d13d9b07bd97b3eee764f2ce73623e15e32d07ddf6825c2a9509afb9",
+  "runtime_bundle_sha256": "17b8f83a33e63d941e3edfb8db2d8381286046b572f94194e576ca394d997597",
+  "runtime_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "schema_version": 2,
+  "scope": "runtime-only",
+  "source_tree_sha256": "5eb598fd991b5447e3bb6a18c92bd2300be1d94a1e169504725bc56127c1d6bc",
+  "weights": {
+    "repo_id": "Synthyra/ANKH2_large",
+    "revision": "392de5ed52bbfd73b45f545e378aaebcff096d0e"
+  },
+  "weights_license_status": "resolved"
+}

special_tokens_map.json

@@ -117,32 +117,8 @@
     "<extra_id_114>",
     "<extra_id_115>"
   ],
-  "eos_token": {
-    "content": "</s>",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  },
-  "mask_token": {
-    "content": "<extra_id_0>",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  },
-  "pad_token": {
-    "content": "<pad>",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  },
-  "unk_token": {
-    "content": "<unk>",
-    "lstrip": false,
-    "normalized": false,
-    "rstrip": false,
-    "single_word": false
-  }
+  "eos_token": "</s>",
+  "mask_token": "<extra_id_0>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>"
 }

tokenizer.json

@@ -1134,8 +1134,7 @@
   "decoder": {
     "type": "Metaspace",
     "replacement": "▁",
-    "prepend_scheme": "always",
-    "split": true
+    "add_prefix_space": true
   },
   "model": {
     "type": "Unigram",
@@ -1717,7 +1716,6 @@
         "<extra_id_0>",
         0.0
       ]
-    ],
-    "byte_fallback": false
+    ]
   }
 }

tokenizer_config.json

@@ -1,959 +1,4 @@
 {
-  "add_prefix_space": null,
-  "added_tokens_decoder": {
-    "0": {
-      "content": "<pad>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "1": {
-      "content": "</s>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "2": {
-      "content": "<unk>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "28": {
-      "content": "<extra_id_115>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "29": {
-      "content": "<extra_id_114>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "30": {
-      "content": "<extra_id_113>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "31": {
-      "content": "<extra_id_112>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "32": {
-      "content": "<extra_id_111>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "33": {
-      "content": "<extra_id_110>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "34": {
-      "content": "<extra_id_109>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "35": {
-      "content": "<extra_id_108>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "36": {
-      "content": "<extra_id_107>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "37": {
-      "content": "<extra_id_106>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "38": {
-      "content": "<extra_id_105>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "39": {
-      "content": "<extra_id_104>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "40": {
-      "content": "<extra_id_103>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "41": {
-      "content": "<extra_id_102>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "42": {
-      "content": "<extra_id_101>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "43": {
-      "content": "<extra_id_100>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "44": {
-      "content": "<extra_id_99>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "45": {
-      "content": "<extra_id_98>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "46": {
-      "content": "<extra_id_97>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "47": {
-      "content": "<extra_id_96>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "48": {
-      "content": "<extra_id_95>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "49": {
-      "content": "<extra_id_94>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "50": {
-      "content": "<extra_id_93>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "51": {
-      "content": "<extra_id_92>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "52": {
-      "content": "<extra_id_91>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "53": {
-      "content": "<extra_id_90>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "54": {
-      "content": "<extra_id_89>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "55": {
-      "content": "<extra_id_88>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "56": {
-      "content": "<extra_id_87>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "57": {
-      "content": "<extra_id_86>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "58": {
-      "content": "<extra_id_85>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "59": {
-      "content": "<extra_id_84>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "60": {
-      "content": "<extra_id_83>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "61": {
-      "content": "<extra_id_82>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "62": {
-      "content": "<extra_id_81>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "63": {
-      "content": "<extra_id_80>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "64": {
-      "content": "<extra_id_79>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "65": {
-      "content": "<extra_id_78>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "66": {
-      "content": "<extra_id_77>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "67": {
-      "content": "<extra_id_76>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "68": {
-      "content": "<extra_id_75>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "69": {
-      "content": "<extra_id_74>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "70": {
-      "content": "<extra_id_73>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "71": {
-      "content": "<extra_id_72>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "72": {
-      "content": "<extra_id_71>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "73": {
-      "content": "<extra_id_70>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "74": {
-      "content": "<extra_id_69>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "75": {
-      "content": "<extra_id_68>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "76": {
-      "content": "<extra_id_67>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "77": {
-      "content": "<extra_id_66>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "78": {
-      "content": "<extra_id_65>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "79": {
-      "content": "<extra_id_64>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "80": {
-      "content": "<extra_id_63>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "81": {
-      "content": "<extra_id_62>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "82": {
-      "content": "<extra_id_61>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "83": {
-      "content": "<extra_id_60>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "84": {
-      "content": "<extra_id_59>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "85": {
-      "content": "<extra_id_58>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "86": {
-      "content": "<extra_id_57>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "87": {
-      "content": "<extra_id_56>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "88": {
-      "content": "<extra_id_55>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "89": {
-      "content": "<extra_id_54>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "90": {
-      "content": "<extra_id_53>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "91": {
-      "content": "<extra_id_52>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "92": {
-      "content": "<extra_id_51>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "93": {
-      "content": "<extra_id_50>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "94": {
-      "content": "<extra_id_49>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "95": {
-      "content": "<extra_id_48>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "96": {
-      "content": "<extra_id_47>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "97": {
-      "content": "<extra_id_46>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "98": {
-      "content": "<extra_id_45>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "99": {
-      "content": "<extra_id_44>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "100": {
-      "content": "<extra_id_43>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "101": {
-      "content": "<extra_id_42>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "102": {
-      "content": "<extra_id_41>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "103": {
-      "content": "<extra_id_40>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "104": {
-      "content": "<extra_id_39>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "105": {
-      "content": "<extra_id_38>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "106": {
-      "content": "<extra_id_37>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "107": {
-      "content": "<extra_id_36>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "108": {
-      "content": "<extra_id_35>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "109": {
-      "content": "<extra_id_34>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "110": {
-      "content": "<extra_id_33>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "111": {
-      "content": "<extra_id_32>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "112": {
-      "content": "<extra_id_31>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "113": {
-      "content": "<extra_id_30>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "114": {
-      "content": "<extra_id_29>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "115": {
-      "content": "<extra_id_28>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "116": {
-      "content": "<extra_id_27>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "117": {
-      "content": "<extra_id_26>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "118": {
-      "content": "<extra_id_25>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "119": {
-      "content": "<extra_id_24>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "120": {
-      "content": "<extra_id_23>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "121": {
-      "content": "<extra_id_22>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "122": {
-      "content": "<extra_id_21>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "123": {
-      "content": "<extra_id_20>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "124": {
-      "content": "<extra_id_19>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "125": {
-      "content": "<extra_id_18>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "126": {
-      "content": "<extra_id_17>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "127": {
-      "content": "<extra_id_16>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "128": {
-      "content": "<extra_id_15>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "129": {
-      "content": "<extra_id_14>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "130": {
-      "content": "<extra_id_13>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "131": {
-      "content": "<extra_id_12>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "132": {
-      "content": "<extra_id_11>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "133": {
-      "content": "<extra_id_10>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "134": {
-      "content": "<extra_id_9>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "135": {
-      "content": "<extra_id_8>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "136": {
-      "content": "<extra_id_7>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "137": {
-      "content": "<extra_id_6>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "138": {
-      "content": "<extra_id_5>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "139": {
-      "content": "<extra_id_4>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "140": {
-      "content": "<extra_id_3>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "141": {
-      "content": "<extra_id_2>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "142": {
-      "content": "<extra_id_1>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    },
-    "143": {
-      "content": "<extra_id_0>",
-      "lstrip": false,
-      "normalized": false,
-      "rstrip": false,
-      "single_word": false,
-      "special": true
-    }
-  },
   "additional_special_tokens": [
     "<extra_id_0>",
     "<extra_id_1>",
@@ -1072,13 +117,12 @@
     "<extra_id_114>",
     "<extra_id_115>"
   ],
-  "clean_up_tokenization_spaces": false,
   "eos_token": "</s>",
   "extra_ids": 116,
-  "extra_special_tokens": {},
-  "mask_token": "<extra_id_0>",
   "model_max_length": 1000000000000000019884624838656,
+  "name_or_path": "agemagician/protx-large-1gspan-partreconst-15mlmp-encl48-decl24-ramd128-ranb64-dmodel1536-silu-seq768-tpuv4",
   "pad_token": "<pad>",
-  "tokenizer_class": "T5TokenizerFast",
+  "special_tokens_map_file": "tokenizer/protein_t5_fast_tokenizer_v4/special_tokens_map.json",
+  "tokenizer_class": "T5Tokenizer",
   "unk_token": "<unk>"
 }