Update FastPLMs runtime files

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

LICENSES/FastPLMs-Apache-2.0.txt

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

LICENSES/e1/ATTRIBUTION

@@ -0,0 +1,34 @@
+Profluent-E1 Attribution Guidelines
+-----------------------------------
+
+These Profluent-E1 Attribution Guidelines (the “Guidelines”) set forth certain terms, conditions and restrictions relating to how You can use, share and distribute Profluent-E1 and any Derivative Works thereof in a manner that ensures transparent acknowledgement of Profluent-E1. The terms, conditions and restrictions set forth in these Guidelines are in addition to those set forth in the Profluent-E1 Clickthrough License Agreement, available at https://github.com/Profluent-AI/E1/blob/main/LICENSE, as may be updated or amended from time to time (the “Agreement”) and in the event of any conflict between any provision of these Guidelines and any provision of the Agreement, the provision that is more protective of Profluent Bio, Inc. (“Profluent”) and Profluent-E1 shall control. Any capitalized terms used but not defined herein shall have the meanings set forth in the Agreement. Profluent may update, modify or amend these Guidelines from time to time and Profluent will use reasonable efforts to provide You with notice of any material changes that may negatively impact Your use of Profluent-E1 by posting an updated version of these Guidelines, for example. Any violation of these Guidelines may result in the suspension or termination of Your access to and use of Profluent-E1.
+
+1. Attribution Requirements. You may use and redistribute Profluent-E1 and Derivative Works in Source and Object forms, with or without modification, in accordance with the terms of the Agreement provided that the following conditions are met:
+
+    1.1 You must include a prominently displayed attribution to “Profluent-E1” in any distribution, documentation or public use of Profluent-E1 or any Derivative Works built, developed, improved, or modified through the use of Profluent-E1 by any Commercial Entity, including, without limitation, in any publication, report, white paper, article, product documentation, slide deck or other media that incorporates Profluent-E1 or any Derivative Works; and
+
+    1.2 You must identify any drug, drug candidate, medication or pharmaceutical of any kind and any molecular or biological target, hit or lead identification that is created, identified, discovered, developed, modified or improved using Profluent-E1 as “Built with Profluent-E1,” including, without limitation, in public disclosures to the United States Food and Drug Administration and similar regulatory authorities.
+
+The term “Commercial Entity” as used in this Section 1 shall mean any entity engaged in any activity intended for or directed toward commercial advantage or monetary compensation, including, without limitation, the development of any product or service intended to be sold or made available for a fee, expressly excluding any university, non-profit organization, research institute, educational or government body and any not-for-profit Legal Entity.
+
+2. No Endorsement. You may not use or allow the use of the Profluent name nor any Profluent trade name, trademark, service mark or product name to endorse or promote, or imply that Profluent endorses or promotes, any products or other offerings derived from Profluent-E1 without Profluent's specific prior written permission.
+
+3. Format of Derivative Works. All attribution notices should be easy to find and easy to understand. The attribution requirements set forth in these Guidelines applies to both textual and visual materials where Profluent-E1 contributes materially to the resulting product or output. Any uses, distributions or documentation in Source form must incorporate a prominent display of text that reads “Profluent-E1” in verifiable form, while redistributions in Object form must be accompanied by a prominent display of “Profluent-E1” text in any documentation and each time the resulting executable program or a program dependent thereon is launched, You must also launch a prominent display (e.g., splash screen or banner text) of such text.
+
+4. Placement and Prominence. any entity engaged in any activity intended for or directed toward commercial advantage or monetary compensation, including, without limitation, the development of any product or service intended to be sold or made available for a fee, expressly excluding any university, non-profit organization, research institute, educational or government body and any not-for-profit Legal Entity. The attribution must appear within the main body of any publication, report, white paper or article. For product documentation or software tool or interface references, the attribution must be listed on the title page or in the introductory section, near other third-party license notices. In slide decks or conference presentations, the attribution must appear on a title slide, or a slide within the first three slides of the presentation. For webpages, software interfaces, or interactive demos, attribution must be visible on-screen in a footer, “About” section, or a prominently accessible credits area without requiring extra actions like scrolling or menu navigation.
+
+5. Formatting. You must use the exact “Profluent-E1” model name as provided by Profluent, and the font size must be at least equal to the main body text or at least 75% of the largest text on the page where the attribution is incorporated (whichever is greater). The text must be visually legible and in a color that contrasts sufficiently with its background. Graphic or visual logos provided by the Profluent may be used in place of plain-text attribution if approved, but they must remain clearly readable at normal viewing sizes.
+
+6. Illustrative Examples. The following are examples of how attribution to Profluent-E1 may be implemented in accordance with the requirements of these guidelines:
+
+    6.1 Example 1: A research article that incorporates data derived from outputs of Profluent-E1 includes reference to “Profluent-E1” in or beneath the paragraph of text first introducing these outputs.
+
+    6.2 Example 2: A company releases a model fine-tuned from Profluent-E1 and the website's footer includes “Built with Profluent-E1.”
+
+    6.3 Example 3: A presentation at a conference includes “This work incorporates works derived from use of Profluent-E1” on the title slide.
+
+    6.4 Example 4: An open-source repository “README” file includes a visible statement near the top: “This project builds upon Profluent-E1 under the terms of the Profluent-E1 Clickthrough License Agreement.”
+
+7. Non-Obscuration. The attribution must not be hidden or disguised by other visual or textual elements. It should remain readable at standard display resolutions and must not be abbreviated, truncated or replaced with unofficial branding.
+
+8. Exceptions. Written exceptions to these attribution requirements may be granted by the Licensor upon request, provided that alternative acknowledgment conveys equivalent visibility and transparency.

LICENSES/e1/Apache-2.0.txt

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

LICENSES/e1/BSD-3-Clause.txt

@@ -0,0 +1,29 @@
+BSD 3-Clause License
+
+Copyright (c) 2022, Tri Dao.
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice,
+   this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from this
+   software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
+ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
+ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

LICENSES/e1/LICENSE

@@ -0,0 +1,86 @@
+Profluent-E1 Clickthrough License Agreement
+-------------------------------------------
+
+Please read this Profluent-E1 Clickthrough License Agreement (the “Agreement”) carefully before using Profluent-E1 (as defined below), which is offered by Profluent Bio Inc. (“Profluent”).
+
+By downloading Profluent-E1, or otherwise using Profluent-E1 in any manner, You agree that You have read and agree to be bound by the terms of this Agreement. If You are accessing Profluent-E1 on behalf of an organization or entity, You represent and warrant that You are authorized to enter into this Agreement on that organization's or entity's behalf and bind them to the terms of this Agreement (in which case, the references to “You” and “Your” in this Agreement, except for in this sentence, refer to that organization or entity). Use of Profluent-E1 and all other Profluent-E1 IP is expressly conditioned upon Your assent to all terms of this Agreement, including the Attribution Guidelines incorporated herein, to the exclusion of all other terms.
+
+1. Definitions.
+
+    1.1 “AAA Rules” shall mean the Commercial Arbitration Rules and Mediation Procedures of the American Arbitration Association (“AAA”).
+
+    1.2 “Claim” shall mean any claim (including any tort claim), cause of action, and/or dispute under, arising out of, or relating to this Agreement.
+
+    1.3 “Contribution” shall mean any work of authorship, including the original version of Profluent-E1 and any modifications or additions to that Profluent-E1 or Derivative Works thereof, that is intentionally submitted to Profluent for inclusion in Profluent-E1 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 Profluent 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, Profluent for the purpose of discussing and improving Profluent-E1, but excluding communication that is conspicuously marked or otherwise designated in writing by the copyright owner as "Not a Contribution."
+
+    1.4 “Contributor” shall mean Profluent and any individual or Legal Entity on behalf of whom a Contribution has been received by Profluent and subsequently incorporated within Profluent-E1.
+
+    1.5 “Derivative Works” shall mean any work, whether in Source or Object form, that is based on (or derived from) Profluent-E1 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 Agreement, Derivative Works shall not include works that remain separable from, or merely link (or bind by name) to the interfaces of, Profluent-E1 and Derivative Works thereof.
+
+    1.6 “GitHub Page” shall mean the page made available at https://github.com/Profluent-AI/E1, as may be updated and amended from time to time.
+
+    1.7 “Hugging Face Pages” shall mean the pages made available at  https://huggingface.co/Profluent-Bio/E1-600m, https://huggingface.co/Profluent-Bio/E1-300m and https://huggingface.co/Profluent-Bio/E1-150m, each as may be updated and amended from time to time.
+
+    1.8 “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.
+
+    1.9 “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.
+
+    1.10 “Profluent-E1” shall mean the Profluent-E1 Model Code, Profluent-E1 Model Weights and all software, algorithms, machine-learning model code, inference-enabling code, training-enabling code, fine-tuning enabling code and other elements of the foregoing distributed or made available to You by Profluent on the GitHub Page, each as may be updated and amended from time to time, whether in Source or Object form.
+
+    1.11 “Profluent-E1 Model Code” shall mean the code and data for Profluent-E1 made available to You at https://github.com/Profluent-AI/E1, as may be updated and amended from time to time.
+
+    1.12 “Profluent-E1 Model Weights” shall mean the trained model weights for Profluent-E1 made available to You on one or more of the Hugging Face Pages, as may be updated and amended from time to time, including all model weights that are directly or indirectly accessed or copied from the Hugging Face Pages by You or any third party.
+
+    1.13 “Source” form shall mean the preferred form for making modifications, including but not limited to software source code, documentation source, and configuration files.
+
+    1.14 “You” or “Your” shall mean an individual entering into this Agreement or the organization or Legal Entity on whose behalf such individual is entering into this Agreement.
+
+2. Grant of Copyright License. Subject to the terms and conditions of this Agreement, 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 Profluent-E1 and such Derivative Works in Source or Object form.
+
+3. Grant of Patent License. Subject to the terms and conditions of this Agreement, 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 Profluent-E1, 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 Profluent-E1 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 Profluent-E1 or a Contribution incorporated within Profluent-E1 constitutes direct or contributory patent infringement, then any patent licenses granted to You under this Agreement for Profluent-E1 shall terminate as of the date such litigation is filed.
+
+4. Use of Profluent-E1 Model Code. Your access to and use of the Profluent-E1 Model Code separate and apart from Profluent-E1 and the Profluent-E1 Model Weights is subject to the Apache License, Version 2.0 made available at https://www.apache.org/licenses/LICENSE-2.0, as may be updated and amended from time to time.
+
+5. Redistribution. You may reproduce and distribute copies of Profluent-E1 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 Profluent-E1 or Derivative Works a copy of this Agreement; 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 Profluent-E1, excluding those notices that do not pertain to any part of the Derivative Works; and
+
+    (d) If Profluent-E1 includes a "NOTICE" text file as part of its distribution or there is otherwise a “NOTICE” text file available on the GitHub Page, 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 this Agreement. You may add Your own attribution notices within Derivative Works that You distribute, alongside or as an addendum to the NOTICE text from Profluent-E1, provided that such additional attribution notices cannot be construed as modifying this Agreement; and
+
+    (e) You must comply with the Attribution Guidelines and provide attribution to Profluent-E1 in accordance therewith.
+
+    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 Profluent-E1 otherwise complies with the conditions stated in this License.
+
+6. Submission of Contributions. Unless You explicitly state otherwise, any Contribution intentionally submitted for inclusion in Profluent-E1 by You to Profluent shall be under the terms and conditions of this Agreement, 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 Profluent regarding such Contributions.
+
+7. Trademarks. This Agreement does not grant permission to use any Profluent trade names, trademarks, service marks, or product names, except as required for reasonable and customary use in describing the origin of Profluent-E1, reproducing the content of the NOTICE file and complying with the Attribution Guidelines.
+
+8. Disclaimer of Warranty. Unless required by applicable law or agreed to in writing, Profluent provides Profluent-E1 (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 Profluent-E1 and assume any risks associated with Your exercise of permissions under this Agreement. FOR THE AVOIDANCE OF DOUBT AND NOTWITHSTANDING ANYTHING TO THE CONTRARY, YOU ACKNOWLEDGE AND AGREE THAT PROFLUENT IS NOT RESPONSIBLE OR LIABLE FOR ANYTHING YOU BUILD, CREATE, DEVELOP OR DERIVE FROM YOUR USE OF PROFLUENT-E1, INCLUDING ANY DERIVATIVE WORKS.
+
+9. 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 Agreement or out of the use or inability to use Profluent-E1 (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.
+
+10. Accepting Warranty or Additional Liability. While redistributing Profluent-E1 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 Agreement. 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.
+
+11. Term and Termination. The term of this Agreement will commence upon Your acceptance of this Agreement or access to Profluent-E1 and will continue in full force and effect until terminated in accordance with the terms and conditions herein. Profluent may terminate this Agreement if You are in breach of any term or condition of this Agreement. Upon termination of this Agreement, you shall delete and cease use of Profluent-E1. All provisions of this Agreement which by their nature should survive termination shall survive termination, including, without limitation, ownership provisions, warranty disclaimers, indemnity obligations, limitations of liability and provisions regarding dispute resolution.
+
+12. General.
+
+    12.1 This Agreement constitutes the entire agreement between You and Profluent relating to the subject matter hereof and supersedes all proposals, understandings, or discussions, whether written or oral, relating to the subject matter of this Agreement and all past dealing or industry custom. The failure of either party to enforce its rights under this Agreement at any time for any period shall not be construed as a waiver of such rights.
+
+    12.2 Profluent may amend or modify this Agreement from time to time and will use reasonable efforts to provide You with notice of any material changes that may negatively impact Your use of Profluent-E1 through the GitHub Page or through another means made available to You. No other changes, modifications or waivers to this Agreement will be effective unless in writing and signed by both parties.
+
+    12.3 You and Profluent are independent contractors, and nothing herein shall be deemed to constitute either party as the agent or representative of the other or both parties as joint venturers or partners for any purpose.
+
+    12.4 You shall comply with the U.S. Foreign Corrupt Practices Act and all applicable export laws, restrictions and regulations of the U.S. Department of Commerce, and any other applicable U.S. and foreign authority.
+
+    12.5 This Agreement and the rights and obligations herein may not be assigned or transferred, in whole or in part, by You without the prior written consent of Profluent. Any assignment in violation of this provision is void. Profluent may freely assign or transfer this Agreement, in whole or in part. This Agreement shall be binding upon, and inure to the benefit of, the successors and permitted assigns of the parties.
+
+    12.6 This Agreement shall be governed by and construed under the Federal Arbitration Act, applicable federal law, and the laws of the State of California and the United States without regard to conflicts of laws provisions thereof, and without regard to the Uniform Computer Information Transactions Act.
+
+    12.7 If You have a Claim, You agree to provide Profluent with at least sixty (60) days' written notice so that the parties may attempt to resolve the Claim internally before requesting arbitration. In case of any Claims or disputes between You and Profluent that cannot be resolved through informal internal discussions, You agree that, subject to the provisions of this section, any and all Claims will be settled by final and binding arbitration in accordance with the AAA Rules by an arbitrator selected pursuant to the AAA Rules; provided, however, that Claims arising out of or relating to Your violations of Profluent's intellectual property rights, including copyright infringement, patent infringement, trademark infringement, or efforts to use Profluent-E1 in unauthorized ways will not be subject to such obligation for settlement by final and binding arbitration, and such claims will instead be brought in the state and federal courts of Alameda County, California. The arbitrator's decision and award will be non-appealable and may be entered in, and will be enforceable in, any court of competent jurisdiction. The arbitration will take place in Alameda County, California. Each party will bear its own costs of arbitration unless the arbitrator directs otherwise. The arbitrator may not award relief or damages in excess of or contrary to what this Agreement provides, or order consolidation or arbitration on a class-wide or representative basis.
+
+    12.8 If any provision of this Agreement is held to be invalid, illegal or unenforceable in any respect, that provision shall be limited or eliminated to the minimum extent necessary so that this Agreement otherwise remains in full force and effect and enforceable.

LICENSES/e1/MODIFICATIONS.md

@@ -0,0 +1,27 @@
+# Profluent-E1 modified-file notice
+
+FastPLMs implements Profluent-E1 behavior against the pinned official source at
+revision `bfd2620a602248499f3d2583d85a7ecddf0b6e02`. The FastPLMs files listed
+below are modified or independently reorganized implementations of the
+corresponding E1 interfaces. They are not byte-for-byte copies of the upstream
+files.
+
+| FastPLMs file | Modification notice |
+|---|---|
+| `src/fastplms/models/e1/modeling_e1.py` | Reorganized for Transformers AutoClasses, shared attention selection, sequence and RAG preparation, task heads, and checkpoint-compatible loading. |
+| `src/fastplms/models/e1/attention.py` | Isolated the declared SDPA and Flex Attention paths and their masking contracts. |
+| `src/fastplms/models/e1/preparation.py` | Reorganized raw-sequence, boundary-token, and retrieval-context preparation. |
+| `src/fastplms/models/e1/cache.py` | Adapted the cache interface used by the reorganized Transformers implementation. |
+| `src/fastplms/models/e1/retrieval.py` | Adapted retrieval helpers and their FastPLMs model outputs. |
+| `src/fastplms/models/e1/__init__.py` | Added FastPLMs package exports. |
+| `tools/conversion/state_transforms.py` | Added the deterministic `e1_to_fastplms_v1` checkpoint mapping. |
+
+These changes were present in the FastPLMs 1.0 repository as reviewed on
+2026-07-20. The conversion identifier is `e1_to_fastplms_v1`. Recipients must
+retain the Profluent-E1 agreement, `ATTRIBUTION`, `NOTICE`, this modified-file
+notice, and the applicable Apache-2.0 and BSD-3-Clause texts.
+
+The BSD-3-Clause component is the padding utility identified by the official E1
+repository as adapted from Dao-AILab FlashAttention. FastPLMs does not copy that
+official utility into its production package, but preserves the notice because
+the official source is the parity oracle for the E1 behavior contract.

LICENSES/e1/NOTICE

@@ -0,0 +1,14 @@
+Profluent-E1 Notice File
+------------------------
+
+Copyright 2025 Profluent Bio Inc.
+
+Licensed under the Profluent-E1 Clickthrough License Agreement (the “Agreement”); you may not use
+this file except in compliance with the Agreement. Unless required by applicable law or agreed to
+in writing, software distributed under the Agreement is distributed on an "AS IS" BASIS, WITHOUT
+WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the Agreement for the specific
+language governing permissions and limitations under the Agreement and the Attribution Guidelines
+for more information about attribution requirements for use of Profluent-E1. You may obtain a copy
+of the Agreement at https://github.com/Profluent-AI/E1/blob/main/LICENSE and a copy of the Attribution
+Guidelines at https://github.com/Profluent-AI/E1/blob/main/ATTRIBUTION, each as may be updated or
+amended from time to time.

README.md

@@ -1,255 +1,169 @@
----
-library_name: transformers
-tags: []
----
-
-# NOTE
-The GitHub with the implementation and requirements.txt can be found [here](https://github.com/Synthyra/FastPLMs.git)
-
-# Profluent-E1
-[Synthyra's version of Profluent-E1](https://github.com/Synthyra/Profluent-E1-300M) is a faithful implementation of Profluent's [E1](https://www.profluent.bio/showcase/e1) models ([license](https://github.com/Profluent-AI/E1/tree/main?tab=License-1-ov-file)) that integrates Huggingface AutoModel compatability and nice embedding functionality.
-
-## Attention backends
-
-`sdpa` (PyTorch Scaled Dot Product Attention) is the default. The backend is set via `config.attn_backend` before loading.
-
-| Backend | Key | Notes |
-| :--- | :--- | :--- |
-| PyTorch SDPA | `"sdpa"` | Default. Exact numerics, stable on all hardware. |
-| Flash Attention | `"kernels_flash"` | Fastest on Ampere/Hopper GPUs. Requires `pip install kernels` (pre-built — no hours-long compilation). Outputs are not bitwise identical to SDPA due to online softmax reordering; differences are often small but not guaranteed to be inconsequential — use `"sdpa"` if exact numerics matter. |
-| Flex Attention | `"flex"` | Uses a block-causal mask that skips padding tokens. Near-exact numerics. First use compiles a Triton kernel (30–120 s). Best combined with `torch.compile`. |
-| Auto | `"auto"` | Picks the best available: `kernels_flash` → `flex` → `sdpa`. |
-
-```python
-from transformers import AutoConfig, AutoModelForMaskedLM
-
-config = AutoConfig.from_pretrained("Synthyra/Profluent-E1-150M", trust_remote_code=True)
-config.attn_backend = "flex"  # or "kernels_flash", "sdpa", "auto"
-model = AutoModelForMaskedLM.from_pretrained("Synthyra/Profluent-E1-150M", config=config, trust_remote_code=True)
-```
-
-`torch.compile(model)` is heavily recommended for sustained throughput, especially with Flex Attention.
-
-
-## Use with 🤗 transformers
-### Supported models
-```python
-model_dict = {
-    # Synthyra/Profluent-E1-150M
-    'Profluent-E1-150M': 'Profluent-Bio/E1-150m',
-    # Synthyra/Profluent-E1-150M
-    'Profluent-E1-300M': 'Profluent-Bio/E1-300m',
-    # Synthyra/Profluent-E1-150M
-    'Profluent-E1-600M': 'Profluent-Bio/E1-600m',
-}
-```
-
-```python
-import torch
-from transformers import AutoModelForMaskedLM
-
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model = AutoModelForMaskedLM.from_pretrained('Synthyra/Profluent-E1-150M', trust_remote_code=True, dtype=torch.bfloat16).eval().to(device)
-
-sequences = ['MPRTEIN', 'MSEQWENCE']
-batch = model.prep_tokens.get_batch_kwargs(sequences, device=device)
-
-output = model(**batch) # get all hidden states with output_hidden_states=True
-print(output.logits.shape) # language modeling logits, (batch_size, seq_len, vocab_size), (2, 11, 34)
-print(output.last_hidden_state.shape) # last hidden state of the model, (batch_size, seq_len, hidden_size), (2, 11, 768)
-print(output.loss) # language modeling loss if you passed labels
-#print(output.hidden_states) # all hidden states if you passed output_hidden_states=True (in tuple)
-#print(outout.attentions) # all attention matrices if you passed output_attentions=True (in tuple)
-```
-
-Our E1 implementation also supports sequence and token level classification tasks like ESM2. Simply pass the number of labels during initialization.
-
-```python
-from transformers import AutoModelForSequenceClassification, AutoModelForTokenClassification
-
-model = AutoModelForSequenceClassification.from_pretrained('Synthyra/Profluent-E1-150M', num_labels=2, trust_remote_code=True)
-logits = model(**batch, labels=labels).logits
-print(logits.shape) # (batch_size, num_labels), (2, 2)
-```
-
-E1 weights were trained in bf16 and are in bf16 by default. You can load them in the precision of your choosing by leveraging the dtype parameter:
-```python
-import torch
-model = AutoModelForMaskedLM.from_pretrained('Synthyra/Profluent-E1-150M', trust_remote_code=True, dtype=torch.float) # fp32
-```
-
-## Experimental test-time training
-
-TTT is disabled by default. Normal E1 inference, MSA-context utilities,
-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. It can
-help some difficult proteins, but it adds test-time compute and can degrade
-already confident predictions.
-
-```python
-metrics = model.ttt(
-    seq="MSTNPKPQRKTKRNT",
-    ttt_config={"steps": 3, "ags": 1, "batch_size": 1},
-)
-model.ttt_reset()
-print(metrics["losses"])
-```
-
-## Embed entire datasets with no new code
-To embed a list of protein sequences **fast**, just call embed_dataset. Sequences are sorted to reduce padding tokens, so the initial progress bar estimation is usually much longer than the actual time it will take.
-
-Example:
-```python
-embedding_dict = model.embed_dataset(
-    sequences=[
-        'MALWMRLLPLLALLALWGPDPAAA', ... # list of protein sequences
-    ],
-    batch_size=2, # adjust for your GPU memory
-    max_len=512, # adjust for your needs
-    full_embeddings=False, # if True, no pooling is performed
-    embed_dtype=torch.float32, # cast to what dtype you want
-    pooling_types=['mean', 'cls'], # more than one pooling type will be concatenated together
-    sql=False, # if True, embeddings will be stored in SQLite database
-    sql_db_path='embeddings.db',
-    save=True, # if True, embeddings will be saved as a .pth file
-    save_path='embeddings.pth',
-)
-# embedding_dict is a dictionary mapping sequences to their embeddings as tensors for .pth or numpy arrays for sql
-```
-
-```
-model.embed_dataset()
-Args:
-    sequences: List of protein sequences
-    batch_size: Batch size for processing
-    max_len: Maximum sequence length
-    full_embeddings: Whether to return full residue-wise (True) embeddings or pooled (False)
-    pooling_type: Type of pooling ('mean' or 'cls')
-    sql: Whether to store embeddings in SQLite database - will be stored in float32
-    sql_db_path: Path to SQLite database
-    
-Returns:
-    Dictionary mapping sequences to embeddings, or None if sql=True
-
-Note:
-    - If sql=True, embeddings can only be stored in float32
-    - sql is ideal if you need to stream a very large dataset for training in real-time
-    - save=True is ideal if you can store the entire embedding dictionary in RAM
-    - sql will be used if it is True and save is True or False
-    - If your sql database or .pth file is already present, they will be scanned first for already embedded sequences
-    - Sequences will be truncated to max_len and sorted by length in descending order for faster processing
-```
-
-## MSA context, PPLL scoring, and RAG embeddings
-
-FastPLMs exposes E1 retrieval-augmented MSA context utilities directly on the model object:
-
-```python
-a3m_path = model.search_homologues(
-    sequence="MALWMRLLPLLALLALWGPDPAAA",
-    output_dir="msas",
-    provider="colabfold",
-)
-
-contexts = model.sample_msa_contexts(
-    a3m_path=a3m_path,
-    max_context_tokens=[6144, 12288, 24576],
-    similarity_thresholds=[1.0, 0.95, 0.9, 0.7, 0.5],
-)
-
-scores = model.score_ppll(
-    sequences=["MALWMRLLPLLALLALWGPDPAAA"],
-    a3m_path=a3m_path,
-    ensemble=True,
-)
-
-embeddings = model.embed_with_msa(
-    sequences=["MALWMRLLPLLALLALWGPDPAAA"],
-    a3m_path=a3m_path,
-    pooling_types=["mean"],
-)
-```
-
-The MSA parsing and context sampling follow Profluent's official E1 `msa_sampling` behavior, including A3M insertion stripping, neighbor reweighting, query-similarity filtering, seeded sampling, and context token budgets.
-
-`score_ppll()` is intentionally different from Profluent's official `E1Scorer`. The official scorer computes mutant scores against a parent sequence with wildtype or masked marginal log-probability deltas. FastPLMs uses a PPLL-style mean correct-token probability over each scored sequence, then optionally averages over sampled contexts. We prefer this API because it is much cheaper while remaining comparable for our use cases.
-
-For dataset embeddings with precomputed MSAs:
-
-```python
-embedding_dict = model.embed_dataset_with_msa(
-    sequences=["MALWMRLLPLLALLALWGPDPAAA"],
-    msa_dir="msas",
-    batch_size=2,
-    pooling_types=["mean"],
-)
-```
-
-The standard `embed()` and `embed_dataset()` paths are unchanged. Use `embed_with_msa()` or `embed_dataset_with_msa()` when you want retrieval context included.
-
-## Fine-tuning with 🤗 peft
-```python
-model = AutoModelForSequenceClassification.from_pretrained('Synthyra/Profluent-E1-150M', num_labels=2, trust_remote_code=True)
-# these modules handle E1 attention layers
-target_modules = ["q_proj", "k_proj", "v_proj", "o_proj"]
-
-lora_config = LoraConfig(
-    r=8, # choose lora parameters to your liking
-    lora_alpha=16,
-    lora_dropout=0.01,
-    bias="none",
-    target_modules=target_modules,
-)
-
-# Apply LoRA to the model
-model = get_peft_model(model, lora_config)
-
-# Unfreeze the classifier head
-for param in model.classifier.parameters():
-    param.requires_grad = True
-```
-
-For a more thourough example of fine-tuning, check out our example script [here](https://github.com/Synthyra/FastPLMs/blob/main/fine_tuning_example.py).
-
-
-### Citations
-
-```bibtex
-@misc{FastPLMs,
-  author={Hallee, Logan and Bichara, David and Gleghorn, Jason P.},
-  title={FastPLMs: Fast, efficient, protein language model inference from Huggingface AutoModel.},
-  year={2024},
-  url={https://huggingface.co/Synthyra/ESMplusplus_small},
-  DOI={10.57967/hf/3726},
-  publisher={Hugging Face}
-}
-```
-
-```bibtex
-@article{jain2025e1,
-  title={E1: Retrieval-Augmented Protein Encoder Models},
-  author={Jain, Sarthak and Beazer, Joel and Ruffolo, Jeffrey A and Bhatnagar, Aadyot and Madani, Ali},
-  journal={bioRxiv},
-  DOI={10.1101/2025.11.12.688125},
-  year={2025}
-}
-```
-
-```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: "other"
+license_name: "profluent-e1-clickthrough-license-agreement"
+license_link: "https://github.com/Profluent-AI/E1/blob/bfd2620a602248499f3d2583d85a7ecddf0b6e02/LICENSE"
+tags:
+  - protein-language-model
+  - fastplms
+---
+
+<!-- Generated from src/fastplms/models.toml. Do not edit. -->
+
+# Synthyra/Profluent-E1-300M
+
+This checkpoint packages the FastPLMs `E1` implementation.
+
+Accepted inputs are raw amino-acid sequences prepared by the native E1 adapter.
+Supported Transformers entry points are `AutoConfig`, `AutoModel`,
+`AutoModelForMaskedLM`, `AutoModelForSequenceClassification`,
+`AutoModelForTokenClassification`.
+
+## Install and platform requirements
+
+Install FastPLMs from the exact source revision paired with this model card:
+
+```bash
+python -m pip install \
+  "fastplms @ git+https://github.com/Synthyra/FastPLMs.git@1b9ce023f1e06571cf3e6324be0610ffa53e0a4a"
+```
+
+Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13 are required. The declared CPU gate covers tiny offline contracts; published checkpoint throughput and parity require the documented device tier. The Hub quick start below requires network
+access on first download. For an air-gapped run, first build the manifest-pinned
+local artifact and use the offline form shown in the example.
+
+## Quick start
+
+```python
+from transformers import AutoModel
+
+model_id = "Synthyra/Profluent-E1-300M"
+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/Profluent-E1-300M` path, then pass `local_files_only=True`.
+
+Leave attention unspecified for the Transformers default. Supported explicit
+choices are `sdpa`, `flex_attention`.
+Pass the selected name through `attn_implementation`.
+When an optimized backend cannot return full attention tensors,
+`output_attentions=True` emits one explicit runtime warning and uses a correctly
+masked eager implementation for that call only. The warning identifies the
+configured backend, effective backend, and reason. Configuration and later
+calls are unchanged.
+For BF16 execution, this family uses parameters loaded directly in BF16.
+
+## Dataset embeddings
+
+The shared embedding API accepts sequences, `(id, sequence)` pairs,
+`EmbeddingInput` records, insertion-ordered `{id: sequence}` mappings, or a
+FASTA path. Results preserve order and duplicate identifiers:
+
+```python
+result = model.embed_dataset(
+    ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"],
+    batch_size=2,
+    pooling=("mean", "std"),
+)
+
+for record in result:
+    print(record.id, record.sequence, record.tensor.shape)
+```
+
+Set `full_embeddings=True` for one residue tensor with shape `(l, d)` per
+sequence. Set `output` to a directory for bounded-memory, transactional
+safetensors with ordered-prefix resume, or choose `format="sqlite"` for
+batch-level database commits and exact resume. Pooling excludes boundary,
+padding, and other non-biological positions.
+
+For a long FASTA run, stream completed batches into SQLite:
+
+```python
+persisted = model.embed_dataset(
+    "proteins.fasta",
+    batch_size=64,
+    pooling=("mean",),
+    output="protein-embeddings.sqlite",
+    format="sqlite",
+    resume=True,
+)
+```
+
+Resume verifies the input order, model state, tokenizer policy, backend, dtype,
+and pooling configuration. It never appends incompatible records to an
+existing run.
+
+## Tokenizer-free E1 input
+
+E1 has no tokenizer. The model retains native raw-sequence preparation,
+boundary tokens, sequence positions, and retrieval-augmented context behavior.
+The ordinary representation path accepts sequences directly:
+
+```python
+result = model.embed_dataset(
+    ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"],
+    batch_size=2,
+    pooling=("mean",),
+)
+print(result[0].tensor.shape)
+```
+
+Lower-level masked-language-model calls must use the E1 batch preparer rather
+than an `AutoTokenizer`. E1 launch messages and distributed legal files retain
+the attribution required by the upstream agreement.
+
+## Runtime contract
+
+- Public input: Raw amino-acid sequences prepared by the native E1 adapter
+- Advertised AutoClasses: `AutoConfig`, `AutoModel`, `AutoModelForMaskedLM`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification`
+- AutoClass weight status: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForMaskedLM` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head`
+- Attention implementations: `sdpa`, `flex_attention`
+- Precision policies: `default`
+- BF16 execution: `static_parameters`
+- Generation contract: `not_applicable`
+- Optional dependency group: `core`
+- Weight publication allowed: `true`
+- Weight license status: `resolved`
+- Redistributable: `true`
+- Complete weight publication required: `false`
+
+## Provenance
+
+- FastPLMs weights: `Synthyra/Profluent-E1-300M@5ef52c0ad2ae2578f40622696b763523810e8e26`
+- Runtime revision: `1b9ce023f1e06571cf3e6324be0610ffa53e0a4a`
+- Runtime source-tree SHA-256: `cec0500dd5f6b238474600a6e80b8d3ccd4544ef7daec90441f14273f1b658d2`
+- Runtime bundle SHA-256: `c99634d724e168524f43f56ad1d22af47c75db0cbe657a2a55a237104fe7b833`
+- Generator/schema version and complete/runtime-only attestations: recorded in `provenance.json`
+- Official checkpoint: `Profluent-Bio/E1-300m@5a2871c587eadbcc9237bc686ea45e5b4d28dfb3`
+- Artifact source: `fast`
+- State transform: `e1_to_fastplms_v1`
+- BF16 execution: `static_parameters`
+- Pinned upstreams: `e1`
+- Reference container: `reference-e1`
+- 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: [Profluent-E1 Clickthrough License Agreement](https://github.com/Profluent-AI/E1/blob/bfd2620a602248499f3d2583d85a7ecddf0b6e02/LICENSE). The Hub model-card identifier is
+`other`. Applicable source licenses, notices, attribution,
+and conversion records are distributed with the local artifact. Review them
+before use.

THIRD_PARTY_NOTICES.md

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

config.json

@@ -4,16 +4,26 @@
   ],
   "attn_backend": "sdpa",
   "auto_map": {
-    "AutoConfig": "modeling_e1.E1Config",
-    "AutoModel": "modeling_e1.E1Model",
-    "AutoModelForMaskedLM": "modeling_e1.E1ForMaskedLM",
-    "AutoModelForSequenceClassification": "modeling_e1.E1ForSequenceClassification",
-    "AutoModelForTokenClassification": "modeling_e1.E1ForTokenClassification"
+    "AutoConfig": "modeling_fastplms.E1Config",
+    "AutoModel": "modeling_fastplms.E1Model",
+    "AutoModelForMaskedLM": "modeling_fastplms.E1ForMaskedLM",
+    "AutoModelForSequenceClassification": "modeling_fastplms.E1ForSequenceClassification",
+    "AutoModelForTokenClassification": "modeling_fastplms.E1ForTokenClassification"
   },
   "bos_token_id": 1,
   "clip_qkv": 8,
   "dtype": "bfloat16",
   "eos_token_id": 2,
+  "fastplms_checkpoint_hash": "c29bdfd0241a0cd9685ed347bdd627018ecacff451baf0cd20674b64f44483f9",
+  "fastplms_checkpoint_repo_id": "Synthyra/Profluent-E1-300M",
+  "fastplms_checkpoint_revision": "5ef52c0ad2ae2578f40622696b763523810e8e26",
+  "fastplms_model_id": "e1_300m",
+  "fastplms_release_tool_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "fastplms_release_tool_sha256": "1459b5d7d13d9b07bd97b3eee764f2ce73623e15e32d07ddf6825c2a9509afb9",
+  "fastplms_runtime_bundle_sha256": "c99634d724e168524f43f56ad1d22af47c75db0cbe657a2a55a237104fe7b833",
+  "fastplms_runtime_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "fastplms_source_tree_sha256": "cec0500dd5f6b238474600a6e80b8d3ccd4544ef7daec90441f14273f1b658d2",
+  "fastplms_weights_revision": "5ef52c0ad2ae2578f40622696b763523810e8e26",
   "gated_mlp": true,
   "global_attention_every_n_layers": 3,
   "gradient_checkpointing": false,

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/e1/attention.py

@@ -0,0 +1,441 @@
+"""E1 attention mask, unpadding, FlexAttention, and kernel adapters."""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Callable
+
+import torch
+from torch.nn.attention.flex_attention import _create_sparse_block_from_block_mask
+
+from fastplms.attention import (
+    BlockMask,
+    _ensure_flash_kernels_loaded,
+    _get_flex_attention_fn,
+    _kernels_flash_forward,
+    _kernels_flash_varlen_forward,
+    create_block_mask,
+    flex_attention,
+    index_first_axis,
+    pad_input,
+)
+
+
+@torch.compiler.disable
+def create_block_causal_mask_optimized(sequence_ids: torch.Tensor) -> BlockMask:
+    if create_block_mask is None:
+        raise RuntimeError("Flex Attention block-mask creation is unavailable in this environment.")
+    # Assumes sequence_ids is sorted in increasing order for each batch item, except for
+    # the -1 values, which are used to indicate the padding tokens.
+    def document_mask(b, h, q_idx, kv_idx):  # type: ignore[no-untyped-def]
+        return (
+            (sequence_ids[b, q_idx] >= sequence_ids[b, kv_idx])
+            & (sequence_ids[b, q_idx] != -1)
+            & (sequence_ids[b, kv_idx] != -1)
+        )
+
+    batch_size, seqlen = sequence_ids.shape
+    return create_block_mask(
+        document_mask, batch_size, 1, seqlen, seqlen, device=sequence_ids.device
+    )
+
+
+@torch.compiler.disable
+def create_within_seq_block_mask(sequence_ids: torch.Tensor) -> BlockMask:
+    if create_block_mask is None:
+        raise RuntimeError("Flex Attention block-mask creation is unavailable in this environment.")
+    def document_mask(b, h, q_idx, kv_idx):  # type: ignore[no-untyped-def]
+        return (
+            (sequence_ids[b, q_idx] == sequence_ids[b, kv_idx])
+            & (sequence_ids[b, q_idx] != -1)
+            & (sequence_ids[b, kv_idx] != -1)
+        )
+
+    batch_size, seqlen = sequence_ids.shape
+    return create_block_mask(
+        document_mask, batch_size, 1, seqlen, seqlen, device=sequence_ids.device
+    )
+
+
+def build_within_seq_mask_4d(sequence_ids: torch.Tensor) -> torch.Tensor:
+    not_pad = sequence_ids != -1
+    same_seq = sequence_ids.unsqueeze(-1) == sequence_ids.unsqueeze(-2)
+    valid = not_pad.unsqueeze(-1) & not_pad.unsqueeze(-2)
+    return (same_seq & valid).unsqueeze(1)
+
+
+def build_block_causal_mask_4d(sequence_ids: torch.Tensor) -> torch.Tensor:
+    not_pad = sequence_ids != -1
+    causal = sequence_ids.unsqueeze(-1) >= sequence_ids.unsqueeze(-2)
+    valid = not_pad.unsqueeze(-1) & not_pad.unsqueeze(-2)
+    return (causal & valid).unsqueeze(1)
+
+
+def flex_attention_func(
+    query_states: torch.Tensor,  # Q has shape (b, l, h, d).
+    key_states: torch.Tensor,  # K has shape (b, l, h_kv, d).
+    value_states: torch.Tensor,  # V has shape (b, l, h_kv, d).
+    score_mod: Callable | None = None,
+    block_mask: BlockMask | None = None,
+    sequence_lengths: tuple[int, ...] | None = None,
+    mask_semantics: str = "within_sequence",
+) -> torch.Tensor:
+    if flex_attention is None:
+        raise RuntimeError("Flex Attention is not available in this environment.")
+    if score_mod is not None:
+        raise NotImplementedError("E1 Flex Attention does not support score_mod.")
+    query_states = query_states.transpose(1, 2).contiguous()  # (bs, nh, seqlen, hs)
+    key_states = key_states.transpose(1, 2).contiguous()  # (bs, nkv, seqlen, hs)
+    value_states = value_states.transpose(1, 2).contiguous()  # (bs, nkv, seqlen, hs)
+
+    fn = _get_flex_attention_fn(
+        device=query_states.device,
+        dtype=query_states.dtype,
+        shape=tuple(query_states.shape),
+        sequence_lengths=sequence_lengths,
+        mask_semantics=mask_semantics,
+    )
+    if fn is None:
+        raise RuntimeError("Flex Attention is not available in this environment.")
+    outputs = fn(
+        query_states,
+        key_states,
+        value_states,
+        block_mask=block_mask,
+        score_mod=score_mod,
+        enable_gqa=query_states.shape[1] != key_states.shape[1],  # if nkv != nh
+    )
+
+    outputs = outputs.transpose(1, 2)  # (bs, seqlen, nh, hs)
+    return outputs
+
+
+def kernels_flash_attention_func(
+    query_states: torch.Tensor,  # (bs, seqlen, nh, hs)
+    key_states: torch.Tensor,  # (bs, seqlen, nkv, hs)
+    value_states: torch.Tensor,  # (bs, seqlen, nkv, hs)
+    q_sequence_ids: torch.Tensor,
+    k_sequence_ids: torch.Tensor,
+    causal: bool = False,
+    implementation: str = "flash_attention_3",
+) -> torch.Tensor:  # (bs, seqlen, nh, hs)
+    _ensure_flash_kernels_loaded(implementation)
+
+    if not causal:
+        batch_size, q_len = query_states.shape[0], query_states.shape[1]
+        (
+            query_states,
+            key_states,
+            value_states,
+            indices_q,
+            (cu_seqlens_q, cu_seqlens_k),
+            (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
+        ) = _unpad_input(query_states, key_states, value_states, q_sequence_ids, k_sequence_ids)
+
+        attn_output_unpad = _kernels_flash_varlen_forward(
+            query_states,
+            key_states,
+            value_states,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_k=cu_seqlens_k,
+            max_seqlen_in_batch_q=max_seqlen_in_batch_q,
+            max_seqlen_in_batch_k=max_seqlen_in_batch_k,
+            causal=False,
+            implementation=implementation,
+        )
+        attn_output = pad_input(attn_output_unpad, indices_q, batch_size, q_len)
+
+    else:
+        attn_output = _kernels_flash_forward(
+            query_states, key_states, value_states, causal=True, implementation=implementation
+        )
+
+    return attn_output
+
+
+def block_min_max_seq_ids(
+    sequence_lengths: torch.Tensor,
+    block_size: int = 128,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Map each physical attention block to its first and last sequence."""
+
+    total_tokens = sequence_lengths.sum()
+    block_count = int(
+        torch.div(
+            total_tokens + block_size - 1,
+            block_size,
+            rounding_mode="floor",
+        ).item()
+    )
+    padded_tokens = block_count * block_size - total_tokens
+    lengths_with_tail = torch.cat(
+        (sequence_lengths, padded_tokens.to(sequence_lengths).reshape(1)),
+    )
+    sequence_ends = lengths_with_tail.to(torch.long).cumsum(dim=0)
+    block_starts = torch.arange(
+        start=0,
+        end=block_count * block_size,
+        step=block_size,
+        dtype=torch.long,
+        device=sequence_lengths.device,
+    )
+    block_last_tokens = block_starts + block_size - 1
+    first_sequence = torch.searchsorted(sequence_ends, block_starts, right=True)
+    last_sequence = torch.searchsorted(sequence_ends, block_last_tokens, right=True)
+    return first_sequence, last_sequence
+
+
+def get_overlapping_blocks(
+    q_lengths: torch.Tensor,
+    k_lengths: torch.Tensor,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    """Classify query/key block pairs as full, partial, or disjoint."""
+
+    q_first, q_last = block_min_max_seq_ids(q_lengths)
+    k_first, k_last = block_min_max_seq_ids(k_lengths)
+    intersection_start = torch.maximum(q_first[:, None], k_first[None, :])
+    intersection_end = torch.minimum(q_last[:, None], k_last[None, :])
+    intersects = intersection_start <= intersection_end
+    both_blocks_are_single_sequence = (q_first == q_last)[:, None] & (k_first == k_last)[None, :]
+    full_blocks = intersects & both_blocks_are_single_sequence
+    return full_blocks, intersects & ~both_blocks_are_single_sequence
+
+
+def _document_ids(sequence_lengths: torch.Tensor) -> torch.Tensor:
+    sequence_numbers = torch.arange(
+        sequence_lengths.numel(),
+        device=sequence_lengths.device,
+        dtype=torch.long,
+    )
+    return sequence_numbers.repeat_interleave(sequence_lengths.to(torch.long))
+
+
+@torch.compiler.disable
+def direct_block_mask(q_lengths: torch.Tensor, k_lengths: torch.Tensor) -> BlockMask:
+    """Build a packed-sequence mask from preclassified sparse blocks."""
+
+    full, partial = get_overlapping_blocks(q_lengths, k_lengths)
+    q_document = _document_ids(q_lengths)
+    k_document = _document_ids(k_lengths)
+
+    def same_document(
+        _batch: torch.Tensor,
+        _head: torch.Tensor,
+        q_index: torch.Tensor,
+        k_index: torch.Tensor,
+    ) -> torch.Tensor:
+        return q_document[q_index].eq(k_document[k_index])
+
+    return _create_sparse_block_from_block_mask(
+        (partial[None, None], full[None, None]),
+        same_document,
+        seq_lengths=(q_document.numel(), k_document.numel()),
+        Q_BLOCK_SIZE=128,
+        KV_BLOCK_SIZE=128,
+    )
+
+
+@torch.compiler.disable
+def doc_id_mask(q_lengths: torch.Tensor, k_lengths: torch.Tensor) -> BlockMask:
+    if create_block_mask is None:
+        raise RuntimeError("Flex Attention block-mask creation is unavailable in this environment.")
+    q_document = _document_ids(q_lengths)
+    k_document = _document_ids(k_lengths)
+
+    def same_document(
+        _batch: torch.Tensor,
+        _head: torch.Tensor,
+        q_index: torch.Tensor,
+        k_index: torch.Tensor,
+    ) -> torch.Tensor:
+        return q_document[q_index].eq(k_document[k_index])
+
+    return create_block_mask(
+        same_document,
+        1,
+        1,
+        q_document.numel(),
+        k_document.numel(),
+        BLOCK_SIZE=128,
+        device=q_lengths.device,
+    )
+
+
+def varlen_flex_attention_func(
+    query_states: torch.Tensor,
+    key_states: torch.Tensor,
+    value_states: torch.Tensor,
+    q_sequence_ids: torch.Tensor,
+    k_sequence_ids: torch.Tensor,
+) -> torch.Tensor:
+    if flex_attention is None:
+        raise RuntimeError("Flex Attention is not available in this environment.")
+    batch_size, q_len = query_states.shape[0], query_states.shape[1]
+    (
+        query_states,
+        key_states,
+        value_states,
+        indices_q,
+        (cu_seqlens_q, cu_seqlens_k),
+        (_max_seqlen_in_batch_q, _max_seqlen_in_batch_k),
+    ) = _unpad_input(query_states, key_states, value_states, q_sequence_ids, k_sequence_ids)
+
+    query_states = query_states.unsqueeze(0).transpose(1, 2).contiguous()
+    key_states = key_states.unsqueeze(0).transpose(1, 2).contiguous()
+    value_states = value_states.unsqueeze(0).transpose(1, 2).contiguous()
+
+    seqlens_q = cu_seqlens_q[1:] - cu_seqlens_q[:-1]
+    seqlens_k = cu_seqlens_k[1:] - cu_seqlens_k[:-1]
+    block_mask = block_mask_creator(seqlens_q, seqlens_k)
+
+    packed_lengths = (
+        *(int(length) for length in seqlens_q.tolist()),
+        -1,
+        *(int(length) for length in seqlens_k.tolist()),
+    )
+    fn = _get_flex_attention_fn(
+        device=query_states.device,
+        dtype=query_states.dtype,
+        shape=tuple(query_states.shape) + tuple(key_states.shape),
+        sequence_lengths=packed_lengths,
+        mask_semantics="packed_document_equality",
+    )
+    if fn is None:
+        raise RuntimeError("Flex Attention is not available in this environment.")
+    attn_output_unpad = fn(
+        query_states,
+        key_states,
+        value_states,
+        block_mask=block_mask,
+        enable_gqa=query_states.shape[1] != key_states.shape[1],
+    )
+
+    attn_output = pad_input(
+        attn_output_unpad.transpose(1, 2).squeeze(0), indices_q, batch_size, q_len
+    )
+
+    return attn_output
+
+
+def _get_unpad_data(sequence_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, int]:
+    """Return packed indices and run lengths for the non-padding sequence IDs."""
+
+    flat_ids = sequence_ids.reshape(-1)
+    non_pad_indices = torch.where(flat_ids.ne(-1))[0]
+    if non_pad_indices.numel() == 0:
+        raise ValueError("Packed attention requires at least one non-padding token.")
+
+    valid_ids = flat_ids.index_select(0, non_pad_indices)
+    row_ids = torch.div(
+        non_pad_indices,
+        sequence_ids.shape[1],
+        rounding_mode="floor",
+    )
+    run_starts = torch.ones_like(valid_ids, dtype=torch.bool)
+    run_starts[1:] = (valid_ids[1:] != valid_ids[:-1]) | (row_ids[1:] != row_ids[:-1])
+    start_indices = torch.where(run_starts)[0]
+    end_indices = torch.cat((start_indices[1:], start_indices.new_tensor([valid_ids.numel()])))
+    sequence_lengths = end_indices - start_indices
+    cumulative_lengths = torch.cat(
+        (
+            torch.zeros(1, dtype=torch.int32, device=sequence_ids.device),
+            sequence_lengths.cumsum(dim=0, dtype=torch.int32),
+        ),
+    )
+    return non_pad_indices, cumulative_lengths, int(sequence_lengths.max().item())
+
+
+def _unpad_input(
+    query_layer: torch.Tensor,
+    key_layer: torch.Tensor,
+    value_layer: torch.Tensor,
+    q_sequence_ids: torch.Tensor,
+    k_sequence_ids: torch.Tensor,
+) -> tuple[
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    torch.Tensor,
+    tuple[torch.Tensor, torch.Tensor],
+    tuple[int, int],
+]:
+    for name, layer in (
+        ("query_layer", query_layer),
+        ("key_layer", key_layer),
+        ("value_layer", value_layer),
+    ):
+        if layer.ndim != 4:
+            raise ValueError(
+                f"{name} must have shape (batch, sequence, heads, head_dim); "
+                f"got {tuple(layer.shape)}."
+            )
+    if value_layer.shape != key_layer.shape:
+        raise ValueError(
+            "key_layer and value_layer must have identical shapes; "
+            f"got {tuple(key_layer.shape)} and {tuple(value_layer.shape)}."
+        )
+    if query_layer.shape[0] != key_layer.shape[0]:
+        raise ValueError(
+            "Query and KV batch sizes must match; "
+            f"got {query_layer.shape[0]} and {key_layer.shape[0]}."
+        )
+    if query_layer.shape[-1] != key_layer.shape[-1]:
+        raise ValueError(
+            "Query and KV head dimensions must match; "
+            f"got {query_layer.shape[-1]} and {key_layer.shape[-1]}."
+        )
+    batch_size, kv_seq_len, num_heads, head_dim = key_layer.shape
+    query_length, num_q_heads = query_layer.shape[1], query_layer.shape[2]
+    if query_layer.shape[:2] != q_sequence_ids.shape:
+        raise ValueError(
+            "Shape mismatch between query layer and query sequence ids: "
+            f"{query_layer.shape[:2]} != {q_sequence_ids.shape}"
+        )
+    if key_layer.shape[:2] != k_sequence_ids.shape:
+        raise ValueError(
+            "Shape mismatch between key layer and key sequence ids: "
+            f"{key_layer.shape[:2]} != {k_sequence_ids.shape}"
+        )
+    if query_length > kv_seq_len:
+        raise ValueError(
+            "Query length must be less than or equal to KV sequence length: "
+            f"{query_length} > {kv_seq_len}"
+        )
+
+    indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(k_sequence_ids)
+
+    key_layer = index_first_axis(
+        key_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k
+    )
+    value_layer = index_first_axis(
+        value_layer.reshape(batch_size * kv_seq_len, num_heads, head_dim), indices_k
+    )
+
+    if torch.equal(q_sequence_ids, k_sequence_ids):
+        indices_q = indices_k
+        cu_seqlens_q = cu_seqlens_k
+        max_seqlen_in_batch_q = max_seqlen_in_batch_k
+    else:
+        indices_q, cu_seqlens_q, max_seqlen_in_batch_q = _get_unpad_data(q_sequence_ids)
+
+    query_layer = index_first_axis(
+        query_layer.reshape(batch_size * query_length, num_q_heads, head_dim), indices_q
+    )
+
+    if cu_seqlens_q.shape != cu_seqlens_k.shape:
+        raise ValueError(
+            "Query and KV must have the same number of sequences: "
+            f"{cu_seqlens_q.shape} != {cu_seqlens_k.shape}"
+        )
+
+    return (
+        query_layer,
+        key_layer,
+        value_layer,
+        indices_q,
+        (cu_seqlens_q, cu_seqlens_k),
+        (max_seqlen_in_batch_q, max_seqlen_in_batch_k),
+    )
+
+
+block_mask_creator = direct_block_mask if os.getenv("FAST_BLOCK_MASK", "1") == "1" else doc_id_mask

fastplms/models/e1/cache.py

@@ -0,0 +1,229 @@
+"""Key-value cache implementations used by E1 inference."""
+
+from __future__ import annotations
+
+from typing import Any
+
+import torch
+from transformers.modeling_outputs import ModelOutput
+from transformers.utils import logging
+
+
+def _get_logger():
+    """Resolve the Transformers logger only when a cache path emits a message."""
+
+    return logging.get_logger(__name__)
+
+
+class DynamicCache:
+    """A cache that grows K and V along their sequence dimension.
+
+    Each cached tensor has shape (b, l, h, d).
+
+    Args:
+        key_cache (`list[torch.Tensor]`): The list of key states.
+        value_cache (`list[torch.Tensor]`): The list of value states.
+    """
+
+    def __init__(self) -> None:
+        self.key_cache: list[torch.Tensor] = []
+        self.value_cache: list[torch.Tensor] = []
+
+    def update(
+        self, key_states: torch.Tensor, value_states: torch.Tensor, layer_idx: int
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """
+        Update the key and value caches in-place, and return the necessary keys and value states.
+
+        Args:
+            key_states (`torch.Tensor`): K to cache with shape (b, l, h, d).
+            value_states (`torch.Tensor`): V to cache with shape (b, l, h, d).
+            layer_idx (`int`): The index of the layer to update.
+
+        Returns:
+            tuple[`torch.Tensor`, `torch.Tensor`]: Cached K and V, each with shape
+            (b, l, h, d).
+        """
+        # Lazy initialization
+        if len(self.key_cache) <= layer_idx:
+            # There may be skipped layers, fill them with empty lists
+            for _ in range(len(self.key_cache), layer_idx):
+                self.key_cache.append(torch.tensor([]))
+                self.value_cache.append(torch.tensor([]))
+            self.key_cache.append(key_states)
+            self.value_cache.append(value_states)
+        elif (
+            not self.key_cache[
+                layer_idx
+            ].numel()  # prefers not t.numel() to len(t) == 0 to export the model
+        ):  # fills previously skipped layers; checking for tensor causes errors
+            self.key_cache[layer_idx] = key_states
+            self.value_cache[layer_idx] = value_states
+        else:
+            self.key_cache[layer_idx] = torch.cat([self.key_cache[layer_idx], key_states], dim=1)
+            self.value_cache[layer_idx] = torch.cat(
+                [self.value_cache[layer_idx], value_states], dim=1
+            )
+
+        return self.key_cache[layer_idx], self.value_cache[layer_idx]
+
+    def get_seq_length(self, layer_idx: int = 0) -> int:
+        """Return the cached sequence length for one layer."""
+        is_empty_layer = (
+            len(self.key_cache) == 0  # no cache in any layer
+            or len(self.key_cache)
+            <= layer_idx  # skipped `layer_idx` and hasn't run a layer with cache after it
+            or not self.key_cache[layer_idx].numel()  # the layer has no cache
+        )
+        layer_seq_length = self.key_cache[layer_idx].shape[1] if not is_empty_layer else 0
+        return layer_seq_length
+
+    def crop(self, max_length: int) -> None:
+        """Crop every cached K and V tensor to ``max_length`` tokens."""
+        if max_length <= 0:
+            raise ValueError("max_length must be positive")
+
+        if self.get_seq_length() <= max_length:
+            return
+
+        for layer_idx in range(len(self.key_cache)):
+            if self.key_cache[layer_idx].numel():
+                self.key_cache[layer_idx] = self.key_cache[layer_idx][:, :max_length, ...]
+                self.value_cache[layer_idx] = self.value_cache[layer_idx][:, :max_length, ...]
+
+    def batch_repeat_interleave(self, repeats: int) -> None:
+        """Repeat the cache `repeats` times in the batch dimension. Used in contrastive search."""
+        for layer_idx in range(len(self.key_cache)):
+            if self.key_cache[layer_idx].numel():
+                self.key_cache[layer_idx] = self.key_cache[layer_idx].repeat_interleave(
+                    repeats, dim=0
+                )
+                self.value_cache[layer_idx] = self.value_cache[layer_idx].repeat_interleave(
+                    repeats, dim=0
+                )
+
+    def batch_select_indices(self, indices: torch.Tensor) -> None:
+        """Keep selected rows of the cache batch dimension."""
+        for layer_idx in range(len(self.key_cache)):
+            if self.key_cache[layer_idx].numel():
+                self.key_cache[layer_idx] = self.key_cache[layer_idx][indices, ...]
+                self.value_cache[layer_idx] = self.value_cache[layer_idx][indices, ...]
+
+
+class KVCache:
+    def __init__(self, cache_size: int = 4) -> None:
+        self.cache_size = cache_size
+        self.tensor_input_field_names = [
+            "input_ids",
+            "within_seq_position_ids",
+            "global_position_ids",
+            "sequence_ids",
+            "labels",
+        ]
+        # Upstream E1 called the encoder output ``embeddings``. FastPLMs uses
+        # the standard Transformers ``last_hidden_state`` name, while keeping
+        # the aliases here makes the cache safe for either output contract.
+        self.tensor_output_field_names = [
+            "logits",
+            "last_hidden_state",
+            "embeddings",
+            "token_embeddings",
+        ]
+        self.cache_dict: dict[str, DynamicCache] = {}
+        self.cache_queue: list[str] = []
+
+    def reset(self) -> None:
+        for k in list(self.cache_dict.keys()):
+            del self.cache_dict[k]
+        del self.cache_dict
+        self.cache_dict = {}
+        self.cache_queue = []
+
+        torch.cuda.empty_cache()
+
+    def before_forward(self, batch: dict[str, torch.Tensor]) -> None:
+        contexts: list[str] | None = batch.get("context")
+        if contexts is None or "context_len" not in batch:
+            _get_logger().warning_once(
+                "KVCache requires both `context` and `context_len`; cache setup was skipped."
+            )
+            return
+
+        context_lens: list[int] = list(set(batch["context_len"]))
+        contexts: list[str] = list(set(contexts))  # type: ignore[no-redef]
+        if len(contexts) != 1 or len(context_lens) != 1:
+            _get_logger().warning(
+                "SingleContextKVCache requires a single context and context length. "
+                "Multiple contexts or context lengths found in a single batch. Skipping."
+            )
+            return
+
+        batch_size = batch["input_ids"].shape[0]
+
+        unique_context = contexts[0]
+        unique_context_len = context_lens[0]
+        batch["use_cache"] = True
+
+        if unique_context not in self.cache_dict:
+            return
+
+        self.cache_dict[unique_context].batch_repeat_interleave(batch_size)
+        past_key_values = self.cache_dict[unique_context]
+        batch["past_key_values"] = past_key_values
+
+        # Remove context from the input fields
+        for field_name in self.tensor_input_field_names:
+            if batch.get(field_name) is not None:
+                batch[field_name] = batch[field_name][:, unique_context_len:]
+
+    def after_forward(self, batch: dict[str, Any], outputs: ModelOutput) -> None:
+        contexts = batch.get("context")
+        context_lens = batch.get("context_len", [])
+        if (
+            contexts is None
+            or len(set(contexts)) != 1
+            or len(set(context_lens)) != 1
+            or context_lens[0] == 0
+        ):
+            return
+
+        if not batch.get("use_cache", False):
+            raise ValueError("E1 retrieval cache updates require use_cache=True.")
+        unique_context = contexts[0]
+        unique_context_len = context_lens[0]
+
+        past_key_values = getattr(outputs, "past_key_values", None)
+        if not isinstance(past_key_values, DynamicCache):
+            _get_logger().warning_once(
+                "KVCache is incompatible with models that don't return a DynamicCache. Skipping."
+            )
+            return
+
+        if "past_key_values" not in batch:
+            if len(self.cache_queue) == self.cache_size:
+                last_context = self.cache_queue.pop(0)
+                if last_context not in self.cache_queue:
+                    del self.cache_dict[last_context]
+                    torch.cuda.empty_cache()
+
+            self.cache_dict[unique_context] = past_key_values
+            self.cache_queue.append(unique_context)
+
+            # Remove context from the input fields
+            for field_name in self.tensor_input_field_names:
+                if field_name in batch and batch[field_name] is not None:
+                    batch[field_name] = batch[field_name][:, unique_context_len:]
+
+            # Remove context from the output fields
+            for field_name in self.tensor_output_field_names:
+                if field_name in outputs and outputs[field_name] is not None:
+                    outputs[field_name] = outputs[field_name][:, unique_context_len:]
+            if "hidden_states" in outputs and outputs["hidden_states"] is not None:
+                hidden_states = outputs["hidden_states"]
+                sliced_hidden_states = tuple(
+                    hidden_state[:, unique_context_len:] for hidden_state in hidden_states
+                )
+                outputs["hidden_states"] = sliced_hidden_states
+
+        self.cache_dict[unique_context].crop(unique_context_len)
+        self.cache_dict[unique_context].batch_select_indices([0])

fastplms/models/e1/modeling_e1.py

@@ -0,0 +1,2319 @@
+from __future__ import annotations
+
+import hashlib
+import os
+import sys
+from collections import defaultdict
+from contextvars import ContextVar
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any, ClassVar, TypedDict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from tqdm.auto import tqdm
+from transformers import PretrainedConfig, PreTrainedModel
+from transformers.activations import ACT2FN
+from transformers.modeling_outputs import ModelOutput
+from transformers.utils import logging
+
+try:
+    from fastplms.attention import (
+        AttentionBackend,
+        BlockMask,
+        FastPLMsAttentionMixin,
+        resolve_attention_backend,
+        resolve_attention_backend_for_call,
+    )
+    from fastplms.embeddings import (
+        EmbeddingBatch,
+        EmbeddingMixin,
+        EmbeddingResult,
+        Pooler,
+        embed_dataset,
+        select_hidden_state_embeddings,
+    )
+    from fastplms.models.ttt import FastPLMTestTimeTrainingMixin
+except ModuleNotFoundError as error:
+    _COMPOSITE_REQUIRED_NAMES = (
+        "AttentionBackend",
+        "BlockMask",
+        "EmbeddingBatch",
+        "EmbeddingMixin",
+        "EmbeddingResult",
+        "FastPLMsAttentionMixin",
+        "FastPLMTestTimeTrainingMixin",
+        "Pooler",
+        "embed_dataset",
+        "resolve_attention_backend",
+        "resolve_attention_backend_for_call",
+        "select_hidden_state_embeddings",
+    )
+    if error.name != "fastplms" or any(
+        name not in globals() for name in _COMPOSITE_REQUIRED_NAMES
+    ):
+        raise
+    # Legacy flat Hub composites define every shared symbol above this block.
+
+from .attention import (  # noqa: F401
+    _document_ids,
+    _get_unpad_data,
+    _unpad_input,
+    block_mask_creator,
+    block_min_max_seq_ids,
+    build_block_causal_mask_4d,
+    build_within_seq_mask_4d,
+    create_block_causal_mask_optimized,
+    create_within_seq_block_mask,
+    direct_block_mask,
+    doc_id_mask,
+    flex_attention_func,
+    get_overlapping_blocks,
+    kernels_flash_attention_func,
+    varlen_flex_attention_func,
+)
+from .cache import DynamicCache, KVCache  # noqa: F401
+from .preparation import (  # noqa: F401
+    BOS_TOKEN_ID,
+    E1_TOKENIZER_REPO_ID,
+    E1_VOCAB_SIZE,
+    EOS_TOKEN_ID,
+    PAD_TOKEN_ID,
+    DataPrepConfig,
+    E1BatchPreparer,
+    _load_tokenizer_file,
+    get_context,
+    get_tokenizer,
+)
+from .retrieval import (  # noqa: F401
+    COLABFOLD_HOST,
+    DEFAULT_EMBED_MAX_TOKENS,
+    DEFAULT_EMBED_SIMILARITY,
+    DEFAULT_MAX_CONTEXT_TOKENS,
+    DEFAULT_SIMILARITY_THRESHOLDS,
+    DOCKER_IMAGE,
+    E1_MSA_SAMPLING_SOURCE_REVISION,
+    LOWERCASE_CHARS,
+    ColabFoldSearcher,
+    ContextCache,
+    ContextSpecification,
+    E1Prediction,
+    HomologueSearcher,
+    IdSequence,
+    IndexedSequence,
+    _ColabFoldResponse,
+    _E1ContextPredictor,
+    _forward_for_embedding,
+    _make_homologue_searcher,
+    _pool_hidden_states,
+    _safe_extract_tar,
+    _sequence_output_dir,
+    _strip_a3m_insertions,
+    build_context_specifications,
+    compute_ppll,
+    convert_to_tensor,
+    get_context_id,
+    get_msa_for_sequence,
+    get_num_neighbors,
+    get_query_from_a3m,
+    get_similarity_to_query,
+    load_msa_dir,
+    load_msa_from_hf,
+    parse_msa,
+    read_fasta_sequences,
+    sample_context,
+    sample_contexts_for_msa,
+    sample_multiple_contexts,
+    write_fasta_sequences,
+)
+
+
+def _get_logger():
+    """Resolve the Transformers logger only when a runtime path emits a message."""
+
+    return logging.get_logger(__name__)
+
+
+_TOKENIZER_LOAD_CONTEXT: ContextVar[dict[str, Any] | None] = ContextVar(
+    "fastplms_e1_tokenizer_load_context",
+    default=None,
+)
+
+
+class E1Config(PretrainedConfig):
+    model_type = "E1"
+    keys_to_ignore_at_inference: ClassVar[list[str]] = ["past_key_values"]
+
+    def __init__(  # type: ignore
+        self,
+        # Model architecture/initialization
+        vocab_size=None,
+        hidden_size=4096,
+        intermediate_size=16384,
+        gated_mlp=False,
+        num_hidden_layers=40,
+        num_attention_heads=32,
+        num_key_value_heads=8,
+        hidden_act="silu",
+        rms_norm_eps=1e-5,
+        initializer_range=0.02,
+        dtype="bfloat16",
+        gradient_checkpointing=False,
+        no_ffn_gradient_checkpointing=False,
+        use_cache=False,
+        # Tokenization
+        pad_token_id=None,
+        bos_token_id=None,
+        eos_token_id=None,
+        tie_word_embeddings=False,
+        # Attention implementation & rotary positional embeddings
+        global_attention_every_n_layers=0,
+        max_num_sequences=512,
+        max_num_positions_within_seq=8192,
+        max_num_positions_global=1024 * 128,
+        rope_theta_within_seq=10000.0,
+        rope_theta_global=100000.0,
+        clip_qkv=None,
+        attn_backend=None,
+        **kwargs,
+    ) -> None:
+        super().__init__(
+            pad_token_id=PAD_TOKEN_ID,
+            bos_token_id=BOS_TOKEN_ID,
+            eos_token_id=EOS_TOKEN_ID,
+            tie_word_embeddings=tie_word_embeddings,
+            dtype=dtype,
+            **kwargs,
+        )
+
+        self.hidden_size = hidden_size
+        if intermediate_size is None:
+            intermediate_size = 3 * hidden_size if gated_mlp else 4 * hidden_size
+        self.intermediate_size = intermediate_size
+        self.gated_mlp = gated_mlp
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.max_num_positions_within_seq = max_num_positions_within_seq
+        self.max_num_positions_global = max_num_positions_global
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.hidden_act = hidden_act
+        self.initializer_range = initializer_range
+        self.rms_norm_eps = rms_norm_eps
+        self.rope_theta_within_seq = rope_theta_within_seq
+        self.rope_theta_global = rope_theta_global
+        self.max_num_sequences = max_num_sequences
+        if clip_qkv is not None and clip_qkv <= 0:
+            raise ValueError(f"clip_qkv must be positive when provided, got {clip_qkv}.")
+        self.clip_qkv = clip_qkv
+        self.global_attention_every_n_layers = global_attention_every_n_layers
+
+        self.vocab_size = E1_VOCAB_SIZE
+        self.gradient_checkpointing = gradient_checkpointing
+        self.no_ffn_gradient_checkpointing = no_ffn_gradient_checkpointing
+        if not isinstance(use_cache, bool):
+            raise TypeError("use_cache must be a boolean.")
+        self.use_cache = use_cache
+        self.attn_backend = attn_backend
+
+        if vocab_size is not None:
+            if vocab_size < self.vocab_size:
+                _get_logger().warning(
+                    f"Using vocab_size {vocab_size} smaller than {self.vocab_size} "
+                    "from the tokenizer contract."
+                )
+                self.vocab_size = vocab_size
+            elif vocab_size > self.vocab_size:
+                _get_logger().warning(
+                    f"Using vocab_size {vocab_size} instead of smaller {self.vocab_size} "
+                    "from E1 tokenizer contract."
+                )
+                self.vocab_size = vocab_size
+        if pad_token_id is not None and pad_token_id != self.pad_token_id:
+            _get_logger().warning(
+                f"Ignoring pad_token_id. Using {self.pad_token_id} from E1 tokenizer contract"
+            )
+        if bos_token_id is not None and bos_token_id != self.bos_token_id:
+            _get_logger().warning(
+                f"Ignoring bos_token_id. Using {self.bos_token_id} from E1 tokenizer contract"
+            )
+        if eos_token_id is not None and eos_token_id != self.eos_token_id:
+            _get_logger().warning(
+                f"Ignoring eos_token_id. Using {self.eos_token_id} from E1 tokenizer contract"
+            )
+
+
+class AttentionLayerType(Enum):
+    WITHIN_SEQ = "within_seq"
+    GLOBAL = "global"
+
+
+class AttentionArgs(TypedDict, total=False):
+    within_seq_block_mask: BlockMask | None
+    block_causal_block_mask: BlockMask | None
+    within_seq_mask_4d: torch.Tensor | None
+    block_causal_mask_4d: torch.Tensor | None
+
+
+def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
+    """This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep).
+
+    The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to (batch,
+    num_attention_heads, seqlen, head_dim)
+    """
+    batch, num_key_value_heads, slen, head_dim = hidden_states.shape
+    if n_rep == 1:
+        return hidden_states
+    hidden_states = hidden_states[:, :, None, :, :].expand(
+        batch, num_key_value_heads, n_rep, slen, head_dim
+    )
+    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)
+
+
+class RotaryPositionalEmbedding(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        max_position_embeddings: int = 2048,
+        base: int = 10000,
+        device: torch.device | None = None,
+    ):
+        super().__init__()
+
+        self.dim = dim
+        self.base = base
+        self.max_position_embeddings = max_position_embeddings
+        # Transformers may instantiate modules on the meta device while loading
+        # a checkpoint. Precomputed non-persistent buffers would then be
+        # materialized without values. Empty buffers make initialization lazy
+        # and deterministic on the first real-device forward.
+        empty = torch.empty(0, dtype=torch.float32, device=device)
+        self.register_buffer("inv_freq", empty, persistent=False)
+        self.register_buffer("cos_cached", empty.clone(), persistent=False)
+        self.register_buffer("sin_cached", empty.clone(), persistent=False)
+        self.max_seq_len_cached = 0
+
+    @staticmethod
+    def rotate_half(x: torch.Tensor) -> torch.Tensor:
+        """Rotates half the hidden dims of the input."""
+        x1 = x[..., : x.shape[-1] // 2]
+        x2 = x[..., x.shape[-1] // 2 :]
+        return torch.cat((-x2, x1), dim=-1)
+
+    def _set_sin_cos_cache(self, seq_len: int, device: torch.device) -> None:
+        # Compute angles in FP32, matching the official cache constructed before
+        # the model is converted to its inference dtype.
+        self.max_seq_len_cached = seq_len
+        inv_freq = self.base ** -(
+            torch.arange(0, self.dim, 2, dtype=torch.float32, device=device) / self.dim
+        )
+        self.inv_freq = inv_freq
+        t = torch.arange(seq_len, device=device, dtype=torch.float32)
+        angles = torch.outer(t, inv_freq)
+        angles = torch.cat((angles, angles), dim=1)
+        self.cos_cached = angles.cos()
+        self.sin_cached = angles.sin()
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        k: torch.Tensor,
+        position_ids: torch.LongTensor,
+        seq_len: int | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        # Q and K have shape (b, l, h, d).
+        device, dtype = q.device, q.dtype
+        seq_len = position_ids.max().item() + 1 if seq_len is None else seq_len
+
+        if seq_len > self.max_seq_len_cached:
+            self._set_sin_cos_cache(seq_len=seq_len, device=device)
+
+        # Selecting by position gives C and S shape (b, l, d). Insert a head
+        # axis so they broadcast over Q and K with shape (b, l, h, d).
+        idxs = position_ids.to(device)
+        cos = self.cos_cached.to(device=device, dtype=dtype).unsqueeze(-2)[idxs]
+        sin = self.sin_cached.to(device=device, dtype=dtype).unsqueeze(-2)[idxs]
+
+        # Apply the real and imaginary parts of the rotary transform to Q and K.
+        # Both halves reuse C and S, so rotate_half supplies the cross terms.
+        q_embed = (q * cos) + (self.rotate_half(q) * sin)
+        k_embed = (k * cos) + (self.rotate_half(k) * sin)
+        return q_embed, k_embed
+
+
+class Attention(nn.Module):
+    """Multi-headed attention from 'Attention Is All You Need' paper."""
+
+    def __init__(self, config: E1Config, layer_idx: int):
+        super().__init__()
+        self.config = config
+        self.layer_idx = layer_idx
+
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.num_attention_heads
+        self.head_dim = self.hidden_size // self.num_heads
+        self.num_kv_heads = config.num_key_value_heads
+        self.num_key_value_groups = self.num_heads // self.num_kv_heads
+        self.max_num_seqs = config.max_num_sequences
+        self.clip_qkv = config.clip_qkv
+
+        if (self.head_dim * self.num_heads) != self.hidden_size:
+            raise ValueError(
+                f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"
+                f" and `num_heads`: {self.num_heads})."
+            )
+        self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
+
+        if self.config.global_attention_every_n_layers > 0:
+            self.layer_type = (
+                AttentionLayerType.GLOBAL
+                if (self.layer_idx + 1) % self.config.global_attention_every_n_layers == 0
+                else AttentionLayerType.WITHIN_SEQ
+            )
+        else:
+            self.layer_type = AttentionLayerType.WITHIN_SEQ
+
+        self.rope_theta = (
+            config.rope_theta_within_seq
+            if self.layer_type == AttentionLayerType.WITHIN_SEQ
+            else config.rope_theta_global
+        )
+        self.max_position_embeddings = (
+            config.max_num_positions_within_seq
+            if self.layer_type == AttentionLayerType.WITHIN_SEQ
+            else config.max_num_positions_global
+        )
+
+        self.rotary_emb = RotaryPositionalEmbedding(
+            self.head_dim,
+            max_position_embeddings=self.max_position_embeddings,
+            base=self.rope_theta,
+        )
+
+        self.attn_backend = resolve_attention_backend(config.attn_backend)
+
+    def prepare_qkv(
+        self,
+        hidden_states: torch.Tensor,
+        position_ids: torch.LongTensor,
+        past_key_value: DynamicCache | None = None,
+        use_cache: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        bsz, q_len, _ = hidden_states.size()
+        query_states: torch.Tensor = self.q_proj(hidden_states)
+        key_states: torch.Tensor = self.k_proj(hidden_states)
+        val_states: torch.Tensor = self.v_proj(hidden_states)
+
+        query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim)
+        key_states = key_states.view(bsz, q_len, self.num_kv_heads, self.head_dim)
+        val_states = val_states.view(bsz, q_len, self.num_kv_heads, self.head_dim)
+
+        if self.clip_qkv is not None:
+            query_states = query_states.clamp(-self.clip_qkv, self.clip_qkv)
+            key_states = key_states.clamp(-self.clip_qkv, self.clip_qkv)
+            val_states = val_states.clamp(-self.clip_qkv, self.clip_qkv)
+
+        query_states, key_states = self.rotary_emb(query_states, key_states, position_ids)
+
+        if use_cache and past_key_value is not None:
+            key_states, val_states = past_key_value.update(key_states, val_states, self.layer_idx)
+
+        input_dtype = query_states.dtype
+        if torch.is_autocast_enabled():
+            target_dtype = torch.get_autocast_dtype("cuda")
+        else:
+            target_dtype = self.q_proj.weight.dtype
+        if input_dtype != target_dtype:
+            _get_logger().warning_once(
+                f"The input hidden states seems to be silently casted in {input_dtype}. "
+                f"This might be because you have upcasted embedding or layer norm layers "
+                f"in {input_dtype}. We will cast back the input in {target_dtype}."
+            )
+            query_states = query_states.to(target_dtype)
+            key_states = key_states.to(target_dtype)
+            val_states = val_states.to(target_dtype)
+
+        return query_states, key_states, val_states
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        within_seq_position_ids: torch.LongTensor,
+        global_position_ids: torch.LongTensor,
+        sequence_ids: torch.LongTensor,
+        attention_args: AttentionArgs | None = None,
+        past_key_value: DynamicCache | None = None,
+        output_attentions: bool = False,
+        output_s_max: bool = False,
+        use_cache: bool = False,
+        effective_backend: AttentionBackend | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, DynamicCache | None, list[torch.Tensor] | None]:
+        is_cache_prefilled = (
+            use_cache
+            and past_key_value is not None
+            and past_key_value.get_seq_length(self.layer_idx) > 0
+        )
+
+        query_states, key_states, val_states = self.prepare_qkv(
+            hidden_states=hidden_states,
+            position_ids=within_seq_position_ids
+            if self.layer_type == AttentionLayerType.WITHIN_SEQ
+            else global_position_ids,
+            past_key_value=past_key_value,
+            use_cache=use_cache,
+        )
+
+        attn_output, attn_weights, s_max = self._attn(
+            query_states=query_states,
+            key_states=key_states,
+            val_states=val_states,
+            sequence_ids=sequence_ids,
+            attention_args=attention_args,
+            output_attentions=output_attentions,
+            output_s_max=output_s_max,
+            is_cache_prefilled=is_cache_prefilled,
+            effective_backend=effective_backend,
+        )
+
+        attn_output = self.o_proj(attn_output)
+        return attn_output, attn_weights, past_key_value, s_max
+
+    def _attn(
+        self,
+        query_states: torch.Tensor,
+        key_states: torch.Tensor,
+        val_states: torch.Tensor,
+        sequence_ids: torch.Tensor,
+        attention_args: AttentionArgs | None = None,
+        output_attentions: bool = False,
+        output_s_max: bool = False,
+        is_cache_prefilled: bool = False,
+        effective_backend: AttentionBackend | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]:
+        # A filled cache changes the implementation shape, not the layer's
+        # biological attention contract. Global layers must retain the cached
+        # context, while within-sequence layers consume only the newly appended
+        # sequence. This matches the pinned E1 inference implementation.
+        effective_layer_type = self.layer_type
+
+        if effective_backend is None:
+            effective_backend = resolve_attention_backend_for_call(
+                self.attn_backend,
+                output_attentions=output_attentions,
+            )
+        if output_attentions:
+            return self._manual_attn(
+                query_states,
+                key_states,
+                val_states,
+                sequence_ids=sequence_ids,
+                attention_args=attention_args,
+                effective_layer_type=effective_layer_type,
+                output_s_max=output_s_max,
+                is_cache_prefilled=is_cache_prefilled,
+            )
+
+        if effective_backend == AttentionBackend.EAGER:
+            attn_output, _, s_max = self._manual_attn(
+                query_states,
+                key_states,
+                val_states,
+                sequence_ids=sequence_ids,
+                attention_args=attention_args,
+                effective_layer_type=effective_layer_type,
+                output_s_max=output_s_max,
+                is_cache_prefilled=is_cache_prefilled,
+            )
+            return attn_output, None, s_max
+        if effective_backend.is_flash:
+            if effective_layer_type == AttentionLayerType.WITHIN_SEQ:
+                attn_output, attn_weights = self._kernels_flash_attn(
+                    query_states,
+                    key_states,
+                    val_states,
+                    sequence_ids=sequence_ids,
+                    is_cache_prefilled=is_cache_prefilled,
+                )
+            else:
+                raise ValueError(
+                    "E1 global attention does not support a kernels Flash backend; "
+                    "use eager, sdpa, or flex_attention."
+                )
+        elif effective_backend == AttentionBackend.FLEX:
+            attn_output, attn_weights = self._flex_attn(
+                query_states,
+                key_states,
+                val_states,
+                sequence_ids=sequence_ids,
+                attention_args=attention_args,
+                effective_layer_type=effective_layer_type,
+                is_cache_prefilled=is_cache_prefilled,
+            )
+        elif effective_backend == AttentionBackend.SDPA:
+            attn_output, attn_weights = self._sdpa_attn(
+                query_states,
+                key_states,
+                val_states,
+                sequence_ids=sequence_ids,
+                attention_args=attention_args,
+                effective_layer_type=effective_layer_type,
+                is_cache_prefilled=is_cache_prefilled,
+            )
+        else:
+            raise AssertionError(f"Unsupported resolved backend: {effective_backend}")
+
+        s_max_key_states = key_states
+        if (
+            is_cache_prefilled
+            and effective_layer_type == AttentionLayerType.WITHIN_SEQ
+            and query_states.shape[1] < key_states.shape[1]
+        ):
+            s_max_key_states = key_states[:, -query_states.shape[1] :]
+        s_max = self._compute_s_max(query_states, s_max_key_states) if output_s_max else None
+        return attn_output, attn_weights, s_max
+
+    @torch.no_grad()
+    def _compute_s_max(
+        self,
+        query_states: torch.Tensor,  # Q has shape (b, l, h, d).
+        key_states: torch.Tensor,  # K has shape (b, l, h_kv, d).
+    ) -> list[torch.Tensor]:
+        query_heads = query_states.transpose(1, 2).contiguous()
+        key_heads = key_states.transpose(1, 2).contiguous()
+        key_heads = repeat_kv(key_heads, self.num_key_value_groups)
+        scale = 1.0 / (self.head_dim**0.5)
+        q_norm = torch.linalg.vector_norm(query_heads, dim=-1)
+        k_norm = torch.linalg.vector_norm(key_heads, dim=-1)
+        s_max_bound = (q_norm.max(dim=-1).values * k_norm.max(dim=-1).values).max(
+            dim=0
+        ).values * scale
+        return [s_max_bound[h] for h in range(self.num_heads)]
+
+    def _kernels_flash_attn(
+        self,
+        query_states: torch.Tensor,
+        key_states: torch.Tensor,
+        val_states: torch.Tensor,
+        sequence_ids: torch.Tensor,
+        is_cache_prefilled: bool = False,
+    ) -> tuple[torch.Tensor, None]:
+        bsz, q_len = query_states.shape[0], query_states.shape[1]
+        _, kv_len = key_states.shape[0], key_states.shape[1]
+
+        if self.layer_type == AttentionLayerType.GLOBAL:
+            q_sequence_ids = sequence_ids
+            if q_len < kv_len:
+                first_token_id = sequence_ids[:, 0].unsqueeze(1)
+                k_sequence_ids = torch.cat(
+                    [first_token_id.expand(bsz, kv_len - q_len), sequence_ids], dim=-1
+                )
+            else:
+                k_sequence_ids = sequence_ids
+        else:
+            if q_len < kv_len:
+                key_states = key_states[:, -q_len:]
+                val_states = val_states[:, -q_len:]
+            q_sequence_ids = k_sequence_ids = sequence_ids
+
+        attn_output = kernels_flash_attention_func(
+            query_states,
+            key_states,
+            val_states,
+            q_sequence_ids=q_sequence_ids,
+            k_sequence_ids=k_sequence_ids,
+            causal=False,
+            implementation=self.attn_backend.value,
+        )
+        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()
+        return attn_output, None
+
+    def _flex_attn(
+        self,
+        query_states: torch.Tensor,
+        key_states: torch.Tensor,
+        val_states: torch.Tensor,
+        sequence_ids: torch.Tensor,
+        attention_args: AttentionArgs | None = None,
+        effective_layer_type: AttentionLayerType = AttentionLayerType.WITHIN_SEQ,
+        is_cache_prefilled: bool = False,
+    ) -> tuple[torch.Tensor, None]:
+        bsz, q_len = query_states.shape[0], query_states.shape[1]
+        kv_len = key_states.shape[1]
+        if is_cache_prefilled and q_len < kv_len:
+            if effective_layer_type == AttentionLayerType.WITHIN_SEQ:
+                key_states = key_states[:, -q_len:]
+                val_states = val_states[:, -q_len:]
+                block_mask = create_within_seq_block_mask(sequence_ids)
+                outputs = flex_attention_func(
+                    query_states,
+                    key_states,
+                    val_states,
+                    block_mask=block_mask,
+                    mask_semantics=effective_layer_type.value,
+                )
+            else:
+                q_sequence_ids, k_sequence_ids = self._cached_global_sequence_ids(
+                    sequence_ids,
+                    kv_len,
+                )
+                outputs = varlen_flex_attention_func(
+                    query_states,
+                    key_states,
+                    val_states,
+                    q_sequence_ids=q_sequence_ids,
+                    k_sequence_ids=k_sequence_ids,
+                )
+            outputs = outputs.reshape(bsz, q_len, self.hidden_size).contiguous()
+            return outputs, None
+
+        if effective_layer_type == AttentionLayerType.WITHIN_SEQ:
+            block_mask = (
+                attention_args["within_seq_block_mask"] if attention_args is not None else None
+            )
+        else:
+            block_mask = (
+                attention_args["block_causal_block_mask"] if attention_args is not None else None
+            )
+        outputs = flex_attention_func(
+            query_states,
+            key_states,
+            val_states,
+            block_mask=block_mask,
+            mask_semantics=effective_layer_type.value,
+        )
+        outputs = outputs.reshape(bsz, q_len, self.hidden_size).contiguous()
+        return outputs, None
+
+    @staticmethod
+    def _cached_global_sequence_ids(
+        query_sequence_ids: torch.Tensor,
+        kv_len: int,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Assign cached context to the incoming query sequence.
+
+        E1 retrieval cache hits contain one incoming sequence. The pinned
+        implementation relabels the cached prefix with that sequence ID so its
+        valid query tokens attend the complete cached context, while padding is
+        excluded by the equality mask or packed Flex path.
+        """
+
+        q_len = query_sequence_ids.shape[1]
+        cached_len = kv_len - q_len
+        if cached_len < 0:
+            raise ValueError(f"E1 cached KV length {kv_len} is shorter than query length {q_len}.")
+        first_sequence_id = query_sequence_ids[:, :1]
+        if bool(first_sequence_id.eq(-1).any()):
+            raise ValueError("E1 cached queries must start with a non-padding sequence token.")
+        cached_sequence_ids = first_sequence_id.expand(-1, cached_len)
+        key_sequence_ids = torch.cat((cached_sequence_ids, query_sequence_ids), dim=-1)
+        return query_sequence_ids, key_sequence_ids
+
+    def _cached_attention_mask_4d(
+        self,
+        sequence_ids: torch.Tensor,
+        kv_len: int,
+        effective_layer_type: AttentionLayerType,
+    ) -> torch.Tensor:
+        if effective_layer_type == AttentionLayerType.WITHIN_SEQ:
+            return build_within_seq_mask_4d(sequence_ids)
+        query_sequence_ids, key_sequence_ids = self._cached_global_sequence_ids(
+            sequence_ids,
+            kv_len,
+        )
+        query_valid = query_sequence_ids.ne(-1)
+        key_valid = key_sequence_ids.ne(-1)
+        same_sequence = query_sequence_ids.unsqueeze(-1).eq(key_sequence_ids.unsqueeze(-2))
+        return (same_sequence & query_valid.unsqueeze(-1) & key_valid.unsqueeze(-2)).unsqueeze(1)
+
+    def _sdpa_attn(
+        self,
+        query_states: torch.Tensor,  # Q has shape (b, l, h, d).
+        key_states: torch.Tensor,  # K has shape (b, l, h_kv, d).
+        val_states: torch.Tensor,  # V has shape (b, l, h_kv, d).
+        sequence_ids: torch.Tensor,
+        attention_args: AttentionArgs | None = None,
+        effective_layer_type: AttentionLayerType = AttentionLayerType.WITHIN_SEQ,
+        is_cache_prefilled: bool = False,
+    ) -> tuple[torch.Tensor, None]:
+        bsz, q_len = query_states.shape[:2]
+        kv_len = key_states.shape[1]
+
+        if is_cache_prefilled and q_len < kv_len:
+            if effective_layer_type == AttentionLayerType.WITHIN_SEQ:
+                key_states = key_states[:, -q_len:]
+                val_states = val_states[:, -q_len:]
+            attention_mask_4d = self._cached_attention_mask_4d(
+                sequence_ids,
+                kv_len,
+                effective_layer_type,
+            )
+        elif attention_args is not None:
+            if effective_layer_type == AttentionLayerType.WITHIN_SEQ:
+                attention_mask_4d = attention_args["within_seq_mask_4d"]
+            else:
+                attention_mask_4d = attention_args["block_causal_mask_4d"]
+        else:
+            attention_mask_4d = None
+
+        query_heads = query_states.transpose(1, 2).contiguous()
+        key_heads = key_states.transpose(1, 2).contiguous()
+        value_heads = val_states.transpose(1, 2).contiguous()
+        key_heads = repeat_kv(key_heads, self.num_key_value_groups)
+        value_heads = repeat_kv(value_heads, self.num_key_value_groups)
+        context_heads = F.scaled_dot_product_attention(
+            query_heads, key_heads, value_heads, attn_mask=attention_mask_4d
+        )
+        attn_output = (
+            context_heads.transpose(1, 2).reshape(bsz, q_len, self.hidden_size).contiguous()
+        )
+        return attn_output, None
+
+    def _manual_attn(
+        self,
+        query_states: torch.Tensor,  # Q has shape (b, l, h, d).
+        key_states: torch.Tensor,  # K has shape (b, l, h_kv, d).
+        val_states: torch.Tensor,  # V has shape (b, l, h_kv, d).
+        sequence_ids: torch.Tensor,
+        attention_args: AttentionArgs | None = None,
+        effective_layer_type: AttentionLayerType = AttentionLayerType.WITHIN_SEQ,
+        output_s_max: bool = False,
+        is_cache_prefilled: bool = False,
+    ) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor] | None]:
+        bsz, q_len = query_states.shape[:2]
+        kv_len = key_states.shape[1]
+
+        if is_cache_prefilled and q_len < kv_len:
+            if effective_layer_type == AttentionLayerType.WITHIN_SEQ:
+                key_states = key_states[:, -q_len:]
+                val_states = val_states[:, -q_len:]
+            attention_mask_4d = self._cached_attention_mask_4d(
+                sequence_ids,
+                kv_len,
+                effective_layer_type,
+            )
+        elif attention_args is not None:
+            if effective_layer_type == AttentionLayerType.WITHIN_SEQ:
+                attention_mask_4d = attention_args["within_seq_mask_4d"]
+            else:
+                attention_mask_4d = attention_args["block_causal_mask_4d"]
+        else:
+            attention_mask_4d = None
+
+        query_heads = query_states.transpose(1, 2).contiguous()
+        key_heads = key_states.transpose(1, 2).contiguous()
+        value_heads = val_states.transpose(1, 2).contiguous()
+        key_heads = repeat_kv(key_heads, self.num_key_value_groups)
+        value_heads = repeat_kv(value_heads, self.num_key_value_groups)
+        scale = 1.0 / (self.head_dim**0.5)
+        attn_weights = torch.matmul(query_heads, key_heads.transpose(-2, -1)) * scale
+        if attention_mask_4d is not None:
+            attention_mask_4d = attention_mask_4d.to(dtype=torch.bool)
+            attn_weights = attn_weights.masked_fill(
+                attention_mask_4d.logical_not(),
+                torch.finfo(attn_weights.dtype).min,
+            )
+        attn_weights = F.softmax(attn_weights, dim=-1)
+        if attention_mask_4d is not None:
+            attn_weights = attn_weights.masked_fill(attention_mask_4d.logical_not(), 0.0)
+        context_heads = torch.matmul(attn_weights, value_heads)
+        attn_output = (
+            context_heads.transpose(1, 2).reshape(bsz, q_len, self.hidden_size).contiguous()
+        )
+        s_max = self._compute_s_max(query_states, key_states) if output_s_max else None
+        return attn_output, attn_weights, s_max
+
+
+class MLP(nn.Module):
+    def __init__(self, config: E1Config):
+        super().__init__()
+        self.ffn_dim = config.intermediate_size
+        self.hidden_dim = config.hidden_size
+        self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)
+        self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.w2(self.act_fn(self.w1(hidden_states)))
+
+
+class GLUMLP(nn.Module):
+    def __init__(self, config: E1Config):
+        super().__init__()
+        self.ffn_dim = config.intermediate_size
+        self.hidden_dim = config.hidden_size
+        self.w1 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)
+        self.w2 = nn.Linear(self.ffn_dim, self.hidden_dim, bias=False)
+        self.w3 = nn.Linear(self.hidden_dim, self.ffn_dim, bias=False)
+        self.act_fn = ACT2FN[config.hidden_act]
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.act_fn(self.w1(hidden_states)) * self.w3(hidden_states)
+        hidden_states = self.w2(hidden_states)
+        return hidden_states
+
+
+class FFN(nn.Module):
+    def __init__(self, config: E1Config):
+        super().__init__()
+        mlp_cls = GLUMLP if config.gated_mlp else MLP
+        self.mlp = mlp_cls(config)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return self.mlp(hidden_states)
+
+
+@dataclass
+class E1ModelOutputWithPast(ModelOutput):
+    """E1 encoder outputs.
+
+    ``last_hidden_state`` is H with shape (b, l, d). Optional hidden states use
+    the same shape per layer, while attention tensors have shape (b, h, l, l).
+    ``past_key_values`` stores the reusable K and V tensors for cached decoding.
+    """
+
+    last_hidden_state: torch.FloatTensor | None = None
+    past_key_values: DynamicCache | None = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    attentions: tuple[torch.FloatTensor, ...] | None = None
+    s_max: tuple[list[torch.Tensor], ...] | None = None
+
+
+@dataclass
+class E1MaskedLMOutputWithPast(ModelOutput):
+    """Masked-LM output with the standard HF fields first, then E1 diagnostics."""
+
+    loss: torch.FloatTensor | None = None
+    logits: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    attentions: tuple[torch.FloatTensor, ...] | None = None
+    mlm_loss: torch.FloatTensor | None = None
+    last_hidden_state: torch.FloatTensor | None = None
+    past_key_values: DynamicCache | None = None
+    s_max: tuple[list[torch.Tensor], ...] | None = None
+
+
+@dataclass
+class E1ClassificationOutputWithPast(ModelOutput):
+    """Sequence-classifier output matching HF ``SequenceClassifierOutputWithPast``."""
+
+    loss: torch.FloatTensor | None = None
+    logits: torch.FloatTensor | None = None
+    past_key_values: DynamicCache | None = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    attentions: tuple[torch.FloatTensor, ...] | None = None
+    last_hidden_state: torch.FloatTensor | None = None
+    s_max: tuple[list[torch.Tensor], ...] | None = None
+
+
+@dataclass
+class E1TokenClassificationOutputWithPast(ModelOutput):
+    """Token-classifier output with the standard HF fields before E1 extensions."""
+
+    loss: torch.FloatTensor | None = None
+    logits: torch.FloatTensor | None = None
+    hidden_states: tuple[torch.FloatTensor, ...] | None = None
+    attentions: tuple[torch.FloatTensor, ...] | None = None
+    last_hidden_state: torch.FloatTensor | None = None
+    past_key_values: DynamicCache | None = None
+    s_max: tuple[list[torch.Tensor], ...] | None = None
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.variance_epsilon = eps
+        self.hidden_size = hidden_size
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        return torch.nn.functional.rms_norm(
+            hidden_states, (self.hidden_size,), self.weight, self.variance_epsilon
+        ).to(input_dtype)
+
+
+class NormAttentionNorm(nn.Module):
+    def __init__(self, config: E1Config, layer_idx: int):
+        super().__init__()
+        self.self_attn = Attention(config, layer_idx)
+        self.input_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.post_attention_layernorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        within_seq_position_ids: torch.LongTensor,
+        global_position_ids: torch.LongTensor,
+        sequence_ids: torch.LongTensor,
+        attention_args: AttentionArgs | None = None,
+        past_key_value: DynamicCache | None = None,
+        output_attentions: bool = False,
+        output_s_max: bool = False,
+        use_cache: bool = False,
+        effective_backend: AttentionBackend | None = None,
+    ) -> tuple[
+        torch.Tensor,
+        torch.Tensor,
+        torch.Tensor | None,
+        DynamicCache | None,
+        list[torch.Tensor] | None,
+    ]:
+        residual = hidden_states
+        hidden_states = self.input_layernorm(hidden_states)
+        hidden_states, self_attn_weights, present_key_value, s_max = self.self_attn(
+            hidden_states=hidden_states,
+            within_seq_position_ids=within_seq_position_ids,
+            global_position_ids=global_position_ids,
+            sequence_ids=sequence_ids,
+            attention_args=attention_args,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            output_s_max=output_s_max,
+            use_cache=use_cache,
+            effective_backend=effective_backend,
+        )
+        hidden_states = residual + hidden_states
+
+        residual = hidden_states
+        hidden_states = self.post_attention_layernorm(hidden_states)
+        return hidden_states, residual, self_attn_weights, present_key_value, s_max
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, config: E1Config, layer_idx: int):
+        super().__init__()
+        self.initializer_range = config.initializer_range
+        self.hidden_size = config.hidden_size
+        self.norm_attn_norm = NormAttentionNorm(config, layer_idx)
+        self.ffn = FFN(config)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        within_seq_position_ids: torch.LongTensor,
+        global_position_ids: torch.LongTensor,
+        sequence_ids: torch.LongTensor,
+        attention_args: AttentionArgs | None = None,
+        past_key_value: DynamicCache | None = None,
+        output_attentions: bool = False,
+        output_s_max: bool = False,
+        use_cache: bool = False,
+        effective_backend: AttentionBackend | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor | None, DynamicCache | None, list[torch.Tensor] | None]:
+        hidden_states, residual, self_attn_weights, present_key_value, s_max = self.norm_attn_norm(
+            hidden_states=hidden_states,
+            within_seq_position_ids=within_seq_position_ids,
+            global_position_ids=global_position_ids,
+            sequence_ids=sequence_ids,
+            attention_args=attention_args,
+            past_key_value=past_key_value,
+            output_attentions=output_attentions,
+            output_s_max=output_s_max,
+            use_cache=use_cache,
+            effective_backend=effective_backend,
+        )
+
+        # Fully Connected
+        hidden_states = self.ffn(hidden_states)
+        hidden_states = residual + hidden_states
+
+        return hidden_states, self_attn_weights, present_key_value, s_max
+
+
+class E1PreTrainedModel(FastPLMsAttentionMixin, PreTrainedModel):
+    config_class = E1Config
+    embedding_unsupported_pooling = ("cls", "parti")
+    config: E1Config
+    base_model_prefix = "model"
+    supports_gradient_checkpointing = True
+    _no_split_modules: ClassVar[list[str]] = ["DecoderLayer"]
+    _transformer_layer_cls: ClassVar[list[type[nn.Module]]] = [DecoderLayer]
+    _skip_keys_device_placement = "past_key_values"
+    all_tied_weights_keys: ClassVar[dict[str, str]] = {}
+    _supports_flash_attn_2 = False
+    _supports_flash_attn_3 = False
+    _fastplms_attention_implementations = ("sdpa", "flex_attention")
+    _is_internal_encoder = False
+
+    def __init__(self, config: E1Config, *args: Any, **kwargs: Any) -> None:
+        super().__init__(config, *args, **kwargs)
+        # The E1 agreement requires this exact attribution when an E1 model is
+        # launched. Internal encoder construction is excluded so each public
+        # model launch displays the attribution exactly once.
+        if not self._is_internal_encoder:
+            print("Profluent-E1", file=sys.stderr, flush=True)
+
+    @classmethod
+    def from_pretrained(  # type: ignore[override]
+        cls,
+        pretrained_model_name_or_path: str | os.PathLike,
+        *model_args: Any,
+        **kwargs: Any,
+    ) -> E1PreTrainedModel:
+        tokenizer_token = None
+        if "token" in kwargs:
+            tokenizer_token = kwargs["token"]
+        elif "use_auth_token" in kwargs:
+            tokenizer_token = kwargs["use_auth_token"]
+        load_context_token = _TOKENIZER_LOAD_CONTEXT.set(
+            {
+                "tokenizer_source": pretrained_model_name_or_path,
+                "local_files_only": bool(kwargs.get("local_files_only", False)),
+                "cache_dir": kwargs.get("cache_dir"),
+                "revision": kwargs.get("revision"),
+                "token": tokenizer_token,
+            }
+        )
+        try:
+            return super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs)
+        finally:
+            _TOKENIZER_LOAD_CONTEXT.reset(load_context_token)
+
+    @staticmethod
+    def _tokenizer_kwargs_from_config(config: E1Config) -> dict[str, Any]:
+        load_context = _TOKENIZER_LOAD_CONTEXT.get()
+        resolved_revision = getattr(config, "_commit_hash", None)
+        if not isinstance(resolved_revision, str) or not resolved_revision.strip():
+            resolved_revision = None
+        if load_context is not None:
+            tokenizer_kwargs = dict(load_context)
+            if resolved_revision is not None:
+                tokenizer_kwargs["revision"] = resolved_revision
+            return tokenizer_kwargs
+
+        tokenizer_source = None
+        if isinstance(config._name_or_path, str) and len(config._name_or_path) > 0:
+            tokenizer_source = config._name_or_path
+        return {
+            "tokenizer_source": tokenizer_source,
+            "local_files_only": False,
+            "cache_dir": None,
+            "revision": resolved_revision,
+            "token": None,
+        }
+
+    @property
+    def prep_tokens(self) -> E1BatchPreparer:
+        """Create E1's raw-sequence preparer only when a sequence API uses it."""
+
+        preparer = self.__dict__.get("_fastplms_prep_tokens")
+        if preparer is not None:
+            return preparer
+        encoder = self._modules.get("model")
+        if encoder is not None and encoder is not self:
+            return encoder.prep_tokens
+        tokenizer_kwargs = self.__dict__.get("_fastplms_tokenizer_kwargs")
+        if tokenizer_kwargs is None:
+            raise RuntimeError("E1 tokenizer settings were not initialized.")
+        preparer = E1BatchPreparer(
+            data_prep_config=DataPrepConfig(
+                max_num_sequences=self.config.max_num_sequences,
+                max_num_positions_within_seq=self.config.max_num_positions_within_seq,
+            ),
+            **tokenizer_kwargs,
+        )
+        self.__dict__["_fastplms_prep_tokens"] = preparer
+        return preparer
+
+    @prep_tokens.setter
+    def prep_tokens(self, value: E1BatchPreparer | None) -> None:
+        self.__dict__["_fastplms_prep_tokens"] = value
+
+    def _init_weights(self, module: nn.Module) -> None:
+        if isinstance(module, RMSNorm):
+            nn.init.ones_(module.weight)
+            return
+        if not isinstance(module, (nn.Linear, nn.Embedding)):
+            return
+
+        nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            nn.init.zeros_(module.bias)
+        if isinstance(module, nn.Embedding) and module.padding_idx is not None:
+            with torch.no_grad():
+                module.weight[module.padding_idx].zero_()
+
+    def _backward_compatibility_gradient_checkpointing(self) -> None:
+        if self.supports_gradient_checkpointing and getattr(
+            self.config, "gradient_checkpointing", False
+        ):
+            self.gradient_checkpointing_enable(dict(use_reentrant=False))
+
+    def post_init(self) -> None:
+        super().post_init()
+
+    @property
+    def _device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    @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"E1 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_E1_ENCODER):
+                module._attn_backend = resolved
+            elif isinstance(module, Attention):
+                module.attn_backend = resolved
+
+
+class FAST_E1_ENCODER(E1PreTrainedModel, EmbeddingMixin):
+    config: E1Config
+    config_class = E1Config
+    _is_internal_encoder = True
+
+    def __init__(self, config: E1Config, **kwargs):
+        E1PreTrainedModel.__init__(self, config, **kwargs)
+        self.padding_idx = config.pad_token_id
+        self.vocab_size = config.vocab_size
+        self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)
+        self.embed_seq_id = nn.Embedding(config.max_num_sequences, config.hidden_size)
+        self.layers = nn.ModuleList(
+            [DecoderLayer(config, i) for i in range(config.num_hidden_layers)]
+        )
+        self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.gradient_checkpointing = config.gradient_checkpointing
+        self.__dict__["_fastplms_tokenizer_kwargs"] = (
+            E1PreTrainedModel._tokenizer_kwargs_from_config(config)
+        )
+        self.__dict__["_fastplms_prep_tokens"] = None
+        self._attn_backend = resolve_attention_backend(config.attn_backend)
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.embed_tokens
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.embed_tokens = value
+
+    def _embed(
+        self,
+        sequences: list[str],
+        return_attention_mask: bool = False,
+        hidden_state_index: int = -1,
+        store_all_hidden_states: bool = False,
+        **kwargs,
+    ) -> torch.Tensor:
+        batch = self.prep_tokens.get_batch_kwargs(sequences, device=self._device)
+        # The native preparer also returns training labels plus retrieval
+        # descriptors. The encoder accepts only its aligned model inputs.
+        encoder_batch: dict[str, torch.Tensor] = {}
+        for name in (
+            "input_ids",
+            "within_seq_position_ids",
+            "global_position_ids",
+            "sequence_ids",
+        ):
+            value = batch[name]
+            if not isinstance(value, torch.Tensor):
+                raise TypeError(f"Prepared E1 field {name!r} must be a tensor.")
+            encoder_batch[name] = value
+        output_hidden_states = store_all_hidden_states or hidden_state_index != -1
+        output = self.forward(
+            **encoder_batch,
+            output_hidden_states=output_hidden_states,
+            output_attentions=False,
+            return_dict=True,
+        )
+        embeddings = select_hidden_state_embeddings(
+            output.last_hidden_state,
+            output.hidden_states,
+            hidden_state_index=hidden_state_index,
+            store_all_hidden_states=store_all_hidden_states,
+        )
+        if return_attention_mask:
+            attention_mask = (encoder_batch["sequence_ids"] != -1).long()
+            return embeddings, attention_mask
+        else:
+            return embeddings
+
+    def _prepare_hidden_states(
+        self,
+        input_ids: torch.LongTensor | None,
+        inputs_embeds: torch.FloatTensor | None,
+        within_seq_position_ids: torch.LongTensor | None,
+        global_position_ids: torch.LongTensor | None,
+        sequence_ids: torch.LongTensor | None,
+    ) -> tuple[torch.Tensor, torch.LongTensor, torch.LongTensor, torch.LongTensor]:
+        if (input_ids is None) == (inputs_embeds is None):
+            message = (
+                "Must specify either input_ids or inputs_embeds"
+                if input_ids is None
+                else "Cannot specify both input_ids and inputs_embeds"
+            )
+            raise ValueError(message)
+
+        source = input_ids if input_ids is not None else inputs_embeds
+        if source is None:
+            raise RuntimeError("E1 input validation did not resolve an input tensor.")
+        expected_rank = 2 if input_ids is not None else 3
+        if source.ndim != expected_rank:
+            source_name = "input_ids" if input_ids is not None else "inputs_embeds"
+            raise ValueError(
+                f"{source_name} must have rank {expected_rank}; got shape {tuple(source.shape)}."
+            )
+        batch_size, sequence_length = source.shape[:2]
+        if sequence_length == 0:
+            raise ValueError("E1 inputs must contain at least one token.")
+        if inputs_embeds is not None and inputs_embeds.shape[-1] != self.config.hidden_size:
+            raise ValueError(
+                "inputs_embeds hidden dimension must match config.hidden_size; "
+                f"got {inputs_embeds.shape[-1]} and {self.config.hidden_size}."
+            )
+        if inputs_embeds is not None:
+            default_positions = torch.arange(sequence_length, device=source.device).expand(
+                batch_size,
+                -1,
+            )
+            if within_seq_position_ids is None:
+                within_seq_position_ids = default_positions
+            if global_position_ids is None:
+                global_position_ids = default_positions
+            if sequence_ids is None:
+                sequence_ids = torch.zeros_like(default_positions)
+
+        if within_seq_position_ids is None or global_position_ids is None or sequence_ids is None:
+            raise ValueError("Position and sequence IDs are required when input_ids are provided.")
+        expected_shape = (batch_size, sequence_length)
+        aligned_inputs = {
+            "within_seq_position_ids": within_seq_position_ids,
+            "global_position_ids": global_position_ids,
+            "sequence_ids": sequence_ids,
+        }
+        for name, value in aligned_inputs.items():
+            if tuple(value.shape) != expected_shape:
+                raise ValueError(
+                    f"{name} must have shape {expected_shape}; got {tuple(value.shape)}."
+                )
+        within_positions = within_seq_position_ids.long()
+        global_positions = global_position_ids.long()
+        sequence_numbers = sequence_ids.long()
+        lowest_position, highest_position = torch.aminmax(within_positions)
+        if (
+            lowest_position.item() < -1
+            or highest_position.item() >= self.config.max_num_positions_within_seq
+        ):
+            raise ValueError(
+                "Position ids must be in the range "
+                f"[-1, {self.config.max_num_positions_within_seq}); got max "
+                f"{highest_position.item()} and min {lowest_position.item()}"
+            )
+        lowest_global, highest_global = torch.aminmax(global_positions)
+        if (
+            lowest_global.item() < -1
+            or highest_global.item() >= self.config.max_num_positions_global
+        ):
+            raise ValueError(
+                "Global position ids must be in the range "
+                f"[-1, {self.config.max_num_positions_global}); got max "
+                f"{highest_global.item()} and min {lowest_global.item()}"
+            )
+        lowest_sequence, highest_sequence = torch.aminmax(sequence_numbers)
+        if lowest_sequence.item() < -1 or highest_sequence.item() >= self.config.max_num_sequences:
+            raise ValueError(
+                "Sequence ids must be in the range "
+                f"[-1, {self.config.max_num_sequences}); got max "
+                f"{highest_sequence.item()} and min {lowest_sequence.item()}"
+            )
+
+        if inputs_embeds is None:
+            if input_ids is None:
+                raise RuntimeError("E1 input validation lost the token ID tensor.")
+            token_embeddings = self.embed_tokens(input_ids)
+            inputs_embeds = token_embeddings + self.embed_seq_id(sequence_numbers.clamp_min(0))
+        layer_dtype = self.layers[0].norm_attn_norm.self_attn.q_proj.weight.dtype
+        target_dtype = (
+            torch.get_autocast_dtype("cuda") if torch.is_autocast_enabled() else layer_dtype
+        )
+        return (
+            inputs_embeds.to(target_dtype),
+            within_positions,
+            global_positions,
+            sequence_numbers,
+        )
+
+    def _resolve_forward_cache(
+        self,
+        past_key_values: DynamicCache | None,
+        use_cache: bool,
+    ) -> tuple[DynamicCache | None, bool]:
+        checkpointing = self.gradient_checkpointing and self.training and torch.is_grad_enabled()
+        if checkpointing and use_cache:
+            _get_logger().warning_once(
+                "`use_cache=True` is incompatible with gradient checkpointing; "
+                "setting `use_cache=False`."
+            )
+            use_cache = False
+        if not use_cache:
+            return None, False
+        return past_key_values if past_key_values is not None else DynamicCache(), True
+
+    def _build_forward_attention_args(
+        self,
+        sequence_ids: torch.LongTensor,
+        past_key_values: DynamicCache | None,
+        effective_backend: AttentionBackend,
+    ) -> AttentionArgs | None:
+        if past_key_values is not None and past_key_values.get_seq_length() != 0:
+            return None
+
+        use_flex = effective_backend == AttentionBackend.FLEX
+        use_dense_mask = effective_backend in {
+            AttentionBackend.EAGER,
+            AttentionBackend.SDPA,
+        }
+        return AttentionArgs(
+            block_causal_block_mask=(
+                create_block_causal_mask_optimized(sequence_ids)
+                if use_flex and self.config.global_attention_every_n_layers > 0
+                else None
+            ),
+            within_seq_block_mask=(
+                create_within_seq_block_mask(sequence_ids) if use_flex else None
+            ),
+            within_seq_mask_4d=(build_within_seq_mask_4d(sequence_ids) if use_dense_mask else None),
+            block_causal_mask_4d=(
+                build_block_causal_mask_4d(sequence_ids) if use_dense_mask else None
+            ),
+        )
+
+    def _run_decoder_layers(
+        self,
+        hidden_states: torch.Tensor,
+        within_seq_position_ids: torch.LongTensor,
+        global_position_ids: torch.LongTensor,
+        sequence_ids: torch.LongTensor,
+        attention_args: AttentionArgs | None,
+        past_key_values: DynamicCache | None,
+        use_cache: bool,
+        output_attentions: bool,
+        output_hidden_states: bool,
+        output_s_max: bool,
+        effective_backend: AttentionBackend,
+    ) -> E1ModelOutputWithPast:
+        hidden_history: list[torch.Tensor] | None = [] if output_hidden_states else None
+        attention_history: list[torch.Tensor] | None = [] if output_attentions else None
+        s_max_history: list[list[torch.Tensor]] | None = [] if output_s_max else None
+        next_cache: DynamicCache | None = None
+
+        for layer in self.layers:
+            if hidden_history is not None:
+                hidden_history.append(hidden_states)
+            if self.gradient_checkpointing and self.training and torch.is_grad_enabled():
+                layer_output = self._gradient_checkpointing_func(
+                    layer.__call__,
+                    hidden_states,
+                    within_seq_position_ids,
+                    global_position_ids,
+                    sequence_ids,
+                    attention_args,
+                    past_key_values,
+                    output_attentions,
+                    output_s_max,
+                    use_cache,
+                    effective_backend,
+                )
+            else:
+                layer_output = layer(
+                    hidden_states,
+                    within_seq_position_ids=within_seq_position_ids,
+                    global_position_ids=global_position_ids,
+                    sequence_ids=sequence_ids,
+                    attention_args=attention_args,
+                    past_key_value=past_key_values,
+                    output_attentions=output_attentions,
+                    output_s_max=output_s_max,
+                    use_cache=use_cache,
+                    effective_backend=effective_backend,
+                )
+            hidden_states, attention, layer_cache, s_max = layer_output
+            if use_cache:
+                past_key_values = layer_cache
+                next_cache = layer_cache
+            if attention_history is not None:
+                if attention is None:
+                    raise RuntimeError(
+                        "An E1 layer did not return attention tensors when requested."
+                    )
+                attention_history.append(attention)
+            if s_max_history is not None:
+                if s_max is None:
+                    raise RuntimeError(
+                        "An E1 layer did not return s_max diagnostics when requested."
+                    )
+                s_max_history.append(s_max)
+
+        hidden_states = self.norm(hidden_states)
+        if hidden_history is not None:
+            hidden_history.append(hidden_states)
+        return E1ModelOutputWithPast(
+            last_hidden_state=hidden_states,
+            past_key_values=next_cache,
+            hidden_states=tuple(hidden_history) if hidden_history is not None else None,
+            attentions=tuple(attention_history) if attention_history is not None else None,
+            s_max=tuple(s_max_history) if s_max_history is not None else None,
+        )
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        within_seq_position_ids: torch.LongTensor | None = None,
+        global_position_ids: torch.LongTensor | None = None,
+        sequence_ids: torch.LongTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        past_key_values: DynamicCache | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        output_s_max: bool = False,
+        return_dict: bool | None = None,
+    ) -> E1ModelOutputWithPast | tuple[Any, ...]:
+        """Transform token or soft embeddings H with shape (b, l, d)."""
+
+        use_cache = (
+            use_cache if use_cache is not None else bool(getattr(self.config, "use_cache", False))
+        )
+        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
+        hidden_states, within_positions, global_positions, sequence_numbers = (
+            self._prepare_hidden_states(
+                input_ids,
+                inputs_embeds,
+                within_seq_position_ids,
+                global_position_ids,
+                sequence_ids,
+            )
+        )
+        cache, use_cache = self._resolve_forward_cache(past_key_values, use_cache)
+        effective_backend = resolve_attention_backend_for_call(
+            self._attn_backend,
+            output_attentions=bool(output_attentions),
+        )
+        attention_args = self._build_forward_attention_args(
+            sequence_numbers,
+            cache,
+            effective_backend,
+        )
+        result = self._run_decoder_layers(
+            hidden_states,
+            within_positions,
+            global_positions,
+            sequence_numbers,
+            attention_args,
+            cache,
+            use_cache,
+            output_attentions,
+            output_hidden_states,
+            output_s_max,
+            effective_backend,
+        )
+        if not return_dict:
+            return result.to_tuple()
+        return result
+
+
+class E1Model(E1PreTrainedModel, EmbeddingMixin):
+    config: E1Config
+    config_class = E1Config
+
+    def __init__(self, config: E1Config, **kwargs):
+        E1PreTrainedModel.__init__(self, config, **kwargs)
+        self.model: FAST_E1_ENCODER = FAST_E1_ENCODER(config, **kwargs)
+        self.post_init()
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.model.get_input_embeddings()
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.model.set_input_embeddings(value)
+
+    def _embed(
+        self, sequences: list[str], return_attention_mask: bool = False, **kwargs
+    ) -> torch.Tensor:
+        return self.model._embed(sequences, return_attention_mask=return_attention_mask, **kwargs)
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        within_seq_position_ids: torch.LongTensor | None = None,
+        global_position_ids: torch.LongTensor | None = None,
+        sequence_ids: torch.LongTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        past_key_values: DynamicCache | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        output_s_max: bool = False,
+        return_dict: bool | None = None,
+    ) -> E1ModelOutputWithPast | tuple[Any, ...]:
+        return self.model(
+            input_ids=input_ids,
+            within_seq_position_ids=within_seq_position_ids,
+            global_position_ids=global_position_ids,
+            sequence_ids=sequence_ids,
+            inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            output_s_max=output_s_max,
+            return_dict=return_dict,
+        )
+
+
+class E1ForMaskedLM(FastPLMTestTimeTrainingMixin, E1PreTrainedModel, EmbeddingMixin):
+    config: E1Config
+    config_class = E1Config
+
+    def __init__(self, config: E1Config, **kwargs):
+        E1PreTrainedModel.__init__(self, config, **kwargs)
+        self.model: FAST_E1_ENCODER = FAST_E1_ENCODER(config, **kwargs)
+        self.vocab_size = config.vocab_size
+        self.mlm_head = torch.nn.Sequential(
+            nn.Linear(config.hidden_size, config.hidden_size, bias=True),
+            nn.GELU(),
+            nn.LayerNorm(config.hidden_size, eps=config.rms_norm_eps),
+            nn.Linear(config.hidden_size, config.vocab_size, bias=True),
+        )
+        self.gradient_checkpointing = config.gradient_checkpointing
+        self.post_init()
+        self.init_ttt({"lora_target_replace_module": "Attention"})
+
+    @property
+    def device_mesh(self) -> torch.distributed.device_mesh.DeviceMesh:
+        return self.model.device_mesh
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.model.get_input_embeddings()
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.model.set_input_embeddings(value)
+
+    def _embed(
+        self, sequences: list[str], return_attention_mask: bool = False, **kwargs
+    ) -> torch.Tensor:
+        return self.model._embed(sequences, return_attention_mask=return_attention_mask, **kwargs)
+
+    def get_output_embeddings(self) -> nn.Linear:
+        return self.mlm_head[-1]
+
+    def set_output_embeddings(self, value: nn.Linear) -> None:
+        self.mlm_head[-1] = value
+
+    def _ttt_get_trainable_modules(self) -> list[nn.Module]:
+        return [self.model]
+
+    def _ttt_tokenize(
+        self,
+        seq: str | list[str] | None = None,
+        input_ids: torch.Tensor | None = None,
+        **kwargs,
+    ) -> dict[str, torch.Tensor]:
+        if input_ids is not None:
+            return {
+                "input_ids": input_ids,
+                "within_seq_position_ids": kwargs["within_seq_position_ids"],
+                "global_position_ids": kwargs["global_position_ids"],
+                "sequence_ids": kwargs["sequence_ids"],
+            }
+        if seq is None:
+            raise ValueError("Pass either seq or E1 token tensors for TTT.")
+        sequences = [seq] if isinstance(seq, str) else seq
+        batch = self.prep_tokens.get_batch_kwargs(sequences, device=torch.device("cpu"))
+        return {
+            "input_ids": batch["input_ids"],
+            "within_seq_position_ids": batch["within_seq_position_ids"],
+            "global_position_ids": batch["global_position_ids"],
+            "sequence_ids": batch["sequence_ids"],
+        }
+
+    def _ttt_mask_token(self) -> int:
+        return int(self.prep_tokens.mask_token_id)
+
+    def _ttt_padding_token(self) -> int:
+        return int(self.prep_tokens.pad_token_id)
+
+    def _ttt_replacement_tokens(self, input_ids: torch.Tensor) -> torch.Tensor:
+        amino_acids = "ACDEFGHIKLMNPQRSTVWY"
+        ids = [self.prep_tokens.vocab[aa] for aa in amino_acids]
+        return torch.tensor(ids, device=input_ids.device, dtype=input_ids.dtype)
+
+    def _ttt_non_special_mask(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return ~self.prep_tokens.get_boundary_token_mask(input_ids)
+
+    def _ttt_predict_logits(
+        self,
+        batch: torch.Tensor | dict[str, torch.Tensor],
+        **kwargs,
+    ) -> torch.Tensor:
+        del kwargs
+        if not isinstance(batch, dict):
+            raise TypeError("E1 TTT expects a tensor dictionary.")
+        output = self(
+            input_ids=batch["input_ids"],
+            within_seq_position_ids=batch["within_seq_position_ids"],
+            global_position_ids=batch["global_position_ids"],
+            sequence_ids=batch["sequence_ids"],
+            return_dict=True,
+        )
+        return output.logits
+
+    def search_homologues(
+        self,
+        sequence: str,
+        output_dir: str,
+        provider: str = "colabfold",
+        target_db: str | None = None,
+        seq_id: str | None = None,
+        **kwargs,
+    ) -> str:
+        searcher = _make_homologue_searcher(provider=provider, target_db=target_db, **kwargs)
+        return searcher.search(sequence=sequence, output_dir=output_dir, seq_id=seq_id)
+
+    def batch_search_homologues(
+        self,
+        sequences: list[str],
+        output_dir: str,
+        provider: str = "colabfold",
+        target_db: str | None = None,
+        seq_ids: list[str] | None = None,
+        continue_on_error: bool = True,
+        **kwargs,
+    ) -> dict[str, str]:
+        searcher = _make_homologue_searcher(provider=provider, target_db=target_db, **kwargs)
+        return searcher.batch_search(
+            sequences=sequences,
+            output_dir=output_dir,
+            seq_ids=seq_ids,
+            continue_on_error=continue_on_error,
+        )
+
+    def sample_msa_contexts(
+        self,
+        a3m_path: str,
+        seed: int = 42,
+        max_context_tokens: list[int] | None = None,
+        similarity_thresholds: list[float] | None = None,
+        min_query_similarity: float = 0.3,
+        context_cache_dir: str | None = None,
+    ) -> dict[str, str]:
+        context_specs = build_context_specifications(
+            max_context_tokens=max_context_tokens,
+            similarity_thresholds=similarity_thresholds,
+            min_query_similarity=min_query_similarity,
+        )
+        cache = None
+        if context_cache_dir is not None:
+            key = repr((max_context_tokens, similarity_thresholds, min_query_similarity))
+            specs_hash = hashlib.md5(key.encode()).hexdigest()[:8]
+            cache = ContextCache(context_cache_dir, specs_hash, seed)
+            cached = cache.load(a3m_path)
+            if cached is not None:
+                return cached
+        contexts = sample_contexts_for_msa(a3m_path, context_specs, seed=seed)
+        if cache is not None:
+            cache.store(a3m_path, contexts)
+        return contexts
+
+    @torch.inference_mode()
+    def score_ppll(
+        self,
+        sequences: list[str],
+        a3m_path: str,
+        ensemble: bool = True,
+        seed: int = 42,
+        max_context_tokens: list[int] | None = None,
+        similarity_thresholds: list[float] | None = None,
+        min_query_similarity: float = 0.3,
+        max_batch_tokens: int = 131072,
+        cache_size: int = 1,
+        context_cache_dir: str | None = None,
+        progress: bool = True,
+    ) -> list[float] | list[list[float]]:
+        """Score sequences with FastPLMs PPLL reduction over sampled E1 MSA contexts.
+
+        This intentionally differs from Profluent's official E1Scorer, which scores
+        mutants against a parent sequence with wildtype or masked marginal log-prob
+        deltas. Here each sequence is scored by mean correct-token probability and
+        optionally averaged across sampled contexts.
+        """
+        contexts = self.sample_msa_contexts(
+            a3m_path=a3m_path,
+            seed=seed,
+            max_context_tokens=max_context_tokens,
+            similarity_thresholds=similarity_thresholds,
+            min_query_similarity=min_query_similarity,
+            context_cache_dir=context_cache_dir,
+        )
+        if not contexts:
+            raise ValueError("At least one sampled MSA context is required for PPLL scoring.")
+
+        predictor = _E1ContextPredictor(
+            model=self,
+            data_prep_config=DataPrepConfig(remove_X_tokens=True),
+            max_batch_tokens=max_batch_tokens,
+            fields_to_save=["logits"],
+            save_masked_positions_only=False,
+            keep_predictions_in_gpu=False,
+            use_cache=True,
+            cache_size=cache_size,
+            progress=progress,
+        )
+        vocab = predictor.batch_preparer.vocab
+        seq_token_ids = [
+            torch.tensor([vocab[aa] for aa in seq if aa != "X"], device=self.device)
+            for seq in sequences
+        ]
+        context_ids = list(contexts.keys())
+        all_scores = torch.zeros(len(sequences), len(context_ids), device=self.device)
+
+        iterator = tqdm(context_ids, desc="Scoring with contexts", disable=not progress)
+        for ctx_idx, ctx_id in enumerate(iterator):
+            predictions = list(
+                predictor.predict(
+                    sequences=sequences,
+                    sequence_ids=list(range(len(sequences))),
+                    context_seqs={ctx_id: contexts[ctx_id]},
+                )
+            )
+            for prediction in predictions:
+                seq_idx = prediction["id"]
+                if not isinstance(seq_idx, int):
+                    raise TypeError("Expected integer sequence ids for score aggregation.")
+                all_scores[seq_idx, ctx_idx] = compute_ppll(
+                    prediction["logits"], seq_token_ids[seq_idx]
+                )
+            if predictor.kv_cache is not None:
+                predictor.kv_cache.reset()
+
+        if ensemble:
+            return all_scores.mean(dim=1).tolist()
+        return all_scores.tolist()
+
+    @torch.inference_mode()
+    def embed_with_msa(
+        self,
+        sequences: list[str],
+        a3m_path: str | None = None,
+        context: str | None = None,
+        pooling_types: list[str] | None = None,
+        pooling: str = "mean",
+        matrix_embed: bool = False,
+        seed: int = 42,
+        max_batch_tokens: int = 131072,
+        embed_max_tokens: int = DEFAULT_EMBED_MAX_TOKENS,
+        embed_similarity: float = DEFAULT_EMBED_SIMILARITY,
+        min_query_similarity: float = 0.3,
+        progress: bool = True,
+    ) -> torch.Tensor | list[torch.Tensor]:
+        if a3m_path is not None and context is None:
+            spec = ContextSpecification(
+                max_num_samples=511,
+                max_token_length=embed_max_tokens,
+                max_query_similarity=embed_similarity,
+                min_query_similarity=min_query_similarity,
+            )
+            contexts, _ = sample_multiple_contexts(
+                msa_path=a3m_path,
+                context_specifications=[spec],
+                seed=seed,
+            )
+            context = contexts[0] if contexts else None
+
+        hidden_list = _forward_for_embedding(
+            model=self,
+            sequences=sequences,
+            context=context,
+            max_batch_tokens=max_batch_tokens,
+            progress=progress,
+        )
+        if matrix_embed:
+            return hidden_list
+        if pooling_types is not None:
+            return _pool_hidden_states(hidden_list, pooling_types, self.device)
+        if pooling not in ("mean", "cls"):
+            raise ValueError("pooling must be 'mean' or 'cls' when pooling_types is not provided")
+        embeddings = [
+            hidden.mean(dim=0) if pooling == "mean" else hidden[0] for hidden in hidden_list
+        ]
+        return torch.stack(embeddings)
+
+    @torch.inference_mode()
+    def embed_dataset_with_msa(
+        self,
+        sequences: list[str],
+        msa_lookup: dict[str, str] | None = None,
+        msa_dir: str | None = None,
+        msa_hf_path: str | None = None,
+        batch_size: int = 2,
+        max_len: int = 2048,
+        pooling_types: list[str] | None = None,
+        pooling: str = "mean",
+        matrix_embed: bool = False,
+        embed_dtype: torch.dtype = torch.bfloat16,
+        embed_max_tokens: int = DEFAULT_EMBED_MAX_TOKENS,
+        embed_similarity: float = DEFAULT_EMBED_SIMILARITY,
+        min_query_similarity: float = 0.3,
+        seed: int = 42,
+        progress: bool = True,
+        max_batch_tokens: int = 131072,
+        batch_window_size: int | None = None,
+        max_tokens_per_batch: int | None = None,
+        output: str | os.PathLike[str] | None = None,
+        format: str = "safetensors",
+        resume: bool = True,
+        shard_size: int = 2 * 1024**3,
+        model_state_fingerprint: str | None = None,
+    ) -> EmbeddingResult:
+        """Embed an ordered sequence dataset with optional sampled MSA context.
+
+        Unlike the legacy dictionary return, the result preserves duplicate
+        sequences and input order. ``output`` uses the same transactional,
+        resumable SQLite or safetensors persistence as :meth:`embed_dataset`.
+        ``max_len`` counts biological residues.
+        """
+
+        if not sequences:
+            raise ValueError("sequences must contain at least one protein sequence.")
+        if any(not isinstance(sequence, str) or not sequence for sequence in sequences):
+            raise ValueError("sequences must contain non-empty strings.")
+        if max_len <= 0:
+            raise ValueError("max_len must be positive.")
+        if msa_lookup is None:
+            if msa_dir is not None:
+                msa_lookup = load_msa_dir(msa_dir)
+            elif msa_hf_path is not None:
+                msa_lookup = load_msa_from_hf(msa_hf_path)
+            else:
+                msa_lookup = {}
+
+        truncated_sequences = [sequence[:max_len] for sequence in sequences]
+        unique_seqs = sorted(set(truncated_sequences), key=lambda value: (-len(value), value))
+        context_map: dict[str, str | None] = {}
+        spec = ContextSpecification(
+            max_num_samples=511,
+            max_token_length=embed_max_tokens,
+            max_query_similarity=embed_similarity,
+            min_query_similarity=min_query_similarity,
+        )
+        for seq in unique_seqs:
+            a3m_path = get_msa_for_sequence(seq, msa_lookup)
+            if a3m_path is None:
+                context_map[seq] = None
+                continue
+            contexts, _ = sample_multiple_contexts(
+                msa_path=a3m_path,
+                context_specifications=[spec],
+                seed=seed,
+            )
+            context_map[seq] = contexts[0] if contexts else None
+
+        context_digest = hashlib.sha256()
+        for sequence in unique_seqs:
+            for value in (sequence, context_map[sequence] or ""):
+                encoded = value.encode("utf-8")
+                context_digest.update(len(encoded).to_bytes(8, "big"))
+                context_digest.update(encoded)
+
+        def embed_msa_batch(batch_sequences: list[str]) -> EmbeddingBatch:
+            grouped_positions: dict[str | None, list[int]] = defaultdict(list)
+            for position, sequence in enumerate(batch_sequences):
+                grouped_positions[context_map[sequence]].append(position)
+            hidden_by_position: list[torch.Tensor | None] = [None] * len(batch_sequences)
+            for context, positions in grouped_positions.items():
+                context_sequences = [batch_sequences[position] for position in positions]
+                hidden_states = _forward_for_embedding(
+                    model=self,
+                    sequences=context_sequences,
+                    context=context,
+                    max_batch_tokens=max_batch_tokens,
+                    progress=progress,
+                )
+                for position, hidden in zip(positions, hidden_states, strict=True):
+                    hidden_by_position[position] = hidden
+            resolved = [hidden for hidden in hidden_by_position if hidden is not None]
+            if len(resolved) != len(batch_sequences):
+                raise RuntimeError("E1 MSA embedding did not return every requested sequence.")
+            max_residues = max(hidden.shape[0] for hidden in resolved)
+            hidden_size = resolved[0].shape[-1]
+            X = resolved[0].new_zeros((len(resolved), max_residues, hidden_size))
+            residue_mask = torch.zeros(
+                (len(resolved), max_residues),
+                dtype=torch.bool,
+                device=X.device,
+            )
+            for position, hidden in enumerate(resolved):
+                residue_count = hidden.shape[0]
+                X[position, :residue_count] = hidden
+                residue_mask[position, :residue_count] = True
+            return EmbeddingBatch(X=X, residue_mask=residue_mask)
+
+        resolved_pooling: str | list[str] | None = (
+            None if matrix_embed else pooling_types if pooling_types is not None else pooling
+        )
+        adapter_identity = {
+            "kind": "e1-msa-v1",
+            "sampling_source_revision": E1_MSA_SAMPLING_SOURCE_REVISION,
+            "context_sha256": context_digest.hexdigest(),
+            "context_count": sum(context is not None for context in context_map.values()),
+            "seed": seed,
+            "embed_max_tokens": embed_max_tokens,
+            "embed_similarity": embed_similarity,
+            "min_query_similarity": min_query_similarity,
+            "max_batch_tokens": max_batch_tokens,
+        }
+        return embed_dataset(
+            self,
+            [(str(position), sequence) for position, sequence in enumerate(sequences)],
+            batch_size=batch_size,
+            pooling=resolved_pooling,
+            full_embeddings=matrix_embed,
+            output=output,
+            format=format,
+            resume=resume,
+            max_length=max_len,
+            truncate=True,
+            dtype=embed_dtype,
+            shard_size=shard_size,
+            model_state_fingerprint=model_state_fingerprint,
+            batch_window_size=batch_window_size,
+            max_tokens_per_batch=max_tokens_per_batch,
+            _embedding_batch_fn=embed_msa_batch,
+            _embedding_batch_identity=adapter_identity,
+            _allowed_unsupported_pooling=("cls",),
+        )
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        within_seq_position_ids: torch.LongTensor | None = None,
+        global_position_ids: torch.LongTensor | None = None,
+        sequence_ids: torch.LongTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        past_key_values: DynamicCache | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        output_s_max: bool = False,
+        return_dict: bool | None = None,
+    ) -> E1MaskedLMOutputWithPast | tuple[Any, ...]:
+        """Return hidden states and masked-token logits for E1 inputs.
+
+        Token, position, sequence, and label tensors have shape (b, l).
+        Callers may instead provide precomputed H with shape (b, l, d).
+        """
+        use_cache = (
+            use_cache if use_cache is not None else bool(getattr(self.config, "use_cache", False))
+        )
+        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
+        outputs: E1ModelOutputWithPast = self.model(
+            input_ids=input_ids,
+            within_seq_position_ids=within_seq_position_ids,
+            global_position_ids=global_position_ids,
+            sequence_ids=sequence_ids,
+            inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            output_s_max=output_s_max,
+            return_dict=True,
+        )
+
+        last_hidden_state = outputs.last_hidden_state
+        loss = None
+
+        mlm_logits = self.mlm_head(last_hidden_state).float()
+        mlm_loss = None
+        if labels is not None:
+            mlm_logits_flat = mlm_logits.contiguous().view(-1, self.config.vocab_size)
+            mlm_labels_flat = labels.to(mlm_logits_flat.device).contiguous().view(-1)
+            mlm_loss = F.cross_entropy(
+                mlm_logits_flat,
+                mlm_labels_flat,
+                ignore_index=-100,
+                reduction="none",
+            )
+            mask = mlm_labels_flat.ne(-100) & mlm_labels_flat.ne(self.model.padding_idx)
+            n_mlm = mask.sum().clamp_min(1)
+            mlm_loss = (mlm_loss * mask.to(mlm_loss)).sum() / n_mlm
+            loss = 0.0
+            loss += mlm_loss
+
+        result = E1MaskedLMOutputWithPast(
+            loss=loss,
+            logits=mlm_logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            mlm_loss=mlm_loss,
+            last_hidden_state=last_hidden_state,
+            past_key_values=outputs.past_key_values,
+            s_max=outputs.s_max,
+        )
+        if not return_dict:
+            return result.to_tuple()
+        return result
+
+
+class E1ForSequenceClassification(E1PreTrainedModel, EmbeddingMixin):
+    config: E1Config
+    config_class = E1Config
+
+    def __init__(self, config: E1Config, **kwargs):
+        pooling_types = kwargs.pop("pooling_types", None)
+        if pooling_types is None:
+            pooling_types = ["mean", "var"]
+        elif not isinstance(pooling_types, list):
+            raise TypeError("pooling_types must be a non-empty list of pooling names")
+        elif not pooling_types or not all(isinstance(name, str) for name in pooling_types):
+            raise ValueError("pooling_types must be a non-empty list of pooling names")
+
+        E1PreTrainedModel.__init__(self, config, **kwargs)
+        self.model: FAST_E1_ENCODER = FAST_E1_ENCODER(config, **kwargs)
+        self.vocab_size = config.vocab_size
+        self.num_labels = config.num_labels
+        self.pooler = Pooler(pooling_types)
+        self.classifier = nn.Sequential(
+            nn.Linear(config.hidden_size * len(pooling_types), config.hidden_size * 4),
+            nn.GELU(),
+            nn.LayerNorm(config.hidden_size * 4),
+            nn.Linear(config.hidden_size * 4, config.num_labels),
+        )
+        self.mse = nn.MSELoss()
+        self.ce = nn.CrossEntropyLoss()
+        self.bce = nn.BCEWithLogitsLoss()
+        self.gradient_checkpointing = config.gradient_checkpointing
+        self.post_init()
+
+    @property
+    def device_mesh(self) -> torch.distributed.device_mesh.DeviceMesh:
+        return self.model.device_mesh
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.model.get_input_embeddings()
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.model.set_input_embeddings(value)
+
+    def _embed(
+        self, sequences: list[str], return_attention_mask: bool = False, **kwargs
+    ) -> torch.Tensor:
+        return self.model._embed(sequences, return_attention_mask=return_attention_mask, **kwargs)
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        within_seq_position_ids: torch.LongTensor | None = None,
+        global_position_ids: torch.LongTensor | None = None,
+        sequence_ids: torch.LongTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        past_key_values: DynamicCache | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        output_s_max: bool = False,
+        return_dict: bool | None = None,
+    ) -> E1ClassificationOutputWithPast | tuple[Any, ...]:
+        use_cache = (
+            use_cache if use_cache is not None else bool(getattr(self.config, "use_cache", False))
+        )
+        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
+        outputs: E1ModelOutputWithPast = self.model(
+            input_ids=input_ids,
+            within_seq_position_ids=within_seq_position_ids,
+            global_position_ids=global_position_ids,
+            sequence_ids=sequence_ids,
+            inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            output_s_max=output_s_max,
+            return_dict=True,
+        )
+
+        attention_mask = (
+            (sequence_ids != -1).long()
+            if sequence_ids is not None
+            else torch.ones(
+                outputs.last_hidden_state.shape[:2],
+                device=outputs.last_hidden_state.device,
+                dtype=torch.long,
+            )
+        )
+        x = outputs.last_hidden_state
+        features = self.pooler(x, attention_mask)
+        logits = self.classifier(features)
+        loss = None
+        if labels is not None:
+            labels = labels.to(logits.device)
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (
+                    labels.dtype == torch.long or labels.dtype == torch.int
+                ):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                if self.num_labels == 1:
+                    loss = self.mse(logits.flatten(), labels.flatten())
+                else:
+                    loss = self.mse(logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss = self.ce(logits.view(-1, self.num_labels), labels.view(-1))
+            elif self.config.problem_type == "multi_label_classification":
+                loss = self.bce(logits, labels)
+
+        result = E1ClassificationOutputWithPast(
+            loss=loss,
+            logits=logits,
+            past_key_values=outputs.past_key_values,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            last_hidden_state=x,
+            s_max=outputs.s_max,
+        )
+        if not return_dict:
+            return result.to_tuple()
+        return result
+
+
+class E1ForTokenClassification(E1PreTrainedModel, EmbeddingMixin):
+    config: E1Config
+    config_class = E1Config
+
+    def __init__(self, config: E1Config, **kwargs):
+        E1PreTrainedModel.__init__(self, config, **kwargs)
+        self.model: FAST_E1_ENCODER = FAST_E1_ENCODER(config, **kwargs)
+        self.vocab_size = config.vocab_size
+        self.num_labels = config.num_labels
+        self.classifier = nn.Sequential(
+            nn.Linear(config.hidden_size, config.hidden_size * 4),
+            nn.GELU(),
+            nn.LayerNorm(config.hidden_size * 4),
+            nn.Linear(config.hidden_size * 4, config.num_labels),
+        )
+        self.loss_fct = nn.CrossEntropyLoss()
+        self.gradient_checkpointing = config.gradient_checkpointing
+        self.post_init()
+
+    @property
+    def device_mesh(self) -> torch.distributed.device_mesh.DeviceMesh:
+        return self.model.device_mesh
+
+    def get_input_embeddings(self) -> nn.Embedding:
+        return self.model.get_input_embeddings()
+
+    def set_input_embeddings(self, value: nn.Embedding) -> None:
+        self.model.set_input_embeddings(value)
+
+    def _embed(
+        self, sequences: list[str], return_attention_mask: bool = False, **kwargs
+    ) -> torch.Tensor:
+        return self.model._embed(sequences, return_attention_mask=return_attention_mask, **kwargs)
+
+    def forward(
+        self,
+        input_ids: torch.LongTensor | None = None,
+        within_seq_position_ids: torch.LongTensor | None = None,
+        global_position_ids: torch.LongTensor | None = None,
+        sequence_ids: torch.LongTensor | None = None,
+        inputs_embeds: torch.FloatTensor | None = None,
+        labels: torch.LongTensor | None = None,
+        past_key_values: DynamicCache | None = None,
+        use_cache: bool | None = None,
+        output_attentions: bool | None = None,
+        output_hidden_states: bool | None = None,
+        output_s_max: bool = False,
+        return_dict: bool | None = None,
+    ) -> E1TokenClassificationOutputWithPast | tuple[Any, ...]:
+        use_cache = (
+            use_cache if use_cache is not None else bool(getattr(self.config, "use_cache", False))
+        )
+        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
+        outputs: E1ModelOutputWithPast = self.model(
+            input_ids=input_ids,
+            within_seq_position_ids=within_seq_position_ids,
+            global_position_ids=global_position_ids,
+            sequence_ids=sequence_ids,
+            inputs_embeds=inputs_embeds,
+            past_key_values=past_key_values,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            output_s_max=output_s_max,
+            return_dict=True,
+        )
+
+        x = outputs.last_hidden_state
+        logits = self.classifier(x)
+        loss = None
+        if labels is not None:
+            labels = labels.to(logits.device)
+            loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        result = E1TokenClassificationOutputWithPast(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+            last_hidden_state=x,
+            past_key_values=outputs.past_key_values,
+            s_max=outputs.s_max,
+        )
+        if not return_dict:
+            return result.to_tuple()
+        return result

fastplms/models/e1/preparation.py

@@ -0,0 +1,267 @@
+"""Tokenizer loading and raw-sequence batch preparation for E1."""
+
+from __future__ import annotations
+
+import itertools
+import os
+from dataclasses import dataclass
+
+import torch
+from tokenizers import Tokenizer
+from torch.nn.utils.rnn import pad_sequence
+
+PAD_TOKEN_ID = 0
+BOS_TOKEN_ID = 1
+EOS_TOKEN_ID = 2
+E1_VOCAB_SIZE = 34
+E1_TOKENIZER_REPO_ID = "Synthyra/Profluent-E1-150M"
+
+
+def _load_tokenizer_file(fname: str) -> Tokenizer:
+    tokenizer: Tokenizer = Tokenizer.from_file(fname)
+    padding = tokenizer.padding
+    actual_pad_id = None if padding is None else padding.get("pad_id")
+    if actual_pad_id != PAD_TOKEN_ID:
+        raise ValueError(
+            f"Padding token id must be {PAD_TOKEN_ID}, but got {actual_pad_id}"
+        )
+    return tokenizer
+
+
+def get_tokenizer(
+    pretrained_model_name_or_path: str | os.PathLike | None = None,
+    *,
+    local_files_only: bool = False,
+    cache_dir: str | os.PathLike | None = None,
+    revision: str | None = None,
+    token: str | bool | None = None,
+) -> Tokenizer:
+    source_path = None
+    checked_local_source = False
+    if pretrained_model_name_or_path is not None:
+        source_path = os.fspath(pretrained_model_name_or_path)
+        if os.path.isdir(source_path):
+            checked_local_source = True
+            fname = os.path.join(source_path, "tokenizer.json")
+            if os.path.isfile(fname):
+                return _load_tokenizer_file(fname)
+
+    fname = os.path.join(os.path.dirname(__file__), "tokenizer.json")
+    if os.path.isfile(fname):
+        return _load_tokenizer_file(fname)
+
+    if local_files_only and checked_local_source:
+        raise FileNotFoundError(
+            f"E1 tokenizer.json was not found in {source_path} or next to {__file__}."
+        )
+
+    from huggingface_hub import hf_hub_download
+
+    repo_id = E1_TOKENIZER_REPO_ID
+    if source_path is not None and not checked_local_source:
+        repo_id = source_path
+    try:
+        fname = hf_hub_download(
+            repo_id=repo_id,
+            filename="tokenizer.json",
+            cache_dir=os.fspath(cache_dir) if cache_dir is not None else None,
+            revision=revision,
+            token=token,
+            local_files_only=local_files_only,
+        )
+    except Exception as error:
+        raise FileNotFoundError(
+            f"E1 tokenizer.json was not found locally and could not be loaded from {repo_id}."
+        ) from error
+    return _load_tokenizer_file(fname)
+
+
+@dataclass
+class DataPrepConfig:
+    max_num_sequences: int = 512
+    max_num_positions_within_seq: int = 8192
+    remove_X_tokens: bool = False
+
+
+def get_context(sequence: str) -> str | None:
+    if "," in sequence:
+        return sequence.rsplit(",", 1)[0]
+    return None
+
+
+class E1BatchPreparer:
+    def __init__(
+        self,
+        data_prep_config: DataPrepConfig | None = None,
+        tokenizer: Tokenizer | None = None,
+        tokenizer_source: str | os.PathLike | None = None,
+        local_files_only: bool = False,
+        cache_dir: str | os.PathLike | None = None,
+        revision: str | None = None,
+        token: str | bool | None = None,
+        preserve_context_labels: bool = False,
+    ):
+        self.tokenizer = tokenizer or get_tokenizer(
+            tokenizer_source,
+            local_files_only=local_files_only,
+            cache_dir=cache_dir,
+            revision=revision,
+            token=token,
+        )
+        self.data_prep_config = data_prep_config or DataPrepConfig()
+        self.pad_token_id = self.tokenizer.token_to_id("<pad>")
+        self.preserve_context_labels = preserve_context_labels
+        self.boundary_token_ids = torch.tensor(
+            [self.tokenizer.token_to_id(token) for token in ["<bos>", "<eos>", "1", "2", "<pad>"]]
+        ).long()
+        self.mask_token = "?"  # nosec
+        self.mask_token_id = self.tokenizer.token_to_id(self.mask_token)
+        self.X_token_id = self.tokenizer.token_to_id("X")
+        self.vocab = self.tokenizer.get_vocab()
+
+    def get_batch_kwargs(  # type: ignore[override]
+        self,
+        sequences: list[str],
+        device: torch.device | None = None,
+        non_blocking: bool = False,
+    ) -> dict[str, torch.Tensor | list[str] | list[int]]:
+        device = torch.device("cpu") if device is None else device
+        sequence_encodings = [self.prepare_multiseq(sequence) for sequence in sequences]
+        return self.pad_encodings(sequence_encodings, device, non_blocking)
+
+    def pad_encodings(
+        self,
+        sequence_encodings: list[dict[str, torch.Tensor]],
+        device: torch.device | None = None,
+        non_blocking: bool = False,
+    ) -> dict[str, torch.Tensor | list[str] | list[int]]:
+        device = torch.device("cpu") if device is None else device
+        non_blocking = non_blocking and device.type == "cuda"
+        padded_encodings = {}
+        # Note: We use -1 as the padding value for sequence and position ids because the 0 value
+        # is a valid value for sequence and position ids. -1 is then used to distinguish valid
+        # tokens from padding tokens, for example, when doing padding/unpadding for flash attention.
+        for key, padding_value in {
+            "input_ids": self.pad_token_id,
+            "sequence_ids": -1,
+            "within_seq_position_ids": -1,
+            "global_position_ids": -1,
+            "labels": self.pad_token_id,
+        }.items():
+            padded_encodings[key] = pad_sequence(
+                [enc[key] for enc in sequence_encodings],
+                batch_first=True,
+                padding_value=padding_value,
+            ).to(device=device, dtype=torch.long, non_blocking=non_blocking)
+
+        padded_encodings["context"] = [enc["context"] for enc in sequence_encodings]
+        padded_encodings["context_len"] = [enc["context_len"] for enc in sequence_encodings]
+
+        return padded_encodings
+
+    def prepare_multiseq(self, sequence: str) -> dict[str, torch.Tensor | str | int]:
+        sequences = sequence.split(",")
+        if len(sequences) > self.data_prep_config.max_num_sequences:
+            raise ValueError(
+                f"Number of sequences {len(sequences)} exceeds max number of sequences "
+                f"{self.data_prep_config.max_num_sequences} in the provided multi-sequence "
+                "instance. Please remove some homologous sequences before trying again."
+            )
+
+        encodings = tuple(self.prepare_singleseq(item) for item in sequences)
+        token_counts = torch.tensor(
+            [encoding["input_ids"].numel() for encoding in encodings],
+            dtype=torch.long,
+        )
+        input_ids = torch.cat(tuple(encoding["input_ids"] for encoding in encodings))
+        labels = torch.cat(tuple(encoding["labels"] for encoding in encodings))
+        positions = tuple(encoding["position_ids"] for encoding in encodings)
+        within_seq_position_ids = torch.cat(positions)
+
+        # Offsets preserve gaps left by optional X-token removal.
+        position_spans = torch.tensor(
+            [int(position_ids[-1].item()) + 1 for position_ids in positions],
+            dtype=torch.long,
+        )
+        position_offsets = torch.cat(
+            (torch.zeros(1, dtype=torch.long), position_spans[:-1]),
+        ).cumsum(dim=0)
+        global_position_ids = torch.cat(
+            tuple(
+                position_ids + offset
+                for position_ids, offset in zip(positions, position_offsets, strict=True)
+            )
+        )
+        sequence_ids = torch.arange(len(encodings), dtype=torch.long).repeat_interleave(
+            token_counts
+        )
+
+        context_len = int(token_counts[:-1].sum().item())
+        context = self.tokenizer.decode(input_ids[:context_len].tolist(), skip_special_tokens=False)
+        if not self.preserve_context_labels:
+            labels[:context_len] = self.pad_token_id
+
+        aligned_tensors = (
+            sequence_ids,
+            within_seq_position_ids,
+            global_position_ids,
+            labels,
+        )
+        if any(tensor.shape != input_ids.shape for tensor in aligned_tensors):
+            raise AssertionError(
+                "Input ids, sequence ids, within seq position ids, global position ids, "
+                "and labels must have the same shape"
+            )
+        if input_ids.numel() < context_len:
+            raise AssertionError(
+                "Input ids must have at least as many tokens as the context length"
+            )
+
+        return {
+            "input_ids": input_ids,
+            "sequence_ids": sequence_ids,
+            "within_seq_position_ids": within_seq_position_ids,
+            "global_position_ids": global_position_ids,
+            "labels": labels,
+            "context": context,
+            "context_len": context_len,
+        }
+
+    def prepare_singleseq(self, sequence: str) -> dict[str, torch.Tensor]:
+        if not self.validate_sequence(sequence):
+            raise ValueError(
+                f"Invalid sequence: {sequence}; Input sequence should contain "
+                "[A-Z] or ? characters only"
+            )
+
+        if len(sequence) > self.data_prep_config.max_num_positions_within_seq:
+            raise ValueError(
+                f"Sequence length {len(sequence)} exceeds max length "
+                f"{self.data_prep_config.max_num_positions_within_seq}"
+            )
+
+        symbols = itertools.chain(("<bos>", "1"), sequence, ("2", "<eos>"))
+        tokens = torch.tensor(
+            [self.vocab[symbol] for symbol in symbols],
+            dtype=torch.long,
+        )
+        position_ids = torch.arange(tokens.numel(), dtype=torch.long)
+
+        if self.data_prep_config.remove_X_tokens:
+            keep = tokens.ne(self.X_token_id)
+            tokens = tokens[keep]
+            position_ids = position_ids[keep]
+
+        return {"input_ids": tokens, "labels": tokens, "position_ids": position_ids}
+
+    def get_boundary_token_mask(self, tokens: torch.Tensor) -> torch.BoolTensor:
+        return torch.isin(tokens, self.boundary_token_ids.to(tokens.device))
+
+    def get_mask_positions_mask(self, tokens: torch.Tensor) -> torch.BoolTensor:
+        return tokens == self.mask_token_id
+
+    def validate_sequence(self, sequence: str) -> bool:
+        if not isinstance(sequence, str):
+            raise TypeError("Sequence must be a string.")
+        sequence = sequence.replace(self.mask_token, "")
+        return sequence.isalpha() and sequence.isupper()

fastplms/models/e1/retrieval.py

@@ -0,0 +1,1783 @@
+"""FASTA, MSA, context sampling, and homologue-search utilities for E1."""
+
+from __future__ import annotations
+
+import hashlib
+import itertools
+import json
+import math
+import numbers
+import os
+import platform
+import random
+import re
+import shutil
+import subprocess
+import tarfile
+import tempfile
+import time
+import urllib.error
+import urllib.parse
+import urllib.request
+from collections import defaultdict, namedtuple
+from collections.abc import Iterator, Sequence
+from dataclasses import dataclass
+from datetime import UTC, datetime
+from email.utils import parsedate_to_datetime
+from pathlib import Path, PurePosixPath, PureWindowsPath
+from typing import TYPE_CHECKING, Any, TypedDict
+
+import numpy as np
+import torch
+from tqdm.auto import tqdm
+from transformers import PreTrainedModel
+from transformers.utils import logging
+
+from fastplms.embeddings import Pooler
+
+from .cache import KVCache
+from .preparation import DataPrepConfig, E1BatchPreparer, get_context
+
+if TYPE_CHECKING:
+    from .modeling_e1 import E1MaskedLMOutputWithPast
+
+
+def _get_logger():
+    """Resolve the Transformers logger only when a retrieval path emits a message."""
+
+    return logging.get_logger(__name__)
+
+
+MMSEQS2_IMAGE_REPOSITORY = "ghcr.io/soedinglab/mmseqs2"
+MMSEQS2_VERSION = "18-8cc5c"
+MMSEQS2_CPU_MANIFEST_DIGEST = (
+    "sha256:41b12b0d5f41432fa1b9976123da6e2e06e7fab49a34964f3b54ec038e5845d9"
+)
+MMSEQS2_CPU_ARM64_CHILD_DIGEST = (
+    "sha256:8bec048845f8f20749c2e2ad067a27d67eef839d2bb068e9d6e957113e9a7fba"
+)
+DOCKER_IMAGE = (
+    f"{MMSEQS2_IMAGE_REPOSITORY}:{MMSEQS2_VERSION}@{MMSEQS2_CPU_MANIFEST_DIGEST}"
+)
+DEFAULT_MMSEQS2_PHASE_TIMEOUT = 1800.0
+COLABFOLD_HOST = "https://api.colabfold.com"
+LOWERCASE_CHARS = b"abcdefghijklmnopqrstuvwxyz"
+DEFAULT_MAX_CONTEXT_TOKENS = [6144, 12288, 24576]
+DEFAULT_SIMILARITY_THRESHOLDS = [1.0, 0.95, 0.9, 0.7, 0.5]
+DEFAULT_EMBED_MAX_TOKENS = 8192
+DEFAULT_EMBED_SIMILARITY = 0.95
+E1_MSA_SAMPLING_SOURCE_REVISION = "bfd2620a602248499f3d2583d85a7ecddf0b6e02"
+
+IdSequence = namedtuple("IdSequence", ["id", "sequence"])
+IndexedSequence = tuple[int, str]
+
+_SHA256_DIGEST_RE = re.compile(r"^sha256:[0-9a-f]{64}$")
+_IMAGE_VERSION_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
+_SAFE_SEQUENCE_ID_RE = re.compile(r"^[A-Za-z0-9][A-Za-z0-9._-]*$")
+_SAFE_QUERY_SEQUENCE_RE = re.compile(r"^[A-Za-z*.-]+$")
+
+
+@dataclass(frozen=True, slots=True)
+class _PinnedImageReference:
+    repository: str
+    version: str
+    digest: str
+
+
+@dataclass(frozen=True, slots=True)
+class _DockerImageIdentity:
+    reference: str
+    repository: str
+    version: str
+    manifest_digest: str
+    image_id: str
+    os: str
+    architecture: str
+
+    def to_dict(self) -> dict[str, str]:
+        return {
+            "reference": self.reference,
+            "repository": self.repository,
+            "version": self.version,
+            "manifest_digest": self.manifest_digest,
+            "image_id": self.image_id,
+            "os": self.os,
+            "architecture": self.architecture,
+        }
+
+
+def _parse_pinned_image_reference(reference: str) -> _PinnedImageReference:
+    """Parse ``repository:version@sha256:digest`` and reject mutable images."""
+
+    if not isinstance(reference, str) or not reference or any(char.isspace() for char in reference):
+        raise ValueError(
+            "docker_image must be an immutable repository:version@sha256:digest reference"
+        )
+    try:
+        name_and_version, digest = reference.rsplit("@", maxsplit=1)
+    except ValueError as error:
+        raise ValueError(
+            "docker_image must include an immutable @sha256 digest; mutable tags are rejected"
+        ) from error
+    last_slash = name_and_version.rfind("/")
+    last_colon = name_and_version.rfind(":")
+    if last_colon <= last_slash:
+        raise ValueError("docker_image must include an explicit version tag before its digest")
+    repository = name_and_version[:last_colon]
+    version = name_and_version[last_colon + 1 :]
+    if (
+        not repository
+        or repository.endswith("/")
+        or "@" in repository
+        or _IMAGE_VERSION_RE.fullmatch(version) is None
+    ):
+        raise ValueError("docker_image contains an invalid repository or version tag")
+    if _SHA256_DIGEST_RE.fullmatch(digest) is None:
+        raise ValueError("docker_image must include a lowercase sha256 digest")
+    return _PinnedImageReference(repository=repository, version=version, digest=digest)
+
+
+def _docker_architecture() -> str:
+    machine = platform.machine().lower()
+    aliases = {
+        "aarch64": "arm64",
+        "arm64": "arm64",
+        "amd64": "amd64",
+        "x86_64": "amd64",
+    }
+    try:
+        return aliases[machine]
+    except KeyError as error:
+        raise RuntimeError(f"Unsupported Docker host architecture: {machine!r}") from error
+
+
+def _json_sha256(payload: Any) -> str:
+    encoded = json.dumps(payload, sort_keys=True, separators=(",", ":")).encode("utf-8")
+    return hashlib.sha256(encoded).hexdigest()
+
+
+def _file_sha256(path: str) -> str:
+    hasher = hashlib.sha256()
+    with open(path, "rb") as handle:
+        for block in iter(lambda: handle.read(1024 * 1024), b""):
+            hasher.update(block)
+    return hasher.hexdigest()
+
+
+def _sequence_output_dir(output_dir: str, seq_id: str) -> str:
+    """Return the per-sequence directory after enforcing path containment."""
+
+    if not isinstance(seq_id, str) or _SAFE_SEQUENCE_ID_RE.fullmatch(seq_id) is None:
+        raise ValueError(
+            "seq_id must use only ASCII letters, digits, dot, underscore, and hyphen"
+        )
+    if (
+        seq_id in {".", ".."}
+        or PurePosixPath(seq_id).name != seq_id
+        or PureWindowsPath(seq_id).name != seq_id
+        or PureWindowsPath(seq_id).is_absolute()
+    ):
+        raise ValueError(f"seq_id must be a single relative filename component: {seq_id!r}")
+
+    output_root = Path(output_dir).resolve()
+    sequence_dir = Path(output_dir) / seq_id
+    resolved_sequence_dir = sequence_dir.resolve()
+    if resolved_sequence_dir.parent != output_root:
+        raise ValueError(f"seq_id resolves outside output_dir: {seq_id!r}")
+    return os.fspath(sequence_dir)
+
+
+@dataclass
+class ContextSpecification:
+    max_num_samples: int = 511
+    max_token_length: int = 32768
+    max_query_similarity: float = 1.0
+    min_query_similarity: float = 0.0
+    neighbor_similarity_lower_bound: float = 0.8
+
+
+class E1Prediction(TypedDict, total=False):
+    id: str | int
+    context_id: str | int | None
+    logits: torch.Tensor
+    token_embeddings: torch.Tensor
+    mean_token_embeddings: torch.Tensor
+
+
+def read_fasta_sequences(path: str) -> dict[str, str]:
+    sequences: dict[str, str] = {}
+    header: str | None = None
+    parts: list[str] = []
+    with open(path, encoding="utf-8") as handle:
+        for raw_line in handle:
+            line = raw_line.strip()
+            if not line:
+                continue
+            if line.startswith(">"):
+                if header is not None:
+                    sequences[header] = "".join(parts)
+                header = line[1:].strip()
+                parts = []
+            else:
+                if header is None:
+                    raise ValueError(f"FASTA sequence found before header in {path}")
+                parts.append(line)
+    if header is not None:
+        sequences[header] = "".join(parts)
+    return sequences
+
+
+def write_fasta_sequences(path: str, sequences: dict[str, str]) -> None:
+    os.makedirs(os.path.dirname(path) or ".", exist_ok=True)
+    with open(path, "w", encoding="utf-8") as handle:
+        for header, sequence in sequences.items():
+            handle.write(f">{header}\n{sequence}\n")
+
+
+def parse_msa(path: str) -> list[IdSequence]:
+    records = read_fasta_sequences(path)
+    sequences = []
+    for record_id, record_seq in records.items():
+        sequence = str(record_seq).replace("\x00", "").replace(".", "-")
+        sequences.append(IdSequence(record_id, sequence))
+    if not sequences:
+        raise ValueError(f"No sequences found in MSA file: {path}")
+    return sequences
+
+
+def convert_to_tensor(
+    sequences: list[IdSequence], device: torch.device | None = None
+) -> torch.ByteTensor:
+    if device is None:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    byte_sequences = [
+        sequence.sequence.encode("ascii").translate(None, LOWERCASE_CHARS) for sequence in sequences
+    ]
+    lengths = {len(byte_sequence) for byte_sequence in byte_sequences}
+    if len(lengths) != 1:
+        raise ValueError(
+            "MSA rows must have equal aligned lengths after removing insertions: "
+            f"{sorted(lengths)}"
+        )
+    array = np.vstack(
+        [np.frombuffer(byte_sequence, dtype=np.uint8) for byte_sequence in byte_sequences]
+    )
+    return torch.from_numpy(array).to(device)
+
+
+def get_num_neighbors(byte_seqs: torch.ByteTensor, sim_threshold: float = 0.8) -> list[int]:
+    gap_token_id = np.frombuffer(b"-", np.uint8)[0].item()
+    seq_lens = (byte_seqs != gap_token_id).sum(dim=1)
+    num_neighbors: list[int] = []
+    for i in range(byte_seqs.shape[0]):
+        query_non_gaps = byte_seqs[i] != gap_token_id
+        seqs_sim = (byte_seqs[:, query_non_gaps] == byte_seqs[i, query_non_gaps]).sum(
+            dim=1
+        ) / seq_lens
+        num_neighbors.append(int((seqs_sim >= sim_threshold).sum().item()))
+    return num_neighbors
+
+
+def get_similarity_to_query(byte_seqs: torch.ByteTensor) -> torch.FloatTensor:
+    return (byte_seqs == byte_seqs[0, :]).sum(dim=1) / byte_seqs.shape[1]
+
+
+def sample_context(
+    msa_path: str,
+    max_num_samples: int,
+    max_token_length: int,
+    max_query_similarity: float = 1.0,
+    min_query_similarity: float = 0.0,
+    neighbor_similarity_lower_bound: float = 0.8,
+    use_full_sequences_in_context: bool = False,
+    full_sequences_path: str | None = None,
+    seed: int = 0,
+    device: torch.device | None = None,
+    cache_num_neighbors_path: str | None = None,
+) -> tuple[str, list[str]]:
+    msa_sequences = parse_msa(msa_path)
+    msa_as_byte_tensor = convert_to_tensor(msa_sequences, device)
+    if cache_num_neighbors_path is not None and os.path.exists(cache_num_neighbors_path):
+        num_neighbors = np.load(cache_num_neighbors_path)
+    else:
+        num_neighbors = np.array(
+            get_num_neighbors(msa_as_byte_tensor, neighbor_similarity_lower_bound)
+        )
+        if cache_num_neighbors_path is not None:
+            np.save(cache_num_neighbors_path, num_neighbors)
+
+    sampling_weights = 1.0 / num_neighbors
+    query_similarity = get_similarity_to_query(msa_as_byte_tensor)
+    filtered_mask = (query_similarity <= max_query_similarity) & (
+        query_similarity >= min_query_similarity
+    )
+    if int(filtered_mask.sum()) < 1:
+        raise ValueError(
+            "No sequences found with similarity to query within range "
+            f"{min_query_similarity} <= query_similarity <= {max_query_similarity}."
+        )
+
+    filtered_weights = np.where(filtered_mask.cpu().numpy(), sampling_weights, 0.0)
+    sampled_indices = np.random.default_rng(seed).choice(
+        len(filtered_weights),
+        size=min(max_num_samples, int(filtered_mask.sum())),
+        p=filtered_weights / filtered_weights.sum(),
+        replace=False,
+        shuffle=True,
+    )
+
+    if use_full_sequences_in_context:
+        if full_sequences_path is None:
+            raise ValueError(
+                "full_sequences_path is required when use_full_sequences_in_context=True"
+            )
+        full_sequences = parse_msa(full_sequences_path)
+        if len(full_sequences) != len(msa_sequences):
+            raise ValueError("Number of full sequences must match number of MSA sequences")
+        for i, (full_seq, msa_seq) in enumerate(zip(full_sequences, msa_sequences, strict=True)):
+            if full_seq.id != msa_seq.id:
+                raise ValueError(
+                    "Full sequences and MSA sequences must be in the same order and have the "
+                    f"same ids. Found differing id for sample {i}: "
+                    f"{full_seq.id} != {msa_seq.id}"
+                )
+        sampled_sequences = [full_sequences[int(i)] for i in sampled_indices]
+    else:
+        sampled_sequences = [msa_sequences[int(i)] for i in sampled_indices]
+
+    context_sequences: list[str] = []
+    context_ids: list[str] = []
+    context_length = 0
+    for seq in sampled_sequences:
+        seq_str = seq.sequence.upper().encode("ascii").translate(None, b"-").decode("ascii")
+        if context_length + len(seq_str) > max_token_length:
+            break
+        context_sequences.append(seq_str)
+        context_ids.append(seq.id)
+        context_length += len(seq_str)
+    return ",".join(context_sequences), context_ids
+
+
+def sample_multiple_contexts(
+    msa_path: str,
+    context_specifications: list[ContextSpecification],
+    use_full_sequences_in_context: bool = False,
+    full_sequences_path: str | None = None,
+    seed: int = 0,
+    device: torch.device | None = None,
+    cache_num_neighbors_path: str | None = None,
+) -> tuple[list[str], list[list[str]]]:
+    with tempfile.TemporaryDirectory() as temp_dir:
+        if cache_num_neighbors_path is None:
+            cache_num_neighbors_path = os.path.join(temp_dir, "num_neighbors.npy")
+
+        contexts: list[str] = []
+        context_ids: list[list[str]] = []
+        for i, context_specification in enumerate(context_specifications):
+            context, ids = sample_context(
+                msa_path=msa_path,
+                max_num_samples=context_specification.max_num_samples,
+                max_token_length=context_specification.max_token_length,
+                max_query_similarity=context_specification.max_query_similarity,
+                min_query_similarity=context_specification.min_query_similarity,
+                neighbor_similarity_lower_bound=context_specification.neighbor_similarity_lower_bound,
+                use_full_sequences_in_context=use_full_sequences_in_context,
+                full_sequences_path=full_sequences_path,
+                seed=seed + i,
+                device=device,
+                cache_num_neighbors_path=cache_num_neighbors_path,
+            )
+            contexts.append(context)
+            context_ids.append(ids)
+    return contexts, context_ids
+
+
+def get_context_id(max_tokens: int, sim_threshold: float) -> str:
+    return f"identity_{sim_threshold}_tokens_{max_tokens}"
+
+
+def build_context_specifications(
+    max_context_tokens: list[int] | None = None,
+    similarity_thresholds: list[float] | None = None,
+    min_query_similarity: float = 0.3,
+) -> list[tuple[ContextSpecification, str]]:
+    if max_context_tokens is None:
+        max_context_tokens = DEFAULT_MAX_CONTEXT_TOKENS
+    if similarity_thresholds is None:
+        similarity_thresholds = DEFAULT_SIMILARITY_THRESHOLDS
+
+    specs = []
+    for max_tokens in max_context_tokens:
+        for sim_threshold in similarity_thresholds:
+            spec = ContextSpecification(
+                max_num_samples=511,
+                max_token_length=max_tokens,
+                max_query_similarity=sim_threshold,
+                min_query_similarity=min_query_similarity,
+                neighbor_similarity_lower_bound=0.8,
+            )
+            specs.append((spec, get_context_id(max_tokens, sim_threshold)))
+    return specs
+
+
+def sample_contexts_for_msa(
+    a3m_path: str,
+    context_specs: list[tuple[ContextSpecification, str]],
+    seed: int = 42,
+) -> dict[str, str]:
+    specs_only = [spec for spec, _ in context_specs]
+    context_ids = [context_id for _, context_id in context_specs]
+    contexts, _ = sample_multiple_contexts(
+        msa_path=a3m_path,
+        context_specifications=specs_only,
+        seed=seed,
+    )
+    return dict(zip(context_ids, contexts, strict=True))
+
+
+def _strip_a3m_insertions(sequence: str) -> str:
+    uppercase_or_gap = [char for char in sequence if char.isupper() or char in "-."]
+    return "".join(uppercase_or_gap).replace("-", "").replace(".", "")
+
+
+def get_query_from_a3m(path: str) -> str:
+    header_found = False
+    seq_parts: list[str] = []
+    with open(path, encoding="utf-8") as handle:
+        for raw_line in handle:
+            line = raw_line.strip()
+            if not line:
+                continue
+            if line.startswith(">"):
+                if header_found:
+                    break
+                header_found = True
+                continue
+            if header_found:
+                seq_parts.append(line)
+    if not header_found:
+        raise ValueError(f"No FASTA header found in A3M file: {path}")
+    return _strip_a3m_insertions("".join(seq_parts))
+
+
+def load_msa_dir(msa_dir: str) -> dict[str, str]:
+    msa_lookup: dict[str, str] = {}
+    a3m_files = list(Path(msa_dir).rglob("*.a3m"))
+    if not a3m_files:
+        raise FileNotFoundError(f"No .a3m files found in {msa_dir}")
+    for a3m_path in tqdm(a3m_files, desc="Loading MSAs"):
+        query_seq = get_query_from_a3m(str(a3m_path))
+        msa_lookup[query_seq] = str(a3m_path)
+    _get_logger().info("Loaded %d MSAs from %s", len(msa_lookup), msa_dir)
+    return msa_lookup
+
+
+def _safe_extract_tar(tar: tarfile.TarFile, output_dir: str) -> None:
+    output_root = Path(output_dir).resolve()
+    for member in tar.getmembers():
+        if member.issym() or member.islnk():
+            raise ValueError(f"Tar links are not allowed: {member.name}")
+        if member.isdev():
+            raise ValueError(f"Tar device entries are not allowed: {member.name}")
+        target = (output_root / member.name).resolve()
+        if output_root != target and output_root not in target.parents:
+            raise ValueError(f"Unsafe tar member path: {member.name}")
+    tar.extractall(output_root, filter="data")
+
+
+def load_msa_from_hf(
+    hf_path: str,
+    cache_dir: str | None = None,
+    token: str | None = None,
+) -> dict[str, str]:
+    from huggingface_hub import snapshot_download
+
+    if cache_dir is None:
+        cache_dir = os.path.join(os.path.expanduser("~"), ".cache", "fastplms_msa")
+    os.makedirs(cache_dir, exist_ok=True)
+    local_dir = os.path.join(cache_dir, hf_path.replace("/", "_"))
+    if not os.path.exists(local_dir) or not any(Path(local_dir).rglob("*.a3m")):
+        local_dir = snapshot_download(
+            repo_id=hf_path,
+            repo_type="dataset",
+            local_dir=local_dir,
+            token=token,
+        )
+        for tar_path in Path(local_dir).rglob("*.tar.gz"):
+            with tarfile.open(tar_path) as tar:
+                _safe_extract_tar(tar, str(tar_path.parent))
+    return load_msa_dir(local_dir)
+
+
+def get_msa_for_sequence(
+    sequence: str, msa_lookup: dict[str, str], min_identity: float = 0.95
+) -> str | None:
+    if sequence in msa_lookup:
+        return msa_lookup[sequence]
+
+    best_match_path: str | None = None
+    best_identity = 0.0
+    for query_seq, a3m_path in msa_lookup.items():
+        if abs(len(query_seq) - len(sequence)) > 10:
+            continue
+        min_len = min(len(query_seq), len(sequence))
+        if min_len == 0:
+            continue
+        matches = sum(a == b for a, b in zip(query_seq[:min_len], sequence[:min_len], strict=True))
+        identity = matches / min_len
+        if identity > best_identity:
+            best_identity = identity
+            best_match_path = a3m_path
+
+    if best_identity >= min_identity:
+        return best_match_path
+    return None
+
+
+class ContextCache:
+    """Content-addressed JSON cache for deterministic E1 MSA contexts."""
+
+    _SCHEMA_VERSION = 1
+
+    def __init__(
+        self,
+        cache_dir: str,
+        specs_hash: str,
+        seed: int,
+        source_revision: str = E1_MSA_SAMPLING_SOURCE_REVISION,
+    ) -> None:
+        self.cache_dir = cache_dir
+        self.specs_hash = specs_hash
+        self.seed = seed
+        self.source_revision = source_revision
+        os.makedirs(cache_dir, exist_ok=True)
+
+    def _cache_path(self, key: str) -> str:
+        safe_key = hashlib.sha256(key.encode("utf-8")).hexdigest()[:16]
+        return os.path.join(self.cache_dir, f"{safe_key}_seed{self.seed}_{self.specs_hash}.json")
+
+    def _input_fingerprint(self, key: str) -> str:
+        descriptor: dict[str, Any] = {
+            "key": os.path.abspath(key) if os.path.isfile(key) else key,
+            "seed": self.seed,
+            "source_revision": self.source_revision,
+            "specs_hash": self.specs_hash,
+        }
+        if os.path.isfile(key):
+            hasher = hashlib.sha256()
+            with open(key, "rb") as handle:
+                for block in iter(lambda: handle.read(1024 * 1024), b""):
+                    hasher.update(block)
+            descriptor["content_sha256"] = hasher.hexdigest()
+        else:
+            descriptor["literal_key"] = key
+        payload = json.dumps(descriptor, sort_keys=True, separators=(",", ":"))
+        return hashlib.sha256(payload.encode("utf-8")).hexdigest()
+
+    def load(self, key: str) -> dict[str, str] | None:
+        path = self._cache_path(key)
+        if not os.path.exists(path):
+            return None
+        try:
+            with open(path, encoding="utf-8") as handle:
+                payload = json.load(handle)
+        except (OSError, UnicodeError, json.JSONDecodeError):
+            return None
+        if not isinstance(payload, dict):
+            return None
+        if payload.get("schema_version") != self._SCHEMA_VERSION:
+            return None
+        if payload.get("input_fingerprint") != self._input_fingerprint(key):
+            return None
+        if payload.get("source_revision") != self.source_revision:
+            return None
+        contexts = payload.get("contexts")
+        if not isinstance(contexts, dict) or not all(
+            isinstance(name, str) and isinstance(context, str) for name, context in contexts.items()
+        ):
+            return None
+        return contexts
+
+    def store(self, key: str, contexts: dict[str, str]) -> None:
+        if not all(
+            isinstance(name, str) and isinstance(context, str) for name, context in contexts.items()
+        ):
+            raise TypeError("contexts must map string identifiers to string contexts")
+        path = self._cache_path(key)
+        payload = {
+            "schema_version": self._SCHEMA_VERSION,
+            "source_revision": self.source_revision,
+            "input_fingerprint": self._input_fingerprint(key),
+            "contexts": contexts,
+        }
+        temp_path: str | None = None
+        try:
+            with tempfile.NamedTemporaryFile(
+                mode="w",
+                encoding="utf-8",
+                dir=self.cache_dir,
+                prefix=".context-",
+                suffix=".tmp",
+                delete=False,
+            ) as handle:
+                temp_path = handle.name
+                json.dump(payload, handle, sort_keys=True, separators=(",", ":"))
+                handle.write("\n")
+                handle.flush()
+                os.fsync(handle.fileno())
+            os.replace(temp_path, path)
+            temp_path = None
+        finally:
+            if temp_path is not None:
+                Path(temp_path).unlink(missing_ok=True)
+
+
+def compute_ppll(logits: torch.Tensor, token_ids: torch.Tensor) -> float:
+    if token_ids.numel() == 0:
+        raise ValueError("Cannot score an empty token sequence")
+    if token_ids.device != logits.device:
+        token_ids = token_ids.to(logits.device)
+    if logits.shape[0] != token_ids.shape[0]:
+        raise ValueError(
+            f"Logits length {logits.shape[0]} != token_ids length {token_ids.shape[0]}"
+        )
+    probs = logits.softmax(dim=-1)
+    token_probs = probs.gather(dim=1, index=token_ids.unsqueeze(1)).squeeze(1)
+    return float(token_probs.mean().item())
+
+
+class _E1ContextPredictor:
+    def __init__(
+        self,
+        model: PreTrainedModel,
+        data_prep_config: DataPrepConfig | None = None,
+        max_batch_tokens: int = 65536,
+        use_cache: bool = True,
+        cache_size: int = 4,
+        save_masked_positions_only: bool = False,
+        fields_to_save: list[str] | None = None,
+        keep_predictions_in_gpu: bool = False,
+        progress: bool = True,
+    ) -> None:
+        self.model = model
+        self.max_batch_tokens = max_batch_tokens
+        self.batch_preparer = E1BatchPreparer(data_prep_config=data_prep_config)
+        self.model.eval()
+        self.kv_cache = KVCache(cache_size=cache_size) if use_cache else None
+        self.fields_to_save = fields_to_save or [
+            "logits",
+            "token_embeddings",
+            "mean_token_embeddings",
+        ]
+        self.save_masked_positions_only = save_masked_positions_only
+        self.keep_predictions_in_gpu = keep_predictions_in_gpu
+        self.progress = progress
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.model.parameters()).device
+
+    def group_by_length(
+        self, indexed_sequences: list[IndexedSequence]
+    ) -> list[list[IndexedSequence]]:
+        batches: list[list[IndexedSequence]] = [[]]
+        for idx, seq in sorted(
+            indexed_sequences, key=lambda idx_seq: (len(idx_seq[1]), idx_seq[0])
+        ):
+            if len(batches[-1]) > 0 and len(seq) * (len(batches[-1]) + 1) > self.max_batch_tokens:
+                batches.append([])
+            batches[-1].append((idx, seq))
+        return batches
+
+    def group_by_context(
+        self, indexed_sequences: list[IndexedSequence]
+    ) -> list[list[IndexedSequence]]:
+        batches: dict[str | None, list[IndexedSequence]] = defaultdict(list)
+        for idx, seq in indexed_sequences:
+            batches[get_context(seq)].append((idx, seq))
+        return list(batches.values())
+
+    def batch_sequences(self, sequences: list[str]) -> list[list[int]]:
+        indexed_sequences: list[IndexedSequence] = list(enumerate(sequences))
+        indexed_batches = self.group_by_context(indexed_sequences)
+        indexed_batches = list(
+            itertools.chain.from_iterable(
+                [self.group_by_length(batch) for batch in indexed_batches]
+            )
+        )
+        batches = [[item[0] for item in batch] for batch in indexed_batches]
+        flattened_indices = list(itertools.chain.from_iterable(batches))
+        if sorted(flattened_indices) != list(range(len(sequences))):
+            raise RuntimeError("Batches must contain all indices with no repetition")
+        return batches
+
+    @torch.no_grad()
+    def predict_batch(
+        self, sequences: list[str], sequence_metadata: list[dict[str, str | int]]
+    ) -> list[E1Prediction]:
+        outputs = self.predict_batch_padded(sequences)
+        outputs["logits"] = outputs["logits"].float()
+        outputs["embeddings"] = outputs["embeddings"].float()
+
+        token_mask = outputs["non_boundary_token_mask"] & outputs["last_sequence_mask"]
+        if self.save_masked_positions_only:
+            token_mask = token_mask & outputs["mask_positions_mask"]
+
+        predictions: list[E1Prediction] = []
+        for i in range(len(sequences)):
+            pred: E1Prediction = {"id": sequence_metadata[i]["id"]}
+            if "context_id" in sequence_metadata[i]:
+                pred["context_id"] = sequence_metadata[i]["context_id"]
+            if "logits" in self.fields_to_save:
+                pred["logits"] = outputs["logits"][i, token_mask[i]]
+                if not self.keep_predictions_in_gpu:
+                    pred["logits"] = pred["logits"].to("cpu")
+            if "token_embeddings" in self.fields_to_save:
+                pred["token_embeddings"] = outputs["embeddings"][i, token_mask[i]]
+                if not self.keep_predictions_in_gpu:
+                    pred["token_embeddings"] = pred["token_embeddings"].to("cpu")
+            if "mean_token_embeddings" in self.fields_to_save:
+                pred["mean_token_embeddings"] = outputs["embeddings"][i, token_mask[i]].mean(dim=0)
+                if not self.keep_predictions_in_gpu:
+                    pred["mean_token_embeddings"] = pred["mean_token_embeddings"].to("cpu")
+            predictions.append(pred)
+        return predictions
+
+    @torch.no_grad()
+    def predict_batch_padded(self, sequences: list[str]) -> dict[str, torch.Tensor]:
+        device = self.device
+        autocast_enabled = device.type == "cuda"
+        with torch.autocast(device.type, torch.bfloat16, enabled=autocast_enabled):
+            batch = self.batch_preparer.get_batch_kwargs(sequences, device=device)
+            if self.kv_cache is not None:
+                self.kv_cache.before_forward(batch)
+
+            past_key_values = batch.get("past_key_values")
+            use_cache = bool(batch["use_cache"]) if "use_cache" in batch else False
+            output: E1MaskedLMOutputWithPast = self.model(
+                input_ids=batch["input_ids"],
+                within_seq_position_ids=batch["within_seq_position_ids"],
+                global_position_ids=batch["global_position_ids"],
+                sequence_ids=batch["sequence_ids"],
+                past_key_values=past_key_values,
+                use_cache=use_cache,
+                output_attentions=False,
+                output_hidden_states=False,
+            )
+            if self.kv_cache is not None:
+                self.kv_cache.after_forward(batch, output)
+
+            padding_mask = batch["input_ids"] == self.batch_preparer.pad_token_id
+            last_sequence_mask = (
+                batch["sequence_ids"] == batch["sequence_ids"].max(dim=1).values[:, None]
+            )
+            boundary_token_mask = self.batch_preparer.get_boundary_token_mask(batch["input_ids"])
+            mask_positions_mask = self.batch_preparer.get_mask_positions_mask(batch["input_ids"])
+            return {
+                "logits": output.logits,
+                "embeddings": output.last_hidden_state,
+                "last_sequence_mask": last_sequence_mask,
+                "non_boundary_token_mask": ~boundary_token_mask,
+                "mask_positions_mask": mask_positions_mask,
+                "valid_token_mask": ~padding_mask,
+            }
+
+    @torch.no_grad()
+    def predict(
+        self,
+        sequences: Sequence[str],
+        sequence_ids: Sequence[int | str] | None = None,
+        context_seqs: dict[str, str] | None = None,
+    ) -> Iterator[E1Prediction]:
+        if sequence_ids is None:
+            sequence_ids = list(range(len(sequences)))
+        if context_seqs:
+            sequences_with_context = [
+                (ctx + "," + seq, {"context_id": ctx_id, "id": sequence_id})
+                for ctx_id, ctx in context_seqs.items()
+                for seq, sequence_id in zip(sequences, sequence_ids, strict=True)
+            ]
+        else:
+            sequences_with_context = [
+                (seq, {"id": sequence_id})
+                for seq, sequence_id in zip(sequences, sequence_ids, strict=True)
+            ]
+
+        batched_sequences, sequence_metadata = tuple(zip(*sequences_with_context, strict=True))
+        batches = self.batch_sequences(list(batched_sequences))
+        iterator = tqdm(batches, desc="Predicting batches", disable=not self.progress)
+        for indices in iterator:
+            sequence_batch = [batched_sequences[i] for i in indices]
+            sequence_batch_metadata = [sequence_metadata[i] for i in indices]
+            yield from self.predict_batch(sequence_batch, sequence_batch_metadata)
+
+
+def _pool_hidden_states(
+    hidden_list: list[torch.Tensor],
+    pooling_types: list[str],
+    device: torch.device,
+) -> torch.Tensor:
+    pooler = Pooler(pooling_types)
+    max_len = max(hidden.shape[0] for hidden in hidden_list)
+    hidden_dim = hidden_list[0].shape[1]
+    batch_size = len(hidden_list)
+    padded = torch.zeros(batch_size, max_len, hidden_dim, device=device)
+    attention_mask = torch.zeros(batch_size, max_len, device=device)
+    for i, hidden in enumerate(hidden_list):
+        seq_len = hidden.shape[0]
+        padded[i, :seq_len] = hidden
+        attention_mask[i, :seq_len] = 1.0
+    return pooler(padded, attention_mask)
+
+
+def _forward_for_embedding(
+    model: PreTrainedModel,
+    sequences: list[str],
+    context: str | None,
+    max_batch_tokens: int,
+    progress: bool,
+) -> list[torch.Tensor]:
+    predictor = _E1ContextPredictor(
+        model=model,
+        data_prep_config=DataPrepConfig(remove_X_tokens=True),
+        max_batch_tokens=max_batch_tokens,
+        fields_to_save=["token_embeddings"],
+        keep_predictions_in_gpu=True,
+        use_cache=False,
+        cache_size=1,
+        progress=progress,
+    )
+    context_seqs = {"embed_ctx": context} if context else None
+    predictions = list(
+        predictor.predict(
+            sequences=sequences,
+            sequence_ids=list(range(len(sequences))),
+            context_seqs=context_seqs,
+        )
+    )
+    predictions.sort(key=lambda prediction: prediction["id"])
+    return [prediction["token_embeddings"] for prediction in predictions]
+
+
+class HomologueSearcher:
+    """Run local MMseqs2 searches through one verified, digest-pinned image.
+
+    The default CPU image is multi-architecture and immutable. Pulling and
+    container networking are separate explicit opt-ins. GPU execution requires
+    a caller-supplied digest-pinned GPU image because the official CUDA image is
+    not portable to every supported host architecture.
+    """
+
+    _PROVENANCE_SCHEMA_VERSION = 1
+    _PROVENANCE_FILENAME = "search-provenance.json"
+
+    def __init__(
+        self,
+        target_db: str,
+        docker_image: str = DOCKER_IMAGE,
+        sensitivity: float = 7.5,
+        max_seqs: int = 1000,
+        min_seq_id: float = 0.0,
+        coverage: float = 0.8,
+        split_memory_limit: str | None = None,
+        use_gpu: bool = False,
+        allow_pull: bool = False,
+        allow_network: bool = False,
+        phase_timeout: float = DEFAULT_MMSEQS2_PHASE_TIMEOUT,
+        target_db_identity: str | None = None,
+    ) -> None:
+        image_reference = _parse_pinned_image_reference(docker_image)
+        if not isinstance(target_db, str) or not target_db or "\x00" in target_db:
+            raise ValueError("target_db must be a non-empty path without null bytes")
+        numeric_values = {
+            "sensitivity": sensitivity,
+            "min_seq_id": min_seq_id,
+            "coverage": coverage,
+        }
+        for name, value in numeric_values.items():
+            if (
+                isinstance(value, bool)
+                or not isinstance(value, numbers.Real)
+                or not math.isfinite(float(value))
+            ):
+                raise ValueError(f"{name} must be a finite real number")
+        if sensitivity <= 0:
+            raise ValueError("sensitivity must be positive")
+        if not 0.0 <= min_seq_id <= 1.0:
+            raise ValueError("min_seq_id must be in [0, 1]")
+        if not 0.0 <= coverage <= 1.0:
+            raise ValueError("coverage must be in [0, 1]")
+        if isinstance(max_seqs, bool) or not isinstance(max_seqs, int) or max_seqs < 1:
+            raise ValueError("max_seqs must be an integer >= 1")
+        if split_memory_limit is not None and (
+            not isinstance(split_memory_limit, str)
+            or not split_memory_limit.strip()
+            or "\x00" in split_memory_limit
+        ):
+            raise ValueError("split_memory_limit must be None or a non-empty string")
+        if type(use_gpu) is not bool:
+            raise TypeError("use_gpu must be a boolean")
+        if type(allow_pull) is not bool:
+            raise TypeError("allow_pull must be a boolean")
+        if type(allow_network) is not bool:
+            raise TypeError("allow_network must be a boolean")
+        if (
+            isinstance(phase_timeout, bool)
+            or not isinstance(phase_timeout, (int, float))
+            or not math.isfinite(float(phase_timeout))
+            or phase_timeout <= 0
+        ):
+            raise ValueError("phase_timeout must be a finite positive number")
+        if target_db_identity is not None and (
+            not isinstance(target_db_identity, str) or not target_db_identity.strip()
+        ):
+            raise ValueError("target_db_identity must be None or a non-empty string")
+        if (
+            use_gpu
+            and image_reference.repository == MMSEQS2_IMAGE_REPOSITORY
+            and image_reference.digest == MMSEQS2_CPU_MANIFEST_DIGEST
+        ):
+            raise ValueError(
+                "The default MMseqs2 image is CPU-only. GPU search requires an explicit "
+                "digest-pinned image compatible with the host architecture."
+            )
+        self.target_db = target_db
+        self.docker_image = docker_image
+        self._image_reference = image_reference
+        self.sensitivity = float(sensitivity)
+        self.max_seqs = max_seqs
+        self.min_seq_id = float(min_seq_id)
+        self.coverage = float(coverage)
+        self.split_memory_limit = (
+            split_memory_limit.strip() if split_memory_limit is not None else None
+        )
+        self.use_gpu = use_gpu
+        self.allow_pull = allow_pull
+        self.allow_network = allow_network
+        self.phase_timeout = float(phase_timeout)
+        self.target_db_identity = (
+            target_db_identity.strip() if target_db_identity is not None else None
+        )
+        self._verified_image_identity: _DockerImageIdentity | None = None
+
+    @staticmethod
+    def _seq_hash(sequence: str) -> str:
+        return hashlib.md5(sequence.encode()).hexdigest()[:12]
+
+    def _run_docker_command(
+        self,
+        cmd: list[str],
+        *,
+        phase: str = "docker command",
+        **kwargs,
+    ) -> subprocess.CompletedProcess:
+        kwargs["timeout"] = self.phase_timeout
+        try:
+            return subprocess.run(cmd, **kwargs)
+        except subprocess.TimeoutExpired as error:
+            raise TimeoutError(
+                f"MMseqs2 phase {phase!r} exceeded {self.phase_timeout:g} seconds"
+            ) from error
+
+    @staticmethod
+    def _working_root() -> Path:
+        return Path.cwd().resolve(strict=True)
+
+    def _resolve_path_under_cwd(self, path: str, *, must_exist: bool = False) -> Path:
+        root = self._working_root()
+        candidate = Path(path)
+        if not candidate.is_absolute():
+            candidate = root / candidate
+        try:
+            resolved = candidate.resolve(strict=must_exist)
+        except OSError as error:
+            raise ValueError(f"Path cannot be resolved safely: {path!r}") from error
+        if resolved != root and root not in resolved.parents:
+            raise ValueError(
+                "Path must resolve under the current working directory for the Docker mount. "
+                f"cwd={os.fspath(root)!r}, path={os.fspath(resolved)!r}"
+            )
+        return resolved
+
+    def _validate_paths_under_cwd(self, *paths: str) -> None:
+        for path in paths:
+            self._resolve_path_under_cwd(path)
+
+    def _path_in_container(self, local_path: str) -> str:
+        resolved = self._resolve_path_under_cwd(local_path)
+        relative = resolved.relative_to(self._working_root())
+        return relative.as_posix() or "."
+
+    def _docker_base_cmd(self) -> list[str]:
+        root = self._working_root()
+        cmd = ["docker", "run", "--rm"]
+        if not self.allow_network:
+            cmd.extend(["--network", "none"])
+        cmd.extend(["-v", f"{os.fspath(root)}:/app", "-w", "/app"])
+        if self.use_gpu:
+            if not torch.cuda.is_available():
+                raise RuntimeError(
+                    "use_gpu=True requires CUDA to be available in the FastPLMs host process"
+                )
+            cmd.extend(["--gpus", "all"])
+        cmd.append(self.docker_image)
+        return cmd
+
+    def _inspect_docker_image(self, *, check: bool) -> _DockerImageIdentity | None:
+        inspect = self._run_docker_command(
+            ["docker", "image", "inspect", self.docker_image],
+            phase="image inspection",
+            capture_output=True,
+            text=True,
+            check=check,
+        )
+        if inspect.returncode != 0:
+            stderr = inspect.stderr if isinstance(inspect.stderr, str) else ""
+            if "no such image" in stderr.lower() or "not found" in stderr.lower():
+                return None
+            raise subprocess.CalledProcessError(
+                inspect.returncode,
+                inspect.args,
+                output=inspect.stdout,
+                stderr=inspect.stderr,
+            )
+        try:
+            payload = json.loads(inspect.stdout)
+            if (
+                not isinstance(payload, list)
+                or len(payload) != 1
+                or not isinstance(payload[0], dict)
+            ):
+                raise ValueError("Docker inspect must return exactly one image object")
+            image = payload[0]
+            repo_digests = image.get("RepoDigests")
+            image_id = image.get("Id")
+            image_os = image.get("Os")
+            architecture = image.get("Architecture")
+            if not isinstance(repo_digests, list) or not all(
+                isinstance(value, str) for value in repo_digests
+            ):
+                raise ValueError("Docker inspect did not return RepoDigests")
+            expected_repo_digest = (
+                f"{self._image_reference.repository}@{self._image_reference.digest}"
+            )
+            if expected_repo_digest not in repo_digests:
+                raise ValueError(
+                    "Docker image RepoDigests do not contain the requested repository "
+                    "and manifest digest"
+                )
+            if not isinstance(image_id, str) or _SHA256_DIGEST_RE.fullmatch(image_id) is None:
+                raise ValueError("Docker inspect returned an invalid image ID")
+            if image_os != "linux":
+                raise ValueError(f"MMseqs2 image OS must be 'linux', got {image_os!r}")
+            expected_architecture = _docker_architecture()
+            if architecture != expected_architecture:
+                raise ValueError(
+                    "MMseqs2 image architecture does not match the host: "
+                    f"expected {expected_architecture!r}, got {architecture!r}"
+                )
+        except (KeyError, TypeError, ValueError, json.JSONDecodeError) as error:
+            raise RuntimeError(
+                f"Docker image identity verification failed for {self.docker_image!r}"
+            ) from error
+        return _DockerImageIdentity(
+            reference=self.docker_image,
+            repository=self._image_reference.repository,
+            version=self._image_reference.version,
+            manifest_digest=self._image_reference.digest,
+            image_id=image_id,
+            os=image_os,
+            architecture=architecture,
+        )
+
+    def _ensure_docker_image(self) -> _DockerImageIdentity:
+        if self._verified_image_identity is not None:
+            return self._verified_image_identity
+        self._run_docker_command(
+            ["docker", "version"],
+            phase="Docker availability check",
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+        identity = self._inspect_docker_image(check=False)
+        if identity is None:
+            if not self.allow_pull:
+                raise RuntimeError(
+                    "The pinned MMseqs2 image is not present locally and allow_pull=False. "
+                    "Preload the exact image out of band or opt in with allow_pull=True."
+                )
+            self._run_docker_command(
+                ["docker", "pull", self.docker_image],
+                phase="image pull",
+                check=True,
+                capture_output=True,
+                text=True,
+            )
+            identity = self._inspect_docker_image(check=True)
+        if identity is None:
+            raise RuntimeError("Docker image inspection succeeded without a verified identity")
+        self._verified_image_identity = identity
+        return identity
+
+    def _target_db_descriptor(self) -> dict[str, Any]:
+        prefix = self._resolve_path_under_cwd(self.target_db)
+        files: list[dict[str, Any]] = []
+        for candidate in sorted(prefix.parent.glob(f"{prefix.name}*")):
+            resolved = self._resolve_path_under_cwd(os.fspath(candidate), must_exist=True)
+            if not resolved.is_file():
+                continue
+            stat_result = resolved.stat()
+            files.append(
+                {
+                    "path": resolved.relative_to(self._working_root()).as_posix(),
+                    "size": stat_result.st_size,
+                    "mtime_ns": stat_result.st_mtime_ns,
+                }
+            )
+        if not files:
+            raise FileNotFoundError(
+                f"No MMseqs2 database files found for target_db prefix {self.target_db!r}"
+            )
+        derived_identity = _json_sha256(files)
+        return {
+            "prefix": prefix.relative_to(self._working_root()).as_posix(),
+            "identity": self.target_db_identity or derived_identity,
+            "identity_kind": "explicit" if self.target_db_identity is not None else "file-metadata",
+            "files": files,
+        }
+
+    def _request_provenance(self, sequence: str) -> dict[str, Any]:
+        return {
+            "provider": "mmseqs2",
+            "sequence_sha256": hashlib.sha256(sequence.encode("utf-8")).hexdigest(),
+            "image": {
+                "reference": self.docker_image,
+                "repository": self._image_reference.repository,
+                "version": self._image_reference.version,
+                "manifest_digest": self._image_reference.digest,
+            },
+            "platform": {"os": "linux", "architecture": _docker_architecture()},
+            "target_db": self._target_db_descriptor(),
+            "parameters": {
+                "sensitivity": self.sensitivity,
+                "max_seqs": self.max_seqs,
+                "min_seq_id": self.min_seq_id,
+                "coverage": self.coverage,
+                "split_memory_limit": self.split_memory_limit,
+                "use_gpu": self.use_gpu,
+                "allow_network": self.allow_network,
+            },
+        }
+
+    def _load_cached_result(
+        self,
+        a3m_output: str,
+        provenance_path: str,
+        request_provenance: dict[str, Any],
+    ) -> bool:
+        if not Path(a3m_output).is_file() or not Path(provenance_path).is_file():
+            return False
+        try:
+            with open(provenance_path, encoding="utf-8") as handle:
+                payload = json.load(handle)
+            if not isinstance(payload, dict):
+                return False
+            if payload.get("schema_version") != self._PROVENANCE_SCHEMA_VERSION:
+                return False
+            if payload.get("request") != request_provenance:
+                return False
+            request_identity = _json_sha256(request_provenance)
+            if payload.get("request_identity_sha256") != request_identity:
+                return False
+            runtime = payload.get("runtime")
+            if not isinstance(runtime, dict):
+                return False
+            if runtime.get("reference") != self.docker_image:
+                return False
+            if runtime.get("repository") != self._image_reference.repository:
+                return False
+            if runtime.get("version") != self._image_reference.version:
+                return False
+            if runtime.get("manifest_digest") != self._image_reference.digest:
+                return False
+            if runtime.get("os") != "linux":
+                return False
+            if runtime.get("architecture") != _docker_architecture():
+                return False
+            image_id = runtime.get("image_id")
+            if not isinstance(image_id, str) or _SHA256_DIGEST_RE.fullmatch(image_id) is None:
+                return False
+            cache_identity = _json_sha256(
+                {"request_identity_sha256": request_identity, "runtime": runtime}
+            )
+            if payload.get("cache_identity_sha256") != cache_identity:
+                return False
+            result = payload.get("result")
+            if not isinstance(result, dict):
+                return False
+            if result.get("path") != Path(a3m_output).name:
+                return False
+            if result.get("size") != Path(a3m_output).stat().st_size:
+                return False
+            return result.get("sha256") == _file_sha256(a3m_output)
+        except (OSError, UnicodeError, json.JSONDecodeError, TypeError, ValueError):
+            return False
+
+    def _store_result_provenance(
+        self,
+        provenance_path: str,
+        a3m_output: str,
+        request_provenance: dict[str, Any],
+        identity: _DockerImageIdentity,
+    ) -> None:
+        request_identity = _json_sha256(request_provenance)
+        runtime = identity.to_dict()
+        payload = {
+            "schema_version": self._PROVENANCE_SCHEMA_VERSION,
+            "request": request_provenance,
+            "request_identity_sha256": request_identity,
+            "runtime": runtime,
+            "cache_identity_sha256": _json_sha256(
+                {"request_identity_sha256": request_identity, "runtime": runtime}
+            ),
+            "result": {
+                "path": Path(a3m_output).name,
+                "size": Path(a3m_output).stat().st_size,
+                "sha256": _file_sha256(a3m_output),
+            },
+        }
+        output_dir = os.path.dirname(provenance_path) or "."
+        temporary_path: str | None = None
+        try:
+            with tempfile.NamedTemporaryFile(
+                mode="w",
+                encoding="utf-8",
+                dir=output_dir,
+                prefix=".mmseqs2-provenance-",
+                suffix=".tmp",
+                delete=False,
+            ) as handle:
+                temporary_path = handle.name
+                json.dump(payload, handle, indent=2, sort_keys=True)
+                handle.write("\n")
+                handle.flush()
+                os.fsync(handle.fileno())
+            os.replace(temporary_path, provenance_path)
+            temporary_path = None
+        finally:
+            if temporary_path is not None:
+                Path(temporary_path).unlink(missing_ok=True)
+
+    def create_db(self, fasta_path: str, db_path: str) -> str:
+        self._validate_paths_under_cwd(fasta_path, db_path)
+        os.makedirs(os.path.dirname(db_path) or ".", exist_ok=True)
+        if os.path.exists(f"{db_path}.dbtype"):
+            return db_path
+        self._ensure_docker_image()
+        self._run_docker_command(
+            [
+                *self._docker_base_cmd(),
+                "createdb",
+                self._path_in_container(fasta_path),
+                self._path_in_container(db_path),
+            ],
+            phase="createdb",
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+        return db_path
+
+    def create_index(self, db_path: str, tmp_dir: str | None = None) -> None:
+        if tmp_dir is None:
+            tmp_dir = os.path.join(os.path.dirname(db_path), "tmp_index")
+        self._validate_paths_under_cwd(db_path, tmp_dir)
+        os.makedirs(tmp_dir, exist_ok=True)
+        self._ensure_docker_image()
+        self._run_docker_command(
+            [
+                *self._docker_base_cmd(),
+                "createindex",
+                self._path_in_container(db_path),
+                self._path_in_container(tmp_dir),
+            ],
+            phase="createindex",
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+
+    def search(self, sequence: str, output_dir: str, seq_id: str | None = None) -> str:
+        if (
+            not isinstance(sequence, str)
+            or _SAFE_QUERY_SEQUENCE_RE.fullmatch(sequence) is None
+        ):
+            raise ValueError(
+                "sequence must be a non-empty unaligned ASCII protein sequence"
+            )
+        if seq_id is None:
+            seq_id = self._seq_hash(sequence)
+        seq_output_dir = _sequence_output_dir(output_dir, seq_id)
+        a3m_output = os.path.join(seq_output_dir, f"{seq_id}.a3m")
+        provenance_path = os.path.join(seq_output_dir, self._PROVENANCE_FILENAME)
+        self._validate_paths_under_cwd(seq_output_dir, self.target_db)
+        request_provenance = self._request_provenance(sequence)
+        if self._load_cached_result(a3m_output, provenance_path, request_provenance):
+            return a3m_output
+
+        identity = self._ensure_docker_image()
+        os.makedirs(seq_output_dir, exist_ok=True)
+        Path(a3m_output).unlink(missing_ok=True)
+        Path(provenance_path).unlink(missing_ok=True)
+        query_fasta = os.path.join(seq_output_dir, "query.fasta")
+        write_fasta_sequences(query_fasta, {seq_id: sequence})
+        query_db = os.path.join(seq_output_dir, "queryDB")
+        result_db = os.path.join(seq_output_dir, "resultDB")
+        tmp_dir = os.path.join(seq_output_dir, "tmp")
+        os.makedirs(tmp_dir, exist_ok=True)
+        self._validate_paths_under_cwd(
+            query_fasta, query_db, self.target_db, seq_output_dir, result_db, tmp_dir
+        )
+
+        docker_base = self._docker_base_cmd()
+        self._run_docker_command(
+            [
+                *docker_base,
+                "createdb",
+                self._path_in_container(query_fasta),
+                self._path_in_container(query_db),
+            ],
+            phase="query createdb",
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+        search_cmd = [
+            *docker_base,
+            "search",
+            self._path_in_container(query_db),
+            self._path_in_container(self.target_db),
+            self._path_in_container(result_db),
+            self._path_in_container(tmp_dir),
+            "-s",
+            str(self.sensitivity),
+            "--max-seqs",
+            str(self.max_seqs),
+            "--min-seq-id",
+            str(self.min_seq_id),
+            "-c",
+            str(self.coverage),
+        ]
+        if self.split_memory_limit is not None:
+            search_cmd.extend(["--split-memory-limit", self.split_memory_limit])
+        if self.use_gpu and torch.cuda.is_available():
+            search_cmd.extend(["--gpu", "1"])
+        self._run_docker_command(
+            search_cmd,
+            phase="search",
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+        self._run_docker_command(
+            [
+                *docker_base,
+                "result2msa",
+                self._path_in_container(query_db),
+                self._path_in_container(self.target_db),
+                self._path_in_container(result_db),
+                self._path_in_container(a3m_output),
+                "--msa-format-mode",
+                "6",
+            ],
+            phase="result2msa",
+            check=True,
+            capture_output=True,
+            text=True,
+        )
+        if not Path(a3m_output).is_file():
+            raise RuntimeError("MMseqs2 result2msa did not create a regular A3M file")
+        resolved_a3m = self._resolve_path_under_cwd(a3m_output, must_exist=True)
+        if not resolved_a3m.is_file():
+            raise RuntimeError("MMseqs2 result2msa did not create a regular A3M file")
+        self._store_result_provenance(
+            provenance_path,
+            a3m_output,
+            request_provenance,
+            identity,
+        )
+        for pattern in ["queryDB*", "resultDB*"]:
+            for path in Path(seq_output_dir).glob(pattern):
+                path.unlink(missing_ok=True)
+        tmp_path = Path(tmp_dir)
+        if tmp_path.exists():
+            shutil.rmtree(tmp_path, ignore_errors=True)
+        return a3m_output
+
+    def batch_search(
+        self,
+        sequences: list[str],
+        output_dir: str,
+        seq_ids: list[str] | None = None,
+        continue_on_error: bool = True,
+    ) -> dict[str, str]:
+        if seq_ids is None:
+            seq_ids = [self._seq_hash(seq) for seq in sequences]
+        if len(seq_ids) != len(sequences):
+            raise ValueError("seq_ids must contain exactly one identifier per sequence")
+        self._validate_paths_under_cwd(output_dir)
+        os.makedirs(output_dir, exist_ok=True)
+        results: dict[str, str] = {}
+        for seq, sid in tqdm(
+            list(zip(sequences, seq_ids, strict=True)),
+            desc="Searching homologues",
+        ):
+            try:
+                results[seq] = self.search(seq, output_dir, sid)
+            except Exception as error:
+                if not continue_on_error:
+                    raise
+                _get_logger().warning(
+                    "Homologue search failed and was skipped: "
+                    "provider=mmseqs2 seq_id=%s error_type=%s",
+                    sid,
+                    type(error).__name__,
+                )
+        return results
+
+
+@dataclass(frozen=True)
+class _ColabFoldResponse:
+    """Minimal response surface required by the ColabFold API client."""
+
+    status_code: int
+    headers: dict[str, str]
+    content: bytes
+
+    def json(self) -> dict[str, Any]:
+        value = json.loads(self.content.decode("utf-8"))
+        if not isinstance(value, dict):
+            raise ValueError("ColabFold returned a non-object JSON response")
+        return value
+
+
+class ColabFoldSearcher:
+    def __init__(
+        self,
+        host_url: str = COLABFOLD_HOST,
+        user_agent: str = "",
+        mode: str = "env",
+        timeout: float = 30.0,
+        max_retries: int = 10,
+        base_delay: float = 1.0,
+        max_delay: float = 60.0,
+        inter_request_delay: tuple[float, float] = (1.0, 3.0),
+        max_wait_time: int = 600,
+    ) -> None:
+        self.host_url = host_url.rstrip("/")
+        self.mode = mode
+        self.timeout = timeout
+        self.max_retries = max_retries
+        self.base_delay = base_delay
+        self.max_delay = max_delay
+        self.inter_request_delay = inter_request_delay
+        self.max_wait_time = max_wait_time
+        self.headers = {"User-Agent": user_agent} if user_agent else {}
+
+    @staticmethod
+    def _seq_hash(sequence: str) -> str:
+        return hashlib.md5(sequence.encode()).hexdigest()[:12]
+
+    def _backoff_delay(self, attempt: int) -> float:
+        delay = min(self.base_delay * (2**attempt), self.max_delay)
+        return min(delay + random.uniform(0, delay * 0.5), self.max_delay)
+
+    def _retry_after_delay(self, headers: dict[str, str], attempt: int) -> float:
+        raw_value = next(
+            (value for name, value in headers.items() if name.lower() == "retry-after"),
+            None,
+        )
+        if raw_value is None:
+            return self._backoff_delay(attempt)
+        try:
+            delay = float(raw_value)
+        except (TypeError, ValueError):
+            try:
+                retry_at = parsedate_to_datetime(raw_value)
+                if retry_at.tzinfo is None:
+                    retry_at = retry_at.replace(tzinfo=UTC)
+                delay = (retry_at - datetime.now(UTC)).total_seconds()
+            except (TypeError, ValueError, OverflowError):
+                return self._backoff_delay(attempt)
+        if not math.isfinite(delay):
+            return self._backoff_delay(attempt)
+        return min(max(0.0, delay), self.max_delay)
+
+    def _remaining_timeout(self, deadline: float | None, context: str) -> float:
+        if deadline is None:
+            return self.timeout
+        remaining = deadline - time.monotonic()
+        if remaining <= 0:
+            raise TimeoutError(f"{context} exceeded the {self.max_wait_time}s deadline")
+        return min(self.timeout, remaining)
+
+    def _sleep_with_deadline(
+        self,
+        delay: float,
+        deadline: float | None,
+        context: str,
+    ) -> None:
+        delay = max(0.0, delay)
+        if deadline is None:
+            time.sleep(delay)
+            return
+        remaining = deadline - time.monotonic()
+        if remaining <= 0:
+            raise TimeoutError(f"{context} exceeded the {self.max_wait_time}s deadline")
+        if delay >= remaining:
+            time.sleep(remaining)
+            raise TimeoutError(f"{context} exceeded the {self.max_wait_time}s deadline")
+        time.sleep(delay)
+
+    @staticmethod
+    def _http_error(response: _ColabFoldResponse, url: str) -> RuntimeError:
+        return RuntimeError(f"ColabFold request to {url} returned HTTP {response.status_code}")
+
+    def _request_with_retries(
+        self,
+        method: str,
+        url: str,
+        *,
+        deadline: float | None = None,
+        **kwargs: Any,
+    ) -> _ColabFoldResponse:
+        payload = kwargs.pop("data", None)
+        if kwargs:
+            unexpected = ", ".join(sorted(kwargs))
+            raise TypeError(f"Unsupported HTTP request options: {unexpected}")
+
+        encoded_payload = None
+        headers = dict(self.headers)
+        if payload is not None:
+            encoded_payload = urllib.parse.urlencode(payload).encode("utf-8")
+            headers["Content-Type"] = "application/x-www-form-urlencoded"
+
+        last_error: BaseException | None = None
+        for attempt in range(self.max_retries):
+            try:
+                request = urllib.request.Request(
+                    url,
+                    data=encoded_payload,
+                    headers=headers,
+                    method=method.upper(),
+                )
+                try:
+                    timeout = self._remaining_timeout(deadline, f"Request to {url}")
+                    with urllib.request.urlopen(request, timeout=timeout) as stream:
+                        response = _ColabFoldResponse(
+                            status_code=int(stream.status),
+                            headers={name.lower(): value for name, value in stream.headers.items()},
+                            content=stream.read(),
+                        )
+                except urllib.error.HTTPError as error:
+                    response = _ColabFoldResponse(
+                        status_code=int(error.code),
+                        headers={name.lower(): value for name, value in error.headers.items()},
+                        content=error.read(),
+                    )
+                if response.status_code == 429:
+                    last_error = self._http_error(response, url)
+                    if attempt + 1 >= self.max_retries:
+                        break
+                    self._sleep_with_deadline(
+                        self._retry_after_delay(response.headers, attempt),
+                        deadline,
+                        f"Request to {url}",
+                    )
+                    continue
+                if response.status_code >= 500:
+                    last_error = self._http_error(response, url)
+                    if attempt + 1 >= self.max_retries:
+                        break
+                    self._sleep_with_deadline(
+                        self._backoff_delay(attempt),
+                        deadline,
+                        f"Request to {url}",
+                    )
+                    continue
+                if not 200 <= response.status_code < 300:
+                    raise self._http_error(response, url)
+                return response
+            except (TimeoutError, urllib.error.URLError) as error:
+                last_error = error
+                if deadline is not None and time.monotonic() >= deadline:
+                    raise TimeoutError(
+                        f"Request to {url} exceeded the {self.max_wait_time}s deadline"
+                    ) from error
+                if attempt + 1 >= self.max_retries:
+                    break
+                self._sleep_with_deadline(
+                    self._backoff_delay(attempt),
+                    deadline,
+                    f"Request to {url}",
+                )
+        raise RuntimeError(
+            f"Request to {url} failed after {self.max_retries} attempts"
+        ) from last_error
+
+    def _submit(
+        self,
+        sequence: str,
+        mode: str | None = None,
+        deadline: float | None = None,
+    ) -> dict[str, Any]:
+        mode = mode or self.mode
+        query = f">101\n{sequence}\n"
+        for attempt in range(self.max_retries):
+            response = self._request_with_retries(
+                "POST",
+                f"{self.host_url}/ticket/msa",
+                data={"q": query, "mode": mode},
+                deadline=deadline,
+            )
+            data = response.json()
+            status = data.get("status", "UNKNOWN")
+            if status in ("RATELIMIT", "UNKNOWN"):
+                if attempt + 1 >= self.max_retries:
+                    break
+                self._sleep_with_deadline(
+                    self._backoff_delay(attempt),
+                    deadline,
+                    "ColabFold job submission",
+                )
+                continue
+            return data
+        raise RuntimeError(f"Failed to submit sequence after {self.max_retries} attempts")
+
+    def _poll(self, ticket_id: str, deadline: float | None = None) -> dict[str, Any]:
+        if deadline is None:
+            deadline = time.monotonic() + self.max_wait_time
+        poll_interval = 1.0
+        while True:
+            response = self._request_with_retries(
+                "GET",
+                f"{self.host_url}/ticket/{ticket_id}",
+                deadline=deadline,
+            )
+            data = response.json()
+            status = data.get("status", "ERROR")
+            if status in ("COMPLETE", "ERROR"):
+                return data
+            if status not in ("RUNNING", "PENDING", "UNKNOWN"):
+                return data
+            wait = min(poll_interval + random.uniform(0, 0.5), 5.0)
+            self._sleep_with_deadline(wait, deadline, f"Job {ticket_id}")
+            poll_interval = min(poll_interval + 1.0, 5.0)
+
+    def _download(
+        self,
+        ticket_id: str,
+        output_path: str,
+        deadline: float | None = None,
+    ) -> None:
+        response = self._request_with_retries(
+            "GET",
+            f"{self.host_url}/result/download/{ticket_id}",
+            deadline=deadline,
+        )
+        os.makedirs(os.path.dirname(output_path) or ".", exist_ok=True)
+        with open(output_path, "wb") as handle:
+            handle.write(response.content)
+
+    def _extract_a3m(self, tar_path: str, output_dir: str, seq_id: str) -> str:
+        with tarfile.open(tar_path) as tar:
+            _safe_extract_tar(tar, output_dir)
+
+        uniref_a3m = os.path.join(output_dir, "uniref.a3m")
+        env_a3m = os.path.join(output_dir, "bfd.mgnify30.metaeuk30.smag30.a3m")
+        a3m_files: list[str] = []
+        if os.path.exists(uniref_a3m):
+            a3m_files.append(uniref_a3m)
+        if "env" in self.mode and os.path.exists(env_a3m):
+            a3m_files.append(env_a3m)
+        combined_path = os.path.join(output_dir, f"{seq_id}.a3m")
+        if len(a3m_files) == 1:
+            os.replace(a3m_files[0], combined_path)
+        elif len(a3m_files) > 1:
+            with open(combined_path, "w", encoding="utf-8") as out_handle:
+                for a3m_file in a3m_files:
+                    with open(a3m_file, encoding="utf-8") as in_handle:
+                        out_handle.write(in_handle.read())
+        else:
+            raise RuntimeError("No .a3m files found in downloaded archive")
+        if os.path.exists(tar_path):
+            os.remove(tar_path)
+        for a3m_file in a3m_files:
+            if os.path.exists(a3m_file) and a3m_file != combined_path:
+                os.remove(a3m_file)
+        return combined_path
+
+    def search(self, sequence: str, output_dir: str, seq_id: str | None = None) -> str:
+        if seq_id is None:
+            seq_id = self._seq_hash(sequence)
+        seq_output_dir = _sequence_output_dir(output_dir, seq_id)
+        a3m_output = os.path.join(seq_output_dir, f"{seq_id}.a3m")
+        if os.path.exists(a3m_output):
+            return a3m_output
+        os.makedirs(seq_output_dir, exist_ok=True)
+        deadline = time.monotonic() + self.max_wait_time
+        result = self._submit(sequence, deadline=deadline)
+        status = result.get("status", "UNKNOWN")
+        if status == "ERROR":
+            raise RuntimeError(f"ColabFold API error for {seq_id}")
+        if status == "MAINTENANCE":
+            raise RuntimeError("ColabFold API is under maintenance")
+        ticket_id = result["id"]
+        result = self._poll(ticket_id, deadline=deadline)
+        status = result.get("status", "UNKNOWN")
+        if status != "COMPLETE":
+            raise RuntimeError(f"Job failed for {seq_id}: {status}")
+        tar_path = os.path.join(seq_output_dir, f"{seq_id}.tar.gz")
+        self._download(ticket_id, tar_path, deadline=deadline)
+        return self._extract_a3m(tar_path, seq_output_dir, seq_id)
+
+    def batch_search(
+        self,
+        sequences: list[str],
+        output_dir: str,
+        seq_ids: list[str] | None = None,
+        continue_on_error: bool = True,
+    ) -> dict[str, str]:
+        if seq_ids is None:
+            seq_ids = [self._seq_hash(seq) for seq in sequences]
+        os.makedirs(output_dir, exist_ok=True)
+        results: dict[str, str] = {}
+        pairs = list(zip(sequences, seq_ids, strict=True))
+        for i, (seq, sid) in enumerate(tqdm(pairs, desc="ColabFold search")):
+            try:
+                results[seq] = self.search(seq, output_dir, sid)
+            except Exception as error:
+                if not continue_on_error:
+                    raise
+                _get_logger().warning(
+                    "Homologue search failed and was skipped: "
+                    "provider=colabfold seq_id=%s error_type=%s",
+                    sid,
+                    type(error).__name__,
+                )
+            if i < len(pairs) - 1:
+                time.sleep(random.uniform(*self.inter_request_delay))
+        return results
+
+
+def _make_homologue_searcher(
+    provider: str, target_db: str | None, **kwargs
+) -> HomologueSearcher | ColabFoldSearcher:
+    if provider == "mmseqs2":
+        if target_db is None:
+            raise ValueError("target_db is required for MMseqs2 homologue search")
+        return HomologueSearcher(target_db=target_db, **kwargs)
+    if provider == "colabfold":
+        return ColabFoldSearcher(**kwargs)
+    raise ValueError(f"Unknown homologue search provider: {provider}")

fastplms/models/e1/tokenizer.json

@@ -0,0 +1,134 @@
+{
+  "version": "1.0",
+  "truncation": null,
+  "padding": {
+    "strategy": "BatchLongest",
+    "direction": "Right",
+    "pad_to_multiple_of": null,
+    "pad_id": 0,
+    "pad_type_id": 0,
+    "pad_token": "<pad>"
+  },
+  "added_tokens": [
+    {
+      "id": 0,
+      "content": "<pad>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 1,
+      "content": "<bos>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 2,
+      "content": "<eos>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 3,
+      "content": "<bos_glm>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 4,
+      "content": "<eos_span>",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    },
+    {
+      "id": 5,
+      "content": "?",
+      "single_word": false,
+      "lstrip": false,
+      "rstrip": false,
+      "normalized": false,
+      "special": true
+    }
+  ],
+  "normalizer": null,
+  "pre_tokenizer": {
+    "type": "ByteLevel",
+    "add_prefix_space": false,
+    "trim_offsets": true,
+    "use_regex": true
+  },
+  "post_processor": {
+    "type": "ByteLevel",
+    "add_prefix_space": true,
+    "trim_offsets": true,
+    "use_regex": true
+  },
+  "decoder": {
+    "type": "ByteLevel",
+    "add_prefix_space": true,
+    "trim_offsets": true,
+    "use_regex": true
+  },
+  "model": {
+    "type": "BPE",
+    "dropout": null,
+    "unk_token": "X",
+    "continuing_subword_prefix": null,
+    "end_of_word_suffix": null,
+    "fuse_unk": false,
+    "byte_fallback": false,
+    "ignore_merges": false,
+    "vocab": {
+      "<pad>": 0,
+      "<bos>": 1,
+      "<eos>": 2,
+      "<bos_glm>": 3,
+      "<eos_span>": 4,
+      "?": 5,
+      "1": 6,
+      "2": 7,
+      "A": 8,
+      "B": 9,
+      "C": 10,
+      "D": 11,
+      "E": 12,
+      "F": 13,
+      "G": 14,
+      "H": 15,
+      "I": 16,
+      "J": 17,
+      "K": 18,
+      "L": 19,
+      "M": 20,
+      "N": 21,
+      "O": 22,
+      "P": 23,
+      "Q": 24,
+      "R": 25,
+      "S": 26,
+      "T": 27,
+      "U": 28,
+      "V": 29,
+      "W": 30,
+      "X": 31,
+      "Y": 32,
+      "Z": 33
+    },
+    "merges": []
+  }
+}

fastplms/models/ttt.py

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

fastplms/registry.py

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

fastplms/runtime.py

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

modeling_fastplms.py

@@ -0,0 +1,236 @@
+"""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 != "c99634d724e168524f43f56ad1d22af47c75db0cbe657a2a55a237104fe7b833":
+    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.e1.modeling_e1")
+E1Config = _module_225.E1Config
+E1Config.__module__ = __name__
+E1ForMaskedLM = _module_225.E1ForMaskedLM
+E1ForMaskedLM.__module__ = __name__
+E1ForSequenceClassification = _module_225.E1ForSequenceClassification
+E1ForSequenceClassification.__module__ = __name__
+E1ForTokenClassification = _module_225.E1ForTokenClassification
+E1ForTokenClassification.__module__ = __name__
+E1Model = _module_225.E1Model
+E1Model.__module__ = __name__

runtime-attestation.json

@@ -0,0 +1,52 @@
+{
+  "files": {
+    "LICENSES/FastPLMs-Apache-2.0.txt": "sha256:2d2b50c7b1414bff1189a1db1f0cfb92e3e064b50f4c2b1019827b683e1b629a",
+    "LICENSES/e1/ATTRIBUTION": "sha256:deb22b250f6491b649eda5c63e080dd56486b8d2736cea6a52ef875436214367",
+    "LICENSES/e1/Apache-2.0.txt": "sha256:cfc7749b96f63bd31c3c42b5c471bf756814053e847c10f3eb003417bc523d30",
+    "LICENSES/e1/BSD-3-Clause.txt": "sha256:36e1987f2f17db7f8ad36cd7a37dbb7aeaaf0ab68b97ab4b9d3556f3a7a76ae8",
+    "LICENSES/e1/LICENSE": "sha256:8ef1dd556091544db3044164a8015424a3dcb3450fb3765a81b88463551bbe81",
+    "LICENSES/e1/MODIFICATIONS.md": "sha256:2506f47c0f5475af8e8ff2cff13eb8b79e8e25a08a054cdd617bf336536750ca",
+    "LICENSES/e1/NOTICE": "sha256:6de9db0320b4ee82f665c0951d8fd4cd53701a659c9dbce9bc3e3ea6afc4c6b3",
+    "README.md": "sha256:9f2d8fecae3f233b7546666d2487fca985627e25e8d9bb5277267d5de5a9ec10",
+    "THIRD_PARTY_NOTICES.md": "sha256:25704b3c76404696cae52e7fca13088d329f70f412687340351259e86cd62baa",
+    "config.json": "sha256:5a6ea49bbfefd19984b92448121277dfbecae387bb2fb360ba9cec29a508ed7a",
+    "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/e1/__init__.py": "sha256:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855",
+    "fastplms/models/e1/attention.py": "sha256:8449b2bafba9573dfcfe4ad907e1fce24c55a4f279f1f72cd52227e48dcecf67",
+    "fastplms/models/e1/cache.py": "sha256:b17a0a743505f567b387d1edb8550f028a8770db1ba56ad8ad2ff0d3fce19c46",
+    "fastplms/models/e1/modeling_e1.py": "sha256:96545247491395291060e2cbb7b6b082ed6073e0fbf12d4656e490010aae5ae4",
+    "fastplms/models/e1/preparation.py": "sha256:3c1cc96d36cbfa74dc0f847cc5f44c6edd527d8e223d0a564d4891a041761395",
+    "fastplms/models/e1/retrieval.py": "sha256:d077ef10100cf4f1482fbc7c8a944c241e61061a94374699156388d9e8e5c4db",
+    "fastplms/models/e1/tokenizer.json": "sha256:65d4345dbe908d3deb554891dc01b94e215fb17cd9e0bd0535e062e256737415",
+    "fastplms/models/ttt.py": "sha256:a0df4e98b02120d423e3c7ca9b866a8d0e3748b9076042a0102a838a11aed046",
+    "fastplms/registry.py": "sha256:afca271911b651a882345b74a58366494a1d784e4a48c8683a87f5508f4ba16e",
+    "fastplms/runtime.py": "sha256:110018646d6f248cedab140a030c3065e1b062b61f6aff659c231e538614bc01",
+    "fastplms_bundle.py": "sha256:bbe671e5653b3d3028100278166f0e8a91f72eaeaf3c2ed3b7902cd833020840",
+    "modeling_fastplms.py": "sha256:8fe92d7a9b0b1916575448c4fbba67005fcad75e456f68bba38aff2054d43c1f"
+  },
+  "model_id": "e1_300m",
+  "redistributable": true,
+  "release_tool_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "release_tool_sha256": "1459b5d7d13d9b07bd97b3eee764f2ce73623e15e32d07ddf6825c2a9509afb9",
+  "runtime_bundle_sha256": "c99634d724e168524f43f56ad1d22af47c75db0cbe657a2a55a237104fe7b833",
+  "runtime_revision": "1b9ce023f1e06571cf3e6324be0610ffa53e0a4a",
+  "schema_version": 2,
+  "scope": "runtime-only",
+  "source_tree_sha256": "cec0500dd5f6b238474600a6e80b8d3ccd4544ef7daec90441f14273f1b658d2",
+  "weights": {
+    "repo_id": "Synthyra/Profluent-E1-300M",
+    "revision": "5ef52c0ad2ae2578f40622696b763523810e8e26"
+  },
+  "weights_license_status": "resolved"
+}