diff --git a/LICENSES/FastPLMs-Apache-2.0.txt b/LICENSES/FastPLMs-Apache-2.0.txt new file mode 100644 index 0000000000000000000000000000000000000000..66c354fb5c07a2de923f8de7cfe1e869dac39200 --- /dev/null +++ b/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. diff --git a/LICENSES/biohub-esm/LICENSE.md b/LICENSES/biohub-esm/LICENSE.md new file mode 100644 index 0000000000000000000000000000000000000000..d6021558652f83cd373d2273f16d9cf402134287 --- /dev/null +++ b/LICENSES/biohub-esm/LICENSE.md @@ -0,0 +1,9 @@ +**License (MIT)** + +Copyright 2026 Chan Zuckerberg Biohub, Inc. + +Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the “Software”), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. diff --git a/LICENSES/biohub-esm/THIRD_PARTY_NOTICE.md b/LICENSES/biohub-esm/THIRD_PARTY_NOTICE.md new file mode 100644 index 0000000000000000000000000000000000000000..56d8280e67083eb023fda21b0091f61ea6bdab33 --- /dev/null +++ b/LICENSES/biohub-esm/THIRD_PARTY_NOTICE.md @@ -0,0 +1,13 @@ +The code in this repository depends on the following third-party libraries: + +| Library | License | Link | +|----------|----------|----------| +| flash-attn | BSD | https://github.com/Dao-AILab/flash-attention/blob/main/LICENSE | +| PyTorch | BSD | https://github.com/pytorch/pytorch/blob/main/LICENSE | +| xformers | BSD | https://github.com/facebookresearch/xformers/blob/main/LICENSE | +| jaxtyping | MIT | https://github.com/patrick-kidger/jaxtyping/blob/main/LICENSE | +| einops | MIT | https://github.com/arogozhnikov/einops/blob/main/LICENSE | +| omegaconf | BSD | https://github.com/omry/omegaconf/blob/master/LICENSE | +| attrs | MIT | https://github.com/python-attrs/attrs/blob/main/LICENSE | +| scipy | BSD-3-Clause | https://github.com/scipy/scipy/blob/main/LICENSE.txt
https://github.com/scipy/scipy/blob/main/LICENSES_bundled.txt | +| lightning / torchmetrics | Apache 2.0 | https://github.com/Lightning-AI/torchmetrics/blob/master/LICENSE | diff --git a/LICENSES/biohub-transformers/LICENSE b/LICENSES/biohub-transformers/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..68b7d66c97d66c58de883ed0c451af2b3183e6f3 --- /dev/null +++ b/LICENSES/biohub-transformers/LICENSE @@ -0,0 +1,203 @@ +Copyright 2018- The Hugging Face team. All rights reserved. + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/LICENSES/protein-ttt/LICENSE b/LICENSES/protein-ttt/LICENSE new file mode 100644 index 0000000000000000000000000000000000000000..2e8ce69ee735b7021ac28defffdfecb43d287db6 --- /dev/null +++ b/LICENSES/protein-ttt/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2024 Anton Bushuiev + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/LICENSES/protein-ttt/SOURCE_RECORD.md b/LICENSES/protein-ttt/SOURCE_RECORD.md new file mode 100644 index 0000000000000000000000000000000000000000..b218841986dc236e71f4f124a8f1b0d668ca36b1 --- /dev/null +++ b/LICENSES/protein-ttt/SOURCE_RECORD.md @@ -0,0 +1,8 @@ +# ProteinTTT provenance + +FastPLMs uses `anton-bushuiev/ProteinTTT` revision +`fde2817cd84b936167cc76ccabf31e5c0fe49962` as the official reference for the +optional protein test-time training workflow. The repository is pinned at +`vendor/upstream/protein-ttt/` and is not a production dependency or runtime +image component. The accompanying `LICENSE` is the verbatim MIT text from that +revision. diff --git a/README.md b/README.md new file mode 100644 index 0000000000000000000000000000000000000000..ed68901248fe4853c9afc8c0559f94308275f8fb --- /dev/null +++ b/README.md @@ -0,0 +1,235 @@ +--- +library_name: transformers +license: "mit" +tags: + - protein-language-model + - fastplms +--- + + + +# ESMFold2-600 + +## Model overview + +`Synthyra/ESMFold2-600` packages the +`biohub/ESMFold2-Experimental-Fast-base600M-step1500k` checkpoint with the +FastPLMs runtime for Hugging Face Transformers. It accepts raw amino-acid +sequences or typed molecular-complex specifications; low-level forward accepts +prepared feature tensors. + +The repository uses the standard Transformers loading interface with +`trust_remote_code=True`. See Technical details for each registered class and +whether its weights come from the checkpoint. + +The sequence- and token-classification classes reuse the pretrained backbone, +but their task heads are newly initialized. Fine-tune those heads before +interpreting their logits as predictions. + +## Install and platform requirements + +Install the direct dependencies published with this model: + +```bash +python -m pip install -r \ + "https://huggingface.co/Synthyra/ESMFold2-600/resolve/main/requirements.txt" +``` + +The FastPLMs implementation itself is embedded in the model repository. +Transformers loads it through `trust_remote_code=True`. + +This model requires Python 3.11-3.14, PyTorch 2.13, and Transformers 5.13. + +The artifact requirements include the structure dependencies. + +Validation runs in Docker on any compatible CUDA device. Record the container, +hardware, precision, and inputs; no GPU product or workstation is required. + +The Hub quick start needs network access for the first download. For an +air-gapped run, build the manifest-pinned local artifact first and use the +offline example. + +## Quick start + +```python +from transformers import AutoModel + +model_id = "Synthyra/ESMFold2-600" +model = AutoModel.from_pretrained( + model_id, + trust_remote_code=True, + attn_implementation="sdpa", +).eval() +``` + +For offline validation, replace `model_id` with the manifest-built +`dist/hub/ESMFold2-600` path. Pass `local_files_only=True`. + +## Attention backends + +The quick start uses `sdpa`. + +Available backends are `eager`, `sdpa`, `flex_attention`. Requesting an +unavailable backend raises instead of silently changing implementation. + +`output_attentions=True` can use the documented one-call eager fallback to +materialize attention tensors. The configured backend does not change. + +## Downstream prediction + +The sequence and token prediction AutoClasses use the checkpoint backbone and +create a new, untrained `classifier`. Sequence labels have shape `(b,)`. +Residue labels have shape `(b, l)` and use `-100` outside biological positions. +The folding trunk is skipped. The classifier uses the checkpoint's learned pLM +state mixture and projection, followed by one trainable transformer probe. + +```python +import torch +from transformers import ( + AutoModelForSequenceClassification, + AutoModelForTokenClassification, +) + +model_id = "Synthyra/ESMFold2-600" +sequence_model = AutoModelForSequenceClassification.from_pretrained( + model_id, num_labels=2, trust_remote_code=True +).eval() +token_model = AutoModelForTokenClassification.from_pretrained( + model_id, num_labels=3, trust_remote_code=True +).eval() +sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"] +batch = sequence_model.prepare_classifier_inputs(sequences) +biological = batch["attention_mask"].bool() + +sequence_labels = torch.zeros(len(sequences), dtype=torch.long) +token_labels = torch.full_like(batch["input_ids"], -100) +token_labels[biological] = 0 + +with torch.inference_mode(): + sequence_output = sequence_model(**batch, labels=sequence_labels) + token_output = token_model(**batch, labels=token_labels) +print(sequence_output.logits.shape) # (b, 2) +print(token_output.logits.shape) # (b, l, 3) +``` + +## PEFT fine-tuning + +Install the training dependencies. Then attach LoRA to the loaded checkpoint: + +```bash +python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20" +``` + +```python +from peft import LoraConfig, TaskType, get_peft_model + +peft_model = get_peft_model( + sequence_model, + LoraConfig( + task_type=TaskType.SEQ_CLS, + r=8, + lora_alpha=16, + target_modules="all-linear", + modules_to_save=["classifier"], + ), +) +``` + +This checkpoint advertises a classification head. Save the separately trained +`classifier` with the adapter. +All FastPLMs checkpoints follow the Transformers `PreTrainedModel` contract and +can use PEFT. The ESM2-specific shipped CLI is an example, not a +support boundary. Record the target modules, base revision, data identity, and +trainable parameter scope. + +## Protein folding + +This experimental Fast checkpoint has 24 folding blocks and uses the frozen +`Synthyra/ESMplusplus_large` backbone. The config-declared step-1500000 backbone and the +pinned ESM++ weights are tensor-exact in BF16 after layout conversion. + +```python +import torch + +model = model.cuda().eval() +with torch.inference_mode(): + output = model.infer_protein( + "MQYKLILNGKTLKGETTTEAVDAATAEKVFKQYANDNGVDGEWTYDDATKTFTVTE", + seed=17, + num_diffusion_samples=1, + ) +print(output.sample_atom_coords.shape) +``` + +Folding parameters remain FP32 with CUDA BF16 autocast. The backbone uses +BF16; FP8 requests fail. The 15-step sampler and three folding loops remain +the checkpoint defaults. Protein inputs require `msa=None`. +This checkpoint was trained without MSA conditioning. It rejects +`ProteinInput.msa` and MSA-derived features. Typed multichain and multimolecule +inputs remain supported without MSA conditioning. + +The confidence head is disabled: pLDDT, pTM, iPTM, and PAE are unavailable. +The 300 and 600 suffixes describe backbone scale, not total model parameters. + +## Learned representation and ESMC precision + +The learned projection maps `H: (b, l, 37, 1152) -> Z: (b, l, 256)`. +`embed_dataset` returns one `(l, 256)` residue representation per sequence. +The experimental architecture does not expose folding TTT. + + +## Notes and limitations + +Experimental Fast checkpoint with a frozen 600M ESM++ backbone, tensor-exact in +BF16 with the pinned step-1500000 source, 24 folding blocks, no MSA +conditioning, and no confidence head. BF16 execution uses FP32 folding +parameters with CUDA autocast. FP8 is unsupported. Configuration, weight +identities, and artifact reload are verified; this model is not +inference-validated. Publication is pending. + +## Technical details + +- Inputs: Raw amino-acid sequences or typed molecular-complex specifications; low-level forward accepts prepared feature tensors +- Transformers classes: `AutoConfig`, `AutoModel`, `AutoModelForSequenceClassification`, `AutoModelForTokenClassification` +- Checkpoint weights: `AutoConfig` = `FastPLMs extension`, `AutoModel` = `pretrained`, `AutoModelForSequenceClassification` = `base weights + untrained task head`, `AutoModelForTokenClassification` = `base weights + untrained task head` +- Attention backends: `eager`, `sdpa`, `flex_attention` +- Precision: `auto`, `fp32`, `bf16` +- BF16 execution: `fp32_parameters_autocast` +- Generation contract: `not_applicable` +- Dependencies: `core + structure` +- Weight publication allowed: `true` +- Weight license status: `resolved` +- Redistributable: `true` +- Complete weight publication required: `false` + +## Validation and provenance + +FastPLMs pins the checkpoint, upstream source revisions, state transformation, +and required files in `models.toml`. Built artifacts record exact source +identities and conversion details in `source-record.json`. + +- FastPLMs checkpoint: `Synthyra/ESMFold2-600` +- Runtime revision: recorded separately in the built artifact and published commit +- Runtime source identities: recorded in `source-record.json` +- Official checkpoint: `biohub/ESMFold2-Experimental-Fast-base600M-step1500k` +- Artifact source: `fast` +- State transform: `identity` +- Pinned upstreams: `biohub-esm`, `biohub-transformers`, `protein-ttt` +- Release tiers: `check`, `compliance`, `structure`, `feature`, `artifact`, `benchmark` +- Unresolved required file identities: `0` + +ESMFold2-300 passed a Docker BF16 reference comparison on one compact Protein G +sequence. ESMFold2-600 is not inference-validated. Both have configuration, +weight identity, and artifact loading checks. This is bounded checkpoint +evidence, not a full structure benchmark result. + +Declared tiers compare configuration, tokenizer behavior, state, and +representative inference with the pinned reference. A nonzero unresolved count +blocks release. Metadata alone does not show that a build passed, that a backend +is faster, or that an output is biologically valid. + +## License + +Checkpoint terms: MIT. The Hub model-card identifier is +`mit`. The local artifact contains applicable source +licenses, notices, attribution, and conversion records. Review them before use. diff --git a/THIRD_PARTY_NOTICES.md b/THIRD_PARTY_NOTICES.md new file mode 100644 index 0000000000000000000000000000000000000000..23e1ddc075d5f693ba8c1184427da8c8d070c5e8 --- /dev/null +++ b/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/SOURCE_RECORD.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/SOURCE_RECORD.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 `source-record.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. diff --git a/config.json b/config.json new file mode 100644 index 0000000000000000000000000000000000000000..8e271837cbdada96c4974c8e543f84065e0f06f1 --- /dev/null +++ b/config.json @@ -0,0 +1,104 @@ +{ + "architectures": [ + "ESMFold2ExperimentalModel" + ], + "auto_map": { + "AutoConfig": "modeling_fastplms.ESMFold2Config", + "AutoModel": "modeling_fastplms.ESMFold2ExperimentalModel", + "AutoModelForSequenceClassification": "modeling_fastplms.ESMFold2ExperimentalForSequenceClassification", + "AutoModelForTokenClassification": "modeling_fastplms.ESMFold2ExperimentalForTokenClassification" + }, + "confidence_head": { + "distogram_bins": 128, + "enabled": false, + "folding_trunk": { + "dropout": 0.25, + "n_heads": 8, + "n_layers": 4 + }, + "max_dist": 52.0, + "min_dist": 2.0, + "num_pae_bins": 64, + "num_pde_bins": 64, + "num_plddt_bins": 50 + }, + "d_pair": 256, + "d_single": 384, + "disable_msa_features": true, + "dtype": "float32", + "esmc_id": "Synthyra/ESMplusplus_large", + "fastplms_checkpoint_hash": "5c59bffd38186e80d0a5ebd28e9d73129fd96f462ce145be925d289e945b02d8", + "fastplms_checkpoint_repo_id": "biohub/ESMFold2-Experimental-Fast-base600M-step1500k", + "fastplms_checkpoint_revision": "15cf2d6648692f6c17cee1297d8a285476fffa9b", + "fastplms_model_id": "esmfold2_600", + "fastplms_runtime_bundle_sha256": "ef1f313b70da02019bf0712085e332968c1cecf189633f721fce2554cff6b61f", + "fastplms_weights_revision": "15cf2d6648692f6c17cee1297d8a285476fffa9b", + "folding_trunk": { + "dropout": 0.25, + "n_heads": 8, + "n_layers": 24 + }, + "force_lm_dropout_during_inference": false, + "inputs": { + "atom_encoder": { + "d_atom": 128, + "d_token": 768, + "expansion_ratio": 2, + "n_blocks": 3, + "n_heads": 4, + "n_spatial_rope_pairs_per_axis": 2, + "n_uid_rope_pairs": 10, + "spatial_rope_base_frequency": 20.0, + "swa_window_size": 128, + "uid_rope_base_frequency": 10000.0 + }, + "d_inputs": 451 + }, + "lm_d_model": 1152, + "lm_dropout": 0.0, + "lm_num_layers": 36, + "model_type": "esmfold2", + "msa_conditioning": false, + "msa_encoder": { + "d_hidden": 32, + "d_msa": 128, + "enabled": false, + "msa_head_width": 32, + "n_heads_msa": 8, + "n_layers": 4 + }, + "n_relative_chain_bins": 2, + "n_relative_residx_bins": 32, + "num_diffusion_samples": 16, + "num_loops": 3, + "structure_head": { + "diffusion_module": { + "atom_num_blocks": 3, + "atom_num_heads": 4, + "c_atom": 128, + "c_s_inputs": 451, + "c_token": 768, + "c_z": 256, + "fourier_dim": 256, + "relpos_r_max": 32, + "relpos_s_max": 2, + "sigma_data": 16.0, + "token_num_blocks": 12, + "token_num_heads": 16, + "transition_multiplier": 2 + }, + "distogram_bins": 128, + "gamma_0": 0.605, + "gamma_min": 1.107, + "inference_num_steps": 15, + "inference_p": 8.0, + "inference_s_max": 160.0, + "inference_s_min": 0.0004, + "noise_scale": 0.901, + "step_scale": 1.638, + "train_noise_log_mean": -1.2, + "train_noise_log_std": 1.5 + }, + "transformers_version": "4.57.6", + "type": "experimental" +} diff --git a/fastplms/__init__.py b/fastplms/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c5eebf713f357493cce628d01762dd91279ed6da --- /dev/null +++ b/fastplms/__init__.py @@ -0,0 +1,49 @@ +"""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__)) diff --git a/fastplms/attention/__init__.py b/fastplms/attention/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c54da914ec96b371f45b088091bc5eba9ec042b6 --- /dev/null +++ b/fastplms/attention/__init__.py @@ -0,0 +1,68 @@ +"""Shared attention backends, masks, and optional optimized kernels.""" + +from ._core import ( + LEGACY_CHECKPOINT_ATTENTION_BACKENDS, + 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, + canonical_checkpoint_attention_backend, + 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", + "LEGACY_CHECKPOINT_ATTENTION_BACKENDS", + "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", + "canonical_checkpoint_attention_backend", + "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", +] diff --git a/fastplms/attention/_core.py b/fastplms/attention/_core.py new file mode 100644 index 0000000000000000000000000000000000000000..200d3d93bee777e23dc54e8ed83c399a03ce7d4c --- /dev/null +++ b/fastplms/attention/_core.py @@ -0,0 +1,824 @@ +"""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 +import torch +from collections import OrderedDict +from collections.abc import Callable +from enum import Enum +from threading import RLock +from types import MappingProxyType +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() # mask_pattern.shape + # 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() # mask_pattern.shape + pattern_bytes = host_pattern.view(torch.uint8).numpy().tobytes(order="C") # bytes + 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.""" + + # query_states, key_states, value_states: (b, l, h, d) or (t, h, d) + 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.""" + + # query_states, key_states, value_states: (b, l, h, d) or (t, h, d) + 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. + """ + # query_states, key_states, value_states: (b, l, h, d) + flash_kernel, flash_kernel_variant = _ensure_flash_kernels_loaded(implementation) + if flash_kernel_variant == "flash_attn2": + output = flash_kernel.flash_attn_func( # (b, l, h, d) or tuple with that first + 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 # (b, l, h, d) + if flash_kernel_variant == "flash_attn3": + output = flash_kernel.flash_attn_func( # (b, l, h, d) or tuple with that first + q=query_states, + k=key_states, + v=value_states, + softmax_scale=softmax_scale, + causal=causal, + ) + if isinstance(output, tuple): + return output[0] # (b, l, h, d) + return output # (b, l, h, d) + 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. + """ + # query_states, key_states, value_states: (t, h, d) + # cu_seqlens_q, cu_seqlens_k: (b + 1,) + flash_kernel, flash_kernel_variant = _ensure_flash_kernels_loaded(implementation) + if flash_kernel_variant == "flash_attn2": + output = flash_kernel.flash_attn_varlen_func( # (t, h, d) or tuple with that first + 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 # (t, h, d) + if flash_kernel_variant == "flash_attn3": + output = flash_kernel.flash_attn_varlen_func( # (t, h, d) or tuple with that first + 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] # (t, h, d) + return output # (t, h, d) + 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: + # input: (n, ...); indices: (m,) + 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( # (m, ...) + 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]: + # grad_output: (m, ...) + (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 (...)") # (m, product(...)) + grad_input = torch.zeros( # (n, product(...)) + [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 # (n, ...), None + + +class IndexPutFirstAxis(torch.autograd.Function): + @staticmethod + def forward(ctx, values, indices, first_axis_dim) -> torch.Tensor: + # values: (m, ...); indices: (m,) + 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( # (n, ...) + first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype + ) + output[indices] = values + return output # (n, ...) + + @staticmethod + def backward(ctx, grad_output) -> tuple[torch.Tensor, None, None]: + # grad_output: (n, ...) + (indices,) = ctx.saved_tensors + return grad_output[indices], None, None # (m, ...), 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: + # hidden_states: (t, ...); indices: (t,) + output = index_put_first_axis(hidden_states, indices, batch * seqlen) # (b * l, ...) + return rearrange(output, "(b s) ... -> b s ...", b=batch) # (b, l, ...) + + +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], +]: + # query_layer, key_layer, value_layer: (b, l, h, d); attention_mask_2d: (b, l) + batch_size, seq_len, num_heads, head_dim = query_layer.shape + seqlens = attention_mask_2d.sum(dim=1).int() # (b,) + cu_seqlens = F.pad(seqlens.cumsum(0, dtype=torch.int32), (1, 0)) # (b + 1,) + max_seqlen = int(seqlens.max().item()) + indices = attention_mask_2d.flatten().nonzero(as_tuple=False).flatten() # (t,) + query_layer = index_first_axis( # (t, h, d) + query_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices + ) + key_layer = index_first_axis( # (t, h, d) + key_layer.reshape(batch_size * seq_len, num_heads, head_dim), indices + ) + value_layer = index_first_axis( # (t, h, d) + 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.""" + + # query_states, key_states, value_states: (b, l, h, d); attention_mask_2d: (b, l) + 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) # (b, l) + + +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. + """ + # query_states, key_states, value_states: (b, l, h, d) + # attention_mask_2d: (b, l) or None + _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) # (b, l, h, d) + key_states = key_states.to(dtype=runtime_dtype) # (b, l, h, d) + value_states = value_states.to(dtype=runtime_dtype) # (b, l, h, d) + if attention_mask_2d is not None: + attention_mask_2d = _validate_flash_padding_mask( # (b, l) + 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( # (t, h, d), (t, h, d), (t, h, d), (t,), (b + 1,), scalars + query_states, + key_states, + value_states, + attention_mask_2d, + ) + attn_output_unpad = _kernels_flash_varlen_forward( # (t, h, d) + 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) # (b, l, h, d) + return output.masked_fill(~attention_mask_2d[:, :, None, None], 0) # (b, l, h, d) + else: + return _kernels_flash_forward( # (b, l, h, d) + 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) + +# Official upstream sources predate the Transformers ``flex_attention`` name and +# store ``"flex"`` in their checkpoint configurations. The two names select the +# identical implementation, so translating one to the other is a rename rather +# than a backend substitution. ``"flash"`` is deliberately absent: it does not +# identify a FlashAttention version, and guessing one would be a substitution. +LEGACY_CHECKPOINT_ATTENTION_BACKENDS = MappingProxyType( + {"flex": AttentionBackend.FLEX_ATTENTION.value} +) + + +def canonical_checkpoint_attention_backend(stored_backend: str | None) -> str | None: + """Translate a stored checkpoint backend name to its canonical spelling. + + Call this where a serialized configuration enters FastPLMs. Imperative + selection through ``set_attn_implementation`` stays strict and keeps + rejecting the historical spellings. An unmapped name passes through so the + backend resolver owns the rejection. + """ + + if stored_backend is None: + return None + return LEGACY_CHECKPOINT_ATTENTION_BACKENDS.get(stored_backend, stored_backend) + + +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). + """ + # attention_mask: (b, l) or None + 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) # (b, l) + 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 # (b, l), 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 # (b, l), None, BlockMask + + # 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, :] # (b, 1, 1, l) + return attention_mask_2d, attention_mask_4d, None # (b, l), (b, 1, 1, l), 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. + """ + # bool_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 # (...) diff --git a/fastplms/attention/_kernel_lock.py b/fastplms/attention/_kernel_lock.py new file mode 100644 index 0000000000000000000000000000000000000000..a859f57c9af23eb30743234d333c2f4ffbdf0406 --- /dev/null +++ b/fastplms/attention/_kernel_lock.py @@ -0,0 +1,175 @@ +"""Resolve and validate Hugging Face kernels before importing their binaries.""" + +from __future__ import annotations + +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 requirements/features/flash.in." + ) from error + + +def _kernel_lock_path() -> Path: + """Return the kernel lock from a Hub artifact or source checkout.""" + source_path = Path(__file__).resolve() + for candidate in ( + source_path.parents[1] / "kernels.lock", + source_path.parents[3] / "kernels.lock", + ): + if candidate.is_file(): + return candidate + + raise RuntimeError( + "kernels.lock is missing from the Hugging Face artifact or source checkout." + ) + + +def _locked_entry(lock_path: Path, repository: str) -> dict[str, Any]: + try: + lock_entries = json.loads(lock_path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise RuntimeError(f"Unable to read the kernel lock: {lock_path}") from error + if not isinstance(lock_entries, list): + raise RuntimeError("kernels.lock must contain a JSON list.") + if any(not isinstance(entry, dict) for entry in lock_entries): + raise RuntimeError("Every kernels.lock entry must be a JSON object.") + matches = [entry for entry in lock_entries 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 requirements/features/flash.in." + ) from error + + cached_variants = get_variants_local(build_root) + parsed_names = {variant.variant_str for variant in cached_variants} + 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(cached_variants) + 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 requirements/features/flash.in." + ) 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) diff --git a/fastplms/attention/interfaces.py b/fastplms/attention/interfaces.py new file mode 100644 index 0000000000000000000000000000000000000000..7d219c8145fa8532a2a405f32d176f7bbf90dad5 --- /dev/null +++ b/fastplms/attention/interfaces.py @@ -0,0 +1,249 @@ +"""Transformers-compatible attention selection for FastPLMs models.""" + +from __future__ import annotations + +import torch +from collections.abc import Mapping +from functools import partial +from typing import Any +from transformers import AttentionInterface, AttentionMaskInterface + +from ._core import ( + AttentionBackend, + canonical_checkpoint_attention_backend, + 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. + """ + + # query, key, value: (b, h, l, d); attention_mask: (b, l) or None + 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(), # (b, l, h, d) + key_states=key.transpose(1, 2).contiguous(), # (b, l, h, d) + value_states=value.transpose(1, 2).contiguous(), # (b, l, h, d) + attention_mask_2d=attention_mask, + causal=causal, + softmax_scale=softmax_scale, + implementation=implementation, + ) # (b, l, h, d) + return output, None # (b, l, h, d), 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) + stored = ( + getattr(config, "_attn_implementation", None) if internal is sentinel else internal + ) + legacy = getattr(config, "attn_backend", None) + requested = stored if stored 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__}." + ) + # A serialized configuration can name a backend with the historical + # spelling used by the official source it was converted from. That + # names the same implementation, so translate it here rather than + # rejecting a checkpoint that asked for an implementation FastPLMs has. + canonical = canonical_checkpoint_attention_backend(requested) + self._validate_attention_name(canonical) + # ``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 stored canonical value already agrees and + # is left untouched, including an explicit + # ``attn_implementation=...`` load override. + if canonical != stored: + set_config_attn_implementation(config, canonical) + 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}." + ) diff --git a/fastplms/embeddings/__init__.py b/fastplms/embeddings/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..dd1ddd5abe67511a6fffadd502152ecc619785b1 --- /dev/null +++ b/fastplms/embeddings/__init__.py @@ -0,0 +1,66 @@ +"""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", +] diff --git a/fastplms/embeddings/pooling.py b/fastplms/embeddings/pooling.py new file mode 100644 index 0000000000000000000000000000000000000000..08d577c0e03c8a41d0c5e9ca924a6eb0cfd307a8 --- /dev/null +++ b/fastplms/embeddings/pooling.py @@ -0,0 +1,218 @@ +"""Residue-aware pooling implemented entirely with PyTorch.""" + +from __future__ import annotations + +import math +import torch +from collections.abc import Sequence +from torch import Tensor + + +POOLING_NAMES = frozenset({"mean", "max", "norm", "median", "std", "var", "cls", "parti"}) + + +def _validate_inputs(X: Tensor, M: Tensor) -> Tensor: + # X: (b, l, d); M: (b, l) + 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) # (b, l) + 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 # (b, l) + + +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: (b, h, l, l). + A = torch.stack(tuple(attentions), dim=1) # (b, n, h, l, l) + else: + A = attentions # (b, ..., l, l) + + if A.ndim == 5: + if A.shape[0] != batch_size and A.shape[1] == batch_size: + A = A.transpose(0, 1) # (b, n, h, l, l) + 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) # (b, l, l) + 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) # (b, l, l) + 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 # (b, l, l) + + +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. + """ + + # A: (l, l) + 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) # (l, l) + row_sum = P.sum(dim=-1, keepdim=True) # (l, 1) + uniform = torch.full_like(P, 1.0 / length) # (l, l) + P = torch.where( # (l, l) + 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) # (l,) + teleport = (1.0 - damping) / length + for _ in range(max_iterations): + p_next = teleport + damping * (P.transpose(0, 1) @ p) # (l,) + if torch.linalg.vector_norm(p_next - p, ord=1) <= tolerance: + p = p_next # (l,) + break + p = p_next # (l,) + return p / p.sum() # (l,) + + +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: + # X: (b, l, d); residue_mask: (b, l) + M = _validate_inputs(X, residue_mask) # (b, l) + M_expanded = M.unsqueeze(-1) # (b, l, 1) + count = M_expanded.sum(dim=1).clamp_min(1) # (b, 1) + X_residues = X.masked_fill(~M_expanded, 0) # (b, l, d) + outputs: list[Tensor] = [] + + for name in self.names: + if name == "mean": + Y = X_residues.sum(dim=1) / count # (b, d) + elif name == "max": + Y = X.masked_fill(~M_expanded, -torch.inf).max(dim=1).values # (b, d) + elif name == "norm": + Y = torch.linalg.vector_norm(X_residues, ord=2, dim=1) # (b, d) + elif name == "median": + Y = X.masked_fill(~M_expanded, torch.nan).nanmedian(dim=1).values # (b, d) + elif name in {"var", "std"}: + mean = X_residues.sum(dim=1, keepdim=True) / count.unsqueeze(1) # (b, 1, d) + centered = (X - mean).masked_fill(~M_expanded, 0) # (b, l, d) + variance = (centered**2).sum(dim=1) / count # (b, d) + Y = variance.sqrt() if name == "std" else variance # (b, d) + elif name == "cls": + Y = X[:, 0] # (b, d) + 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) # (b, l, l) + pooled: list[Tensor] = [] + for X_i, M_i, A_i in zip(X, M, A, strict=True): + # X_i: (l, d); M_i: (l,); A_i: (l, l) + indices = M_i.nonzero(as_tuple=True)[0] # (r,) + A_residue = A_i.index_select(0, indices).index_select(1, indices) # (r, r) + w = pagerank_weights(A_residue).to(dtype=X.dtype) # (r,) + pooled.append(w @ X_i.index_select(0, indices)) # (d,) + Y = torch.stack(pooled) # (b, d) + 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) # (b, len(self.names) * d) + + +__all__ = ["POOLING_NAMES", "Pooler", "pagerank_weights"] diff --git a/fastplms/embeddings/runner.py b/fastplms/embeddings/runner.py new file mode 100644 index 0000000000000000000000000000000000000000..200082f7742b04334455cc68ae7d67f7ffd6a589 --- /dev/null +++ b/fastplms/embeddings/runner.py @@ -0,0 +1,1583 @@ +"""Model-independent dataset embedding orchestration.""" + +from __future__ import annotations + +import hashlib +import json +import platform +import sqlite3 +import tempfile +import torch +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from contextlib import contextmanager +from pathlib import Path +from typing import Any, overload +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.""" + + # M: (b, l) + 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.""" + # last_hidden_state and each hidden_states entry: (b, l, d) + 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) # (b, n, l, d) + if hidden_state_index == -1: + return last_hidden_state # (b, l, d) + if not hidden_states: + raise ValueError("hidden_state_index requires model hidden states.") + return hidden_states[hidden_state_index] # (b, l, d) + + +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.""" + + # input_ids, attention_mask: (b, l) + M = attention_mask.to(dtype=torch.bool) # (b, l) + special_ids = tuple(int(token_id) for token_id in getattr(tokenizer, "all_special_ids", ())) + if special_ids: + specials = torch.tensor( # (n_special,) + special_ids, + device=input_ids.device, + dtype=input_ids.dtype, + ) + M = M & ~torch.isin(input_ids, specials) # (b, l) + return M # (b, l) + + +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 # (b, l, d), (b, l) + 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"] # (b, l) + boundary_ids = preparer.boundary_token_ids.to( # (n_boundary,) + 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) # (b, l) + if need_attentions: + raise ValueError("parti is not available for tokenizer-free E1 embedding.") + return EmbeddingBatch( # X: (b, l, d); residue_mask: (b, l) + 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) # (b, l) + attention_mask = encoded.get( # (b, l) + "attention_mask", + input_ids.new_ones(input_ids.shape), + ).to(device) + M = _biological_residue_mask(input_ids, attention_mask, tokenizer) # (b, l) + 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) # (b, l, d) + 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) # each: (b, h, l, l) + if attentions is None: + raise ValueError("The model did not return attentions required by parti.") + return EmbeddingBatch( # X: (b, l, d); M: (b, l) + 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.""" + + # X: (...) + 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] # (chunk_rows, ...) + 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() # chunk.shape + 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: + # decoder_input_ids: (n_records, l_decoder) + indices = torch.tensor( # (b,) + batch_positions, + device=decoder_input_ids.device, + dtype=torch.long, + ) + batch_model_kwargs["decoder_input_ids"] = ( # (b, l_decoder) + decoder_input_ids.index_select(0, indices) + ) + if decoder_attention_mask is not None: + # decoder_attention_mask: (n_records, l_decoder) + indices = torch.tensor( # (b,) + batch_positions, + device=decoder_attention_mask.device, + dtype=torch.long, + ) + batch_model_kwargs["decoder_attention_mask"] = ( + decoder_attention_mask.index_select(0, indices) # (b, l_decoder) + ) + 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 # (b, l, d) or (b, n_states, l, d) + raw_mask = batch.residue_mask # (b, l) + 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) # (b, l) + 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 = ( # X.shape + 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) # unchanged shape + + if full_embeddings: + if X.ndim == 4: + values = [ + X_i[:, M_i, :].detach().cpu() # (n_states, r_i, d) + for X_i, M_i in zip(X, M, strict=True) + ] + else: + values = [ + X_i[M_i].detach().cpu() # (r_i, d) + 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( # (b, n_poolers * d) + X, + M, + attentions=batch.attentions, + attention_backend=attention_backend, + ) + pool_slices = pooler.output_slices(X.shape[-1]) + values = list(Y.detach().cpu().unbind(0)) # each: (n_poolers * d,) + 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", +] diff --git a/fastplms/embeddings/storage.py b/fastplms/embeddings/storage.py new file mode 100644 index 0000000000000000000000000000000000000000..7fc4d1d02284bd5ccf1f7a0e330fba0d4a8449f3 --- /dev/null +++ b/fastplms/embeddings/storage.py @@ -0,0 +1,1601 @@ +"""Lossless, reproducible storage for :mod:`fastplms.embeddings`.""" + +from __future__ import annotations + +import hashlib +import io +import json +import sqlite3 +import struct +import numpy as np +import torch +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 +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 = 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.""" + + # X: (...) + flattened = X.detach().to(device="cpu").reshape(-1) # (n,) + 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] # (n_chunk,) + if chunk.stride(0) != 1: + chunk = chunk.clone(memory_format=torch.contiguous_format) # (n_chunk,) + yield chunk # (n_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() # (n_bytes,) + X = torch.from_numpy(byte_array).view(dtype) # (n_elements,) + return X.reshape(shape).clone() # shape + + +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()): + # X: (...) + 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(" 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 = ( # shape + np.frombuffer(data, dtype=numpy_dtype, offset=offset).copy().reshape(shape) + ) + return torch.from_numpy(array).to(dtype=target_dtype) # shape + + 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 + fallback_shape + ) + return torch.from_numpy(array) # fallback_shape + 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", +] diff --git a/fastplms/embeddings/types.py b/fastplms/embeddings/types.py new file mode 100644 index 0000000000000000000000000000000000000000..5166ebd2149db563a2a305f1d866b535d478b5d5 --- /dev/null +++ b/fastplms/embeddings/types.py @@ -0,0 +1,186 @@ +"""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() # self.shape + 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 # self.shape + + +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", +] diff --git a/fastplms/models.toml b/fastplms/models.toml new file mode 100644 index 0000000000000000000000000000000000000000..21e126f7f4a76085fdb2c9750ac5c81d5e294bb9 --- /dev/null +++ b/fastplms/models.toml @@ -0,0 +1,1275 @@ +schema_version = 1 +legal_files = [ + "LICENSE=sha256:2d2b50c7b1414bff1189a1db1f0cfb92e3e064b50f4c2b1019827b683e1b629a", + "THIRD_PARTY_NOTICES.md=sha256:d35e506b728868b52290672d54a89c6aade09529af99dbc8a4db72d9e9ca3460", +] + +[[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", + "SOURCE_RECORD.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", + "SOURCE_RECORD.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", + "SOURCE_RECORD.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", + "SOURCE_RECORD.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", "fp8"] +experimental_precisions = ["fp8"] +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", AutoModelForSequenceClassification = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esm_plusplus.modeling_esm_plusplus.ESMplusplusForTokenClassification" } + +[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", AutoModelForSequenceClassification = "fastplms.models.esm3.modeling_esm3.FastESM3ForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esm3.modeling_esm3.FastESM3ForTokenClassification" } + +[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/main/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/SOURCE_RECORD.md. Limitation: redistribution remains subject to Apache-2.0 and the pinned source 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/SOURCE_RECORD.md. Limitation: no head exception is permitted by this source 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 = false +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/classification_probe.py", "models/esmfold"] +auto_map = { AutoConfig = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmFoldConfig", AutoModel = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmForProteinFolding", AutoModelForSequenceClassification = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold.modeling_fast_esmfold.FastEsmForTokenClassification" } + +[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, the four 6B-backbone variants, and three BF16-to-FP8 reload cycles on the standard variant. Limitation: only the six 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/classification_probe.py", "models/_esm_rotary.py", "models/esmfold2", "models/esm_plusplus", "models/ttt.py"] +auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2.ESMFold2Model", AutoModelForSequenceClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ForTokenClassification" } + +[[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 = "a3afa1db21c876dff57b3540fa7241e138fb1ed6" +fast_files = [ + "config.json=git-sha1:d1b81bb97129bc75dea04daef1ea2af373018e6b", + "model-00001-of-00001.safetensors=sha256:c943d25cacdafd2c8e3518a74450b5f90f715becf30ceb24c327f1c5a0bc8b5d", + "model.safetensors.index.json=git-sha1:ca251ab9277c06081b33e027f68f5bdc0808b443", + "special_tokens_map.json=git-sha1:55b145827029ae9672e50d4bb368540daacce791", + "tokenizer.json=git-sha1:212c5ef08819fa2463c6289ba4ef7db30e715c0a", + "tokenizer_config.json=git-sha1:a8a872ae3441e7cc85ce19210dff1e4c5d2d7bd0", +] +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 = "92d2403bbe3c32acaa944fbb8dc2beb5f571f008" +fast_files = [ + "config.json=git-sha1:46eef0fff286107820f8ffc523127fa981435aeb", + "model-00001-of-00002.safetensors=sha256:79301f0b6a4fcbfd3b8bd10ca892846d79b1aad6ad06976da7380249e36f5158", + "model-00002-of-00002.safetensors=sha256:20062a5049fcde509030024527665a75062a95d64966558dfcaa9245b441cbec", + "model.safetensors.index.json=git-sha1:6b707ca3ce7255d241a52feeca68c0cbbe2a383f", + "special_tokens_map.json=git-sha1:55b145827029ae9672e50d4bb368540daacce791", + "tokenizer.json=git-sha1:212c5ef08819fa2463c6289ba4ef7db30e715c0a", + "tokenizer_config.json=git-sha1:d7fe02ba6f2b18d9ccfa19ac129c9fdc9ec24d09", +] +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 = "729167c1980316ae61691338838447491926033f" +fast_files = [ + "config.json=git-sha1:dd5d59e6b74bc8afa9fd4a5bda13526c235dabb8", + "generation_config.json=git-sha1:91f792e452403d46e170e206f9e50be5ddef9b9a", + "model-00001-of-00002.safetensors=sha256:7c0c297f60bcf81c732cdfeae6e99e140272807eb52afd70356fc6fdfa94e5a8", + "model-00002-of-00002.safetensors=sha256:f3d425d3e8741ccbdd925446559a9bf317c2c91e328f2eee44924423b56e3a3d", + "model.safetensors.index.json=git-sha1:6b707ca3ce7255d241a52feeca68c0cbbe2a383f", + "special_tokens_map.json=git-sha1:55b145827029ae9672e50d4bb368540daacce791", + "tokenizer.json=git-sha1:212c5ef08819fa2463c6289ba4ef7db30e715c0a", + "tokenizer_config.json=git-sha1:854e5db75dae8b1e9dd39c5bae80dae5508b3e25", +] +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 = "c6d16ca2a1b3b27a27bcf3875e816a059029d264" +fast_files = [ + "config.json=git-sha1:813ffc6c319549a2c1f3503e1309c36110202d65", + "generation_config.json=git-sha1:5767cc0cacebfd06884eb27ae1c796d3ca829fd2", + "model-00001-of-00002.safetensors=sha256:7f1f5c5dcff4b6bc6b8464fe9a7eebdd99b0789ee8da895f42a41bdb04191654", + "model-00002-of-00002.safetensors=sha256:c1a67cef9b76202362ff00c9d2b2dc4b5fc7acd1f22d30c8b3f2e3d2597d0f22", + "model.safetensors.index.json=git-sha1:a20cfc1f8517ef47d12d08604dc93c064f1e6736", + "special_tokens_map.json=git-sha1:d596919b7fa2a197edd441ec3ec4685ecacd2de4", + "spiece.model=sha256:f2b5e1bbd110b71ca9b2878e1fcd3265610076ecc97bd696e8a745c9bacc54e0", + "tokenizer.json=git-sha1:90f0c94b43c81496b3ca81e3ec1c092ef2dd7fca", + "tokenizer_config.json=git-sha1:0e699eebfa778698473b4faf1e66ef363b93fb21", +] +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 = "d2856892e7535af2f55c2c4de043b1b272a29ed8" +fast_files = [ + "config.json=git-sha1:791460a5c0d6c03bebbac1d7eec7e35805eaf7b7", + "generation_config.json=git-sha1:91f792e452403d46e170e206f9e50be5ddef9b9a", + "model-00001-of-00005.safetensors=sha256:f6b841f6b800e436b08e362d04f8442fd044839b1e32ba6ac01ecb30a9d2bae5", + "model-00002-of-00005.safetensors=sha256:ea556d511d4747ada49b9d0c24ef503774410093e53c0d003ffb9407efc2be31", + "model-00003-of-00005.safetensors=sha256:125884f5dcb5b44435e3b76330582f31f74547b67c9b6cadf9a4d7cf38748eb7", + "model-00004-of-00005.safetensors=sha256:e776ef6c5d6a50d4b3fcf0bbb2431de7ce813eddd2903290e54809c801ddb241", + "model-00005-of-00005.safetensors=sha256:58ee3b065cfcccd179fdbecef9827dfb91feb33e0d8385b692e6309d56ec530e", + "model.safetensors.index.json=git-sha1:74d149f64234c3f43eb85971f40b4a1c6d05a407", + "special_tokens_map.json=git-sha1:d596919b7fa2a197edd441ec3ec4685ecacd2de4", + "spiece.model=sha256:f2b5e1bbd110b71ca9b2878e1fcd3265610076ecc97bd696e8a745c9bacc54e0", + "tokenizer.json=git-sha1:90f0c94b43c81496b3ca81e3ec1c092ef2dd7fca", + "tokenizer_config.json=git-sha1:0e699eebfa778698473b4faf1e66ef363b93fb21", +] +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", AutoModelForSequenceClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForTokenClassification" } + +[[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", AutoModelForSequenceClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForTokenClassification" } + +[[models]] +id = "esmfold2_300" +family = "esmfold2" +size_category = "structure" +generation_contract = "not_applicable" +msa_conditioning = false +publication_status = "pending" +fast_repo = "Synthyra/ESMFold2-300" +fast_revision = "unpublished" +fast_files = ["config.json=git-sha1:8c9a04fe22b0e5fca77bc4e2861a12c9494ef4d4", "model.safetensors=sha256:44d6797c5efebf24753d502b40950e0874871c96ceea14f2d7f7e39cebac67fd"] +official_repo = "biohub/ESMFold2-Experimental-Fast-base300M-step1500k" +official_revision = "21531e59002c9205284715e28ee802dafb430637" +official_files = ["config.json=git-sha1:8c9a04fe22b0e5fca77bc4e2861a12c9494ef4d4", "model.safetensors=sha256:44d6797c5efebf24753d502b40950e0874871c96ceea14f2d7f7e39cebac67fd"] +notes = "Experimental Fast checkpoint with a frozen 300M ESM++ backbone, tensor-exact in BF16 with the pinned step-1500000 source, 24 folding blocks, no MSA conditioning, and no confidence head. BF16 execution uses FP32 folding parameters with CUDA autocast. FP8 is unsupported. Docker BF16 inference validation passed on the compact Protein G case; publication is pending." +auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2_experimental.ESMFold2ExperimentalModel", AutoModelForSequenceClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForTokenClassification" } +backbone_model = "esmc_small" +backbone = { repo = "biohub/ESMC-300M-1500000", revision = "56803b6378b82e16c3b24aac49d1fce4445540b7", files = ["config.json=git-sha1:7fe728a0eb3fb81b24491d6cc1de816bf7797c27", "model.safetensors=sha256:8bd6cacf9b5a92d51954b64b20407f1f9f564a7e4849b8663470784d2a8b7ed2", "tokenizer.json=git-sha1:81c797f56768b22dec0301fa771f018b7e43e98c", "tokenizer_config.json=git-sha1:f49f57b24a1c93bd544974811e8ecbd61b7fae89", "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b"] } + +[[models]] +id = "esmfold2_600" +family = "esmfold2" +size_category = "structure" +generation_contract = "not_applicable" +msa_conditioning = false +publication_status = "pending" +fast_repo = "Synthyra/ESMFold2-600" +fast_revision = "unpublished" +fast_files = ["config.json=git-sha1:95517e555f17a1eca4b68866c89033a1f6916a5d", "model.safetensors=sha256:11a53c1b4700b6c62a5a464fc3ec7076160c19e8584157f88e13275116cfd602"] +official_repo = "biohub/ESMFold2-Experimental-Fast-base600M-step1500k" +official_revision = "15cf2d6648692f6c17cee1297d8a285476fffa9b" +official_files = ["config.json=git-sha1:95517e555f17a1eca4b68866c89033a1f6916a5d", "model.safetensors=sha256:11a53c1b4700b6c62a5a464fc3ec7076160c19e8584157f88e13275116cfd602"] +notes = "Experimental Fast checkpoint with a frozen 600M ESM++ backbone, tensor-exact in BF16 with the pinned step-1500000 source, 24 folding blocks, no MSA conditioning, and no confidence head. BF16 execution uses FP32 folding parameters with CUDA autocast. FP8 is unsupported. Configuration, weight identities, and artifact reload are verified; this model is not inference-validated. Publication is pending." +auto_map = { AutoConfig = "fastplms.models.esmfold2.configuration_esmfold2.ESMFold2Config", AutoModel = "fastplms.models.esmfold2.modeling_esmfold2_experimental.ESMFold2ExperimentalModel", AutoModelForSequenceClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForSequenceClassification", AutoModelForTokenClassification = "fastplms.models.esmfold2.modeling_esmfold2_classification.ESMFold2ExperimentalForTokenClassification" } +backbone_model = "esmc_large" +backbone = { repo = "biohub/ESMC-600M-1500000", revision = "21af9cc429af76ebda6c48074fb624db4735aaaf", files = ["config.json=git-sha1:ec29f6009b21d710f64bf1c058f3a9710833d692", "model.safetensors=sha256:d6869f5ae0f11e5dc829b195e062e87cfcc2f851a08a5edbaf5d1083ae7f76cc", "tokenizer.json=git-sha1:81c797f56768b22dec0301fa771f018b7e43e98c", "tokenizer_config.json=git-sha1:f49f57b24a1c93bd544974811e8ecbd61b7fae89", "special_tokens_map.json=git-sha1:c907ee1dc19b24241749b32d665c291c7e6e8e4b"] } diff --git a/fastplms/models/__init__.py b/fastplms/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..340b1a15b14b256253ce91ab3eb3114a394ac8e9 --- /dev/null +++ b/fastplms/models/__init__.py @@ -0,0 +1,11 @@ +"""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, ...] = () diff --git a/fastplms/models/_esm_rotary.py b/fastplms/models/_esm_rotary.py new file mode 100644 index 0000000000000000000000000000000000000000..57b3c5c29efb2a4a26755c24bc49cf9872cc5eea --- /dev/null +++ b/fastplms/models/_esm_rotary.py @@ -0,0 +1,92 @@ +"""Stable ESM rotary embeddings independent of Transformers internals. + +Transformers 5 changed both the name and call contract of its private ESM +rotary helper. FastPLMs checkpoints use the earlier two-tensor contract, so +the small mathematical primitive lives here instead of importing a private +Transformers implementation. +""" + +from __future__ import annotations + +import torch +from torch import nn + + +def _rotate_half(tensor: torch.Tensor) -> torch.Tensor: + """Rotate the final dimension of X by 90 degrees in paired subspaces.""" + + # tensor: (..., d) + first, second = tensor.chunk(2, dim=-1) # (..., d / 2), (..., d / 2) + return torch.cat((-second, first), dim=-1) # (..., d) + + +def apply_rotary_pos_emb( + tensor: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> torch.Tensor: + """Apply cached rotary factors to X with shape ``(b, h, l, d)``.""" + + # tensor: (b, h, l, d); cos, sin: (1, 1, l_cache, d) + cos = cos[:, :, : tensor.shape[-2], :] # (1, 1, l, d) + sin = sin[:, :, : tensor.shape[-2], :] # (1, 1, l, d) + return tensor * cos + _rotate_half(tensor) * sin # (b, h, l, d) + + +class RotaryEmbedding(nn.Module): + """Apply rotary position embeddings to query and key tensors.""" + + inv_freq: torch.Tensor + + def __init__(self, dim: int) -> None: + super().__init__() + frequencies = 1.0 / ( # (d / 2,) + 10_000 ** (torch.arange(0, dim, 2, dtype=torch.int64).float() / dim) + ) + # Keep this persistent to preserve the historical checkpoint schema. + self.register_buffer("inv_freq", frequencies) + self._seq_len_cached: int | None = None + self._cos_cached: torch.Tensor | None = None + self._sin_cached: torch.Tensor | None = None + + def _update_cos_sin_tables( + self, + tensor: torch.Tensor, + seq_dimension: int = 2, + ) -> tuple[torch.Tensor, torch.Tensor]: + # tensor: (..., l, d) + seq_len = tensor.shape[seq_dimension] + cache_stale = ( + self._cos_cached is None + or self._sin_cached is None + or self._seq_len_cached != seq_len + or self._cos_cached.device != tensor.device + ) + if cache_stale: + self._seq_len_cached = seq_len + positions = torch.arange(seq_len, device=tensor.device).type_as( # (l,) + self.inv_freq + ) + angles = torch.outer(positions, self.inv_freq) # (l, d / 2) + angles = torch.cat((angles, angles), dim=-1).to(tensor.device) # (l, d) + self._cos_cached = angles.cos()[None, None, :, :] # (1, 1, l, d) + self._sin_cached = angles.sin()[None, None, :, :] # (1, 1, l, d) + + assert self._cos_cached is not None + assert self._sin_cached is not None + return self._cos_cached, self._sin_cached # (1, 1, l, d), (1, 1, l, d) + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + # query, key: (b, h, l, d) + cos, sin = self._update_cos_sin_tables( # (1, 1, l, d), (1, 1, l, d) + key, + seq_dimension=-2, + ) + return ( + apply_rotary_pos_emb(query, cos, sin).to(dtype=query.dtype), # (b, h, l, d) + apply_rotary_pos_emb(key, cos, sin).to(dtype=key.dtype), # (b, h, l, d) + ) diff --git a/fastplms/models/classification_probe.py b/fastplms/models/classification_probe.py new file mode 100644 index 0000000000000000000000000000000000000000..61ab5d77ddcadbdc87e3b7d200f4541b28389a3e --- /dev/null +++ b/fastplms/models/classification_probe.py @@ -0,0 +1,542 @@ +"""Shared transformer probes for residue and sequence prediction tasks.""" + +from __future__ import annotations + +import math +from typing import Any + +import torch +from torch import nn +from torch.nn import functional as F +from transformers.modeling_outputs import ( + BaseModelOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) + +try: + from fastplms.attention import ( + AttentionBackend, + _get_flex_attention_fn, + flex_attention, + get_attention_mask, + resolve_attention_backend, + ) + from fastplms.embeddings.pooling import Pooler + from fastplms.models._esm_rotary import RotaryEmbedding +except ModuleNotFoundError as error: + _COMPOSITE_REQUIRED_NAMES = ( + "AttentionBackend", + "Pooler", + "RotaryEmbedding", + "_get_flex_attention_fn", + "flex_attention", + "get_attention_mask", + "resolve_attention_backend", + ) + if error.name != "fastplms" or any( + name not in globals() for name in _COMPOSITE_REQUIRED_NAMES + ): + raise + # Flat Hub composites define every shared symbol above this source. + + +_SUPPORTED_BACKENDS = frozenset( + { + AttentionBackend.EAGER, + AttentionBackend.SDPA, + AttentionBackend.FLEX_ATTENTION, + } +) +_SUPPORTED_PROBLEM_TYPES = frozenset( + { + "regression", + "single_label_classification", + "multi_label_classification", + } +) +_UNSUPPORTED_POOLING = frozenset({"cls", "parti"}) + + +def _config_value(config: Any, name: str, default: Any) -> Any: + value = getattr(config, name, None) + return default if value is None else value + + +def _attention_backend(config: Any) -> AttentionBackend: + requested = getattr(config, "_attn_implementation", None) + if requested is None: + requested = getattr(config, "attn_backend", "sdpa") + backend = resolve_attention_backend(requested) + if backend not in _SUPPORTED_BACKENDS: + expected = ", ".join(sorted(item.value for item in _SUPPORTED_BACKENDS)) + raise ValueError( + f"Classification probes support only {expected}; received {backend.value!r}." + ) + return backend + + +def resolve_problem_type( + config: Any, + labels: torch.Tensor, + *, + num_labels: int, +) -> str: + """Resolve and persist the standard Transformers classification problem type.""" + + problem_type = getattr(config, "problem_type", None) + if problem_type is None: + if num_labels == 1: + problem_type = "regression" + elif labels.dtype in {torch.long, torch.int}: + problem_type = "single_label_classification" + else: + problem_type = "multi_label_classification" + config.problem_type = problem_type + if problem_type not in _SUPPORTED_PROBLEM_TYPES: + raise ValueError( + f"Unsupported problem_type {problem_type!r}; expected one of " + f"{sorted(_SUPPORTED_PROBLEM_TYPES)}." + ) + return problem_type + + +def sequence_classification_loss( + logits: torch.Tensor, + labels: torch.Tensor, + *, + problem_type: str, + num_labels: int, +) -> torch.Tensor: + """Compute a Hugging Face-compatible sequence task loss.""" + + labels = labels.to(logits.device) + if problem_type == "regression": + if num_labels == 1: + return F.mse_loss(logits.squeeze(-1), labels.squeeze(-1).to(logits.dtype)) + return F.mse_loss(logits, labels.to(logits.dtype)) + if problem_type == "single_label_classification": + return F.cross_entropy(logits.reshape(-1, num_labels), labels.reshape(-1).long()) + if problem_type == "multi_label_classification": + return F.binary_cross_entropy_with_logits(logits, labels.to(logits.dtype)) + raise ValueError(f"Unsupported problem_type {problem_type!r}.") + + +def _masked_elementwise_loss( + losses: torch.Tensor, + labels: torch.Tensor, +) -> torch.Tensor: + valid = labels.ne(-100) + if not bool(valid.any()): + return losses.sum() * 0 + return losses.masked_select(valid).mean() + + +def token_classification_loss( + logits: torch.Tensor, + labels: torch.Tensor, + *, + problem_type: str, + num_labels: int, +) -> torch.Tensor: + """Compute a token task loss, excluding every label element equal to ``-100``.""" + + labels = labels.to(logits.device) + if problem_type == "regression": + targets = labels.to(logits.dtype) + if num_labels == 1 and targets.ndim == logits.ndim - 1: + targets = targets.unsqueeze(-1) + if targets.shape != logits.shape: + raise ValueError( + "Token regression labels must match logits, except that the final " + "singleton dimension may be omitted when num_labels=1." + ) + return _masked_elementwise_loss(F.mse_loss(logits, targets, reduction="none"), targets) + if problem_type == "single_label_classification": + if not bool(labels.ne(-100).any()): + return logits.sum() * 0 + return F.cross_entropy( + logits.reshape(-1, num_labels), + labels.reshape(-1).long(), + ignore_index=-100, + ) + if problem_type == "multi_label_classification": + if labels.shape != logits.shape: + raise ValueError("Multilabel token labels must have the same shape as logits.") + losses = F.binary_cross_entropy_with_logits( + logits, + labels.to(logits.dtype), + reduction="none", + ) + return _masked_elementwise_loss(losses, labels) + raise ValueError(f"Unsupported problem_type {problem_type!r}.") + + +class SwiGLU(nn.Module): + """SwiGLU activation used by the Protify-aligned feed-forward layer.""" + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + gate, values = inputs.chunk(2, dim=-1) + return F.silu(gate) * values + + +class ProbeSelfAttention(nn.Module): + """Four-head RoPE self-attention with explicit, fail-closed dispatch.""" + + def __init__( + self, + hidden_size: int, + num_heads: int, + dropout: float, + backend: AttentionBackend, + use_bias: bool, + ) -> None: + super().__init__() + if hidden_size % num_heads: + raise ValueError("classifier_probe_hidden_size must be divisible by its head count.") + self.hidden_size = hidden_size + self.num_heads = num_heads + self.head_size = hidden_size // num_heads + self.dropout = dropout + self.backend = backend + self.qkv = nn.Linear(hidden_size, 3 * hidden_size, bias=use_bias) + self.output = nn.Linear(hidden_size, hidden_size, bias=use_bias) + self.rotary = RotaryEmbedding(self.head_size) + + def _reshape(self, tensor: torch.Tensor) -> torch.Tensor: + batch_size, sequence_length, _ = tensor.shape + return tensor.view( + batch_size, + sequence_length, + self.num_heads, + self.head_size, + ).transpose(1, 2) + + def forward( + self, + hidden_states: torch.Tensor, + *, + attention_mask: torch.Tensor | None, + output_attentions: bool, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + batch_size, sequence_length, _ = hidden_states.shape + query, key, value = self.qkv(hidden_states).chunk(3, dim=-1) + query = self._reshape(query) + key = self._reshape(key) + value = self._reshape(value) + query, key = self.rotary(query, key) + if output_attentions and self.backend != AttentionBackend.EAGER: + raise ValueError( + f"output_attentions=True is unavailable for {self.backend.value!r}; " + "select 'eager' explicitly." + ) + _, attention_mask_4d, flex_block_mask = get_attention_mask( + self.backend, + batch_size, + sequence_length, + hidden_states.device, + attention_mask, + hidden_states.dtype, + ) + dropout = self.dropout if self.training else 0.0 + attention_weights = None + if self.backend == AttentionBackend.EAGER: + scores = query @ key.transpose(-2, -1) / math.sqrt(self.head_size) + if attention_mask_4d is not None: + scores = scores.masked_fill(~attention_mask_4d, float("-inf")) + attention_weights = scores.softmax(dim=-1) + context = F.dropout(attention_weights, p=dropout, training=self.training) @ value + elif self.backend == AttentionBackend.SDPA: + context = F.scaled_dot_product_attention( + query, + key, + value, + attn_mask=attention_mask_4d, + dropout_p=dropout, + ) + elif self.backend == AttentionBackend.FLEX_ATTENTION: + if flex_attention is None: + raise RuntimeError("'flex_attention' was requested but is unavailable.") + flex_fn = _get_flex_attention_fn( + device=query.device, + dtype=query.dtype, + shape=tuple(query.shape), + mask_semantics="padding", + ) + if flex_fn is None: + raise RuntimeError("'flex_attention' was requested but is unavailable.") + context = flex_fn( + query, + key, + value, + block_mask=flex_block_mask, + scale=1.0 / math.sqrt(self.head_size), + kernel_options={"PRESCALE_QK": True, "BLOCK_N": 32}, + ) + else: + raise AssertionError(f"Unhandled attention backend {self.backend.value!r}.") + context = context.transpose(1, 2).contiguous().view( + batch_size, + sequence_length, + self.hidden_size, + ) + return self.output(context), attention_weights + + +class ProteinTransformerProbe(nn.Module): + """Project residue embeddings and refine them with exactly one pre-LN block.""" + + def __init__(self, config: Any, input_size: int) -> None: + super().__init__() + hidden_size = int(_config_value(config, "classifier_probe_hidden_size", 512)) + num_heads = int(_config_value(config, "classifier_probe_num_heads", 4)) + dropout = float(_config_value(config, "classifier_probe_dropout", 0.1)) + use_bias = bool( + _config_value( + config, + "classifier_use_bias", + _config_value(config, "use_bias", False), + ) + ) + if hidden_size != 512 or num_heads != 4 or hidden_size // num_heads != 128: + raise ValueError( + "The folding classification probe requires a 512-wide projection with " + "four 128-wide attention heads." + ) + self.hidden_size = hidden_size + self.input_norm = nn.LayerNorm(input_size) + self.input_projection = nn.Linear(input_size, hidden_size, bias=use_bias) + self.attention_norm = nn.LayerNorm(hidden_size) + self.attention = ProbeSelfAttention( + hidden_size, + num_heads, + dropout, + _attention_backend(config), + use_bias, + ) + intermediate_size = int(math.ceil((8 / 3) * hidden_size / 256) * 256) + self.feed_forward_norm = nn.LayerNorm(hidden_size) + self.feed_forward = nn.Sequential( + nn.Linear(hidden_size, 2 * intermediate_size, bias=use_bias), + SwiGLU(), + nn.Dropout(dropout), + nn.Linear(intermediate_size, hidden_size, bias=use_bias), + ) + self.residual_dropout = nn.Dropout(dropout) + + @property + def attn_backend(self) -> str: + return self.attention.backend.value + + def forward( + self, + embeddings: torch.Tensor, + attention_mask: torch.Tensor | None = None, + *, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = True, + ) -> BaseModelOutput | tuple[torch.Tensor, ...]: + if embeddings.ndim != 3: + raise ValueError("embeddings must have shape (batch, residue, channel).") + embeddings = embeddings.to(dtype=self.input_projection.weight.dtype) + hidden_states = self.input_projection(self.input_norm(embeddings)) + attention_output, attention_weights = self.attention( + self.attention_norm(hidden_states), + attention_mask=attention_mask, + output_attentions=output_attentions, + ) + hidden_states = hidden_states + self.residual_dropout(attention_output) + hidden_states = hidden_states + self.residual_dropout( + self.feed_forward(self.feed_forward_norm(hidden_states)) + ) + output = BaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=(hidden_states,) if output_hidden_states else None, + attentions=(attention_weights,) if output_attentions else None, + ) + return output if return_dict else output.to_tuple() + + +class _ClassificationProbe(nn.Module): + def __init__(self, config: Any, input_size: int, *, sequence_task: bool) -> None: + super().__init__() + self.config = config + self.num_labels = int(_config_value(config, "num_labels", 2)) + self.transformer = ProteinTransformerProbe(config, input_size) + self.sequence_task = sequence_task + pooling_types = _config_value(config, "classifier_pooling_types", ["mean"]) + self.pooler = Pooler(pooling_types) if sequence_task else None + if self.pooler is not None: + unsupported = sorted(set(self.pooler.names) & _UNSUPPORTED_POOLING) + if unsupported: + raise ValueError( + "Classification probes consume residue-only representations and do not " + f"support pooling operation(s) {unsupported}." + ) + hidden_size = self.transformer.hidden_size + classifier_input = hidden_size * (len(self.pooler.names) if self.pooler else 1) + classifier_hidden = int(_config_value(config, "classifier_hidden_size", 4096)) + classifier_dropout = float(_config_value(config, "classifier_dropout", 0.2)) + use_bias = bool( + _config_value( + config, + "classifier_use_bias", + _config_value(config, "use_bias", False), + ) + ) + projection_size = int(math.ceil((2 * self.num_labels) / 256) * 256) + classifier_layers: list[nn.Module] = [ + nn.LayerNorm(classifier_input), + nn.Linear(classifier_input, classifier_hidden, bias=use_bias), + nn.ReLU(), + nn.Dropout(classifier_dropout), + nn.Linear(classifier_hidden, projection_size, bias=use_bias), + nn.ReLU(), + nn.Dropout(classifier_dropout), + ] + if not sequence_task: + classifier_layers.extend( + [ + nn.Linear(projection_size, projection_size, bias=use_bias), + nn.ReLU(), + ] + ) + classifier_layers.append(nn.Linear(projection_size, self.num_labels, bias=use_bias)) + self.classifier = nn.Sequential(*classifier_layers) + + def _forward_transformer( + self, + embeddings: torch.Tensor, + attention_mask: torch.Tensor | None, + output_attentions: bool | None, + output_hidden_states: bool | None, + ) -> BaseModelOutput: + output_attentions = ( + bool(output_attentions) + if output_attentions is not None + else bool(getattr(self.config, "output_attentions", False)) + ) + output_hidden_states = ( + bool(output_hidden_states) + if output_hidden_states is not None + else bool(getattr(self.config, "output_hidden_states", False)) + ) + return self.transformer( + embeddings, + attention_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=True, + ) + + +class SequenceClassificationProbe(_ClassificationProbe): + """Protify-style sequence classifier over externally supplied residue embeddings.""" + + def __init__(self, config: Any, input_size: int) -> None: + super().__init__(config, input_size, sequence_task=True) + + def forward( + self, + embeddings: torch.Tensor, + attention_mask: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ) -> SequenceClassifierOutput | tuple[torch.Tensor, ...]: + if attention_mask is None: + attention_mask = torch.ones( + embeddings.shape[:2], + device=embeddings.device, + dtype=torch.bool, + ) + outputs = self._forward_transformer( + embeddings, + attention_mask, + output_attentions, + output_hidden_states, + ) + if self.pooler is None: + raise AssertionError("Sequence classification requires a configured pooler.") + pooled = self.pooler(outputs.last_hidden_state, attention_mask) + logits = self.classifier(pooled) + loss = None + if labels is not None: + problem_type = resolve_problem_type(self.config, labels, num_labels=self.num_labels) + loss = sequence_classification_loss( + logits, + labels, + problem_type=problem_type, + num_labels=self.num_labels, + ) + result = SequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + use_return_dict = ( + bool(return_dict) + if return_dict is not None + else bool(getattr(self.config, "use_return_dict", True)) + ) + return result if use_return_dict else result.to_tuple() + + +class TokenClassificationProbe(_ClassificationProbe): + """Protify-style residue classifier or regressor over supplied embeddings.""" + + def __init__(self, config: Any, input_size: int) -> None: + super().__init__(config, input_size, sequence_task=False) + + def forward( + self, + embeddings: torch.Tensor, + attention_mask: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ) -> TokenClassifierOutput | tuple[torch.Tensor, ...]: + outputs = self._forward_transformer( + embeddings, + attention_mask, + output_attentions, + output_hidden_states, + ) + logits = self.classifier(outputs.last_hidden_state) + loss = None + if labels is not None: + problem_type = resolve_problem_type(self.config, labels, num_labels=self.num_labels) + loss = token_classification_loss( + logits, + labels, + problem_type=problem_type, + num_labels=self.num_labels, + ) + result = TokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + use_return_dict = ( + bool(return_dict) + if return_dict is not None + else bool(getattr(self.config, "use_return_dict", True)) + ) + return result if use_return_dict else result.to_tuple() + + +__all__ = [ + "ProbeSelfAttention", + "ProteinTransformerProbe", + "SequenceClassificationProbe", + "SwiGLU", + "TokenClassificationProbe", + "resolve_problem_type", + "sequence_classification_loss", + "token_classification_loss", +] diff --git a/fastplms/models/esm_plusplus/__init__.py b/fastplms/models/esm_plusplus/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/fastplms/models/esm_plusplus/modeling_esm_plusplus.py b/fastplms/models/esm_plusplus/modeling_esm_plusplus.py new file mode 100644 index 0000000000000000000000000000000000000000..d07fa56f11d7d534bd07532c76fe548a5fe57b5e --- /dev/null +++ b/fastplms/models/esm_plusplus/modeling_esm_plusplus.py @@ -0,0 +1,2056 @@ +"""Hugging Face-compatible ESMC models implemented by FastPLMs.""" + +from __future__ import annotations + +import importlib +import importlib.metadata +import math +import os +from collections.abc import Sequence +from contextlib import contextmanager +from dataclasses import asdict, dataclass +from functools import partial +from typing import Any, ClassVar + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange +from tokenizers import Tokenizer +from tokenizers.models import BPE +from tokenizers.processors import TemplateProcessing +from transformers import PretrainedConfig, PreTrainedModel, PreTrainedTokenizerFast +from transformers.modeling_outputs import ( + MaskedLMOutput, + ModelOutput, + SequenceClassifierOutput, + TokenClassifierOutput, +) + +from .modeling_esm_plusplus_sae import ESMplusplusSAELayer, load_esmc_sae_layers + + +try: + from fastplms.attention import ( + AttentionBackend, + BlockMask, + FastPLMsAttentionMixin, + _get_flex_attention_fn, + _get_flex_block_mask, + flex_attention, + get_attention_mask, + kernels_flash_attention_func, + resolve_attention_backend, + resolve_attention_backend_for_call, + ) + from fastplms.embeddings import EmbeddingMixin, Pooler, select_hidden_state_embeddings + from fastplms.models.ttt import FastPLMTestTimeTrainingMixin +except ModuleNotFoundError as error: + _COMPOSITE_REQUIRED_NAMES = ( + "AttentionBackend", + "BlockMask", + "EmbeddingMixin", + "FastPLMsAttentionMixin", + "FastPLMTestTimeTrainingMixin", + "Pooler", + "_get_flex_attention_fn", + "_get_flex_block_mask", + "flex_attention", + "get_attention_mask", + "kernels_flash_attention_func", + "resolve_attention_backend", + "resolve_attention_backend_for_call", + "select_hidden_state_embeddings", + ) + if error.name != "fastplms" or any( + name not in globals() for name in _COMPOSITE_REQUIRED_NAMES + ): + raise + # Legacy flat Hub composites define every shared symbol above this block. + + +_ESMC_FP8_LINEAR_SUFFIX = ".attn.out_proj" +_ESMC_FP8_ALIGNMENT = 16 + + +@dataclass(frozen=True, slots=True) +class ESMplusplusFP8Status: + """Resolved state of the explicit Transformer Engine ESMC FP8 path.""" + + enabled: bool + reason: str + device: str + transformer_engine_version: str | None + converted_projections: int + + def as_dict(self) -> dict[str, str | int | bool | None]: + return asdict(self) + + +def _transformer_engine_version() -> str | None: + try: + return importlib.metadata.version("transformer-engine") + except importlib.metadata.PackageNotFoundError: + return None + + +def _load_transformer_engine() -> tuple[Any, Any]: + """Load Transformer Engine lazily so ordinary ESM++ imports stay portable.""" + + try: + te = importlib.import_module("transformer_engine.pytorch") + recipe = importlib.import_module("transformer_engine.common.recipe") + except (ImportError, OSError, RuntimeError) as error: + raise RuntimeError( + f"Transformer Engine could not be imported: {type(error).__name__}: {error}" + ) from error + if not hasattr(recipe, "Float8CurrentScaling"): + raise RuntimeError( + "Transformer Engine does not expose Float8CurrentScaling, which is " + "required by the validated ESMC FP8 path." + ) + return te, recipe + + +def _te_fp8_capability(device: torch.device) -> tuple[bool, str]: + """Return whether the strict Transformer Engine FP8 path can run.""" + + if device.type != "cuda": + return False, "FP8 requires ESM++ on a CUDA device." + if not torch.cuda.is_available(): + return False, "CUDA is unavailable." + try: + major, minor = torch.cuda.get_device_capability(device) + except (AssertionError, RuntimeError, ValueError) as error: + return False, f"CUDA capability query failed: {error}" + if not (major >= 9 or (major == 8 and minor >= 9)): + return False, f"CUDA capability {major}.{minor} does not support FP8." + try: + te, _ = _load_transformer_engine() + except RuntimeError as error: + return False, str(error) + + probe = getattr(te, "is_fp8_available", None) + if probe is None: + try: + probe = importlib.import_module("transformer_engine.pytorch.fp8").is_fp8_available + except (ImportError, AttributeError, OSError, RuntimeError) as error: + return False, f"Transformer Engine has no usable FP8 probe: {error}" + try: + try: + result = probe(return_reason=True) + except TypeError: + result = probe() + except (OSError, RuntimeError) as error: + return False, f"Transformer Engine FP8 probe failed: {error}" + if isinstance(result, tuple): + available = bool(result[0]) + detail = str(result[1]) if len(result) > 1 and result[1] else "" + else: + available = bool(result) + detail = "" + if not available: + return False, detail or "Transformer Engine reports FP8 unavailable." + return True, "Transformer Engine reports FP8 availability." + + +def _convert_esmc_attention_outputs_to_te( + module: nn.Module, + *, + expected_projections: int, +) -> tuple[str, ...]: + """Replace exactly one attention output projection per ESMC block.""" + + targets = [ + (path, child) + for path, child in module.named_modules() + if isinstance(child, nn.Linear) and path.endswith(_ESMC_FP8_LINEAR_SUFFIX) + ] + if len(targets) != expected_projections: + raise RuntimeError( + "ESMC FP8 conversion expected exactly " + f"{expected_projections} attention output projections, found {len(targets)}." + ) + + te, _ = _load_transformer_engine() + modules = dict(module.named_modules()) + converted: list[str] = [] + for path, child in targets: + owner_path, name = path.rsplit(".", 1) + owner = modules[owner_path] + replacement = te.Linear( + child.in_features, + child.out_features, + bias=child.bias is not None, + params_dtype=child.weight.dtype, + device=child.weight.device, + ) + with torch.no_grad(): + replacement.weight.copy_(child.weight) + if child.bias is not None: + replacement.bias.copy_(child.bias) + replacement.eval().requires_grad_(False) + setattr(owner, name, replacement) + converted.append(path) + return tuple(converted) + + +@contextmanager +def _esmplusplus_fp8_context(enabled: bool, device: torch.device): + """Enter the validated BF16-storage, Transformer Engine FP8 context.""" + + if not enabled: + yield + return + if torch.is_grad_enabled(): + raise RuntimeError("ESM++ FP8 is inference-only; use torch.inference_mode() or no_grad().") + te, recipe = _load_transformer_engine() + fp8_recipe = recipe.Float8CurrentScaling( + use_power_2_scales=False, + fp8_format=recipe.Format.HYBRID, + ) + with torch.autocast(device_type=device.type, dtype=torch.bfloat16): + with te.autocast(enabled=True, recipe=fp8_recipe): + yield + + +class ESMplusplusConfig(PretrainedConfig): + """Configuration class for ESM++ model. + + Args: + vocab_size: Size of the vocabulary + hidden_size: Dimension of hidden layers + num_attention_heads: Number of attention heads + num_hidden_layers: Number of transformer layers + num_labels: Number of output labels for classification + problem_type: Type of problem - regression, single/multi label classification + """ + + model_type = "ESMplusplus" + + def __init__( + self, + vocab_size: int = 64, + hidden_size: int = 960, + num_attention_heads: int = 15, + num_hidden_layers: int = 30, + num_labels: int | None = None, + problem_type: str | None = None, + dropout: float = 0.0, + initializer_range: float = 0.02, + classifier_dropout: float = 0.1, + classifier_pooling_types: list[str] | None = None, + attn_backend: str | None = None, + pad_token_id: int = 1, + mask_token_id: int = 32, + **kwargs, + ): + if num_labels is None: + configured_labels = kwargs.get("id2label") + num_labels = len(configured_labels) if configured_labels else 2 + super().__init__( + pad_token_id=pad_token_id, + mask_token_id=mask_token_id, + num_labels=num_labels, + **kwargs, + ) + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.num_attention_heads = num_attention_heads + self.num_hidden_layers = num_hidden_layers + self.problem_type = problem_type + self.dropout = dropout + self.initializer_range = initializer_range + self.classifier_dropout = classifier_dropout + self.classifier_pooling_types = ( + list(classifier_pooling_types) if classifier_pooling_types is not None else None + ) + self.tie_word_embeddings = False + self.attn_backend = attn_backend + + +### Rotary Embeddings +def rotate_half(x: torch.Tensor, interleaved: bool = False) -> torch.Tensor: + """Rotate the final axis of X by 90 degrees in each two-dimensional plane.""" + if interleaved: + paired = x.unflatten(-1, (-1, 2)) + return torch.stack((-paired[..., 1], paired[..., 0]), dim=-1).flatten(-2) + + # torch.chunk assigns an odd remainder to the first half. Express the same + # public behavior explicitly while keeping the ESMC path branch-free. + midpoint = (x.shape[-1] + 1) // 2 + return torch.cat((-x[..., midpoint:], x[..., :midpoint]), dim=-1) + + +def apply_rotary_emb_torch( + x: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + interleaved: bool = False, + _inplace: bool = False, +) -> torch.Tensor: + """Apply cached rotary angles to X while preserving any unrotated features.""" + del _inplace # Kept in the signature for checkpoint remote-code compatibility. + rotary_width = 2 * cos.shape[-1] + if rotary_width > x.shape[-1]: + raise AssertionError("rotary width exceeds the attention head dimension") + + token_count = x.shape[1] + cos_full = torch.cat((cos[:token_count], cos[:token_count]), dim=-1).unsqueeze(1) + sin_full = torch.cat((sin[:token_count], sin[:token_count]), dim=-1).unsqueeze(1) + x_rotary = x[..., :rotary_width] + y_rotary = x_rotary * cos_full + rotate_half(x_rotary, interleaved) * sin_full + if rotary_width == x.shape[-1]: + return y_rotary + return torch.cat((y_rotary, x[..., rotary_width:]), dim=-1) + + +class RotaryEmbedding(torch.nn.Module): + """Rotary position embeddings. + + Based on the paper "RoFormer: Enhanced Transformer with Rotary Position Embedding" + + Args: + dim: Dimension of the embedding + base: Base for computing angular frequencies + interleaved: Whether to use interleaved rotations + scale_base: Base for scaling + scaling_factor: Factor for scaling positions + pos_idx_in_fp32: Whether to compute position indices in fp32 + device: Computation device + """ + + def __init__( + self, + dim: int, + base: float = 10000.0, + interleaved: bool = False, + scale_base: float | None = None, + scaling_factor: float = 1.0, + pos_idx_in_fp32: bool = True, + device: torch.device | None = None, + ) -> None: + super().__init__() + self.dim, self.base = dim, float(base) + self.interleaved, self.scale_base = interleaved, scale_base + self.scaling_factor, self.pos_idx_in_fp32 = scaling_factor, pos_idx_in_fp32 + self.device = device + self._clear_cache() + self.reset_parameters() + + def _clear_cache(self) -> None: + self._seq_len_cached = 0 + self._cos_cached: torch.Tensor | None = None + self._sin_cached: torch.Tensor | None = None + self._cos_k_cached: torch.Tensor | None = None + self._sin_k_cached: torch.Tensor | None = None + + def reset_parameters(self, device: torch.device | str | None = None) -> None: + """Rebuild the non-persistent frequency buffers on ``device``.""" + if device is not None: + buffer_device = torch.device(device) + elif "inv_freq" in self._buffers and isinstance(self._buffers["inv_freq"], torch.Tensor): + buffer_device = self._buffers["inv_freq"].device + else: + buffer_device = self.device + inv_freq = self._compute_inv_freq(buffer_device) + self._clear_cache() + self.register_buffer("inv_freq", inv_freq, persistent=False) + arange = torch.arange(0, self.dim, 2, device=buffer_device, dtype=torch.float32) + scale = ( + (arange + 0.4 * self.dim) / (1.4 * self.dim) if self.scale_base is not None else None + ) + self.register_buffer("scale", scale) + + def _compute_inv_freq(self, device: torch.device | None = None) -> torch.Tensor: + """Compute inverse frequency bands on their execution device.""" + return 1 / ( + self.base + ** (torch.arange(0, self.dim, 2, device=device, dtype=torch.float32) / self.dim) + ) + + def _apply(self, fn, recurse: bool = True): + """Move the module, then regenerate device-specific RoPE frequencies.""" + if self.inv_freq.is_meta: + self.reset_parameters(device="cpu") + result = super()._apply(fn, recurse=recurse) + self.register_buffer( + "inv_freq", + self._compute_inv_freq(self.inv_freq.device), + persistent=False, + ) + self._clear_cache() + return result + + def _cache_is_current( + self, + token_count: int, + device: torch.device | None, + dtype: torch.dtype | None, + ) -> bool: + cached = self._cos_cached + return ( + cached is not None + and self._seq_len_cached >= token_count + and cached.device == device + and cached.dtype == dtype + and not (self.training and cached.is_inference()) + ) + + def _rotary_angles( + self, + token_count: int, + device: torch.device | None, + ) -> torch.Tensor: + position_dtype = torch.float32 if self.pos_idx_in_fp32 else self.inv_freq.dtype + positions = torch.arange(token_count, device=device, dtype=position_dtype) # (l,) + positions.div_(self.scaling_factor) + frequencies = ( + self.inv_freq.to(torch.float32) + if self.pos_idx_in_fp32 and self.inv_freq.dtype != torch.float32 + else self.inv_freq + ) + return torch.outer(positions, frequencies) # (l, d / 2) + + def _update_cos_sin_cache( + self, seqlen: int, device: torch.device | None = None, dtype: torch.dtype | None = None + ) -> None: + """Build angle tables when the requested cache identity has changed.""" + if self._cache_is_current(seqlen, device, dtype): + return + + self._seq_len_cached = seqlen + angles = self._rotary_angles(seqlen, device) # (l, d / 2) + cos_angles = torch.cos(angles) # (l, d / 2) + sin_angles = torch.sin(angles) # (l, d / 2) + if self.scale is None: + self._cos_cached = cos_angles.to(dtype) + self._sin_cached = sin_angles.to(dtype) + return + + centered_positions = ( + torch.arange(seqlen, dtype=self.scale.dtype, device=self.scale.device) - seqlen // 2 + ) / self.scale_base + scale = self.scale ** centered_positions.unsqueeze(-1) + self._cos_cached = (cos_angles * scale).to(dtype) + self._sin_cached = (sin_angles * scale).to(dtype) + self._cos_k_cached = (cos_angles / scale).to(dtype) + self._sin_k_cached = (sin_angles / scale).to(dtype) + + def forward(self, q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Apply rotary embeddings to queries and keys. + + Args: + q: Query tensor Q with shape (b, l, h, d). + k: Key tensor K with shape (b, l, h, d). + + Returns: + Tuple of rotated query and key tensors + """ + # The pinned Biohub Transformers oracle recomputes inverse frequencies + # on the execution device. CPU and CUDA differ by about one FP32 ULP in + # some bands, which is immaterial in BF16 but accumulates measurably in + # deep FP32 execution. + self._update_cos_sin_cache(q.shape[1], device=q.device, dtype=q.dtype) + if self._cos_cached is None or self._sin_cached is None: + raise RuntimeError( + "Rotary cache initialization did not produce cosine and sine values." + ) + if self.scale is not None: + raise AssertionError("Scaled rotary embeddings are unsupported for ESMC.") + + cos_angles = self._cos_cached + sin_angles = self._sin_cached + return ( + apply_rotary_emb_torch(q, cos_angles, sin_angles, self.interleaved, True), + apply_rotary_emb_torch(k, cos_angles, sin_angles, self.interleaved, True), + ) + + +def swiglu_correction_fn(expansion_ratio: float, d_model: int) -> int: + """Compute corrected dimension for SwiGLU.""" + return int(((expansion_ratio * d_model) + 255) // 256 * 256) + + +class SwiGLU(nn.Module): + """SwiGLU activation function.""" + + def __init__(self) -> None: + super().__init__() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x1, x2 = x.chunk(2, dim=-1) + return F.silu(x1) * x2 + + +def swiglu_ln_ffn(d_model: int, expansion_ratio: float) -> nn.Sequential: + """Create SwiGLU feedforward network with layer normalization.""" + return nn.Sequential( + nn.LayerNorm(d_model), + nn.Linear(d_model, swiglu_correction_fn(expansion_ratio, d_model) * 2, bias=False), + SwiGLU(), + nn.Linear(swiglu_correction_fn(expansion_ratio, d_model), d_model, bias=False), + ) + + +class MultiHeadAttention(nn.Module): + """Multi-head attention with rotary embeddings and configurable backend. + + Args: + d_model: Model dimension + n_heads: Number of attention heads + attn_backend: One of "eager", "sdpa", or "flex_attention". + """ + + def __init__( + self, + d_model: int, + n_heads: int, + attn_backend: str = "sdpa", + ) -> None: + super().__init__() + self.d_model = d_model + self.n_heads = n_heads + self.d_head = self.d_model // self.n_heads + self.scale = 1.0 / math.sqrt(self.d_head) + self.attn_backend = resolve_attention_backend(attn_backend) + self.layernorm_qkv = nn.Sequential( + nn.LayerNorm(d_model), nn.Linear(d_model, d_model * 3, bias=False) + ) + self.out_proj = nn.Linear(d_model, d_model, bias=False) + self.q_ln = nn.LayerNorm(d_model, bias=False) + self.k_ln = nn.LayerNorm(d_model, bias=False) + self.reshaper = partial(rearrange, pattern="b s (h d) -> b h s d", h=n_heads) + self.rotary = RotaryEmbedding(d_model // n_heads) + + def _apply_rotary(self, q: torch.Tensor, k: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + # q, k: (b, l, d) + q = q.unflatten(-1, (self.n_heads, self.d_head)) # (b, l, h, d_h) + k = k.unflatten(-1, (self.n_heads, self.d_head)) # (b, l, h, d_h) + q, k = self.rotary(q, k) + q = q.flatten(-2, -1) + k = k.flatten(-2, -1) + return q, k + + def forward( + self, + x: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + attention_mask_4d: torch.Tensor | None = None, + flex_block_mask: BlockMask | None = None, + output_attentions: bool = False, + output_s_max: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]: + # x: (b, l, d) + qkv = self.layernorm_qkv(x) # (b, l, 3 * d) + query_sequence, key_sequence, value_sequence = torch.chunk(qkv, 3, dim=-1) + query_sequence, key_sequence = ( + self.q_ln(query_sequence).to(query_sequence.dtype), + self.k_ln(key_sequence).to(query_sequence.dtype), + ) + query_sequence, key_sequence = self._apply_rotary(query_sequence, key_sequence) + query_heads, key_heads, value_heads = map( + self.reshaper, (query_sequence, key_sequence, value_sequence) + ) # each (b, h, l, d_h) + + attn_output, attn_weights, s_max = self._attn( + query_heads, + key_heads, + value_heads, + attention_mask_2d=attention_mask_2d, + attention_mask_4d=attention_mask_4d, + flex_block_mask=flex_block_mask, + output_attentions=output_attentions, + output_s_max=output_s_max, + ) + + output = self.out_proj(attn_output) + return output, attn_weights, s_max + + def _attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + attention_mask_4d: torch.Tensor | None = None, + flex_block_mask: BlockMask | None = None, + output_attentions: bool = False, + output_s_max: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]: + if output_attentions: + return self._manual_attn( + query_heads, key_heads, value_heads, attention_mask_4d, output_s_max + ) + + if self.attn_backend == AttentionBackend.EAGER: + attn_output, _, s_max = self._manual_attn( + query_heads, key_heads, value_heads, attention_mask_4d, output_s_max + ) + return attn_output, None, s_max + if self.attn_backend.is_flash: + attn_output, attn_weights = self._kernels_flash_attn( + query_heads, key_heads, value_heads, attention_mask_2d + ) + elif self.attn_backend == AttentionBackend.FLEX: + attn_output, attn_weights = self._flex_attn( + query_heads, + key_heads, + value_heads, + flex_block_mask, + attention_mask_2d, + ) + elif self.attn_backend == AttentionBackend.SDPA: + attn_output, attn_weights = self._sdpa_attn( + query_heads, key_heads, value_heads, attention_mask_4d + ) + else: + raise AssertionError(f"Unsupported resolved backend: {self.attn_backend}") + + s_max = self._compute_s_max(query_heads, key_heads) if output_s_max else None + return attn_output, attn_weights, s_max + + @torch.no_grad() + def _compute_s_max( + self, query_heads: torch.Tensor, key_heads: torch.Tensor + ) -> list[torch.Tensor]: + q_norm = torch.linalg.vector_norm(query_heads, dim=-1) # (b, h, l) + k_norm = torch.linalg.vector_norm(key_heads, dim=-1) # (b, h, l) + s_max_bound = (q_norm.max(dim=-1).values * k_norm.max(dim=-1).values).max( + dim=0 + ).values * self.scale + return [s_max_bound[h] for h in range(self.n_heads)] + + def _manual_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_4d: torch.Tensor | None = None, + output_s_max: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor, list[torch.Tensor] | None]: + # query_heads, key_heads, value_heads: (b, h, l, d_h) + attn_weights = ( + torch.matmul(query_heads, key_heads.transpose(-2, -1)) * self.scale + ) # (b, h, l, l) + if attention_mask_4d is not None: + attn_weights = attn_weights.masked_fill(attention_mask_4d.logical_not(), float("-inf")) + attn_weights = F.softmax(attn_weights, dim=-1) + context_heads = torch.matmul(attn_weights, value_heads) # (b, h, l, d_h) + attn_output = rearrange(context_heads, "b h s d -> b s (h d)") # (b, l, d) + s_max = self._compute_s_max(query_heads, key_heads) if output_s_max else None + return attn_output, attn_weights, s_max + + def _kernels_flash_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + query_tokens = query_heads.transpose(1, 2).contiguous() + key_tokens = key_heads.transpose(1, 2).contiguous() + value_tokens = value_heads.transpose(1, 2).contiguous() + attn_output = kernels_flash_attention_func( + query_states=query_tokens, + key_states=key_tokens, + value_states=value_tokens, + attention_mask_2d=attention_mask_2d, + causal=False, + implementation=self.attn_backend.value, + ) + return rearrange(attn_output, "b s h d -> b s (h d)"), None + + def _flex_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + flex_block_mask: BlockMask | None = None, + attention_mask_2d: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + if flex_attention is None: + raise RuntimeError("Flex attention is not available in this environment.") + fn = _get_flex_attention_fn( + device=query_heads.device, + dtype=query_heads.dtype, + shape=tuple(query_heads.shape), + mask_semantics="padding", + ) + context_heads = fn( + query_heads, + key_heads, + value_heads, + block_mask=flex_block_mask, + scale=self.scale, + kernel_options={"PRESCALE_QK": True, "BLOCK_N": 32}, + ) + return rearrange(context_heads, "b h s d -> b s (h d)"), None + + def _sdpa_attn( + self, + query_heads: torch.Tensor, + key_heads: torch.Tensor, + value_heads: torch.Tensor, + attention_mask_4d: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, None]: + context_heads = F.scaled_dot_product_attention( + query_heads, + key_heads, + value_heads, + attn_mask=attention_mask_4d, + scale=self.scale, + ) + return rearrange(context_heads, "b h s d -> b s (h d)"), None + + +def RegressionHead(d_model: int, output_dim: int, hidden_dim: int | None = None) -> nn.Module: + """Create a regression head with optional hidden dimension. + + Args: + d_model: Input dimension + output_dim: Output dimension + hidden_dim: Optional hidden dimension (defaults to d_model) + """ + hidden_dim = hidden_dim if hidden_dim is not None else d_model + return nn.Sequential( + nn.Linear(d_model, hidden_dim), + nn.GELU(), + nn.LayerNorm(hidden_dim), + nn.Linear(hidden_dim, output_dim), + ) + + +class UnifiedTransformerBlock(nn.Module): + """Transformer block with attention and feedforward layers.""" + + def __init__( + self, + d_model: int, + n_heads: int, + residue_scaling_factor: float = 1, + expansion_ratio: float = 8 / 3, + dropout: float = 0.0, + attn_backend: str = "sdpa", + ) -> None: + super().__init__() + self.attn = MultiHeadAttention(d_model=d_model, n_heads=n_heads, attn_backend=attn_backend) + self.ffn = swiglu_ln_ffn(d_model, expansion_ratio) + self.scaling_factor = residue_scaling_factor + self.dropout = nn.Dropout(dropout) + + def forward( + self, + x: torch.Tensor, + attention_mask_2d: torch.Tensor | None = None, + attention_mask_4d: torch.Tensor | None = None, + flex_block_mask: BlockMask | None = None, + output_attentions: bool = False, + output_s_max: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None, list[torch.Tensor] | None]: + attn_output, attn_weights, s_max = self.attn( + x, + attention_mask_2d=attention_mask_2d, + attention_mask_4d=attention_mask_4d, + flex_block_mask=flex_block_mask, + output_attentions=output_attentions, + output_s_max=output_s_max, + ) + x = x + self.dropout(attn_output) / self.scaling_factor + x = x + self.dropout(self.ffn(x)) / self.scaling_factor + return x, attn_weights, s_max + + +@dataclass +class TransformerOutput(ModelOutput): + """Output type for transformer encoder.""" + + last_hidden_state: torch.Tensor | None = None + hidden_states: tuple[torch.Tensor] | None = None + attentions: tuple[torch.Tensor] | None = None + s_max: tuple[list[torch.Tensor], ...] | None = None + sae_outputs: dict[str, torch.Tensor] | None = None + sae_hidden_states: dict[int, torch.Tensor] | None = None + + +@dataclass +class ESMplusplusOutput(MaskedLMOutput): + """Masked-LM output with FastPLMs fields after the HF contract.""" + + s_max: tuple[list[torch.Tensor], ...] | None = None + last_hidden_state: torch.Tensor | None = None + sae_outputs: dict[str, torch.Tensor] | None = None + + +@dataclass +class ESMplusplusSequenceClassifierOutput(SequenceClassifierOutput): + """Sequence-classification output with optional attention diagnostics.""" + + s_max: tuple[list[torch.Tensor], ...] | None = None + sae_outputs: dict[str, torch.Tensor] | None = None + + +@dataclass +class ESMplusplusTokenClassifierOutput(TokenClassifierOutput): + """Token-classification output with optional attention diagnostics.""" + + s_max: tuple[list[torch.Tensor], ...] | None = None + sae_outputs: dict[str, torch.Tensor] | None = None + + +class TransformerStack(nn.Module): + """Stack of transformer blocks.""" + + def __init__( + self, + d_model: int, + n_heads: int, + n_layers: int, + dropout: float = 0.0, + attn_backend: str = "sdpa", + ) -> None: + super().__init__() + self.attention_backend = resolve_attention_backend(attn_backend) + self.blocks = nn.ModuleList( + [ + UnifiedTransformerBlock( + d_model, + n_heads, + residue_scaling_factor=math.sqrt(n_layers / 36), + dropout=dropout, + attn_backend=attn_backend, + ) + for i in range(n_layers) + ] + ) + self.norm = nn.LayerNorm(d_model, bias=False) + self.gradient_checkpointing = False + + @property + def attn_backend(self) -> AttentionBackend: + return self.attention_backend + + @attn_backend.setter + def attn_backend(self, backend: str) -> None: + resolved = resolve_attention_backend(backend) + self.attention_backend = resolved + for block in self.blocks: + block.attn.attn_backend = resolved + + def forward( + self, + x: torch.Tensor, + attention_mask: torch.Tensor | None = None, + sequence_id: torch.Tensor | None = None, + output_hidden_states: bool | None = False, + output_attentions: bool | None = False, + output_s_max: bool | None = False, + esmfold2_hidden_states: bool = False, + sae_layers: tuple[int, ...] = (), + ) -> TransformerOutput: + # x: (b, l, d); attention_mask, sequence_id: (b, l) + hidden_states = () if output_hidden_states else None + attentions = () if output_attentions else None + full_s_max = () if output_s_max else None + sae_layer_set = set(sae_layers) + sae_hidden_states = {} if sae_layer_set else None + # Match the pinned Biohub Transformers contract: a supplied sequence_id + # is authoritative and must encode padding as -1. attention_mask is + # ignored in that mode rather than intersected with the chain mask. + attention_mask_2d, attention_mask_4d, flex_block_mask = ( + self._prepare_attention_masks( + attention_mask=attention_mask, + sequence_id=sequence_id, + batch_size=x.shape[0], + seq_len=x.shape[1], + device=x.device, + dtype=x.dtype, + output_attentions=bool(output_attentions), + ) + ) + + for layer_index, block in enumerate(self.blocks): + if output_hidden_states: + if hidden_states is None: + raise RuntimeError( + "Hidden-state collection was not initialized for an enabled request." + ) + # Biohub Transformers records the input to each block followed + # by the final normalized state. This gives n_layers + 1 states + # and, for ESMC-6B, the 81-state order consumed by ESMFold2. + hidden_states += (x,) + if sae_hidden_states is not None and layer_index in sae_layer_set: + sae_hidden_states[layer_index] = x + if self.gradient_checkpointing and self.training: + x, attn_weights, s_max = self._gradient_checkpointing_func( + block.__call__, + x=x, + attention_mask_2d=attention_mask_2d, + attention_mask_4d=attention_mask_4d, + flex_block_mask=flex_block_mask, + output_attentions=output_attentions, + output_s_max=output_s_max, + ) + else: + x, attn_weights, s_max = block( + x=x, + attention_mask_2d=attention_mask_2d, + attention_mask_4d=attention_mask_4d, + flex_block_mask=flex_block_mask, + output_attentions=output_attentions, + output_s_max=output_s_max, + ) + + if attentions is not None: + attentions += (attn_weights,) + if full_s_max is not None: + full_s_max += (s_max,) + + last_hidden_state = self.norm(x) + if output_hidden_states: + hidden_states += (last_hidden_state,) + final_layer_index = len(self.blocks) + if sae_hidden_states is not None and final_layer_index in sae_layer_set: + sae_hidden_states[final_layer_index] = last_hidden_state + + return TransformerOutput( + last_hidden_state=last_hidden_state, + hidden_states=hidden_states, + attentions=attentions, + s_max=full_s_max, + sae_hidden_states=sae_hidden_states, + ) + + @torch.compiler.disable + def _prepare_attention_masks( + self, + attention_mask: torch.Tensor | None, + sequence_id: torch.Tensor | None, + batch_size: int, + seq_len: int, + device: torch.device, + dtype: torch.dtype | None = None, + effective_backend: AttentionBackend | None = None, + output_attentions: bool = False, + ) -> tuple[torch.Tensor | None, torch.Tensor | None, BlockMask | None]: + mask_name = "sequence_id" if sequence_id is not None else "attention_mask" + mask_pattern = sequence_id if sequence_id is not None else attention_mask + if mask_pattern is None: + backend = resolve_attention_backend_for_call( + self.attention_backend, + output_attentions=output_attentions, + ) + return get_attention_mask( + effective_backend=backend, + batch_size=batch_size, + seq_len=seq_len, + device=device, + attention_mask=None, + dtype=dtype, + mask_semantics="padding", + ) + + expected_shape = (batch_size, seq_len) + if mask_pattern.ndim != 2 or tuple(mask_pattern.shape) != expected_shape: + raise ValueError( + f"{mask_name} must have shape {expected_shape}; " + f"received {tuple(mask_pattern.shape)}." + ) + if mask_pattern.device != device: + mask_pattern = mask_pattern.to(device=device) + if sequence_id is None: + mask_pattern = mask_pattern.to(dtype=torch.bool) + attention_mask_2d = ( + mask_pattern if mask_pattern.dtype == torch.bool else mask_pattern != -1 + ) + if not bool(attention_mask_2d.any(dim=1).all()): + raise ValueError("attention_mask must keep at least one valid key per batch row.") + + if mask_pattern.dtype == torch.bool: + # Biohub's boolean single-chain form groups biological positions + # together and padding positions together. Padding queries remain + # finite without allowing their states to enter residue attention. + attention_mask_4d = ( + mask_pattern[:, None, :, None] == mask_pattern[:, None, None, :] + ) + else: + attention_mask_4d = ( + mask_pattern.unsqueeze(-1) == mask_pattern.unsqueeze(-2) + ).unsqueeze(1) + backend = ( + resolve_attention_backend_for_call( + self.attention_backend, + output_attentions=output_attentions, + ) + if effective_backend is None + else resolve_attention_backend(effective_backend) + ) + + if backend.is_flash: + if mask_pattern.dtype != torch.bool: + raise ValueError( + "ESM++ FlashAttention only supports boolean sequence_id padding masks. " + "Use eager, sdpa, or flex_attention for chain-aware integer sequence_id " + "masks." + ) + return attention_mask_2d, attention_mask_4d, None + + if backend == AttentionBackend.FLEX: + if mask_pattern.dtype == torch.bool: + + def mask_mod(batch_idx, head_idx, q_idx, kv_idx): + del head_idx + return mask_pattern[batch_idx, q_idx] == mask_pattern[batch_idx, kv_idx] + + else: + + def mask_mod(batch_idx, head_idx, q_idx, kv_idx): + del head_idx + q_id = mask_pattern[batch_idx, q_idx] + kv_id = mask_pattern[batch_idx, kv_idx] + return q_id == kv_id + + flex_block_mask = _get_flex_block_mask( + mask_pattern=mask_pattern, + batch_size=batch_size, + query_length=seq_len, + key_value_length=seq_len, + device=device, + dtype=dtype, + mask_semantics=( + "boolean_sequence_id" + if mask_pattern.dtype == torch.bool + else "integer_sequence_id" + ), + mask_mod=mask_mod, + ) + return attention_mask_2d, attention_mask_4d, flex_block_mask + + return attention_mask_2d, attention_mask_4d, None + + +class PreTrainedESMplusplusModel(FastPLMsAttentionMixin, PreTrainedModel): + """ + init weights for ESM++ models + """ + + config_class = ESMplusplusConfig + base_model_prefix = "esm++" + supports_gradient_checkpointing = True + all_tied_weights_keys: ClassVar[dict[str, str]] = {} + _supports_flash_attn = True + _supports_flash_attn_2 = True + _supports_flash_attn_3 = True + _fastplms_attention_implementations = ( + "eager", + "sdpa", + "flex_attention", + "flash_attention_2", + "flash_attention_3", + ) + + def __init__(self, config: ESMplusplusConfig, *args: object, **kwargs: object) -> None: + super().__init__(config, *args, **kwargs) + self._sae_models = nn.ModuleDict() + self._esmc_fp8 = False + self._esmc_fp8_module_paths: tuple[str, ...] = () + self._esmc_precision_status = ESMplusplusFP8Status( + enabled=False, + reason="FP8 has not been enabled; canonical checkpoint precision is unchanged.", + device="cpu", + transformer_engine_version=_transformer_engine_version(), + converted_projections=0, + ) + + @property + def esmc_precision_status(self) -> ESMplusplusFP8Status: + """Return the explicit FP8 conversion status for this ESM++ instance.""" + + return self._esmc_precision_status + + def enable_fp8(self) -> ESMplusplusFP8Status: + """Enable the strict inference-only Transformer Engine FP8 path. + + Canonical parameters must already be BF16 on one supported CUDA device. + Exactly one attention output projection per transformer block is replaced; + all other operations and all SAE weights remain BF16. + """ + + if self.config.hidden_size != 2560 or self.config.num_hidden_layers != 80: + raise ValueError("FP8 is supported only for the ESMC-6B backbone.") + if self._esmc_fp8: + return self._esmc_precision_status + if self.training: + raise RuntimeError("ESM++ FP8 is inference-only; call eval() before enable_fp8().") + parameter_devices = {parameter.device for parameter in self.parameters()} + if len(parameter_devices) != 1: + raise RuntimeError( + "ESM++ FP8 requires every parameter on one CUDA device; sharded device maps " + f"are unsupported, found {sorted(map(str, parameter_devices))}." + ) + device = next(iter(parameter_devices), self.device) + available, reason = _te_fp8_capability(device) + if not available: + raise RuntimeError(f"ESM++ FP8 is unavailable: {reason}") + non_bf16 = [ + name + for name, parameter in self.named_parameters() + if parameter.is_floating_point() and parameter.dtype != torch.bfloat16 + ] + if non_bf16: + examples = ", ".join(non_bf16[:3]) + raise RuntimeError( + "ESM++ FP8 requires canonical BF16 parameters before conversion; " + f"found {len(non_bf16)} non-BF16 parameters (for example: {examples})." + ) + paths = _convert_esmc_attention_outputs_to_te( + self, + expected_projections=self.config.num_hidden_layers, + ) + self._esmc_fp8 = True + self._esmc_fp8_module_paths = paths + self._esmc_precision_status = ESMplusplusFP8Status( + enabled=True, + reason=( + f"{reason} Converted {len(paths)} attention output projections; " + "canonical checkpoint and SAE weights remain BF16." + ), + device=str(device), + transformer_engine_version=_transformer_engine_version(), + converted_projections=len(paths), + ) + return self._esmc_precision_status + + def load_sae_models( + self, + repository: str | os.PathLike[str], + layers: Sequence[int], + *, + revision: str | None = None, + cache_dir: str | os.PathLike[str] | None = None, + token: str | bool | None = None, + local_files_only: bool = False, + dtype: torch.dtype | None = None, + ) -> dict[int, ESMplusplusSAELayer]: + """Load hidden-state SAE layers from a Hub repository or local directory, then attach them. + + The layers land on this model's device and, unless ``dtype`` says otherwise, in this + model's parameter dtype, so they consume its hidden states without a dtype mismatch. + """ + + sae_layers = load_esmc_sae_layers( + repository, + layers, + revision=revision, + cache_dir=cache_dir, + token=token, + local_files_only=local_files_only, + device=self.device, + dtype=self.dtype if dtype is None else dtype, + ) + self.add_sae_models(list(sae_layers.values())) + return sae_layers + + def add_sae_models(self, sae_models: list[nn.Module]) -> None: + """Attach hidden-state SAE layers to this ESM++ model. + + Accepts layers from ``load_sae_models`` and official Biohub ``ESMCSAEModel.layers`` + entries, which share one attachment contract. + """ + + for sae_model in sae_models: + if not isinstance(sae_model, nn.Module): + raise TypeError( + "Each SAE must be an nn.Module exposing the hidden-state SAE contract, such " + "as a load_sae_models layer or an official Biohub ESMCSAEModel.layers entry." + ) + layer = getattr(sae_model, "layer", None) + if isinstance(layer, bool) or not isinstance(layer, int): + raise TypeError("Each SAE layer must expose an integer .layer attribute.") + if not 0 <= layer <= self.config.num_hidden_layers: + raise ValueError( + f"SAE target layer {layer} is outside the ESM++ hidden-state range " + f"0..{self.config.num_hidden_layers}." + ) + params = getattr(sae_model, "params", None) + d_model = getattr(params, "d_model", None) + if d_model != self.config.hidden_size: + raise ValueError( + f"SAE layer {layer} expects d_model={d_model!r}, but this ESM++ " + f"checkpoint has hidden_size={self.config.hidden_size}." + ) + if not callable(getattr(sae_model, "get_sae_output", None)): + raise TypeError("Each SAE layer must expose get_sae_output(layer_states, token_mask).") + for name in ("idf", "max"): + if not isinstance(getattr(sae_model, name, None), torch.Tensor): + raise TypeError(f"Each SAE layer must expose a tensor {name!r} buffer.") + key = f"layer{layer}" + if key in self._sae_models: + raise ValueError( + f"An SAE is already registered at {key!r}; only one SAE per layer " + "can be active." + ) + self._sae_models[key] = sae_model + + def _prepare_sae_forward( + self, + *, + compute_sae: bool, + input_ids: torch.Tensor | None, + attention_mask: torch.Tensor | None, + sequence_id: torch.Tensor | None, + ) -> tuple[tuple[int, ...], torch.Tensor | None]: + if not compute_sae or not self._sae_models: + return (), None + if input_ids is None and attention_mask is None and sequence_id is None: + # Embedding inputs carry no token identities, so the caller must supply the mask. + # Rejecting masked tokens is then the caller's precondition rather than ours. + raise ValueError( + "SAE computation from inputs_embeds requires an explicit attention_mask or " + "sequence_id, because the token mask cannot be recovered from embeddings." + ) + if input_ids is not None and torch.any(input_ids == self.config.mask_token_id): + raise ValueError("SAE inputs must not contain mask tokens; SAEs were trained unmasked.") + if sequence_id is not None: + token_mask = sequence_id >= 0 + elif attention_mask is not None: + token_mask = attention_mask.to(dtype=torch.bool) + else: + token_mask = input_ids != self.config.pad_token_id + layers = tuple(sorted(int(name.removeprefix("layer")) for name in self._sae_models)) + return layers, token_mask + + def _get_sae_outputs( + self, + hidden_states: dict[int, torch.Tensor] | None, + token_mask: torch.Tensor | None, + *, + normalize_sae: bool, + differentiable_sae: bool = False, + ) -> dict[str, torch.Tensor] | None: + """Encode collected hidden states with every attached SAE. + + The default detaches and sparsifies, which is right for interpretation and matches the + official implementation. ``differentiable_sae`` instead keeps the result attached to the + graph and dense, so a gradient-based sequence designer can optimize an objective built on + SAE features. The arithmetic is identical either way; only the tape and the layout differ. + """ + if not self._sae_models: + return None + if hidden_states is None or token_mask is None: + raise RuntimeError("SAE hidden-state collection was not initialized.") + outputs: dict[str, torch.Tensor] = {} + for key, sae_model in self._sae_models.items(): + layer = int(key.removeprefix("layer")) + if layer not in hidden_states: + raise RuntimeError(f"ESM++ did not collect the requested SAE layer {layer}.") + layer_states = hidden_states[layer] + sae_output = sae_model.get_sae_output( + layer_states if differentiable_sae else layer_states.clone(), token_mask + ) + features = getattr(sae_output, "feature_magnitudes", None) + if not isinstance(features, torch.Tensor): + raise TypeError("SAE get_sae_output must return tensor feature_magnitudes.") + if not differentiable_sae: + features = features.detach() + if normalize_sae: + features = (features / sae_model.max) * sae_model.idf + outputs[key] = features if differentiable_sae else features.to_sparse() + return outputs + + def _pad_fp8_inputs( + self, + input_ids: torch.Tensor | None, + attention_mask: torch.Tensor | None, + sequence_id: torch.Tensor | None, + inputs_embeds: torch.Tensor | None, + ) -> tuple[ + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + int | None, + ]: + if not self._esmc_fp8: + return input_ids, attention_mask, sequence_id, inputs_embeds, None + source = input_ids if input_ids is not None else inputs_embeds + if source is None: + return input_ids, attention_mask, sequence_id, inputs_embeds, None + sequence_length = source.shape[1] + padded_length = ( + (sequence_length + _ESMC_FP8_ALIGNMENT - 1) // _ESMC_FP8_ALIGNMENT + ) * _ESMC_FP8_ALIGNMENT + padding = padded_length - sequence_length + if padding == 0: + return input_ids, attention_mask, sequence_id, inputs_embeds, None + if input_ids is not None: + input_ids = F.pad(input_ids, (0, padding), value=self.config.pad_token_id) + if inputs_embeds is not None: + inputs_embeds = F.pad(inputs_embeds, (0, 0, 0, padding), value=0.0) + if sequence_id is not None: + sequence_id = F.pad(sequence_id.to(dtype=torch.long), (0, padding), value=-1) + else: + if attention_mask is None: + attention_mask = ( + source != self.config.pad_token_id + if input_ids is not None + else torch.ones( + source.shape[:2], + dtype=torch.bool, + device=source.device, + ) + ) + attention_mask = F.pad(attention_mask, (0, padding), value=0) + return input_ids, attention_mask, sequence_id, inputs_embeds, sequence_length + + @staticmethod + def _trim_transformer_output( + output: TransformerOutput, + sequence_length: int | None, + ) -> TransformerOutput: + if sequence_length is None: + return output + hidden_states = ( + tuple(state[:, :sequence_length] for state in output.hidden_states) + if output.hidden_states is not None + else None + ) + attentions = ( + tuple( + attention[..., :sequence_length, :sequence_length] + for attention in output.attentions + ) + if output.attentions is not None + else None + ) + return TransformerOutput( + last_hidden_state=output.last_hidden_state[:, :sequence_length], + hidden_states=hidden_states, + attentions=attentions, + s_max=output.s_max, + sae_hidden_states=output.sae_hidden_states, + ) + + @property + def tokenizer(self) -> EsmSequenceTokenizer: + """Construct the sequence tokenizer only when a raw-sequence API needs it.""" + + tokenizer = self.__dict__.get("_fastplms_tokenizer") + if tokenizer is None: + tokenizer = EsmSequenceTokenizer() + self.__dict__["_fastplms_tokenizer"] = tokenizer + return tokenizer + + @tokenizer.setter + def tokenizer(self, value: EsmSequenceTokenizer | None) -> None: + self.__dict__["_fastplms_tokenizer"] = value + + def _init_weights(self, module): + """Initialize the weights""" + # HF from_pretrained marks loaded parameters with `_is_hf_initialized`. + # Skip this module if any local parameter is already marked as loaded. + for parameter in module.parameters(recurse=False): + if parameter.__dict__.get("_is_hf_initialized"): + return + + if isinstance(module, nn.Linear): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.bias is not None: + nn.init.zeros_(module.bias) + elif isinstance(module, nn.Embedding): + nn.init.normal_(module.weight, mean=0.0, std=self.config.initializer_range) + if module.padding_idx is not None: + with torch.no_grad(): + module.weight[module.padding_idx].zero_() + elif isinstance(module, nn.LayerNorm): + if module.bias is not None: + nn.init.zeros_(module.bias) + nn.init.ones_(module.weight) + + @property + def attn_backend(self) -> str: + return self.config.attn_backend + + @attn_backend.setter + def attn_backend(self, backend: str) -> None: + if backend not in self._fastplms_attention_implementations: + raise ValueError( + f"{type(self).__name__} does not support {backend!r}; expected one of " + f"{self._fastplms_attention_implementations}." + ) + self.set_attn_implementation(backend) + + def _reset_rotary_embeddings(self): + """Refresh non-persistent rotary buffers after checkpoint loading.""" + for module in self.modules(): + if isinstance(module, RotaryEmbedding): + module.reset_parameters() + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs): + output_loading_info = ( + bool(kwargs["output_loading_info"]) if "output_loading_info" in kwargs else False + ) + loaded = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) + if output_loading_info: + model, loading_info = loaded + model._reset_rotary_embeddings() + return model, loading_info + loaded._reset_rotary_embeddings() + return loaded + + +### ESM++ Models +class ESMplusplusModel(PreTrainedESMplusplusModel, EmbeddingMixin): + """ + ESM++ transformer backbone. + + Official ESM++ checkpoints contain the sequence head even when loaded through + ``AutoModel``. Keep that module in the base class so the checkpoint has one + exact state-dict contract across ``AutoModel`` and ``AutoModelForMaskedLM``; + the base forward path intentionally does not compute or return logits. + """ + + config_class = ESMplusplusConfig + + def __init__(self, config: ESMplusplusConfig, **kwargs) -> None: + PreTrainedESMplusplusModel.__init__(self, config, **kwargs) + self.config = config + self.vocab_size = config.vocab_size + self.embed = nn.Embedding(self.vocab_size, config.hidden_size) + self.transformer = TransformerStack( + d_model=config.hidden_size, + n_heads=config.num_attention_heads, + n_layers=config.num_hidden_layers, + dropout=config.dropout, + attn_backend=config.attn_backend, + ) + self.sequence_head = RegressionHead(config.hidden_size, self.vocab_size) + self.init_weights() + + def get_input_embeddings(self): + return self.embed + + def set_input_embeddings(self, value): + self.embed = value + + def get_output_embeddings(self): + return self.sequence_head[-1] + + def set_output_embeddings(self, new_embeddings): + self.sequence_head[-1] = new_embeddings + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + if attention_mask is None: + attention_mask = input_ids.ne(self.config.pad_token_id) + x = self.embed(input_ids) + output_hidden_states = store_all_hidden_states or hidden_state_index != -1 + output = self.transformer( + x=x, + attention_mask=attention_mask, + output_hidden_states=output_hidden_states, + output_attentions=False, + ) + return select_hidden_state_embeddings( + output.last_hidden_state, + output.hidden_states, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + sequence_id: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + esmfold2_hidden_states: bool = False, + return_dict: bool | None = None, + compute_sae: bool = True, + normalize_sae: bool = False, + differentiable_sae: bool = False, + ) -> TransformerOutput | tuple[torch.Tensor, ...]: + """Run ESMC inference with the pinned Biohub mask precedence. + + ``sequence_id`` is authoritative when supplied: non-negative integers + identify chains and ``-1`` identifies padding. In that mode + ``attention_mask`` is ignored, matching the official implementation. + Without ``sequence_id``, ``attention_mask`` is the ordinary padding + mask and defaults to ``input_ids != pad_token_id``. + """ + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + output_attentions = ( + output_attentions if output_attentions is not None else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + + input_ids, attention_mask, sequence_id, inputs_embeds, original_length = ( + self._pad_fp8_inputs(input_ids, attention_mask, sequence_id, inputs_embeds) + ) + if attention_mask is None and sequence_id is None and input_ids is not None: + attention_mask = input_ids.ne(self.config.pad_token_id) + sae_layers, sae_token_mask = self._prepare_sae_forward( + compute_sae=compute_sae, + input_ids=input_ids, + attention_mask=attention_mask, + sequence_id=sequence_id, + ) + + x = self.embed(input_ids) if inputs_embeds is None else inputs_embeds + + with _esmplusplus_fp8_context(self._esmc_fp8, self.device): + transformer_output = self.transformer( + x=x, + attention_mask=attention_mask, + sequence_id=sequence_id, + output_hidden_states=output_hidden_states, + output_attentions=output_attentions, + output_s_max=output_s_max, + esmfold2_hidden_states=esmfold2_hidden_states, + sae_layers=sae_layers, + ) + sae_outputs = ( + self._get_sae_outputs( + transformer_output.sae_hidden_states, + sae_token_mask, + normalize_sae=normalize_sae, + differentiable_sae=differentiable_sae, + ) + if sae_layers + else None + ) + transformer_output = self._trim_transformer_output(transformer_output, original_length) + result = TransformerOutput( + last_hidden_state=transformer_output.last_hidden_state, + hidden_states=transformer_output.hidden_states, + attentions=transformer_output.attentions, + s_max=transformer_output.s_max, + sae_outputs=sae_outputs, + ) + return result if return_dict else result.to_tuple() + + +class ESMplusplusForMaskedLM( + FastPLMTestTimeTrainingMixin, PreTrainedESMplusplusModel, EmbeddingMixin +): + """ + ESM++ model for masked language modeling. + Implements the base ESM++ architecture with a masked language modeling head. + """ + + config_class = ESMplusplusConfig + + def __init__(self, config: ESMplusplusConfig, **kwargs) -> None: + PreTrainedESMplusplusModel.__init__(self, config, **kwargs) + self.config = config + self.vocab_size = config.vocab_size + self.embed = nn.Embedding(self.vocab_size, config.hidden_size) + self.transformer = TransformerStack( + d_model=config.hidden_size, + n_heads=config.num_attention_heads, + n_layers=config.num_hidden_layers, + dropout=config.dropout, + attn_backend=config.attn_backend, + ) + self.sequence_head = RegressionHead(config.hidden_size, self.vocab_size) + self.ce_loss = nn.CrossEntropyLoss() + self.init_weights() + self.init_ttt({"lora_target_replace_module": "MultiHeadAttention"}) + + def get_input_embeddings(self): + return self.embed + + def set_input_embeddings(self, value): + self.embed = value + + def get_output_embeddings(self): + return self.sequence_head[-1] + + def set_output_embeddings(self, new_embeddings): + self.sequence_head[-1] = new_embeddings + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + if attention_mask is None: + attention_mask = input_ids.ne(self.config.pad_token_id) + x = self.embed(input_ids) + output_hidden_states = store_all_hidden_states or hidden_state_index != -1 + output = self.transformer( + x=x, + attention_mask=attention_mask, + output_hidden_states=output_hidden_states, + output_attentions=False, + ) + return select_hidden_state_embeddings( + output.last_hidden_state, + output.hidden_states, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def _ttt_get_trainable_modules(self) -> list[nn.Module]: + return [self.transformer] + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + sequence_id: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + esmfold2_hidden_states: bool = False, + return_dict: bool | None = None, + compute_logits: bool = True, + compute_sae: bool = True, + normalize_sae: bool = False, + differentiable_sae: bool = False, + ) -> ESMplusplusOutput | tuple[torch.Tensor, ...]: + if input_ids is None and inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + if input_ids is not None and inputs_embeds is not None: + raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time") + if labels is not None and not compute_logits: + raise ValueError("labels require compute_logits=True.") + output_attentions = ( + output_attentions if output_attentions is not None else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + input_ids, attention_mask, sequence_id, inputs_embeds, original_length = ( + self._pad_fp8_inputs(input_ids, attention_mask, sequence_id, inputs_embeds) + ) + if attention_mask is None and sequence_id is None and input_ids is not None: + attention_mask = input_ids.ne(self.config.pad_token_id) + sae_layers, sae_token_mask = self._prepare_sae_forward( + compute_sae=compute_sae, + input_ids=input_ids, + attention_mask=attention_mask, + sequence_id=sequence_id, + ) + + x = self.embed(input_ids) if inputs_embeds is None else inputs_embeds + + with _esmplusplus_fp8_context(self._esmc_fp8, self.device): + output = self.transformer( + x=x, + attention_mask=attention_mask, + sequence_id=sequence_id, + output_hidden_states=output_hidden_states, + output_attentions=output_attentions, + output_s_max=output_s_max, + esmfold2_hidden_states=esmfold2_hidden_states, + sae_layers=sae_layers, + ) + sae_outputs = ( + self._get_sae_outputs( + output.sae_hidden_states, + sae_token_mask, + normalize_sae=normalize_sae, + differentiable_sae=differentiable_sae, + ) + if sae_layers + else None + ) + output = self._trim_transformer_output(output, original_length) + + last_hidden_state = output.last_hidden_state + logits = self.sequence_head(last_hidden_state) if compute_logits else None + loss = None + if labels is not None: + if logits is None: + raise ValueError("labels require compute_logits=True.") + labels = labels.to(logits.device) + loss = self.ce_loss(logits.view(-1, self.vocab_size), labels.view(-1)) + + result = ESMplusplusOutput( + loss=loss, + logits=logits, + hidden_states=output.hidden_states, + attentions=output.attentions, + s_max=output.s_max, + last_hidden_state=last_hidden_state, + sae_outputs=sae_outputs, + ) + return result if return_dict else result.to_tuple() + + +class ESMplusplusForSequenceClassification(ESMplusplusForMaskedLM, EmbeddingMixin): + """ + ESM++ model for sequence classification. + Extends the base ESM++ model with a classification head. + """ + + def __init__(self, config: ESMplusplusConfig, **kwargs) -> None: + pooling_types = kwargs.pop("pooling_types", None) + if pooling_types is None: + pooling_types = config.classifier_pooling_types or ["mean", "var"] + elif not isinstance(pooling_types, list): + raise TypeError("pooling_types must be a non-empty list of strings.") + elif not pooling_types: + raise ValueError("pooling_types must contain at least one pooling operation.") + elif not all(isinstance(pooling_type, str) for pooling_type in pooling_types): + raise TypeError("pooling_types must be a non-empty list of strings.") + if "parti" in pooling_types: + raise ValueError( + "pooling_types cannot contain 'parti' for sequence classification " + "because the classifier does not expose layer attentions to its pooler." + ) + config.classifier_pooling_types = list(pooling_types) + + ESMplusplusForMaskedLM.__init__(self, config, **kwargs) + self.config = config + self.num_labels = config.num_labels + self.classifier = RegressionHead( + config.hidden_size * len(pooling_types), + config.num_labels, + config.hidden_size * 4, + ) + # Large intermediate projections help with sequence classification tasks (*4) + self.mse = nn.MSELoss() + self.ce = nn.CrossEntropyLoss() + self.bce = nn.BCEWithLogitsLoss() + self.pooler = Pooler(pooling_types) + self.init_weights() + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + x = self.embed(input_ids) + output_hidden_states = store_all_hidden_states or hidden_state_index != -1 + output = self.transformer( + x=x, + attention_mask=attention_mask, + output_hidden_states=output_hidden_states, + output_attentions=False, + ) + return select_hidden_state_embeddings( + output.last_hidden_state, + output.hidden_states, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + sequence_id: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + return_dict: bool | None = None, + compute_sae: bool = True, + normalize_sae: bool = False, + ) -> ESMplusplusSequenceClassifierOutput | tuple[torch.Tensor, ...]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + pooling_mask = attention_mask + if pooling_mask is None: + if sequence_id is not None: + pooling_mask = ( + sequence_id if sequence_id.dtype == torch.bool else sequence_id.ne(-1) + ) + elif input_ids is not None: + pooling_mask = input_ids.ne(self.config.pad_token_id) + else: + if inputs_embeds is None: + raise ValueError("You have to specify either input_ids or inputs_embeds") + pooling_mask = torch.ones( + inputs_embeds.shape[:2], + dtype=torch.bool, + device=inputs_embeds.device, + ) + + output = super().forward( + input_ids=input_ids, + attention_mask=attention_mask, + sequence_id=sequence_id, + inputs_embeds=inputs_embeds, + labels=None, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_s_max=output_s_max, + return_dict=True, + compute_logits=False, + compute_sae=compute_sae, + normalize_sae=normalize_sae, + ) + + last_hidden_state = output.last_hidden_state + features = self.pooler(last_hidden_state, pooling_mask) + logits = self.classifier(features) + + loss = None + if labels is not None: + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and ( + labels.dtype == torch.long or labels.dtype == torch.int + ): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + if self.num_labels == 1: + loss = self.mse(logits.flatten(), labels.flatten()) + else: + loss = self.mse(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss = self.ce(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss = self.bce(logits, labels) + + result = ESMplusplusSequenceClassifierOutput( + loss=loss, + logits=logits, + hidden_states=output.hidden_states, + attentions=output.attentions, + s_max=output.s_max, + sae_outputs=output.sae_outputs, + ) + return result if return_dict else result.to_tuple() + + +class ESMplusplusForTokenClassification(ESMplusplusForMaskedLM, EmbeddingMixin): + """ + ESM++ model for token classification. + Extends the base ESM++ model with a token classification head. + """ + + def __init__(self, config: ESMplusplusConfig, **kwargs) -> None: + ESMplusplusForMaskedLM.__init__(self, config, **kwargs) + self.config = config + self.num_labels = config.num_labels + self.classifier = RegressionHead( + config.hidden_size, config.num_labels, config.hidden_size * 4 + ) + # Large intermediate projections help with sequence classification tasks (*4) + self.loss_fct = nn.CrossEntropyLoss() + self.init_weights() + + def _embed( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor | None = None, + hidden_state_index: int = -1, + store_all_hidden_states: bool = False, + ) -> torch.Tensor: + x = self.embed(input_ids) + output_hidden_states = store_all_hidden_states or hidden_state_index != -1 + output = self.transformer( + x, + attention_mask, + output_hidden_states=output_hidden_states, + output_attentions=False, + ) + return select_hidden_state_embeddings( + output.last_hidden_state, + output.hidden_states, + hidden_state_index=hidden_state_index, + store_all_hidden_states=store_all_hidden_states, + ) + + def forward( + self, + input_ids: torch.Tensor | None = None, + attention_mask: torch.Tensor | None = None, + sequence_id: torch.Tensor | None = None, + inputs_embeds: torch.Tensor | None = None, + labels: torch.Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + output_s_max: bool | None = False, + return_dict: bool | None = None, + compute_sae: bool = True, + normalize_sae: bool = False, + ) -> ESMplusplusTokenClassifierOutput | tuple[torch.Tensor, ...]: + return_dict = return_dict if return_dict is not None else self.config.use_return_dict + output = super().forward( + input_ids=input_ids, + attention_mask=attention_mask, + sequence_id=sequence_id, + inputs_embeds=inputs_embeds, + labels=None, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + output_s_max=output_s_max, + return_dict=True, + compute_logits=False, + compute_sae=compute_sae, + normalize_sae=normalize_sae, + ) + + last_hidden_state = output.last_hidden_state + logits = self.classifier(last_hidden_state) + loss = None + if labels is not None: + labels = labels.to(logits.device) + loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + + result = ESMplusplusTokenClassifierOutput( + loss=loss, + logits=logits, + hidden_states=output.hidden_states, + attentions=output.attentions, + s_max=output.s_max, + sae_outputs=output.sae_outputs, + ) + return result if return_dict else result.to_tuple() + + +### Tokenization +SEQUENCE_VOCAB = [ + "", + "", + "", + "", + "L", + "A", + "G", + "V", + "S", + "E", + "R", + "T", + "I", + "D", + "P", + "K", + "Q", + "N", + "F", + "Y", + "M", + "H", + "W", + "C", + "X", + "B", + "U", + "Z", + "O", + ".", + "-", + "|", + "", +] + + +def _build_sequence_tokenizer_backend( + *, + unk_token: str, + cls_token: str, + pad_token: str, + mask_token: str, + eos_token: str, + chain_break_token: str, +) -> Tokenizer: + """Build the fixed ESMC character vocabulary and boundary-token policy.""" + vocabulary = dict(zip(SEQUENCE_VOCAB, range(len(SEQUENCE_VOCAB)), strict=True)) + backend = Tokenizer(BPE(vocabulary, merges=[], unk_token=unk_token)) + backend.add_special_tokens([cls_token, pad_token, mask_token, eos_token, chain_break_token]) + backend.post_processor = TemplateProcessing( + single=" $A ", + pair=":0 $A:0 :0 $B:1 :1", + special_tokens=[ + ("", backend.token_to_id("")), + ("", backend.token_to_id("")), + ], + ) + return backend + + +class EsmSequenceTokenizer(PreTrainedTokenizerFast): + model_input_names: ClassVar[list[str]] = ["input_ids", "attention_mask"] + + def __init__( + self, + unk_token="", + cls_token="", + pad_token="", + mask_token="", + eos_token="", + chain_break_token="|", + **kwargs, + ): + backend = _build_sequence_tokenizer_backend( + unk_token=unk_token, + cls_token=cls_token, + pad_token=pad_token, + mask_token=mask_token, + eos_token=eos_token, + chain_break_token=chain_break_token, + ) + self.cb_token = chain_break_token + super().__init__( + tokenizer_object=backend, + unk_token=unk_token, + cls_token=cls_token, + pad_token=pad_token, + mask_token=mask_token, + eos_token=eos_token, + additional_special_tokens=[chain_break_token], + **kwargs, + ) + + # ESMC does not use BOS, so expose the sequence-start token through the HF BOS fields. + @property + def bos_token(self): + return self.cls_token + + @property + def bos_token_id(self): + return self.cls_token_id + + @property + def chain_break_token(self): + return self.cb_token + + @property + def chain_break_token_id(self): + return self.convert_tokens_to_ids(self.chain_break_token) + + @property + def all_token_ids(self): + return list(range(self.vocab_size)) + + @property + def special_token_ids(self): + return self.all_special_ids diff --git a/fastplms/models/esm_plusplus/modeling_esm_plusplus_sae.py b/fastplms/models/esm_plusplus/modeling_esm_plusplus_sae.py new file mode 100644 index 0000000000000000000000000000000000000000..a74514f3c9e8d8974cd470c4a3d9700033dc087f --- /dev/null +++ b/fastplms/models/esm_plusplus/modeling_esm_plusplus_sae.py @@ -0,0 +1,266 @@ +"""Hidden-state sparse autoencoders for ESM++ (ESMC) models. + +FastPLMs implements the published Biohub hidden-state SAE contract directly, so attaching an SAE +needs only PyTorch, Transformers, and the checkpoint itself. Biohub still owns the SAE weights: +this module reads their published repository layout, one shared ``config.json`` plus one +``layer_{index}.safetensors`` shard per backbone layer, and never redistributes those tensors. + +A layer produced here satisfies the same attachment contract as an official Biohub +``ESMCSAEModel.layers`` entry, so ``PreTrainedESMplusplusModel.add_sae_models`` accepts either. +""" + +from __future__ import annotations + +import json +import os +import torch +import torch.nn as nn +import torch.nn.functional as F + +from collections.abc import Sequence +from dataclasses import dataclass +from pathlib import Path +from huggingface_hub import hf_hub_download +from safetensors.torch import load_file +from torch import Tensor + + +SAE_CONFIG_FILE = "config.json" +SAE_REQUIRED_CONFIG_FIELDS = ("d_model", "codebook_dim", "k") +# Trained per-feature statistics. Shards that never trained them omit them, and the ones default +# makes the (features / max) * idf normalization an identity. +SAE_OPTIONAL_STATE_NAMES = ("idf", "max") +STANDARDIZATION_EPSILON = 1e-5 + + +@dataclass(frozen=True, slots=True) +class ESMplusplusSAEParams: + """Shape contract of one hidden-state SAE. + + ``layer`` indexes the ESM++ hidden state the SAE reads, where ``0`` is the embedding output and + ``num_hidden_layers`` is the final normalized state. + """ + + d_model: int + codebook_dim: int + k: int + layer: int + + +@dataclass(slots=True) +class ESMplusplusSAEOutput: + """Sparse features, and the optional reconstruction error, for one batch of residues.""" + + feature_magnitudes: Tensor + reconstruction_loss: Tensor | None = None + + +def _standardize_residue_states(residue_states: Tensor) -> Tensor: + """Center and scale each residue vector the way the SAEs were trained.""" + + # residue_states: (n, d) + centered = residue_states - residue_states.mean(dim=-1, keepdim=True) # (n, d) + return centered / (centered.std(dim=-1, keepdim=True) + STANDARDIZATION_EPSILON) # (n, d) + + +class ESMplusplusSAELayer(nn.Module): + """Top-k sparse autoencoder over one ESM++ hidden state. + + Encoding standardizes each residue vector, projects it into a wide codebook, and keeps only the + ``k`` largest activations. The decoder is used only to measure reconstruction error, which is + why it is opt-in: interpretation and gradient-based design read ``feature_magnitudes`` alone. + """ + + idf: Tensor + max: Tensor + + def __init__(self, params: ESMplusplusSAEParams) -> None: + super().__init__() + self.params = params + self.W_enc = nn.Parameter(torch.empty(params.d_model, params.codebook_dim)) # (d, c) + self.W_dec = nn.Parameter(torch.empty(params.codebook_dim, params.d_model)) # (c, d) + self.b_dec = nn.Parameter(torch.zeros(params.d_model)) # (d,) + self.register_buffer("idf", torch.ones(params.codebook_dim)) # (c,) + self.register_buffer("max", torch.ones(params.codebook_dim)) # (c,) + + @property + def layer(self) -> int: + """ESM++ hidden-state index this SAE was trained against.""" + + return self.params.layer + + def forward( + self, + residue_states: Tensor, + *, + compute_reconstruction_loss: bool = False, + ) -> ESMplusplusSAEOutput: + # residue_states: (n, d) for n residues; c is the codebook width, k the retained features + standardized = _standardize_residue_states(residue_states) # (n, d) + preactivations = F.relu((standardized - self.b_dec) @ self.W_enc) # (n, c) + retained = torch.topk(preactivations, self.params.k, dim=-1) # values, indices: (n, k) + feature_magnitudes = torch.zeros_like(preactivations).scatter( + -1, retained.indices, retained.values + ) # (n, c) + if not compute_reconstruction_loss: + return ESMplusplusSAEOutput(feature_magnitudes=feature_magnitudes) + + reconstructed = feature_magnitudes @ self.W_dec + self.b_dec # (n, d) + return ESMplusplusSAEOutput( + feature_magnitudes=feature_magnitudes, + reconstruction_loss=(reconstructed - standardized).pow(2).mean(dim=-1), # (n,) + ) + + def get_sae_output(self, layer_states: Tensor, token_mask: Tensor) -> ESMplusplusSAEOutput: + """Encode the unpadded residues of one ESM++ hidden state. + + This name and signature are the attachment contract that + ``PreTrainedESMplusplusModel.add_sae_models`` validates and calls. + """ + + # layer_states: (b, l, d); token_mask: (b, l) + residue_states = layer_states[token_mask] # (n, d) for n valid tokens + encoded: ESMplusplusSAEOutput = self(residue_states) + return encoded + + +def _repository_file( + repository: str | os.PathLike[str], + filename: str, + *, + revision: str | None, + cache_dir: str | os.PathLike[str] | None, + token: str | bool | None, + local_files_only: bool, +) -> Path: + """Resolve one repository file, from a local directory or the Hub cache.""" + + # A directory counts as local only when it holds the shared config, so a stale directory named + # like a Hub identifier cannot shadow the download. + local_directory = Path(repository) + if (local_directory / SAE_CONFIG_FILE).is_file(): + path = local_directory / filename + if not path.is_file(): + raise FileNotFoundError(f"SAE repository {local_directory} has no {filename}.") + return path + + return Path( + hf_hub_download( + repo_id=str(repository), + filename=filename, + revision=revision, + cache_dir=None if cache_dir is None else str(cache_dir), + token=token, + local_files_only=local_files_only, + ) + ) + + +def _load_sae_layer( + shard_path: Path, + *, + params: ESMplusplusSAEParams, + device: torch.device, + dtype: torch.dtype | None, +) -> ESMplusplusSAELayer: + state = load_file(str(shard_path), device=str(device)) + encoder = state.get("W_enc") + if encoder is None: + raise ValueError(f"{shard_path} is not an ESMC SAE shard; it has no 'W_enc' entry.") + + # Build on the meta device so the shard tensors are the only materialized copy of a codebook + # that reaches roughly one gigabyte for the widest published SAEs. + with torch.device("meta"): + sae_layer = ESMplusplusSAELayer(params) + sae_layer.to(dtype=encoder.dtype if dtype is None else dtype) + sae_layer.to_empty(device=device) + # to_empty leaves the statistics buffers uninitialized, so restore the identity defaults that + # shards without trained statistics rely on. + sae_layer.idf.fill_(1.0) + sae_layer.max.fill_(1.0) + + incompatible = sae_layer.load_state_dict(state, strict=False) + missing = tuple( + name for name in incompatible.missing_keys if name not in SAE_OPTIONAL_STATE_NAMES + ) + unexpected = tuple(incompatible.unexpected_keys) + if missing or unexpected: + raise ValueError( + f"{shard_path} does not match the ESMC SAE state contract; " + f"missing {list(missing)}, unexpected {list(unexpected)}." + ) + return sae_layer + + +def load_esmc_sae_layers( + repository: str | os.PathLike[str], + layers: Sequence[int], + *, + revision: str | None = None, + cache_dir: str | os.PathLike[str] | None = None, + token: str | bool | None = None, + local_files_only: bool = False, + device: torch.device | str = "cpu", + dtype: torch.dtype | None = None, +) -> dict[int, ESMplusplusSAELayer]: + """Load the requested hidden-state SAE layers from a Hub repository or a local directory. + + Only the shared config and the requested shards are read, so a repository that publishes every + backbone layer costs one shard per requested layer. ``dtype`` defaults to the dtype stored in + the shard; pass the ESM++ model dtype when the SAE has to consume its hidden states directly. + """ + + requested = tuple(dict.fromkeys(int(layer) for layer in layers)) + if not requested: + raise ValueError("Loading SAE layers requires at least one backbone layer index.") + + config_path = _repository_file( + repository, + SAE_CONFIG_FILE, + revision=revision, + cache_dir=cache_dir, + token=token, + local_files_only=local_files_only, + ) + config = json.loads(config_path.read_text(encoding="utf-8")) + missing_fields = [name for name in SAE_REQUIRED_CONFIG_FIELDS if name not in config] + if missing_fields: + raise ValueError(f"{config_path} is not an ESMC SAE config; it omits {missing_fields}.") + available = tuple(int(index) for index in config.get("available_layers", ())) + + target_device = torch.device(device) + sae_layers: dict[int, ESMplusplusSAELayer] = {} + for layer in requested: + if available and layer not in available: + raise ValueError( + f"SAE repository {repository} does not publish layer {layer}; " + f"available layers are {list(available)}." + ) + shard_path = _repository_file( + repository, + f"layer_{layer}.safetensors", + revision=revision, + cache_dir=cache_dir, + token=token, + local_files_only=local_files_only, + ) + sae_layers[layer] = _load_sae_layer( + shard_path, + params=ESMplusplusSAEParams( + d_model=int(config["d_model"]), + codebook_dim=int(config["codebook_dim"]), + k=int(config["k"]), + layer=layer, + ), + device=target_device, + dtype=dtype, + ) + return sae_layers + + +__all__ = [ + "ESMplusplusSAELayer", + "ESMplusplusSAEOutput", + "ESMplusplusSAEParams", + "load_esmc_sae_layers", +] diff --git a/fastplms/models/esmfold2/__init__.py b/fastplms/models/esmfold2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..6868ae5312cab54842dd34ff5a6c649d4974a7bd --- /dev/null +++ b/fastplms/models/esmfold2/__init__.py @@ -0,0 +1,49 @@ +"""ESMFold2 public classes, imported lazily to keep optional extras isolated.""" + +from __future__ import annotations + +from importlib import import_module +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from .configuration_esmfold2 import ESMFold2Config as ESMFold2Config + from .modeling_esmfold2 import ESMFold2Model as ESMFold2Model + from .modeling_esmfold2 import ESMFold2Output as ESMFold2Output + from .modeling_esmfold2_classification import ( + ESMFold2ExperimentalForSequenceClassification as ESMFold2ExperimentalForSequenceClassification, + ESMFold2ExperimentalForTokenClassification as ESMFold2ExperimentalForTokenClassification, + ESMFold2ForSequenceClassification as ESMFold2ForSequenceClassification, + ESMFold2ForTokenClassification as ESMFold2ForTokenClassification, + ) + from .modeling_esmfold2_experimental import ( + ESMFold2ExperimentalModel as ESMFold2ExperimentalModel, + ) + from .reproducibility import seed_context as seed_context + +_EXPORT_MODULES = { + "ESMFold2Config": ".configuration_esmfold2", + "ESMFold2ExperimentalModel": ".modeling_esmfold2_experimental", + "ESMFold2ExperimentalForSequenceClassification": ".modeling_esmfold2_classification", + "ESMFold2ExperimentalForTokenClassification": ".modeling_esmfold2_classification", + "ESMFold2ForSequenceClassification": ".modeling_esmfold2_classification", + "ESMFold2ForTokenClassification": ".modeling_esmfold2_classification", + "ESMFold2Model": ".modeling_esmfold2", + "ESMFold2Output": ".modeling_esmfold2", + "seed_context": ".reproducibility", +} + + +def __getattr__(name: str) -> Any: + module_name = _EXPORT_MODULES.get(name) + if module_name is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + value = getattr(import_module(module_name, __name__), name) + globals()[name] = value + return value + + +def __dir__() -> list[str]: + return sorted(set(globals()) | set(_EXPORT_MODULES)) + + +__all__ = list(_EXPORT_MODULES) diff --git a/fastplms/models/esmfold2/attention.py b/fastplms/models/esmfold2/attention.py new file mode 100644 index 0000000000000000000000000000000000000000..5823fc67550fd0fdbb1bcb7d5f0aa143d1a78c79 --- /dev/null +++ b/fastplms/models/esmfold2/attention.py @@ -0,0 +1,48 @@ +"""Transformers-compatible attention selection for ESMFold2's ESMC backbone.""" + +from __future__ import annotations + +from collections.abc import Mapping + +from ...attention import FastPLMsAttentionMixin, get_attn_implementation + + +class ESMFold2AttentionMixin(FastPLMsAttentionMixin): + """Route the outer Transformers attention API into the loaded ESMC model.""" + + _supports_attention_backend = True + _supports_sdpa = True + _supports_flex_attn = True + _supports_flash_attn_2 = False + _supports_flash_attn_3 = False + _fastplms_attention_implementations = ( + "eager", + "sdpa", + "flex_attention", + ) + + def __init__(self, config, *args, **kwargs) -> None: + super().__init__(config, *args, **kwargs) + config.esmc_attn_backend = get_attn_implementation(config) + + def set_attn_implementation( + self, + attn_implementation: str | Mapping[str, str], + allow_all_kernels: bool = False, + ) -> None: + """Set one canonical backend on ESMFold2 and its loaded ESMC model.""" + + if allow_all_kernels: + raise ValueError( + "ESMFold2 accepts only its declared built-in attention backends; " + "external attention kernels are not supported." + ) + super().set_attn_implementation(attn_implementation) + resolved = get_attn_implementation(self.config) + self.config.esmc_attn_backend = resolved + esmc = getattr(self, "_esmc", None) + if esmc is not None: + esmc.set_attn_implementation(resolved) + + +__all__ = ["ESMFold2AttentionMixin"] diff --git a/fastplms/models/esmfold2/configuration_esmfold2.py b/fastplms/models/esmfold2/configuration_esmfold2.py new file mode 100644 index 0000000000000000000000000000000000000000..405a809c64cd16de990b147cd5b8117b9df16f5e --- /dev/null +++ b/fastplms/models/esmfold2/configuration_esmfold2.py @@ -0,0 +1,336 @@ +# Copyright 2026 Biohub. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Configuration schema for release and experimental ESMFold2 checkpoints.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +from typing import Any, TypeVar, cast + +from transformers.configuration_utils import PretrainedConfig + +from fastplms.attention import canonical_checkpoint_attention_backend + + +_ESMC_ATTENTION_IMPLEMENTATIONS = frozenset({"eager", "flex_attention", "sdpa"}) +_ESMC_PRECISIONS = frozenset({"auto", "bf16", "fp32", "fp8"}) + + +def _esmc_backbone_checkpoint_ids() -> tuple[str, str]: + """Return the manifest-pinned official and FastPLMs ESMC repositories.""" + + from fastplms.registry import RegistryError, get_model_registry + + registry = get_model_registry() + family = registry.families["esmfold2"] + if family.backbone_model is None: + raise RegistryError("families.esmfold2 must declare backbone_model.") + backbone = registry[family.backbone_model] + return backbone.official.repo_id, backbone.fast.repo_id + + +def normalize_esmc_id(esmc_id: str) -> str: + """Resolve an official ESMC identifier to its FastPLMs checkpoint mirror.""" + + from fastplms.registry import get_model_registry + + official_repo, fast_repo = _esmc_backbone_checkpoint_ids() + if esmc_id == official_repo: + return fast_repo + registry = get_model_registry() + for spec in registry.by_family("esmfold2"): + if spec.backbone is not None and spec.backbone_model is not None: + backbone = registry[spec.backbone_model] + if esmc_id in {spec.backbone.repo_id, backbone.official.repo_id}: + return backbone.fast.repo_id + return esmc_id + + +def normalize_esmc_attention_implementation( + implementation: str | dict[str, str] | None, +) -> str | None: + """Validate the ESMC backend and translate the historical ``flex`` name.""" + + if isinstance(implementation, dict): + if tuple(implementation) != ("",): + raise ValueError( + "ESMFold2 has one ESMC attention backbone; use a string or {'': implementation}." + ) + implementation = implementation[""] + canonical = canonical_checkpoint_attention_backend(implementation) + if canonical is not None and canonical not in _ESMC_ATTENTION_IMPLEMENTATIONS: + expected = sorted(_ESMC_ATTENTION_IMPLEMENTATIONS) + raise ValueError( + f"Unsupported ESMFold2 attention implementation {canonical!r}; " + f"expected one of {expected}." + ) + return canonical + + +NestedConfig = TypeVar("NestedConfig") + + +def _nested_config(value: Any, config_type: type[NestedConfig]) -> NestedConfig: + if isinstance(value, config_type): + return value + return config_type(**value) if isinstance(value, dict) else config_type() + + +def _coerce_nested_field( + value: NestedConfig | dict[str, Any], config_type: type[NestedConfig] +) -> NestedConfig: + """Convert serialized nested dictionaries while retaining supplied objects.""" + + return config_type(**value) if isinstance(value, dict) else value + + +@dataclass +class AtomAttentionConfig: + """Sliding-window atom attention and three-dimensional RoPE settings.""" + + d_atom: int = field(default=128) + d_token: int = field(default=768) + n_blocks: int = field(default=3) + n_heads: int = field(default=4) + swa_window_size: int = field(default=128) + expansion_ratio: int = field(default=2) + spatial_rope_base_frequency: float = field(default=20.0) + n_spatial_rope_pairs_per_axis: int = field(default=2) + n_uid_rope_pairs: int = field(default=10) + uid_rope_base_frequency: float = field(default=10000.0) + + +@dataclass +class DiffusionModuleConfig: + """Dimensions and depth of the coordinate diffusion network.""" + + sigma_data: float = field(default=16.0) + c_atom: int = field(default=128) + c_token: int = field(default=768) + c_z: int = field(default=256) + c_s_inputs: int = field(default=451) + fourier_dim: int = field(default=256) + relpos_r_max: int = field(default=32) + relpos_s_max: int = field(default=2) + atom_num_blocks: int = field(default=3) + atom_num_heads: int = field(default=4) + token_num_blocks: int = field(default=12) + token_num_heads: int = field(default=16) + transition_multiplier: int = field(default=2) + + +@dataclass +class FoldingTrunkConfig: + """Iterative pair/single trunk dimensions.""" + + n_layers: int = field(default=24) + n_heads: int = field(default=8) + dropout: float = field(default=0.0) + + +@dataclass +class InputsEmbedderConfig: + """Input feature width and atom encoder settings.""" + + d_inputs: int = field(default=451) + atom_encoder: AtomAttentionConfig = field(default_factory=AtomAttentionConfig) + + def __post_init__(self) -> None: + self.atom_encoder = _coerce_nested_field(self.atom_encoder, AtomAttentionConfig) + + +@dataclass +class DiffusionStructureHeadConfig: + """Training and inference schedules for coordinate denoising.""" + + diffusion_module: DiffusionModuleConfig = field(default_factory=DiffusionModuleConfig) + distogram_bins: int = field(default=128) + train_noise_log_mean: float = field(default=-1.2) + train_noise_log_std: float = field(default=1.5) + gamma_0: float = field(default=0.605) + gamma_min: float = field(default=1.107) + noise_scale: float = field(default=0.0) + step_scale: float = field(default=1.0) + inference_s_max: float = field(default=160.0) + inference_s_min: float = field(default=4e-4) + inference_p: float = field(default=8.0) + inference_num_steps: int = field(default=68) + + def __post_init__(self) -> None: + self.diffusion_module = _coerce_nested_field(self.diffusion_module, DiffusionModuleConfig) + + +@dataclass +class ConfidenceHeadConfig: + """Confidence-bin definitions and the compact confidence trunk.""" + + enabled: bool = field(default=True) + num_plddt_bins: int = field(default=50) + num_pde_bins: int = field(default=64) + num_pae_bins: int = field(default=64) + min_dist: float = field(default=2.0) + max_dist: float = field(default=52.0) + distogram_bins: int = field(default=128) + folding_trunk: FoldingTrunkConfig = field( + default_factory=lambda: FoldingTrunkConfig(n_layers=4) + ) + + def __post_init__(self) -> None: + self.folding_trunk = _coerce_nested_field(self.folding_trunk, FoldingTrunkConfig) + + +@dataclass +class MSAEncoderConfig: + """Optional multiple-sequence-alignment encoder settings.""" + + enabled: bool = field(default=False) + d_msa: int = field(default=128) + d_hidden: int = field(default=32) + n_layers: int = field(default=4) + n_heads_msa: int = field(default=8) + msa_head_width: int = field(default=32) + + +@dataclass +class LMEncoderConfig: + """Release-model pair encoder derived from language-model states.""" + + enabled: bool = field(default=True) + n_layers: int = field(default=4) + lm_dropout: float = field(default=0.25) + per_loop_lm_dropout: bool = field(default=True) + + +@dataclass +class ParcaeConfig: + """Release-model diffusion-loop scheduler settings.""" + + enabled: bool = field(default=True) + poisson_mean: float = field(default=3.0) + min_steps: int = field(default=1) + max_steps: int | None = field(default=6) + coda_n_layers: int = field(default=2) + + +_SCALAR_DEFAULTS: tuple[tuple[str, Any], ...] = ( + ("d_single", 384), + ("d_pair", 256), + ("n_relative_residx_bins", 32), + ("n_relative_chain_bins", 2), + ("num_loops", 10), + ("num_diffusion_samples", 8), + ("disable_msa_features", False), + ("lm_dropout", 0.0), + ("force_lm_dropout_during_inference", False), + ("lm_mask_pct", 0.0), + ("lm_d_model", 2560), + ("lm_num_layers", 80), +) +_NESTED_CONFIGS = ( + ("inputs", InputsEmbedderConfig), + ("folding_trunk", FoldingTrunkConfig), + ("structure_head", DiffusionStructureHeadConfig), + ("confidence_head", ConfidenceHeadConfig), + ("msa_encoder", MSAEncoderConfig), + ("parcae", ParcaeConfig), + ("lm_encoder", LMEncoderConfig), +) + + +class ESMFold2Config(PretrainedConfig): + """Serializable ESMFold2 architecture, runtime, and precision settings.""" + + model_type = "esmfold2" + has_no_defaults_at_init = True + + def __init__(self, **kwargs: Any) -> None: + legacy_backend = normalize_esmc_attention_implementation(kwargs.get("esmc_attn_backend")) + requested_backend = normalize_esmc_attention_implementation( + kwargs.get("attn_implementation") + ) + resolved_backend = requested_backend or legacy_backend + kwargs["attn_implementation"] = resolved_backend + super().__init__(**kwargs) + + self.type = kwargs.get("type", "release") + if self.type not in {"experimental", "release"}: + raise ValueError( + f"ESMFold2Config.type must be 'release' or 'experimental', got {self.type!r}" + ) + + for name, default in _SCALAR_DEFAULTS: + setattr(self, name, kwargs.get(name, default)) + + _official_esmc_repo, default_esmc_repo = _esmc_backbone_checkpoint_ids() + self.esmc_id = normalize_esmc_id(kwargs.get("esmc_id", default_esmc_repo)) + self.esmc_attn_backend = resolved_backend + self.esmc_precision = str(kwargs.get("esmc_precision", "auto")) + if self.esmc_precision not in _ESMC_PRECISIONS: + raise ValueError( + "esmc_precision must be 'auto', 'bf16', 'fp32', or 'fp8', " + f"got {self.esmc_precision!r}." + ) + + for name, config_type in _NESTED_CONFIGS: + setattr(self, name, _nested_config(kwargs.get(name), config_type)) + if not isinstance(self.msa_encoder.enabled, bool): + raise TypeError("msa_encoder.enabled must be a boolean.") + declared_msa_conditioning = kwargs.get("msa_conditioning") + if "msa_conditioning" in kwargs and not isinstance(declared_msa_conditioning, bool): + raise TypeError("msa_conditioning must be a boolean when provided.") + self.msa_conditioning = ( + self.msa_encoder.enabled + if "msa_conditioning" not in kwargs + else declared_msa_conditioning + ) + if self.msa_conditioning != self.msa_encoder.enabled: + raise ValueError( + "msa_conditioning must match msa_encoder.enabled; received " + f"{self.msa_conditioning!r} and {self.msa_encoder.enabled!r}." + ) + self.msa_encoder_overwrite = bool(kwargs.get("msa_encoder_overwrite", True)) + + self.classifier_train_scope = str(kwargs.get("classifier_train_scope", "probe")) + if self.classifier_train_scope not in {"probe", "projection"}: + raise ValueError( + "classifier_train_scope must be 'probe' or 'projection', " + f"got {self.classifier_train_scope!r}." + ) + self.classifier_probe_hidden_size = int( + kwargs.get("classifier_probe_hidden_size", 512) + ) + self.classifier_probe_num_heads = int(kwargs.get("classifier_probe_num_heads", 4)) + self.classifier_probe_dropout = float(kwargs.get("classifier_probe_dropout", 0.1)) + self.classifier_hidden_size = int(kwargs.get("classifier_hidden_size", 4096)) + self.classifier_dropout = float(kwargs.get("classifier_dropout", 0.2)) + self.classifier_pooling_types = list( + kwargs.get("classifier_pooling_types", ["mean"]) + ) + + def to_dict(self) -> dict[str, Any]: + output = cast(dict[str, Any], super().to_dict()) + for name, _config_type in _NESTED_CONFIGS: + output[name] = asdict(getattr(self, name)) + return output + + +__all__ = [ + "ESMFold2Config", + "LMEncoderConfig", + "MSAEncoderConfig", + "ParcaeConfig", + "normalize_esmc_attention_implementation", + "normalize_esmc_id", +] diff --git a/fastplms/models/esmfold2/embedding.py b/fastplms/models/esmfold2/embedding.py new file mode 100644 index 0000000000000000000000000000000000000000..e22160f1f45ae92c08ebb42d81409f59874d036e --- /dev/null +++ b/fastplms/models/esmfold2/embedding.py @@ -0,0 +1,115 @@ +"""ESMFold2 integration for the shared FastPLMs embedding API.""" + +from __future__ import annotations + +import torch +from typing import Any, ClassVar +from torch import Tensor + +from ...embeddings import EmbeddingBatch, EmbeddingResult, embed_dataset +from .esmfold2_constants_esm3 import SEQUENCE_PAD_TOKEN, SEQUENCE_VOCAB + + +_TOKEN_TO_ID = {token: index for index, token in enumerate(SEQUENCE_VOCAB)} +_VALID_RESIDUES = frozenset(SEQUENCE_VOCAB[4:31]) - {".", "-", "|"} + + +def _encode_single_chain(sequence: str) -> list[int]: + normalized = sequence.upper() + if not normalized: + raise ValueError("ESMFold2 dataset embedding requires at least one protein residue.") + invalid = sorted(set(normalized) - _VALID_RESIDUES) + if invalid: + raise ValueError( + "ESMFold2 dataset embedding accepts one ungapped protein chain; " + f"invalid symbols: {invalid}." + ) + return [_TOKEN_TO_ID[residue] for residue in normalized] + + +class ESMFold2EmbeddingMixin: + """Learned ESMC sequence summaries for ESMFold2 models.""" + + embedding_unsupported_pooling = frozenset({"cls", "parti"}) + embedding_layer = "all_81_esmc_states" + embedding_projection = "esmfold2_learned_sequence_summary" + embedding_token_policy: ClassVar[dict[str, object]] = { + "unit": "residue", + "normalization": "uppercase", + "include": ["single-chain protein residues"], + "exclude": [ + "BOS", + "EOS", + "padding", + "chain delimiters", + "non-protein tokens", + ], + } + + def project_esmc_hidden_states( + self, + hidden_states: Tensor, + residue_mask: Tensor | None = None, + ) -> Tensor: + """Project H from ``(b, l, 81, 2560)`` to Z with shape ``(b, l, 256)``.""" + + # hidden_states: (b, l, 81, d_model); residue_mask: (b, l) or None. + # d_z is the learned projection width: 256 for released ESMFold2 checkpoints. + if hidden_states.ndim != 4 or hidden_states.shape[-2] != 81: + raise ValueError( + "ESMFold2 projection requires the official ordered 81-state " + "ESMC tensor H with shape (b, l, 81, d_model)." + ) + return self.language_model.project_sequence( + hidden_states, residue_mask + ) # (b, l, d_z) + + def _embedding_batch(self, sequences: list[str], **kwargs: Any) -> EmbeddingBatch: + if kwargs: + raise TypeError(f"Unexpected ESMFold2 embedding options: {', '.join(sorted(kwargs))}.") + if self._esmc is None: + raise RuntimeError("ESMFold2 embeddings require load_esmc=True.") + encoded = [_encode_single_chain(sequence) for sequence in sequences] + sequence_length = max(map(len, encoded)) # l + b = len(encoded) # b + device = self.device + input_ids = torch.full( + (b, sequence_length), + SEQUENCE_PAD_TOKEN, + dtype=torch.long, + device=device, + ) # (b, l) + residue_mask = torch.zeros( + (b, sequence_length), dtype=torch.bool, device=device + ) # (b, l) + for batch_index, token_ids in enumerate(encoded): + length = len(token_ids) # l_i + input_ids[batch_index, :length] = torch.tensor( + token_ids, dtype=torch.long, device=device + ) # (l_i,) -> input_ids[batch_index, :l_i]: (l_i,) + residue_mask[batch_index, :length] = True # (l_i,) + + residue_index = torch.arange(sequence_length, device=device).expand(b, -1) # (b, l) + asym_id = torch.zeros_like(input_ids) # (b, l) + mol_type = torch.zeros_like(input_ids) # (b, l) + hidden_states = self._compute_lm_hidden_states( + input_ids, + asym_id, + residue_index, + mol_type, + residue_mask, + ) # (b, l, 81, d_model) + projected = self.project_esmc_hidden_states( + hidden_states, residue_mask + ) # (b, l, d_z) + return EmbeddingBatch( + X=projected, residue_mask=residue_mask + ) # X: (b, l, d_z); residue_mask: (b, l) + + def embed_dataset(self, inputs: Any, **kwargs: Any) -> EmbeddingResult: + """Embed single-chain proteins using the learned 256-wide ESMFold2 summary.""" + + return embed_dataset(self, inputs, **kwargs) + + +__all__ = ["ESMFold2EmbeddingMixin"] diff --git a/fastplms/models/esmfold2/esmfold2_affine3d.py b/fastplms/models/esmfold2/esmfold2_affine3d.py new file mode 100644 index 0000000000000000000000000000000000000000..17c47f2fbbbd90bf8006b58f0b0a537499b4fd23 --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_affine3d.py @@ -0,0 +1,605 @@ +"""Differentiable rigid rotations and affine transforms for ESMFold2.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Self + +import torch +from torch.nn import functional as F + +from .esmfold2_misc import fp32_autocast_context + + +def _index_tuple(index: Any) -> tuple[Any, ...]: + if isinstance(index, int) or index is None: + return (index,) + return tuple(index) + + +def _sqrt_subgradient(values: torch.Tensor) -> torch.Tensor: + """Square root with a zero subgradient for non-positive inputs.""" + + result = torch.zeros_like(values) + positive = values > 0 + result[positive] = torch.sqrt(values[positive]) + return result + + +def _quat_invert(quaternion: torch.Tensor) -> torch.Tensor: + conjugate_sign = torch.tensor([1, -1, -1, -1], device=quaternion.device) + return quaternion * conjugate_sign + + +def _quat_mult(left: torch.Tensor, right: torch.Tensor) -> torch.Tensor: + """Hamilton product for real-first quaternion tensors.""" + + aw, ax, ay, az = torch.unbind(left, -1) + bw, bx, by, bz = torch.unbind(right, -1) + return torch.stack( + ( + aw * bw - ax * bx - ay * by - az * bz, + aw * bx + ax * bw + ay * bz - az * by, + aw * by - ax * bz + ay * bw + az * bx, + aw * bz + ax * by - ay * bx + az * bw, + ), + -1, + ) + + +def _quat_rotation( + quaternion: torch.Tensor, + points: torch.Tensor, +) -> torch.Tensor: + """Rotate points using normalized real-first quaternions.""" + + aw, ax, ay, az = torch.unbind(quaternion, -1) + bx, by, bz = torch.unbind(points, -1) + product = torch.stack( + ( + -ax * bx - ay * by - az * bz, + aw * bx + ay * bz - az * by, + aw * by - ax * bz + az * bx, + aw * bz + ax * by - ay * bx, + ), + -1, + ) + return _quat_mult(product, _quat_invert(quaternion))[..., 1:] + + +def _graham_schmidt( + x_axis: torch.Tensor, + xy_plane: torch.Tensor, + eps: float = 1e-12, +) -> torch.Tensor: + """Construct a right-handed orthonormal frame from two directions.""" + + with fp32_autocast_context(x_axis.device.type): + e1 = xy_plane + denominator = torch.sqrt((x_axis**2).sum(dim=-1, keepdim=True) + eps) + x_axis = x_axis / denominator + projection = (x_axis * e1).sum(dim=-1, keepdim=True) + e1 = e1 - x_axis * projection + denominator = torch.sqrt((e1**2).sum(dim=-1, keepdim=True) + eps) + e1 = e1 / denominator + e2 = torch.cross(x_axis, e1, dim=-1) + return torch.stack([x_axis, e1, e2], dim=-1) + + +class Rotation: + """Common interface for matrix-backed and quaternion-backed rotations.""" + + @classmethod + def identity(cls, shape: tuple[int, ...], **tensor_kwargs) -> Self: ... + + @classmethod + def random(cls, shape: tuple[int, ...], **tensor_kwargs) -> Self: ... + + def __getitem__(self, idx: Any) -> Self: ... + + @property + def tensor(self) -> torch.Tensor: ... + + @property + def shape(self) -> torch.Size: ... + + def as_matrix(self) -> RotationMatrix: ... + + def as_quat(self, normalize: bool = False) -> RotationQuat: ... + + def compose(self, other: Self) -> Self: ... + + def convert_compose(self, other: Self) -> Self: ... + + def apply(self, points: torch.Tensor) -> torch.Tensor: ... + + def invert(self) -> Self: ... + + @property + def dtype(self) -> torch.dtype: + return self.tensor.dtype + + @property + def device(self) -> torch.device: + return self.tensor.device + + @property + def requires_grad(self) -> bool: + return self.tensor.requires_grad + + @classmethod + def _from_tensor(cls, tensor: torch.Tensor) -> Self: + return cls(tensor) # type: ignore[call-arg] + + def to(self, **kwargs) -> Self: + return self._from_tensor(self.tensor.to(**kwargs)) + + def detach(self, *args, **kwargs) -> Self: + return self._from_tensor(self.tensor.detach(**kwargs)) + + def tensor_apply(self, func) -> Self: + transformed = [func(component) for component in self.tensor.unbind(dim=-1)] + return self._from_tensor(torch.stack(transformed, dim=-1)) + + +class RotationQuat(Rotation): + """A rotation represented by a real-first quaternion.""" + + def __init__(self, quats: torch.Tensor, normalized: bool = False): + if not isinstance(quats, torch.Tensor): + raise TypeError("quats must be a Torch tensor.") + if quats.ndim == 0 or quats.shape[-1] != 4: + raise ValueError( + f"quats must have trailing dimension 4, got shape {tuple(quats.shape)}." + ) + if not isinstance(normalized, bool): + raise TypeError("normalized must be a boolean.") + self._normalized = normalized + if normalized: + quats = F.normalize(quats.to(torch.float32), dim=-1) + self._quats = quats.where(quats[..., :1] >= 0, -quats) + else: + self._quats = quats.to(torch.float32) + + @property + def tensor(self) -> torch.Tensor: + return self._quats + + @property + def shape(self) -> torch.Size: + return self._quats.shape[:-1] + + @classmethod + def identity(cls, shape, **tensor_kwargs) -> RotationQuat: + quaternions = torch.ones((*shape, 4), **tensor_kwargs) + selector = torch.tensor([1, 0, 0, 0], device=quaternions.device) + return cls(quaternions * selector) + + @classmethod + def random(cls, shape, **tensor_kwargs) -> RotationQuat: + return cls(torch.randn((*shape, 4), **tensor_kwargs), normalized=True) + + def __getitem__(self, idx: Any) -> RotationQuat: + indices = _index_tuple(idx) + return RotationQuat(self._quats[(*indices, slice(None))]) + + def normalized(self) -> RotationQuat: + if self._normalized: + return self + return RotationQuat(self._quats, normalized=True) + + def as_quat(self, normalize: bool = False) -> RotationQuat: + return self + + def as_matrix(self) -> RotationMatrix: + quaternion = self.normalized().tensor + r, i, j, k = torch.unbind(quaternion, -1) + scale = 2.0 / torch.linalg.norm(quaternion, dim=-1) + elements = torch.stack( + ( + 1 - scale * (j * j + k * k), + scale * (i * j - k * r), + scale * (i * k + j * r), + scale * (i * j + k * r), + 1 - scale * (i * i + k * k), + scale * (j * k - i * r), + scale * (i * k - j * r), + scale * (j * k + i * r), + 1 - scale * (i * i + j * j), + ), + -1, + ) + return RotationMatrix(elements.reshape((*quaternion.shape[:-1], 3, 3))) + + def compose(self, other: RotationQuat) -> RotationQuat: + with fp32_autocast_context(self.device.type): + return RotationQuat(_quat_mult(self._quats, other._quats)) + + def convert_compose(self, other: Rotation) -> RotationQuat: + return self.compose(other.as_quat()) + + def apply(self, points: torch.Tensor) -> torch.Tensor: + return _quat_rotation(self.normalized()._quats, points) + + def invert(self) -> RotationQuat: + return RotationQuat(_quat_invert(self._quats)) + + +class RotationMatrix(Rotation): + """A rotation represented by a dense FP32 matrix.""" + + def __init__(self, rots: torch.Tensor): + if not isinstance(rots, torch.Tensor): + raise TypeError("rots must be a Torch tensor.") + if rots.ndim > 0 and rots.shape[-1] == 9: + rots = rots.unflatten(-1, (3, 3)) + if rots.ndim < 2 or rots.shape[-2:] != (3, 3): + raise ValueError( + "rots must have trailing shape (3, 3) or flattened width 9, got " + f"shape {tuple(rots.shape)}." + ) + self._rots = rots.to(torch.float32) + + @property + def tensor(self) -> torch.Tensor: + return self._rots.flatten(-2) + + @property + def shape(self) -> torch.Size: + return self._rots.shape[:-2] + + @classmethod + def identity(cls, shape, **tensor_kwargs) -> RotationMatrix: + matrix = torch.eye(3, **tensor_kwargs) + matrix = matrix.view(*(1 for _ in shape), 3, 3) + return cls(matrix.expand(*shape, -1, -1)) + + @classmethod + def random(cls, shape, **tensor_kwargs) -> RotationMatrix: + return RotationQuat.random(shape, **tensor_kwargs).as_matrix() + + @staticmethod + def from_graham_schmidt( + x_axis: torch.Tensor, + xy_plane: torch.Tensor, + eps: float = 1e-12, + ) -> RotationMatrix: + return RotationMatrix(_graham_schmidt(x_axis, xy_plane, eps)) + + def __getitem__(self, idx: Any) -> RotationMatrix: + indices = _index_tuple(idx) + return RotationMatrix(self._rots[(*indices, slice(None), slice(None))]) + + def as_matrix(self) -> RotationMatrix: + return self + + def to_3x3(self) -> torch.Tensor: + return self._rots + + def as_quat(self, normalize: bool = False) -> RotationQuat: + m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind( + self._rots.flatten(-2), + dim=-1, + ) + q_abs = _sqrt_subgradient( + torch.stack( + ( + 1.0 + m00 + m11 + m22, + 1.0 + m00 - m11 - m22, + 1.0 - m00 + m11 - m22, + 1.0 - m00 - m11 + m22, + ), + dim=-1, + ) + ) + products = torch.stack( + ( + q_abs[..., 0] ** 2, + m21 - m12, + m02 - m20, + m10 - m01, + m21 - m12, + q_abs[..., 1] ** 2, + m10 + m01, + m02 + m20, + m02 - m20, + m10 + m01, + q_abs[..., 2] ** 2, + m12 + m21, + m10 - m01, + m20 + m02, + m21 + m12, + q_abs[..., 3] ** 2, + ), + dim=-1, + ).unflatten(-1, (4, 4)) + floor = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device) + candidates = products / (2.0 * q_abs[..., None].max(floor)) + best = torch.zeros_like(q_abs, dtype=torch.bool) + best.scatter_(-1, q_abs.argmax(dim=-1, keepdim=True), True) + quaternion = candidates[best, :].reshape(q_abs.shape) + return RotationQuat(quaternion) + + def compose(self, other: RotationMatrix) -> RotationMatrix: + with fp32_autocast_context(self.device.type): + return RotationMatrix(self._rots @ other._rots) + + def convert_compose(self, other: Rotation) -> RotationMatrix: + return self.compose(other.as_matrix()) + + def apply(self, points: torch.Tensor) -> torch.Tensor: + with fp32_autocast_context(self.device.type): + if self._rots.shape[-3] == 1: + return points @ self._rots.transpose(-1, -2).squeeze(-3) + return torch.einsum("...ij,...j", self._rots, points) + + def invert(self) -> RotationMatrix: + return RotationMatrix(self._rots.transpose(-1, -2)) + + +@dataclass(frozen=True) +class Affine3D: + """A rigid transform with translation and rotation components.""" + + trans: torch.Tensor + rot: Rotation + + def __post_init__(self) -> None: + if not isinstance(self.trans, torch.Tensor): + raise TypeError("trans must be a Torch tensor.") + if not isinstance(self.rot, Rotation): + raise TypeError("rot must implement the ESMFold2 Rotation interface.") + if self.trans.ndim == 0 or self.trans.shape[-1] != 3: + raise ValueError( + "trans must have trailing dimension 3, got " + f"shape {tuple(self.trans.shape)}." + ) + if self.trans.shape[:-1] != self.rot.shape: + raise ValueError( + "translation and rotation batch shapes must match, got " + f"{tuple(self.trans.shape[:-1])} and {tuple(self.rot.shape)}." + ) + + @property + def shape(self) -> torch.Size: + return self.trans.shape[:-1] + + @property + def dtype(self) -> torch.dtype: + return self.trans.dtype + + @property + def device(self) -> torch.device: + return self.trans.device + + @property + def requires_grad(self) -> bool: + return self.trans.requires_grad + + @property + def tensor(self) -> torch.Tensor: + return torch.cat((self.rot.tensor, self.trans), dim=-1) + + @staticmethod + def identity( + shape_or_affine: tuple[int, ...] | Affine3D, + rotation_type: type[Rotation] = RotationMatrix, + **tensor_kwargs, + ) -> Affine3D: + if isinstance(shape_or_affine, Affine3D): + kwargs = { + "dtype": shape_or_affine.dtype, + "device": shape_or_affine.device, + } + kwargs.update(tensor_kwargs) + shape = shape_or_affine.shape + rotation_type = type(shape_or_affine.rot) + else: + kwargs = tensor_kwargs + shape = shape_or_affine + return Affine3D( + torch.zeros((*shape, 3), **kwargs), + rotation_type.identity(shape, **kwargs), + ) + + @staticmethod + def random( + shape: tuple[int, ...], + std: float = 1, + rotation_type: type[Rotation] = RotationMatrix, + **tensor_kwargs, + ) -> Affine3D: + translation = torch.randn((*shape, 3), **tensor_kwargs).mul(std) + rotation = rotation_type.random(shape, **tensor_kwargs) + return Affine3D(trans=translation, rot=rotation) + + @staticmethod + def from_tensor(tensor: torch.Tensor) -> Affine3D: + if not isinstance(tensor, torch.Tensor): + raise TypeError("tensor must be a Torch tensor.") + if tensor.ndim == 0: + raise ValueError("tensor must have at least one dimension.") + width = tensor.shape[-1] + if width == 4: + if tensor.ndim < 2 or tensor.shape[-2] not in (3, 4): + raise ValueError( + "matrix-form affine tensors must have trailing shape (3, 4) or " + f"(4, 4), got {tuple(tensor.shape)}." + ) + translation = tensor[..., :3, 3] + rotation: Rotation = RotationMatrix(tensor[..., :3, :3]) + elif width == 6: + translation = tensor[..., -3:] + rotation = RotationQuat(F.pad(tensor[..., :3], (1, 0), value=1)) + elif width == 7: + translation = tensor[..., -3:] + rotation = RotationQuat(tensor[..., :4]) + elif width == 12: + translation = tensor[..., -3:] + rotation = RotationMatrix(tensor[..., :-3].unflatten(-1, (3, 3))) + else: + raise RuntimeError( + f"Cannot detect rotation format from {tensor.shape[-1] - 3}-d flat vector" + ) + return Affine3D(translation, rotation) + + @staticmethod + def from_tensor_pair( + translation: torch.Tensor, + rotation: torch.Tensor, + ) -> Affine3D: + return Affine3D(translation, RotationMatrix(rotation)) + + @staticmethod + def from_graham_schmidt( + neg_x_axis: torch.Tensor, + origin: torch.Tensor, + xy_plane: torch.Tensor, + eps: float = 1e-10, + ) -> Affine3D: + x_axis = origin - neg_x_axis + plane_direction = xy_plane - origin + rotation = RotationMatrix.from_graham_schmidt( + x_axis, + plane_direction, + eps, + ) + return Affine3D(trans=origin, rot=rotation) + + @staticmethod + def cat(affines: list[Affine3D], dim: int = 0) -> Affine3D: + if not affines: + raise ValueError("affines must contain at least one transform.") + if any(not isinstance(affine, Affine3D) for affine in affines): + raise TypeError("affines must contain only Affine3D instances.") + if dim < 0: + dim = len(affines[0].shape) + dim + return Affine3D.from_tensor(torch.cat([affine.tensor for affine in affines], dim=dim)) + + def __getitem__(self, idx: Any) -> Affine3D: + indices = _index_tuple(idx) + translation = self.trans[(*indices, slice(None))] + return Affine3D(trans=translation, rot=self.rot[idx]) + + def to(self, **kwargs) -> Affine3D: + return Affine3D(self.trans.to(**kwargs), self.rot.to(**kwargs)) + + def detach(self, *args, **kwargs) -> Affine3D: + return Affine3D( + self.trans.detach(**kwargs), + self.rot.detach(**kwargs), + ) + + def tensor_apply(self, func) -> Affine3D: + components = [func(value) for value in self.tensor.unbind(dim=-1)] + return Affine3D.from_tensor(torch.stack(components, dim=-1)) + + def as_matrix(self) -> Affine3D: + return Affine3D(trans=self.trans, rot=self.rot.as_matrix()) + + def as_quat(self, normalize: bool = False) -> Affine3D: + return Affine3D( + trans=self.trans, + rot=self.rot.as_quat(normalize), + ) + + def compose( + self, + other: Affine3D, + autoconvert: bool = False, + ) -> Affine3D: + compose_rotation = self.rot.convert_compose if autoconvert else self.rot.compose + rotation = compose_rotation(other.rot) + translation = self.rot.apply(other.trans) + self.trans + return Affine3D(trans=translation, rot=rotation) + + def compose_rotation( + self, + other: Rotation, + autoconvert: bool = False, + ) -> Affine3D: + compose = self.rot.convert_compose if autoconvert else self.rot.compose + return Affine3D(trans=self.trans, rot=compose(other)) + + def scale(self, value: torch.Tensor | float) -> Affine3D: + return Affine3D(self.trans * value, self.rot) + + def mask(self, mask: torch.Tensor, with_zero: bool = False) -> Affine3D: + if with_zero: + masked = torch.zeros_like(self.tensor).where( + mask[..., None], + self.tensor, + ) + return Affine3D.from_tensor(masked) + identity = self.identity( + self.shape, + rotation_type=type(self.rot), + device=self.device, + dtype=self.dtype, + ).tensor + return Affine3D.from_tensor(identity.where(mask[..., None], self.tensor)) + + def apply(self, points: torch.Tensor) -> torch.Tensor: + return self.rot.apply(points) + self.trans + + def invert(self) -> Affine3D: + rotation = self.rot.invert() + return Affine3D(trans=-rotation.apply(self.trans), rot=rotation) + + +def build_affine3d_from_coordinates( + coords: torch.Tensor, +) -> tuple[Affine3D, torch.Tensor]: + """Build residue frames from X with shape (b, l, 3, 3).""" + + if not isinstance(coords, torch.Tensor): + raise TypeError("coords must be a Torch tensor.") + if coords.ndim != 4 or coords.shape[-2:] != (3, 3): + raise ValueError( + "coords must have shape (batch, length, 3, 3), got " + f"{tuple(coords.shape)}." + ) + + maximum_distance = 1e6 + coord_mask = torch.all( + torch.all( + torch.isfinite(coords) & (coords < maximum_distance), + dim=-1, + ), + dim=-1, + ) + + def backbone_affine(positions: torch.Tensor) -> Affine3D: + n, ca, c = positions.unbind(dim=-2) + return Affine3D.from_graham_schmidt(c, ca, n) + + coords = coords.clone().float() + coords[~coord_mask] = 0 + average = coords.masked_fill(~coord_mask[..., None, None], 0).sum(1) / ( + coord_mask.sum(-1)[..., None, None] + 1e-8 + ) + average_affine = backbone_affine(average.float()).as_matrix() + + b, length, _, _ = coords.shape + rotation = average_affine.rot.tensor[..., None, :].expand(b, length, 9) + translation = average_affine.trans[..., None, :].expand(b, length, 3) + identity = RotationMatrix.identity( + (b, length), + dtype=torch.float32, + device=coords.device, + requires_grad=False, + ) + rotation = rotation.where( + coord_mask.any(-1)[..., None, None], + identity.tensor, + ) + missing_frame = Affine3D(translation, RotationMatrix(rotation)) + + residue_frame = backbone_affine(coords.float()) + residue_frame = Affine3D.from_tensor( + residue_frame.tensor.where( + coord_mask[..., None], + missing_frame.tensor, + ) + ) + return residue_frame, coord_mask diff --git a/fastplms/models/esmfold2/esmfold2_aligner.py b/fastplms/models/esmfold2/esmfold2_aligner.py new file mode 100644 index 0000000000000000000000000000000000000000..a23b9d6d6b727ba020cef08f9208281e65b02393 --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_aligner.py @@ -0,0 +1,87 @@ +"""Rigid alignment for structure dataclasses.""" + +from __future__ import annotations + +from dataclasses import Field, replace +from typing import Any, ClassVar, Protocol, TypeVar + +import numpy as np +import torch +from torch import Tensor + +from .esmfold2_protein_structure import compute_affine_and_rmsd + + +class Alignable(Protocol): + """Minimum structure interface accepted by :class:`Aligner`.""" + + __dataclass_fields__: ClassVar[dict[str, Field[Any]]] + + @property + def atom37_positions(self) -> np.ndarray: ... + + @property + def atom37_mask(self) -> np.ndarray: ... + + def __len__(self) -> int: ... + + +AlignableT = TypeVar("AlignableT", bound=Alignable) + + +def _coordinate_batch(structure: Alignable) -> Tensor: + return torch.as_tensor(structure.atom37_positions, dtype=torch.double).unsqueeze(0) + + +def _shared_atom_mask(mobile: Alignable, target: Alignable, backbone_only: bool) -> Tensor: + shared = np.asarray(mobile.atom37_mask, dtype=bool) & np.asarray( + target.atom37_mask, + dtype=bool, + ) + if backbone_only: + shared = shared.copy() + shared[:, 3:] = False + return torch.from_numpy(shared).unsqueeze(0) + + +class Aligner: + """Fit a mobile structure onto a target with masked Kabsch alignment.""" + + def __init__( + self, + mobile: Alignable, + target: Alignable, + only_use_backbone: bool = False, + use_reflection: bool = False, + ) -> None: + if len(mobile) != len(target): + raise AssertionError("mobile and target must contain the same residue count") + + mobile_coordinates = _coordinate_batch(mobile) + target_coordinates = _coordinate_batch(target) + if use_reflection: + target_coordinates = -target_coordinates + atom_mask = _shared_atom_mask(mobile, target, only_use_backbone) + self._affine3D, rmsd = compute_affine_and_rmsd( + mobile_coordinates, + target_coordinates, + atom_exists_mask=atom_mask, + ) + self._rmsd = rmsd.item() + + @property + def rmsd(self) -> float: + return self._rmsd + + def apply(self, mobile: AlignableT) -> AlignableT: + """Return a dataclass copy with all present atom coordinates aligned.""" + + present = np.asarray(mobile.atom37_mask, dtype=bool) + packed = torch.as_tensor( + mobile.atom37_positions[present], + dtype=torch.float32, + ).unsqueeze(0) + aligned = self._affine3D.apply(packed).squeeze(0).cpu().numpy() + atom37_positions = np.full_like(mobile.atom37_positions, np.nan) + atom37_positions[present] = aligned + return replace(mobile, atom37_positions=atom37_positions) diff --git a/fastplms/models/esmfold2/esmfold2_atom_indexer.py b/fastplms/models/esmfold2/esmfold2_atom_indexer.py new file mode 100644 index 0000000000000000000000000000000000000000..676cbde76427c6c64672c0fe4e991457cdf57d70 --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_atom_indexer.py @@ -0,0 +1,30 @@ +"""Name-based views into an atom-axis property.""" + +from __future__ import annotations + +from operator import attrgetter +from typing import Any + +import numpy as np + +from .esmfold2_protein_structure import index_by_atom_name + + +class AtomIndexer: + """Select named atoms from one property of a structure-like object. + + The wrapper intentionally remains small because ``ProteinChain.atom37`` and + related public properties expose it directly. + """ + + __slots__ = ("_get_property", "dim", "property", "structure") + + def __init__(self, structure: Any, property: str, dim: int): + self.structure = structure + self.property = property + self.dim = dim + self._get_property = attrgetter(property) + + def __getitem__(self, atom_names: str | list[str]) -> np.ndarray: + values = self._get_property(self.structure) + return index_by_atom_name(values, atom_names, dim=self.dim) diff --git a/fastplms/models/esmfold2/esmfold2_conformers.py b/fastplms/models/esmfold2/esmfold2_conformers.py new file mode 100644 index 0000000000000000000000000000000000000000..ddd0e9b55e4175aa0d144c7330ecfcd1d3d902d4 --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_conformers.py @@ -0,0 +1,402 @@ +"""Lazy access to Chemical Component Dictionary conformers. + +The feature pipeline depends on atom names, formal charges, bonds, leaving-atom +flags, and one preferred reference conformer. Asset resolution is explicit at +``load_ccd`` time; importing this module performs no download or file access. +""" + +from __future__ import annotations + +import os +import pickle +import stat +import tempfile +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from hashlib import file_digest +from pathlib import Path +from typing import Any, BinaryIO + +import numpy as np +from huggingface_hub import hf_hub_download +from huggingface_hub.constants import HF_HUB_CACHE + +from fastplms.registry import RuntimeAsset, get_model_registry + +from .esmfold2_constants import RES_TYPE_TO_CCD + +_CCD_ENVIRONMENT_VARIABLE = "ESMCFOLD_CCD_PATH" +_CCD_ASSET_ID = "esmfold2_ccd" + + +@dataclass(frozen=True) +class _ResolvedAsset: + path: Path + trusted_hub_cache_root: Path | None = None + + +def _asset_contract() -> RuntimeAsset: + """Return the manifest-owned identity of the trusted CCD pickle.""" + + try: + asset = get_model_registry().runtime_assets[_CCD_ASSET_ID] + except KeyError as error: + raise RuntimeError( + f"The package manifest does not declare runtime asset {_CCD_ASSET_ID!r}." + ) from error + if asset.trust_kind != "hash_pinned_pickle": + raise RuntimeError( + f"Runtime asset {_CCD_ASSET_ID!r} must use the hash_pinned_pickle trust policy." + ) + return asset + + +@contextmanager +def _open_verified_asset( + asset_path: Path, + contract: RuntimeAsset, + *, + trusted_hub_cache_root: Path | None = None, +) -> Iterator[BinaryIO]: + """Yield a private snapshot containing exactly the verified pickle bytes.""" + + try: + path_state = asset_path.lstat() + except FileNotFoundError as error: + raise FileNotFoundError(f"CCD asset does not exist: {asset_path}") from error + opened_path = asset_path + if stat.S_ISLNK(path_state.st_mode): + if trusted_hub_cache_root is None: + raise ValueError(f"CCD asset must not be a symlink: {asset_path}") + opened_path = _resolve_trusted_hub_snapshot_link( + asset_path, + contract, + trusted_hub_cache_root, + ) + path_state = opened_path.lstat() + if not stat.S_ISREG(path_state.st_mode): + raise ValueError(f"CCD asset must be a regular file: {asset_path}") + + flags = os.O_RDONLY | getattr(os, "O_BINARY", 0) | getattr(os, "O_CLOEXEC", 0) + flags |= getattr(os, "O_NOFOLLOW", 0) + descriptor: int | None = None + try: + descriptor = os.open(opened_path, flags) + opened_state = os.fstat(descriptor) + if not stat.S_ISREG(opened_state.st_mode): + raise ValueError(f"CCD asset must be a regular file: {asset_path}") + if (path_state.st_dev, path_state.st_ino) != ( + opened_state.st_dev, + opened_state.st_ino, + ): + raise ValueError(f"CCD asset changed while it was being opened: {asset_path}") + + source = os.fdopen(descriptor, "rb") + descriptor = None + with source, tempfile.TemporaryFile(mode="w+b") as snapshot: + actual_size = opened_state.st_size + if actual_size != contract.size: + raise ValueError( + "CCD asset size mismatch: " + f"expected {contract.size} bytes, received {actual_size}." + ) + # Copy into a loader-owned OS temporary file. Hashing and + # deserialization then consume the same immutable snapshot, so a + # path replacement or in-place source write cannot substitute + # unverified pickle bytes after validation. + remaining = contract.size + while remaining: + chunk = source.read(min(1024 * 1024, remaining)) + if not chunk: + break + snapshot.write(chunk) + remaining -= len(chunk) + copied_size = snapshot.tell() + extra_byte = source.read(1) + if remaining or extra_byte: + observed_size = copied_size if remaining else copied_size + len(extra_byte) + raise ValueError( + "CCD asset size changed while it was being copied: " + f"expected {contract.size} bytes, received at least {observed_size}." + ) + snapshot.flush() + snapshot.seek(0) + actual_hash = file_digest(snapshot, "sha256").hexdigest() + if actual_hash != contract.sha256: + raise ValueError( + "CCD asset SHA256 mismatch; refusing to cross the " + "trusted-pickle boundary." + ) + snapshot.seek(0) + yield snapshot + finally: + if descriptor is not None: + os.close(descriptor) + + +def _resolve_trusted_hub_snapshot_link( + asset_path: Path, + contract: RuntimeAsset, + cache_root: Path, +) -> Path: + """Resolve only the immutable Hub snapshot link declared by the manifest.""" + + root = cache_root.expanduser().resolve(strict=True) + if len(contract.revision) != 40 or any( + character not in "0123456789abcdef" for character in contract.revision.lower() + ): + raise ValueError("CCD Hub asset revision must be an immutable 40-character commit.") + relative_asset = Path(contract.path) + if relative_asset.is_absolute() or ".." in relative_asset.parts: + raise ValueError(f"CCD Hub asset path is unsafe: {contract.path!r}") + repository_cache = root / f"models--{contract.repository.replace('/', '--')}" + try: + repository_cache.resolve(strict=True).relative_to(root) + except (FileNotFoundError, ValueError) as error: + raise ValueError( + f"CCD Hub repository cache escapes the effective Hub cache root: {repository_cache}" + ) from error + snapshot_root = repository_cache / "snapshots" / contract.revision + expected_path = snapshot_root / relative_asset + lexical_path = Path(os.path.abspath(asset_path)) + if lexical_path != Path(os.path.abspath(expected_path)): + raise ValueError( + "CCD Hub symlink is not the manifest-owned immutable snapshot path: " + f"{asset_path}" + ) + + try: + asset_path.parent.resolve(strict=True).relative_to(root) + except (FileNotFoundError, ValueError) as error: + raise ValueError( + f"CCD Hub snapshot path escapes the effective Hub cache root: {asset_path}" + ) from error + + resolved = asset_path.resolve(strict=True) + blob_root = (repository_cache / "blobs").resolve(strict=True) + try: + blob_root.relative_to(root) + resolved.relative_to(blob_root) + except ValueError as error: + raise ValueError( + f"CCD Hub snapshot link escapes its repository blob cache: {asset_path}" + ) from error + if not resolved.is_file() or resolved.is_symlink(): + raise ValueError(f"CCD Hub snapshot target must be a regular file: {resolved}") + return resolved + + +class _ChemicalComponentStore: + def __init__(self) -> None: + self.molecules: dict[str, Any] | None = None + self.conformers: dict[str, dict[str, np.ndarray]] = {} + self.atoms: dict[str, list[tuple[str, str, int]]] = {} + self.bonds: dict[str, list[tuple[str, str]]] = {} + self.leaving_atoms: dict[str, set[str]] = {} + self.standard_positions: dict[tuple[int, str], np.ndarray | None] = {} + self.ligand_positions: dict[tuple[str, str], np.ndarray | None] = {} + + def load(self, cache_dir: Path | str | None = None) -> dict[str, Any]: + if self.molecules is not None: + return self.molecules + contract = _asset_contract() + resolved = self._resolve_asset_location(cache_dir, contract) + asset = resolved.path + try: + # SECURITY: the private snapshot is both hash-validated and + # deserialized, closing path-replacement and in-place-write races. + with _open_verified_asset( + asset, + contract, + trusted_hub_cache_root=resolved.trusted_hub_cache_root, + ) as handle: + loaded = pickle.load(handle) + except FileNotFoundError: + raise + except Exception as error: + raise ValueError(f"Could not read the CCD asset at {asset}: {error}") from error + if loaded is not None and not isinstance(loaded, dict): + raise TypeError("The CCD asset must contain a component dictionary.") + self.molecules = loaded or {} + return self.molecules + + @staticmethod + def _resolve_asset(cache_dir: Path | str | None) -> Path: + contract = _asset_contract() + return _ChemicalComponentStore._resolve_asset_location(cache_dir, contract).path + + @staticmethod + def _resolve_asset_location( + cache_dir: Path | str | None, + contract: RuntimeAsset, + ) -> _ResolvedAsset: + configured = os.environ.get(_CCD_ENVIRONMENT_VARIABLE) + if configured: + asset = Path(configured).expanduser() + elif cache_dir is not None: + asset = Path(cache_dir).expanduser() / contract.path + else: + try: + asset = Path( + hf_hub_download( + repo_id=contract.repository, + filename=contract.path, + revision=contract.revision, + ) + ) + except Exception as error: + raise FileNotFoundError( + "Could not resolve the ESMFold2 CCD asset. Set " + f"{_CCD_ENVIRONMENT_VARIABLE} or populate the Hugging Face cache." + ) from error + return _ResolvedAsset( + path=asset, + trusted_hub_cache_root=Path(HF_HUB_CACHE), + ) + return _ResolvedAsset(path=asset) + + def _component_with_conformer(self, component_id: str): + molecule = self.load().get(component_id) + if molecule is None or molecule.GetNumConformers() == 0: + return None, None + + conformers = list(molecule.GetConformers()) + priority = {"Computed": 0, "Ideal": 1} + selected_index = min( + range(len(conformers)), + key=lambda index: priority.get(conformers[index].GetPropsAsDict().get("name"), 2), + ) + + from rdkit import Chem + + heavy_molecule = Chem.RemoveHs(molecule, sanitize=False) + if heavy_molecule.GetNumConformers() == 0: + return None, None + conformer_index = min(selected_index, heavy_molecule.GetNumConformers() - 1) + return heavy_molecule, heavy_molecule.GetConformer(conformer_index) + + def conformer(self, component_id: str) -> dict[str, np.ndarray] | None: + if component_id not in self.conformers: + molecule, conformer = self._component_with_conformer(component_id) + positions: dict[str, np.ndarray] = {} + if molecule is not None and conformer is not None: + for atom in molecule.GetAtoms(): + atom_name = atom.GetPropsAsDict().get("name") + if not isinstance(atom_name, str) or not atom_name: + continue + point = conformer.GetAtomPosition(atom.GetIdx()) + positions[atom_name] = np.asarray((point.x, point.y, point.z), dtype=np.float32) + self.conformers[component_id] = positions + result = self.conformers[component_id] + return result or None + + def atom_records(self, component_id: str) -> list[tuple[str, str, int]] | None: + if component_id not in self.atoms: + molecule, _conformer = self._component_with_conformer(component_id) + records: list[tuple[str, str, int]] = [] + if molecule is not None: + for atom in molecule.GetAtoms(): + atom_name = atom.GetPropsAsDict().get("name") + if isinstance(atom_name, str) and atom_name: + records.append((atom_name, atom.GetSymbol(), atom.GetFormalCharge())) + self.atoms[component_id] = records + result = self.atoms[component_id] + return result or None + + def bond_records(self, component_id: str) -> list[tuple[str, str]] | None: + if component_id not in self.bonds: + molecule, _conformer = self._component_with_conformer(component_id) + records: list[tuple[str, str]] = [] + if molecule is not None: + names = { + atom.GetIdx(): atom.GetPropsAsDict().get("name") for atom in molecule.GetAtoms() + } + for bond in molecule.GetBonds(): + first = names.get(bond.GetBeginAtomIdx()) + second = names.get(bond.GetEndAtomIdx()) + if isinstance(first, str) and first and isinstance(second, str) and second: + records.append((first, second)) + self.bonds[component_id] = records + result = self.bonds[component_id] + return result or None + + def component_leaving_atoms(self, component_id: str) -> set[str]: + if component_id not in self.leaving_atoms: + molecule = self.load().get(component_id) + names: set[str] = set() + if molecule is not None: + for atom in molecule.GetAtoms(): + if atom.HasProp("leaving_atom") and atom.GetProp("leaving_atom") == "1": + name = atom.GetProp("name") if atom.HasProp("name") else "" + if name: + names.add(name) + self.leaving_atoms[component_id] = names + return self.leaving_atoms[component_id] + + +_STORE = _ChemicalComponentStore() + + +def load_ccd(cache_dir: Path | str | None = None) -> dict[str, Any]: + """Load and cache the CCD asset, resolving it only when called.""" + + return _STORE.load(cache_dir) + + +def get_ccd_conformer(component_id: str) -> dict[str, np.ndarray] | None: + """Return the preferred heavy-atom conformer by atom name.""" + + return _STORE.conformer(component_id) + + +def get_idealized_atom_pos(res_type: int, atom_name: str) -> np.ndarray | None: + """Return one standard-residue atom position from the preferred conformer.""" + + key = (res_type, atom_name) + if key not in _STORE.standard_positions: + component_id = RES_TYPE_TO_CCD.get(res_type) + conformer = _STORE.conformer(component_id) if component_id is not None else None + _STORE.standard_positions[key] = None if conformer is None else conformer.get(atom_name) + return _STORE.standard_positions[key] + + +def get_ligand_idealized_atom_pos(residue_name: str, atom_name: str) -> np.ndarray | None: + """Return one ligand atom position from the preferred conformer.""" + + key = (residue_name, atom_name) + if key not in _STORE.ligand_positions: + conformer = _STORE.conformer(residue_name) + _STORE.ligand_positions[key] = None if conformer is None else conformer.get(atom_name) + return _STORE.ligand_positions[key] + + +def get_ligand_ccd_atoms_with_charges( + component_id: str, +) -> list[tuple[str, str, int]] | None: + """Return heavy-atom name, element, and formal-charge records.""" + + return _STORE.atom_records(component_id) + + +def get_ligand_ccd_bonds(component_id: str) -> list[tuple[str, str]] | None: + """Return bonds as component atom-name pairs.""" + + return _STORE.bond_records(component_id) + + +def get_ccd_leaving_atoms(component_id: str) -> set[str]: + """Return atoms removed when a CCD component is polymerized.""" + + return _STORE.component_leaving_atoms(component_id) + + +__all__ = [ + "get_ccd_conformer", + "get_ccd_leaving_atoms", + "get_idealized_atom_pos", + "get_ligand_ccd_atoms_with_charges", + "get_ligand_ccd_bonds", + "get_ligand_idealized_atom_pos", + "load_ccd", +] diff --git a/fastplms/models/esmfold2/esmfold2_constants.py b/fastplms/models/esmfold2/esmfold2_constants.py new file mode 100644 index 0000000000000000000000000000000000000000..df0709620fda4f5cdf91b69f6eff2d3cd61e9b19 --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_constants.py @@ -0,0 +1,156 @@ +"""Declarative molecular schema for ESMFold2 feature preparation. + +The package manifest owns the upstream revision and license provenance. This +module expresses the corresponding checkpoint-facing integer schema as compact +ordered records, then derives lookup tables from those records. The generated +tables are validated at import without reading files, downloading assets, or +mutating process state. +""" + +from __future__ import annotations + +SCHEMA_PROVENANCE = { + "manifest_family": "esmfold2", + "contract": "biohub_esmfold2_input_v1", +} + + +def _words(value: str) -> list[str]: + return value.split() + + +MOL_TYPE_PROTEIN = 0 +MOL_TYPE_DNA = 1 +MOL_TYPE_RNA = 2 +MOL_TYPE_NONPOLYMER = 3 + +# The record order is part of the checkpoint input contract. Residue indices +# start at two because zero and one are reserved by the model feature schema. +_PROTEIN_SCHEMA = tuple( + tuple(record.split(":")) + for record in ( + "ALA:A:N CA C O CB", + "ARG:R:N CA C O CB CG CD NE CZ NH1 NH2", + "ASN:N:N CA C O CB CG OD1 ND2", + "ASP:D:N CA C O CB CG OD1 OD2", + "CYS:C:N CA C O CB SG", + "GLN:Q:N CA C O CB CG CD OE1 NE2", + "GLU:E:N CA C O CB CG CD OE1 OE2", + "GLY:G:N CA C O", + "HIS:H:N CA C O CB CG ND1 CD2 CE1 NE2", + "ILE:I:N CA C O CB CG1 CG2 CD1", + "LEU:L:N CA C O CB CG CD1 CD2", + "LYS:K:N CA C O CB CG CD CE NZ", + "MET:M:N CA C O CB CG SD CE", + "PHE:F:N CA C O CB CG CD1 CD2 CE1 CE2 CZ", + "PRO:P:N CA C O CB CG CD", + "SER:S:N CA C O CB OG", + "THR:T:N CA C O CB OG1 CG2", + "TRP:W:N CA C O CB CG CD1 CD2 NE1 CE2 CE3 CZ2 CZ3 CH2", + "TYR:Y:N CA C O CB CG CD1 CD2 CE1 CE2 CZ OH", + "VAL:V:N CA C O CB CG1 CG2", + ) +) + +PROTEIN_RESIDUE_TO_RES_TYPE = { + residue: index for index, (residue, _letter, _atoms) in enumerate(_PROTEIN_SCHEMA, 2) +} +PROTEIN_RESIDUE_TO_RES_TYPE["MSE"] = PROTEIN_RESIDUE_TO_RES_TYPE["MET"] +PROTEIN_UNK_RES_TYPE = 22 + +RNA_RESIDUE_TO_RES_TYPE = dict(zip("AGCU", range(23, 27), strict=True)) +RNA_UNK_RES_TYPE = 27 +DNA_RESIDUE_TO_RES_TYPE = dict(zip(("DA", "DG", "DC", "DT"), range(28, 32), strict=True)) +DNA_UNK_RES_TYPE = 32 +GAP_RES_TYPE = DNA_UNK_RES_TYPE + +PROTEIN_3TO1 = {residue: letter for residue, letter, _atoms in _PROTEIN_SCHEMA} +PROTEIN_3TO1["MSE"] = "M" +PROTEIN_1TO3 = {letter: residue for residue, letter, _atoms in _PROTEIN_SCHEMA} +PROTEIN_1TO3["X"] = "UNK" +DNA_1TO3 = dict(zip("ATCG", ("DA", "DT", "DC", "DG"), strict=True)) +RNA_1TO3 = {letter: letter for letter in "AUCG"} + +_ESM_RESIDUE_ORDER = "LAGVSERTIDPKQNFYM HWC".replace(" ", "") +ESM_PROTEIN_VOCAB = {residue: token_id for token_id, residue in enumerate(_ESM_RESIDUE_ORDER, 4)} +ESM_PROTEIN_VOCAB["X"] = 3 +DNA_RNA_LIGAND_INPUT_ID = 24 +MSA_PAD_TOKEN_ID = 0 +MSA_GAP_TOKEN_ID = 1 + +RES_TYPE_TO_CCD = { + **{ + index: residue for residue, index in PROTEIN_RESIDUE_TO_RES_TYPE.items() if residue != "MSE" + }, + 22: "UNK", + **dict(zip(range(23, 28), ("A", "G", "C", "U", "N"), strict=True)), + **dict(zip(range(28, 33), ("DA", "DG", "DC", "DT", "DN"), strict=True)), +} + +_CHARGE_SCHEMA = _words( + "LYS:NZ:1 ARG:NH2:1 HIS:ND1:1 PO4:O2:-1 PO4:O3:-1 PO4:O4:-1 " + "SO4:O3:-1 SO4:O4:-1 MG:MG:2 ZN:ZN:2 CA:CA:2 FE2:FE:2 MN:MN:2 " + "CO:CO:2 NCO:CO:3 CU:CU:2 NI:NI:2 K:K:1 NA:NA:1 CD:CD:2 CL:CL:-1 " + "ACT:OXT:-1 NAD:O2N:-1 NAD:N1N:1 NAP:O2N:-1 NAP:N1N:1 IMD:N3:1 " + "SAM:SD:1 FE:FE:3 A1BH3:N3:1" +) +CHARGED_ATOMS = { + (component, atom): int(charge) + for component, atom, charge in (record.split(":") for record in _CHARGE_SCHEMA) +} + +_PERIODIC_SYMBOLS = _words( + "H HE LI BE B C N O F NE NA MG AL SI P S CL AR K CA SC TI V CR MN FE CO NI CU ZN " + "GA GE AS SE BR KR RB SR Y ZR NB MO TC RU RH PD AG CD IN SN SB TE I XE CS BA LA CE " + "PR ND PM SM EU GD TB DY HO ER TM YB LU HF TA W RE OS IR PT AU HG TL PB BI PO AT RN " + "FR RA AC TH PA U" +) +ELEMENT_TO_ATOMIC_NUM = { + symbol: atomic_number + for atomic_number, symbol in enumerate(_PERIODIC_SYMBOLS, 1) + if symbol != "HE" +} +ELEMENT_NUMBER_TO_SYMBOL = { + atomic_number: symbol for symbol, atomic_number in ELEMENT_TO_ATOMIC_NUM.items() +} + +PROTEIN_HEAVY_ATOMS = { + residue: atom_string.split() for residue, _letter, atom_string in _PROTEIN_SCHEMA +} +PROTEIN_HEAVY_ATOMS["MSE"] = PROTEIN_HEAVY_ATOMS["MET"].copy() +PROTEIN_HEAVY_ATOMS["UNK"] = _words("N CA C O") + +DNA_BACKBONE_ATOMS = _words("P OP1 OP2 O5' C5' C4' O4' C3' O3' C2' C1'") +RNA_BACKBONE_ATOMS = _words("P OP1 OP2 O5' C5' C4' O4' C3' O3' C2' O2' C1'") +_NUCLEOBASE_ATOMS = { + "A": _words("N9 C8 N7 C5 C6 N6 N1 C2 N3 C4"), + "G": _words("N9 C8 N7 C5 C6 O6 N1 C2 N2 N3 C4"), + "C": _words("N1 C2 O2 N3 C4 N4 C5 C6"), + "U": _words("N1 C2 O2 N3 C4 O4 C5 C6"), + "T": _words("N1 C2 O2 N3 C4 O4 C5 C7 C6"), +} +DNA_HEAVY_ATOMS = { + "DA": DNA_BACKBONE_ATOMS + _NUCLEOBASE_ATOMS["A"], + "DG": DNA_BACKBONE_ATOMS + _NUCLEOBASE_ATOMS["G"], + "DC": DNA_BACKBONE_ATOMS + _NUCLEOBASE_ATOMS["C"], + "DT": DNA_BACKBONE_ATOMS + _NUCLEOBASE_ATOMS["T"], +} +RNA_HEAVY_ATOMS = {residue: RNA_BACKBONE_ATOMS + _NUCLEOBASE_ATOMS[residue] for residue in "AGCU"} + + +def _validate_schema() -> None: + if sorted(set(PROTEIN_RESIDUE_TO_RES_TYPE.values())) != list(range(2, 22)): + raise RuntimeError("Protein residue indices must cover the checkpoint interval 2..21.") + if len(_ESM_RESIDUE_ORDER) != 20 or len(set(_ESM_RESIDUE_ORDER)) != 20: + raise RuntimeError("The ESM residue vocabulary must contain 20 canonical residues.") + if RES_TYPE_TO_CCD[14] != "MET" or PROTEIN_RESIDUE_TO_RES_TYPE["MSE"] != 14: + raise RuntimeError("Selenomethionine must share the methionine residue index.") + if ELEMENT_TO_ATOMIC_NUM.get("U") != 92 or 2 in ELEMENT_NUMBER_TO_SYMBOL: + raise RuntimeError("The element schema must preserve the training-time atomic-number map.") + if set(DNA_HEAVY_ATOMS) != {"DA", "DG", "DC", "DT"}: + raise RuntimeError("The DNA atom schema is incomplete.") + + +_validate_schema() + +__all__ = [name for name in globals() if name.isupper()] diff --git a/fastplms/models/esmfold2/esmfold2_constants_esm3.py b/fastplms/models/esmfold2/esmfold2_constants_esm3.py new file mode 100644 index 0000000000000000000000000000000000000000..4c044111bb9e5f18d11b207a3b9d2c0a75f83fc5 --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_constants_esm3.py @@ -0,0 +1,98 @@ +"""Token schemas needed by the ESMC encoder inside ESMFold2. + +The values implement the published Biohub ESM sequence-token contract pinned by +``models.toml``. They are generated from ordered schemas so token position and +special-token relationships are explicit and independently testable. This +module performs no downloads and resolves no model assets at import time. +""" + +from __future__ import annotations + +from types import MappingProxyType + + +def _words(value: str) -> list[str]: + return value.split() + + +SEQUENCE_VOCAB = _words( + " L A G V S E R T I D P K Q N F Y M H W C X B U Z O . - | " +) + +_sequence_token_ids = MappingProxyType({token: index for index, token in enumerate(SEQUENCE_VOCAB)}) +SEQUENCE_BOS_TOKEN = _sequence_token_ids[""] +SEQUENCE_PAD_TOKEN = _sequence_token_ids[""] +SEQUENCE_EOS_TOKEN = _sequence_token_ids[""] +SEQUENCE_CHAINBREAK_TOKEN = _sequence_token_ids["|"] +SEQUENCE_MASK_TOKEN = _sequence_token_ids[""] +SEQUENCE_STANDARD_AA_MIN_TOKEN = _sequence_token_ids["L"] +SEQUENCE_STANDARD_AA_MAX_TOKEN = _sequence_token_ids["X"] + +VQVAE_CODEBOOK_SIZE = 4096 +VQVAE_SPECIAL_TOKENS = { + name: VQVAE_CODEBOOK_SIZE + offset + for offset, name in enumerate(("MASK", "EOS", "BOS", "PAD", "CHAINBREAK")) +} +VQVAE_DIRECTION_LOSS_BINS = 16 +VQVAE_PAE_BINS = 64 +VQVAE_MAX_PAE_BIN = 31.0 +VQVAE_PLDDT_BINS = 50 + +STRUCTURE_MASK_TOKEN = VQVAE_SPECIAL_TOKENS["MASK"] +STRUCTURE_EOS_TOKEN = VQVAE_SPECIAL_TOKENS["EOS"] +STRUCTURE_BOS_TOKEN = VQVAE_SPECIAL_TOKENS["BOS"] +STRUCTURE_PAD_TOKEN = VQVAE_SPECIAL_TOKENS["PAD"] +STRUCTURE_CHAINBREAK_TOKEN = VQVAE_SPECIAL_TOKENS["CHAINBREAK"] +STRUCTURE_UNDEFINED_TOKEN = 955 + +SASA_PAD_TOKEN = 0 +SS8_PAD_TOKEN = 0 +INTERPRO_PAD_TOKEN = 0 +RESIDUE_PAD_TOKEN = 0 + +CHAIN_BREAK_STR = "|" +SEQUENCE_BOS_STR = "" +SEQUENCE_EOS_STR = "" +MASK_STR_SHORT = "_" +SEQUENCE_MASK_STR = "" +SASA_MASK_STR = "" +SS8_MASK_STR = "" + +SSE_8CLASS_VOCAB = "GHITEBSC" +SSE_3CLASS_VOCAB = "HEC" +SSE_8CLASS_TO_3CLASS_MAP = dict(zip(SSE_8CLASS_VOCAB, "HHHCEECC", strict=True)) + +SASA_DISCRETIZATION_BOUNDARIES = [ + 0.8, + 4.0, + 9.6, + 16.4, + 24.5, + 32.9, + 42.0, + 51.5, + 61.2, + 70.9, + 81.6, + 93.3, + 107.2, + 125.4, + 151.4, +] +MAX_RESIDUE_ANNOTATIONS = 16 +TFIDF_VECTOR_SIZE = 58_641 +FUNCTION_TOKENS_DEPTH = 8 + + +def _validate_schema() -> None: + if len(SEQUENCE_VOCAB) != len(set(SEQUENCE_VOCAB)): + raise RuntimeError("The ESM sequence vocabulary contains duplicate tokens.") + if SEQUENCE_STANDARD_AA_MAX_TOKEN - SEQUENCE_STANDARD_AA_MIN_TOKEN != 20: + raise RuntimeError("The canonical residue interval must contain 20 tokens.") + if tuple(VQVAE_SPECIAL_TOKENS.values()) != tuple(range(4096, 4101)): + raise RuntimeError("The structure special-token interval is not contiguous.") + + +_validate_schema() + +__all__ = [name for name in globals() if name.isupper()] diff --git a/fastplms/models/esmfold2/esmfold2_input_builder.py b/fastplms/models/esmfold2/esmfold2_input_builder.py new file mode 100644 index 0000000000000000000000000000000000000000..ca51bc2c91fd4c1e4139f7bcca71898e38e23065 --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_input_builder.py @@ -0,0 +1,244 @@ +"""Typed, JSON-safe inputs for ESMFold2 feature preparation.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any, TypeAlias + +import numpy as np + +from .esmfold2_msa import MSA + +MSAInput: TypeAlias = MSA | None + + +@dataclass +class Modification: + """A zero-indexed residue substitution using a CCD component.""" + + position: int + ccd: str + smiles: str | None = None + + +@dataclass +class ProteinInput: + id: str | list[str] + sequence: str + modifications: list[Modification] | None = None + msa: MSAInput = None + + +@dataclass +class RNAInput: + id: str | list[str] + sequence: str + modifications: list[Modification] | None = None + + +@dataclass +class DNAInput: + id: str | list[str] + sequence: str + modifications: list[Modification] | None = None + + +@dataclass +class LigandInput: + id: str | list[str] + smiles: str | None = None + ccd: list[str] | None = None + + +@dataclass +class DistogramConditioning: + chain_id: str + distogram: np.ndarray + + +@dataclass +class PocketConditioning: + binder_chain_id: str + contacts: list[tuple[str, int]] + + +@dataclass +class CovalentBond: + chain_id1: str + res_idx1: int + atom_idx1: int + chain_id2: str + res_idx2: int + atom_idx2: int + + +SequenceInput: TypeAlias = ProteinInput | RNAInput | DNAInput | LigandInput + + +@dataclass +class StructurePredictionInput: + sequences: Sequence[SequenceInput] + pocket: PocketConditioning | None = None + distogram_conditioning: list[DistogramConditioning] | None = None + covalent_bonds: list[CovalentBond] | None = None + + +_CHAIN_TYPE = { + ProteinInput: "protein", + RNAInput: "rna", + DNAInput: "dna", +} + + +def _serialize_modifications( + modifications: list[Modification] | None, +) -> list[dict[str, Any]] | None: + if not modifications: + return None + return [{"position": item.position, "ccd": item.ccd} for item in modifications] + + +def _serialize_chain(chain: SequenceInput) -> dict[str, Any]: + if isinstance(chain, LigandInput): + return { + "smiles": chain.smiles, + "id": chain.id, + "ccd": chain.ccd, + "type": "ligand", + } + + chain_type = _CHAIN_TYPE.get(type(chain)) + if chain_type is None: + raise ValueError(f"Unsupported sequence input type: {type(chain)}") + serialized: dict[str, Any] = { + "sequence": chain.sequence, + "id": chain.id, + "type": chain_type, + } + if modifications := _serialize_modifications(chain.modifications): + serialized["modifications"] = modifications + if isinstance(chain, ProteinInput): + if chain.msa is not None and not isinstance(chain.msa, MSA): + raise AttributeError(f"MSA must be None or MSA. Got {chain.msa} instead.") + serialized["msa"] = None if chain.msa is None else {"sequences": chain.msa.sequences} + return serialized + + +def serialize_structure_prediction_input( + structure_input: StructurePredictionInput, +) -> dict[str, Any]: + """Convert an input object to a JSON-safe mapping.""" + + serialized: dict[str, Any] = { + "sequences": [_serialize_chain(chain) for chain in structure_input.sequences] + } + if structure_input.covalent_bonds is not None: + serialized["covalent_bonds"] = [ + vars(bond).copy() for bond in structure_input.covalent_bonds + ] + if structure_input.pocket is not None: + serialized["pocket"] = { + "binder_chain_id": structure_input.pocket.binder_chain_id, + "contacts": structure_input.pocket.contacts, + } + if structure_input.distogram_conditioning is not None: + serialized["distogram_conditioning"] = [ + {"chain_id": item.chain_id, "distogram": item.distogram.tolist()} + for item in structure_input.distogram_conditioning + ] + return serialized + + +def _deserialize_modifications(chain: dict[str, Any]) -> list[Modification] | None: + raw = chain.get("modifications") + if not raw: + return None + return [Modification(position=item["position"], ccd=item["ccd"]) for item in raw] + + +def _deserialize_msa(chain: dict[str, Any]) -> MSAInput: + raw = chain.get("msa") + if raw is None: + return None + if not isinstance(raw, dict) or not isinstance(raw.get("sequences"), list): + raise ValueError(f"Unexpected MSA value: {raw!r}") + return MSA.from_sequences(raw["sequences"]) + + +def _deserialize_chain(chain: dict[str, Any]) -> SequenceInput: + chain_type = chain.get("type") + common = {"id": chain["id"]} + if chain_type == "protein": + return ProteinInput( + **common, + sequence=chain["sequence"], + modifications=_deserialize_modifications(chain), + msa=_deserialize_msa(chain), + ) + if chain_type == "rna": + return RNAInput( + **common, + sequence=chain["sequence"], + modifications=_deserialize_modifications(chain), + ) + if chain_type == "dna": + return DNAInput( + **common, + sequence=chain["sequence"], + modifications=_deserialize_modifications(chain), + ) + if chain_type == "ligand": + return LigandInput(**common, smiles=chain.get("smiles"), ccd=chain.get("ccd")) + raise ValueError(f"Unsupported sequence type: {chain_type!r}") + + +def deserialize_structure_prediction_input(data: dict[str, Any]) -> StructurePredictionInput: + """Reconstruct the typed input represented by a serialized mapping.""" + + pocket_data = data.get("pocket") + pocket = None + if pocket_data is not None: + pocket = PocketConditioning( + binder_chain_id=pocket_data["binder_chain_id"], + contacts=[tuple(contact) for contact in pocket_data["contacts"]], + ) + + distogram_data = data.get("distogram_conditioning") + distograms = None + if distogram_data is not None: + distograms = [ + DistogramConditioning( + chain_id=item["chain_id"], distogram=np.asarray(item["distogram"]) + ) + for item in distogram_data + ] + + bond_data = data.get("covalent_bonds") + bonds = None + if bond_data is not None: + bonds = [CovalentBond(**item) for item in bond_data] + + return StructurePredictionInput( + sequences=[_deserialize_chain(chain) for chain in data["sequences"]], + pocket=pocket, + distogram_conditioning=distograms, + covalent_bonds=bonds, + ) + + +__all__ = [ + "CovalentBond", + "DNAInput", + "DistogramConditioning", + "LigandInput", + "MSAInput", + "Modification", + "PocketConditioning", + "ProteinInput", + "RNAInput", + "SequenceInput", + "StructurePredictionInput", + "deserialize_structure_prediction_input", + "serialize_structure_prediction_input", +] diff --git a/fastplms/models/esmfold2/esmfold2_metrics.py b/fastplms/models/esmfold2/esmfold2_metrics.py new file mode 100644 index 0000000000000000000000000000000000000000..fcd665167967a7524b7de15bff5afcc9e4c63bf6 --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_metrics.py @@ -0,0 +1,235 @@ +"""Contact, lDDT, RMSD, and GDT-TS metrics for structure validation.""" + +from __future__ import annotations + +import numpy as np +import torch +import torch.nn.functional as F +from torch import Tensor +from torch.amp import autocast # type: ignore + +from . import esmfold2_residue_constants as residue_constants +from .esmfold2_misc import binpack, unbinpack +from .esmfold2_protein_structure import ( + compute_alignment_tensors, + compute_gdt_ts_no_alignment, + compute_rmsd_no_alignment, +) + + +def _distance_matrix(positions: Tensor, eps: float) -> Tensor: + displacement = positions[..., None, :] - positions[..., None, :, :] + return torch.sqrt(eps + torch.sum(displacement**2, dim=-1)) + + +def compute_lddt_from_dmat( + dmat_pred: Tensor, + dmat_true: Tensor, + pairwise_mask: Tensor, + cutoff: float | Tensor = 15.0, + eps: float = 1e-10, + per_residue: bool = True, +) -> Tensor: + """Score distance matrices ``D_pred`` and ``D_true`` with shape (..., l, l).""" + + sequence_length = dmat_true.size(-1) + identity = torch.eye(sequence_length, device=dmat_true.device) + scored_pairs = (dmat_true < cutoff) * pairwise_mask * (1.0 - identity) + absolute_error = torch.abs(dmat_true - dmat_pred) + score = ( + (absolute_error < 0.5).type(absolute_error.dtype) + + (absolute_error < 1.0).type(absolute_error.dtype) + + (absolute_error < 2.0).type(absolute_error.dtype) + + (absolute_error < 4.0).type(absolute_error.dtype) + ) * 0.25 + dimensions = (-1,) if per_residue else (-2, -1) + normalization = 1.0 / (eps + scored_pairs.sum(dim=dimensions)) + return normalization * (eps + (scored_pairs * score).sum(dim=dimensions)) + + +def compute_lddt( + all_atom_pred_pos: Tensor, + all_atom_positions: Tensor, + all_atom_mask: Tensor, + pairwise_all_atom_mask: Tensor | None = None, + cutoff: float | Tensor = 15.0, + eps: float = 1e-10, + per_residue: bool = True, + sequence_id: Tensor | None = None, +) -> Tensor: + """Compute lDDT from coordinate tensors and atom masks.""" + + expanded_mask = all_atom_mask[..., None] + true_distances = _distance_matrix(all_atom_positions, eps) + predicted_distances = _distance_matrix(all_atom_pred_pos, eps) + pair_mask = expanded_mask * expanded_mask.transpose(-2, -1) + if pairwise_all_atom_mask is not None: + pair_mask = pair_mask * pairwise_all_atom_mask + if sequence_id is not None: + same_sequence = sequence_id[..., None] == sequence_id[..., None, :] + pair_mask = pair_mask * same_sequence.type_as(pair_mask) + return compute_lddt_from_dmat( + predicted_distances, + true_distances, + pair_mask, + cutoff=cutoff, + eps=eps, + per_residue=per_residue, + ) + + +def compute_lddt_ca( + all_atom_pred_pos: Tensor, + all_atom_positions: Tensor, + all_atom_mask: Tensor, + cutoff: float = 15.0, + eps: float = 1e-10, + per_residue: bool = True, + sequence_id: Tensor | None = None, +) -> Tensor: + """Compute lDDT using only C-alpha coordinates.""" + + ca_index = residue_constants.atom_order["CA"] + predicted_ca = ( + all_atom_pred_pos if all_atom_pred_pos.dim() == 3 else all_atom_pred_pos[..., ca_index, :] + ) + return compute_lddt( + predicted_ca, + all_atom_positions[..., ca_index, :], + all_atom_mask[..., ca_index], + cutoff=cutoff, + eps=eps, + per_residue=per_residue, + sequence_id=sequence_id, + ) + + +@torch.no_grad() +@autocast("cuda", enabled=False) +def compute_rmsd( + mobile: Tensor, + target: Tensor, + atom_exists_mask: Tensor | None = None, + sequence_id: Tensor | None = None, + reduction: str = "batch", +) -> Tensor: + """Align ``X`` to ``Y`` and compute RMSD.""" + + centered_mobile, _, centered_target, _, rotation, counts = compute_alignment_tensors( + mobile, + target, + atom_exists_mask, + sequence_id, + ) + rmsd = compute_rmsd_no_alignment( + torch.matmul(centered_mobile, rotation), + centered_target, + counts, + reduction=reduction, + ) + if reduction == "per_residue" and sequence_id is not None: + return binpack(rmsd, sequence_id, pad_value=0) + return rmsd + + +def compute_gdt_ts( + mobile: Tensor, + target: Tensor, + atom_exists_mask: Tensor | None = None, + sequence_id: Tensor | None = None, + reduction: str = "per_sample", +) -> Tensor: + """Align ``X`` to ``Y`` and compute GDT-TS.""" + + if atom_exists_mask is None: + atom_exists_mask = torch.isfinite(target).all(dim=-1) + centered_mobile, _, centered_target, _, rotation, _ = compute_alignment_tensors( + mobile, + target, + atom_exists_mask, + sequence_id, + ) + if sequence_id is not None: + atom_exists_mask = unbinpack(atom_exists_mask, sequence_id, pad_value=False) + return compute_gdt_ts_no_alignment( + torch.matmul(centered_mobile, rotation), + centered_target, + atom_exists_mask, + reduction, + ) + + +def _batched_contacts(predictions: Tensor, targets: Tensor) -> tuple[Tensor, Tensor]: + if predictions.dim() == 2: + predictions = predictions.unsqueeze(0) + if targets.dim() == 2: + targets = targets.unsqueeze(0) + if predictions.size() != targets.size(): + raise ValueError( + f"Size mismatch. Received predictions of size {predictions.size()}, " + f"targets of size {targets.size()}" + ) + return predictions, targets + + +def _valid_contact_mask( + targets: Tensor, + src_lengths: Tensor, + minsep: int, + maxsep: int | None, +) -> Tensor: + sequence_length = targets.shape[-1] + positions = torch.arange(sequence_length, device=targets.device) + separation = (positions.unsqueeze(0) - positions.unsqueeze(1)).unsqueeze(0) + valid = (separation >= minsep) & (targets >= 0) + if maxsep is not None: + valid &= separation < maxsep + within_length = positions.unsqueeze(0) < src_lengths.unsqueeze(1) + return valid & within_length.unsqueeze(1) & within_length.unsqueeze(2) + + +def contact_precision( + predictions: Tensor, + targets: Tensor, + src_lengths: Tensor | None = None, + minsep: int = 6, + maxsep: int | None = None, + override_length: int | None = None, +) -> dict[str, Tensor]: + """Compute P@L, P@L/5, and binned area for contact probabilities.""" + + predictions, targets = _batched_contacts(predictions, targets) + batch_size, sequence_length, _ = predictions.shape + if src_lengths is None: + src_lengths = torch.full( + (batch_size,), + sequence_length, + dtype=torch.long, + device=predictions.device, + ) + valid = _valid_contact_mask(targets, src_lengths, minsep, maxsep) + masked_predictions = predictions.masked_fill(~valid, float("-inf")) + row_index, column_index = np.triu_indices(sequence_length, minsep) + upper_predictions = masked_predictions[:, row_index, column_index] + upper_targets = targets[:, row_index, column_index] + + topk = sequence_length if override_length is None else max(sequence_length, override_length) + ranked_indices = upper_predictions.argsort(dim=-1, descending=True)[:, :topk] + batch_indices = torch.arange(batch_size, device=ranked_indices.device).unsqueeze(1) + ranked_targets = upper_targets[batch_indices, ranked_indices] + if ranked_targets.size(1) < topk: + ranked_targets = F.pad(ranked_targets, [0, topk - ranked_targets.size(1)]) + cumulative_contacts = ranked_targets.type_as(predictions).cumsum(dim=-1) + + gather_lengths = src_lengths.unsqueeze(1) + if override_length is not None: + gather_lengths = override_length * torch.ones_like(gather_lengths) + fractions = torch.arange(0.1, 1.1, 0.1, device=predictions.device).unsqueeze(0) + gather_indices = (fractions * gather_lengths).type(torch.long).sub(1).clamp_min(0) + cumulative_bins = cumulative_contacts.gather(1, gather_indices) + precisions = cumulative_bins / (gather_indices + 1).type_as(cumulative_bins) + return { + "AUC": precisions.mean(dim=-1), + "P@L": precisions[:, 9], + "P@L5": precisions[:, 1], + } diff --git a/fastplms/models/esmfold2/esmfold2_misc.py b/fastplms/models/esmfold2/esmfold2_misc.py new file mode 100644 index 0000000000000000000000000000000000000000..022edaf3916ab22b21363018e3cbf54963440c0e --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_misc.py @@ -0,0 +1,400 @@ +"""Small tensor, sequence, and annotation utilities used by ESMFold2. + +The helpers in this module are deliberately free of model state. Importing the +module therefore performs no device selection, compilation, or remote access. +""" + +from __future__ import annotations + +from collections import defaultdict +from collections.abc import Generator, Iterable, Sequence +from contextlib import AbstractContextManager, nullcontext +from dataclasses import is_dataclass +from io import BytesIO +from typing import Any, Protocol, TypeVar, runtime_checkable +from warnings import warn + +import numpy as np +import torch +import zstandard + +from .esmfold2_constants_esm3 import CHAIN_BREAK_STR +from .esmfold2_utils_types import FunctionAnnotation + +MAX_SUPPORTED_DISTANCE = 1e6 + +TSequence = TypeVar("TSequence", bound=Sequence) + + +@runtime_checkable +class Concatable(Protocol): + """Protocol for sequence-like records with a class-level concatenator.""" + + @classmethod + def concat(cls, objs: list[Concatable]) -> Concatable: ... + + +def fp32_autocast_context( + device_type: str, +) -> AbstractContextManager[Any]: # type: ignore + """Return a context that keeps numerically sensitive work in FP32.""" + + if device_type == "mps": + return nullcontext() + if device_type == "cpu": + return torch.amp.autocast(device_type, enabled=False) # type: ignore + if device_type == "cuda": + return torch.amp.autocast(device_type, dtype=torch.float32) # type: ignore + raise ValueError(f"Unsupported device type: {device_type}") + + +def maybe_tensor(value, convert_none_to_nan: bool = False) -> torch.Tensor | None: + """Convert an optional array-like value to a tensor.""" + + if value is None: + return None + if isinstance(value, torch.Tensor): + return value + if isinstance(value, list) and all(isinstance(element, torch.Tensor) for element in value): + return torch.stack(value) + if convert_none_to_nan: + value = np.asarray(value, dtype=np.float32) + value = np.where(value is None, np.nan, value) + return torch.tensor(value) + + +def maybe_list(value, convert_nan_to_none: bool = False) -> list | None: + """Convert an optional tensor or NumPy array to nested Python lists.""" + + if value is None: + return None + if not convert_nan_to_none: + return value.tolist() + if isinstance(value, torch.Tensor): + nan_mask = torch.isnan(value).cpu().numpy() + array = value.cpu().numpy().astype(object) + elif isinstance(value, np.ndarray): + nan_mask = np.isnan(value) + array = value.astype(object) + else: + raise TypeError("maybe_list can only work with torch.tensor or np.ndarray.") + array[nan_mask] = None + return array.tolist() + + +def replace_inf(data): + """Replace infinite array values by the ESM API sentinel value.""" + + if data is None: + return None + array = np.asarray(data, dtype=np.float32) + return np.where(np.isinf(array), 1000, array).tolist() + + +def slice_python_object_as_numpy( + obj: TSequence, + idx: int | list[int] | slice | np.ndarray, +) -> TSequence: + """Apply NumPy-style scalar, mask, or index-array slicing to Python data.""" + + normalized_idx: list[int] | slice | np.ndarray = ( + [int(idx)] if np.isscalar(idx) else idx # type: ignore[arg-type] + ) + + if isinstance(normalized_idx, np.ndarray) and normalized_idx.dtype == bool: + selected = [obj[position] for position in np.flatnonzero(normalized_idx)] + elif isinstance(normalized_idx, slice): + selected = obj[normalized_idx] + else: + selected = [obj[position] for position in normalized_idx] + + if isinstance(obj, str) and isinstance(selected, list): + return "".join(selected) # type: ignore[return-value] + return obj.__class__(selected) # type: ignore[call-arg,return-value] + + +def slice_any_object( + obj: TSequence, + idx: int | list[int] | slice | np.ndarray, +) -> TSequence: + """Slice tensors, arrays, dataclasses, and ordinary Python sequences.""" + + if isinstance(obj, (np.ndarray, torch.Tensor)) or is_dataclass(obj): + return obj[idx] # type: ignore[index,return-value] + return slice_python_object_as_numpy(obj, idx) + + +def join_lists( + lists: Sequence[Sequence[Any]], + separator: Sequence[Any] | None = None, +) -> list[Any]: + """Join lists, inserting all elements of ``separator`` between inputs.""" + + if len(lists) == 0: + return [] + joined = list(lists[0]) + for values in lists[1:]: + if separator: + joined.extend(separator) + joined.extend(values) + return joined + + +def iterate_with_intermediate( + lists: Iterable, + intermediate, +) -> Generator[Any, None, None]: + """Yield an intermediate value between consecutive input values.""" + + iterator = iter(lists) + yield next(iterator) + for value in iterator: + yield intermediate + yield value + + +def concat_objects(objs: Sequence[Any], separator: Any | None = None): + """Concatenate one supported homogeneous collection.""" + + if not objs: + raise ValueError("objs must contain at least one value.") + first = objs[0] + if isinstance(first, Concatable): + return first.__class__.concat(objs) + if isinstance(first, str): + if not isinstance(separator, str): + raise TypeError("separator must be a string when joining strings.") + return separator.join(objs) + if isinstance(first, list): + return join_lists(objs, None if separator is None else [separator]) + if isinstance(first, np.ndarray): + pieces = ( + objs + if separator is None + else list(iterate_with_intermediate(objs, np.array([separator]))) + ) + return np.concatenate(pieces) + if isinstance(first, torch.Tensor): + pieces = ( + objs + if separator is None + else list(iterate_with_intermediate(objs, torch.tensor([separator]))) + ) + return torch.cat(pieces) # type: ignore[arg-type] + raise TypeError(type(first)) + + +def rbf(values, v_min, v_max, n_bins=16): + """Encode values against evenly spaced radial basis centers.""" + + centers = torch.linspace( + v_min, + v_max, + n_bins, + dtype=values.dtype, + device=values.device, + ) + centers = centers.reshape((1,) * values.ndim + (-1,)) + standardized = (values.unsqueeze(-1) - centers) / ((v_max - v_min) / n_bins) + return torch.exp(-(standardized**2)) + + +def batched_gather(data, inds, dim=0, no_batch_dims=0): + """Gather along one data dimension while retaining leading batch axes.""" + + batch_indices = [] + index_rank = len(inds.shape) + for axis, size in enumerate(data.shape[:no_batch_dims]): + shape = (1,) * axis + (-1,) + (1,) * (index_rank - axis - 1) + batch_indices.append(torch.arange(size).view(*shape)) + tail = [slice(None)] * (len(data.shape) - no_batch_dims) + tail[dim - no_batch_dims if dim >= 0 else dim] = inds + return data[tuple(batch_indices + tail)] + + +def node_gather(s: torch.Tensor, edges: torch.Tensor) -> torch.Tensor: + """Gather node features for each row of an edge-index tensor.""" + + return batched_gather( + s.unsqueeze(-3), + edges, + -2, + no_batch_dims=len(s.shape) - 1, + ) + + +def knn_graph( + coords: torch.Tensor, + coord_mask: torch.Tensor, + padding_mask: torch.Tensor, + sequence_id: torch.Tensor, + *, + no_knn: int, +): + """Build nearest-neighbor edges, using sequence distance for missing geometry.""" + + length = coords.shape[-2] + coords = coords.nan_to_num() + missing_pair = ~(coord_mask[..., None, :] & coord_mask[..., :, None]) + excluded_pair = padding_mask[..., None, :] | padding_mask[..., :, None] + if sequence_id is not None: + excluded_pair |= sequence_id.unsqueeze(1) != sequence_id.unsqueeze(2) + + distances = (coords.unsqueeze(-2) - coords.unsqueeze(-3)).norm(dim=-1) + residue_index = torch.arange(length, device=coords.device) + sequence_distance = (residue_index.unsqueeze(-1) - residue_index.unsqueeze(-2)).abs() + if not (distances[~missing_pair] < MAX_SUPPORTED_DISTANCE).all(): + raise ValueError( + "Coordinate pairwise distances exceed max supported distance " + f"({MAX_SUPPORTED_DISTANCE}). " + ) + + rank_distance = sequence_distance.to(distances.dtype).mul(1e2).add(MAX_SUPPORTED_DISTANCE) + rank_distance = rank_distance.where(missing_pair, distances) + rank_distance = rank_distance.masked_fill(excluded_pair, torch.inf) + sorted_distance, sorted_edge = rank_distance.sort(dim=-1, descending=False) + width = min(no_knn, length) + return sorted_edge[..., :width], sorted_distance[..., :width].isfinite() + + +def stack_variable_length_tensors( + sequences: Sequence[torch.Tensor], + constant_value: int | float = 0, + dtype: torch.dtype | None = None, +) -> torch.Tensor: + """Pad arbitrary tensor dimensions to their maxima, then stack.""" + + output_shape = [ + len(sequences), + *np.max([sequence.shape for sequence in sequences], axis=0).tolist(), + ] + output = torch.full( + output_shape, + constant_value, + dtype=sequences[0].dtype if dtype is None else dtype, + device=sequences[0].device, + ) + for destination, source in zip(output, sequences, strict=True): + destination[tuple(slice(size) for size in source.shape)] = source + return output + + +def binpack( + tensor: torch.Tensor, + sequence_id: torch.Tensor | None, + pad_value: int | float, +): + """Scatter a sequence-major tensor into the packed layout described by IDs.""" + + if sequence_id is None: + return tensor + sequence_counts = sequence_id.max(dim=-1).values + 1 + output = torch.full( + sequence_id.shape + tensor.shape[2:], + fill_value=pad_value, + dtype=tensor.dtype, + device=tensor.device, + ) + source_index = 0 + for batch_index, (batch_ids, count) in enumerate( + zip(sequence_id, sequence_counts, strict=True) + ): + for seqid in range(count): + selection = batch_ids == seqid + output[batch_index, selection] = tensor[source_index, : selection.sum()] + source_index += 1 + return output + + +def unbinpack( + tensor: torch.Tensor, + sequence_id: torch.Tensor | None, + pad_value: int | float, +): + """Restore sequence-major rows from a packed tensor and its sequence IDs.""" + + if sequence_id is None: + return tensor + rows = [] + sequence_counts = sequence_id.max(dim=-1).values + 1 + for batch_index, (batch_ids, count) in enumerate( + zip(sequence_id, sequence_counts, strict=True) + ): + for seqid in range(count): + rows.append(tensor[batch_index, batch_ids == seqid]) + return stack_variable_length_tensors(rows, pad_value) + + +def merge_ranges( + ranges: list[range], + merge_gap_max: int | None = None, +) -> list[range]: + """Merge overlapping or sufficiently close ranges in positional order.""" + + maximum_gap = 0 if merge_gap_max is None else merge_gap_max + if not isinstance(maximum_gap, int) or isinstance(maximum_gap, bool): + raise TypeError("merge_gap_max must be an integer or None.") + if maximum_gap < 0: + raise ValueError(f"merge_gap_max must be non-negative, got {maximum_gap}.") + merged: list[range] = [] + for current in sorted(ranges, key=lambda item: item.start): + if not merged or merged[-1].stop + maximum_gap < current.start: + merged.append(current) + continue + previous = merged[-1] + merged[-1] = range(previous.start, max(previous.stop, current.stop)) + return merged + + +def merge_annotations( + annotations: list[FunctionAnnotation], + merge_gap_max: int | None = None, +) -> list[FunctionAnnotation]: + """Merge overlapping annotations independently for each label.""" + + grouped: dict[str, list[range]] = defaultdict(list) + for annotation in annotations: + grouped[annotation.label].append(range(annotation.start, annotation.end + 1)) + result = [] + for label, spans in grouped.items(): + result.extend( + FunctionAnnotation(label=label, start=span.start, end=span.stop - 1) + for span in merge_ranges(spans, merge_gap_max=merge_gap_max) + ) + return result + + +def get_chainbreak_boundaries_from_sequence( + sequence: Sequence[str], +) -> np.ndarray: + """Return half-open chain intervals split by chain-break tokens.""" + + boundaries = [0] + final_index = len(sequence) - 1 + for index, residue in enumerate(sequence): + if residue != CHAIN_BREAK_STR: + continue + if index == final_index: + raise ValueError( + "Encountered chain break token at end of sequence, this is unexpected." + ) + if index == final_index - 1: + warn( + "Encountered chain break token at penultimate position, this is unexpected.", + stacklevel=2, + ) + boundaries.extend((index, index + 1)) + boundaries.append(len(sequence)) + assert len(boundaries) % 2 == 0 + return np.asarray(boundaries).reshape(-1, 2) + + +def deserialize_tensors(data: bytes) -> Any: + """Decompress a tensor-only Torch payload onto CPU.""" + + decompressed = zstandard.ZstdDecompressor().decompress(data) + return torch.load( + BytesIO(decompressed), + map_location="cpu", + weights_only=True, + ) diff --git a/fastplms/models/esmfold2/esmfold2_mmcif_parsing.py b/fastplms/models/esmfold2/esmfold2_mmcif_parsing.py new file mode 100644 index 0000000000000000000000000000000000000000..caa4a04d76bf48212f6476008d360640ec6202f4 --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_mmcif_parsing.py @@ -0,0 +1,469 @@ +"""Biotite-backed mmCIF parsing used by ESMFold2 structure records.""" + +from __future__ import annotations + +import functools +import io +import os +from contextlib import suppress +from dataclasses import dataclass +from datetime import datetime + +import biotite.structure as bs +import biotite.structure.io.pdbx as pdbx +import numpy as np +from biotite.structure.io.pdbx import CIFColumn, CIFData, CIFFile + +from . import esmfold2_residue_constants as residue_constants + +PathOrBuffer = str | os.PathLike | io.StringIO + +PLDDT_B_FACTOR_SCALE = 100.0 +_MMCIF_COLUMN_DECIMALS = { + "Cartn_x": 3, + "Cartn_y": 3, + "Cartn_z": 3, + "B_iso_or_equiv": 2, +} +_NONPOLYMER_ENTITY_TYPES = frozenset({"NON-POLYMER", "WATER", "BRANCHED"}) + + +class NoProteinError(Exception): + """Raised internally when an mmCIF block contains no model-one atoms.""" + + +@dataclass +class Residue: + residue_number: int | None = None + insertion_code: str = "" + hetflag: bool = False + + +@dataclass +class MmcifHeader: + release_date: datetime | None = None + resolution: float | None = None + structure_method: str = "UNKNOWN" + + +def round_mmcif_columns(cif_file: CIFFile) -> None: + """Round coordinate and confidence columns in place for stable exports.""" + + if "atom_site" not in cif_file.block: + return + atom_site = cif_file.block["atom_site"] + for name, decimals in _MMCIF_COLUMN_DECIMALS.items(): + if name not in atom_site: + continue + original = atom_site[name] + values = original.as_array(np.float64) + strings = np.asarray( + [f"{value:.{decimals}f}" for value in values], + dtype=np.str_, + ) + atom_site[name] = CIFColumn( + data=CIFData(array=strings, dtype=np.str_), + mask=original.mask, + ) + + +def _clean_chain_list(value: str) -> list[str]: + return [chain.strip() for chain in value.split(",") if chain.strip()] + + +def _empty_residue() -> Residue: + return Residue(residue_number=None, insertion_code="", hetflag=False) + + +def _header_from_block( + block, + header: MmcifHeader | None = None, +) -> MmcifHeader: + header = MmcifHeader() if header is None else header + try: + if "pdbx_database_status" in block: + category = block["pdbx_database_status"] + if "recvd_initial_deposition_date" in category: + value = category["recvd_initial_deposition_date"].as_item() + if value and value != "?": + with suppress(ValueError): + header.release_date = datetime.strptime(value, "%Y-%m-%d") + if "refine" in block: + category = block["refine"] + if "ls_d_res_high" in category: + value = category["ls_d_res_high"].as_item() + if value and value != "?": + with suppress(ValueError): + header.resolution = float(value) + if "exptl" in block: + category = block["exptl"] + if "method" in category: + value = category["method"].as_item() + if value and value != "?": + header.structure_method = value.upper() + except Exception: + pass + return header + + +def _entities_from_block( + block, + entities: dict[int, list[str]] | None = None, +) -> dict[int, list[str]]: + entities = {} if entities is None else entities + if "entity" in block: + category = block["entity"] + ids = category["id"].as_array(str) + types = category["type"].as_array(str) + for entity_id, _ in zip(ids, types, strict=False): + entities[int(entity_id)] = [] + if "entity_poly" in block: + category = block["entity_poly"] + ids = category["entity_id"].as_array(str) + chain_lists = category["pdbx_strand_id"].as_array(str) + for raw_id, raw_chains in zip(ids, chain_lists, strict=False): + entity_id = int(raw_id) + if entity_id in entities: + entities[entity_id] = _clean_chain_list(raw_chains) + if "struct_asym" in block: + category = block["struct_asym"] + asym_ids = category["id"].as_array(str) + entity_ids = category["entity_id"].as_array(str) + for asym_id, raw_id in zip(asym_ids, entity_ids, strict=False): + entity_id = int(raw_id) + if entity_id in entities and not entities[entity_id]: + entities[entity_id].append(asym_id) + return entities + + +def _polymer_sequences(block) -> dict[str, str]: + sequences: dict[str, str] = {} + if "entity_poly" not in block: + return sequences + category = block["entity_poly"] + entity_ids = category["entity_id"].as_array(str) + raw_sequences = category["pdbx_seq_one_letter_code_can"].as_array(str) + chain_lists = category["pdbx_strand_id"].as_array(str) + for _, raw_sequence, raw_chains in zip( + entity_ids, + raw_sequences, + chain_lists, + strict=False, + ): + sequence = "".join(raw_sequence.split()) + for chain_id in _clean_chain_list(raw_chains): + sequences[chain_id] = sequence + return sequences + + +def _scheme_columns(category): + asym_ids = category["asym_id"].as_array(str) + insertion_codes = ( + category["pdb_ins_code"].as_array(str) + if "pdb_ins_code" in category + else [""] * len(asym_ids) + ) + hetflags = category["hetflag"].as_array(str) if "hetflag" in category else ["N"] * len(asym_ids) + author_chains = ( + category["pdb_strand_id"].as_array(str) if "pdb_strand_id" in category else asym_ids + ) + return ( + asym_ids, + category["seq_id"].as_array(str), + category["auth_seq_num"].as_array(str), + insertion_codes, + hetflags, + author_chains, + ) + + +def _scheme_residue_map(category): + ( + asym_ids, + sequence_positions, + author_numbers, + insertion_codes, + hetflags, + author_chains, + ) = _scheme_columns(category) + asym_to_author = { + asym_id: author_id for asym_id, author_id in zip(asym_ids, author_chains, strict=False) + } + per_chain: dict[str, dict[int, Residue]] = {} + for asym_id, raw_position, raw_number, raw_code, raw_hetflag in zip( + asym_ids, + sequence_positions, + author_numbers, + insertion_codes, + hetflags, + strict=False, + ): + residues = per_chain.setdefault(asym_id, {}) + try: + position = int(raw_position) - 1 + residue_number = int(raw_number) if raw_number != "?" else None + except ValueError: + continue + if residue_number is None: + insertion_code = "" + else: + insertion_code = "" if raw_code in (".", "?") else raw_code + residues[position] = Residue( + residue_number=residue_number, + insertion_code=insertion_code, + hetflag=raw_hetflag.upper() == "Y", + ) + return per_chain, asym_to_author + + +def _renumber_duplicate_residues( + per_chain: dict[str, dict[int, Residue]], +) -> None: + for residues in per_chain.values(): + positions_by_number: dict[int, list[int]] = {} + for position, residue in residues.items(): + if residue.residue_number is not None: + positions_by_number.setdefault(residue.residue_number, []).append(position) + for number, positions in positions_by_number.items(): + if len(positions) <= 1: + continue + positions.sort() + for offset, position in enumerate(positions): + previous = residues[position] + residues[position] = Residue( + residue_number=number + offset, + insertion_code=previous.insertion_code, + hetflag=previous.hetflag, + ) + + +def _ordered_scheme_mapping( + per_chain: dict[str, dict[int, Residue]], + asym_to_author: dict[str, str], + chain_sequences: dict[str, str], +) -> dict[str, dict[int, Residue]]: + result: dict[str, dict[int, Residue]] = {} + for asym_id, residues in per_chain.items(): + author_chain = asym_to_author.get(asym_id, asym_id) + if author_chain in chain_sequences: + result[author_chain] = { + position: residues.get(position, _empty_residue()) + for position in range(len(chain_sequences[author_chain])) + } + elif residues: + result[author_chain] = { + index: residues[position] for index, position in enumerate(sorted(residues)) + } + return result + + +def _complete_polymer_mappings( + mappings: dict[str, dict[int, Residue]], + chain_sequences: dict[str, str], +) -> None: + for chain_id, sequence in chain_sequences.items(): + mapping = mappings.setdefault(chain_id, {}) + for position in range(len(sequence)): + if position not in mapping: + mapping[position] = _empty_residue() + + +def _add_structure_fallbacks( + mappings: dict[str, dict[int, Residue]], + structure: bs.AtomArray, +) -> None: + if not ( + structure + and hasattr(structure, "chain_id") + and structure.chain_id is not None + and hasattr(structure.chain_id, "__iter__") + ): + return + for chain_id in set(structure.chain_id): + if chain_id in mappings: + continue + chain = structure[structure.chain_id == chain_id] + if not ( + hasattr(chain, "res_id") + and chain.res_id is not None + and hasattr(chain.res_id, "__iter__") + ): + continue + residue_ids = sorted(set(chain.res_id)) + mappings[chain_id] = { + index: Residue( + residue_number=residue_id, + insertion_code="", + hetflag=False, + ) + for index, residue_id in enumerate(residue_ids) + } + + +def _nonpolymer_entity_ids(block) -> set[str]: + result = set() + if "entity" not in block: + return result + category = block["entity"] + ids = category["id"].as_array(str) + types = category["type"].as_array(str) + for entity_id, entity_type in zip(ids, types, strict=False): + if entity_type.upper() in _NONPOLYMER_ENTITY_TYPES: + result.add(entity_id) + return result + + +def _nonpolymer_component_map(block, entity_ids: set[str]) -> dict[str, str]: + result = {} + if "pdbx_entity_nonpoly" not in block: + return result + category = block["pdbx_entity_nonpoly"] + ids = category["entity_id"].as_array(str) + components = category["comp_id"].as_array(str) + for entity_id, component in zip(ids, components, strict=False): + if entity_id in entity_ids: + result[entity_id] = component + return result + + +class MmcifWrapper: + """Parsed model-one structure, metadata, sequences, and residue mappings.""" + + def __init__(self, id: str | None = None): + self.id = id or "" + self.raw: pdbx.CIFFile | None = None + self.structure: bs.AtomArray + self.header = MmcifHeader() + self.entities: dict[int, list[str]] = {} + self.chain_to_seqres: dict[str, str] = {} + self.seqres_to_structure: dict[str, dict[int, Residue]] = {} + + @classmethod + def read(cls, path: PathOrBuffer, id: str | None = None) -> MmcifWrapper: + wrapper = cls(id=id) + wrapper._load(path) + return wrapper + + def _load(self, path: PathOrBuffer, fileid: str | None = None) -> None: + self.raw = pdbx.CIFFile.read(path) + self._parse_structure() + self._parse_header() + self._parse_entities() + self._parse_sequences() + + def _parse_structure(self) -> None: + try: + structure = pdbx.get_structure(self.raw, model=1) + if structure is None or not isinstance(structure, bs.AtomArray): + raise NoProteinError("No structure found in mmCIF file") + if len(structure) == 0: + raise NoProteinError("Empty structure in mmCIF file") + self.structure = structure + except Exception as error: + raise ValueError(f"Failed to parse structure: {error}") from error + + def _parse_header(self) -> None: + if self.raw: + self.header = _header_from_block(self.raw.block, self.header) + + def _parse_entities(self) -> None: + if not self.raw: + return + try: + self.entities = _entities_from_block(self.raw.block, self.entities) + except Exception: + if ( + self.structure + and hasattr(self.structure, "chain_id") + and self.structure.chain_id is not None + and hasattr(self.structure.chain_id, "__iter__") + ): + self.entities = {1: list(set(self.structure.chain_id))} + + def _parse_sequences(self) -> None: + if not self.raw: + return + block = self.raw.block + self.chain_to_seqres.update(_polymer_sequences(block)) + if "pdbx_poly_seq_scheme" in block: + per_chain, asym_to_author = _scheme_residue_map(block["pdbx_poly_seq_scheme"]) + _renumber_duplicate_residues(per_chain) + self.seqres_to_structure.update( + _ordered_scheme_mapping( + per_chain, + asym_to_author, + self.chain_to_seqres, + ) + ) + _complete_polymer_mappings( + self.seqres_to_structure, + self.chain_to_seqres, + ) + _add_structure_fallbacks(self.seqres_to_structure, self.structure) + + def _parse_nonpoly_from_mmcif(self) -> dict[tuple, bs.AtomArray]: + assert self.raw is not None + block = self.raw.block + entity_ids = _nonpolymer_entity_ids(block) + _nonpolymer_component_map(block, entity_ids) + groups: dict[tuple[str, str], list[int]] = {} + if "atom_site" in block: + category = block["atom_site"] + chain_ids = category["label_asym_id"].as_array(str) + atom_entity_ids = category["label_entity_id"].as_array(str) + component_ids = category["label_comp_id"].as_array(str) + for index, (chain_id, entity_id, component_id) in enumerate( + zip(chain_ids, atom_entity_ids, component_ids, strict=False) + ): + if entity_id in entity_ids: + groups.setdefault((component_id, chain_id), []).append(index) + + coordinates = {} + for component_id, chain_id in groups: + selection = (self.structure.chain_id == chain_id) & ( + self.structure.res_name == component_id + ) + if not selection.any(): + continue + atoms = self.structure[selection] + if isinstance(atoms, (bs.AtomArray, bs.AtomArrayStack)) and len(atoms) > 0: + coordinates[(component_id, chain_id)] = atoms + return coordinates + + def _parse_nonpoly_fallback(self) -> dict[tuple, bs.AtomArray]: + result = {} + if not (self.structure and hasattr(self.structure, "chain_id")): + return result + standard_residues = set(residue_constants.resnames[:-1]) + standard_residues.update({"A", "C", "G", "T", "U"}) + if self.structure.chain_id is None: + return result + for chain_id in set(self.structure.chain_id): + chain = self.structure[self.structure.chain_id == chain_id] + if not ( + hasattr(chain, "res_name") + and chain.res_name is not None + and hasattr(chain.res_name, "__iter__") + ): + continue + for residue_name in set(chain.res_name): + if residue_name in standard_residues: + continue + selection = (chain.chain_id == chain_id) & (chain.res_name == residue_name) + if selection.any() and isinstance( + chain, + (bs.AtomArray, bs.AtomArrayStack), + ): + result[(residue_name, chain_id)] = chain[selection] + return result + + @functools.cached_property + def non_polymer_coords(self) -> dict[tuple, bs.AtomArray]: + """Map each non-polymer component and chain to its atoms.""" + + if not self.structure or not self.raw: + return {} + try: + return self._parse_nonpoly_from_mmcif() + except Exception: + return self._parse_nonpoly_fallback() diff --git a/fastplms/models/esmfold2/esmfold2_molecular_complex.py b/fastplms/models/esmfold2/esmfold2_molecular_complex.py new file mode 100644 index 0000000000000000000000000000000000000000..3c8921534302e524a9cc3e11242a6c54f16a53e5 --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_molecular_complex.py @@ -0,0 +1,1016 @@ +"""Flat molecular-complex records used by the ESMFold2 public API. + +The folding model operates on tokens and a single atom table. This module owns +that representation, its protein-only bridge, mmCIF I/O, structure metrics, and +the compact wire format. It deliberately has no dependency on the upstream +Biohub package; the pinned submodule is used only by differential tests. +""" + +from __future__ import annotations + +import io +import os +import re +from dataclasses import asdict, dataclass +from pathlib import Path +from subprocess import check_output +from tempfile import TemporaryDirectory +from typing import TYPE_CHECKING, Any + +import biotite.structure as bs +import biotite.structure.io.pdbx as pdbx +import brotli +import msgpack +import numpy as np +import torch +from biotite.structure.io.pdbx import ( + CIFCategory, + CIFColumn, + CIFData, + CIFFile, + set_structure, +) + +from . import esmfold2_residue_constants as residue_constants +from .esmfold2_metrics import compute_lddt, compute_rmsd +from .esmfold2_mmcif_parsing import PLDDT_B_FACTOR_SCALE, round_mmcif_columns +from .esmfold2_protein_complex import ProteinComplex, ProteinComplexMetadata + + +@dataclass +class MolecularComplexResult: + """One folded complex and the optional model outputs associated with it.""" + + complex: MolecularComplex + plddt: torch.Tensor | None = None + ptm: float | None = None + iptm: float | None = None + pae: torch.Tensor | None = None + distogram: torch.Tensor | None = None + pair_chains_iptm: torch.Tensor | None = None + output_embedding_sequence: torch.Tensor | None = None + output_embedding_pair_pooled: torch.Tensor | None = None + residue_index: torch.Tensor | None = None + entity_id: torch.Tensor | None = None + sae_features: np.ndarray | None = None # X has shape (l, n_features). + ttt_metrics: dict[str, Any] | None = None + + +@dataclass +class MolecularComplexMetadata: + """Entity and chain labels carried with a molecular complex.""" + + entity_lookup: dict[int, str] + chain_lookup: dict[int, str] + assembly_composition: dict[str, list[str]] | None = None + + +@dataclass +class Molecule: + """The atom slice represented by one model token.""" + + token: str + token_idx: int + atom_positions: np.ndarray # P has shape (n_atoms, 3). + atom_elements: np.ndarray # E has shape (n_atoms,). + atom_names: np.ndarray | None = None # N has shape (n_atoms,) when present. + atom_hetero: np.ndarray | None = None # M has shape (n_atoms,) when present. + residue_type: int = 0 + molecule_type: int = 0 + confidence: float = 0.0 + + +_NUCLEOTIDE_NAMES = frozenset({"A", "T", "G", "C", "U", "DA", "DT", "DG", "DC"}) +_SERIALIZED_ARRAYS = frozenset( + { + "atom_positions", + "atom_elements", + "atom_names", + "atom_hetero", + "token_to_atoms", + "chain_id", + "entity_id", + "sym_id", + "plddt", + } +) + + +def _assert_table_lengths(complex_value: MolecularComplex) -> None: + """Check that token and atom annotations align with their tables.""" + if not isinstance(complex_value.sequence, list) or any( + not isinstance(token, str) for token in complex_value.sequence + ): + raise TypeError("sequence must be a list of token strings.") + n_tokens = len(complex_value.sequence) + if not isinstance(complex_value.atom_positions, np.ndarray): + raise TypeError("atom_positions must be a NumPy array.") + if complex_value.atom_positions.ndim != 2 or complex_value.atom_positions.shape[1:] != ( + 3, + ): + raise ValueError( + "atom_positions must have shape (n_atoms, 3), got " + f"{complex_value.atom_positions.shape}." + ) + if not np.issubdtype(complex_value.atom_positions.dtype, np.number): + raise TypeError("atom_positions must use a numeric dtype.") + n_atoms = len(complex_value.atom_positions) + if not isinstance(complex_value.atom_elements, np.ndarray): + raise TypeError("atom_elements must be a NumPy array.") + if complex_value.atom_elements.shape != (n_atoms,): + raise ValueError( + f"atom_elements shape {complex_value.atom_elements.shape} != {n_atoms} atoms" + ) + token_tables = { + "token_to_atoms": complex_value.token_to_atoms, + "chain_id": complex_value.chain_id, + "plddt": complex_value.plddt, + } + if complex_value.entity_id is not None: + token_tables["entity_id"] = complex_value.entity_id + if complex_value.sym_id is not None: + token_tables["sym_id"] = complex_value.sym_id + for label, values in token_tables.items(): + if not isinstance(values, np.ndarray): + raise TypeError(f"{label} must be a NumPy array, got {type(values).__name__}.") + if values.ndim == 0 or values.shape[0] != n_tokens: + raise ValueError(f"{label} shape {values.shape} != {n_tokens} tokens") + if complex_value.token_to_atoms.shape != (n_tokens, 2): + raise ValueError( + "token_to_atoms must have shape " + f"({n_tokens}, 2), got {complex_value.token_to_atoms.shape}." + ) + if not np.issubdtype(complex_value.token_to_atoms.dtype, np.integer): + raise TypeError("token_to_atoms must use an integer dtype.") + if complex_value.chain_id.shape != (n_tokens,): + raise ValueError(f"chain_id must have shape ({n_tokens},).") + for label, values in ( + ("chain_id", complex_value.chain_id), + ("entity_id", complex_value.entity_id), + ("sym_id", complex_value.sym_id), + ): + if values is not None and values.shape != (n_tokens,): + raise ValueError(f"{label} must have shape ({n_tokens},).") + if values is not None and not np.issubdtype(values.dtype, np.integer): + raise TypeError(f"{label} must use an integer dtype.") + if complex_value.plddt.shape != (n_tokens,): + raise ValueError(f"plddt must have shape ({n_tokens},).") + if not np.issubdtype(complex_value.plddt.dtype, np.number): + raise TypeError("plddt must use a numeric dtype.") + if n_tokens: + starts = complex_value.token_to_atoms[:, 0] + stops = complex_value.token_to_atoms[:, 1] + if np.any(starts < 0) or np.any(stops < starts) or np.any(stops > n_atoms): + raise ValueError("token_to_atoms contains an invalid or out-of-bounds atom span.") + for label, values in ( + ("atom_names", complex_value.atom_names), + ("atom_hetero", complex_value.atom_hetero), + ): + if values is not None and not isinstance(values, np.ndarray): + raise TypeError(f"{label} must be a NumPy array, got {type(values).__name__}.") + if isinstance(values, np.ndarray) and values.shape != (n_atoms,): + raise ValueError(f"{label} shape {values.shape} != {n_atoms} atoms") + + +def _flat_protein_atoms( + protein: ProteinComplex, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Flatten the populated atom37 entries of a protein complex.""" + positions: list[np.ndarray] = [] + elements: list[str] = [] + names: list[str] = [] + hetero: list[bool] = [] + spans: list[tuple[int, int]] = [] + + for sequence_index, residue in enumerate(protein.sequence): + if residue == "|": + continue + start = len(positions) + mask = protein.atom37_mask[sequence_index] + residue_positions = protein.atom37_positions[sequence_index] + for atom_index in np.flatnonzero(mask): + atom_name = residue_constants.atom_types[int(atom_index)] + positions.append(residue_positions[atom_index]) + elements.append(atom_name[0] if atom_name else "C") + names.append(atom_name) + hetero.append(False) + spans.append((start, len(positions))) + + return ( + np.asarray(positions, dtype=np.float32), + np.asarray(elements, dtype=object), + np.asarray(names, dtype=object), + np.asarray(hetero, dtype=bool), + np.asarray(spans, dtype=np.int32), + ) + + +def _protein_sequence_and_indices( + complex_value: MolecularComplex, +) -> tuple[list[int], str, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + protein_indices = [ + index + for index, token in enumerate(complex_value.sequence) + if token in residue_constants.restype_3to1 + ] + if not protein_indices: + raise ValueError("No protein tokens found in MolecularComplex") + + chain_ids = complex_value.chain_id[protein_indices] + entity_ids = ( + chain_ids + if complex_value.entity_id is None + else complex_value.entity_id[protein_indices] + ) + sym_ids = ( + np.zeros_like(chain_ids) + if complex_value.sym_id is None + else complex_value.sym_id[protein_indices] + ) + confidences = complex_value.plddt[protein_indices] + sequence: list[str] = [] + previous_instance: Any = None + preserve_instances = complex_value.sym_id is not None + for index, chain_id, sym_id in zip( + protein_indices, chain_ids, sym_ids, strict=True + ): + instance = (int(chain_id), int(sym_id)) if preserve_instances else int(chain_id) + if previous_instance is not None and instance != previous_instance: + sequence.append("|") + sequence.append(residue_constants.restype_3to1[complex_value.sequence[index]]) + previous_instance = instance + return protein_indices, "".join(sequence), chain_ids, entity_ids, sym_ids, confidences + + +def _protein_entity_metadata_value(value: int | str) -> int | str: + """Restore the numeric entity labels used by ProteinComplex metadata.""" + if isinstance(value, str): + try: + return int(value) + except ValueError: + pass + return value + + +def _atom37_from_flat( + complex_value: MolecularComplex, protein_indices: list[int] +) -> tuple[np.ndarray, np.ndarray]: + n_residues = len(protein_indices) + positions = np.full((n_residues, 37, 3), np.nan, dtype=np.float32) + mask = np.zeros((n_residues, 37), dtype=bool) + if complex_value.atom_names is None: + return positions, mask + + for residue_index, token_index in enumerate(protein_indices): + start, stop = complex_value.token_to_atoms[token_index] + seen: set[str] = set() + for atom_name, atom_position in zip( + complex_value.atom_names[start:stop], + complex_value.atom_positions[start:stop], + strict=True, + ): + normalized = str(atom_name).upper().strip() + if normalized in seen: + continue + seen.add(normalized) + atom37_index = residue_constants.atom_order.get(normalized) + if atom37_index is not None: + positions[residue_index, atom37_index] = atom_position + mask[residue_index, atom37_index] = True + return positions, mask + + +def _expand_protein_rows( + sequence: str, + protein_chain_ids: np.ndarray, + protein_entity_ids: np.ndarray, + protein_sym_ids: np.ndarray, + confidences: np.ndarray, + compact_positions: np.ndarray, + compact_mask: np.ndarray, +) -> dict[str, np.ndarray]: + """Insert empty rows at chain separators in a protein representation.""" + n_positions = len(sequence) + expanded = { + "chain_id": np.full(n_positions, -1, dtype=np.int64), + "entity_id": np.full(n_positions, -1, dtype=np.int64), + "sym_id": np.zeros(n_positions, dtype=np.int64), + "residue_index": np.zeros(n_positions, dtype=np.int64), + "insertion_code": np.asarray([""] * n_positions, dtype=object), + "confidence": np.zeros(n_positions, dtype=np.float32), + "atom37_positions": np.full((n_positions, 37, 3), np.nan, dtype=np.float32), + "atom37_mask": np.zeros((n_positions, 37), dtype=bool), + } + residue_number = 0 + compact_index = 0 + for sequence_index, residue in enumerate(sequence): + if residue == "|": + residue_number = 0 + continue + chain_id = protein_chain_ids[compact_index] + residue_number += 1 + expanded["chain_id"][sequence_index] = chain_id + expanded["entity_id"][sequence_index] = protein_entity_ids[compact_index] + expanded["sym_id"][sequence_index] = protein_sym_ids[compact_index] + expanded["residue_index"][sequence_index] = residue_number + expanded["confidence"][sequence_index] = confidences[compact_index] + expanded["atom37_positions"][sequence_index] = compact_positions[compact_index] + expanded["atom37_mask"][sequence_index] = compact_mask[compact_index] + compact_index += 1 + return expanded + + +def _read_cif(source: str) -> CIFFile: + if os.path.exists(source): + return pdbx.CIFFile.read(source) + return pdbx.CIFFile.read(io.StringIO(source)) + + +def _read_structure(cif_file: CIFFile) -> Any: + try: + return pdbx.get_structure(cif_file, model=1, extra_fields=["b_factor"]) + except (KeyError, ValueError): + try: + return pdbx.get_structure(cif_file) + except Exception: + return pdbx.get_structure(cif_file, model=None) + + +def _column_array(category: Any, name: str) -> np.ndarray: + column = category[name] + if hasattr(column, "as_array"): + return column.as_array(str) + return np.asarray(list(column), dtype=str) + + +def _label_asym_ids(cif_file: CIFFile, n_structure_atoms: int) -> list[str] | None: + """Return label-asym identifiers after applying Biohub's atom filters.""" + block = cif_file.block + if "atom_site" not in block or "label_asym_id" not in block["atom_site"]: + return None + atom_site = block["atom_site"] + labels = _column_array(atom_site, "label_asym_id") + keep = np.ones(len(labels), dtype=bool) + if "pdbx_PDB_model_num" in atom_site: + keep &= _column_array(atom_site, "pdbx_PDB_model_num") == "1" + if "label_alt_id" in atom_site: + keep &= np.isin(_column_array(atom_site, "label_alt_id"), [".", "?", "", "A"]) + filtered = labels[keep] + return filtered.tolist() if len(filtered) == n_structure_atoms else None + + +def _entity_metadata(cif_file: CIFFile) -> dict[Any, Any]: + result: dict[Any, Any] = {} + try: + category = cif_file.block["entity"] + if "id" not in category or "type" not in category: + return result + for entity_id, entity_type in zip(category["id"], category["type"], strict=False): + result[entity_id] = entity_type + except Exception: + return {} + return result + + +def _group_structure_atoms( + structure: Any, labels: list[str] | None +) -> dict[str, dict[tuple[int, str], dict[str, Any]]]: + grouped: dict[str, dict[tuple[int, str], dict[str, Any]]] = {} + for atom_index, atom in enumerate(structure): + chain = labels[atom_index] if labels is not None else atom.chain_id + residues = grouped.setdefault(chain, {}) + key = (atom.res_id, atom.res_name) + record = residues.setdefault( + key, + {"atoms": [], "res_name": atom.res_name, "is_hetero": atom.hetero}, + ) + record["atoms"].append(atom) + return grouped + + +def _flatten_structure_groups( + grouped: dict[str, dict[tuple[int, str], dict[str, Any]]], +) -> tuple[ + list[str], + list[np.ndarray], + list[str], + list[str], + list[bool], + list[tuple[int, int]], + list[float], + list[int], + dict[str, int], +]: + tokens: list[str] = [] + positions: list[np.ndarray] = [] + elements: list[str] = [] + names: list[str] = [] + hetero: list[bool] = [] + spans: list[tuple[int, int]] = [] + confidences: list[float] = [] + token_chains: list[int] = [] + chain_numbers = {chain: index for index, chain in enumerate(sorted(grouped))} + + for chain in sorted(grouped): + for residue_key in sorted(grouped[chain]): + record = grouped[chain][residue_key] + if record["res_name"] == "HOH": + continue + atoms = record["atoms"] + tokens.append(record["res_name"]) + token_chains.append(chain_numbers[chain]) + start = len(positions) + positions.extend(atom.coord for atom in atoms) + elements.extend(atom.element for atom in atoms) + names.extend(atom.atom_name for atom in atoms) + hetero.extend(atom.hetero for atom in atoms) + spans.append((start, len(positions))) + b_factor = getattr(atoms[0], "b_factor", 50.0) if atoms else 50.0 + confidences.append(min(b_factor / PLDDT_B_FACTOR_SCALE, 1.0)) + return ( + tokens, + positions, + elements, + names, + hetero, + spans, + confidences, + token_chains, + chain_numbers, + ) + + +def _chain_entity_maps( + complex_value: MolecularComplex, +) -> tuple[dict[str, list[str]], dict[str, int], dict[int, tuple[str, ...]]]: + chains: dict[str, list[str]] = {} + for token_index, numeric_chain in enumerate(complex_value.chain_id): + numeric = int(numeric_chain) + label = complex_value.metadata.chain_lookup.get(numeric, chr(65 + numeric)) + chains.setdefault(label, []).append(complex_value.sequence[token_index]) + + sequence_entities: dict[tuple[str, ...], int] = {} + chain_entities: dict[str, int] = {} + entity_sequences: dict[int, tuple[str, ...]] = {} + for label, sequence in chains.items(): + key = tuple(sequence) + entity_id = sequence_entities.get(key) + if entity_id is None: + entity_id = len(sequence_entities) + 1 + sequence_entities[key] = entity_id + entity_sequences[entity_id] = key + chain_entities[label] = entity_id + return chains, chain_entities, entity_sequences + + +def _cif_column(values: list[str]) -> CIFColumn: + return CIFColumn(data=CIFData(array=np.asarray(values), dtype=np.str_)) + + +def _add_entity_categories( + cif_file: CIFFile, + complex_value: MolecularComplex, + entity_sequences: dict[int, tuple[str, ...]], +) -> None: + ids: list[str] = [] + types: list[str] = [] + descriptions: list[str] = [] + for entity_id in sorted(entity_sequences): + sequence = entity_sequences[entity_id] + protein = any(token in residue_constants.restype_3to1 for token in sequence) + nucleic = any(token in _NUCLEOTIDE_NAMES for token in sequence) + ids.append(str(entity_id)) + types.append("polymer" if protein or nucleic else "non-polymer") + if protein: + descriptions.append(f"Polymer entity {entity_id} (protein)") + elif nucleic: + descriptions.append(f"Polymer entity {entity_id} (nucleic acid)") + else: + descriptions.append(f"Non-polymer entity {entity_id}") + + if ids: + cif_file.block["entity"] = CIFCategory( + name="entity", + columns={ + "id": _cif_column(ids), + "type": _cif_column(types), + "pdbx_description": _cif_column(descriptions), + }, + ) + + _, chain_entities, _ = _chain_entity_maps(complex_value) + if chain_entities: + labels = sorted(chain_entities) + cif_file.block["struct_asym"] = CIFCategory( + name="struct_asym", + columns={ + "id": _cif_column(labels), + "entity_id": _cif_column([str(chain_entities[label]) for label in labels]), + }, + ) + + entity_chains: dict[int, list[str]] = {} + for chain, entity_id in chain_entities.items(): + entity_chains.setdefault(entity_id, []).append(chain) + polymer_rows: list[tuple[str, str, str, str]] = [] + residue_rows: list[tuple[str, str, str, str]] = [] + for entity_id in sorted(entity_sequences): + sequence = entity_sequences[entity_id] + protein = any(token in residue_constants.restype_3to1 for token in sequence) + nucleic = any(token in _NUCLEOTIDE_NAMES for token in sequence) + if not (protein or nucleic): + continue + if protein: + polymer_type = "polypeptide(L)" + canonical = "".join( + residue_constants.restype_3to1.get(token, "(X)") for token in sequence + ) + else: + polymer_type = ( + "polyribonucleotide" + if "U" in sequence + else ( + "polydeoxyribonucleotide" + if any(token in {"DA", "DT", "DG", "DC"} for token in sequence) + else "polyribonucleotide" + ) + ) + nucleotide_letters = {"DA": "A", "DT": "T", "DG": "G", "DC": "C"} + canonical = "".join(nucleotide_letters.get(token, token) for token in sequence) + strand_ids = ",".join(sorted(entity_chains.get(entity_id, []))) or "?" + polymer_rows.append((str(entity_id), polymer_type, strand_ids, canonical)) + residue_rows.extend( + (str(entity_id), str(number), token, "n") + for number, token in enumerate(sequence, start=1) + ) + + if polymer_rows: + columns = list(zip(*polymer_rows, strict=True)) + cif_file.block["entity_poly"] = CIFCategory( + name="entity_poly", + columns={ + "entity_id": _cif_column(list(columns[0])), + "type": _cif_column(list(columns[1])), + "pdbx_strand_id": _cif_column(list(columns[2])), + "pdbx_seq_one_letter_code_can": _cif_column(list(columns[3])), + }, + ) + if residue_rows: + columns = list(zip(*residue_rows, strict=True)) + cif_file.block["entity_poly_seq"] = CIFCategory( + name="entity_poly_seq", + columns={ + "entity_id": _cif_column(list(columns[0])), + "num": _cif_column(list(columns[1])), + "mon_id": _cif_column(list(columns[2])), + "hetero": _cif_column(list(columns[3])), + }, + ) + + +def _fallback_atom_names(token: str, count: int) -> list[str]: + if token in residue_constants.restype_3to1: + names = list(residue_constants.residue_atoms.get(token, ["N", "CA", "C", "O"]))[:count] + names.extend(f"X{index + 1}" for index in range(len(names), count)) + return names + return [f"C{index + 1}" for index in range(count)] + + +def _as_atom_array(complex_value: MolecularComplex, chain_entities: dict[str, int]) -> bs.AtomArray: + n_atoms = len(complex_value.atom_positions) + atom_array = bs.AtomArray(length=n_atoms) + atom_array.coord = complex_value.atom_positions + residue_ids = np.zeros(n_atoms, dtype=np.int32) + chain_labels = np.empty(n_atoms, dtype=object) + residue_names = np.empty(n_atoms, dtype=object) + hetero = np.zeros(n_atoms, dtype=bool) + b_factors = np.zeros(n_atoms, dtype=np.float32) + atom_names = np.empty(n_atoms, dtype=object) + entity_ids = np.zeros(n_atoms, dtype=np.int32) + next_residue: dict[Any, int] = {} + + for token_index, (start, stop) in enumerate(complex_value.token_to_atoms): + token = complex_value.sequence[token_index] + numeric_chain = complex_value.chain_id[token_index] + numeric = int(numeric_chain) + chain = complex_value.metadata.chain_lookup.get(numeric, chr(65 + numeric)) + residue_id = next_residue.get(numeric_chain, 0) + 1 + next_residue[numeric_chain] = residue_id + count = int(stop - start) + names = ( + list(complex_value.atom_names[start:stop]) + if complex_value.atom_names is not None + else _fallback_atom_names(token, count) + ) + residue_ids[start:stop] = residue_id + chain_labels[start:stop] = chain + residue_names[start:stop] = token + hetero[start:stop] = ( + complex_value.atom_hetero[start:stop] + if complex_value.atom_hetero is not None + else token not in residue_constants.restype_3to1 + ) + b_factors[start:stop] = complex_value.plddt[token_index] * PLDDT_B_FACTOR_SCALE + atom_names[start:stop] = names + entity_ids[start:stop] = chain_entities.get(chain, 1) + + atom_array.res_id = residue_ids + atom_array.chain_id = np.asarray(chain_labels, dtype="U16") + atom_array.res_name = np.asarray(residue_names, dtype="U8") + atom_array.hetero = hetero + atom_array.atom_name = np.asarray(atom_names, dtype="U4") + atom_array.add_annotation("b_factor", dtype=float) + atom_array.b_factor = b_factors + atom_array.add_annotation("occupancy", dtype=float) + atom_array.occupancy = np.ones(n_atoms, dtype=np.float32) + atom_array.add_annotation("entity_id", dtype=int) + atom_array.entity_id = entity_ids + if complex_value.atom_elements is not None and len(complex_value.atom_elements) == n_atoms: + atom_array.element = np.asarray(complex_value.atom_elements, dtype="U4") + else: + atom_array.element = bs.infer_elements(atom_array) + return atom_array + + +def _repair_label_entity_ids(cif_file: CIFFile, chain_entities: dict[str, int]) -> None: + if "atom_site" not in cif_file.block: + return + atom_site = cif_file.block["atom_site"] + if "label_asym_id" not in atom_site or "label_entity_id" not in atom_site: + return + labels = _column_array(atom_site, "label_asym_id").tolist() + if labels: + atom_site["label_entity_id"] = _cif_column( + [str(chain_entities.get(label, 1)) for label in labels] + ) + + +def _centroid_tensors( + mobile: MolecularComplex, + target: MolecularComplex, + *, + retain_missing: bool, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if len(mobile) != len(target): + raise ValueError( + f"Complexes must have the same number of tokens: {len(mobile)} vs {len(target)}" + ) + mobile_centers: list[np.ndarray] = [] + target_centers: list[np.ndarray] = [] + valid: list[bool] = [] + for token_index in range(len(mobile)): + mobile_start, mobile_stop = mobile.token_to_atoms[token_index] + target_start, target_stop = target.token_to_atoms[token_index] + mobile_atoms = mobile.atom_positions[mobile_start:mobile_stop] + target_atoms = target.atom_positions[target_start:target_stop] + present = len(mobile_atoms) > 0 and len(target_atoms) > 0 + if not present and not retain_missing: + continue + if present: + mobile_centers.append(mobile_atoms.mean(axis=0)) + target_centers.append(target_atoms.mean(axis=0)) + else: + mobile_centers.append(np.full(3, np.nan)) + target_centers.append(np.full(3, np.nan)) + valid.append(present) + if not any(valid): + metric = "LDDT" if retain_missing else "RMSD" + raise ValueError(f"No valid atoms found for {metric} computation") + return ( + torch.from_numpy(np.stack(mobile_centers)).unsqueeze(0), + torch.from_numpy(np.stack(target_centers)).unsqueeze(0), + torch.as_tensor(valid, dtype=torch.bool).unsqueeze(0), + ) + + +@dataclass(frozen=True) +class MolecularComplex: + """A token sequence backed by one contiguous atom table. + + P stores atom coordinates with shape (n_atoms, 3). Token span ``i`` is + ``P[token_to_atoms[i, 0]:token_to_atoms[i, 1]]``. ``chain_id`` identifies + the author chain, while optional ``entity_id`` and ``sym_id`` distinguish + biological entities and repeated chain instances. + """ + + id: str + sequence: list[str] + atom_positions: np.ndarray # P has shape (n_atoms, 3). + atom_elements: np.ndarray # E has shape (n_atoms,). + token_to_atoms: np.ndarray # I has shape (n_tokens, 2). + chain_id: np.ndarray # C has shape (n_tokens,). + plddt: np.ndarray # S has shape (n_tokens,). + metadata: MolecularComplexMetadata + atom_names: np.ndarray | None = None # N has shape (n_atoms,) when present. + atom_hetero: np.ndarray | None = None # M has shape (n_atoms,) when present. + # These token-aligned IDs are optional for compatibility with older blobs. + # ProteinComplex adapters populate them so homomers and repeated author-chain + # labels survive a MolecularComplex round trip. + entity_id: np.ndarray | None = None + sym_id: np.ndarray | None = None + + def __post_init__(self) -> None: + _assert_table_lengths(self) + + def __len__(self) -> int: + return len(self.sequence) + + def __getitem__(self, idx: int) -> Molecule: + if idx < 0 or idx >= len(self): + raise IndexError(f"Token index {idx} out of range for {len(self)} tokens") + start, stop = self.token_to_atoms[idx] + return Molecule( + token=self.sequence[idx], + token_idx=idx, + atom_positions=self.atom_positions[start:stop], + atom_elements=self.atom_elements[start:stop], + atom_names=None if self.atom_names is None else self.atom_names[start:stop], + atom_hetero=(None if self.atom_hetero is None else self.atom_hetero[start:stop]), + residue_type=0, + molecule_type=0, + confidence=self.plddt[idx], + ) + + @property + def atom_coordinates(self) -> np.ndarray: + """Return P, the flat atom-coordinate table with shape (n_atoms, 3).""" + return self.atom_positions + + @classmethod + def from_protein_complex(cls, pc: ProteinComplex) -> MolecularComplex: + positions, elements, names, hetero, spans = _flat_protein_atoms(pc) + residue_positions = [index for index, value in enumerate(pc.sequence) if value != "|"] + metadata = MolecularComplexMetadata( + entity_lookup={key: str(value) for key, value in pc.metadata.entity_lookup.items()}, + chain_lookup=dict(pc.metadata.chain_lookup), + assembly_composition=pc.metadata.assembly_composition, + ) + return cls( + id=pc.id, + sequence=[ + residue_constants.restype_1to3.get(pc.sequence[index], "UNK") + for index in residue_positions + ], + atom_positions=positions, + atom_elements=elements, + token_to_atoms=spans, + chain_id=np.asarray(pc.chain_id[residue_positions], dtype=np.int64), + plddt=np.asarray(pc.confidence[residue_positions], dtype=np.float32), + metadata=metadata, + atom_names=names, + atom_hetero=hetero, + entity_id=np.asarray(pc.entity_id[residue_positions], dtype=np.int64), + sym_id=np.asarray(pc.sym_id[residue_positions], dtype=np.int64), + ) + + def to_protein_complex(self) -> ProteinComplex: + ( + protein_indices, + sequence, + chain_ids, + entity_ids, + sym_ids, + confidences, + ) = _protein_sequence_and_indices(self) + compact_positions, compact_mask = _atom37_from_flat(self, protein_indices) + arrays = _expand_protein_rows( + sequence, + chain_ids, + entity_ids, + sym_ids, + confidences, + compact_positions, + compact_mask, + ) + unique_chains = np.unique(chain_ids) + unique_entities = np.unique(entity_ids) + metadata = ProteinComplexMetadata( + entity_lookup={ + int(entity): _protein_entity_metadata_value( + self.metadata.entity_lookup.get(int(entity), int(entity)) + ) + for entity in unique_entities + }, + chain_lookup={ + int(chain): self.metadata.chain_lookup.get(int(chain), chr(65 + int(chain))) + for chain in unique_chains + }, + assembly_composition=self.metadata.assembly_composition, + ) + return ProteinComplex( + id=self.id, + sequence=sequence, + entity_id=arrays["entity_id"], + chain_id=arrays["chain_id"], + sym_id=arrays["sym_id"], + residue_index=arrays["residue_index"], + insertion_code=arrays["insertion_code"], + atom37_positions=arrays["atom37_positions"], + atom37_mask=arrays["atom37_mask"], + confidence=arrays["confidence"], + metadata=metadata, + ) + + @classmethod + def from_mmcif(cls, inp: str, id: str | None = None) -> MolecularComplex: + cif_file = _read_cif(inp) + structure = _read_structure(cif_file) + if TYPE_CHECKING: + structure: Any = structure + labels = _label_asym_ids(cif_file, len(structure)) + grouped = _group_structure_atoms(structure, labels) + ( + tokens, + positions, + elements, + names, + hetero, + spans, + confidences, + token_chains, + chain_numbers, + ) = _flatten_structure_groups(grouped) + n_tokens = len(tokens) + if positions: + position_array = np.asarray(positions, dtype=np.float32) + element_array = np.asarray(elements, dtype=object) + name_array = np.asarray(names, dtype=object) + hetero_array = np.asarray(hetero, dtype=bool) + span_array = np.asarray(spans, dtype=np.int32) + chain_array = np.asarray(token_chains, dtype=np.int64) + else: + position_array = np.zeros((0, 3), dtype=np.float32) + element_array = np.zeros(0, dtype=object) + name_array = np.zeros(0, dtype=object) + hetero_array = np.zeros(0, dtype=bool) + span_array = np.zeros((n_tokens, 2), dtype=np.int32) + chain_array = ( + np.asarray(token_chains, dtype=np.int64) + if token_chains + else np.zeros(n_tokens, dtype=np.int64) + ) + complex_id = id or (Path(inp).stem if os.path.exists(inp) else "complex_from_string") + return cls( + id=complex_id, + sequence=tokens, + atom_positions=position_array, + atom_elements=element_array, + token_to_atoms=span_array, + chain_id=chain_array, + plddt=np.asarray(confidences, dtype=np.float32), + metadata=MolecularComplexMetadata( + entity_lookup=_entity_metadata(cif_file), + chain_lookup={number: chain for chain, number in chain_numbers.items()}, + assembly_composition=None, + ), + atom_names=name_array, + atom_hetero=hetero_array, + ) + + def _get_entity_mapping( + self, + ) -> tuple[dict[str, list[str]], dict[str, int], dict[int, tuple[str, ...]]]: + return _chain_entity_maps(self) + + def _add_entity_information( + self, cif_file: CIFFile, entity_sequences: dict[int, tuple[str, ...]] + ) -> None: + _add_entity_categories(cif_file, self, entity_sequences) + + def to_mmcif(self) -> str: + _, chain_entities, entity_sequences = _chain_entity_maps(self) + atom_array = _as_atom_array(self, chain_entities) + cif_file = CIFFile() + set_structure(cif_file, atom_array, data_block=self.id) + _repair_label_entity_ids(cif_file, chain_entities) + _add_entity_categories(cif_file, self, entity_sequences) + round_mmcif_columns(cif_file) + output = io.StringIO() + cif_file.write(output) + return output.getvalue() + + def dockq(self, native: MolecularComplex) -> Any: + try: + mobile = self.to_protein_complex().normalize_chain_ids_for_pdb() + target = native.to_protein_complex().normalize_chain_ids_for_pdb() + except ValueError as error: + raise ValueError( + f"Cannot convert MolecularComplex to ProteinComplex for DockQ: {error}" + ) from None + try: + return mobile.dockq(target) + except Exception: + return self._compute_dockq_manual(native) + + def _compute_dockq_manual(self, native: MolecularComplex) -> Any: + try: + mobile = self.to_protein_complex().normalize_chain_ids_for_pdb() + target = native.to_protein_complex().normalize_chain_ids_for_pdb() + except ValueError as error: + raise ValueError( + f"Cannot convert MolecularComplex to ProteinComplex for DockQ: {error}" + ) from None + with TemporaryDirectory() as directory: + mobile_path = Path(directory) / "self.pdb" + target_path = Path(directory) / "native.pdb" + mobile.to_pdb(mobile_path) + target.to_pdb(target_path) + try: + raw_output = check_output(["DockQ", str(mobile_path), str(target_path)]) + output = raw_output.decode() + score: float | None = None + for line in output.split("\n"): + if "Total DockQ" in line: + match = re.search(r"Total DockQ.*: ([\d.]+)", line) + if match: + score = float(match.group(1)) + break + if score is None: + for line in output.split("\n"): + if line.startswith("DockQ") and ":" in line: + try: + score = float(line.split(":")[1].strip()) + break + except (ValueError, IndexError): + continue + if score is None: + raise ValueError("Could not parse DockQ score from output") + return {"total_dockq": score, "raw_output": output, "aligned": self} + except FileNotFoundError: + raise RuntimeError( + "DockQ is not installed. Please install DockQ to use this method." + ) from None + except Exception as error: + raise RuntimeError(f"DockQ computation failed: {error}") from error + + def rmsd(self, target: MolecularComplex, **kwargs: Any) -> float: + mobile, reference, mask = _centroid_tensors(self, target, retain_missing=False) + value = compute_rmsd( + mobile=mobile, + target=reference, + atom_exists_mask=mask, + reduction="batch", + **kwargs, + ) + return float(value) + + def lddt_ca(self, target: MolecularComplex, **kwargs: Any) -> float: + mobile, reference, mask = _centroid_tensors(self, target, retain_missing=True) + value = compute_lddt( + all_atom_pred_pos=mobile, + all_atom_positions=reference, + all_atom_mask=mask, + per_residue=False, + **kwargs, + ) + return float(value) + + def state_dict(self) -> dict[str, Any]: + state = dict(vars(self)) + for optional_identity in ("entity_id", "sym_id"): + if state[optional_identity] is None: + state.pop(optional_identity) + for key, value in tuple(state.items()): + if isinstance(value, MolecularComplexMetadata): + state[key] = asdict(value) + elif isinstance(value, np.ndarray): + if value.dtype == np.int64: + value = value.astype(np.int32) + elif value.dtype in (np.dtype(np.float64), np.dtype(np.float32)): + value = value.astype(np.float16) + state[key] = value.tolist() + return state + + def to_blob(self) -> bytes: + return brotli.compress(msgpack.dumps(self.state_dict()), quality=5) + + @classmethod + def from_state_dict(cls, dct: dict[str, Any]) -> MolecularComplex: + dct = dict(dct) + for key, value in tuple(dct.items()): + if isinstance(value, list) and key in _SERIALIZED_ARRAYS: + dct[key] = np.asarray(value) + for key, value in tuple(dct.items()): + if not isinstance(value, np.ndarray): + continue + if key in {"atom_positions", "plddt"}: + dct[key] = value.astype(np.float32) + elif key == "token_to_atoms": + dct[key] = value.astype(np.int32) + elif key in {"chain_id", "entity_id", "sym_id"}: + dct[key] = value.astype(np.int64) + dct["metadata"] = MolecularComplexMetadata(**dct["metadata"]) + if "chain_id" not in dct: + dct["chain_id"] = np.zeros(len(dct["sequence"]), dtype=np.int64) + return cls(**dct) + + @classmethod + def from_blob(cls, input: Path | str | io.BytesIO | bytes) -> MolecularComplex: + if isinstance(input, (Path, str)): + payload = Path(input).read_bytes() + elif isinstance(input, io.BytesIO): + payload = input.getvalue() + else: + payload = input + state = msgpack.loads(brotli.decompress(payload), strict_map_key=False) + return cls.from_state_dict(state) diff --git a/fastplms/models/esmfold2/esmfold2_msa.py b/fastplms/models/esmfold2/esmfold2_msa.py new file mode 100644 index 0000000000000000000000000000000000000000..56c625556d3096e5ac060e1f4571577043d9eb15 --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_msa.py @@ -0,0 +1,577 @@ +"""Multiple-sequence-alignment value objects and lossless encodings.""" + +from __future__ import annotations + +import dataclasses +import string +from collections.abc import Sequence +from dataclasses import dataclass +from functools import cached_property +from itertools import islice +from typing import Any + +import numpy as np +from Bio import SeqIO +from scipy.spatial.distance import cdist + +from .esmfold2_misc import slice_any_object +from .esmfold2_msa_filter_sequences import greedy_select_indices, hhfilter +from .esmfold2_parsing import FastaEntry, read_sequences, write_sequences +from .esmfold2_sequential_dataclass import SequentialDataclass +from .esmfold2_system import PathOrBuffer + +_A3M_INSERTION_DELETE_TABLE = str.maketrans( + dict.fromkeys(string.ascii_lowercase + ".") +) +_SERIALIZATION_VERSION = 1 +_UINT32_BYTES = 4 + + +def is_a3m_insertion(character: str) -> bool: + """Return whether a character is an A3M insertion marker.""" + + return character == "." or character.islower() + + +def remove_insertions_from_sequence(sequence: str) -> str: + """Remove lowercase residues and dot insertion markers from an A3M row.""" + + return sequence.translate(_A3M_INSERTION_DELETE_TABLE) + + +def a3m_deletion_counts(sequence: str) -> np.ndarray: + """Count insertions preceding each A3M match column.""" + + codes = np.frombuffer(sequence.encode("ascii"), dtype=np.uint8) + lowercase = (codes >= ord("a")) & (codes <= ord("z")) + insertion_mask = lowercase | (codes == ord(".")) + prefix_counts = np.concatenate(([0], np.cumsum(insertion_mask))) + match_positions = np.flatnonzero(~insertion_mask) + return np.diff(prefix_counts[match_positions], prepend=0) + + +def _parse_full_payload(data: bytes) -> tuple[np.ndarray, list[str]]: + version = int.from_bytes(data[:1], "little") + if version != _SERIALIZATION_VERSION: + raise ValueError(f"Unsupported version: {version}") + seqlen = int.from_bytes(data[1:5], "little") + depth = int.from_bytes(data[5:9], "little") + body = data[9:] + split = seqlen * depth + array = np.frombuffer(body[:split], dtype="|S1").reshape(depth, seqlen) + headers = [header for header in body[split:].decode().split("\n") if header] + if not headers and depth > 0: + headers = [""] * depth + return array, headers + + +def _parse_sequence_payload(data: bytes) -> np.ndarray: + seqlen = int.from_bytes(data[:_UINT32_BYTES], "little") + return np.frombuffer(data[_UINT32_BYTES:], dtype="|S1").reshape(-1, seqlen) + + +def _full_payload(array: np.ndarray, headers: Sequence[str]) -> bytes: + depth, seqlen = array.shape + prefix = b"".join( + ( + _SERIALIZATION_VERSION.to_bytes(1, "little"), + seqlen.to_bytes(_UINT32_BYTES, "little"), + depth.to_bytes(_UINT32_BYTES, "little"), + ) + ) + return prefix + array.tobytes() + "\n".join(headers).encode() + + +def _sequence_payload(array: np.ndarray) -> bytes: + return array.shape[1].to_bytes(_UINT32_BYTES, "little") + array.tobytes() + + +def _random_row_indices(depth: int, count: int) -> np.ndarray: + sampled = np.random.choice(depth - 1, count - 1, replace=False) + 1 + return np.sort(np.append(0, sampled)) + + +@dataclass(frozen=True) +class FastMSA(SequentialDataclass): + """An MSA stored as a two-dimensional NumPy byte array.""" + + array: np.ndarray + headers: list[str] | None = None + + def __post_init__(self) -> None: + if not isinstance(self.array, np.ndarray): + raise TypeError("FastMSA array must be a NumPy array.") + if self.array.ndim != 2 or self.array.shape[0] == 0 or self.array.shape[1] == 0: + raise ValueError( + f"FastMSA array must have non-empty shape (depth, length), got {self.array.shape}." + ) + if self.headers is not None and len(self.headers) != self.depth: + raise ValueError("Number of headers must match depth.") + + @property + def depth(self) -> int: + return self.array.shape[0] + + @property + def seqlen(self) -> int: + return self.array.shape[1] + + def __len__(self) -> int: + return self.seqlen + + @classmethod + def from_bytes(cls, data: bytes) -> FastMSA: + array, headers = _parse_full_payload(data) + return cls(array, headers) + + @classmethod + def from_sequence_bytes(cls, data: bytes) -> FastMSA: + return cls(_parse_sequence_payload(data)) + + def __getitem__( + self, + indices: int | list[int] | slice | np.ndarray, + ) -> FastMSA: + column_indices = [indices] if isinstance(indices, int) else indices + return dataclasses.replace(self, array=self.array[:, column_indices]) + + def select_sequences( + self, + indices: Sequence[int] | np.ndarray, + ) -> FastMSA: + headers = None + if self.headers is not None: + headers = [self.headers[index] for index in indices] + return dataclasses.replace( + self, + array=self.array[indices], + headers=headers, + ) + + def select_random_sequences(self, num_seqs: int) -> FastMSA: + if num_seqs >= self.depth: + return self + return self.select_sequences(_random_row_indices(self.depth, num_seqs)) + + def pad_to_depth(self, depth: int) -> FastMSA: + if depth < self.depth: + raise ValueError(f"Cannot pad to depth {depth} when depth is {self.depth}") + if depth == self.depth: + return self + row_count = depth - self.depth + pad_value = ord("-") if self.array.dtype == np.uint8 else b"-" + array = np.pad( + self.array, + ((0, row_count), (0, 0)), + constant_values=pad_value, + ) + headers = None if self.headers is None else self.headers + [""] * row_count + return dataclasses.replace(self, array=array, headers=headers) + + @classmethod + def concat( + cls, + msas: Sequence[FastMSA], + join_token: str | None = None, + allow_depth_mismatch: bool = False, + ) -> FastMSA: + if not msas: + raise ValueError("Cannot concatenate an empty list of MSAs") + if join_token not in (None, ""): + raise NotImplementedError("join_token is not supported for FastMSA") + depths = [msa.depth for msa in msas] + if len(set(depths)) != 1: + if not allow_depth_mismatch: + raise ValueError("Depth mismatch in concatenating MSAs") + maximum_depth = max(depths) + msas = [msa.pad_to_depth(maximum_depth) for msa in msas] + header_columns = ( + msa.headers if msa.headers is not None else [""] * msa.depth for msa in msas + ) + headers = [ + "|".join(str(header) for header in row) for row in zip(*header_columns, strict=False) + ] + return cls( + np.concatenate([msa.array for msa in msas], axis=1), + headers, + ) + + @classmethod + def stack( + cls, + msas: Sequence[FastMSA], + remove_query_from_later_msas: bool = True, + ) -> FastMSA: + if not msas: + raise ValueError("Cannot stack an empty list of MSAs") + arrays: list[np.ndarray] = [] + headers: list[str] | None = [] if any(msa.headers is not None for msa in msas) else None + for index, msa in enumerate(msas): + start = 1 if index > 0 and remove_query_from_later_msas else 0 + arrays.append(msa.array[start:]) + if headers is not None: + source_headers = msa.headers or [""] * msa.depth + headers.extend(source_headers[start:]) + return cls(np.concatenate(arrays, axis=0), headers) + + def to_msa(self) -> MSA: + headers = self.headers + if headers is None: + headers = [f"seq{index}" for index in range(self.depth)] + entries = [ + FastaEntry(header, b"".join(row).decode()) + for header, row in zip(headers, self.array, strict=False) + ] + return MSA(entries) + + +@dataclass(frozen=True) +class MSA(SequentialDataclass): + """An ordered set of aligned protein sequences and optional A3M metadata.""" + + entries: list[FastaEntry] + deletions: np.ndarray | None = dataclasses.field(default=None, compare=False) + + def __post_init__(self) -> None: + if not isinstance(self.entries, list): + raise TypeError("MSA entries must be a list of FastaEntry rows.") + if not self.entries: + raise ValueError("MSA requires at least one aligned sequence.") + if any(not isinstance(entry, FastaEntry) for entry in self.entries): + raise TypeError("Every MSA entry must be a FastaEntry.") + expected_length = len(self.entries[0].sequence) + if expected_length == 0: + raise ValueError("MSA sequences must be non-empty.") + for row, entry in enumerate(self.entries[1:], start=1): + if len(entry.sequence) != expected_length: + raise ValueError( + "MSA row length mismatch: " + f"row 0 has {expected_length} columns, row {row} has " + f"{len(entry.sequence)}." + ) + deletions = self.deletions + if deletions is not None and not isinstance(deletions, np.ndarray): + raise TypeError("MSA deletions must be a NumPy array when provided.") + if isinstance(deletions, np.ndarray) and deletions.shape != ( + len(self.entries), + expected_length, + ): + raise ValueError( + "MSA deletion matrix must have shape " + f"({len(self.entries)}, {expected_length}), got {deletions.shape}." + ) + + @cached_property + def sequences(self) -> list[str]: + return [entry.sequence for entry in self.entries] + + @cached_property + def headers(self) -> list[str]: + return [entry.header for entry in self.entries] + + @property + def depth(self) -> int: + return len(self.entries) + + @property + def seqlen(self) -> int: + return len(self.entries[0].sequence) + + @property + def query(self) -> str: + return self.entries[0].sequence + + @cached_property + def array(self) -> np.ndarray: + return np.array([list(sequence) for sequence in self.sequences], dtype="|S1") + + @cached_property + def seqid(self) -> np.ndarray: + byte_array = self.array.view(np.uint8) + return (1 - cdist(byte_array[0][None], byte_array, "hamming"))[0] + + def __len__(self) -> int: + return self.seqlen + + def __repr__(self) -> str: + return f"MSA({self.entries[0].header}: Depth={self.depth}, Length={self.seqlen})" + + @classmethod + def from_a3m( + cls, + path: PathOrBuffer, + remove_insertions: bool = True, + max_sequences: int | None = None, + ) -> MSA: + entries = [] + deletion_rows = [] + for header, raw_sequence in islice(read_sequences(path), max_sequences): + if remove_insertions: + deletion_rows.append(a3m_deletion_counts(raw_sequence)) + sequence = ( + remove_insertions_from_sequence(raw_sequence) if remove_insertions else raw_sequence + ) + if entries: + expected_length = len(entries[0].sequence) + if len(sequence) != expected_length: + raise ValueError( + "Sequence length mismatch. " + f"Expected: {expected_length}, Received: {len(sequence)}" + ) + entries.append(FastaEntry(header, sequence)) + deletions = None + if remove_insertions and deletion_rows: + deletions = np.stack(deletion_rows).astype(np.float32) + return cls(entries, deletions=deletions) + + @classmethod + def from_stockholm( + cls, + path: PathOrBuffer, + remove_insertions: bool = True, + max_sequences: int | None = None, + ) -> MSA: + entries = [] + for record in islice(SeqIO.parse(path, "stockholm"), max_sequences): + sequence = str(record.seq) + if entries: + expected_length = len(entries[0].sequence) + if len(sequence) != expected_length: + raise ValueError( + "Sequence length mismatch. " + f"Expected: {expected_length}, Received: {len(sequence)}" + ) + entries.append(FastaEntry(f"{record.id} {record.description}", sequence)) + msa = cls(entries) + if remove_insertions: + msa = msa.select_positions( + [index for index, residue in enumerate(msa.query) if residue != "-"] + ) + return msa + + @classmethod + def from_sequences( + cls, + sequences: list[str], + remove_insertions: bool = False, + ) -> MSA: + transform = ( + remove_insertions_from_sequence if remove_insertions else lambda sequence: sequence + ) + return cls([FastaEntry("", transform(sequence)) for sequence in sequences]) + + @classmethod + def from_bytes(cls, data: bytes) -> MSA: + array, headers = _parse_full_payload(data) + return cls( + [ + FastaEntry(header, b"".join(row).decode()) + for header, row in zip(headers, array, strict=False) + ] + ) + + @classmethod + def from_sequence_bytes(cls, data: bytes) -> MSA: + array = _parse_sequence_payload(data) + return cls([FastaEntry("", b"".join(row).decode()) for row in array]) + + @classmethod + def from_state_dict(cls, dct: dict[str, Any]) -> MSA: + deletions = dct.get("deletions") + return cls( + [FastaEntry("", sequence) for sequence in dct["sequences"]], + deletions=(None if deletions is None else np.asarray(deletions, dtype=np.float32)), + ) + + def to_a3m(self, path: PathOrBuffer) -> None: + write_sequences(self.entries, path) + + def to_fast_msa(self) -> FastMSA: + return FastMSA(self.array, self.headers) + + def to_bytes(self) -> bytes: + return _full_payload(self.array, self.headers) + + def to_sequence_bytes(self) -> bytes: + """Serialize aligned sequences without their headers.""" + + return _sequence_payload(self.array) + + def state_dict(self, json_serializable: bool = False) -> dict[str, Any]: + result: dict[str, Any] = {"sequences": self.sequences} + if self.deletions is not None: + result["deletions"] = self.deletions.tolist() if json_serializable else self.deletions + return result + + def _aligned_deletions(self) -> np.ndarray | None: + if self.deletions is None: + return None + if self.deletions.shape != (self.depth, self.seqlen): + return None + return self.deletions + + def _select_deletion_columns(self, indices) -> np.ndarray | None: + if self.deletions is None or self.deletions.shape[1] != self.seqlen: + return None + return self.deletions[:, indices] + + def select_sequences( + self, + indices: Sequence[int] | np.ndarray, + ) -> MSA: + deletions = None if self.deletions is None else self.deletions[np.asarray(indices)] + return dataclasses.replace( + self, + entries=[self.entries[index] for index in indices], + deletions=deletions, + ) + + def select_positions( + self, + indices: Sequence[int] | np.ndarray, + ) -> MSA: + entries = [ + FastaEntry( + entry.header, + "".join(entry.sequence[index] for index in indices), + ) + for entry in self.entries + ] + return dataclasses.replace( + self, + entries=entries, + deletions=self._select_deletion_columns(indices), + ) + + def __getitem__( + self, + indices: int | list[int] | slice | np.ndarray, + ) -> MSA: + column_indices = [indices] if isinstance(indices, int) else indices + entries = [ + FastaEntry( + entry.header, + slice_any_object(entry.sequence, column_indices), + ) + for entry in self.entries + ] + return dataclasses.replace( + self, + entries=entries, + deletions=self._select_deletion_columns(column_indices), + ) + + def greedy_select(self, num_seqs: int, mode: str = "max") -> MSA: + if mode not in ("max", "min"): + raise ValueError(f"Unsupported MSA selection mode: {mode!r}.") + if self.depth <= num_seqs: + return self + return self.select_sequences(greedy_select_indices(self.array, num_seqs, mode)) + + def hhfilter( + self, + seqid: int = 90, + diff: int = 0, + cov: int = 0, + qid: int = 0, + qsc: float = -20.0, + binary: str = "hhfilter", + ) -> MSA: + indices = hhfilter( + self.sequences, + seqid=seqid, + diff=diff, + cov=cov, + qid=qid, + qsc=qsc, + binary=binary, + ) + return self.select_sequences(indices) + + def select_random_sequences(self, num_seqs: int) -> MSA: + if num_seqs >= self.depth: + return self + return self.select_sequences(_random_row_indices(self.depth, num_seqs)) + + def select_diverse_sequences(self, num_seqs: int) -> MSA: + if num_seqs >= self.depth: + return self + filtered = self.hhfilter(diff=num_seqs) + if num_seqs < filtered.depth: + filtered = filtered.select_random_sequences(num_seqs) + return filtered + + def pad_to_depth(self, depth: int) -> MSA: + if depth < self.depth: + raise ValueError(f"Cannot pad to depth {depth} when depth is {self.depth}") + if depth == self.depth: + return self + count = depth - self.depth + extra = [FastaEntry("", "-" * self.seqlen) for _ in range(count)] + deletions = self._aligned_deletions() + if deletions is not None: + zero_rows = np.zeros((count, self.seqlen), dtype=deletions.dtype) + deletions = np.concatenate((deletions, zero_rows), axis=0) + return dataclasses.replace( + self, + entries=self.entries + extra, + deletions=deletions, + ) + + @classmethod + def stack( + cls, + msas: Sequence[MSA], + remove_query_from_later_msas: bool = True, + ) -> MSA: + entries = [] + deletion_arrays = [] + for index, msa in enumerate(msas): + start = 1 if index > 0 and remove_query_from_later_msas else 0 + entries.extend(msa.entries[start:]) + aligned = msa._aligned_deletions() + if aligned is not None: + deletion_arrays.append(aligned[start:]) + deletions = None + if ( + len(deletion_arrays) == len(msas) + and len({array.shape[1] for array in deletion_arrays}) == 1 + ): + deletions = np.concatenate(deletion_arrays, axis=0) + return cls(entries=entries, deletions=deletions) + + @classmethod + def concat( + cls, + msas: Sequence[MSA], + join_token: str | None = "|", + allow_depth_mismatch: bool = False, + ) -> MSA: + if not msas: + raise ValueError("Cannot concatenate an empty list of MSAs") + depths = [msa.depth for msa in msas] + if len(set(depths)) != 1: + if not allow_depth_mismatch: + raise ValueError("Depth mismatch in concatenating MSAs") + maximum_depth = max(depths) + msas = [msa.pad_to_depth(maximum_depth) for msa in msas] + headers = [ + "|".join(str(header) for header in row) + for row in zip(*(msa.headers for msa in msas), strict=False) + ] + separator = "" if join_token is None else join_token + sequences = [ + separator.join(row) for row in zip(*(msa.sequences for msa in msas), strict=False) + ] + deletions = None + if separator == "": + arrays = [msa._aligned_deletions() for msa in msas] + if all(array is not None for array in arrays): + deletions = np.concatenate(arrays, axis=1) # type: ignore[arg-type] + return cls( + [ + FastaEntry(header, sequence) + for header, sequence in zip(headers, sequences, strict=False) + ], + deletions=deletions, + ) diff --git a/fastplms/models/esmfold2/esmfold2_msa_filter_sequences.py b/fastplms/models/esmfold2/esmfold2_msa_filter_sequences.py new file mode 100644 index 0000000000000000000000000000000000000000..bc13db7bee2398d291265e6f5a2e95752b2be61e --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_msa_filter_sequences.py @@ -0,0 +1,111 @@ +"""Sequence selection for multiple-sequence alignments.""" + +from __future__ import annotations + +import os +import tempfile +from pathlib import Path + +import numpy as np + +from .esmfold2_system import run_subprocess_with_errorcheck + + +def _byte_matrix(array: np.ndarray) -> np.ndarray: + """Return a two-dimensional byte view used for Hamming comparisons.""" + + matrix = np.asarray(array).view(np.uint8) + return matrix.reshape(matrix.shape[0], -1) + + +def _hamming_to_all(query: np.ndarray, sequences: np.ndarray) -> np.ndarray: + return np.not_equal(sequences, query).mean(axis=1, dtype=np.float64) + + +def greedy_select_indices(array: np.ndarray, num_seqs: int, mode: str = "max") -> list[int]: + """Select MSA rows by greedy mean Hamming distance from the query row. + + Row zero is always retained. At each step the selector chooses the remaining + row with greatest distance for ``mode="max"`` or least distance for + ``mode="min"``. Returned indices follow source order. + """ + + if not isinstance(array, np.ndarray): + raise TypeError("array must be a NumPy array") + if array.ndim != 2 or array.shape[0] == 0 or array.shape[1] == 0: + raise ValueError( + f"array must have non-empty shape (depth, length), got {array.shape}" + ) + if isinstance(num_seqs, bool) or not isinstance(num_seqs, int): + raise TypeError("num_seqs must be an integer") + if num_seqs <= 0: + raise ValueError("num_seqs must be greater than zero") + if not isinstance(mode, str): + raise TypeError("mode must be a string") + if mode not in {"max", "min"}: + raise ValueError(f"unsupported selection mode: {mode}") + depth = array.shape[0] + if depth <= num_seqs: + return list(range(depth)) + + sequences = _byte_matrix(array) + selected = [0] + available = np.ones(depth, dtype=bool) + available[0] = False + distance_sum = _hamming_to_all(sequences[0], sequences) + choose = np.argmax if mode == "max" else np.argmin + + while len(selected) < num_seqs: + candidates = np.flatnonzero(available) + candidate_scores = distance_sum[candidates] / len(selected) + next_index = int(candidates[int(choose(candidate_scores))]) + selected.append(next_index) + available[next_index] = False + distance_sum += _hamming_to_all(sequences[next_index], sequences) + return sorted(selected) + + +def _temporary_root() -> str | None: + shared_memory = Path("/dev/shm") + return os.fspath(shared_memory) if shared_memory.is_dir() else None + + +def hhfilter( + sequences: list[str], + seqid: int = 90, + diff: int = 0, + cov: int = 0, + qid: int = 0, + qsc: float = -20.0, + binary: str = "hhfilter", +) -> list[int]: + """Run HH-suite filtering and return source indices from its FASTA headers.""" + + with tempfile.TemporaryDirectory(dir=_temporary_root()) as directory: + work = Path(directory) + source_path = work / "input.fasta" + result_path = work / "output.fasta" + records = (f">{index}\n{sequence}" for index, sequence in enumerate(sequences)) + source_path.write_text("\n".join(records), encoding="utf-8") + command = [ + binary, + "-i", + os.fspath(source_path), + "-M", + "a3m", + "-o", + os.fspath(result_path), + "-id", + str(seqid), + "-diff", + str(diff), + "-cov", + str(cov), + "-qid", + str(qid), + "-qsc", + str(qsc), + ] + run_subprocess_with_errorcheck(command, capture_output=True) + headers = result_path.read_text(encoding="utf-8").splitlines() + return [int(line[1:].strip()) for line in headers if line.startswith(">")] diff --git a/fastplms/models/esmfold2/esmfold2_normalize_coordinates.py b/fastplms/models/esmfold2/esmfold2_normalize_coordinates.py new file mode 100644 index 0000000000000000000000000000000000000000..45e065247f9b1ac048c697d9b10c0a14605ee8e9 --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_normalize_coordinates.py @@ -0,0 +1,67 @@ +"""Rigid-frame normalization for atom37 coordinates.""" + +from __future__ import annotations + +from typing import TypeVar + +import numpy as np +import torch +from torch import Tensor + +from . import esmfold2_residue_constants as residue_constants +from .esmfold2_affine3d import Affine3D + +ArrayOrTensor = TypeVar("ArrayOrTensor", np.ndarray, Tensor) + + +def atom3_to_backbone_frames(bb_positions: Tensor) -> Affine3D: + """Construct a frame from N, C-alpha, and C positions in ``X``.""" + + n_position, ca_position, c_position = bb_positions.unbind(dim=-2) + return Affine3D.from_graham_schmidt(c_position, ca_position, n_position) + + +def index_by_atom_name( + atom37: ArrayOrTensor, + atom_names: str | list[str], + dim: int = -2, +) -> ArrayOrTensor: + """Select one or more named atoms along an atom37 axis.""" + + single_atom = isinstance(atom_names, str) + names = [atom_names] if single_atom else atom_names + indices = [residue_constants.atom_order[name] for name in names] + axis = dim % atom37.ndim + if isinstance(atom37, Tensor): + index = torch.tensor(indices, dtype=torch.long, device=atom37.device) + selected = torch.index_select(atom37, axis, index) + else: + selected = np.take(atom37, indices, axis=axis) + return selected.squeeze(axis) if single_atom else selected # type: ignore[return-value] + + +def get_protein_normalization_frame(coords: Tensor) -> Affine3D: + """Build one frame from backbone coordinates ``X`` with shape (l, 37, 3).""" + + backbone = index_by_atom_name(coords, ["N", "CA", "C"], dim=-2) + residue_is_valid = torch.isfinite(backbone).all(dim=-1).all(dim=-1) + weights = residue_is_valid[..., None, None] + coordinate_sum = backbone.masked_fill(~weights, 0).sum(dim=-3) + count = residue_is_valid.sum(dim=-1)[..., None, None] + mean_backbone = coordinate_sum / (count + 1e-8) + return atom3_to_backbone_frames(mean_backbone.float()) + + +def apply_frame_to_coords(coords: Tensor, frame: Affine3D) -> Tensor: + """Express atom coordinates ``X`` in the inverse of ``frame``.""" + + transformed = frame[..., None, None].invert().apply(coords) + frame_is_valid = frame.trans.norm(dim=-1) > 0 + normalized = torch.where(frame_is_valid[..., None, None, None], transformed, coords) + return normalized.masked_fill(torch.isinf(coords), torch.inf) + + +def normalize_coordinates(coords: Tensor) -> Tensor: + """Normalize ``X`` with shape (..., l, 37, 3) to its backbone frame.""" + + return apply_frame_to_coords(coords, get_protein_normalization_frame(coords)) diff --git a/fastplms/models/esmfold2/esmfold2_output.py b/fastplms/models/esmfold2/esmfold2_output.py new file mode 100644 index 0000000000000000000000000000000000000000..97c846a73358dc5ab25506037ec77f6786441ef4 --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_output.py @@ -0,0 +1,201 @@ +"""Convert ESMFold2 coordinate tensors into molecular-complex records.""" + +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import dataclass, field +from itertools import groupby +from typing import Any + +import numpy as np +import torch + +from .esmfold2_constants import ELEMENT_NUMBER_TO_SYMBOL, MOL_TYPE_NONPOLYMER +from .esmfold2_molecular_complex import MolecularComplex, MolecularComplexMetadata + + +def get_element_symbol(atomic_number: int) -> str: + """Map a training-time atomic number to a chemical symbol.""" + + return ELEMENT_NUMBER_TO_SYMBOL.get(atomic_number, "X") + + +def _decode_atom_name(encoded_name: Any) -> str: + values = encoded_name.tolist() if hasattr(encoded_name, "tolist") else encoded_name + return "".join(chr(int(value) + 32) for value in values if int(value)).strip() + + +@dataclass +class _ComplexRecords: + sequence: list[str] = field(default_factory=list) + chain_ids: list[int] = field(default_factory=list) + token_to_atoms: list[list[int]] = field(default_factory=list) + confidence: list[float] = field(default_factory=list) + positions: list[list[float]] = field(default_factory=list) + elements: list[str] = field(default_factory=list) + atom_names: list[str] = field(default_factory=list) + atom_hetero: list[bool] = field(default_factory=list) + chain_lookup: dict[int, str] = field(default_factory=dict) + entity_lookup: dict[int, str] = field(default_factory=dict) + + def add_token( + self, + *, + residue_name: str, + asym_id: int, + plddt: float, + atoms: Iterable[tuple[list[float], str, str]], + hetero: bool, + ) -> None: + atom_start = len(self.positions) + for position, element, atom_name in atoms: + self.positions.append(position) + self.elements.append(element) + self.atom_names.append(atom_name) + self.atom_hetero.append(hetero) + self.sequence.append(residue_name) + self.chain_ids.append(asym_id) + self.confidence.append(plddt) + self.token_to_atoms.append([atom_start, len(self.positions)]) + + def build(self, complex_id: str) -> MolecularComplex: + return MolecularComplex( + id=complex_id, + sequence=self.sequence, + atom_positions=np.asarray(self.positions, dtype=np.float32).reshape(-1, 3), + atom_elements=np.asarray(self.elements, dtype=object), + token_to_atoms=np.asarray(self.token_to_atoms, dtype=np.int32).reshape(-1, 2), + chain_id=np.asarray(self.chain_ids, dtype=np.int64), + plddt=np.asarray(self.confidence, dtype=np.float32), + atom_names=np.asarray(self.atom_names, dtype=object), + atom_hetero=np.asarray(self.atom_hetero, dtype=bool), + metadata=MolecularComplexMetadata( + entity_lookup=self.entity_lookup, + chain_lookup=self.chain_lookup, + assembly_composition=None, + ), + ) + + +def build_molecular_complex_from_features( + coords: torch.Tensor, + plddt: torch.Tensor, + atom_mask: torch.Tensor, + ref_element: torch.Tensor, + ref_atom_name_chars: torch.Tensor, + chain_infos: list[Any], + complex_id: str, +) -> MolecularComplex: + """Decode model features into one complex without intermediate structure files. + + Protein, DNA, and RNA tokens are grouped by residue index. Ligand atom + tokens are collapsed into one non-polymer residue per chain. + """ + + M = atom_mask.bool().cpu().numpy() + X = coords.float().cpu().numpy() + atom_names = ref_atom_name_chars.cpu().numpy() + elements = ref_element.cpu().numpy() + confidence = plddt.float().cpu().numpy() + records = _ComplexRecords() + + def decode_atoms(tokens: Iterable[Any]): + for token in tokens: + for atom_index in range(token.atom_start, token.atom_start + token.atom_count): + if M[atom_index]: + yield ( + X[atom_index].tolist(), + get_element_symbol(int(elements[atom_index])), + _decode_atom_name(atom_names[atom_index]), + ) + + for chain in chain_infos: + is_nonpolymer = chain.mol_type == MOL_TYPE_NONPOLYMER + records.chain_lookup[chain.asym_id] = chain.chain_id + records.entity_lookup[chain.entity_id] = "non-polymer" if is_nonpolymer else "polymer" + + if is_nonpolymer: + mean_confidence = ( + float(np.mean([confidence[token.token_index] for token in chain.tokens])) + if chain.tokens + else 0.0 + ) + records.add_token( + residue_name=chain.tokens[0].residue_name if chain.tokens else "LIG", + asym_id=chain.asym_id, + plddt=mean_confidence, + atoms=decode_atoms(chain.tokens), + hetero=True, + ) + continue + + residue_groups = groupby(chain.tokens, key=lambda token: token.residue_index) + for _residue_index, group in residue_groups: + residue_tokens = list(group) + records.add_token( + residue_name=residue_tokens[0].residue_name, + asym_id=chain.asym_id, + plddt=float(np.mean([confidence[token.token_index] for token in residue_tokens])), + atoms=decode_atoms(residue_tokens), + hetero=False, + ) + + return records.build(complex_id) + + +def build_molecular_complex( + structure: Any, + coords: torch.Tensor, + plddt: torch.Tensor, + complex_id: str, +) -> MolecularComplex: + """Decode coordinates using the atom and residue arrays of a prepared structure.""" + + records = _ComplexRecords() + coordinate_index = 0 + confidence_index = 0 + + for chain in structure.chains: + asym_id = int(chain["asym_id"]) + mol_type = int(chain["mol_type"]) + is_nonpolymer = mol_type == MOL_TYPE_NONPOLYMER + records.chain_lookup[asym_id] = str(chain["name"]) + records.entity_lookup[int(chain["entity_id"])] = ( + "non-polymer" if is_nonpolymer else "polymer" + ) + + residue_start = int(chain["res_idx"]) + residue_stop = residue_start + int(chain["res_num"]) + for residue in structure.residues[residue_start:residue_stop]: + atom_start = int(residue["atom_idx"]) + atom_stop = atom_start + int(residue["atom_num"]) + decoded_atoms: list[tuple[list[float], str, str]] = [] + for atom in structure.atoms[atom_start:atom_stop]: + if not atom["is_present"]: + continue + decoded_atoms.append( + ( + coords[coordinate_index].tolist(), + get_element_symbol(int(atom["element"].item())), + _decode_atom_name(atom["name"]), + ) + ) + coordinate_index += 1 + + records.add_token( + residue_name=str(residue["name"]), + asym_id=asym_id, + plddt=float(plddt[confidence_index].item()), + atoms=decoded_atoms, + hetero=is_nonpolymer, + ) + confidence_index += 1 + + return records.build(complex_id) + + +__all__ = [ + "build_molecular_complex", + "build_molecular_complex_from_features", + "get_element_symbol", +] diff --git a/fastplms/models/esmfold2/esmfold2_paired_msa.py b/fastplms/models/esmfold2/esmfold2_paired_msa.py new file mode 100644 index 0000000000000000000000000000000000000000..44f050b37c83ae70cf70b532380eb11c94a608a3 --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_paired_msa.py @@ -0,0 +1,282 @@ +"""Construct taxonomy-paired MSA features for multichain folding.""" + +from __future__ import annotations + +import re +from dataclasses import dataclass + +import numpy as np + +from .esmfold2_constants import ( + MSA_GAP_TOKEN_ID, + PROTEIN_3TO1, + PROTEIN_RESIDUE_TO_RES_TYPE, + PROTEIN_UNK_RES_TYPE, +) +from .esmfold2_msa import MSA + +_TAXONOMY_PATTERN = re.compile(r"key=(-?\d+)") + + +def protein_letter_to_res_type() -> dict[str, int]: + """Return the one-letter residue vocabulary used by the MSA encoder.""" + + vocabulary = { + one_letter: PROTEIN_RESIDUE_TO_RES_TYPE[three_letter] + for three_letter, one_letter in PROTEIN_3TO1.items() + if three_letter in PROTEIN_RESIDUE_TO_RES_TYPE + } + vocabulary.update({"-": MSA_GAP_TOKEN_ID, "X": PROTEIN_UNK_RES_TYPE}) + return vocabulary + + +def _taxonomy_from_header(header: str) -> int: + match = _TAXONOMY_PATTERN.search(header) if header else None + return int(match.group(1)) if match is not None else -1 + + +def _emitted_length(sequence: str) -> int: + return sum(character != "." and not character.islower() for character in sequence) + + +def _decode_a3m_row( + sequence: str, + sequence_length: int, + vocabulary: dict[str, int], +) -> tuple[np.ndarray, np.ndarray]: + residues = np.full(sequence_length, MSA_GAP_TOKEN_ID, dtype=np.int64) + deletions = np.zeros(sequence_length, dtype=np.float32) + column = 0 + insertion_count = 0 + for character in sequence: + if character == "." or character.islower(): + insertion_count += 1 + continue + if column == sequence_length: + break + residues[column] = ( + MSA_GAP_TOKEN_ID + if character == "-" + else vocabulary.get(character.upper(), PROTEIN_UNK_RES_TYPE) + ) + if insertion_count: + deletions[column] = float(insertion_count) + insertion_count = 0 + column += 1 + return residues, deletions + + +def msa_to_res_type_and_deletions( + msa: MSA, + letter_to_res_type: dict[str, int], +) -> tuple[np.ndarray, np.ndarray]: + """Decode an A3M alignment into arrays ``X`` and ``D`` with shape (m, l).""" + + sequence_length = _emitted_length(msa.entries[0].sequence) + residue_rows: list[np.ndarray] = [] + deletion_rows: list[np.ndarray] = [] + for entry in msa.entries: + residues, deletions = _decode_a3m_row( + entry.sequence, + sequence_length, + letter_to_res_type, + ) + residue_rows.append(residues) + deletion_rows.append(deletions) + return np.stack(residue_rows), np.stack(deletion_rows) + + +@dataclass(frozen=True) +class _ChainAlignment: + residues: np.ndarray + deletions: np.ndarray + taxonomies: list[int] + + +def _chain_alignment( + msa: MSA | None, + query_res_types: np.ndarray, + vocabulary: dict[str, int], +) -> _ChainAlignment: + if msa is None or msa.depth == 0: + return _ChainAlignment( + residues=query_res_types[None, :], + deletions=np.zeros((1, query_res_types.shape[0]), dtype=np.float32), + taxonomies=[-1], + ) + residues, deletions = msa_to_res_type_and_deletions(msa, vocabulary) + taxonomies = [_taxonomy_from_header(entry.header) for entry in msa.entries] + return _ChainAlignment(residues, deletions, taxonomies) + + +def _taxonomy_groups( + chain_ids: list[int], + alignments: dict[int, _ChainAlignment], +) -> dict[int, list[tuple[int, int]]]: + groups: dict[int, list[tuple[int, int]]] = {} + for chain_id in chain_ids: + for row, taxonomy in enumerate(alignments[chain_id].taxonomies): + if row and taxonomy != -1: + groups.setdefault(taxonomy, []).append((chain_id, row)) + return {taxonomy: rows for taxonomy, rows in groups.items() if len(rows) > 1} + + +def _available_rows( + chain_ids: list[int], + alignments: dict[int, _ChainAlignment], + groups: dict[int, list[tuple[int, int]]], +) -> dict[int, list[int]]: + used = {row for group in groups.values() for row in group} + return { + chain_id: [ + row + for row in range(1, len(alignments[chain_id].taxonomies)) + if (chain_id, row) not in used + ] + for chain_id in chain_ids + } + + +def _append_taxonomy_rows( + rows: list[dict[int, int]], + paired_flags: list[dict[int, int]], + chain_ids: list[int], + groups: dict[int, list[tuple[int, int]]], + available: dict[int, list[int]], + max_pairs: int, +) -> None: + ordered_groups = sorted( + groups.values(), + key=lambda group: len({chain_id for chain_id, _row in group}), + reverse=True, + ) + for group in ordered_groups: + rows_by_chain: dict[int, list[int]] = {} + for chain_id, row in group: + rows_by_chain.setdefault(chain_id, []).append(row) + for occurrence in range(max(map(len, rows_by_chain.values()))): + selected: dict[int, int] = {} + flags: dict[int, int] = {} + for chain_id, candidates in rows_by_chain.items(): + selected[chain_id] = candidates[occurrence % len(candidates)] + flags[chain_id] = 1 + for chain_id in chain_ids: + if chain_id not in selected: + flags[chain_id] = 0 + selected[chain_id] = available[chain_id].pop(0) if available[chain_id] else -1 + rows.append(selected) + paired_flags.append(flags) + if len(rows) >= max_pairs: + break + if len(rows) >= max_pairs: + break + + +def _append_unpaired_rows( + rows: list[dict[int, int]], + paired_flags: list[dict[int, int]], + chain_ids: list[int], + available: dict[int, list[int]], + max_total: int, +) -> None: + max_remaining = max((len(indices) for indices in available.values()), default=0) + for _ in range(min(max_total - len(rows), max_remaining)): + rows.append( + { + chain_id: available[chain_id].pop(0) if available[chain_id] else -1 + for chain_id in chain_ids + } + ) + paired_flags.append({chain_id: 0 for chain_id in chain_ids}) + if len(rows) >= max_total: + break + + +def _pairing_plan( + chain_ids: list[int], + alignments: dict[int, _ChainAlignment], + max_pairs: int, + max_total: int, + max_seqs: int, +) -> tuple[list[dict[int, int]], list[dict[int, int]]]: + groups = _taxonomy_groups(chain_ids, alignments) + available = _available_rows(chain_ids, alignments, groups) + rows = [{chain_id: 0 for chain_id in chain_ids}] + flags = [{chain_id: 1 for chain_id in chain_ids}] + _append_taxonomy_rows(rows, flags, chain_ids, groups, available, max_pairs) + _append_unpaired_rows(rows, flags, chain_ids, available, max_total) + return rows[:max_seqs], flags[:max_seqs] + + +def _project_alignment_rows( + chain_ids: list[int], + alignments: dict[int, _ChainAlignment], + rows: list[dict[int, int]], + flags: list[dict[int, int]], + token_asym_ids: np.ndarray, + token_res_ids: np.ndarray, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + m, t = len(rows), len(token_asym_ids) + residues = np.full((m, t), MSA_GAP_TOKEN_ID, dtype=np.int64) + deletions = np.zeros((m, t), dtype=np.float32) + paired_mask = np.zeros((m, t), dtype=np.float32) + for chain_id in chain_ids: + alignment = alignments[chain_id] + selected_rows = np.asarray([row[chain_id] for row in rows], dtype=np.int64) + chain_flags = np.asarray([row[chain_id] for row in flags], dtype=np.float32) + token_mask = token_asym_ids == chain_id + if not token_mask.any(): + continue + columns = np.minimum(token_res_ids[token_mask], alignment.residues.shape[1] - 1) + valid_rows = selected_rows >= 0 + if valid_rows.any(): + output_rows = np.flatnonzero(valid_rows) + output_columns = np.flatnonzero(token_mask) + residues[np.ix_(output_rows, output_columns)] = alignment.residues[ + selected_rows[valid_rows] + ][:, columns] + deletions[np.ix_(output_rows, output_columns)] = alignment.deletions[ + selected_rows[valid_rows] + ][:, columns] + paired_mask[:, token_mask] = chain_flags[:, None] + return residues, deletions, paired_mask + + +def construct_paired_msa( + chain_msas: dict[int, MSA | None], + chain_query_res_types: dict[int, np.ndarray], + token_asym_ids: np.ndarray, + token_res_ids: np.ndarray, + letter_to_res_type: dict[str, int] | None = None, + *, + max_pairs: int = 8192, + max_total: int = 16384, + max_seqs: int = 16384, +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Return residue, deletion, and pairing arrays with shape (m, t).""" + + vocabulary = protein_letter_to_res_type() if letter_to_res_type is None else letter_to_res_type + chain_ids = sorted(chain_msas) + alignments = { + chain_id: _chain_alignment( + chain_msas[chain_id], + chain_query_res_types[chain_id], + vocabulary, + ) + for chain_id in chain_ids + } + rows, flags = _pairing_plan( + chain_ids, + alignments, + max_pairs, + max_total, + max_seqs, + ) + return _project_alignment_rows( + chain_ids, + alignments, + rows, + flags, + token_asym_ids, + token_res_ids, + ) diff --git a/fastplms/models/esmfold2/esmfold2_parsing.py b/fastplms/models/esmfold2/esmfold2_parsing.py new file mode 100644 index 0000000000000000000000000000000000000000..f906d1ff2804c726daf283d6b323f44b43d2a662 --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_parsing.py @@ -0,0 +1,126 @@ +"""FASTA parsing and writing with explicit stream ownership.""" + +from __future__ import annotations + +import gzip +import io +from collections.abc import Generator, Iterable +from contextlib import nullcontext +from pathlib import Path +from typing import NamedTuple, TextIO + +from .esmfold2_utils_types import PathOrBuffer + + +class FastaEntry(NamedTuple): + """One FASTA record in source order.""" + + header: str + sequence: str + + +def parse_fasta(text: str) -> Generator[FastaEntry, None, None]: + """Yield records from FASTA text without normalizing sequence symbols.""" + + header: str | None = None + sequence_lines: list[str] = [] + found_record = False + + for line in text.splitlines(): + if not line or line.startswith("#"): + continue + if line.startswith(">"): + if header is not None: + found_record = True + yield FastaEntry(header, "".join(sequence_lines)) + header = line[1:].strip() + sequence_lines.clear() + elif header is not None: + sequence_lines.append(line) + + if header is not None: + found_record = True + yield FastaEntry(header, "".join(sequence_lines)) + if not found_record: + raise ValueError("Found no sequences in input") + + +def _open_reader(source: PathOrBuffer): + if isinstance(source, io.TextIOBase): + return nullcontext(source) + path = Path(source) + if path.suffix.lower() == ".gz": + return gzip.open(path, mode="rt", encoding="utf-8") + return path.open(mode="r", encoding="utf-8") + + +def read_sequences(source: PathOrBuffer) -> Generator[FastaEntry, None, None]: + """Read FASTA records while leaving caller-owned streams open.""" + + with _open_reader(source) as handle: + yield from parse_fasta(handle.read()) + + +def read_first_sequence(source: PathOrBuffer) -> FastaEntry: + """Return the first FASTA record from a path or text stream.""" + + return next(read_sequences(source)) + + +def count_fasta_sequences(path: str | Path) -> int: + """Count FASTA headers without parsing sequence bodies.""" + + source = Path(path) + if not source.exists(): + return 0 + with source.open(encoding="utf-8") as handle: + return sum(line.startswith(">") for line in handle) + + +def append_fasta_sequence(header: str, sequence: str, path: str | Path) -> None: + """Append one record, inserting a separator if the file lacks a final newline.""" + + destination = Path(path) + destination.parent.mkdir(parents=True, exist_ok=True) + needs_separator = ( + destination.exists() + and destination.stat().st_size > 0 + and destination.read_bytes()[-1:] != b"\n" + ) + with destination.open(mode="a", encoding="utf-8") as handle: + if needs_separator: + handle.write("\n") + handle.write(f">{header}\n{sequence}\n") + + +def _open_writer(destination: PathOrBuffer): + if isinstance(destination, io.TextIOBase): + return nullcontext(destination) + path = Path(destination) + path.parent.mkdir(parents=True, exist_ok=True) + return path.open(mode="w", encoding="utf-8") + + +def write_sequences(sequences: Iterable[tuple[str, str]], destination: PathOrBuffer) -> None: + """Write records with one blank-line-free separator between entries.""" + + with _open_writer(destination) as handle: + _write_records(handle, sequences) + + +def _write_records(handle: TextIO, sequences: Iterable[tuple[str, str]]) -> None: + for index, (header, sequence) in enumerate(sequences): + if index: + handle.write("\n") + handle.write(f">{header}\n{sequence}") + + +__all__ = [ + "FastaEntry", + "append_fasta_sequence", + "count_fasta_sequences", + "parse_fasta", + "read_first_sequence", + "read_sequences", + "write_sequences", +] diff --git a/fastplms/models/esmfold2/esmfold2_predicted_aligned_error.py b/fastplms/models/esmfold2/esmfold2_predicted_aligned_error.py new file mode 100644 index 0000000000000000000000000000000000000000..d6ab3e74beb4ae86d22fb7d879126f5b23a1096e --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_predicted_aligned_error.py @@ -0,0 +1,127 @@ +"""Predicted-aligned-error scores and training loss.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F +from torch import Tensor + +from .esmfold2_affine3d import Affine3D + +_CPU_DEVICE = torch.device("cpu") + + +def _compute_pae_masks(mask: Tensor) -> Tensor: + residue_mask = mask.bool() + return residue_mask.unsqueeze(-1) & residue_mask.unsqueeze(-2) + + +def _pae_bins( + max_bin: float = 31, + num_bins: int = 64, + device: torch.device = _CPU_DEVICE, +) -> Tensor: + """Return the representative distance for each PAE probability bin.""" + + boundaries = torch.linspace(0, max_bin, steps=num_bins - 1, device=device) + width = max_bin / (num_bins - 2) + centers = boundaries + width / 2 + overflow_center = centers[-1:] + width + return torch.cat((centers, overflow_center)) + + +def _masked_probabilities(logits: Tensor, pair_mask: Tensor) -> Tensor: + masked_logits = logits.masked_fill( + ~pair_mask.unsqueeze(-1), + torch.finfo(logits.dtype).min, + ) + return masked_logits.softmax(dim=-1) + + +def masked_mean( + mask: Tensor, + value: Tensor, + dim: int | tuple[int, ...] | None = None, + eps: float = 1e-10, +) -> Tensor: + """Average values over true entries of a broadcast-compatible mask.""" + + weights = mask.expand_as(value) + weighted_sum = torch.sum(weights * value, dim=dim) + weight_sum = torch.sum(weights, dim=dim) + return weighted_sum / (weight_sum + eps) + + +def compute_predicted_aligned_error( + logits: Tensor, + aa_mask: Tensor, + sequence_id: Tensor | None = None, + max_bin: float = 31, +) -> Tensor: + """Convert PAE logits ``X`` with shape (..., l, l, n) to distances.""" + + del sequence_id + pair_mask = _compute_pae_masks(aa_mask) + probabilities = _masked_probabilities(logits, pair_mask) + centers = _pae_bins(max_bin, logits.shape[-1], logits.device) + return torch.sum(probabilities * centers, dim=-1) + + +@torch.no_grad() +def compute_tm(logits: Tensor, aa_mask: Tensor, max_bin: float = 31.0) -> Tensor: + """Estimate TM score from pairwise PAE logits.""" + + pair_mask = _compute_pae_masks(aa_mask) + sequence_lengths = aa_mask.sum(dim=-1, keepdim=True) + centers = _pae_bins(max_bin, logits.shape[-1], logits.device) + distance_scale = 1.24 * (sequence_lengths.clamp_min(19) - 15) ** (1 / 3) - 1.8 + tm_weights = 1.0 / (1 + (centers / distance_scale.unsqueeze(-1)) ** 2) + probabilities = _masked_probabilities(logits, pair_mask) + score_per_pair = torch.sum(probabilities * tm_weights.unsqueeze(-2), dim=-1) + score_per_anchor = masked_mean(pair_mask, score_per_pair, dim=-1) + return score_per_anchor.max(dim=-1).values + + +def _local_coordinates(frames: Affine3D) -> Tensor: + origins = frames.trans[..., None, :, :] + return frames.invert()[..., None].apply(origins) + + +def tm_loss( + logits: Tensor, + pred_affine: Tensor, + targ_affine: Tensor, + targ_mask: Tensor, + tm_mask: Tensor | None = None, + sequence_id: Tensor | None = None, + max_bin: float = 31, +) -> Tensor: + """Cross-entropy loss for discretized aligned-position errors.""" + + del sequence_id + predicted_frames = Affine3D.from_tensor(pred_affine) + target_frames = Affine3D.from_tensor(targ_affine) + with torch.no_grad(): + squared_error = ( + (_local_coordinates(predicted_frames) - _local_coordinates(target_frames)) + .square() + .sum(dim=-1) + ) + boundaries = torch.linspace( + 0, + max_bin, + logits.shape[-1] - 1, + device=logits.device, + ).square() + target_bins = (squared_error[..., None] > boundaries).sum(dim=-1).long() + + cross_entropy = F.cross_entropy( + logits.movedim(3, 1), + target_bins, + reduction="none", + ) + pair_mask = _compute_pae_masks(targ_mask) + loss_per_sample = masked_mean(pair_mask, cross_entropy, dim=(-1, -2)) + if tm_mask is None: + return loss_per_sample.mean() + return masked_mean(tm_mask, loss_per_sample) diff --git a/fastplms/models/esmfold2/esmfold2_prepare_input.py b/fastplms/models/esmfold2/esmfold2_prepare_input.py new file mode 100644 index 0000000000000000000000000000000000000000..870c45754d5e3d41d1f0db90ffb56eec2a5eb3fd --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_prepare_input.py @@ -0,0 +1,1130 @@ +"""Translate typed sequence inputs into the tensors consumed by ESMFold2. + +The conversion has four explicit stages: entity and chain assignment, residue +tokenization, structural feature construction, and atom-table padding. Keeping +those stages separate makes the biological indexing rules testable without +loading model weights. +""" + +from __future__ import annotations + +import math +import warnings +from collections import defaultdict +from contextlib import suppress +from dataclasses import dataclass, field +from itertools import combinations +from typing import Any + +import numpy as np +import torch + +from .esmfold2_conformers import ( + get_ccd_leaving_atoms, + get_idealized_atom_pos, + get_ligand_ccd_atoms_with_charges, + get_ligand_ccd_bonds, + get_ligand_idealized_atom_pos, +) +from .esmfold2_constants import ( + CHARGED_ATOMS, + DNA_1TO3, + DNA_BACKBONE_ATOMS, + DNA_HEAVY_ATOMS, + DNA_RESIDUE_TO_RES_TYPE, + DNA_RNA_LIGAND_INPUT_ID, + DNA_UNK_RES_TYPE, + ELEMENT_TO_ATOMIC_NUM, + ESM_PROTEIN_VOCAB, + MOL_TYPE_DNA, + MOL_TYPE_NONPOLYMER, + MOL_TYPE_PROTEIN, + MOL_TYPE_RNA, + MSA_GAP_TOKEN_ID, + PROTEIN_1TO3, + PROTEIN_3TO1, + PROTEIN_HEAVY_ATOMS, + PROTEIN_RESIDUE_TO_RES_TYPE, + PROTEIN_UNK_RES_TYPE, + RNA_1TO3, + RNA_BACKBONE_ATOMS, + RNA_HEAVY_ATOMS, + RNA_RESIDUE_TO_RES_TYPE, + RNA_UNK_RES_TYPE, +) +from .esmfold2_types import ( + MSA, + DNAInput, + LigandInput, + Modification, + ProteinInput, + RNAInput, + StructurePredictionInput, +) + +_ZERO_POS = np.zeros(3, dtype=np.float32) +_ENCODE_ATOM_NAME_CACHE: dict[str, list[int]] = {} +_ELEMENT_ATOMIC_NUM_CACHE: dict[str, int] = {} +_TWO_LETTER_ELEMENTS = frozenset({"FE", "ZN", "MG", "MN", "CO", "NI", "CU", "SE", "BR"}) + + +@dataclass +class AtomInfo: + """One row in the unpadded atom table.""" + + name: str + element: str + charge: int + ref_pos: np.ndarray # R has shape (3,). + pos: np.ndarray # X has shape (3,). + token_index: int = -1 + atom_index: int = -1 + space_uid: int = -1 + is_valid: bool = True + + +@dataclass +class TokenInfo: + """Biological and atom-span annotations for one model token.""" + + token_index: int + residue_index: int + residue_name: str + mol_type: int + res_type: int + input_id: int + asym_id: int + sym_id: int + entity_id: int + atom_start: int + atom_count: int + + +@dataclass +class ChainInfo: + """One input chain after entity and symmetry assignment.""" + + chain_id: str + asym_id: int + entity_id: int + sym_id: int + mol_type: int + tokens: list[TokenInfo] = field(default_factory=list) + ligand_bonds: list[tuple[str, str]] = field(default_factory=list) + + +@dataclass +class _TokenizationState: + """Mutable cursor shared by residue tokenizers.""" + + token_index: int + atom_index: int + space_uid: int + tokens: list[TokenInfo] = field(default_factory=list) + atoms: list[AtomInfo] = field(default_factory=list) + + def _append_atom( + self, + name: str, + element: str, + charge: int, + ref_pos: np.ndarray | None, + ) -> None: + self.atoms.append( + AtomInfo( + name=name, + element=element, + charge=charge, + ref_pos=(ref_pos.copy() if ref_pos is not None else _ZERO_POS.copy()), + pos=_ZERO_POS.copy(), + token_index=self.token_index, + atom_index=self.atom_index, + space_uid=self.space_uid, + ) + ) + self.atom_index += 1 + + def _append_token( + self, + *, + residue_index: int, + residue_name: str, + mol_type: int, + res_type: int, + input_id: int, + asym_id: int, + sym_id: int, + entity_id: int, + atom_start: int, + atom_count: int, + ) -> None: + self.tokens.append( + TokenInfo( + token_index=self.token_index, + residue_index=residue_index, + residue_name=residue_name, + mol_type=mol_type, + res_type=res_type, + input_id=input_id, + asym_id=asym_id, + sym_id=sym_id, + entity_id=entity_id, + atom_start=atom_start, + atom_count=atom_count, + ) + ) + self.token_index += 1 + + def add_residue_token( + self, + atom_specs: list[tuple[str, str, int, np.ndarray | None]], + **token_fields: Any, + ) -> None: + start = self.atom_index + for atom_spec in atom_specs: + self._append_atom(*atom_spec) + self._append_token( + atom_start=start, + atom_count=len(atom_specs), + **token_fields, + ) + self.space_uid += 1 + + def add_atom_tokens( + self, + atom_specs: list[tuple[str, str, int, np.ndarray | None]], + **token_fields: Any, + ) -> None: + for atom_spec in atom_specs: + start = self.atom_index + self._append_atom(*atom_spec) + self._append_token(atom_start=start, atom_count=1, **token_fields) + self.space_uid += 1 + + +def encode_atom_name(name: str) -> list[int]: + """Encode a four-character atom name with the model's ASCII offset.""" + cached = _ENCODE_ATOM_NAME_CACHE.get(name) + if cached is None: + cached = [0 if char == " " else ord(char) - 32 for char in name.ljust(4)[:4]] + _ENCODE_ATOM_NAME_CACHE[name] = cached + return cached + + +def get_element_atomic_num(element: str) -> int: + """Map an element symbol to the model's atomic-number vocabulary.""" + cached = _ELEMENT_ATOMIC_NUM_CACHE.get(element) + if cached is None: + cached = ELEMENT_TO_ATOMIC_NUM.get(element.upper(), 0) + _ELEMENT_ATOMIC_NUM_CACHE[element] = cached + return cached + + +def _infer_element(atom_name: str) -> str: + normalized = atom_name.strip() + if not normalized: + return "C" + if normalized[0].isdigit(): + return normalized[1] if len(normalized) > 1 else "H" + if len(normalized) == 2 and normalized in _TWO_LETTER_ELEMENTS: + return normalized + return normalized[0] + + +def _compute_res_type(name: str, mol_type: int) -> int: + if mol_type == MOL_TYPE_PROTEIN: + return PROTEIN_RESIDUE_TO_RES_TYPE.get(name, PROTEIN_UNK_RES_TYPE) + if mol_type == MOL_TYPE_DNA: + return DNA_RESIDUE_TO_RES_TYPE.get( + name, RNA_RESIDUE_TO_RES_TYPE.get(name, DNA_UNK_RES_TYPE) + ) + if mol_type == MOL_TYPE_RNA: + return RNA_RESIDUE_TO_RES_TYPE.get( + name, DNA_RESIDUE_TO_RES_TYPE.get(name, RNA_UNK_RES_TYPE) + ) + return PROTEIN_UNK_RES_TYPE + + +def _compute_esm_input_id(name: str, mol_type: int) -> int: + if mol_type != MOL_TYPE_PROTEIN: + return DNA_RNA_LIGAND_INPUT_ID + letter = PROTEIN_3TO1.get(name) + return ( + DNA_RNA_LIGAND_INPUT_ID + if letter is None + else ESM_PROTEIN_VOCAB.get(letter, ESM_PROTEIN_VOCAB["X"]) + ) + + +def _apply_modifications(residues: list[str], modifications: list[Modification] | None) -> set[int]: + changed: set[int] = set() + for modification in modifications or (): + residues[modification.position] = modification.ccd + changed.add(modification.position) + return changed + + +def _ideal_atom_specs( + residue_name: str, + residue_type: int, + atom_names: list[str], + *, + charges: bool = True, +) -> list[tuple[str, str, int, np.ndarray | None]]: + return [ + ( + atom_name, + _infer_element(atom_name), + CHARGED_ATOMS.get((residue_name, atom_name), 0) if charges else 0, + get_idealized_atom_pos(residue_type, atom_name), + ) + for atom_name in atom_names + ] + + +def _ccd_atom_specs( + residue_name: str, + atoms: list[tuple[str, str, int]], + excluded: set[str], + *, + force_zero: bool = False, +) -> list[tuple[str, str, int, np.ndarray | None]]: + return [ + ( + atom_name, + element, + charge, + None if force_zero else get_ligand_idealized_atom_pos(residue_name, atom_name), + ) + for atom_name, element, charge in atoms + if atom_name not in excluded + ] + + +def tokenize_protein( + sequence: str, + modifications: list[Modification] | None, + entity_id: int, + asym_id: int, + sym_id: int, + token_offset: int, + atom_offset: int, + space_uid_offset: int, +) -> tuple[list[TokenInfo], list[AtomInfo]]: + """Tokenize protein residues, atom-tokenizing modified CCD components.""" + residues = [PROTEIN_1TO3.get(letter, "UNK") for letter in sequence] + modified = _apply_modifications(residues, modifications) + state = _TokenizationState(token_offset, atom_offset, space_uid_offset) + + for residue_index, residue_name in enumerate(residues): + canonical_name = "MET" if residue_name == "MSE" else residue_name + common_fields = { + "residue_index": residue_index, + "mol_type": MOL_TYPE_PROTEIN, + "asym_id": asym_id, + "sym_id": sym_id, + "entity_id": entity_id, + } + if residue_index not in modified and canonical_name in PROTEIN_HEAVY_ATOMS: + residue_type = _compute_res_type(canonical_name, MOL_TYPE_PROTEIN) + state.add_residue_token( + _ideal_atom_specs( + canonical_name, + residue_type, + PROTEIN_HEAVY_ATOMS[canonical_name], + ), + residue_name=canonical_name, + res_type=residue_type, + input_id=_compute_esm_input_id(canonical_name, MOL_TYPE_PROTEIN), + **common_fields, + ) + continue + + ccd_atoms = get_ligand_ccd_atoms_with_charges(residue_name) + if ccd_atoms is None: + ccd_atoms = [ + (_infer_element(name), _infer_element(name), 0) for name in ("N", "CA", "C", "O") + ] + excluded = ( + set() if residue_index == len(residues) - 1 else get_ccd_leaving_atoms(residue_name) + ) + retained = [atom for atom in ccd_atoms if atom[0] not in excluded] + state.add_atom_tokens( + _ccd_atom_specs( + residue_name, + retained, + set(), + force_zero=len(retained) == 1, + ), + residue_name=residue_name, + res_type=PROTEIN_UNK_RES_TYPE, + input_id=DNA_RNA_LIGAND_INPUT_ID, + **common_fields, + ) + return state.tokens, state.atoms + + +def tokenize_nucleotide( + sequence: str, + modifications: list[Modification] | None, + mol_type: int, + entity_id: int, + asym_id: int, + sym_id: int, + token_offset: int, + atom_offset: int, + space_uid_offset: int, +) -> tuple[list[TokenInfo], list[AtomInfo]]: + """Tokenize DNA or RNA, retaining backbone atoms for unknown bases.""" + dna = mol_type == MOL_TYPE_DNA + letter_map = DNA_1TO3 if dna else RNA_1TO3 + heavy_atoms = DNA_HEAVY_ATOMS if dna else RNA_HEAVY_ATOMS + backbone_atoms = DNA_BACKBONE_ATOMS if dna else RNA_BACKBONE_ATOMS + unknown_type = DNA_UNK_RES_TYPE if dna else RNA_UNK_RES_TYPE + residues = [letter_map.get(letter, "UNK") for letter in sequence] + modified = _apply_modifications(residues, modifications) + state = _TokenizationState(token_offset, atom_offset, space_uid_offset) + + for residue_index, residue_name in enumerate(residues): + common_fields = { + "residue_index": residue_index, + "residue_name": residue_name, + "mol_type": mol_type, + "asym_id": asym_id, + "sym_id": sym_id, + "entity_id": entity_id, + "input_id": DNA_RNA_LIGAND_INPUT_ID, + } + if residue_index not in modified and residue_name in heavy_atoms: + residue_type = _compute_res_type(residue_name, mol_type) + state.add_residue_token( + _ideal_atom_specs(residue_name, residue_type, heavy_atoms[residue_name]), + res_type=residue_type, + **common_fields, + ) + continue + if residue_index not in modified and residue_name == "UNK": + state.add_residue_token( + [(atom_name, _infer_element(atom_name), 0, None) for atom_name in backbone_atoms], + res_type=unknown_type, + **common_fields, + ) + continue + + ccd_atoms = get_ligand_ccd_atoms_with_charges(residue_name) + if ccd_atoms is None: + ccd_atoms = [(_infer_element(name), _infer_element(name), 0) for name in backbone_atoms] + excluded = ( + set() if residue_index == len(residues) - 1 else get_ccd_leaving_atoms(residue_name) + ) + state.add_atom_tokens( + _ccd_atom_specs(residue_name, ccd_atoms, excluded), + res_type=PROTEIN_UNK_RES_TYPE, + **common_fields, + ) + return state.tokens, state.atoms + + +def tokenize_ligand_ccd( + ccd_codes: list[str], + entity_id: int, + asym_id: int, + sym_id: int, + token_offset: int, + atom_offset: int, + space_uid_offset: int, + has_covalent_bond: bool, +) -> tuple[list[TokenInfo], list[AtomInfo]]: + """Tokenize CCD ligands with one model token per retained atom.""" + state = _TokenizationState(token_offset, atom_offset, space_uid_offset) + for residue_index, code in enumerate(ccd_codes): + ccd_atoms = get_ligand_ccd_atoms_with_charges(code) + if ccd_atoms is None: + raise ValueError(f"CCD component {code} not found") + excluded = get_ccd_leaving_atoms(code) if has_covalent_bond else set() + state.add_atom_tokens( + _ccd_atom_specs(code, ccd_atoms, excluded), + residue_index=residue_index, + residue_name=code, + mol_type=MOL_TYPE_NONPOLYMER, + res_type=PROTEIN_UNK_RES_TYPE, + input_id=DNA_RNA_LIGAND_INPUT_ID, + asym_id=asym_id, + sym_id=sym_id, + entity_id=entity_id, + ) + return state.tokens, state.atoms + + +def tokenize_ligand_smiles( + smiles: str, + entity_id: int, + asym_id: int, + sym_id: int, + token_offset: int, + atom_offset: int, + space_uid_offset: int, + seed: int | None = None, +) -> tuple[list[TokenInfo], list[AtomInfo], list[tuple[str, str]]]: + """Generate a conformer and tokenize each heavy atom of a SMILES ligand.""" + from rdkit import Chem + from rdkit.Chem import AllChem + + molecule = Chem.MolFromSmiles(smiles) + if molecule is None: + raise ValueError(f"Failed to parse SMILES: {smiles}") + molecule = Chem.AddHs(molecule) + canonical_order = AllChem.CanonicalRankAtoms(molecule) # type: ignore[attr-defined] + for atom, canonical_index in zip(molecule.GetAtoms(), canonical_order, strict=True): + name = atom.GetSymbol().upper() + str(canonical_index + 1) + if len(name) > 4: + raise ValueError(f"SMILES {smiles} has atom name longer than 4 chars: {name}") + atom.SetProp("name", name) + + options = AllChem.ETKDGv3() # type: ignore[attr-defined] + options.clearConfs = False + if seed is not None: + options.randomSeed = seed + conformer_id = AllChem.EmbedMolecule(molecule, options) # type: ignore[attr-defined] + if conformer_id == -1: + options.useRandomCoords = True + conformer_id = AllChem.EmbedMolecule(molecule, options) # type: ignore[attr-defined] + if conformer_id != -1: + with suppress(RuntimeError, ValueError): + AllChem.UFFOptimizeMolecule( # type: ignore[attr-defined] + molecule, confId=conformer_id, maxIters=1000 + ) + + molecule = Chem.RemoveHs(molecule) + if molecule.GetNumConformers() == 0: + raise ValueError(f"Failed to generate conformer for SMILES: {smiles}") + conformer = molecule.GetConformer(0) + atom_specs: list[tuple[str, str, int, np.ndarray | None]] = [] + for atom in molecule.GetAtoms(): + position = conformer.GetAtomPosition(atom.GetIdx()) + atom_specs.append( + ( + atom.GetProp("name"), + atom.GetSymbol(), + atom.GetFormalCharge(), + np.asarray([position.x, position.y, position.z], dtype=np.float32), + ) + ) + state = _TokenizationState(token_offset, atom_offset, space_uid_offset) + state.add_atom_tokens( + atom_specs, + residue_index=0, + residue_name="LIG", + mol_type=MOL_TYPE_NONPOLYMER, + res_type=PROTEIN_UNK_RES_TYPE, + input_id=DNA_RNA_LIGAND_INPUT_ID, + asym_id=asym_id, + sym_id=sym_id, + entity_id=entity_id, + ) + bonds = [ + ( + bond.GetBeginAtom().GetProp("name"), + bond.GetEndAtom().GetProp("name"), + ) + for bond in molecule.GetBonds() + ] + return state.tokens, state.atoms, bonds + + +def _get_sequence_key(item: Any) -> str: + if isinstance(item, ProteinInput): + return f"PROTEIN:{item.sequence}" + if isinstance(item, DNAInput): + return f"DNA:{item.sequence}" + if isinstance(item, RNAInput): + return f"RNA:{item.sequence}" + if isinstance(item, LigandInput): + return f"LIGAND_CCD:{','.join(item.ccd)}" if item.ccd else f"LIGAND_SMILES:{item.smiles}" + raise ValueError(f"Unknown input type: {type(item)}") + + +def _tokenize_chain( + item: Any, + chain_id: str, + *, + entity_id: int, + asym_id: int, + sym_id: int, + token_offset: int, + atom_offset: int, + space_uid_offset: int, + covalent_chains: set[str], + seed: int | None, +) -> tuple[list[TokenInfo], list[AtomInfo], list[tuple[str, str]]]: + common = { + "entity_id": entity_id, + "asym_id": asym_id, + "sym_id": sym_id, + "token_offset": token_offset, + "atom_offset": atom_offset, + "space_uid_offset": space_uid_offset, + } + if isinstance(item, ProteinInput): + if item.msa is None: + warnings.warn( + f"No MSA provided for {item.id}, using single sequence mode", + stacklevel=2, + ) + tokens, atoms = tokenize_protein(item.sequence, item.modifications, **common) + return tokens, atoms, [] + if isinstance(item, (DNAInput, RNAInput)): + mol_type = MOL_TYPE_DNA if isinstance(item, DNAInput) else MOL_TYPE_RNA + tokens, atoms = tokenize_nucleotide( + item.sequence, item.modifications, mol_type=mol_type, **common + ) + return tokens, atoms, [] + if not isinstance(item, LigandInput): + raise ValueError(f"Unknown input type: {type(item)}") + if item.ccd is not None: + if item.smiles is not None: + warnings.warn("Both ccd and smiles provided, using ccd", stacklevel=2) + tokens, atoms = tokenize_ligand_ccd( + item.ccd, + has_covalent_bond=chain_id in covalent_chains, + **common, + ) + return tokens, atoms, [] + if item.smiles is not None: + return tokenize_ligand_smiles(item.smiles, seed=seed, **common) + raise ValueError("LigandInput must have either ccd or smiles") + + +def build_chains_from_input( + input: StructurePredictionInput, seed: int | None = None +) -> tuple[list[ChainInfo], list[TokenInfo], list[AtomInfo]]: + """Assign entities and symmetry copies, then tokenize every input chain.""" + chains: list[ChainInfo] = [] + tokens: list[TokenInfo] = [] + atoms: list[AtomInfo] = [] + entity_for_sequence: dict[str, int] = {} + next_symmetry: dict[int, int] = {} + covalent_chains = { + chain_id + for bond in input.covalent_bonds or () + for chain_id in (bond.chain_id1, bond.chain_id2) + } + space_uid_offset = 0 + + for item in input.sequences: + key = _get_sequence_key(item) + entity_id = entity_for_sequence.setdefault(key, len(entity_for_sequence)) + chain_ids = [item.id] if isinstance(item.id, str) else item.id + for chain_id in chain_ids: + sym_id = next_symmetry.get(entity_id, 0) + next_symmetry[entity_id] = sym_id + 1 + asym_id = len(chains) + new_tokens, new_atoms, ligand_bonds = _tokenize_chain( + item, + chain_id, + entity_id=entity_id, + asym_id=asym_id, + sym_id=sym_id, + token_offset=len(tokens), + atom_offset=len(atoms), + space_uid_offset=space_uid_offset, + covalent_chains=covalent_chains, + seed=seed, + ) + chains.append( + ChainInfo( + chain_id=chain_id, + asym_id=asym_id, + entity_id=entity_id, + sym_id=sym_id, + mol_type=(new_tokens[0].mol_type if new_tokens else MOL_TYPE_PROTEIN), + tokens=new_tokens, + ligand_bonds=ligand_bonds, + ) + ) + tokens.extend(new_tokens) + atoms.extend(new_atoms) + space_uid_offset += len({atom.space_uid for atom in new_atoms}) + return chains, tokens, atoms + + +def _atom_indices_by_name(atoms: list[AtomInfo]) -> dict[int, dict[str, int]]: + result: dict[int, dict[str, int]] = defaultdict(dict) + for atom in atoms: + if atom.is_valid: + result[atom.token_index][atom.name] = atom.atom_index + return result + + +def _ligand_frames( + tokens: list[TokenInfo], + atoms: list[AtomInfo], + atom_indices: dict[int, dict[str, int]], +) -> dict[int, tuple[int, int, int]]: + atom_for_token: dict[int, int] = {} + tokens_by_residue: dict[tuple[int, int], list[int]] = defaultdict(list) + for token in tokens: + if token.mol_type != MOL_TYPE_NONPOLYMER: + continue + named_atoms = atom_indices.get(token.token_index) + if named_atoms: + atom_for_token[token.token_index] = next(iter(named_atoms.values())) + tokens_by_residue[(token.asym_id, token.residue_index)].append(token.token_index) + + frames: dict[int, tuple[int, int, int]] = {} + for residue_tokens in tokens_by_residue.values(): + residue_atoms = [ + atom_for_token[token] for token in residue_tokens if token in atom_for_token + ] + if len(residue_atoms) < 3: + for token in residue_tokens: + if token in atom_for_token: + atom_index = atom_for_token[token] + frames[token] = (atom_index, atom_index, atom_index) + continue + R = np.asarray([atoms[index].ref_pos for index in residue_atoms]) + distances = np.sqrt(((R[:, None] - R[None]) ** 2).sum(-1)) + nearest = np.argsort(distances, axis=1) + local = np.column_stack((nearest[:, 1], nearest[:, 0], nearest[:, 2])) + local_index = {atom_index: index for index, atom_index in enumerate(residue_atoms)} + for token in residue_tokens: + atom_index = atom_for_token.get(token) + if atom_index is None: + continue + selected = local[local_index[atom_index]] + frames[token] = tuple(residue_atoms[int(index)] for index in selected) + return frames + + +def _frame_for_token( + token: TokenInfo, + named_atoms: dict[str, int], + ligand_frames: dict[int, tuple[int, int, int]], +) -> tuple[int, int, int]: + fallback = next(iter(named_atoms.values()), 0) + if token.mol_type == MOL_TYPE_PROTEIN: + return ( + (fallback, fallback, fallback) + if token.res_type == PROTEIN_UNK_RES_TYPE + else ( + named_atoms.get("N", 0), + named_atoms.get("CA", 0), + named_atoms.get("C", 0), + ) + ) + if token.mol_type in (MOL_TYPE_DNA, MOL_TYPE_RNA): + return ( + (fallback, fallback, fallback) + if token.res_type == PROTEIN_UNK_RES_TYPE + else ( + named_atoms.get("C1'", 0), + named_atoms.get("C3'", 0), + named_atoms.get("C4'", 0), + ) + ) + if token.mol_type == MOL_TYPE_NONPOLYMER: + return ligand_frames.get(token.token_index, (fallback, fallback, fallback)) + return fallback, fallback, fallback + + +def _resolved_frames( + frames: np.ndarray, tokens: list[TokenInfo], atoms: list[AtomInfo] +) -> np.ndarray: + if not tokens: + return np.zeros(0, dtype=bool) + X = ( + np.asarray([atom.pos for atom in atoms], dtype=np.float32) + if atoms + else np.zeros((0, 3), dtype=np.float32) + ) + valid_atoms = ( + np.asarray([atom.is_valid for atom in atoms], dtype=bool) + if atoms + else np.zeros(0, dtype=bool) + ) + resolved_atoms = valid_atoms & np.any(X != 0, axis=1) + origin = X[frames[:, 1]] + left = X[frames[:, 0]] - origin + right = X[frames[:, 2]] - origin + left_norm = np.linalg.norm(left, axis=1) + right_norm = np.linalg.norm(right, axis=1) + valid_norms = (left_norm >= 1e-6) & (right_norm >= 1e-6) + cosine = np.zeros(len(tokens), dtype=np.float32) + if np.any(valid_norms): + cosine[valid_norms] = np.sum(left[valid_norms] * right[valid_norms], axis=1) / ( + left_norm[valid_norms] * right_norm[valid_norms] + ) + angle = np.degrees(np.arccos(np.abs(np.clip(cosine, -1, 1)))) + all_resolved = resolved_atoms[frames].all(axis=1) + repeated = (frames[:, 0] == frames[:, 1]) & (frames[:, 1] == frames[:, 2]) + return all_resolved & ~repeated & valid_norms & (angle >= 25) + + +def compute_frame_indices( + tokens: list[TokenInfo], atoms: list[AtomInfo] +) -> tuple[np.ndarray, np.ndarray]: + """Return frame atom indices F with shape (l, 3) and validity M with shape (l,).""" + named_atoms = _atom_indices_by_name(atoms) + ligand_frames = _ligand_frames(tokens, atoms, named_atoms) + frames = np.asarray( + [ + _frame_for_token(token, named_atoms.get(token.token_index, {}), ligand_frames) + for token in tokens + ], + dtype=np.int64, + ) + return frames, _resolved_frames(frames, tokens, atoms) + + +def _atom_tokenized_residues( + tokens: list[TokenInfo], atoms: list[AtomInfo] +) -> dict[tuple[int, int], list[tuple[str, int]]]: + grouped: dict[tuple[int, int], list[tuple[str, int]]] = defaultdict(list) + for atom in atoms: + if not atom.is_valid or atom.token_index >= len(tokens): + continue + token = tokens[atom.token_index] + if token.mol_type == MOL_TYPE_NONPOLYMER or token.res_type == PROTEIN_UNK_RES_TYPE: + grouped[(token.asym_id, token.residue_index)].append((atom.name, atom.token_index)) + return grouped + + +def _backbone_token( + residue_tokens: list[TokenInfo], atom_name: str, atoms: list[AtomInfo] +) -> int | None: + if len(residue_tokens) == 1 and residue_tokens[0].res_type != PROTEIN_UNK_RES_TYPE: + return residue_tokens[0].token_index + for token in residue_tokens: + for atom_index in range(token.atom_start, token.atom_start + token.atom_count): + if atom_index < len(atoms) and atoms[atom_index].name == atom_name: + return token.token_index + return residue_tokens[0].token_index if residue_tokens else None + + +def compute_token_bonds( + tokens: list[TokenInfo], + atoms: list[AtomInfo], + input: StructurePredictionInput, + chains: list[ChainInfo], +) -> torch.Tensor: + """Build the symmetric token-bond matrix M with shape (l, l, 1).""" + edges: set[tuple[int, int]] = set() + + def connect(left: int | None, right: int | None) -> None: + if left is not None and right is not None and left != right: + edges.add((min(left, right), max(left, right))) + + explicit_bonds = { + (chain.asym_id, 0): chain.ligand_bonds for chain in chains if chain.ligand_bonds + } + for residue_key, atom_list in _atom_tokenized_residues(tokens, atoms).items(): + if not atom_list: + continue + residue_name = tokens[atom_list[0][1]].residue_name + token_for_name = {name: token_index for name, token_index in atom_list} + bonds = explicit_bonds.get(residue_key) + if bonds is None: + bonds = get_ligand_ccd_bonds(residue_name) + if bonds: + for left_name, right_name in bonds: + if left_name in token_for_name and right_name in token_for_name: + connect(token_for_name[left_name], token_for_name[right_name]) + else: + for left, right in combinations([token_index for _, token_index in atom_list], 2): + connect(left, right) + + if input.covalent_bonds: + chain_for_id = {chain.chain_id: chain for chain in chains} + residue_atoms: dict[tuple[int, int], list[AtomInfo]] = defaultdict(list) + for atom in atoms: + if atom.is_valid and atom.token_index < len(tokens): + token = tokens[atom.token_index] + residue_atoms[(token.asym_id, token.residue_index)].append(atom) + for bond in input.covalent_bonds: + left_chain = chain_for_id.get(bond.chain_id1) + right_chain = chain_for_id.get(bond.chain_id2) + if left_chain is None or right_chain is None: + continue + left_atoms = residue_atoms.get((left_chain.asym_id, bond.res_idx1), []) + right_atoms = residue_atoms.get((right_chain.asym_id, bond.res_idx2), []) + if bond.atom_idx1 < len(left_atoms) and bond.atom_idx2 < len(right_atoms): + connect( + left_atoms[bond.atom_idx1].token_index, + right_atoms[bond.atom_idx2].token_index, + ) + + protein_residues: dict[tuple[int, int], list[TokenInfo]] = defaultdict(list) + for token in tokens: + if token.mol_type == MOL_TYPE_PROTEIN: + protein_residues[(token.asym_id, token.residue_index)].append(token) + for (asym_id, residue_index), residue_tokens in protein_residues.items(): + if not any(token.res_type == PROTEIN_UNK_RES_TYPE for token in residue_tokens): + continue + previous = protein_residues.get((asym_id, residue_index - 1)) + following = protein_residues.get((asym_id, residue_index + 1)) + if previous: + connect( + _backbone_token(previous, "C", atoms), + _backbone_token(residue_tokens, "N", atoms), + ) + if following: + connect( + _backbone_token(residue_tokens, "C", atoms), + _backbone_token(following, "N", atoms), + ) + + matrix = torch.zeros(len(tokens), len(tokens), 1, dtype=torch.float32) + for left, right in edges: + matrix[left, right, 0] = 1.0 + matrix[right, left, 0] = 1.0 + return matrix + + +def compute_representative_atoms(tokens: list[TokenInfo], atoms: list[AtomInfo]) -> torch.Tensor: + """Choose one distogram atom per token and return indices I with shape (l,).""" + named_atoms = _atom_indices_by_name(atoms) + representatives = torch.zeros(len(tokens), dtype=torch.int64) + for token in tokens: + names = named_atoms.get(token.token_index, {}) + fallback = next(iter(names.values()), 0) + if token.mol_type == MOL_TYPE_PROTEIN: + representative = names.get("CB", names.get("CA", fallback)) + elif token.mol_type in (MOL_TYPE_DNA, MOL_TYPE_RNA): + if token.res_type in (27, 32): + representative = names.get("C1'", fallback) + elif token.res_type in (23, 24, 28, 29): + representative = names.get("C4", names.get("C1'", fallback)) + else: + representative = names.get("C2", names.get("C1'", fallback)) + else: + representative = fallback + representatives[token.token_index] = representative + return representatives + + +def _msa_assignments( + input: StructurePredictionInput, chains: list[ChainInfo] +) -> dict[int, MSA | None]: + chain_msas: dict[int, MSA | None] = {} + chain_index = 0 + for item in input.sequences: + chain_ids = [item.id] if isinstance(item.id, str) else list(item.id) + for _ in chain_ids: + chain = chains[chain_index] + if isinstance(item, ProteinInput): + chain_msas[chain.asym_id] = ( + MSA.from_sequences([item.sequence]) if item.msa is None else item.msa + ) + else: + chain_msas[chain.asym_id] = None + chain_index += 1 + return chain_msas + + +def compute_msa_features( + input: StructurePredictionInput, + chains: list[ChainInfo], + tokens: list[TokenInfo], + max_seqs: int = 16384, +) -> dict[str, torch.Tensor]: + """Pair per-chain MSAs and return row features with shape (m, l).""" + from .esmfold2_paired_msa import ( + construct_paired_msa, + protein_letter_to_res_type, + ) + + chain_msas = _msa_assignments(input, chains) + query_types = { + chain.asym_id: np.asarray( + [token.res_type for token in tokens if token.asym_id == chain.asym_id], + dtype=np.int64, + ) + for chain in chains + } + msa_residues, deletion_counts, _ = construct_paired_msa( + chain_msas, + query_types, + np.asarray([token.asym_id for token in tokens], dtype=np.int64), + np.asarray([token.residue_index for token in tokens], dtype=np.int64), + letter_to_res_type=protein_letter_to_res_type(), + max_seqs=max_seqs, + ) + for token in tokens: + if chain_msas.get(token.asym_id) is None: + msa_residues[:, token.token_index] = MSA_GAP_TOKEN_ID + msa_residues[0, token.token_index] = token.res_type + if msa_residues.shape[0] == 0: + msa_residues = np.full((1, len(tokens)), MSA_GAP_TOKEN_ID, dtype=np.int64) + deletion_counts = np.zeros((1, len(tokens)), dtype=np.float32) + + msa = torch.from_numpy(msa_residues) + deletion_count = torch.from_numpy(deletion_counts) + deletion_value = (np.pi / 2) * torch.arctan(deletion_count / 3) + return { + "msa": msa, + "deletion_value": deletion_value, + "has_deletion": deletion_count > 0, + "deletion_mean": deletion_value.mean(dim=0), + "msa_attention_mask": torch.ones_like(msa, dtype=torch.bool), + } + + +def compute_distogram_conditioning( + input: StructurePredictionInput, + chains: list[ChainInfo], + tokens: list[TokenInfo], + disto_center: torch.Tensor, + min_dist: float = 2.0, + max_dist: float = 22.0, + num_bins: int = 64, +) -> tuple[torch.Tensor, torch.Tensor]: + """Bin user distances into D and return D plus its Boolean mask M.""" + del disto_center + n_tokens = len(tokens) + bins = torch.zeros((n_tokens, n_tokens), dtype=torch.long) + mask = torch.zeros((n_tokens, n_tokens), dtype=torch.bool) + if not input.distogram_conditioning: + return bins, mask + asym_for_chain = {chain.chain_id: chain.asym_id for chain in chains} + tokens_for_asym: dict[int, list[int]] = defaultdict(list) + for token in tokens: + tokens_for_asym[token.asym_id].append(token.token_index) + boundaries = torch.linspace(min_dist, max_dist, num_bins + 1) + + for conditioning in input.distogram_conditioning: + asym_id = asym_for_chain.get(conditioning.chain_id) + if asym_id is None: + continue + indices = tokens_for_asym[asym_id] + distances = torch.as_tensor(conditioning.distogram, dtype=torch.float32) + expected_shape = (len(indices), len(indices)) + if distances.shape != expected_shape: + raise ValueError( + f"Distogram shape {distances.shape} doesn't match chain length {len(indices)}" + ) + selected = torch.bucketize(distances, boundaries[:-1]).sub(1).clamp(0, num_bins - 1) + token_indices_tensor = torch.as_tensor(indices, dtype=torch.long) + bins[token_indices_tensor[:, None], token_indices_tensor[None, :]] = selected + mask[token_indices_tensor[:, None], token_indices_tensor[None, :]] = True + return bins, mask + + +def _padded_atoms(atoms: list[AtomInfo]) -> list[AtomInfo]: + target = math.ceil(len(atoms) / 32) * 32 if atoms else 32 + padding = [ + AtomInfo( + name="", + element="", + charge=0, + ref_pos=_ZERO_POS.copy(), + pos=_ZERO_POS.copy(), + token_index=0, + atom_index=index, + space_uid=0, + is_valid=False, + ) + for index in range(len(atoms), target) + ] + return [*atoms, *padding] + + +def _token_tensors(tokens: list[TokenInfo]) -> dict[str, torch.Tensor]: + fields = { + "token_index": "token_index", + "residue_index": "residue_index", + "asym_id": "asym_id", + "sym_id": "sym_id", + "entity_id": "entity_id", + "mol_type": "mol_type", + "res_type": "res_type", + "input_ids": "input_id", + } + return { + output_name: torch.from_numpy( + np.asarray([getattr(token, attribute) for token in tokens], dtype=np.int64) + ) + for output_name, attribute in fields.items() + } + + +def _atom_tensors(atoms: list[AtomInfo]) -> dict[str, torch.Tensor]: + n_atoms = len(atoms) + ref_pos = np.zeros((n_atoms, 3), dtype=np.float32) + ref_element = np.zeros(n_atoms, dtype=np.int64) + ref_charge = np.zeros(n_atoms, dtype=np.int8) + ref_name = np.zeros((n_atoms, 4), dtype=np.int64) + ref_space = np.zeros(n_atoms, dtype=np.int64) + atom_mask = np.zeros(n_atoms, dtype=np.bool_) + atom_to_token = np.zeros(n_atoms, dtype=np.int64) + positions = np.zeros((n_atoms, 3), dtype=np.float64) + valid = np.zeros(n_atoms, dtype=np.bool_) + for index, atom in enumerate(atoms): + if atom.ref_pos is not None: + ref_pos[index] = atom.ref_pos + ref_charge[index] = atom.charge + ref_space[index] = atom.space_uid if atom.space_uid >= 0 else atom.token_index + atom_mask[index] = atom.is_valid + valid[index] = atom.is_valid + positions[index] = atom.pos + if atom.is_valid: + ref_element[index] = get_element_atomic_num(atom.element) + ref_name[index] = encode_atom_name(atom.name) + atom_to_token[index] = atom.token_index + + resolved = valid & np.any(positions != 0, axis=1) + X = torch.from_numpy(positions) + resolved_mask = torch.from_numpy(resolved) + valid_mask = torch.from_numpy(valid) + if resolved_mask.any(): + X = X - X[resolved_mask].mean(dim=0, keepdim=True) + X[~valid_mask] = 0.0 + return { + "ref_pos": torch.from_numpy(ref_pos), + "ref_element": torch.from_numpy(ref_element), + "ref_charge": torch.from_numpy(ref_charge), + "ref_atom_name_chars": torch.from_numpy(ref_name), + "ref_space_uid": torch.from_numpy(ref_space), + "gt_coords": X.float().unsqueeze(0), + "atom_attention_mask": torch.from_numpy(atom_mask), + "atom_to_token": torch.from_numpy(atom_to_token), + "is_resolved": torch.tensor(resolved, dtype=torch.bool), + } + + +def build_feature_tensors( + chains: list[ChainInfo], + tokens: list[TokenInfo], + atoms: list[AtomInfo], + input: StructurePredictionInput, +) -> dict[str, torch.Tensor]: + """Assemble the complete unbatched ESMFold2 feature dictionary.""" + token_features = _token_tensors(tokens) + atom_features = _atom_tensors(_padded_atoms(atoms)) + frames, _ = compute_frame_indices(tokens, atoms) + msa_features = compute_msa_features(input, chains, tokens) + distogram, distogram_mask = compute_distogram_conditioning( + input, + chains, + tokens, + torch.zeros(len(tokens), 3, dtype=torch.float32), + ) + return { + **token_features, + "token_bonds": compute_token_bonds(tokens, atoms, input, chains), + "token_attention_mask": torch.ones(len(tokens), dtype=torch.bool), + "pocket_feature": torch.zeros(len(tokens), dtype=torch.long), + **atom_features, + "distogram_atom_idx": compute_representative_atoms(tokens, atoms), + "frames_idx": torch.from_numpy(frames).to(torch.int64), + "disto_cond": distogram, + "disto_cond_mask": distogram_mask, + **msa_features, + } + + +def prepare_esmfold2_input( + input: StructurePredictionInput, seed: int | None = None +) -> tuple[dict[str, torch.Tensor], list[ChainInfo]]: + """Convert one typed request to model features and output-chain metadata.""" + chains, tokens, atoms = build_chains_from_input(input, seed) + return build_feature_tensors(chains, tokens, atoms, input), chains diff --git a/fastplms/models/esmfold2/esmfold2_processor.py b/fastplms/models/esmfold2/esmfold2_processor.py new file mode 100644 index 0000000000000000000000000000000000000000..ba2192ac5bcbab2a0e6bb955a6d10914929b7e66 --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_processor.py @@ -0,0 +1,338 @@ +"""Input preparation and output decoding for ESMFold2 inference.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from torch import Tensor + +from .esmfold2_conformers import load_ccd +from .esmfold2_molecular_complex import MolecularComplexResult +from .esmfold2_output import build_molecular_complex_from_features +from .esmfold2_prepare_input import ChainInfo, prepare_esmfold2_input +from .esmfold2_types import MSA, Modification, ProteinInput, StructurePredictionInput +from .modeling_esmfold2_common import MSA_CONDITIONING_INPUT_NAMES +from .reproducibility import seed_context + +# Backward-compatible private alias for the pinned parity helpers. New callers +# should import ``seed_context`` from the public ``fastplms.models.esmfold2`` +# package instead of reaching into implementation modules. +_seed_context = seed_context + + +@dataclass(frozen=True) +class _SplitProteinState: + ids: dict[str, list[str]] + modifications: dict[str, list[Modification]] + msas: dict[str, MSA | None] + + +@dataclass(frozen=True) +class _PendingProtein: + source: ProteinInput + sequence: str + state: _SplitProteinState + + +def _chain_starts(chains: list[str]) -> list[int]: + starts: list[int] = [] + position = 0 + for chain in chains: + starts.append(position) + position += len(chain) + 1 + return starts + + +def _split_modifications( + item: ProteinInput, + chains: list[str], + starts: list[int], +) -> dict[str, list[Modification]]: + grouped: dict[str, list[Modification]] = {} + if item.modifications is None: + return grouped + for chain, start in zip(chains, starts, strict=True): + end = start + len(chain) + adjusted = [ + Modification(position=modification.position - start, ccd=modification.ccd) + for modification in item.modifications + if start <= modification.position < end + ] + grouped.setdefault(chain, []).extend(adjusted) + return grouped + + +def _split_msas( + item: ProteinInput, + chains: list[str], + starts: list[int], +) -> dict[str, MSA | None]: + grouped: dict[str, MSA | None] = {} + if item.msa is None: + return grouped + for chain, start in zip(chains, starts, strict=True): + if chain not in grouped: + grouped[chain] = item.msa.select_positions(np.arange(start, start + len(chain))) + return grouped + + +def _split_protein(item: ProteinInput) -> tuple[list[_PendingProtein], _SplitProteinState]: + chains = ":".join(item.sequence.split("|")).split(":") + starts = _chain_starts(chains) + base_id = item.id[0] if isinstance(item.id, list) else item.id + ids: dict[str, list[str]] = {} + for index, chain in enumerate(chains): + chain_ids = ids.setdefault(chain, []) + chain_ids.append(f"{base_id}_{index}") + state = _SplitProteinState( + ids=ids, + modifications=_split_modifications(item, chains, starts), + msas=_split_msas(item, chains, starts), + ) + pending = [ + _PendingProtein(item, chain, state) + for chain, chain_ids in ids.items() + if chain_ids + ] + return pending, state + + +def _resolve_pending(pending: _PendingProtein) -> ProteinInput: + item = pending.source + sequence = pending.sequence + state = pending.state + return ProteinInput( + id=state.ids[sequence], + sequence=sequence, + msa=state.msas.get(sequence) if item.msa else None, + modifications=(state.modifications.get(sequence) if item.modifications else None), + ) + + +def clean_esmfold2_input(input: StructurePredictionInput) -> StructurePredictionInput: + """Expand chain delimiters and group repeated protein sequences by entity.""" + + if input.pocket is not None: + raise NotImplementedError( + "ESMFold2 pocket conditioning is present in the upstream input schema but " + "the published ESMFold2 feature pipeline drops it. FastPLMs refuses this " + "input instead of silently emitting an all-zero pocket feature." + ) + if input.distogram_conditioning is not None: + raise NotImplementedError( + "ESMFold2 distogram conditioning is present in the upstream input schema but " + "the published ESMFold2 forward does not consume it. FastPLMs refuses this " + "input instead of silently ignoring the supplied distogram." + ) + + cleaned: list[Any] = [] + for item in input.sequences: + if not isinstance(item, ProteinInput): + cleaned.append(item) + continue + sequence = ":".join(item.sequence.split("|")) + if ":" not in sequence: + cleaned.append(item) + continue + if input.covalent_bonds is not None: + raise ValueError( + "Covalent bonds are not supported when using chainbreaks. " + "Chains must be separated into multiple ProteinInput objects." + ) + pending, _state = _split_protein(item) + cleaned.extend(pending) + + resolved = [ + _resolve_pending(item) if isinstance(item, _PendingProtein) else item + for item in cleaned + ] + return StructurePredictionInput( + sequences=resolved, + pocket=input.pocket, + distogram_conditioning=input.distogram_conditioning, + covalent_bonds=input.covalent_bonds, + ) + + +def _batch_features( + features: dict[str, Any], + device: torch.device | str | None, +) -> dict[str, Any]: + return { + name: (value[None].to(device) if device is not None else value[None]) + if isinstance(value, Tensor) + else value + for name, value in features.items() + } + + +def _sampler_overrides( + noise_scale: float | None, + step_scale: float | None, + max_inference_sigma: int | None, +) -> dict[str, Any]: + values = { + "noise_scale": noise_scale, + "step_scale": step_scale, + "max_inference_sigma": max_inference_sigma, + } + return {name: value for name, value in values.items() if value is not None} + + +class ESMFold2InputBuilder: + """Prepare public input objects, run folding, and decode model tensors.""" + + def __init__(self, ccd_cache: Path | None = None) -> None: + load_ccd(ccd_cache) + + def prepare_input( + self, + input: StructurePredictionInput, + seed: int | None = None, + device: torch.device | str | None = None, + ) -> tuple[dict[str, Any], list[ChainInfo]]: + cleaned = clean_esmfold2_input(input) + with seed_context(seed): + features, chain_infos = prepare_esmfold2_input(cleaned, seed=seed) + return _batch_features(features, device), chain_infos + + def prepare_model_input( + self, + model: Any, + input: StructurePredictionInput, + seed: int | None = None, + device: torch.device | str | None = None, + ) -> tuple[dict[str, Any], list[ChainInfo]]: + """Prepare features while enforcing the checkpoint's MSA contract.""" + + msa_conditioning = getattr(model.config, "msa_conditioning", None) + if not isinstance(msa_conditioning, bool): + raise RuntimeError("The ESMFold2 config has no Boolean msa_conditioning contract.") + if not msa_conditioning: + explicit_msa_ids = [ + item.id + for item in input.sequences + if isinstance(item, ProteinInput) and item.msa is not None + ] + if explicit_msa_ids: + raise ValueError( + "This ESMFold2 checkpoint was trained without MSA conditioning and " + f"rejects explicit MSAs for protein inputs {explicit_msa_ids!r}." + ) + features, chain_infos = self.prepare_input(input, seed=seed, device=device) + if not msa_conditioning: + for name in MSA_CONDITIONING_INPUT_NAMES: + features.pop(name, None) + return features, chain_infos + + def __call__( + self, + input: StructurePredictionInput, + seed: int | None = None, + device: torch.device | str | None = None, + ) -> tuple[dict[str, Any], list[ChainInfo]]: + return self.prepare_input(input, seed=seed, device=device) + + def _decode_sample( + self, + output: Mapping[str, Tensor], + features: dict[str, Tensor], + chain_infos: list[ChainInfo], + sample: int, + complex_id: str, + ) -> MolecularComplexResult: + plddt = output["plddt"][sample] + molecular_complex = build_molecular_complex_from_features( + coords=output["sample_atom_coords"][sample], + plddt=plddt, + atom_mask=features["atom_attention_mask"][0], + ref_element=features["ref_element"][0], + ref_atom_name_chars=features["ref_atom_name_chars"][0], + chain_infos=chain_infos, + complex_id=complex_id, + ) + + def sample_tensor(name: str) -> Tensor | None: + value = output.get(name) + return None if value is None else value[sample].detach().cpu() + + def shared_tensor(name: str) -> Tensor | None: + value = output.get(name) + return None if value is None else value[0].detach().cpu() + + ptm = output.get("ptm") + iptm = output.get("iptm") + return MolecularComplexResult( + complex=molecular_complex, + plddt=plddt.detach().cpu(), + ptm=float(ptm[sample].item()) if ptm is not None else None, + iptm=float(iptm[sample].item()) if iptm is not None else None, + pae=sample_tensor("pae"), + distogram=shared_tensor("distogram_logits"), + pair_chains_iptm=sample_tensor("pair_chains_iptm"), + residue_index=shared_tensor("residue_index"), + entity_id=shared_tensor("entity_id"), + ) + + def decode( + self, + output: Mapping[str, Tensor], + features: dict[str, Tensor], + chain_infos: list[ChainInfo], + *, + num_diffusion_samples: int = 1, + complex_id: str = "pred", + ) -> MolecularComplexResult | list[MolecularComplexResult]: + results = [ + self._decode_sample(output, features, chain_infos, sample, complex_id) + for sample in range(output["sample_atom_coords"].shape[0]) + ] + return results[0] if num_diffusion_samples == 1 and len(results) == 1 else results + + def fold( + self, + model: Any, + input: StructurePredictionInput, + *, + num_loops: int = 3, + num_sampling_steps: int = 200, + num_diffusion_samples: int = 1, + seed: int | None = None, + noise_scale: float | None = None, + step_scale: float | None = None, + max_inference_sigma: int | None = None, + early_exit: bool = False, + complex_id: str = "pred", + ) -> MolecularComplexResult | list[MolecularComplexResult]: + features, chain_infos = self.prepare_model_input( + model, + input, + seed=seed, + device=model.device, + ) + overrides = _sampler_overrides(noise_scale, step_scale, max_inference_sigma) + with torch.no_grad(), seed_context(seed): + output = model( + **features, + num_loops=num_loops, + num_sampling_steps=num_sampling_steps, + num_diffusion_samples=num_diffusion_samples, + early_exit=early_exit, + return_dict=True, + **overrides, + ) + return self.decode( + output, + features, + chain_infos, + num_diffusion_samples=num_diffusion_samples, + complex_id=complex_id, + ) + + +__all__ = ["ESMFold2InputBuilder", "clean_esmfold2_input", "seed_context"] diff --git a/fastplms/models/esmfold2/esmfold2_protein_chain.py b/fastplms/models/esmfold2/esmfold2_protein_chain.py new file mode 100644 index 0000000000000000000000000000000000000000..acaaaf585e354204bb99916f9d2b2f515c58ea29 --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_protein_chain.py @@ -0,0 +1,1450 @@ +"""Protein-chain data, geometry, and serialization for ESMFold2.""" + +from __future__ import annotations + +import io +import warnings +from collections.abc import Mapping, Sequence +from dataclasses import asdict, dataclass, replace +from functools import cached_property +from pathlib import Path +from typing import Any + +import biotite.structure as bs +import brotli +import msgpack +import msgpack_numpy +import numpy as np +import torch +from biotite.database import rcsb +from biotite.structure.io.pdb import PDBFile +from biotite.structure.io.pdbx import CIFCategory, CIFColumn, CIFData, CIFFile +from biotite.structure.io.pdbx import set_structure as set_structure_pdbx +from scipy.spatial import ConvexHull, KDTree +from scipy.spatial.distance import cdist, pdist, squareform + +from . import esmfold2_residue_constants as residue_constants +from .esmfold2_affine3d import Affine3D +from .esmfold2_aligner import Aligner +from .esmfold2_atom_indexer import AtomIndexer +from .esmfold2_metrics import compute_gdt_ts, compute_lddt_ca +from .esmfold2_misc import slice_python_object_as_numpy +from .esmfold2_mmcif_parsing import ( + PLDDT_B_FACTOR_SCALE, + MmcifWrapper, + Residue, + round_mmcif_columns, +) +from .esmfold2_normalize_coordinates import ( + apply_frame_to_coords, + get_protein_normalization_frame, +) +from .esmfold2_protein_structure import index_by_atom_name +from .esmfold2_utils_types import PathOrBuffer + +CHAIN_ID_CONST = "A" + + +def _str_key_to_int_key(values: dict, ignore_keys: list[str] | None = None) -> dict: + """Restore integer dictionary keys after JSON-compatible serialization.""" + ignored = frozenset(ignore_keys or ()) + restored = {} + for key, value in values.items(): + if isinstance(value, dict) and key not in ignored: + value = _str_key_to_int_key(value, ignore_keys=ignore_keys) + restored_key = int(key) if isinstance(key, str) and key.isdigit() else key + restored[restored_key] = value + return restored + + +def _num_non_null_residues(seqres_to_structure_chain: Mapping[int, Residue]) -> int: + return sum(residue.residue_number is not None for residue in seqres_to_structure_chain.values()) + + +def infer_cb( + C, + N, + Ca, + bond_length: float = 1.522, + bond_angle: float = 1.927, + dihedral: float = -2.143, +): + """Infer C-beta coordinates from C, N, and C-alpha coordinates.""" + + def normalize(X: np.ndarray) -> np.ndarray: + return X / np.sqrt(np.square(X).sum(-1, keepdims=True) + 1e-8) + + with np.errstate(invalid="ignore"): + n_to_ca = N - Ca + n_to_c = N - C + axis = normalize(n_to_ca) + normal = normalize(np.cross(n_to_c, axis)) + basis = (axis, np.cross(normal, axis), normal) + offsets = ( + bond_length * np.cos(bond_angle), + bond_length * np.sin(bond_angle) * np.cos(dihedral), + -bond_length * np.sin(bond_angle) * np.sin(dihedral), + ) + return Ca + sum(vector * offset for vector, offset in zip(basis, offsets, strict=False)) + + +def chain_to_ndarray( + atom_array: bs.AtomArray, mmcif: MmcifWrapper, chain_id: str, is_predicted=False +): + if not isinstance(atom_array, bs.AtomArray): + raise TypeError("atom_array must be a biotite AtomArray.") + if not isinstance(mmcif, MmcifWrapper): + raise TypeError("mmcif must be an MmcifWrapper.") + if not isinstance(chain_id, str) or not chain_id: + raise ValueError("chain_id must be a non-empty string.") + if chain_id not in mmcif.chain_to_seqres or chain_id not in mmcif.seqres_to_structure: + raise ValueError(f"mmCIF data does not contain sequence mappings for chain {chain_id!r}.") + entity_id = None + for entity, chains in mmcif.entities.items(): + if chain_id in chains: + entity_id = entity + num_res = len(mmcif.chain_to_seqres[chain_id]) + sequence = mmcif.chain_to_seqres[chain_id] + + atom_positions = np.full([num_res, residue_constants.atom_type_num, 3], np.nan) + atom_mask = np.full([num_res, residue_constants.atom_type_num], False, dtype=bool) + residue_index = np.full([num_res], -1, dtype=np.int64) + insertion_code = np.full([num_res], "", dtype=" ProteinChain: + """Return a ProteinChain object from an pdb file. NOTE: prefer mmcif for rcsb PDB files. + This function is mostly to interface with old PDB files and predicted structures - + it will not fill out the entity id correctly + + Args: + path: PDB path or text buffer. + id: Optional structure identifier. + is_predicted (bool): If True, reads b factor as the confidence readout. Default: False. + chain_id: Author chain identifier. ``"detect"`` selects the first chain. + """ + + if id is not None: + file_id = id + else: + match path: + case Path() | str(): + file_id = Path(path).with_suffix("").name + case _: + file_id = "null" + + atom_array = PDBFile.read(path).get_structure(model=1, extra_fields=["b_factor"]) + if len(atom_array) == 0: + raise ValueError("PDB contains no atoms.") + if chain_id == "detect": + chain_id = atom_array.chain_id[0] + atom_array = atom_array[ + bs.filter_amino_acids(atom_array) + & ~atom_array.hetero + & (atom_array.chain_id == chain_id) + ] + if len(atom_array) == 0: + raise ValueError(f"PDB contains no amino-acid atoms for chain {chain_id!r}.") + + entity_id = 1 # Not supplied in PDBfiles + + sequence = "".join( + residue_constants.restype_3to1.get(monomer[0].res_name, "X") + for monomer in bs.residue_iter(atom_array) + ) + num_res = len(sequence) + + atom_positions = np.full( + [num_res, residue_constants.atom_type_num, 3], np.nan, dtype=np.float32 + ) + atom_mask = np.full([num_res, residue_constants.atom_type_num], False, dtype=bool) + residue_index = np.full([num_res], -1, dtype=np.int64) + insertion_code = np.full([num_res], "", dtype=" ProteinChain: + return cls( + id=data["id"], + chain_id=data["chain_id"], + entity_id=data["entity_id"], + sequence=data["sequence"], + residue_index=data["residue_index"], + insertion_code=np.asarray(data["insertion_code"]), + atom37_positions=data["atom37_positions"], + atom37_mask=data["atom37_mask"].astype(bool), + confidence=data["confidence"], + mmcif=None, + ) + + @classmethod + def from_rcsb( + cls, + pdb_id: str, + chain_id: str | None = None, + entity_id: int | None = None, + keep_source: bool = False, + ) -> ProteinChain: + f: io.StringIO = rcsb.fetch(pdb_id, "cif") # type: ignore + return cls.from_mmcif( + f, + id=pdb_id, + chain_id=chain_id, + entity_id=entity_id, + keep_source=keep_source, + is_predicted=False, + ) + + @classmethod + def from_atomarray( + cls, atom_array: bs.AtomArray, id: str | None = None, is_predicted: bool = False + ) -> ProteinChain: + """A simple converter from bs.AtomArray -> ProteinChain. + Uses PDB file format as intermediate.""" + atom_array = atom_array.copy() + atom_array.box = None # remove surrounding box, from_pdb won't handle this + pdb_file = PDBFile() # pyright: ignore + pdb_file.set_structure(atom_array) + + buf = io.StringIO() + pdb_file.write(buf) + buf.seek(0) + return cls.from_pdb(buf, id=id, is_predicted=is_predicted) + + # Object invariants and atom views + def __post_init__(self): + if not isinstance(self.id, str): + raise TypeError("id must be a string.") + if not isinstance(self.sequence, str): + raise TypeError("sequence must be a string.") + if not isinstance(self.chain_id, str) or not self.chain_id: + raise ValueError("chain_id must be a non-empty string.") + if self.entity_id is not None and ( + not isinstance(self.entity_id, int) or isinstance(self.entity_id, bool) + ): + raise TypeError("entity_id must be an integer or None.") + sequence_length = len(self.sequence) + aligned = { + "atom37_positions": self.atom37_positions, + "atom37_mask": self.atom37_mask, + "residue_index": self.residue_index, + "insertion_code": self.insertion_code, + "confidence": self.confidence, + } + for name, values in aligned.items(): + if not isinstance(values, np.ndarray): + raise TypeError(f"{name} must be a NumPy array, got {type(values).__name__}.") + if values.ndim == 0 or values.shape[0] != sequence_length: + raise ValueError( + f"{name} shape {values.shape} does not align with " + f"sequence length {sequence_length}." + ) + if self.atom37_positions.shape != (sequence_length, 37, 3): + raise ValueError( + "atom37_positions must have shape " + f"({sequence_length}, 37, 3), got {self.atom37_positions.shape}." + ) + if self.atom37_mask.shape != (sequence_length, 37): + raise ValueError( + "atom37_mask must have shape " + f"({sequence_length}, 37), got {self.atom37_mask.shape}." + ) + if self.atom37_mask.dtype != bool: + raise TypeError(f"atom37_mask must have Boolean dtype, got {self.atom37_mask.dtype}.") + if not np.issubdtype(self.atom37_positions.dtype, np.number): + raise TypeError("atom37_positions must use a numeric dtype.") + if not np.issubdtype(self.residue_index.dtype, np.integer): + raise TypeError("residue_index must use an integer dtype.") + if self.insertion_code.dtype.kind not in {"U", "S", "O"}: + raise TypeError("insertion_code must use a string-compatible dtype.") + if any(not isinstance(value, str) for value in self.insertion_code.tolist()): + raise TypeError("insertion_code must contain only strings.") + for name, values in ( + ("residue_index", self.residue_index), + ("insertion_code", self.insertion_code), + ("confidence", self.confidence), + ): + if values.shape != (sequence_length,): + raise ValueError( + f"{name} must have shape ({sequence_length},), got {values.shape}." + ) + if not np.issubdtype(self.confidence.dtype, np.number): + raise TypeError("confidence must use a numeric dtype.") + atom37_confidence = self.atom37_confidence + if atom37_confidence is not None and not isinstance(atom37_confidence, np.ndarray): + raise TypeError("atom37_confidence must be a NumPy array when provided.") + if ( + isinstance(atom37_confidence, np.ndarray) + and atom37_confidence.shape != self.atom37_mask.shape + ): + raise ValueError( + "atom37_confidence shape must match atom37_mask: " + f"{atom37_confidence.shape} != {self.atom37_mask.shape}." + ) + if isinstance(atom37_confidence, np.ndarray) and not np.issubdtype( + atom37_confidence.dtype, np.number + ): + raise TypeError("atom37_confidence must use a numeric dtype.") + + @cached_property + def atoms(self) -> AtomIndexer: + return AtomIndexer(self, property="atom37_positions", dim=-2) + + @cached_property + def atom_mask(self) -> AtomIndexer: + return AtomIndexer(self, property="atom37_mask", dim=-1) + + @cached_property + def atom_array(self) -> bs.AtomArray: + atoms = [] + for res_idx_i, ( + res_name, + res_idx, + ins_code, + positions, + mask, + conf, + ) in enumerate( + zip( + self.sequence, + self.residue_index, + self.insertion_code, + self.atom37_positions, + self.atom37_mask.astype(bool), + self.confidence, + strict=False, + ) + ): + for i, pos in zip(np.where(mask)[0], positions[mask], strict=False): + b_factor = ( + self.atom37_confidence[res_idx_i, i] + if self.atom37_confidence is not None + else conf + ) + atom = bs.Atom( + coord=pos, + chain_id="A" if self.chain_id is None else self.chain_id, + res_id=res_idx, + ins_code=ins_code, + res_name=residue_constants.restype_1to3.get(res_name, "UNK"), + hetero=False, + atom_name=residue_constants.atom_types[i], + element=residue_constants.atom_types[i][0], + b_factor=float(b_factor) * PLDDT_B_FACTOR_SCALE, + occupancy=1.0, + ) + atoms.append(atom) + return bs.array(atoms) + + # Coordinate transformations and dataset adapters + def get_normalization_frame(self) -> Affine3D: + """Given a set of coordinates, compute a single frame. + The frame is built from the mean N, C-alpha, and C coordinates with + Gram-Schmidt orthogonalization. Its origin is the mean C-alpha position. + + Returns: + Affine3D: [] tensor of Affine3D frame + """ + coords = torch.from_numpy(self.atom37_positions) + frame = get_protein_normalization_frame(coords) + + return frame + + def apply_frame(self, frame: Affine3D) -> ProteinChain: + """Given a frame, apply the frame to the protein's coordinates. + + Args: + frame (Affine3D): [] tensor of Affine3D frame + + Returns: + ProteinChain: Transformed protein chain + """ + coords = torch.from_numpy(self.atom37_positions).to(frame.trans.dtype) + coords = apply_frame_to_coords(coords, frame) + atom37_positions = coords.numpy() + return replace(self, atom37_positions=atom37_positions) + + def normalize_coordinates(self) -> ProteinChain: + """Normalize the coordinates of the protein chain.""" + return self.apply_frame(self.get_normalization_frame()) + + def infer_oxygen(self) -> ProteinChain: + """Oxygen position is fixed given N, CA, C atoms. Infer it if not provided.""" + O_missing_indices = np.argwhere(~np.isfinite(self.atoms["O"]).all(axis=1)).squeeze() + + O_vector = torch.tensor([0.6240, -1.0613, 0.0103], dtype=torch.float32) + N, CA, C = torch.from_numpy(self.atoms[["N", "CA", "C"]]).float().unbind(dim=1) + N = torch.roll(N, -3) + N[..., -1, :] = torch.nan + + # Get the frame defined by the CA-C-N atom + frames = Affine3D.from_graham_schmidt(CA, C, N) + oxygen_coordinates = frames.apply(O_vector) + atom37_positions = self.atom37_positions.copy() + atom37_mask = self.atom37_mask.copy() + + atom37_positions[O_missing_indices, residue_constants.atom_order["O"]] = oxygen_coordinates[ + O_missing_indices + ].numpy() + atom37_mask[O_missing_indices, residue_constants.atom_order["O"]] = ~np.isnan( + atom37_positions[O_missing_indices, residue_constants.atom_order["O"]] + ).any(-1) + new_chain = replace(self, atom37_positions=atom37_positions, atom37_mask=atom37_mask) + return new_chain + + @cached_property + def inferred_cbeta(self) -> np.ndarray: + """Infer cbeta positions based on N, C, CA.""" + N, CA, C = np.moveaxis(self.atoms[["N", "CA", "C"]], 1, 0) + # See usage in trDesign codebase. + # https://github.com/gjoni/trDesign/blob/f2d5930b472e77bfacc2f437b3966e7a708a8d37/02-GD/utils.py#L140 + CB = infer_cb(C, N, CA, 1.522, 1.927, -2.143) + return CB + + def infer_cbeta(self, infer_cbeta_for_glycine: bool = False) -> ProteinChain: + """Return a new chain with inferred CB atoms at all residues except GLY. + + Args: + infer_cbeta_for_glycine (bool): If True, infers a beta carbon for glycine + residues, even though that residue doesn't have one. Default off. + + NOTE(rverkuil): The reason for having this switch in the first place + is that sometimes we want a (inferred) CB coordinate for every residue, + for example for making a pairwise distance matrix, or doing an RMSD + calculation between two designs for a given structural template, w/ + CB atoms. + """ + atom37_positions = self.atom37_positions.copy() + atom37_mask = self.atom37_mask.copy() + + inferred_cbeta_positions = self.inferred_cbeta + if not infer_cbeta_for_glycine: + inferred_cbeta_positions[np.array(list(self.sequence)) == "G", :] = np.nan + + atom37_positions[:, residue_constants.atom_order["CB"]] = inferred_cbeta_positions + atom37_mask[:, residue_constants.atom_order["CB"]] = ~np.isnan( + atom37_positions[:, residue_constants.atom_order["CB"]] + ).any(-1) + new_chain = replace(self, atom37_positions=atom37_positions, atom37_mask=atom37_mask) + return new_chain + + @cached_property + def pdist_CA(self) -> np.ndarray: + CA = self.atoms["CA"] + pdist_CA = squareform(pdist(CA)) + return pdist_CA + + @cached_property + def pdist_CB(self) -> np.ndarray: + pdist_CB = squareform(pdist(self.inferred_cbeta)) + return pdist_CB + + @classmethod + def as_complex(cls, chains: Sequence[ProteinChain]): + raise RuntimeError( + ".as_complex() has been deprecated in favor of .concat(). " + ".concat() will eventually be deprecated in favor of ProteinComplex..." + ) + + @classmethod + def concat(cls, chains: Sequence[ProteinChain], use_chainbreak: bool = True): + if not chains: + raise ValueError("chains must contain at least one ProteinChain.") + if any(not isinstance(chain, ProteinChain) for chain in chains): + raise TypeError("chains must contain only ProteinChain instances.") + sep_tokens = { + "residue_index": np.array([-1]), + "insertion_code": np.array([""]), + "atom37_positions": np.full([1, 37, 3], np.inf), + "atom37_mask": np.zeros([1, 37], dtype=bool), + "confidence": np.array([0]), + } + + def join_arrays(arrays: Sequence[np.ndarray], sep: np.ndarray): + if use_chainbreak: + full_array = [] + for array in arrays: + full_array.append(array) + full_array.append(sep) + full_array = full_array[:-1] + return np.concatenate(full_array, 0) + else: + return np.concatenate(arrays, 0) + + array_args: dict[str, np.ndarray] = { + name: join_arrays([getattr(chain, name) for chain in chains], sep) + for name, sep in sep_tokens.items() + } + + chain_break = residue_constants.CHAIN_BREAK_TOKEN if use_chainbreak else "" + return cls( + id=chains[0].id, + sequence=chain_break.join(chain.sequence for chain in chains), + chain_id="A", + entity_id=None, + mmcif=None, + **array_args, + ) + + def find_nonpolymer_contacts(self): + if self.mmcif is None: + raise ValueError( + "find_nonpolymer_contacts requires a chain loaded with keep_source=True." + ) + nonpolymer_and_chain_id_to_array = self.mmcif.non_polymer_coords + + results = [] + for ( + nonpolymer, + _, + ), nonpolymer_array in nonpolymer_and_chain_id_to_array.items(): + if nonpolymer_array.coord is None: + raise ValueError( + f"Non-polymer {nonpolymer.comp_id!r} has no coordinate table." + ) + chain_coords = self.atom37_positions[self.atom37_mask] + distance = cdist(nonpolymer_array.coord, chain_coords) + + is_contact = distance < 5 + if not is_contact.any(): + continue + contacting_atoms = np.where(is_contact.any(0))[0] + chain_index = np.where(self.atom37_mask)[0] + contacting_residues = np.unique(chain_index[contacting_atoms]) + + result = { + "ligand": nonpolymer.name, + "ligand_id": nonpolymer.comp_id, + "contacting_residues": contacting_residues.tolist(), + } + results.append(result) + return results + + def select_residue_indices( + self, indices: list[int | str], ignore_x_mismatch: bool = False + ) -> ProteinChain: + numeric_indices = [idx if isinstance(idx, int) else int(idx[1:]) for idx in indices] + mask = np.isin(self.residue_index, numeric_indices) + new = self[mask] + mismatches = [] + for aa, idx in zip(new.sequence, indices, strict=False): + if isinstance(idx, int): + continue + if aa == "X" and ignore_x_mismatch: + continue + if aa != idx[0]: + mismatches.append((aa, idx)) + if mismatches: + mismatch_str = "; ".join( + f"Position {idx[1:]}, Expected: {idx[0]}, Received: {aa}" for aa, idx in mismatches + ) + raise RuntimeError(mismatch_str) + + return new + + def to_structure_encoder_inputs( + self, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Convert protein chain to structure encoder inputs. + + Returns: + tuple: (coordinates, plddt, residue_index) where: + - coordinates: X with shape (1, l, 37, 3), containing atom positions + - plddt: P with shape (1, l), containing confidence scores + - residue_index: R with shape (1, l), containing residue indices + """ + # Convert to tensors and add batch dimension + coordinates = ( + torch.from_numpy(self.atom37_positions).float().unsqueeze(0) + ) # X has shape (1, l, 37, 3). + plddt = torch.from_numpy(self.confidence).float().unsqueeze(0) # P: (1, l) + residue_index = ( + torch.from_numpy(self.residue_index).long().unsqueeze(0) + ) # R has shape (1, l). + + return coordinates, plddt, residue_index + + # Sequence access, interchange, and compact storage + def __getitem__(self, idx: int | list[int] | slice | np.ndarray | torch.Tensor): + if isinstance(idx, int): + idx = [idx] + if isinstance(idx, torch.Tensor): + idx = idx.cpu().numpy() + + sequence = slice_python_object_as_numpy(self.sequence, idx) + return replace( + self, + sequence=sequence, + residue_index=self.residue_index[..., idx], + insertion_code=self.insertion_code[..., idx], + atom37_positions=self.atom37_positions[..., idx, :, :], + atom37_mask=self.atom37_mask[..., idx, :], + confidence=self.confidence[..., idx], + atom37_confidence=self.atom37_confidence[..., idx, :] + if self.atom37_confidence is not None + else None, + ) + + def __len__(self): + return len(self.sequence) + + def cbeta_contacts(self, distance_threshold: float = 8.0) -> np.ndarray: + distance = self.pdist_CB + contacts = (distance < distance_threshold).astype(np.int64) + contacts[np.isnan(distance)] = -1 + np.fill_diagonal(contacts, -1) + return contacts + + def to_pdb(self, path: PathOrBuffer, include_insertions: bool = True): + """Dssp works better w/o insertions.""" + f = PDBFile() + if not include_insertions: + f.set_structure(self.atom_array_no_insertions) + else: + f.set_structure(self.atom_array) + f.write(path) + + def to_pdb_string(self, include_insertions: bool = True) -> str: + buf = io.StringIO() + self.to_pdb(buf, include_insertions=include_insertions) + buf.seek(0) + return buf.read() + + def to_mmcif(self, path: PathOrBuffer): + f = CIFFile() + set_structure_pdbx(f, self.atom_array, data_block=self.id) + + # incantations molstar needs to render pLDDT / confidence onto + # the structure with "alphafold-view" + f.block["ma_qa_metric"] = CIFCategory( + name="ma_qa_metric", + columns={ + "id": CIFColumn(data=CIFData(array=np.array([1, 2]), dtype=np.int64)), + "mode": CIFColumn(data=CIFData(array=np.array(["global", "local"]), dtype=np.str_)), + "name": CIFColumn(data=CIFData(array=np.array(["pLDDT", "pLDDT"]), dtype=np.str_)), + }, + ) + + # table is a duplicate of data already in the atom array, but + # needed by molstar to render pLDDT / confidence + resid_pldd_table = { + # hard coded to as we currently only support single chain structures + "label_asym_id": CIFColumn( + data=CIFData(array=[CHAIN_ID_CONST] * len(self.residue_index), dtype=np.str_) + ), + "label_comp_id": CIFColumn( + data=CIFData( + array=[residue_constants.restype_1to3.get(c, "UNK") for c in self.sequence], + dtype=np.str_, + ) + ), + "label_seq_id": CIFColumn(data=CIFData(array=self.residue_index, dtype=np.int64)), + "ordinal_id": CIFColumn(data=CIFData(array=self.residue_index, dtype=np.int64)), + # hard coded to show these are all local plDDT values + "metric_id": CIFColumn( + data=CIFData(array=["2"] * len(self.residue_index), dtype=np.str_) + ), + "metric_value": CIFColumn( + data=CIFData( + array=self.confidence * PLDDT_B_FACTOR_SCALE, + dtype=np.float32, + ) + ), + # hard coded to show there are the initial version, there are no revisions + "model_id": CIFColumn( + data=CIFData(array=["1"] * len(self.residue_index), dtype=np.str_) + ), + } + f.block["ma_qa_metric_local"] = CIFCategory( + name="ma_qa_metric_local", columns=resid_pldd_table + ) + round_mmcif_columns(f) + f.write(path) + + def to_mmcif_string(self) -> str: + buf = io.StringIO() + self.to_mmcif(buf) + buf.seek(0) + return buf.read() + + def state_dict(self, backbone_only=False, json_serializable=False): + """This state dict is optimized for storage, so it turns things to fp16 whenever + possible. Note that we also only support int32 residue indices, I'm hoping we don't + need more than 2**32 residues...""" + dct = {k: v for k, v in asdict(self).items() if k not in ["mmcif"]} + if backbone_only: + dct["atom37_mask"][:, 3:] = False + dct["atom37_positions"] = dct["atom37_positions"][dct["atom37_mask"]] + if dct.get("atom37_confidence") is not None: + dct["atom37_confidence"] = dct["atom37_confidence"][dct["atom37_mask"]] + else: + dct.pop("atom37_confidence", None) + + for k, v in dct.items(): + if isinstance(v, np.ndarray): + match v.dtype: + case np.int64: + dct[k] = v.astype(np.int32) + case np.float64 | np.float32: + dct[k] = v.astype(np.float16) + case _: + pass + if json_serializable: + dct[k] = v.tolist() + return dct + + def to_blob(self, backbone_only=False) -> bytes: + payload = msgpack.dumps(self.state_dict(backbone_only), default=msgpack_numpy.encode) + return brotli.compress(payload, quality=5) + + @classmethod + def from_open_source(cls, pc: ProteinChain): + return cls(**vars(pc)) + + @classmethod + def from_state_dict(cls, dct): + # Note: assembly_composition is *supposed* to have string keys. + dct = _str_key_to_int_key(dct, ignore_keys=["assembly_composition"]) + + for k, v in dct.items(): + if isinstance(v, list): + dct[k] = np.array(v) + + atom37 = np.full((*dct["atom37_mask"].shape, 3), np.nan) + atom37[dct["atom37_mask"]] = dct["atom37_positions"] + dct["atom37_positions"] = atom37 + if "atom37_confidence" in dct: + atom37_conf = np.full(dct["atom37_mask"].shape, np.nan, dtype=np.float32) + atom37_conf[dct["atom37_mask"]] = dct["atom37_confidence"] + dct["atom37_confidence"] = atom37_conf + dct = { + k: ( + v.astype(np.float32) + if k in ["atom37_positions", "confidence", "atom37_confidence"] + else v + ) + for k, v in dct.items() + if not (k == "atom37_confidence" and v is None) + } + return cls(**dct, mmcif=None) + + @classmethod + def from_blob(cls, input: Path | str | io.BytesIO | bytes): + """NOTE(@zlin): blob + sparse coding + brotli + fp16 reduces memory + of chains from 52G/1M chains to 20G/1M chains, I think this is a good first + shot at compressing and dumping chains to disk. I'm sure there's better ways.""" + match input: + case Path() | str(): + bytes = Path(input).read_bytes() + case io.BytesIO(): + bytes = input.getvalue() + case _: + bytes = input + state = msgpack.loads(brotli.decompress(bytes), object_hook=msgpack_numpy.decode) + return cls.from_state_dict(state) + + # Surface and structural comparison metrics + def sasa(self, by_residue: bool = True): + arr = self.atom_array_no_insertions + if len(arr) == 0: + raise ValueError("SASA requires at least one resolved atom.") + sasa_per_atom = bs.sasa(arr) # type: ignore + if by_residue: + # Sum per-atom SASA into residue "bins", with np.bincount. + if arr.res_id is None: + raise RuntimeError("Biotite AtomArray is missing residue identifiers.") + # Residue IDs are one-indexed, so discard the unused zero bin. + # NOTE(aderry): We compute only for residues with coordinates, return NaN otherwise. + num_trailing_residues = len(self) - arr.res_id.max() + sasa_per_residue = np.concatenate( + [ + np.bincount(arr.res_id, weights=sasa_per_atom)[1:], + np.zeros(num_trailing_residues), + ] + ) + sasa_per_residue[~self.atom37_mask.any(-1)] = np.nan + if len(sasa_per_residue) != len(self): + raise RuntimeError("Residue SASA output does not align with the protein chain.") + return sasa_per_residue + return sasa_per_atom + + def sap_score(self, aggregation: str = "atom") -> np.ndarray: + """Compute per-atom spatial aggregation propensity (SAP). + + Residue aggregation averages resolved atoms and omits unresolved residues. + Protein aggregation sums positive atom scores, following Lauer et al. 2011. + """ + sap_radius = 5.0 + arr = self.atom_array_no_insertions + if len(arr) == 0: + raise ValueError("SAP requires at least one resolved atom.") + + for name in ("res_id", "res_name", "atom_name", "coord"): + if getattr(arr, name) is None: + raise RuntimeError(f"Biotite AtomArray is missing required {name!r} data.") + + # compute SASA and residue-specific properties + sasa_per_atom = self.sasa(by_residue=False) + resid_to_resname = dict(zip(arr.res_id, arr.res_name, strict=False)) + + max_side_chain_asa = np.full(len(self), np.nan) + res_hydrophobicity = np.full(len(self), np.nan) + resolved_res_mask = self.atom37_mask.any(-1) + num_trailing_residues = len(self) - arr.res_id.max() + + max_side_chain_asa[resolved_res_mask] = np.array( + [residue_constants.side_chain_asa[resid_to_resname[i]] for i in np.unique(arr.res_id)] + ) + res_hydrophobicity[resolved_res_mask] = np.array( + [residue_constants.hydrophobicity[resid_to_resname[i]] for i in np.unique(arr.res_id)] + ) + + # compute SAP score + is_side_chain = ~bs.filter_peptide_backbone(arr) + sasa_per_atom[is_side_chain] = 0 + kdtree = KDTree(arr.coord) + neighbors = kdtree.query_ball_tree(kdtree, sap_radius, p=2.0) + sap_by_atom = np.zeros_like(sasa_per_atom) + for i, nn_list in enumerate(neighbors): + saa_nn = np.zeros_like(sasa_per_atom) + saa_nn[nn_list] = sasa_per_atom[nn_list] + sasa_within_r = np.concatenate( + [ + np.bincount(arr.res_id, weights=saa_nn)[1:], + np.zeros(num_trailing_residues), + ] + ) + sap = np.nansum((sasa_within_r / max_side_chain_asa) * res_hydrophobicity) + sap_by_atom[i] = sap + + match aggregation: + case "atom": + return sap_by_atom + case "residue": + sap_by_residue = np.concatenate( + [ + np.bincount(arr.res_id, weights=sap_by_atom)[1:], + np.zeros(num_trailing_residues), + ] + ) / ( + np.concatenate([np.bincount(arr.res_id)[1:], np.zeros(num_trailing_residues)]) + + 1e-8 + ) + sap_by_residue[~resolved_res_mask] = np.nan + if len(sap_by_residue) != len(self): + raise RuntimeError("Residue SAP output does not align with the protein chain.") + return sap_by_residue + case "protein": + return sum(sap_by_atom[sap_by_atom > 0]) # pyright: ignore[reportReturnType] + case _: + raise ValueError( + f"Invalid aggregation method: {aggregation}. Must be one of " + "'atom', 'residue', or 'protein'" + ) + + def globularity(self) -> float: + # Computes globularity using total volumes divided by MVEE. + # We make the simplifying approximation that atoms never overlap. + # The globularity is only computed where structure exists. + # Besides the approximation above, this is inspired by: + + # https://www.mdpi.com/2073-4352/11/12/1539 + # The non-overlapping-atom approximation can produce globularity above one. + mask = self.atom37_mask.any(-1) + points = self.atom37_positions[self.atom37_mask] + sequence = [aa for aa, m in zip(self.sequence, mask, strict=False) if m] # type: ignore + A, _ = self._mvee(points, tol=1e-3) + mvee_volume = (4 * np.pi) / (3 * np.sqrt(np.linalg.det(A))) + volume = sum(residue_constants.amino_acid_volumes[x] for x in sequence) + ratio = volume / mvee_volume + + # The paper compares the ellipsoidal profile with scalar t, a measurement + # of elongation. We want a single number, so we multiply by 1/(2t), so + # that value is normalized between 0-1 + eigenvalues = np.linalg.eigvals(A) + R = 1 / np.sqrt(eigenvalues) + # ellipsoid radii length triangle inequality coefficient + t = max(R[0] / (R[1] + R[2]), R[1] / (R[0] + R[2]), R[2] / (R[0] + R[1])) + elongation_metric = 1 / max(t, 1) + return ratio * elongation_metric + + @staticmethod + def _mvee(P: np.ndarray, tol, max_iter=10000): + # Finds minimum volume enclosing ellipsoid of a set of points. + # Returns A, c where the ellipse is defined as: + # (x-c).T @ A @ (x-c) = 1 + hull = ConvexHull(P) + P = P[hull.vertices] + P = P.T + + # Data points + d, n = P.shape + Q = np.zeros((d + 1, n)) + Q[:d, :] = P[:d, :n] + Q[d, :] = np.ones((1, n)) + + # Initializations + count = 1 + err = 1.0 + u = np.full((n, 1), 1 / n) # First iteration. + + # Khachiyan Algorithm + for _ in range(max_iter): + X = Q.dot(np.diag(u.squeeze())) @ Q.T + M = np.diag(Q.T @ np.linalg.inv(X) @ Q) + maximum, j = np.max(M), np.argmax(M) + step_size = (maximum - d - 1) / ((d + 1) * (maximum - 1)) + new_u = (1 - step_size) * u + new_u[j] += step_size + count += 1 + err = np.linalg.norm(new_u - u) + u = new_u + if err < tol: + break + else: + raise ValueError("MVEE did not converge") + + d = P.shape[0] # Fixed: use P.shape[0] instead of P.shape + U = np.diag(u.squeeze()) + + # The A matrix for the ellipse + A = (1 / d) * np.linalg.inv(P @ U @ P.T - (P @ u) @ (P @ u).T) + + # Center of the ellipse + c = P @ u + + return A, c + + def radius_of_gyration(self): + arr = self.atom_array_no_insertions + return bs.gyration_radius(arr) + + def align( + self, + target: ProteinChain, + mobile_inds: list[int] | np.ndarray | None = None, + target_inds: list[int] | np.ndarray | None = None, + only_use_backbone: bool = False, + ): + """ + Aligns the current protein to the provided target. + + Args: + target (ProteinChain): The target protein to align to. + mobile_inds: Mobile atom indices, not residue indices. + target_inds: Target atom indices, not residue indices. + only_use_backbone (bool, optional): If True, only align the backbone atoms. + """ + aligner = Aligner( + self if mobile_inds is None else self[mobile_inds], + target if target_inds is None else target[target_inds], + only_use_backbone, + ) + + return aligner.apply(self) + + def rmsd( + self, + target: ProteinChain, + also_check_reflection: bool = False, + mobile_inds: list[int] | np.ndarray | None = None, + target_inds: list[int] | np.ndarray | None = None, + only_compute_backbone_rmsd: bool = False, + ): + """ + Compute the RMSD between this protein chain and another. + + Args: + target (ProteinChain): The target (other) protein chain to compare to. + also_check_reflection: Compare the reflected mobile coordinates too. + mobile_inds: Mobile atom indices, not residue indices. + target_inds: Target atom indices, not residue indices. + only_compute_backbone_rmsd: Restrict the score to backbone atoms. + """ + if isinstance(target, bs.AtomArray): + raise ValueError( + "Support for bs.AtomArray removed, use ProteinChain.from_atomarry for ProteinChain." + ) + aligner = Aligner( + self if mobile_inds is None else self[mobile_inds], + target if target_inds is None else target[target_inds], + only_compute_backbone_rmsd, + ) + avg_rmsd = aligner.rmsd + + if not also_check_reflection: + return avg_rmsd + + aligner = Aligner( + self if mobile_inds is None else self[mobile_inds], + target if target_inds is None else target[target_inds], + only_compute_backbone_rmsd, + use_reflection=True, + ) + avg_rmsd_neg = aligner.rmsd + + return min(avg_rmsd, avg_rmsd_neg) + + def lddt_ca( + self, + native: ProteinChain, + mobile_inds: list[int] | np.ndarray | None = None, + target_inds: list[int] | np.ndarray | None = None, + **kwargs, + ) -> float | np.ndarray: + """Compute the LDDT between this protein chain and another. NOTE: LDDT IS NOT SYMMETRIC. + The call should always be prediction.lddt_ca(native). + + Arguments: + native (ProteinChain): The ground truth protein chain + mobile_inds: Mobile atom indices, not residue indices. + target_inds: Target atom indices, not residue indices. + + Returns: + float | np.ndarray: The LDDT score between the two protein chains, either + a single float or per-residue LDDT scores if `per_residue` is True. + """ + lddt = compute_lddt_ca( + torch.tensor(self.atom37_positions[mobile_inds]).unsqueeze(0), + torch.tensor(native.atom37_positions[target_inds]).unsqueeze(0), + torch.tensor(native.atom37_mask[mobile_inds]).unsqueeze(0), + **kwargs, + ) + return float(lddt) if lddt.numel() == 1 else lddt.numpy().flatten() + + def gdt_ts( + self, + target: ProteinChain, + mobile_inds: list[int] | np.ndarray | None = None, + target_inds: list[int] | np.ndarray | None = None, + **kwargs, + ) -> float | np.ndarray: + """Compute the GDT_TS between this protein chain and another. + + Arguments: + target (ProteinChain): The other protein chain to compare to. + mobile_inds: Mobile atom indices, not residue indices. + target_inds: Target atom indices, not residue indices. + + Returns: + float: The GDT_TS score between the two protein chains. + """ + gdt_ts = compute_gdt_ts( + mobile=torch.tensor( + index_by_atom_name(self.atom37_positions[mobile_inds], "CA"), + dtype=torch.float32, + ).unsqueeze(0), + target=torch.tensor( + index_by_atom_name(target.atom37_positions[target_inds], "CA"), + dtype=torch.float32, + ).unsqueeze(0), + atom_exists_mask=torch.tensor( + index_by_atom_name(self.atom37_mask[mobile_inds], "CA", dim=-1) + & index_by_atom_name(target.atom37_mask[target_inds], "CA", dim=-1) + ).unsqueeze(0), + **kwargs, + ) + return float(gdt_ts) if gdt_ts.numel() == 1 else gdt_ts.numpy().flatten() + + @cached_property + def residue_index_no_insertions(self) -> np.ndarray: + return self.residue_index + np.cumsum(self.insertion_code != "") + + @cached_property + def atom_array_no_insertions(self) -> bs.AtomArray: + atoms = [] + for res_idx, (res_name, positions, mask, conf) in enumerate( + zip( + self.sequence, + self.atom37_positions, + self.atom37_mask.astype(bool), + self.confidence, + strict=False, + ) + ): + for i, pos in zip(np.where(mask)[0], positions[mask], strict=False): + b_factor = ( + self.atom37_confidence[res_idx, i] + if self.atom37_confidence is not None + else conf + ) + atom = bs.Atom( + coord=pos, + # hard coded to as we currently only support single chain structures + chain_id=CHAIN_ID_CONST, + res_id=res_idx + 1, + res_name=residue_constants.restype_1to3.get(res_name, "UNK"), + hetero=False, + atom_name=residue_constants.atom_types[i], + element=residue_constants.atom_types[i][0], + b_factor=float(b_factor) * PLDDT_B_FACTOR_SCALE, + occupancy=1.0, + ) + atoms.append(atom) + return bs.array(atoms) diff --git a/fastplms/models/esmfold2/esmfold2_protein_complex.py b/fastplms/models/esmfold2/esmfold2_protein_complex.py new file mode 100644 index 0000000000000000000000000000000000000000..688ea5ce68efecdcc522e1ec4015d60d12bf782a --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_protein_complex.py @@ -0,0 +1,1240 @@ +"""Protein-complex data, assembly expansion, and geometry for ESMFold2.""" + +from __future__ import annotations + +import io +import itertools +import random +import re +import warnings +from collections.abc import Iterable, Sequence +from dataclasses import asdict, dataclass, replace +from functools import cached_property +from pathlib import Path +from subprocess import check_output +from tempfile import TemporaryDirectory +from typing import Any + +import biotite.structure as bs +import brotli +import msgpack +import msgpack_numpy +import numpy as np +import torch +from biotite.database import rcsb +from biotite.file import InvalidFileError +from biotite.structure.io.pdb import PDBFile +from biotite.structure.io.pdbx import CIFCategory, CIFColumn, CIFData, CIFFile +from biotite.structure.io.pdbx import set_structure as set_structure_pdbx +from biotite.structure.io.pdbx.convert import _get_transformations, get_structure +from biotite.structure.util import matrix_rotate +from scipy.spatial import KDTree + +from . import esmfold2_residue_constants as residue_constants +from .esmfold2_affine3d import Affine3D +from .esmfold2_aligner import Aligner +from .esmfold2_atom_indexer import AtomIndexer +from .esmfold2_metrics import compute_gdt_ts, compute_lddt_ca +from .esmfold2_misc import slice_python_object_as_numpy +from .esmfold2_mmcif_parsing import ( + MmcifWrapper, + NoProteinError, + round_mmcif_columns, +) +from .esmfold2_protein_chain import ( + ProteinChain, + _str_key_to_int_key, + chain_to_ndarray, + index_by_atom_name, + infer_cb, +) +from .esmfold2_utils_types import PathOrBuffer + +SINGLE_LETTER_CHAIN_IDS = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789" + + +def _parse_operation_expression(expression: str) -> list[tuple[str, ...]]: + """Expand an mmCIF operation expression in application order.""" + + def expand_group(group: str) -> list[str]: + operation_ids: list[str] = [] + for term in group.split(","): + if "-" not in term: + operation_ids.append(term) + continue + first, last = (int(value) for value in term.split("-")) + operation_ids.extend(str(value) for value in range(first, last + 1)) + return operation_ids + + groups = [group for group in expression.replace(")", "").split("(") if group] + groups.reverse() + return list(itertools.product(*(expand_group(group) for group in groups))) + + +def _apply_transformations_fast(chains, transformation_dict, operations): + """Return transformed copies of each affected protein chain.""" + transformed_chains = [] + for chain in chains: + for operation in operations: + coordinates = chain.atom37_positions.copy() + for op_step in operation: + transform = transformation_dict[op_step] + coordinates = matrix_rotate(coordinates, transform.rotation) + coordinates += transform.target_translation + transformed_chains.append(replace(chain, atom37_positions=coordinates)) + return transformed_chains + + +@dataclass +class ProteinComplexMetadata: + entity_lookup: dict[int, int | str] + chain_lookup: dict[int, str] + mmcif: MmcifWrapper | None = None + # This is a dictionary that maps assembly ids to the list of unique chains + # in that assembly. Allows for usage of `switch_assembly`. + assembly_composition: dict[str, list[str]] | None = None + + +@dataclass +class DockQSingleScore: + native_chains: tuple[str, str] + DockQ: float + interface_rms: float + ligand_rms: float + fnat: float + fnonnat: float + clashes: float + F1: float + DockQ_F1: float + + +@dataclass +class DockQResult: + total_dockq: float + native_interfaces: int + chain_mapping: dict[str, str] + interfaces: dict[tuple[str, str], DockQSingleScore] + # zip(aligned.chain_iter(), native.chain_iter()) gives you the pairing + # aligned.rmsd(native) should give you a low rmsd irrespective of shuffling + aligned: ProteinComplex + aligned_rmsd: float + + +@dataclass(frozen=True) +class ProteinComplex: + """Dataclass with atom37 representation of an entire protein complex.""" + + id: str + sequence: str + entity_id: np.ndarray # entities map to unique sequences + chain_id: np.ndarray # multiple chains might share an entity id + sym_id: np.ndarray # complexes might be copies of the same chain + residue_index: np.ndarray + insertion_code: np.ndarray + atom37_positions: np.ndarray + atom37_mask: np.ndarray + confidence: np.ndarray + # This metadata is parsed from the MMCIF file. For synthetic data, we do a best effort. + metadata: ProteinComplexMetadata + atom37_confidence: np.ndarray | None = None # P has shape (l, 37). + + # Coordinate completion, concatenation, and comparison + def infer_oxygen(self) -> ProteinComplex: + """Oxygen position is fixed given N, CA, C atoms. Infer it if not provided.""" + O_missing_indices = np.argwhere(~np.isfinite(self.atoms["O"]).all(axis=1)).squeeze() + + O_vector = torch.tensor([0.6240, -1.0613, 0.0103], dtype=torch.float32) + N, CA, C = torch.from_numpy(self.atoms[["N", "CA", "C"]]).float().unbind(dim=1) + N = torch.roll(N, -3) + N[..., -1, :] = torch.nan + + # Get the frame defined by the CA-C-N atom + frames = Affine3D.from_graham_schmidt(CA, C, N) + oxygen_coordinates = frames.apply(O_vector) + atom37_positions = self.atom37_positions.copy() + atom37_mask = self.atom37_mask.copy() + + atom37_positions[O_missing_indices, residue_constants.atom_order["O"]] = oxygen_coordinates[ + O_missing_indices + ].numpy() + atom37_mask[O_missing_indices, residue_constants.atom_order["O"]] = ~np.isnan( + atom37_positions[O_missing_indices, residue_constants.atom_order["O"]] + ).any(-1) + new_chain = replace(self, atom37_positions=atom37_positions, atom37_mask=atom37_mask) + return new_chain + + def infer_cbeta(self, infer_cbeta_for_glycine: bool = False) -> ProteinComplex: + """Return a new chain with inferred CB atoms at all residues except GLY. + + Args: + infer_cbeta_for_glycine (bool): If True, infers a beta carbon for glycine + residues, even though that residue doesn't have one. Default off. + + NOTE(rverkuil): The reason for having this switch in the first place + is that sometimes we want a (inferred) CB coordinate for every residue, + for example for making a pairwise distance matrix, or doing an RMSD + calculation between two designs for a given structural template, w/ + CB atoms. + """ + atom37_positions = self.atom37_positions.copy() + atom37_mask = self.atom37_mask.copy() + + N, CA, C = np.moveaxis(self.atoms[["N", "CA", "C"]], 1, 0) + # See usage in trDesign codebase. + # https://github.com/gjoni/trDesign/blob/f2d5930b472e77bfacc2f437b3966e7a708a8d37/02-GD/utils.py#L140 + inferred_cbeta_positions = infer_cb(C, N, CA, 1.522, 1.927, -2.143) + if not infer_cbeta_for_glycine: + inferred_cbeta_positions[np.array(list(self.sequence)) == "G", :] = np.nan + + atom37_positions[:, residue_constants.atom_order["CB"]] = inferred_cbeta_positions + atom37_mask[:, residue_constants.atom_order["CB"]] = ~np.isnan( + atom37_positions[:, residue_constants.atom_order["CB"]] + ).any(-1) + new_chain = replace(self, atom37_positions=atom37_positions, atom37_mask=atom37_mask) + return new_chain + + @classmethod + def from_open_source(cls, pc: ProteinComplex): + # TODO(@zeming): deprecated, should delete + return pc + + @classmethod + def concat(cls, objs: list[ProteinComplex]) -> ProteinComplex: + pdb_ids = [obj.id for obj in objs] + if len(set(pdb_ids)) > 1: + raise RuntimeError( + "Concatention of protein complexes across different PDB ids is unsupported" + ) + return ProteinComplex.from_chains( + list(itertools.chain.from_iterable(obj.chain_iter() for obj in objs)) + ) + + def _sanity_check_complexes_are_comparable(self, other: ProteinComplex): + if len(self) != len(other): + raise ValueError("Protein complexes must have the same length") + if len(list(self.chain_iter())) != len(list(other.chain_iter())): + raise ValueError("Protein complexes must have the same number of chains") + + def rmsd( + self, + target: ProteinComplex, + also_check_reflection: bool = False, + mobile_inds: list[int] | np.ndarray | None = None, + target_inds: list[int] | np.ndarray | None = None, + only_compute_backbone_rmsd: bool = False, + compute_chain_assignment: bool = True, + ): + """ + Compute the RMSD between this protein chain and another. + + Args: + target (ProteinComplex): The target (other) protein complex to compare to. + also_check_reflection: Compare the reflected mobile coordinates too. + mobile_inds: Mobile atom indices, not residue indices. + target_inds: Target atom indices, not residue indices. + only_compute_backbone_rmsd: Restrict the score to backbone atoms. + """ + aligned = self.dockq(target).aligned if compute_chain_assignment else self + + aligner = Aligner( + aligned if mobile_inds is None else aligned[mobile_inds], + target if target_inds is None else target[target_inds], + only_compute_backbone_rmsd, + ) + avg_rmsd = aligner.rmsd + + if not also_check_reflection: + return avg_rmsd + + aligner = Aligner( + aligned if mobile_inds is None else aligned[mobile_inds], + target if target_inds is None else target[target_inds], + only_compute_backbone_rmsd, + use_reflection=True, + ) + avg_rmsd_neg = aligner.rmsd + + return min(avg_rmsd, avg_rmsd_neg) + + def lddt_ca( + self, + target: ProteinComplex, + mobile_inds: list[int] | np.ndarray | None = None, + target_inds: list[int] | np.ndarray | None = None, + compute_chain_assignment: bool = True, + **kwargs, + ) -> float | np.ndarray: + """Compute the LDDT between this protein complex and another. + + Arguments: + target (ProteinComplex): The other protein complex to compare to. + mobile_inds: Mobile atom indices, not residue indices. + target_inds: Target atom indices, not residue indices. + + Returns: + float | np.ndarray: The LDDT score between the two protein chains, either + a single float or per-residue LDDT scores if `per_residue` is True. + """ + aligned = self.dockq(target).aligned if compute_chain_assignment else self + lddt = compute_lddt_ca( + torch.tensor(aligned.atom37_positions[mobile_inds]).unsqueeze(0), + torch.tensor(target.atom37_positions[target_inds]).unsqueeze(0), + torch.tensor(aligned.atom37_mask[mobile_inds]).unsqueeze(0), + **kwargs, + ) + return float(lddt) if lddt.numel() == 1 else lddt.numpy().flatten() + + def gdt_ts( + self, + target: ProteinComplex, + mobile_inds: list[int] | np.ndarray | None = None, + target_inds: list[int] | np.ndarray | None = None, + compute_chain_assignment: bool = True, + **kwargs, + ) -> float | np.ndarray: + """Compute the GDT_TS between this protein complex and another. + + Arguments: + target (ProteinComplex): The other protein complex to compare to. + mobile_inds: Mobile atom indices, not residue indices. + target_inds: Target atom indices, not residue indices. + + Returns: + float: The GDT_TS score between the two protein chains. + """ + aligned = self.dockq(target).aligned if compute_chain_assignment else self + gdt_ts = compute_gdt_ts( + mobile=torch.tensor( + index_by_atom_name(aligned.atom37_positions[mobile_inds], "CA"), + dtype=torch.float32, + ).unsqueeze(0), + target=torch.tensor( + index_by_atom_name(target.atom37_positions[target_inds], "CA"), + dtype=torch.float32, + ).unsqueeze(0), + atom_exists_mask=torch.tensor( + index_by_atom_name(aligned.atom37_mask[mobile_inds], "CA", dim=-1) + & index_by_atom_name(target.atom37_mask[target_inds], "CA", dim=-1) + ).unsqueeze(0), + **kwargs, + ) + return float(gdt_ts) if gdt_ts.numel() == 1 else gdt_ts.numpy().flatten() + + def dockq(self, native: ProteinComplex): + # This function uses dockqv2 to compute the DockQ score. Because it does a mapping + # over all possible chains, it's quite slow. Be careful not to use this in an inference loop + # or something that requires fast scoring. It defaults to 8 CPUs. + # + # TODO(@zeming): Because we haven't properly implemented protein complexes for mmcif, + # if your protein has multi-letter or repeated chain IDs, this will fail. Please call + # Normalize chain IDs before DockQ when IDs repeat or use multiple letters. + + try: + pass + except BaseException: + raise RuntimeError("DockQ is not installed. Please update your environment.") from None + self._sanity_check_complexes_are_comparable(native) + + def sanity_check_chain_ids(pc: ProteinComplex): + ids = [] + for i, chain in enumerate(pc.chain_iter()): + if i > len(SINGLE_LETTER_CHAIN_IDS): + raise ValueError("Too many chains to write to PDB file") + if len(chain.chain_id) > 1: + raise ValueError("We only supports single letter chain IDs for DockQ") + ids.append(chain.chain_id) + if len(set(ids)) != len(ids): + raise ValueError(f"Duplicate chain IDs in protein complex: {ids}") + return ids + + sanity_check_chain_ids(self) + sanity_check_chain_ids(native) + + with TemporaryDirectory() as tdir: + dir = Path(tdir) + self.to_pdb(dir / "self.pdb") + native.to_pdb(dir / "native.pdb") + + output = check_output(["DockQ", dir / "self.pdb", dir / "native.pdb"]) + lines = output.decode().split("\n") + + # Remove the header comments + start_index = next(i for i, line in enumerate(lines) if line.startswith("Model")) + lines = lines[start_index:] + + result = {} + interfaces = [] + current_interface: dict = {} + + for line in lines: + line = line.strip() + if not line: + continue + + if line.startswith(("Model :", "Native :")): + pass # Tmp pdb file location, it's useless... + elif line.startswith("Total DockQ"): + total_dockq_match = re.search( + r"Total DockQ over (\d+) native interfaces: ([\d.]+) with " + r"(.*) model:native mapping", + line, + ) + if total_dockq_match: + result["value"] = float(total_dockq_match.group(2)) + result["native interfaces"] = int(total_dockq_match.group(1)) + native_chains, self_chains = total_dockq_match.group(3).split(":") + result["mapping"] = dict(zip(native_chains, self_chains, strict=False)) + else: + raise RuntimeError( + "Failed to parse DockQ output, maybe your DockQ version is wrong?" + ) + elif line.startswith("Native chains:"): + if current_interface: + interfaces.append(current_interface) + current_interface = {"Native chains": line.split(":")[1].strip().split(", ")} + elif line.startswith("Model chains:"): + current_interface["Model chains"] = line.split(":")[1].strip().split(", ") + elif ":" in line: + key, value = line.split(":", 1) + current_interface[key.strip()] = float(value.strip()) + + if current_interface: + interfaces.append(current_interface) + + def parse_dict(d: dict[str, Any]) -> DockQSingleScore: + return DockQSingleScore( + native_chains=tuple(d["Native chains"]), # type: ignore + DockQ=float(d["DockQ"]), + interface_rms=float(d["irms"]), + ligand_rms=float(d["Lrms"]), # Note the capitalization difference + fnat=float(d["fnat"]), + fnonnat=float(d["fnonnat"]), + clashes=float(d["clashes"]), + F1=float(d["F1"]), + DockQ_F1=float(d["DockQ_F1"]), + ) + + inv_mapping = {v: k for k, v in result["mapping"].items()} + + self_chain_map = {c.chain_id: c for c in self.chain_iter()} + realigned = [] + for chain in native.chain_iter(): + realigned.append(self_chain_map[inv_mapping[chain.chain_id]]) + + realigned = ProteinComplex.from_chains(realigned) + aligner = Aligner(realigned, native) + realigned = aligner.apply(realigned) + + result = DockQResult( + total_dockq=result["value"], + native_interfaces=result["native interfaces"], + chain_mapping=result["mapping"], + interfaces={ + (i["Model chains"][0], i["Model chains"][1]): parse_dict(i) for i in interfaces + }, + aligned=realigned, + aligned_rmsd=aligner.rmsd, + ) + + return result + + # Object invariants, slicing, and chain views + def __post_init__(self): + if not isinstance(self.sequence, str): + raise TypeError("sequence must be a string.") + sequence_length = len(self.sequence) + aligned = { + "atom37_positions": self.atom37_positions, + "atom37_mask": self.atom37_mask, + "residue_index": self.residue_index, + "insertion_code": self.insertion_code, + "confidence": self.confidence, + "entity_id": self.entity_id, + "chain_id": self.chain_id, + "sym_id": self.sym_id, + } + for name, values in aligned.items(): + if not isinstance(values, np.ndarray): + raise TypeError(f"{name} must be a NumPy array, got {type(values).__name__}.") + if values.ndim == 0 or values.shape[0] != sequence_length: + raise ValueError( + f"{name} shape {values.shape} does not align with " + f"sequence length {sequence_length}." + ) + if self.atom37_positions.shape != (sequence_length, 37, 3): + raise ValueError( + "atom37_positions must have shape " + f"({sequence_length}, 37, 3), got {self.atom37_positions.shape}." + ) + if self.atom37_mask.shape != (sequence_length, 37): + raise ValueError( + "atom37_mask must have shape " + f"({sequence_length}, 37), got {self.atom37_mask.shape}." + ) + if self.atom37_mask.dtype != bool: + raise TypeError(f"atom37_mask must have Boolean dtype, got {self.atom37_mask.dtype}.") + if not np.issubdtype(self.atom37_positions.dtype, np.number): + raise TypeError("atom37_positions must use a numeric dtype.") + for name, values in ( + ("residue_index", self.residue_index), + ("insertion_code", self.insertion_code), + ("confidence", self.confidence), + ("entity_id", self.entity_id), + ("chain_id", self.chain_id), + ("sym_id", self.sym_id), + ): + if values.shape != (sequence_length,): + raise ValueError( + f"{name} must have shape ({sequence_length},), got {values.shape}." + ) + if not np.issubdtype(self.confidence.dtype, np.number): + raise TypeError("confidence must use a numeric dtype.") + atom37_confidence = self.atom37_confidence + if atom37_confidence is not None and not isinstance(atom37_confidence, np.ndarray): + raise TypeError("atom37_confidence must be a NumPy array when provided.") + if ( + isinstance(atom37_confidence, np.ndarray) + and atom37_confidence.shape != self.atom37_mask.shape + ): + raise ValueError( + "atom37_confidence shape must match atom37_mask: " + f"{atom37_confidence.shape} != {self.atom37_mask.shape}." + ) + + def __getitem__(self, idx: int | list[int] | slice | np.ndarray): + """This function slices protein complexes without consideration of chain breaks + NOTE: When slicing with a boolean mask, it's possible that the output array won't + be the expected length. This is because we do our best to preserve chainbreak tokens. + """ + + if isinstance(idx, int): + idx = [idx] + if isinstance(idx, list): + raise ValueError("ProteinComplex doesn't supports indexing with lists of indices") + + if isinstance(idx, np.ndarray): + is_chainbreak = np.asarray([s == "|" for s in self.sequence]) + idx = idx.astype(bool) | is_chainbreak + + complex = self._unsafe_slice(idx) + if len(complex) == 0: + return complex + + # detect runs of chainbreaks by searching for instances of '||' in complex.sequence + chainbreak_runs = np.asarray( + [complex.sequence[i : i + 2] == "||" for i in range(len(complex.sequence) - 1)] + + [complex.sequence[-1] == "|"] + ) + # We should remove as many chainbreaks as possible from the start of the sequence + for i in range(len(chainbreak_runs)): + if complex.sequence[i] == "|": + chainbreak_runs[i] = True + else: + break + complex = complex._unsafe_slice(~chainbreak_runs) + return complex + + def _unsafe_slice(self, idx: int | list[int] | slice | np.ndarray): + sequence = slice_python_object_as_numpy(self.sequence, idx) + return replace( + self, + sequence=sequence, + entity_id=self.entity_id[..., idx], + chain_id=self.chain_id[..., idx], + sym_id=self.sym_id[..., idx], + residue_index=self.residue_index[..., idx], + insertion_code=self.insertion_code[..., idx], + atom37_positions=self.atom37_positions[..., idx, :, :], + atom37_mask=self.atom37_mask[..., idx, :], + confidence=self.confidence[..., idx], + atom37_confidence=self.atom37_confidence[..., idx, :] + if self.atom37_confidence is not None + else None, + ) + + def __len__(self): + return len(self.sequence) + + @property + def num_chains(self): + return len(self.chain_boundaries) + + @cached_property + def atoms(self) -> AtomIndexer: + return AtomIndexer(self, property="atom37_positions", dim=-2) + + @cached_property + def atom_mask(self) -> AtomIndexer: + return AtomIndexer(self, property="atom37_mask", dim=-1) + + @cached_property + def chain_lengths(self) -> np.ndarray: + return np.diff(self.chain_boundaries, axis=1).flatten() + + @cached_property + def chain_boundaries(self) -> list[tuple[int, int]]: + cb = [-1] + for i, s in enumerate(self.sequence): + if s == "|": + cb.append(i) + cb.append(len(self)) + return [(cb[i] + 1, cb[i + 1]) for i in range(len(cb) - 1)] + + def get_chain_by_index(self, index: int) -> ProteinChain: + try: + start, end = self.chain_boundaries[index] + return self[start:end].as_chain() + except IndexError: + raise IndexError(f"Chain index {index} out of bounds") from None + + def get_chain_by_id( + self, chain_id: str, sample_chain_if_duplicate: bool = True + ) -> ProteinChain: + valid_indices = [ + index + for index, id_of_index in self.metadata.chain_lookup.items() + if id_of_index == chain_id + ] + if not valid_indices: + raise KeyError(f"Chain ID {chain_id} not found") + if sample_chain_if_duplicate: + index_to_return = random.choice(valid_indices) + return self.get_chain_by_index(index_to_return) + else: + if len(valid_indices) > 1: + raise ValueError(f"Multiple chains with chain ID {chain_id} found") + return self.get_chain_by_index(valid_indices[0]) + + def chain_iter(self) -> Iterable[ProteinChain]: + for start, end in self.chain_boundaries: + c = self[start:end] + yield c.as_chain() + + def as_chain(self, force_conversion: bool = False) -> ProteinChain: + """Convert the ProteinComplex to a ProteinChain. + + Args: + force_conversion: Flatten multiple chains to access chain-only utilities. + + """ + if not force_conversion: + if len(np.unique(self.chain_id)) != 1: + raise ValueError( + f"Protein complex {self.id!r} has multiple chains; " + "pass force_conversion=True to flatten it." + ) + if len(np.unique(self.entity_id)) != 1: + raise ValueError( + f"Protein complex {self.id!r} has multiple entities; " + "pass force_conversion=True to flatten it." + ) + if self.chain_id[0] not in self.metadata.chain_lookup: + warnings.warn( + "Chain ID not found in metadata, using 'A' as default", + stacklevel=2, + ) + if self.entity_id[0] not in self.metadata.entity_lookup: + warnings.warn( + "Entity ID not found in metadata, using None as default", + stacklevel=2, + ) + chain_id = self.metadata.chain_lookup.get(self.chain_id[0], "A") + entity_id = self.metadata.entity_lookup.get(self.entity_id[0], None) + else: + chain_id = "A" + entity_id = None + + return ProteinChain( + id=self.id, + sequence=self.sequence, + chain_id=chain_id, + entity_id=entity_id, + atom37_positions=self.atom37_positions, + atom37_mask=self.atom37_mask, + residue_index=self.residue_index, + insertion_code=self.insertion_code, + confidence=self.confidence, + mmcif=self.metadata.mmcif, + atom37_confidence=self.atom37_confidence, + ) + + # Contact topology and mmCIF export + @cached_property + def per_chain_kd_trees(self): + # Iterate over chains, build KDTree for each chain + kdtrees = [] + + CA = self.atoms["CA"] + + for start, end in self.chain_boundaries: + chain_CA = CA[start:end] + chain_CA = chain_CA[np.isfinite(chain_CA).all(axis=-1)] + kdtrees.append(KDTree(chain_CA)) + + return kdtrees + + def chain_adjacency(self, cutoff: float = 8.0) -> np.ndarray: + # Compute adjacency matrix for protein complex + num_chains = self.num_chains + adjacency = np.zeros((num_chains, num_chains), dtype=bool) + for (i, kdtree), (j, kdtree2) in itertools.combinations( + enumerate(self.per_chain_kd_trees), 2 + ): + adj = kdtree.query_ball_tree(kdtree2, cutoff) + any_is_adjacent = any(len(a) > 0 for a in adj) + adjacency[i, j] = any_is_adjacent + adjacency[j, i] = any_is_adjacent + return adjacency + + def chain_adjacency_by_index(self, index: int, cutoff: float = 8.0) -> np.ndarray: + num_chains = len(self.chain_boundaries) + adjacency = np.zeros(num_chains, dtype=bool) + for i, kdtree in enumerate(self.per_chain_kd_trees): + if i == index: + continue + adj = kdtree.query_ball_tree(self.per_chain_kd_trees[index], cutoff) + adjacency[i] = any(len(a) > 0 for a in adj) + return adjacency + + def add_prefix_to_chain_ids(self, prefix: str) -> ProteinComplex: + """Rename all chains in the complex with a given prefix. + + Args: + prefix (str): The prefix to use for the new chain IDs. Each chain will be + named as "{prefix}_{chain_id}". + + Returns: + ProteinComplex: A new protein complex with renamed chains. + """ + new_chains = [] + for chain in self.chain_iter(): + # Create new chain with updated chain_id + new_chain = replace(chain, chain_id=f"{prefix}_{chain.chain_id}") + new_chains.append(new_chain) + return ProteinComplex.from_chains(new_chains) + + def sasa(self, by_residue: bool = True): + chain = self.as_chain(force_conversion=True) + return chain.sasa(by_residue=by_residue) + + def to_mmcif_string(self) -> str: + """Convert the ProteinComplex to mmCIF format. + + Returns: + str: The mmCIF content as a string. + """ + # Convert the ProteinComplex to a biotite AtomArray + # Collect all atoms from all chains + all_atoms = [] + for chain in self.chain_iter(): + chain_atom_array = chain.atom_array + # Convert AtomArray to list of atoms and add to collection + all_atoms.extend(chain_atom_array) + + # Create combined AtomArray from all atoms + if not all_atoms: + raise ValueError("No atoms found in protein complex") + + atom_array = bs.array(all_atoms) + + # Create CIF file + f = CIFFile() + set_structure_pdbx(f, atom_array, data_block=self.id) + + # Add entity information for proper mmCIF structure + self._add_entity_information(f) + round_mmcif_columns(f) + + # Write to string + output = io.StringIO() + f.write(output) + return output.getvalue() + + def _add_entity_information(self, cif_file: CIFFile) -> None: + """Add entity, entity_poly, and struct_asym sections to CIF file.""" + + # Group chains by sequence to create unique entities + entity_map = {} # sequence -> entity_id + chain_to_entity = {} # chain_id -> entity_id + entity_sequences = {} # entity_id -> sequence + entity_id_counter = 1 + + for chain in self.chain_iter(): + sequence = chain.sequence + if sequence not in entity_map: + entity_map[sequence] = entity_id_counter + entity_sequences[entity_id_counter] = sequence + entity_id_counter += 1 + chain_to_entity[chain.chain_id] = entity_map[sequence] + + # Create _entity section + entity_ids = [] + entity_types = [] + entity_descriptions = [] + + for entity_id in sorted(entity_sequences.keys()): + entity_ids.append(str(entity_id)) + entity_types.append("polymer") + entity_descriptions.append(f"Protein chain (entity {entity_id})") + + cif_file.block["entity"] = CIFCategory( + name="entity", + columns={ + "id": CIFColumn(data=CIFData(array=np.array(entity_ids), dtype=np.str_)), + "type": CIFColumn(data=CIFData(array=np.array(entity_types), dtype=np.str_)), + "pdbx_description": CIFColumn( + data=CIFData(array=np.array(entity_descriptions), dtype=np.str_) + ), + }, + ) + + # Create _entity_poly section + poly_entity_ids = [] + poly_types = [] + poly_nstd_linkages = [] + poly_sequences = [] + + for entity_id in sorted(entity_sequences.keys()): + poly_entity_ids.append(str(entity_id)) + poly_types.append("polypeptide(L)") + poly_nstd_linkages.append("no") + poly_sequences.append(entity_sequences[entity_id]) + + cif_file.block["entity_poly"] = CIFCategory( + name="entity_poly", + columns={ + "entity_id": CIFColumn( + data=CIFData(array=np.array(poly_entity_ids), dtype=np.str_) + ), + "type": CIFColumn(data=CIFData(array=np.array(poly_types), dtype=np.str_)), + "nstd_linkage": CIFColumn( + data=CIFData(array=np.array(poly_nstd_linkages), dtype=np.str_) + ), + "pdbx_seq_one_letter_code": CIFColumn( + data=CIFData(array=np.array(poly_sequences), dtype=np.str_) + ), + }, + ) + + # Create _struct_asym section + asym_ids = [] + asym_entity_ids = [] + asym_details = [] + + for chain in self.chain_iter(): + asym_ids.append(chain.chain_id) + asym_entity_ids.append(str(chain_to_entity[chain.chain_id])) + asym_details.append("") + + cif_file.block["struct_asym"] = CIFCategory( + name="struct_asym", + columns={ + "id": CIFColumn(data=CIFData(array=np.array(asym_ids), dtype=np.str_)), + "entity_id": CIFColumn( + data=CIFData(array=np.array(asym_entity_ids), dtype=np.str_) + ), + "details": CIFColumn(data=CIFData(array=np.array(asym_details), dtype=np.str_)), + }, + ) + + # Construction, PDB interchange, and compact storage + @classmethod + def from_pdb( + cls, path: PathOrBuffer, id: str | None = None, is_predicted: bool = False + ) -> ProteinComplex: + atom_array = PDBFile.read(path).get_structure(model=1, extra_fields=["b_factor"]) + + chains = [] + for chain in bs.chain_iter(atom_array): + chain = chain[~chain.hetero] + if len(chain) == 0: + continue + chains.append(ProteinChain.from_atomarray(chain, id, is_predicted)) + return ProteinComplex.from_chains(chains) + + def to_pdb(self, path: PathOrBuffer, include_insertions: bool = True): + atom_array = None + for chain in self.chain_iter(): + carr = chain.atom_array if include_insertions else chain.atom_array_no_insertions + atom_array = carr if atom_array is None else atom_array + carr + f = PDBFile() + f.set_structure(atom_array) + f.write(path) + + def to_pdb_string(self, include_insertions: bool = True) -> str: + buf = io.StringIO() + self.to_pdb(buf, include_insertions=include_insertions) + buf.seek(0) + return buf.read() + + def normalize_chain_ids_for_pdb(self): + # Since PDB files have 1-letter chain IDs and don't support the idea of a symmetric index, + # we can normalize it instead which might be necessary for DockQ and to_pdb. + ids = SINGLE_LETTER_CHAIN_IDS + chains = [] + for i, chain in enumerate(self.chain_iter()): + chain = replace(chain, chain_id=ids[i]) + if i > len(ids): + raise RuntimeError("Too many chains to write to PDB file") + chains.append(chain) + + return ProteinComplex.from_chains(chains) + + def find_assembly_ids_with_chain(self, id: str) -> list[str]: + good_chains = [] + if (comp := self.metadata.assembly_composition) is not None: + for assembly_id, chain_ids in comp.items(): + if id in chain_ids: + good_chains.append(assembly_id) + else: + raise ValueError( + "Cannot switch assemblies on this ProteinComplex; construct it from " + "mmCIF to retain assembly metadata" + ) + return good_chains + + def switch_assembly(self, id: str): + if self.metadata.mmcif is None: + raise ValueError( + "Cannot switch assemblies without retained mmCIF source metadata." + ) + return get_assembly_fast(self.metadata.mmcif, assembly_id=id) + + def state_dict(self, backbone_only=False, json_serializable=False): + """This state dict is optimized for storage, so it turns things to fp16 whenever + possible. Note that we also only support int32 residue indices, I'm hoping we don't + need more than 2**32 residues...""" + dct = {k: v for k, v in vars(self).items()} + if backbone_only: + # Frozen dataclasses do not make their NumPy members immutable. Work on a + # private mask so requesting a compact backbone payload cannot clear the + # caller's side-chain atoms in-place. + atom37_mask = dct["atom37_mask"].copy() + atom37_mask[:, 3:] = False + dct["atom37_mask"] = atom37_mask + dct["atom37_positions"] = dct["atom37_positions"][dct["atom37_mask"]] + if dct.get("atom37_confidence") is not None: + dct["atom37_confidence"] = dct["atom37_confidence"][dct["atom37_mask"]] + else: + dct.pop("atom37_confidence", None) + for k, v in dct.items(): + if isinstance(v, np.ndarray): + match v.dtype: + case np.int64: + dct[k] = v.astype(np.int32) + case np.float64 | np.float32: + dct[k] = v.astype(np.float16) + case _: + pass + if json_serializable: + dct[k] = v.tolist() + elif isinstance(v, ProteinComplexMetadata): + dct[k] = asdict(v) + dct["metadata"]["mmcif"] = None + # These can be populated with non-serializable objects and are not needed for reconstruction + dct.pop("atoms", None) + dct.pop("atom_mask", None) + dct.pop("per_chain_kd_trees", None) + return dct + + def to_blob(self, backbone_only=False) -> bytes: + payload = msgpack.dumps(self.state_dict(backbone_only), default=msgpack_numpy.encode) + return brotli.compress(payload, quality=5) + + @classmethod + def from_state_dict(cls, dct): + # Note: assembly_composition is *supposed* to have string keys. + dct = _str_key_to_int_key(dct, ignore_keys=["assembly_composition"]) + + for k, v in dct.items(): + if isinstance(v, list): + dct[k] = np.array(v) + + atom37 = np.full((*dct["atom37_mask"].shape, 3), np.nan) + atom37[dct["atom37_mask"]] = dct["atom37_positions"] + dct["atom37_positions"] = atom37 + if "atom37_confidence" in dct: + atom37_conf = np.full(dct["atom37_mask"].shape, np.nan, dtype=np.float32) + atom37_conf[dct["atom37_mask"]] = dct["atom37_confidence"] + dct["atom37_confidence"] = atom37_conf + dct = { + k: ( + v.astype(np.float32) + if k in ["atom37_positions", "confidence", "atom37_confidence"] + else v + ) + for k, v in dct.items() + } + if "chain_boundaries" in dct: + del dct["chain_boundaries"] + if "chain_boundaries" in dct["metadata"]: + del dct["metadata"]["chain_boundaries"] + dct["metadata"] = ProteinComplexMetadata(**dct["metadata"]) + return cls(**dct) + + @classmethod + def from_blob(cls, input: Path | str | io.BytesIO | bytes): + """NOTE(@zlin): blob + sparse coding + brotli + fp16 reduces memory + of chains from 52G/1M chains to 20G/1M chains, I think this is a good first + shot at compressing and dumping chains to disk. I'm sure there's better ways.""" + match input: + case Path() | str(): + bytes = Path(input).read_bytes() + case io.BytesIO(): + bytes = input.getvalue() + case _: + bytes = input + state = msgpack.loads( + brotli.decompress(bytes), + object_hook=msgpack_numpy.decode, + strict_map_key=False, + ) + return cls.from_state_dict(state) + + @classmethod + def from_rcsb(cls, pdb_id: str, keep_source: bool = False) -> ProteinComplex: + f: io.StringIO = rcsb.fetch(pdb_id, "cif") # type: ignore + return cls.from_mmcif(f, id=pdb_id, keep_source=keep_source, is_predicted=False) + + @classmethod + def from_mmcif( + cls, + path: PathOrBuffer, + id: str | None = None, + assembly_id: str | None = None, + is_predicted: bool = False, + keep_source: bool = False, + ): + """Return a ProteinComplex object from an mmcif file. + TODO(@zeming): there's actually multiple complexes per file, but for ease of implementation, + we only consider the first defined complex! + + Args: + path: Uncompressed mmCIF path or text buffer. + id: Optional structure identifier. + is_predicted (bool): If True, reads b factor as the confidence readout. Default: False. + chain_id (str, optional): Select a chain corresponding to (author) chain id. + """ + mmcif = MmcifWrapper.read(path, id) + return get_assembly_fast(mmcif, assembly_id=assembly_id) + + @classmethod + def from_chains( + cls, + chains: Sequence[ProteinChain], + mmcif: MmcifWrapper | None = None, + all_assembly_metadata_dictionary: dict[str, list[str]] | None = None, + ): + if not chains: + raise ValueError("Cannot create a ProteinComplex from an empty list of chains") + + # TODO(roshan): Make a proper protein complex class + def join_arrays(arrays: Sequence[np.ndarray], sep: np.ndarray): + full_array = [] + for array in arrays: + full_array.append(array) + full_array.append(sep) + full_array = full_array[:-1] + return np.concatenate(full_array, 0) + + sep_tokens = { + "residue_index": np.array([-1]), + "insertion_code": np.array([""]), + "atom37_positions": np.full([1, 37, 3], np.nan), + "atom37_mask": np.zeros([1, 37], dtype=bool), + "confidence": np.array([0]), + } + + any_has_atom37_conf = any(c.atom37_confidence is not None for c in chains) + if any_has_atom37_conf: + sep_tokens["atom37_confidence"] = np.full([1, 37], np.nan, dtype=np.float32) + + def _get_chain_attr(chain: ProteinChain, name: str) -> np.ndarray: + val = getattr(chain, name) + if val is None and name == "atom37_confidence": + return np.full([len(chain), 37], np.nan, dtype=np.float32) + return val + + array_args: dict[str, np.ndarray] = { + name: join_arrays([_get_chain_attr(chain, name) for chain in chains], sep) + for name, sep in sep_tokens.items() + } + + multimer_arrays = [] + chain2num_max = -1 + chain2num = {} + ent2num_max = -1 + ent2num = {} + total_index = 0 + for i, c in enumerate(chains): + num_res = c.residue_index.shape[0] + if c.chain_id not in chain2num: + chain2num[c.chain_id] = (chain2num_max := chain2num_max + 1) + chain_id_array = np.full([num_res], chain2num[c.chain_id], dtype=np.int64) + + if c.entity_id is None: + entity_num = (ent2num_max := ent2num_max + 1) + else: + if c.entity_id not in ent2num: + ent2num[c.entity_id] = (ent2num_max := ent2num_max + 1) + entity_num = ent2num[c.entity_id] + entity_id_array = np.full([num_res], entity_num, dtype=np.int64) + + sym_id_array = np.full([num_res], i, dtype=np.int64) + + multimer_arrays.append( + { + "chain_id": chain_id_array, + "entity_id": entity_id_array, + "sym_id": sym_id_array, + } + ) + + total_index += num_res + 1 + + sep = np.array([-1]) + update = { + name: join_arrays([dct[name] for dct in multimer_arrays], sep=sep) + for name in ["chain_id", "entity_id", "sym_id"] + } + array_args.update(update) + + metadata = ProteinComplexMetadata( + mmcif=mmcif, + chain_lookup={v: k for k, v in chain2num.items()}, + entity_lookup={v: k for k, v in ent2num.items()}, + assembly_composition=all_assembly_metadata_dictionary, + ) + + return cls( + id=chains[0].id, + sequence=residue_constants.CHAIN_BREAK_TOKEN.join(chain.sequence for chain in chains), + metadata=metadata, + **array_args, + ) + + +# Biological-assembly expansion +def get_assembly_fast( + mmcif: MmcifWrapper, + assembly_id=None, + model=None, + data_block=None, + altloc="first", + use_author_fields=True, +): + pdbx_file = mmcif.raw + if pdbx_file is None: + raise InvalidFileError("No mmCIF data loaded") + assembly_gen_category = pdbx_file.block["pdbx_struct_assembly_gen"] + if assembly_gen_category is None: + raise InvalidFileError("File has no 'pdbx_struct_assembly_gen' category") + + struct_oper_category = pdbx_file.block["pdbx_struct_oper_list"] + if struct_oper_category is None: + raise InvalidFileError("File has no 'pdbx_struct_oper_list' category") + + if assembly_id is None: + assembly_id = assembly_gen_category["assembly_id"].data.array[0] + elif assembly_id not in assembly_gen_category["assembly_id"].data.array: + raise KeyError(f"File has no Assembly ID '{assembly_id}'") + + ### Calculate all possible transformations + transformations = _get_transformations(struct_oper_category) + + ### Get structure according to additional parameters + structure = get_structure( + pdbx_file, model, data_block, altloc, ["label_asym_id"], use_author_fields + )[0] # type: ignore + # TODO(@zeming) This line will remove all non-protein structural elements, + # we should remove this when we want to parse these too. + structure: bs.AtomArray = structure[ + bs.filter_amino_acids(structure) & ~structure.hetero # type: ignore + ] + if len(structure) == 0: + raise NoProteinError + unique_asym_ids = np.unique(structure.label_asym_id) # type: ignore + asym2chain = {} + asym2auth = {} + for asym_id in unique_asym_ids: + sub_structure: bs.AtomArray = structure[structure.label_asym_id == asym_id] # type: ignore + chain_id: str = sub_structure[0].chain_id # type: ignore + ( + sequence, + atom_positions, + atom_mask, + residue_index, + insertion_code, + confidence, + entity_id, + ) = chain_to_ndarray(sub_structure, mmcif, chain_id, False) + + asym2chain[asym_id] = ProteinChain( + id=mmcif.id or "unknown", + sequence=sequence, + chain_id=chain_id, + entity_id=entity_id, + atom37_positions=atom_positions, + atom37_mask=atom_mask, + residue_index=residue_index, + insertion_code=insertion_code, + confidence=confidence, + mmcif=None, + ) + asym2auth[asym_id] = chain_id + + ### Get transformations and apply them to the affected asym IDs + assembly = [] + assembly_id_dict: dict[str, list[str]] = {} + + # Process the target assembly ID + for aid, op_expr, asym_id_expr in zip( + assembly_gen_category["assembly_id"].data.array, + assembly_gen_category["oper_expression"].data.array, + assembly_gen_category["asym_id_list"].data.array, + strict=False, + ): + if aid == assembly_id: + # Parse operations and asym IDs for this specific entry + operations = _parse_operation_expression(op_expr) + asym_ids = asym_id_expr.split(",") + + # Filter affected asym IDs to only protein chains, preserving order + sub_structures = [asym2chain[asym_id] for asym_id in asym_ids if asym_id in asym2chain] + + # Apply transformations + sub_assembly = _apply_transformations_fast(sub_structures, transformations, operations) + assembly.extend(sub_assembly) + + # Build assembly_id_dict for this entry + assembly_id_dict[aid] = assembly_id_dict.get(aid, []) + [ + asym2auth[id_] for id_ in asym_ids if id_ in asym2auth + ] + + if len(assembly) == 0: + raise NoProteinError + return ProteinComplex.from_chains(assembly, mmcif, assembly_id_dict) + + +def protein_chain_to_protein_complex(chain: ProteinChain) -> ProteinComplex: + if "|" not in chain.sequence: + return ProteinComplex.from_chains([chain]) + chain_breaks = np.array(list(chain.sequence)) == "|" + chain_break_inds = np.where(chain_breaks)[0] + chain_break_inds = np.concatenate([[0], chain_break_inds, [len(chain)]]) + chain_break_inds = np.array(list(itertools.pairwise(chain_break_inds))) + complex_chains = [] + for start, end in chain_break_inds: + if start != 0: + start += 1 + complex_chains.append(chain[start:end]) + complex_chains = [ + ProteinChain.from_atom37( + chain.atom37_positions, + sequence=chain.sequence, + chain_id=SINGLE_LETTER_CHAIN_IDS[i], + entity_id=i, + ) + for i, chain in enumerate(complex_chains) + ] + return ProteinComplex.from_chains(complex_chains) diff --git a/fastplms/models/esmfold2/esmfold2_protein_structure.py b/fastplms/models/esmfold2/esmfold2_protein_structure.py new file mode 100644 index 0000000000000000000000000000000000000000..1d41330753a1d2d82194f77004645d162bda3617 --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_protein_structure.py @@ -0,0 +1,251 @@ +"""Atom selection, rigid alignment, RMSD, and GDT-TS primitives.""" + +from __future__ import annotations + +from collections.abc import Callable +from typing import TypeVar + +import numpy as np +import torch +import torch.nn.functional as F +from torch import Tensor +from torch.amp import autocast # type: ignore + +from .esmfold2_affine3d import Affine3D +from .esmfold2_misc import unbinpack +from .esmfold2_normalize_coordinates import index_by_atom_name + +ArrayOrTensor = TypeVar("ArrayOrTensor", np.ndarray, Tensor) + + +def _coordinate_operations( + coordinates: ArrayOrTensor, +) -> tuple[Callable[[ArrayOrTensor], ArrayOrTensor], Callable[..., ArrayOrTensor]]: + if isinstance(coordinates, np.ndarray): + + def normalize(X: ArrayOrTensor) -> ArrayOrTensor: + return X / np.linalg.norm(X, axis=-1, keepdims=True) + + return normalize, np.cross + return F.normalize, torch.cross # type: ignore[return-value] + + +def infer_cbeta_from_atom37( + atom37: ArrayOrTensor, + bond_length: float = 1.522, + bond_angle: float = 1.927, + dihedral: float = -2.143, +) -> ArrayOrTensor: + """Infer C-beta coordinates from backbone tensor ``X``. + + The scalar keyword arguments encode the bond length, bond angle, and + dihedral in radians used by the checkpoint's training geometry. + """ + + n_position = index_by_atom_name(atom37, "N", dim=-2) + ca_position = index_by_atom_name(atom37, "CA", dim=-2) + c_position = index_by_atom_name(atom37, "C", dim=-2) + normalize, cross = _coordinate_operations(atom37) + with np.errstate(invalid="ignore"): + n_to_ca = n_position - ca_position + n_to_c = n_position - c_position + unit_n_to_ca = normalize(n_to_ca) + normal = normalize(cross(n_to_c, unit_n_to_ca)) + basis = [unit_n_to_ca, cross(normal, unit_n_to_ca), normal] + coefficients = [ + bond_length * np.cos(bond_angle), + bond_length * np.sin(bond_angle) * np.cos(dihedral), + -bond_length * np.sin(bond_angle) * np.sin(dihedral), + ] + offset = sum( + vector * coefficient for vector, coefficient in zip(basis, coefficients, strict=True) + ) + return ca_position + offset + + +def _unpack_alignment_inputs( + mobile: Tensor, + target: Tensor, + atom_mask: Tensor | None, + sequence_id: Tensor | None, +) -> tuple[Tensor, Tensor, Tensor | None]: + if sequence_id is None: + return mobile, target, atom_mask + unpacked_mobile = unbinpack(mobile, sequence_id, pad_value=torch.nan) + unpacked_target = unbinpack(target, sequence_id, pad_value=torch.nan) + if atom_mask is None: + unpacked_mask = torch.isfinite(unpacked_target).all(dim=-1) + else: + unpacked_mask = unbinpack(atom_mask, sequence_id, pad_value=0) + return unpacked_mobile, unpacked_target, unpacked_mask + + +def _flatten_atom_axes( + mobile: Tensor, + target: Tensor, + atom_mask: Tensor | None, +) -> tuple[Tensor, Tensor, Tensor | None]: + b = mobile.shape[0] + flat_mobile = mobile.view(b, -1, 3) if mobile.dim() == 4 else mobile + flat_target = target.view(b, -1, 3) if target.dim() == 4 else target + flat_mask = atom_mask + if flat_mask is not None and flat_mask.dim() == 3: + flat_mask = flat_mask.view(b, -1) + return flat_mobile, flat_target, flat_mask + + +def _masked_coordinates( + mobile: Tensor, + target: Tensor, + atom_mask: Tensor | None, +) -> tuple[Tensor, Tensor, Tensor]: + if atom_mask is None: + atom_mask = torch.ones( + mobile.shape[:2], + dtype=torch.bool, + device=mobile.device, + ) + return mobile, target, atom_mask + expanded_mask = atom_mask.unsqueeze(-1) + return ( + mobile.masked_fill(~expanded_mask, 0), + target.masked_fill(~expanded_mask, 0), + atom_mask, + ) + + +@torch.no_grad() +@autocast("cuda", enabled=False) +def compute_alignment_tensors( + mobile: Tensor, + target: Tensor, + atom_exists_mask: Tensor | None = None, + sequence_id: Tensor | None = None, +) -> tuple[Tensor, Tensor, Tensor, Tensor, Tensor, Tensor]: + """Center and align coordinate tensors ``X`` and ``Y``. + + Inputs have shape (b, n, 3), or (b, l, n_atoms, 3). The returned rotation + tensor ``R`` has shape (b, 3, 3), and atom counts have shape (b, 1). + """ + + mobile, target, atom_exists_mask = _unpack_alignment_inputs( + mobile, + target, + atom_exists_mask, + sequence_id, + ) + if mobile.shape != target.shape: + raise AssertionError("Batch structure shapes do not match!") + mobile, target, atom_exists_mask = _flatten_atom_axes( + mobile, + target, + atom_exists_mask, + ) + mobile, target, atom_exists_mask = _masked_coordinates( + mobile, + target, + atom_exists_mask, + ) + + num_valid_atoms = atom_exists_mask.sum(dim=-1, keepdim=True) + centroid_mobile = mobile.sum(dim=-2, keepdim=True) / num_valid_atoms.unsqueeze(-1) + centroid_target = target.sum(dim=-2, keepdim=True) / num_valid_atoms.unsqueeze(-1) + centroid_mobile[num_valid_atoms == 0] = 0 + centroid_target[num_valid_atoms == 0] = 0 + + expanded_mask = atom_exists_mask.unsqueeze(-1) + centered_mobile = (mobile - centroid_mobile).masked_fill(~expanded_mask, 0) + centered_target = (target - centroid_target).masked_fill(~expanded_mask, 0) + covariance = torch.matmul(centered_mobile.transpose(1, 2), centered_target) + left_vectors, _, right_vectors = torch.svd(covariance) + rotation = torch.matmul(left_vectors, right_vectors.transpose(1, 2)) + return ( + centered_mobile, + centroid_mobile, + centered_target, + centroid_target, + rotation, + num_valid_atoms, + ) + + +def _validate_reduction(reduction: str, allowed: tuple[str, ...]) -> None: + if reduction not in allowed: + raise ValueError("Unrecognized reduction: '{reduction}'") + + +@torch.no_grad() +@autocast("cuda", enabled=False) +def compute_rmsd_no_alignment( + aligned: Tensor, + target: Tensor, + num_valid_atoms: Tensor, + reduction: str = "batch", +) -> Tensor: + """Measure RMSD after alignment using a declared reduction.""" + + _validate_reduction(reduction, ("per_residue", "per_sample", "batch")) + difference = aligned - target + if reduction == "per_residue": + mean_squared_error = difference.square().view(difference.size(0), -1, 9).mean(-1) + else: + mean_squared_error = difference.square().sum(dim=(1, 2)) / num_valid_atoms.squeeze(-1) + rmsd = torch.sqrt(mean_squared_error) + if reduction in {"per_residue", "per_sample"}: + return rmsd + valid_samples = num_valid_atoms.squeeze(-1) > 0 + return rmsd.masked_fill(~valid_samples, 0).sum() / (valid_samples.sum() + 1e-8) + + +@torch.no_grad() +@autocast("cuda", enabled=False) +def compute_affine_and_rmsd( + mobile: Tensor, + target: Tensor, + atom_exists_mask: Tensor | None = None, + sequence_id: Tensor | None = None, +) -> tuple[Affine3D, Tensor]: + """Fit ``X`` onto ``Y`` and return the rigid transform and batch RMSD.""" + + ( + centered_mobile, + centroid_mobile, + centered_target, + centroid_target, + rotation, + num_valid_atoms, + ) = compute_alignment_tensors(mobile, target, atom_exists_mask, sequence_id) + translation = torch.matmul(-centroid_mobile, rotation) + centroid_target + affine = Affine3D.from_tensor_pair( + translation, + rotation.unsqueeze(dim=-3).transpose(-2, -1), + ) + rotated_mobile = torch.matmul(centered_mobile, rotation) + rmsd = compute_rmsd_no_alignment( + rotated_mobile, + centered_target, + num_valid_atoms, + reduction="batch", + ) + return affine, rmsd + + +def compute_gdt_ts_no_alignment( + aligned: Tensor, + target: Tensor, + atom_exists_mask: Tensor, + reduction: str = "batch", +) -> Tensor: + """Compute GDT-TS for already aligned coordinate tensors.""" + + _validate_reduction(reduction, ("per_sample", "batch")) + if atom_exists_mask is None: + atom_exists_mask = torch.isfinite(target).all(dim=-1) + deviation = torch.linalg.vector_norm(aligned - target, dim=-1) + counts = atom_exists_mask.sum(dim=-1) + score_1 = ((deviation < 1) * atom_exists_mask).sum(dim=-1) / counts + score_2 = ((deviation < 2) * atom_exists_mask).sum(dim=-1) / counts + score_4 = ((deviation < 4) * atom_exists_mask).sum(dim=-1) / counts + score_8 = ((deviation < 8) * atom_exists_mask).sum(dim=-1) / counts + score = (score_1 + score_2 + score_4 + score_8) * 0.25 + return score.mean() if reduction == "batch" else score diff --git a/fastplms/models/esmfold2/esmfold2_residue_constants.py b/fastplms/models/esmfold2/esmfold2_residue_constants.py new file mode 100644 index 0000000000000000000000000000000000000000..42655c58c794c2862036bc993a4970989afbbf05 --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_residue_constants.py @@ -0,0 +1,1017 @@ +# Copyright 2025 EvolutionaryScale +# Copyright 2021 AlQuraishi Laboratory +# Copyright 2021 DeepMind Technologies Limited +# +# 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. + +"""Canonical amino-acid geometry tables used by ESMFold2. + +The literal chemistry measurements are kept visible for review. Derived masks, +indices, frames, and atom mappings are built locally below and are checked +exactly against the pinned Biohub implementation. +""" + +from __future__ import annotations + +import functools +from collections import defaultdict, namedtuple +from collections.abc import Mapping +from pathlib import Path +from typing import Any, cast + +import numpy as np + +ca_ca = 3.80209737096 +chi_angles_atoms = { + "ALA": [], + "ARG": [ + ["N", "CA", "CB", "CG"], + ["CA", "CB", "CG", "CD"], + ["CB", "CG", "CD", "NE"], + ["CG", "CD", "NE", "CZ"], + ], + "ASN": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "OD1"]], + "ASP": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "OD1"]], + "CYS": [["N", "CA", "CB", "SG"]], + "GLN": [ + ["N", "CA", "CB", "CG"], + ["CA", "CB", "CG", "CD"], + ["CB", "CG", "CD", "OE1"], + ], + "GLU": [ + ["N", "CA", "CB", "CG"], + ["CA", "CB", "CG", "CD"], + ["CB", "CG", "CD", "OE1"], + ], + "GLY": [], + "HIS": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "ND1"]], + "ILE": [["N", "CA", "CB", "CG1"], ["CA", "CB", "CG1", "CD1"]], + "LEU": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]], + "LYS": [ + ["N", "CA", "CB", "CG"], + ["CA", "CB", "CG", "CD"], + ["CB", "CG", "CD", "CE"], + ["CG", "CD", "CE", "NZ"], + ], + "MET": [ + ["N", "CA", "CB", "CG"], + ["CA", "CB", "CG", "SD"], + ["CB", "CG", "SD", "CE"], + ], + "PHE": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]], + "PRO": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD"]], + "SER": [["N", "CA", "CB", "OG"]], + "THR": [["N", "CA", "CB", "OG1"]], + "TRP": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]], + "TYR": [["N", "CA", "CB", "CG"], ["CA", "CB", "CG", "CD1"]], + "VAL": [["N", "CA", "CB", "CG1"]], + "UNK": [], +} +chi_angles_mask = [ + [1.0] * len(groups) + [0.0] * (4 - len(groups)) for groups in chi_angles_atoms.values() +] +_PI_PERIODIC_CHI = {"ASP": {1}, "GLU": {2}, "PHE": {1}, "TYR": {1}} +chi_pi_periodic = [ + [1.0 if chi_index in _PI_PERIODIC_CHI.get(residue_name, ()) else 0.0 for chi_index in range(4)] + for residue_name in chi_angles_atoms +] +rigid_group_atom_positions: dict[str, list[list[Any]]] = { + "ALA": [ + ["N", 0, (-0.525, 1.363, 0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.526, -0.0, -0.0)], + ["CB", 0, (-0.529, -0.774, -1.205)], + ["O", 3, (0.627, 1.062, 0.0)], + ], + "ARG": [ + ["N", 0, (-0.524, 1.362, -0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.525, -0.0, -0.0)], + ["CB", 0, (-0.524, -0.778, -1.209)], + ["O", 3, (0.626, 1.062, 0.0)], + ["CG", 4, (0.616, 1.39, -0.0)], + ["CD", 5, (0.564, 1.414, 0.0)], + ["NE", 6, (0.539, 1.357, -0.0)], + ["NH1", 7, (0.206, 2.301, 0.0)], + ["NH2", 7, (2.078, 0.978, -0.0)], + ["CZ", 7, (0.758, 1.093, -0.0)], + ], + "ASN": [ + ["N", 0, (-0.536, 1.357, 0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.526, -0.0, -0.0)], + ["CB", 0, (-0.531, -0.787, -1.2)], + ["O", 3, (0.625, 1.062, 0.0)], + ["CG", 4, (0.584, 1.399, 0.0)], + ["ND2", 5, (0.593, -1.188, 0.001)], + ["OD1", 5, (0.633, 1.059, 0.0)], + ], + "ASP": [ + ["N", 0, (-0.525, 1.362, -0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.527, 0.0, -0.0)], + ["CB", 0, (-0.526, -0.778, -1.208)], + ["O", 3, (0.626, 1.062, -0.0)], + ["CG", 4, (0.593, 1.398, -0.0)], + ["OD1", 5, (0.61, 1.091, 0.0)], + ["OD2", 5, (0.592, -1.101, -0.003)], + ], + "CYS": [ + ["N", 0, (-0.522, 1.362, -0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.524, 0.0, 0.0)], + ["CB", 0, (-0.519, -0.773, -1.212)], + ["O", 3, (0.625, 1.062, -0.0)], + ["SG", 4, (0.728, 1.653, 0.0)], + ], + "GLN": [ + ["N", 0, (-0.526, 1.361, -0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.526, 0.0, 0.0)], + ["CB", 0, (-0.525, -0.779, -1.207)], + ["O", 3, (0.626, 1.062, -0.0)], + ["CG", 4, (0.615, 1.393, 0.0)], + ["CD", 5, (0.587, 1.399, -0.0)], + ["NE2", 6, (0.593, -1.189, -0.001)], + ["OE1", 6, (0.634, 1.06, 0.0)], + ], + "GLU": [ + ["N", 0, (-0.528, 1.361, 0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.526, -0.0, -0.0)], + ["CB", 0, (-0.526, -0.781, -1.207)], + ["O", 3, (0.626, 1.062, 0.0)], + ["CG", 4, (0.615, 1.392, 0.0)], + ["CD", 5, (0.6, 1.397, 0.0)], + ["OE1", 6, (0.607, 1.095, -0.0)], + ["OE2", 6, (0.589, -1.104, -0.001)], + ], + "GLY": [ + ["N", 0, (-0.572, 1.337, 0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.517, -0.0, -0.0)], + ["O", 3, (0.626, 1.062, -0.0)], + ], + "HIS": [ + ["N", 0, (-0.527, 1.36, 0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.525, 0.0, 0.0)], + ["CB", 0, (-0.525, -0.778, -1.208)], + ["O", 3, (0.625, 1.063, 0.0)], + ["CG", 4, (0.6, 1.37, -0.0)], + ["CD2", 5, (0.889, -1.021, 0.003)], + ["ND1", 5, (0.744, 1.16, -0.0)], + ["CE1", 5, (2.03, 0.851, 0.002)], + ["NE2", 5, (2.145, -0.466, 0.004)], + ], + "ILE": [ + ["N", 0, (-0.493, 1.373, -0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.527, -0.0, -0.0)], + ["CB", 0, (-0.536, -0.793, -1.213)], + ["O", 3, (0.627, 1.062, -0.0)], + ["CG1", 4, (0.534, 1.437, -0.0)], + ["CG2", 4, (0.54, -0.785, -1.199)], + ["CD1", 5, (0.619, 1.391, 0.0)], + ], + "LEU": [ + ["N", 0, (-0.52, 1.363, 0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.525, -0.0, -0.0)], + ["CB", 0, (-0.522, -0.773, -1.214)], + ["O", 3, (0.625, 1.063, -0.0)], + ["CG", 4, (0.678, 1.371, 0.0)], + ["CD1", 5, (0.53, 1.43, -0.0)], + ["CD2", 5, (0.535, -0.774, 1.2)], + ], + "LYS": [ + ["N", 0, (-0.526, 1.362, -0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.526, 0.0, 0.0)], + ["CB", 0, (-0.524, -0.778, -1.208)], + ["O", 3, (0.626, 1.062, -0.0)], + ["CG", 4, (0.619, 1.39, 0.0)], + ["CD", 5, (0.559, 1.417, 0.0)], + ["CE", 6, (0.56, 1.416, 0.0)], + ["NZ", 7, (0.554, 1.387, 0.0)], + ], + "MET": [ + ["N", 0, (-0.521, 1.364, -0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.525, 0.0, 0.0)], + ["CB", 0, (-0.523, -0.776, -1.21)], + ["O", 3, (0.625, 1.062, -0.0)], + ["CG", 4, (0.613, 1.391, -0.0)], + ["SD", 5, (0.703, 1.695, 0.0)], + ["CE", 6, (0.32, 1.786, -0.0)], + ], + "PHE": [ + ["N", 0, (-0.518, 1.363, 0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.524, 0.0, -0.0)], + ["CB", 0, (-0.525, -0.776, -1.212)], + ["O", 3, (0.626, 1.062, -0.0)], + ["CG", 4, (0.607, 1.377, 0.0)], + ["CD1", 5, (0.709, 1.195, -0.0)], + ["CD2", 5, (0.706, -1.196, 0.0)], + ["CE1", 5, (2.102, 1.198, -0.0)], + ["CE2", 5, (2.098, -1.201, -0.0)], + ["CZ", 5, (2.794, -0.003, -0.001)], + ], + "PRO": [ + ["N", 0, (-0.566, 1.351, -0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.527, -0.0, 0.0)], + ["CB", 0, (-0.546, -0.611, -1.293)], + ["O", 3, (0.621, 1.066, 0.0)], + ["CG", 4, (0.382, 1.445, 0.0)], + ["CD", 5, (0.477, 1.424, 0.0)], + ], + "SER": [ + ["N", 0, (-0.529, 1.36, -0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.525, -0.0, -0.0)], + ["CB", 0, (-0.518, -0.777, -1.211)], + ["O", 3, (0.626, 1.062, -0.0)], + ["OG", 4, (0.503, 1.325, 0.0)], + ], + "THR": [ + ["N", 0, (-0.517, 1.364, 0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.526, 0.0, -0.0)], + ["CB", 0, (-0.516, -0.793, -1.215)], + ["O", 3, (0.626, 1.062, 0.0)], + ["CG2", 4, (0.55, -0.718, -1.228)], + ["OG1", 4, (0.472, 1.353, 0.0)], + ], + "TRP": [ + ["N", 0, (-0.521, 1.363, 0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.525, -0.0, 0.0)], + ["CB", 0, (-0.523, -0.776, -1.212)], + ["O", 3, (0.627, 1.062, 0.0)], + ["CG", 4, (0.609, 1.37, -0.0)], + ["CD1", 5, (0.824, 1.091, 0.0)], + ["CD2", 5, (0.854, -1.148, -0.005)], + ["CE2", 5, (2.186, -0.678, -0.007)], + ["CE3", 5, (0.622, -2.53, -0.007)], + ["NE1", 5, (2.14, 0.69, -0.004)], + ["CH2", 5, (3.028, -2.89, -0.013)], + ["CZ2", 5, (3.283, -1.543, -0.011)], + ["CZ3", 5, (1.715, -3.389, -0.011)], + ], + "TYR": [ + ["N", 0, (-0.522, 1.362, 0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.524, -0.0, -0.0)], + ["CB", 0, (-0.522, -0.776, -1.213)], + ["O", 3, (0.627, 1.062, -0.0)], + ["CG", 4, (0.607, 1.382, -0.0)], + ["CD1", 5, (0.716, 1.195, -0.0)], + ["CD2", 5, (0.713, -1.194, -0.001)], + ["CE1", 5, (2.107, 1.2, -0.002)], + ["CE2", 5, (2.104, -1.201, -0.003)], + ["OH", 5, (4.168, -0.002, -0.005)], + ["CZ", 5, (2.791, -0.001, -0.003)], + ], + "VAL": [ + ["N", 0, (-0.494, 1.373, -0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.527, -0.0, -0.0)], + ["CB", 0, (-0.533, -0.795, -1.213)], + ["O", 3, (0.627, 1.062, -0.0)], + ["CG1", 4, (0.54, 1.429, -0.0)], + ["CG2", 4, (0.533, -0.776, 1.203)], + ], + "UNK": [ + ["N", 0, (-0.525, 1.363, 0.0)], + ["CA", 0, (0.0, 0.0, 0.0)], + ["C", 0, (1.526, -0.0, -0.0)], + ], +} +residue_atoms = { + "ALA": ["C", "CA", "CB", "N", "O"], + "ARG": ["C", "CA", "CB", "CG", "CD", "CZ", "N", "NE", "O", "NH1", "NH2"], + "ASP": ["C", "CA", "CB", "CG", "N", "O", "OD1", "OD2"], + "ASN": ["C", "CA", "CB", "CG", "N", "ND2", "O", "OD1"], + "CYS": ["C", "CA", "CB", "N", "O", "SG"], + "GLU": ["C", "CA", "CB", "CG", "CD", "N", "O", "OE1", "OE2"], + "GLN": ["C", "CA", "CB", "CG", "CD", "N", "NE2", "O", "OE1"], + "GLY": ["C", "CA", "N", "O"], + "HIS": ["C", "CA", "CB", "CG", "CD2", "CE1", "N", "ND1", "NE2", "O"], + "ILE": ["C", "CA", "CB", "CG1", "CG2", "CD1", "N", "O"], + "LEU": ["C", "CA", "CB", "CG", "CD1", "CD2", "N", "O"], + "LYS": ["C", "CA", "CB", "CG", "CD", "CE", "N", "NZ", "O"], + "MET": ["C", "CA", "CB", "CG", "CE", "N", "O", "SD"], + "PHE": ["C", "CA", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "N", "O"], + "PRO": ["C", "CA", "CB", "CG", "CD", "N", "O"], + "SER": ["C", "CA", "CB", "N", "O", "OG"], + "THR": ["C", "CA", "CB", "CG2", "N", "O", "OG1"], + "TRP": [ + "C", + "CA", + "CB", + "CG", + "CD1", + "CD2", + "CE2", + "CE3", + "CZ2", + "CZ3", + "CH2", + "N", + "NE1", + "O", + ], + "TYR": ["C", "CA", "CB", "CG", "CD1", "CD2", "CE1", "CE2", "CZ", "N", "O", "OH"], + "VAL": ["C", "CA", "CB", "CG1", "CG2", "N", "O"], + "UNK": ["C", "CA", "N"], +} +residue_atom_renaming_swaps = { + "ASP": {"OD1": "OD2"}, + "GLU": {"OE1": "OE2"}, + "PHE": {"CD1": "CD2", "CE1": "CE2"}, + "TYR": {"CD1": "CD2", "CE1": "CE2"}, +} +van_der_waals_radius = {"C": 1.7, "N": 1.55, "O": 1.52, "S": 1.8} +Bond = namedtuple("Bond", ["atom1_name", "atom2_name", "length", "stddev"]) +BondAngle = namedtuple( + "BondAngle", ["atom1_name", "atom2_name", "atom3name", "angle_rad", "stddev"] +) + +_STEREO_CHEMICAL_PROPS_PATH = Path("evolutionaryscale/structure/stereo_chemical_props.txt") + + +def _bond_key(atom_a: str, atom_b: str) -> tuple[str, str]: + """Return an order-independent key for a covalent bond.""" + + return (atom_a, atom_b) if atom_a <= atom_b else (atom_b, atom_a) + + +def _read_stereo_sections(text: str) -> tuple[list[str], list[str]]: + """Split the two tabular sections while discarding their headers.""" + + lines = iter(text.splitlines()) + next(lines) + bond_rows = list(iter(lambda: next(lines).strip(), "-")) + next(lines) + next(lines) + angle_rows = list(iter(lambda: next(lines).strip(), "-")) + return bond_rows, angle_rows + + +@functools.cache +def load_stereo_chemical_props() -> tuple[ + dict[str, list[Any]], dict[str, list[Any]], dict[str, list[Any]] +]: + """Load covalent geometry and derive virtual bonds for bond angles. + + The returned dictionaries are keyed by three-letter residue name. Virtual + bond lengths and uncertainties use the same operation order as the + checkpoint's reference feature pipeline so their floating-point values are + bitwise reproducible. + """ + + bond_rows, angle_rows = _read_stereo_sections(_STEREO_CHEMICAL_PROPS_PATH.read_text()) + residue_bonds: dict[str, list[Any]] = {} + for row in bond_rows: + atom_pair, residue_name, length, stddev = row.split() + atom_a, atom_b = atom_pair.split("-") + residue_bonds.setdefault(residue_name, []).append( + Bond(atom_a, atom_b, float(length), float(stddev)) + ) + residue_bonds["UNK"] = [] + + residue_bond_angles: dict[str, list[Any]] = {} + for row in angle_rows: + atom_triple, residue_name, angle_degrees, stddev_degrees = row.split() + atom_a, atom_b, atom_c = atom_triple.split("-") + residue_bond_angles.setdefault(residue_name, []).append( + BondAngle( + atom_a, + atom_b, + atom_c, + float(angle_degrees) / 180.0 * np.pi, + float(stddev_degrees) / 180.0 * np.pi, + ) + ) + residue_bond_angles["UNK"] = [] + + residue_virtual_bonds: dict[str, list[Any]] = {} + for residue_name, angles in residue_bond_angles.items(): + lookup = { + _bond_key(bond.atom1_name, bond.atom2_name): bond + for bond in residue_bonds[residue_name] + } + derived = residue_virtual_bonds.setdefault(residue_name, []) + for angle in angles: + left = lookup[_bond_key(angle.atom1_name, angle.atom2_name)] + right = lookup[_bond_key(angle.atom2_name, angle.atom3name)] + theta = angle.angle_rad + length = np.sqrt( + left.length**2 + right.length**2 - 2 * left.length * right.length * np.cos(theta) + ) + dl_outer = 0.5 / length + dl_dgamma = 2 * left.length * right.length * np.sin(theta) * dl_outer + dl_db1 = (2 * left.length - 2 * right.length * np.cos(theta)) * dl_outer + dl_db2 = (2 * right.length - 2 * left.length * np.cos(theta)) * dl_outer + stddev = np.sqrt( + (dl_dgamma * angle.stddev) ** 2 + + (dl_db1 * left.stddev) ** 2 + + (dl_db2 * right.stddev) ** 2 + ) + derived.append(Bond(angle.atom1_name, angle.atom3name, length, stddev)) + return residue_bonds, residue_virtual_bonds, residue_bond_angles + + +between_res_bond_length_c_n = [1.329, 1.341] +between_res_bond_length_stddev_c_n = [0.014, 0.016] +between_res_cos_angles_c_n_ca = [-0.5203, 0.0353] +between_res_cos_angles_ca_c_n = [-0.4473, 0.0311] +atom_types = [ + "N", + "CA", + "C", + "CB", + "O", + "CG", + "CG1", + "CG2", + "OG", + "OG1", + "SG", + "CD", + "CD1", + "CD2", + "ND1", + "ND2", + "OD1", + "OD2", + "SD", + "CE", + "CE1", + "CE2", + "CE3", + "NE", + "NE1", + "NE2", + "OE1", + "OE2", + "CH2", + "NH1", + "NH2", + "OH", + "CZ", + "CZ2", + "CZ3", + "NZ", + "OXT", +] +atom_order = {atom_type: i for (i, atom_type) in enumerate(atom_types)} +atom_type_num = len(atom_types) +restype_name_to_atom14_names = { + "ALA": ["N", "CA", "C", "O", "CB", "", "", "", "", "", "", "", "", ""], + "ARG": [ + "N", + "CA", + "C", + "O", + "CB", + "CG", + "CD", + "NE", + "CZ", + "NH1", + "NH2", + "", + "", + "", + ], + "ASN": ["N", "CA", "C", "O", "CB", "CG", "OD1", "ND2", "", "", "", "", "", ""], + "ASP": ["N", "CA", "C", "O", "CB", "CG", "OD1", "OD2", "", "", "", "", "", ""], + "CYS": ["N", "CA", "C", "O", "CB", "SG", "", "", "", "", "", "", "", ""], + "GLN": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "NE2", "", "", "", "", ""], + "GLU": ["N", "CA", "C", "O", "CB", "CG", "CD", "OE1", "OE2", "", "", "", "", ""], + "GLY": ["N", "CA", "C", "O", "", "", "", "", "", "", "", "", "", ""], + "HIS": [ + "N", + "CA", + "C", + "O", + "CB", + "CG", + "ND1", + "CD2", + "CE1", + "NE2", + "", + "", + "", + "", + ], + "ILE": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "CD1", "", "", "", "", "", ""], + "LEU": ["N", "CA", "C", "O", "CB", "CG", "CD1", "CD2", "", "", "", "", "", ""], + "LYS": ["N", "CA", "C", "O", "CB", "CG", "CD", "CE", "NZ", "", "", "", "", ""], + "MET": ["N", "CA", "C", "O", "CB", "CG", "SD", "CE", "", "", "", "", "", ""], + "PHE": [ + "N", + "CA", + "C", + "O", + "CB", + "CG", + "CD1", + "CD2", + "CE1", + "CE2", + "CZ", + "", + "", + "", + ], + "PRO": ["N", "CA", "C", "O", "CB", "CG", "CD", "", "", "", "", "", "", ""], + "SER": ["N", "CA", "C", "O", "CB", "OG", "", "", "", "", "", "", "", ""], + "THR": ["N", "CA", "C", "O", "CB", "OG1", "CG2", "", "", "", "", "", "", ""], + "TRP": [ + "N", + "CA", + "C", + "O", + "CB", + "CG", + "CD1", + "CD2", + "NE1", + "CE2", + "CE3", + "CZ2", + "CZ3", + "CH2", + ], + "TYR": [ + "N", + "CA", + "C", + "O", + "CB", + "CG", + "CD1", + "CD2", + "CE1", + "CE2", + "CZ", + "OH", + "", + "", + ], + "VAL": ["N", "CA", "C", "O", "CB", "CG1", "CG2", "", "", "", "", "", "", ""], + "UNK": ["N", "CA", "C", "", "", "", "", "", "", "", "", "", "", ""], +} +restypes = [ + "A", + "R", + "N", + "D", + "C", + "Q", + "E", + "G", + "H", + "I", + "L", + "K", + "M", + "F", + "P", + "S", + "T", + "W", + "Y", + "V", +] +restype_order = {restype: i for (i, restype) in enumerate(restypes)} +restype_num = len(restypes) +unk_restype_index = restype_num +restypes_with_x = [*restypes, "X"] +restype_order_with_x = {restype: i for (i, restype) in enumerate(restypes_with_x)} +bb_atoms = ["N", "CA", "C", "O"] +hydrophobicity = { + "ALA": 0.116, + "ARG": -0.5, + "ASN": -0.264, + "ASP": -0.472, + "CYS": 0.18, + "GLN": -0.249, + "GLU": -0.457, + "GLY": 0.001, + "HIS": -0.335, + "ILE": 0.443, + "LEU": 0.443, + "LYS": -0.217, + "MET": 0.238, + "PHE": 0.5, + "PRO": 0.211, + "SER": -0.141, + "THR": -0.05, + "TRP": 0.378, + "TYR": 0.38, + "VAL": 0.325, +} +side_chain_asa = { + "ALA": 64.7809, + "ARG": 210.02, + "ASN": 113.187, + "ASP": 110.209, + "CYS": 95.2439, + "GLN": 147.855, + "GLU": 143.924, + "GLY": 23.1338, + "HIS": 146.449, + "ILE": 151.242, + "LEU": 139.524, + "LYS": 177.366, + "MET": 164.674, + "PHE": 186.7, + "PRO": 111.533, + "SER": 81.2159, + "THR": 111.597, + "TRP": 229.619, + "TYR": 200.306, + "VAL": 124.237, +} +amino_acid_volumes = { + "A": 88.6, + "R": 173.4, + "N": 114.1, + "D": 111.1, + "C": 108.5, + "Q": 143.8, + "E": 138.4, + "G": 60.1, + "H": 153.2, + "I": 166.7, + "L": 166.7, + "K": 168.6, + "M": 162.9, + "F": 189.9, + "P": 112.7, + "S": 89.0, + "T": 116.1, + "W": 227.8, + "Y": 193.6, + "V": 140.0, + "X": 88.6, +} + + +def sequence_to_onehot( + sequence: str, mapping: Mapping[str, int], map_unknown_to_x: bool = False +) -> np.ndarray: + """Encode a sequence as X with shape ``(l, n_alphabet)``. + + Unknown uppercase letters map to ``X`` only when ``map_unknown_to_x`` is + enabled. Non-alphabetic or lowercase input remains invalid in that mode. + """ + + n_alphabet = max(mapping.values()) + 1 + observed_indices = sorted(set(mapping.values())) + if observed_indices != list(range(n_alphabet)): + raise ValueError( + "The mapping must have values from 0 to num_unique_aas-1 without any gaps. " + f"Got: {sorted(mapping.values())}" + ) + + encoded: np.ndarray = np.empty(len(sequence), dtype=np.intp) + for position, symbol in enumerate(sequence): + if map_unknown_to_x: + if not symbol.isalpha() or not symbol.isupper(): + raise ValueError(f"Invalid character in the sequence: {symbol}") + encoded[position] = mapping.get(symbol, mapping["X"]) + else: + encoded[position] = mapping[symbol] + + one_hot: np.ndarray = np.zeros((len(sequence), n_alphabet), dtype=np.int32) + one_hot[np.arange(len(sequence)), encoded] = 1 + return one_hot + + +restype_1to3 = { + "A": "ALA", + "R": "ARG", + "N": "ASN", + "D": "ASP", + "C": "CYS", + "Q": "GLN", + "E": "GLU", + "G": "GLY", + "H": "HIS", + "I": "ILE", + "L": "LEU", + "K": "LYS", + "M": "MET", + "F": "PHE", + "P": "PRO", + "S": "SER", + "T": "THR", + "W": "TRP", + "Y": "TYR", + "V": "VAL", + "X": "UNK", +} +restype_3to1 = {v: k for (k, v) in restype_1to3.items()} +unk_restype = "UNK" +resnames = [restype_1to3[r] for r in restypes] + [unk_restype] +resname_to_idx = {resname: i for (i, resname) in enumerate(resnames)} +hydrophobic_resnames = {"VAL", "ILE", "LEU", "PHE", "MET", "TRP"} +HHBLITS_AA_TO_ID = { + "A": 0, + "B": 2, + "C": 1, + "D": 2, + "E": 3, + "F": 4, + "G": 5, + "H": 6, + "I": 7, + "J": 20, + "K": 8, + "L": 9, + "M": 10, + "N": 11, + "O": 20, + "P": 12, + "Q": 13, + "R": 14, + "S": 15, + "T": 16, + "U": 1, + "V": 17, + "W": 18, + "X": 20, + "Y": 19, + "Z": 3, + "-": 21, +} +ID_TO_HHBLITS_AA = { + 0: "A", + 1: "C", + 2: "D", + 3: "E", + 4: "F", + 5: "G", + 6: "H", + 7: "I", + 8: "K", + 9: "L", + 10: "M", + 11: "N", + 12: "P", + 13: "Q", + 14: "R", + 15: "S", + 16: "T", + 17: "V", + 18: "W", + 19: "Y", + 20: "X", + 21: "-", +} +restypes_with_x_and_gap = [*restypes, "X", "-"] +MAP_HHBLITS_AATYPE_TO_OUR_AATYPE = tuple( + restypes_with_x_and_gap.index(ID_TO_HHBLITS_AA[i]) for i in range(len(restypes_with_x_and_gap)) +) + + +def _make_standard_atom_mask() -> np.ndarray: + """Return M with shape ``(n_residue_types, n_atom_types)``.""" + + mask: np.ndarray = np.zeros((restype_num + 1, atom_type_num), dtype=np.int32) + for residue_index, residue_code in enumerate(restypes): + residue_name = restype_1to3[residue_code] + atom_indices = [atom_order[name] for name in residue_atoms[residue_name]] + mask[residue_index, atom_indices] = 1 + return mask + + +STANDARD_ATOM_MASK = _make_standard_atom_mask() + + +def chi_angle_atom(atom_index: int) -> np.ndarray: + """Return chi-group selectors X with shape ``(21, 37, 4)``.""" + + selectors: list[np.ndarray] = [] + identity = np.eye(atom_type_num) + for residue_code in restypes: + groups = chi_angles_atoms[restype_1to3[residue_code]] + indices = [atom_order[group[atom_index]] for group in groups] + indices += [-1] * (4 - len(indices)) + selectors.append(identity[indices]) + selectors.append(np.zeros((4, atom_type_num))) + return cast(np.ndarray, np.stack(selectors).transpose(0, 2, 1)) + + +chi_atom_1_one_hot = chi_angle_atom(1) +chi_atom_2_one_hot = chi_angle_atom(2) +chi_angles_atom_indices = [chi_angles_atoms[restype_1to3[r]] for r in restypes] +chi_angles_atom_indices = np.array( + [chi_atoms + [[0, 0, 0, 0]] * (4 - len(chi_atoms)) for chi_atoms in chi_angles_atom_indices] +) +_chi_groups_for_atom: defaultdict[tuple[str, str], list[tuple[int, int]]] = defaultdict(list) +for res_name, chi_angle_atoms_for_res in chi_angles_atoms.items(): + for chi_group_i, chi_group in enumerate(chi_angle_atoms_for_res): + for atom_i, atom in enumerate(chi_group): + _chi_groups_for_atom[res_name, atom].append((chi_group_i, atom_i)) +chi_groups_for_atom = dict(_chi_groups_for_atom) + + +def _make_rigid_transformation_4x4( + ex: np.ndarray, ey: np.ndarray, translation: np.ndarray +) -> np.ndarray: + """Construct homogeneous transform M from two basis vectors and an origin.""" + + unit_x = ex / np.linalg.norm(ex) + orthogonal_y = ey - np.dot(ey, unit_x) * unit_x + unit_y = orthogonal_y / np.linalg.norm(orthogonal_y) + unit_z = np.cross(unit_x, unit_y) + rotation_and_origin = np.stack((unit_x, unit_y, unit_z, translation), axis=0).T + homogeneous_row = np.asarray(((0.0, 0.0, 0.0, 1.0),)) + return cast(np.ndarray, np.concatenate((rotation_and_origin, homogeneous_row), axis=0)) + + +restype_atom37_to_rigid_group: np.ndarray = np.zeros((21, 37), dtype=int) +restype_atom37_mask: np.ndarray = np.zeros((21, 37), dtype=np.float32) +restype_atom37_rigid_group_positions: np.ndarray = np.zeros((21, 37, 3), dtype=np.float32) +restype_atom14_to_rigid_group: np.ndarray = np.zeros((21, 14), dtype=int) +restype_atom14_mask: np.ndarray = np.zeros((21, 14), dtype=np.float32) +restype_atom14_rigid_group_positions: np.ndarray = np.zeros((21, 14, 3), dtype=np.float32) +restype_rigid_group_default_frame: np.ndarray = np.zeros((21, 8, 4, 4), dtype=np.float32) + + +def _make_rigid_group_constants() -> None: + """Populate atom-to-frame assignments and default rigid transforms.""" + + for residue_index, residue_code in enumerate(restypes_with_x): + residue_name = restype_1to3[residue_code] + atom14_names = restype_name_to_atom14_names[residue_name] + for atom_name, group_index, coordinates in rigid_group_atom_positions[residue_name]: + atom37_index = atom_order[atom_name] + atom14_index = atom14_names.index(atom_name) + restype_atom37_to_rigid_group[residue_index, atom37_index] = group_index + restype_atom37_mask[residue_index, atom37_index] = 1 + restype_atom37_rigid_group_positions[residue_index, atom37_index] = coordinates + restype_atom14_to_rigid_group[residue_index, atom14_index] = group_index + restype_atom14_mask[residue_index, atom14_index] = 1 + restype_atom14_rigid_group_positions[residue_index, atom14_index] = coordinates + + positions = { + atom_name: np.asarray(coordinates) + for atom_name, _group_index, coordinates in rigid_group_atom_positions[residue_name] + } + restype_rigid_group_default_frame[residue_index, :2] = np.eye(4) + restype_rigid_group_default_frame[residue_index, 2] = _make_rigid_transformation_4x4( + positions["N"] - positions["CA"], + np.asarray((1.0, 0.0, 0.0)), + positions["N"], + ) + restype_rigid_group_default_frame[residue_index, 3] = _make_rigid_transformation_4x4( + positions["C"] - positions["CA"], + positions["CA"] - positions["N"], + positions["C"], + ) + + groups = chi_angles_atoms[residue_name] + if groups: + first_group = [positions[name] for name in groups[0]] + restype_rigid_group_default_frame[residue_index, 4] = _make_rigid_transformation_4x4( + first_group[2] - first_group[1], + first_group[0] - first_group[1], + first_group[2], + ) + for chi_index, group in enumerate(groups[1:], start=1): + axis_end = positions[group[2]] + restype_rigid_group_default_frame[residue_index, 4 + chi_index] = ( + _make_rigid_transformation_4x4( + axis_end, + np.asarray((-1.0, 0.0, 0.0)), + axis_end, + ) + ) + + +_make_rigid_group_constants() + + +def make_atom14_dists_bounds( + overlap_tolerance: float = 1.5, + bond_length_tolerance_factor: float = 15.0, +) -> dict[str, np.ndarray]: + """Build lower, upper, and uncertainty tensors with shape ``(21, 14, 14)``.""" + + lower_bounds: np.ndarray = np.zeros((21, 14, 14), np.float32) + upper_bounds: np.ndarray = np.zeros((21, 14, 14), np.float32) + stddevs: np.ndarray = np.zeros((21, 14, 14), np.float32) + residue_bonds, residue_virtual_bonds, _angles = load_stereo_chemical_props() + for residue_index, residue_code in enumerate(restypes): + residue_name = restype_1to3[residue_code] + atom_names = restype_name_to_atom14_names[residue_name] + for atom_a_index, atom_a_name in enumerate(atom_names): + if not atom_a_name: + continue + radius_a = van_der_waals_radius[atom_a_name[0]] + for atom_b_index, atom_b_name in enumerate(atom_names): + if not atom_b_name or atom_a_index == atom_b_index: + continue + clash_floor = radius_a + van_der_waals_radius[atom_b_name[0]] - overlap_tolerance + lower_bounds[residue_index, atom_a_index, atom_b_index] = clash_floor + lower_bounds[residue_index, atom_b_index, atom_a_index] = clash_floor + upper_bounds[residue_index, atom_a_index, atom_b_index] = 1e10 + upper_bounds[residue_index, atom_b_index, atom_a_index] = 1e10 + + for bond in residue_bonds[residue_name] + residue_virtual_bonds[residue_name]: + atom_a_index = atom_names.index(bond.atom1_name) + atom_b_index = atom_names.index(bond.atom2_name) + lower = bond.length - bond_length_tolerance_factor * bond.stddev + upper = bond.length + bond_length_tolerance_factor * bond.stddev + for row, column in ( + (atom_a_index, atom_b_index), + (atom_b_index, atom_a_index), + ): + lower_bounds[residue_index, row, column] = lower + upper_bounds[residue_index, row, column] = upper + stddevs[residue_index, row, column] = bond.stddev + return { + "lower_bound": lower_bounds, + "upper_bound": upper_bounds, + "stddev": stddevs, + } + + +restype_atom14_ambiguous_atoms: np.ndarray = np.zeros((21, 14), dtype=np.float32) +restype_atom14_ambiguous_atoms_swap_idx = np.tile(np.arange(14, dtype=int), (21, 1)) + + +def _make_atom14_ambiguity_feats() -> None: + """Mark symmetry-equivalent atom names and their exchange indices.""" + + for residue_name, swaps in residue_atom_renaming_swaps.items(): + residue_index = restype_order[restype_3to1[residue_name]] + atom_names = restype_name_to_atom14_names[residue_name] + for atom_a, atom_b in swaps.items(): + atom_a_index = atom_names.index(atom_a) + atom_b_index = atom_names.index(atom_b) + restype_atom14_ambiguous_atoms[residue_index, (atom_a_index, atom_b_index)] = 1 + restype_atom14_ambiguous_atoms_swap_idx[residue_index, (atom_a_index, atom_b_index)] = ( + atom_b_index, + atom_a_index, + ) + + +_make_atom14_ambiguity_feats() + + +def aatype_to_str_sequence(aatype: np.ndarray) -> str: + """Decode residue-type indices without changing their order.""" + + return "".join(restypes_with_x[index] for index in aatype) + + +CA_TO_N_NORM = 1.4591 +CA_TO_C_NORM = 1.5252 + + +def _make_restype_atom37_to_atom14() -> np.ndarray: + """Return atom37-to-atom14 lookup M with shape ``(21, 37)``.""" + + rows: list[list[int]] = [] + for residue_code in restypes: + names = restype_name_to_atom14_names[restype_1to3[residue_code]] + atom14_index = {name: index for index, name in enumerate(names)} + rows.append([atom14_index.get(name, 0) for name in atom_types]) + rows.append([0] * atom_type_num) + return np.asarray(rows, dtype=np.int32) + + +def _make_restype_atom14_to_atom37() -> np.ndarray: + """Return atom14-to-atom37 lookup M with shape ``(21, 14)``.""" + + rows = [ + [atom_order.get(name, 0) for name in restype_name_to_atom14_names[residue_name]] + for residue_name in resnames[:-1] + ] + rows.append([0] * 14) + return np.asarray(rows, dtype=np.int32) + + +RESTYPE_ATOM14_TO_ATOM37 = _make_restype_atom14_to_atom37() +RESTYPE_ATOM37_TO_ATOM14 = _make_restype_atom37_to_atom14() +CHAIN_BREAK_TOKEN = "|" diff --git a/fastplms/models/esmfold2/esmfold2_sequential_dataclass.py b/fastplms/models/esmfold2/esmfold2_sequential_dataclass.py new file mode 100644 index 0000000000000000000000000000000000000000..97b1a47c3c65d51625541f41abfbd0821617301a --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_sequential_dataclass.py @@ -0,0 +1,112 @@ +"""Dataclass support for aligned residue-level fields.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from collections.abc import Iterable +from dataclasses import Field, dataclass, fields, replace +from typing import Any, Self + +import numpy as np + +from .esmfold2_misc import concat_objects, slice_any_object + +Index = int | list[int] | slice | np.ndarray + + +def _is_sequential(field: Field[Any]) -> bool: + return bool(field.metadata.get("sequence", False)) + + +def _sequence_axis(field: Field[Any]) -> int: + axis = int(field.metadata.get("sequence_dim", 0)) + if axis not in (0, 1): + raise NotImplementedError("SequentialDataclass supports sequence_dim values zero and one.") + return axis + + +def _slice_value(value: Any, index: Index, axis: int) -> Any: + if axis == 0: + return slice_any_object(value, index) + sliced = [slice_any_object(track, index) for track in value] + return value.__class__(sliced) + + +def _iter_sequence_lengths(value: Any, axis: int) -> Iterable[int]: + if axis == 0: + yield len(value) + else: + yield from (len(track) for track in value) + + +@dataclass(frozen=True) +class SequentialDataclass(ABC): + """Keep dataclass fields aligned along a shared residue dimension. + + A subclass marks aligned fields with ``metadata={"sequence": True}``. + ``sequence_dim`` may be zero for a direct sequence or one for a collection + of aligned tracks. ``join_token`` is passed to the package concatenation + helper when instances are joined. + """ + + def __post_init__(self) -> None: + expected = len(self) + for field in fields(self): + if not _is_sequential(field) or field.name == "complex": + continue + value = getattr(self, field.name) + if value is None: + continue + for actual in _iter_sequence_lengths(value, _sequence_axis(field)): + if actual != expected: + raise ValueError( + f"Mismatch in sequence length for field: {field.name}. " + f"Expected {expected}, received {actual}" + ) + + @abstractmethod + def __len__(self) -> int: + """Return the shared sequence length.""" + + raise NotImplementedError + + def __getitem__(self, index: Index) -> Self: + """Apply one sequence index to every aligned field.""" + + normalized_index: Index = [index] if isinstance(index, int) else index + updates: dict[str, Any] = {} + for field in fields(self): + if not _is_sequential(field): + continue + value = getattr(self, field.name) + if value is not None: + updates[field.name] = _slice_value(value, normalized_index, _sequence_axis(field)) + return replace(self, **updates) + + @classmethod + def concat(cls, items: list[Self], **overrides: Any) -> Self: + """Join aligned fields and retain non-sequential values from the first item.""" + + if not items: + raise ValueError("SequentialDataclass.concat requires at least one item.") + + updates: dict[str, Any] = {} + for field in fields(cls): + if not _is_sequential(field): + continue + first_value = getattr(items[0], field.name) + if first_value is None: + continue + values = [getattr(item, field.name) for item in items] + join_token = field.metadata.get("join_token") + if _sequence_axis(field) == 0: + updates[field.name] = concat_objects(values, join_token) + else: + tracks = [concat_objects(track, join_token) for track in zip(*values, strict=True)] + updates[field.name] = first_value.__class__(tracks) + + updates.update(overrides) + return replace(items[0], **updates) + + +__all__ = ["SequentialDataclass"] diff --git a/fastplms/models/esmfold2/esmfold2_system.py b/fastplms/models/esmfold2/esmfold2_system.py new file mode 100644 index 0000000000000000000000000000000000000000..23b69dc10fae713e3b7be76fd9475ab054a58913 --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_system.py @@ -0,0 +1,56 @@ +"""Filesystem and subprocess types used by optional structure utilities.""" + +from __future__ import annotations + +import io +import subprocess +from pathlib import Path +from typing import Any, TypeAlias + +PathLike: TypeAlias = str | Path +PathOrBuffer: TypeAlias = PathLike | io.StringIO + + +def _stdout_destination(*, capture_output: bool, quiet: bool) -> int | None: + if capture_output: + return subprocess.PIPE + if quiet: + return subprocess.DEVNULL + return None + + +def _stderr_text(error: subprocess.CalledProcessError) -> str: + stderr = error.stderr + if stderr is None: + return "" + if isinstance(stderr, bytes): + return stderr.decode(errors="replace") + return str(stderr) + + +def run_subprocess_with_errorcheck( + *popenargs: Any, + capture_output: bool = False, + quiet: bool = False, + env: dict[str, str] | None = None, + shell: bool = False, + executable: str | None = None, + **kwargs: Any, +) -> subprocess.CompletedProcess[Any]: + """Run a command and include captured standard error in failures.""" + + stdout = _stdout_destination(capture_output=capture_output, quiet=quiet) + try: + return subprocess.run( + *popenargs, + check=True, + env=env, + executable=executable, + shell=shell, + stderr=subprocess.PIPE, + stdout=stdout, + **kwargs, + ) + except subprocess.CalledProcessError as error: + message = f"Command failed with errorcode {error.returncode}.\n\n{_stderr_text(error)}" + raise RuntimeError(message) from error diff --git a/fastplms/models/esmfold2/esmfold2_types.py b/fastplms/models/esmfold2/esmfold2_types.py new file mode 100644 index 0000000000000000000000000000000000000000..38d581ec65c1f918063c271f15b5a71ae1fe7c99 --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_types.py @@ -0,0 +1,31 @@ +"""Stable namespace for ESMFold2 input schema types.""" + +from __future__ import annotations + +from . import esmfold2_input_builder as _input_schema +from .esmfold2_msa import MSA +from .esmfold2_parsing import FastaEntry + +Modification = _input_schema.Modification +ProteinInput = _input_schema.ProteinInput +RNAInput = _input_schema.RNAInput +DNAInput = _input_schema.DNAInput +LigandInput = _input_schema.LigandInput +DistogramConditioning = _input_schema.DistogramConditioning +PocketConditioning = _input_schema.PocketConditioning +CovalentBond = _input_schema.CovalentBond +StructurePredictionInput = _input_schema.StructurePredictionInput + +__all__ = [ + "MSA", + "CovalentBond", + "DNAInput", + "DistogramConditioning", + "FastaEntry", + "LigandInput", + "Modification", + "PocketConditioning", + "ProteinInput", + "RNAInput", + "StructurePredictionInput", +] diff --git a/fastplms/models/esmfold2/esmfold2_utils_types.py b/fastplms/models/esmfold2/esmfold2_utils_types.py new file mode 100644 index 0000000000000000000000000000000000000000..13590243d386bc83fa3c01bb28f8ed6fb79762f2 --- /dev/null +++ b/fastplms/models/esmfold2/esmfold2_utils_types.py @@ -0,0 +1,38 @@ +"""Small public types shared by the ESMFold2 structure utilities. + +These definitions are intentionally independent of cloud-storage packages. Any +path object implementing :class:`os.PathLike` is accepted by the runtime file +helpers, including cloud-path implementations installed by an application. +""" + +from __future__ import annotations + +import io +import os +from dataclasses import dataclass +from typing import TypeAlias + +PathLike: TypeAlias = str | os.PathLike[str] +PathOrBuffer: TypeAlias = PathLike | io.TextIOBase + + +@dataclass(slots=True) +class FunctionAnnotation: + """A residue-range annotation using one-based inclusive coordinates.""" + + label: str + start: int + end: int + + def to_tuple(self) -> tuple[str, int, int]: + """Return the serialization order used by annotation tokenizers.""" + + return (self.label, self.start, self.end) + + def __len__(self) -> int: + """Return the number of annotated residues.""" + + return self.end - self.start + 1 + + +__all__ = ["FunctionAnnotation", "PathLike", "PathOrBuffer"] diff --git a/fastplms/models/esmfold2/modeling_esmfold2.py b/fastplms/models/esmfold2/modeling_esmfold2.py new file mode 100644 index 0000000000000000000000000000000000000000..2afc02b24c7b2124b2fe360c114d4279f7f7ca47 --- /dev/null +++ b/fastplms/models/esmfold2/modeling_esmfold2.py @@ -0,0 +1,2040 @@ +"""PyTorch ESMFold2 model: the standard released architecture. + +Quickstart:: + + from transformers import ESMFold2Model + + model = ESMFold2Model.from_pretrained("biohub/ESMFold2").cuda().eval() + open("ubq.pdb", "w").write(model.infer_protein_as_pdb("MQIFVKTLTGKT...")) + +For multi-chain, ligand, and MSA inputs, use ``model.input_types`` together +with ``model.fold(...)`` or ``model.prepare_structure_input(...)``. +""" + +from __future__ import annotations + +import gc +import importlib +import importlib.metadata +import math +from contextlib import contextmanager +from dataclasses import asdict, dataclass +from pathlib import Path +from typing import Any, ClassVar, Literal, cast + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from transformers.modeling_outputs import ModelOutput +from transformers.modeling_utils import PreTrainedModel + +from ...attention import get_attn_implementation, set_config_attn_implementation + +try: + from fastplms.models.ttt import FastPLMTestTimeTrainingMixin, TTTConfig +except ModuleNotFoundError as error: + if error.name != "fastplms": + raise + from ..ttt import FastPLMTestTimeTrainingMixin, TTTConfig + +from .attention import ESMFold2AttentionMixin +from .configuration_esmfold2 import ESMFold2Config, normalize_esmc_id +from .embedding import ESMFold2EmbeddingMixin +from .esmfold2_constants_esm3 import ( + SEQUENCE_BOS_TOKEN, + SEQUENCE_EOS_TOKEN, + SEQUENCE_MASK_TOKEN, + SEQUENCE_PAD_TOKEN, + SEQUENCE_STANDARD_AA_MAX_TOKEN, + SEQUENCE_STANDARD_AA_MIN_TOKEN, + SEQUENCE_VOCAB, +) +from .modeling_esmfold2_common import ( + CHAR_VOCAB_SIZE, + MAX_ATOMIC_NUMBER, + MSA_CONDITIONING_INPUT_NAMES, + NUM_RES_TYPES, + DiffusionStructureHead, + FoldingTrunk, + InputsEmbedder, + LanguageModelShim, + MSAPairWeightedAveraging, + OuterProductMean, + ResIdxAsymIdSymIdEntityIdEncoding, + RowAttentionPooling, + SwiGLUMLP, + TriangleMultiplicativeUpdate, + _categorical_mean, + _compute_intra_token_idx, + compute_lm_hidden_states, + gather_rep_atom_coords, + gather_token_to_atom, + maybe_apply_msa_column_masking, + maybe_subsample_msa, + validate_kernel_backend, + validate_msa_conditioning_inputs, + validate_prepared_auxiliary_inputs, +) + +_ESMC_FP8_LINEAR_SUFFIX = ".attn.out_proj" +_ESMC_FP8_EXPECTED_PROJECTIONS = 80 +_ESMC_BF16_ROUNDED_SHAPES = frozenset({(960, 30), (1152, 36)}) +_EPS = 1e-6 +_NONPOLYMER_ID = 4 + +# Default for the triangle, OPM, and pair-transition l^2 operations. Caps peak +# memory so l around 2k folds on an 80 GB GPU (about 76 GB at chunk=128 for +# l=1438; +# chunk=64 leaves headroom for the largest foldbench targets). Override via +# ``model.set_chunk_size(...)``; pass None to disable chunking (faster for +# short l but OOM-prone past approximately 600). +_DEFAULT_CHUNK_SIZE = 64 + + +@dataclass +class ESMFold2Output(ModelOutput): + """Transformers-compatible output shared by released and experimental folds. + + ``last_hidden_state`` is the final pair representation. When requested, + ``hidden_states`` contains the token-input representation followed by the + final pair representation. The structure trunks do not expose normalized + post-softmax attention tensors, so ``output_attentions=True`` fails + explicitly instead of returning incomplete data. + """ + + last_hidden_state: Tensor | None = None + hidden_states: tuple[Tensor, ...] | None = None + attentions: tuple[Tensor, ...] | None = None + distogram_logits: Tensor | None = None + sample_atom_coords: Tensor | None = None + representative_atom_coords: Tensor | None = None + atom_pad_mask: Tensor | None = None + residue_index: Tensor | None = None + entity_id: Tensor | None = None + plddt_logits: Tensor | None = None + plddt: Tensor | None = None + plddt_per_atom: Tensor | None = None + plddt_ca: Tensor | None = None + complex_plddt: Tensor | None = None + complex_iplddt: Tensor | None = None + pae_logits: Tensor | None = None + pae: Tensor | None = None + pde_logits: Tensor | None = None + pde: Tensor | None = None + resolved_logits: Tensor | None = None + ptm: Tensor | None = None + iptm: Tensor | None = None + pair_chains_iptm: Tensor | None = None + + +def _resolve_structure_output_controls( + config: ESMFold2Config, + *, + output_attentions: bool | None, + output_hidden_states: bool | None, + return_dict: bool | None, +) -> tuple[bool, bool]: + resolved_attentions = ( + config.output_attentions if output_attentions is None else output_attentions + ) + if resolved_attentions: + raise NotImplementedError( + "ESMFold2 does not expose normalized attention tensors from its structure " + "trunk. output_attentions=True is unsupported." + ) + resolved_hidden_states = ( + config.output_hidden_states + if output_hidden_states is None + else output_hidden_states + ) + resolved_return_dict = config.use_return_dict if return_dict is None else return_dict + return bool(resolved_hidden_states), bool(resolved_return_dict) + + +def _finalize_structure_output( + output: dict[str, Tensor], + *, + token_input_state: Tensor, + pair_state: Tensor, + output_hidden_states: bool, + return_dict: bool, +) -> ESMFold2Output | tuple[Any, ...]: + model_output = ESMFold2Output( + last_hidden_state=pair_state, + hidden_states=(token_input_state, pair_state) if output_hidden_states else None, + **output, + ) + return model_output if return_dict else model_output.to_tuple() + + +class _ESMFold2ESMplusplusAdapter(nn.Module): + def __init__(self, model: nn.Module) -> None: + super().__init__() + self.model = model + + @property + def config(self): + return self.model.config + + def set_attn_implementation(self, attn_implementation: str) -> None: + """Update ESMC through its Transformers-compatible attention API.""" + + self.model.set_attn_implementation(attn_implementation) + + def forward( + self, + input_ids: Tensor, + attention_mask: Tensor | None = None, + sequence_id: Tensor | None = None, + output_hidden_states: bool | None = None, + output_attentions: bool | None = None, + return_dict: bool | None = None, + compute_sae: bool = True, + normalize_sae: bool = False, + ): + del return_dict, compute_sae, normalize_sae + output = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + sequence_id=sequence_id, + output_hidden_states=output_hidden_states, + output_attentions=output_attentions, + return_dict=True, + esmfold2_hidden_states=True, + ) + if output_hidden_states: + hidden_states = output.hidden_states + if hidden_states is None: + raise RuntimeError("ESM++ did not return requested hidden states.") + if isinstance(hidden_states, torch.Tensor): + output.hidden_states = hidden_states + else: + output.hidden_states = torch.stack(tuple(hidden_states), dim=0) + return output + + +def _load_fastplms_esmplusplus_for_esmfold2( + esmc_model_path: str, + attn_backend: str, + device: torch.device, + dtype: torch.dtype, + local_files_only: bool = False, +) -> _ESMFold2ESMplusplusAdapter: + from fastplms.models.esm_plusplus.modeling_esm_plusplus import ( + ESMplusplusConfig, + ESMplusplusModel, + ) + + normalized_path = normalize_esmc_id(esmc_model_path) + source_revision, _ = _manifest_esmc_checkpoint_contract(normalized_path) + revision_kwargs: dict[str, Any] = { + "local_files_only": local_files_only, + } + if source_revision is not None: + revision_kwargs["revision"] = source_revision + esmc_config = ESMplusplusConfig.from_pretrained(normalized_path, **revision_kwargs) + set_config_attn_implementation(esmc_config, attn_backend) + # The smaller folding checkpoints pin BF16-rounded frozen backbone values. + # Retain that rounding if a caller explicitly requests FP32 execution. + small_backbone = ( + esmc_config.hidden_size, + esmc_config.num_hidden_layers, + ) in _ESMC_BF16_ROUNDED_SHAPES + checkpoint_dtype = torch.bfloat16 if small_backbone else dtype + load_kwargs: dict[str, Any] = { + "config": esmc_config, + "torch_dtype": checkpoint_dtype, + **revision_kwargs, + } + if device.type == "cuda": + # Device mapping constructs parameters on the destination GPU instead + # of materializing the 6B backbone in host memory first. + load_kwargs["device_map"] = {"": str(device)} + esmc = ESMplusplusModel.from_pretrained(normalized_path, **load_kwargs) + if device.type != "cuda" or checkpoint_dtype != dtype: + esmc = esmc.to(device=device, dtype=dtype) + else: + loaded_device = next(esmc.parameters()).device + if loaded_device != device: + raise RuntimeError( + f"ESMC loaded on {loaded_device}, expected direct loading on {device}." + ) + return _ESMFold2ESMplusplusAdapter(esmc).eval() + + +def _manifest_esmc_checkpoint_contract( + esmc_model_path: str, +) -> tuple[str | None, dict[str, str]]: + """Return the immutable manifest identity for a registered ESMC source. + + Local checkpoint directories deliberately return no Hub revision. A known + Hub repository is always loaded at the revision and file identities in + ``models.toml`` instead of following a mutable branch. + """ + + normalized_path = normalize_esmc_id(esmc_model_path) + try: + if Path(normalized_path).exists(): + return None, {} + except OSError: + # A repository ID may be too long or otherwise invalid as a local path. + pass + + from fastplms.registry import get_model_registry + + registry = get_model_registry() + backbone_model = registry.families["esmfold2"].backbone_model + if backbone_model is None: + raise RuntimeError("families.esmfold2 must declare backbone_model.") + spec = registry[backbone_model] + checkpoints = ( + spec.fast, + spec.official, + *( + checkpoint + for model in registry.by_family("esmfold2") + if model.backbone_model is not None + for checkpoint in ( + registry[model.backbone_model].fast, + registry[model.backbone_model].official, + ) + ), + ) + for checkpoint in checkpoints: + if checkpoint.repo_id == normalized_path: + return checkpoint.revision, { + item.path: item.encoded for item in checkpoint.files + } + if "/" in normalized_path: + raise ValueError( + f"Remote ESMC source {normalized_path!r} is not a manifest-declared " + "ESMFold2 backbone." + ) + return None, {} + + +ESMCPrecision = Literal["auto", "bf16", "fp32", "fp8"] + + +@dataclass(frozen=True, slots=True) +class ESMCPrecisionStatus: + """Resolved ESMC precision and the evidence used to choose it.""" + + requested: str + resolved: str + reason: str + device: str + transformer_engine_version: str | None + + def as_dict(self) -> dict[str, str | None]: + return asdict(self) + + +def _transformer_engine_version() -> str | None: + try: + return importlib.metadata.version("transformer-engine") + except importlib.metadata.PackageNotFoundError: + return None + + +def _load_transformer_engine() -> tuple[Any, Any]: + """Load Transformer Engine lazily so core imports stay dependency-free.""" + + try: + te = importlib.import_module("transformer_engine.pytorch") + recipe = importlib.import_module("transformer_engine.common.recipe") + except (ImportError, OSError, RuntimeError) as error: + raise RuntimeError( + f"Transformer Engine could not be imported: {type(error).__name__}: {error}" + ) from error + if not hasattr(recipe, "Float8CurrentScaling"): + raise RuntimeError( + "Transformer Engine does not expose Float8CurrentScaling, which is " + "required by the validated ESMC FP8 path." + ) + return te, recipe + + +def _te_fp8_capability(device: torch.device) -> tuple[bool, str]: + """Return whether the validated Transformer Engine FP8 path can run.""" + + if device.type != "cuda": + return False, "FP8 requires direct ESMC loading onto a CUDA device." + if not torch.cuda.is_available(): + return False, "CUDA is unavailable." + try: + major, minor = torch.cuda.get_device_capability(device) + except (AssertionError, RuntimeError, ValueError) as error: + return False, f"CUDA capability query failed: {error}" + if not (major >= 9 or (major == 8 and minor >= 9)): + return False, f"CUDA capability {major}.{minor} does not support FP8." + try: + te, _ = _load_transformer_engine() + except RuntimeError as error: + return False, str(error) + + probe = getattr(te, "is_fp8_available", None) + if probe is None: + try: + probe = importlib.import_module("transformer_engine.pytorch.fp8").is_fp8_available + except (ImportError, AttributeError, OSError, RuntimeError) as error: + return False, f"Transformer Engine has no usable FP8 probe: {error}" + try: + try: + result = probe(return_reason=True) + except TypeError: + result = probe() + except (OSError, RuntimeError) as error: + return False, f"Transformer Engine FP8 probe failed: {error}" + if isinstance(result, tuple): + available = bool(result[0]) + detail = str(result[1]) if len(result) > 1 and result[1] else "" + else: + available = bool(result) + detail = "" + if not available: + return False, detail or "Transformer Engine reports FP8 unavailable." + return True, ( + "Transformer Engine reports FP8 availability; FastPLMs will convert " + "the validated ESMC attention output projections." + ) + + +def _resolve_esmc_precision(requested: str, device: torch.device) -> ESMCPrecisionStatus: + allowed = {"auto", "bf16", "fp32", "fp8"} + if requested not in allowed: + raise ValueError(f"precision must be one of {sorted(allowed)}, got {requested!r}.") + if requested in {"auto", "bf16", "fp32"}: + resolved = "bf16" if requested == "auto" else requested + reason = ( + "Automatic precision defaults to BF16; select esmc_precision='fp8' " + "explicitly to opt in to the validated Transformer Engine path." + if requested == "auto" + else "Precision was selected explicitly." + ) + return ESMCPrecisionStatus( + requested=requested, + resolved=resolved, + reason=reason, + device=str(device), + transformer_engine_version=_transformer_engine_version(), + ) + available, reason = _te_fp8_capability(device) + if not available: + raise RuntimeError(f"esmc_precision='fp8' is unavailable: {reason}") + return ESMCPrecisionStatus( + requested=requested, + resolved="fp8", + reason=reason, + device=str(device), + transformer_engine_version=_transformer_engine_version(), + ) + + +def _install_esmc_backbone( + model: Any, + esmc_model_path: str, + *, + precision: str, + device: str | torch.device | None = None, + local_files_only: bool = False, +) -> None: + if precision == "fp8" and ( + model.config.lm_num_layers != 80 or model.config.lm_d_model != 2560 + ): + raise ValueError("ESMFold2 FP8 is supported only for the ESMC-6B backbone.") + target_device = torch.device(device) if device is not None else model.device + if target_device.type == "cuda" and target_device.index is None and torch.cuda.is_available(): + target_device = torch.device("cuda", torch.cuda.current_device()) + model_device = torch.device(model.device) + if model_device.type == "cuda" and model_device.index is None and torch.cuda.is_available(): + model_device = torch.device("cuda", torch.cuda.current_device()) + if target_device != model_device: + raise ValueError( + f"ESMC target device {target_device} must match the ESMFold2 device " + f"{model_device}. Move ESMFold2 before loading or reloading ESMC." + ) + status = _resolve_esmc_precision(precision, target_device) + normalized_source = normalize_esmc_id(esmc_model_path) + source_revision, source_files = _manifest_esmc_checkpoint_contract(normalized_source) + attention_implementation = get_attn_implementation(model.config) + model.config.esmc_attn_backend = attention_implementation + dtype = torch.float32 if status.resolved == "fp32" else torch.bfloat16 + esmc = _load_fastplms_esmplusplus_for_esmfold2( + esmc_model_path=esmc_model_path, + attn_backend=attention_implementation, + device=target_device, + dtype=dtype, + local_files_only=local_files_only, + ) + if esmc.config.hidden_size != model.config.lm_d_model: + raise ValueError( + f"ESMFold2 expected lm_d_model={model.config.lm_d_model}, " + f"but loaded ESMC hidden_size={esmc.config.hidden_size}." + ) + if esmc.config.num_hidden_layers != model.config.lm_num_layers: + raise ValueError( + f"ESMFold2 expected lm_num_layers={model.config.lm_num_layers}, " + f"but loaded ESMC num_hidden_layers={esmc.config.num_hidden_layers}." + ) + esmc.eval().requires_grad_(False) + fp8_module_paths: tuple[str, ...] = () + if status.resolved == "fp8": + fp8_module_paths = _convert_esmc_attention_outputs_to_te(esmc) + status = ESMCPrecisionStatus( + requested=status.requested, + resolved=status.resolved, + reason=( + f"{status.reason} Converted {len(fp8_module_paths)} projections; " + "canonical checkpoint weights remain BF16." + ), + device=status.device, + transformer_engine_version=status.transformer_engine_version, + ) + model._esmc_source = normalized_source + model._esmc_source_revision = source_revision + model._esmc_source_files = source_files + model._esmc_local_files_only = local_files_only + model._esmc_precision_policy = precision + model._esmc_precision_status = status + model._esmc_fp8 = status.resolved == "fp8" + model._esmc_fp8_module_paths = fp8_module_paths + model.config.esmc_precision = precision + model._esmc = esmc + model._ttt_lm_head = None + + +def _drop_transient_esmc_state( + module: nn.Module, + state_dict: dict[str, Tensor], + prefix: str, + local_metadata: dict[str, Any], +) -> None: + """Exclude runtime ESMC/TTT modules from canonical folding checkpoints.""" + + del module, local_metadata + transient_prefixes = (f"{prefix}_esmc.", f"{prefix}_ttt_lm_head.") + for key in tuple(state_dict): + if key.startswith(transient_prefixes): + del state_dict[key] + + +def _reload_esmc_bf16_for_gradients(model: Any, *, reason: str) -> None: + """Use BF16 temporarily without overwriting the persisted serving policy.""" + + policy = model._esmc_precision_policy + model.reload_esmc(precision="bf16", device=model.device) + model._esmc_precision_policy = policy + model.config.esmc_precision = policy + status = model._esmc_precision_status + model._esmc_precision_status = ESMCPrecisionStatus( + requested=policy, + resolved="bf16", + reason=reason, + device=status.device, + transformer_engine_version=status.transformer_engine_version, + ) + + +class PairTransition(nn.Module): + """LayerNorm + SwiGLU feed-forward residual block on the pair representation.""" + + def __init__(self, d_model: int, expansion_ratio: int = 4) -> None: + super().__init__() + self.norm = nn.LayerNorm(d_model) + self.ffn = SwiGLUMLP(d_model, expansion_ratio=expansion_ratio, bias=False) + self._chunk_size: int | None = _DEFAULT_CHUNK_SIZE + + def set_chunk_size(self, chunk_size: int | None) -> None: + self._chunk_size = chunk_size + + def forward(self, x: Tensor) -> Tensor: + if self._chunk_size is None or x.shape[1] <= self._chunk_size: + return self.ffn(self.norm(x)) + out: list[Tensor] = [] + for s in range(0, x.shape[1], self._chunk_size): + e = min(s + self._chunk_size, x.shape[1]) + sl = x[:, s:e] + out.append(self.ffn(self.norm(sl))) + return torch.cat(out, dim=1) + + +class ConfidenceHead(nn.Module): + """Predicts pLDDT, PAE, PDE, resolved-atom probability and distogram bins.""" + + boundaries: Tensor + + def __init__(self, config: ESMFold2Config) -> None: + super().__init__() + ch = config.confidence_head + d_single = config.d_single + d_pair = config.d_pair + d_inputs = config.inputs.d_inputs + + boundaries = torch.linspace(ch.min_dist, ch.max_dist, ch.distogram_bins - 1) + self.register_buffer("boundaries", boundaries) + self.dist_bin_pairwise_embed = nn.Embedding(ch.distogram_bins, d_pair) + + self.s_norm = nn.LayerNorm(d_single) + self.s_inputs_to_single = nn.Linear(d_inputs, d_single, bias=False) + self.s_to_z = nn.Linear(d_inputs, d_pair, bias=False) + self.s_to_z_transpose = nn.Linear(d_inputs, d_pair, bias=False) + self.s_to_z_prod_in1 = nn.Linear(d_inputs, d_pair, bias=False) + self.s_to_z_prod_in2 = nn.Linear(d_inputs, d_pair, bias=False) + self.s_to_z_prod_out = nn.Linear(d_pair, d_pair, bias=False) + self.s_input_to_s = nn.Linear(d_inputs, d_single, bias=False) + self.s_inputs_norm = nn.LayerNorm(d_inputs) + self.z_norm = nn.LayerNorm(d_pair) + + self.row_attention_pooling = RowAttentionPooling(d_pair=d_pair, d_single=d_single) + + pf = ch.folding_trunk + self.folding_trunk = FoldingTrunk(n_layers=pf.n_layers, d_pair=d_pair, expansion_ratio=4) + + # Heads. + self.plddt_ln = nn.LayerNorm(d_single) + max_atoms_per_token = 23 + self.plddt_weight = nn.Parameter( + torch.zeros(max_atoms_per_token, d_single, ch.num_plddt_bins) + ) + + self.pae_ln = nn.LayerNorm(d_pair) + self.pae_head = nn.Linear(d_pair, ch.num_pae_bins, bias=False) + + self.pde_ln = nn.LayerNorm(d_pair) + self.pde_head = nn.Linear(d_pair, ch.num_pde_bins, bias=False) + + self.resolved_ln = nn.LayerNorm(d_single) + # 2 = resolved logits ([unresolved, resolved]). + self.resolved_weight = nn.Parameter(torch.zeros(max_atoms_per_token, d_single, 2)) + + def set_kernel_backend(self, backend: str | None) -> None: + self.folding_trunk.set_kernel_backend(backend) + + def set_chunk_size(self, chunk_size: int | None) -> None: + self.folding_trunk.set_chunk_size(chunk_size) + + @staticmethod + def _repeat_batch(x: Tensor, num_diffusion_samples: int) -> Tensor: + return x if num_diffusion_samples == 1 else x.repeat_interleave(num_diffusion_samples, 0) + + @staticmethod + def _flatten_sample_axis(x: Tensor) -> Tensor: + if x.ndim == 4: + b, mult, n, c = x.shape + return x.reshape(b * mult, n, c) + return x + + def forward( + self, + s_inputs: Tensor, + z: Tensor, + x_pred: Tensor, + distogram_atom_idx: Tensor, + token_attention_mask: Tensor, + atom_to_token: Tensor, + atom_attention_mask: Tensor, + asym_id: Tensor, + mol_type: Tensor, + num_diffusion_samples: int = 1, + relative_position_encoding: Tensor | None = None, + token_bonds_encoding: Tensor | None = None, + ) -> dict[str, Tensor]: + s_inputs_normed = self.s_inputs_norm(s_inputs) + + z_base = self.z_norm(z) + if relative_position_encoding is not None: + z_base = z_base + relative_position_encoding + if token_bonds_encoding is not None: + z_base = z_base + token_bonds_encoding + z_base = z_base + self.s_to_z(s_inputs_normed).unsqueeze(2) + z_base = z_base + self.s_to_z_transpose(s_inputs_normed).unsqueeze(1) + z_base = z_base + self.s_to_z_prod_out( + self.s_to_z_prod_in1(s_inputs_normed)[:, :, None, :] + * self.s_to_z_prod_in2(s_inputs_normed)[:, None, :, :] + ) + + pair = self._repeat_batch(z_base, num_diffusion_samples) + x_pred_flat = self._flatten_sample_axis(x_pred) + atom_to_token_m = self._repeat_batch(atom_to_token, num_diffusion_samples) + atom_mask_m = self._repeat_batch(atom_attention_mask, num_diffusion_samples) + rep_idx_m = self._repeat_batch(distogram_atom_idx, num_diffusion_samples).long() + mask = self._repeat_batch(token_attention_mask, num_diffusion_samples) + expanded_batch_size = pair.shape[0] + + rep_coords = gather_rep_atom_coords(x_pred_flat, rep_idx_m) + rep_distances = torch.cdist( + rep_coords, rep_coords, compute_mode="donot_use_mm_for_euclid_dist" + ) + distogram_bins = (rep_distances.unsqueeze(-1) > self.boundaries).sum(dim=-1).long() + pair = pair + self.dist_bin_pairwise_embed(distogram_bins) + + pair_mask = mask[:, :, None].float() * mask[:, None, :].float() + + # FoldingTrunk handles the bf16 cast internally during inference so + # each block's fused trimul engages. In-place residual avoids an + # extra fp32 pair allocation. + with torch.amp.autocast("cuda", enabled=pair.is_cuda, dtype=torch.bfloat16): + pair_delta = self.folding_trunk(pair, pair_attention_mask=pair_mask) + pair.add_(pair_delta.float()) + del pair_delta + single = self.row_attention_pooling(pair, mask) + + atom_mask_f = atom_mask_m.float() + s_at_atoms = gather_token_to_atom(single, atom_to_token_m) + s_at_atoms_ln = self.plddt_ln(s_at_atoms) + + intra_idx = _compute_intra_token_idx(atom_to_token_m) + intra_idx = intra_idx.clamp(max=self.plddt_weight.shape[0] - 1) + w_plddt = self.plddt_weight[intra_idx] + plddt_logits = torch.einsum("...c,...cb->...b", s_at_atoms_ln, w_plddt) + plddt_per_atom = _categorical_mean(plddt_logits, start=0.0, end=1.0) + + sequence_length = single.shape[1] + plddt_sum = torch.zeros( + expanded_batch_size, + sequence_length, + device=single.device, + dtype=plddt_per_atom.dtype, + ) + atom_count = torch.zeros( + expanded_batch_size, + sequence_length, + device=single.device, + dtype=plddt_per_atom.dtype, + ) + atom_mask_t = atom_mask_f.to(plddt_per_atom.dtype) + plddt_sum.scatter_add_(1, atom_to_token_m, plddt_per_atom * atom_mask_t) + atom_count.scatter_add_(1, atom_to_token_m, atom_mask_t) + plddt = plddt_sum / atom_count.clamp(min=1e-6) + + complex_plddt = (plddt_per_atom * atom_mask_f).sum(dim=-1) / ( + atom_mask_f.sum(dim=-1) + _EPS + ) + + expanded_type = self._repeat_batch(mol_type, num_diffusion_samples) + expanded_asym = self._repeat_batch(asym_id, num_diffusion_samples) + is_ligand = (expanded_type == _NONPOLYMER_ID).float() + inter_chain = (expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2)).float() + near_contact = (rep_distances < 8).float() + interface_per_token = (near_contact * inter_chain * (1.0 - is_ligand).unsqueeze(-1)).amax( + dim=-1 + ) + iplddt_weight = torch.where( + is_ligand.bool(), + torch.full_like(interface_per_token, 2.0), + interface_per_token, + ) + iplddt_weight_atoms = gather_token_to_atom( + iplddt_weight.unsqueeze(-1), atom_to_token_m + ).squeeze(-1) + atom_iplddt_w = atom_mask_f * iplddt_weight_atoms + complex_iplddt = (plddt_per_atom * atom_iplddt_w).sum(dim=-1) / ( + atom_iplddt_w.sum(dim=-1) + _EPS + ) + + plddt_ca = plddt_per_atom.gather(1, rep_idx_m) + + # PAE + pae_logits = self.pae_head(self.pae_ln(pair)) + pae = _categorical_mean(pae_logits, start=0.0, end=32.0).detach() + + # PDE + pde_logits = self.pde_head(self.pde_ln(pair)) + pde = _categorical_mean(pde_logits, start=0.0, end=32.0).detach() + + # Resolved (per-atom binary). + s_at_atoms_res = self.resolved_ln(s_at_atoms) + w_res = self.resolved_weight[intra_idx] + resolved_logits = torch.einsum("...c,...cb->...b", s_at_atoms_res, w_res) + + # pTM / ipTM from pae_logits. + n_bins = pae_logits.shape[-1] + bin_width = 32.0 / n_bins + bin_centers = torch.arange(0.5 * bin_width, 32.0, bin_width, device=pae_logits.device) + mask_f = mask.float() + n_residues = mask_f.sum(dim=-1, keepdim=True) + d0 = 1.24 * (n_residues.clamp(min=19) - 15) ** (1 / 3) - 1.8 + tm_per_bin = 1 / (1 + (bin_centers / d0) ** 2) + pae_probs = F.softmax(pae_logits, dim=-1) + tm_expected = (pae_probs * tm_per_bin[:, None, None, :]).sum(dim=-1) + + pair_mask_2d = mask_f.unsqueeze(-1) * mask_f.unsqueeze(-2) + ptm_per_row = (tm_expected * pair_mask_2d).sum(dim=-1) / (pair_mask_2d.sum(dim=-1) + _EPS) + ptm = ptm_per_row.max(dim=-1).values + + inter_chain_mask = ( + expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2) + ).float() * pair_mask_2d + iptm_per_row = (tm_expected * inter_chain_mask).sum(dim=-1) / ( + inter_chain_mask.sum(dim=-1) + _EPS + ) + iptm = iptm_per_row.max(dim=-1).values + + max_chain_id = int(expanded_asym.max().item()) if expanded_batch_size > 0 else 0 + n_chains = max_chain_id + 1 + pair_chains_iptm = torch.zeros( + expanded_batch_size, + n_chains, + n_chains, + device=tm_expected.device, + dtype=tm_expected.dtype, + ) + for c1 in range(n_chains): + chain_c1 = (expanded_asym == c1).float() * mask_f + if chain_c1.sum() == 0: + continue + for c2 in range(n_chains): + chain_c2 = (expanded_asym == c2).float() * mask_f + pair_m = chain_c1.unsqueeze(-1) * chain_c2.unsqueeze(-2) + denom = pair_m.sum(dim=(-1, -2)) + _EPS + pair_chains_iptm[:, c1, c2] = (tm_expected * pair_m).sum(dim=(-1, -2)) / denom + + return { + "plddt_logits": plddt_logits, + "plddt": plddt.detach(), + "plddt_per_atom": plddt_per_atom.detach(), + "plddt_ca": plddt_ca.detach(), + "complex_plddt": complex_plddt.detach(), + "complex_iplddt": complex_iplddt.detach(), + "pae_logits": pae_logits, + "pae": pae, + "pde_logits": pde_logits, + "pde": pde, + "resolved_logits": resolved_logits, + "ptm": ptm.detach(), + "iptm": iptm.detach(), + "pair_chains_iptm": pair_chains_iptm.detach(), + } + + +def _inverse_softplus(value: float) -> float: + return value + math.log(-math.expm1(-value)) + + +def _convert_esmc_attention_outputs_to_te(module: nn.Module) -> tuple[str, ...]: + """Replace the 80 ESMC attention output projections with TE linears. + + Converting every ESMC linear compounds FP8 error across the 80-layer + network. The validated inference path limits FP8 GEMMs to each layer's + attention output projection. Transformer Engine retains canonical BF16 + parameters and creates runtime quantization workspaces during autocast. + """ + + te, _ = _load_transformer_engine() + converted: list[str] = [] + + def walk(owner: nn.Module, prefix: str = "") -> None: + for name, child in tuple(owner.named_children()): + path = f"{prefix}.{name}" if prefix else name + if isinstance(child, nn.Linear) and path.endswith(_ESMC_FP8_LINEAR_SUFFIX): + replacement = te.Linear( + child.in_features, + child.out_features, + bias=child.bias is not None, + params_dtype=child.weight.dtype, + device=child.weight.device, + ) + with torch.no_grad(): + replacement.weight.copy_(child.weight) + if child.bias is not None: + replacement.bias.copy_(child.bias) + replacement.eval().requires_grad_(False) + setattr(owner, name, replacement) + converted.append(path) + else: + walk(child, path) + + walk(module) + if len(converted) != _ESMC_FP8_EXPECTED_PROJECTIONS: + raise RuntimeError( + "ESMC FP8 conversion expected exactly " + f"{_ESMC_FP8_EXPECTED_PROJECTIONS} attention output projections, " + f"found {len(converted)}." + ) + return tuple(converted) + + +@contextmanager +def _lm_precision_context(precision: str, device: torch.device): + """Apply the resolved ESMC inference precision.""" + + if device.type != "cuda" or precision == "fp32": + yield + return + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + if precision == "fp8": + te, recipe = _load_transformer_engine() + fp8_recipe = recipe.Float8CurrentScaling( + use_power_2_scales=False, + fp8_format=recipe.Format.HYBRID, + ) + with te.autocast(enabled=True, recipe=fp8_recipe): + yield + else: + yield + + +class ESMFold2Model( + FastPLMTestTimeTrainingMixin, + ESMFold2EmbeddingMixin, + ESMFold2AttentionMixin, + PreTrainedModel, +): + """ESMFold2: all-atom structure prediction with an ESMC PLM backbone. + + This is the standard released ESMFold2 architecture (uses a linear- + recurrent trunk, internally referred to as "parcae"). + + Forward kwargs that callers commonly override: + + * ``num_loops`` (default ``config.num_loops``): trunk refinement + loops. + * ``num_diffusion_samples`` (default ``config.num_diffusion_samples``): + parallel structure samples; the confidence head re-runs once per + sample, so memory scales linearly. Pass ``1`` for cheap inference. + * ``num_sampling_steps`` (default ``config.structure_head.inference_num_steps``): + diffusion ODE solver steps. Lower for speed, higher for quality. + + Memory / perf knobs: + + * ``model.set_chunk_size(int|None)``: caps l^2 ops (triangle / OPM / + pair transition) at this token-axis chunk. Default 64: fits + l approximately 2k on an 80 GB GPU. Pass ``None`` for faster inference + when l is below 600. + * ``model.set_kernel_backend(None | "fused" | "cuequivariance")``: + select kernel backend (None = reference path). + """ + + config_class = ESMFold2Config + _keys_to_ignore_on_load_unexpected: ClassVar[list[str]] = [r"\._extra_state$"] + + def __init__(self, config: ESMFold2Config) -> None: + super().__init__(config) + d_inputs = config.inputs.d_inputs + d_pair = config.d_pair + + self.inputs_embedder = InputsEmbedder(config) + self.z_init_1 = nn.Linear(d_inputs, d_pair, bias=False) + self.z_init_2 = nn.Linear(d_inputs, d_pair, bias=False) + self.rel_pos = ResIdxAsymIdSymIdEntityIdEncoding( + n_relative_residx_bins=config.n_relative_residx_bins, + n_relative_chain_bins=config.n_relative_chain_bins, + d_pair=d_pair, + ) + self.token_bonds = nn.Linear(1, d_pair, bias=False) + self.language_model = LanguageModelShim( + d_z=d_pair, d_model=config.lm_d_model, num_layers=config.lm_num_layers + ) + self._esmc: nn.Module | None = None + self._esmc_fp8: bool = False + self._esmc_fp8_module_paths: tuple[str, ...] = () + self._esmc_source: str = config.esmc_id + self._esmc_source_revision: str | None = None + self._esmc_source_files: dict[str, str] = {} + self._esmc_local_files_only = False + self._esmc_precision_policy: str = str(getattr(config, "esmc_precision", "auto")) + self._esmc_precision_status = ESMCPrecisionStatus( + requested=self._esmc_precision_policy, + resolved="unloaded", + reason="ESMC has not been loaded.", + device=str(self.device), + transformer_engine_version=_transformer_engine_version(), + ) + self._ttt_lm_head: nn.Module | None = None + self._esmfold2_input_builder: Any | None = None + self._kernel_backend: str | None = None + + pf = config.folding_trunk + self.folding_trunk = FoldingTrunk(n_layers=pf.n_layers, d_pair=d_pair, expansion_ratio=4) + if config.lm_encoder.enabled: + self.lm_encoder: FoldingTrunk | None = FoldingTrunk( + n_layers=config.lm_encoder.n_layers, d_pair=d_pair, expansion_ratio=4 + ) + else: + self.lm_encoder = None + + self.parcae_input_norm = nn.LayerNorm(d_pair) + self.parcae_log_a = nn.Parameter(torch.zeros(d_pair)) + parcae_decay_init = math.sqrt(1.0 / 5.0) + parcae_delta_init = -math.log(parcae_decay_init) + self.parcae_log_delta = nn.Parameter( + torch.full((d_pair,), _inverse_softplus(parcae_delta_init), dtype=torch.float32) + ) + self.parcae_b_cont = nn.Parameter(torch.eye(d_pair)) + self.parcae_readout = nn.Linear(d_pair, d_pair, bias=False) + nn.init.eye_(self.parcae_readout.weight) + self.parcae_coda = FoldingTrunk( + n_layers=config.parcae.coda_n_layers, d_pair=d_pair, expansion_ratio=4 + ) + + # Heads -------------------------------------------------------------- + self.structure_head = DiffusionStructureHead(config) + self.distogram_head = nn.Linear(d_pair, config.structure_head.distogram_bins, bias=True) + self.confidence_head = ConfidenceHead(config) + + msa_cfg = config.msa_encoder + self.msa_encoder = None + if msa_cfg.enabled: + self.msa_encoder = MSAEncoder( + d_msa=msa_cfg.d_msa, + d_pair=d_pair, + d_inputs=d_inputs, + d_hidden=msa_cfg.d_hidden, + n_layers=msa_cfg.n_layers, + n_heads_msa=msa_cfg.n_heads_msa, + msa_head_width=msa_cfg.msa_head_width, + ) + + self.post_init() + self._register_state_dict_hook(_drop_transient_esmc_state) + self.init_ttt({"lora_target_replace_module": "MultiHeadAttention"}) + + @property + def esmc_precision_status(self) -> ESMCPrecisionStatus: + return self._esmc_precision_status + + def load_esmc( + self, + esmc_model_path: str, + precision: ESMCPrecision = "auto", + device: str | torch.device | None = None, + local_files_only: bool = False, + ) -> None: + """Load canonical ESMC weights and resolve the inference precision.""" + + _install_esmc_backbone( + self, + esmc_model_path, + precision=precision, + device=device, + local_files_only=local_files_only, + ) + + def reload_esmc( + self, + precision: ESMCPrecision = "auto", + device: str | torch.device | None = None, + local_files_only: bool | None = None, + ) -> None: + """Reload canonical weights with the requested precision policy.""" + + source = self._esmc_source or self.config.esmc_id + old_esmc = self._esmc + old_head = self._ttt_lm_head + self._esmc = None + self._esmc_fp8 = False + self._esmc_fp8_module_paths = () + self._ttt_lm_head = None + del old_esmc, old_head + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + self.load_esmc( + source, + precision=precision, + device=device, + local_files_only=( + self._esmc_local_files_only + if local_files_only is None + else local_files_only + ), + ) + + def _ensure_ttt_bf16(self) -> None: + if self._esmc_fp8: + _reload_esmc_bf16_for_gradients( + self, + reason="TTT requires BF16; the persisted serving policy is unchanged.", + ) + + def _ensure_ttt_lm_head(self) -> None: + self._ensure_ttt_bf16() + if self._esmc is None: + raise RuntimeError("ESMFold2 TTT requires load_esmc=True.") + if self._ttt_lm_head is not None: + return + from fastplms.models.esm_plusplus.modeling_esm_plusplus import ( + ESMplusplusConfig, + ESMplusplusForMaskedLM, + ) + + source = self._esmc_source or self.config.esmc_id + source_revision = self._esmc_source_revision + if source_revision is None: + source_revision, _ = _manifest_esmc_checkpoint_contract(source) + revision_kwargs: dict[str, Any] = { + "local_files_only": self._esmc_local_files_only, + } + if source_revision is not None: + revision_kwargs["revision"] = source_revision + esmc_config = ESMplusplusConfig.from_pretrained( + source, + **revision_kwargs, + ) + set_config_attn_implementation(esmc_config, get_attn_implementation(self.config)) + mlm, loading_info = ESMplusplusForMaskedLM.from_pretrained( + source, + config=esmc_config, + output_loading_info=True, + **revision_kwargs, + ) + missing_head_keys = [ + key for key in loading_info["missing_keys"] if key.startswith("sequence_head") + ] + if missing_head_keys: + raise RuntimeError( + "ESMFold2 TTT could not load a pretrained ESM++ MLM head from " + f"{source}: missing {missing_head_keys}" + ) + dtype = next(self._esmc.parameters()).dtype + mlm = mlm.to(device=self.device, dtype=dtype).eval() + self._ttt_lm_head = mlm.sequence_head + self._ttt_lm_head.requires_grad_(False) + del mlm + + def _ttt_get_trainable_modules(self) -> list[nn.Module]: + self._ensure_ttt_bf16() + if self._esmc is None: + raise RuntimeError("ESMFold2 TTT requires load_esmc=True.") + return [self._esmc] + + def _ttt_tokenize( + self, + seq: str | list[str] | None = None, + input_ids: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + del kwargs + if input_ids is not None: + return input_ids + if seq is None: + raise ValueError("Pass either seq or input_ids for ESMFold2 TTT.") + sequences = [seq] if isinstance(seq, str) else seq + if not sequences: + raise ValueError("ESMFold2 TTT requires at least one protein sequence.") + token_to_id = {token: idx for idx, token in enumerate(SEQUENCE_VOCAB)} + encoded = [] + for sequence in sequences: + token_ids = [SEQUENCE_BOS_TOKEN] + for amino_acid in sequence: + token_ids.append(token_to_id[amino_acid if amino_acid in token_to_id else "X"]) + token_ids.append(SEQUENCE_EOS_TOKEN) + encoded.append(token_ids) + max_len = max(len(token_ids) for token_ids in encoded) + input_tensor = torch.full( + (len(encoded), max_len), + SEQUENCE_PAD_TOKEN, + dtype=torch.long, + ) + for row, token_ids in enumerate(encoded): + input_tensor[row, : len(token_ids)] = torch.tensor( + token_ids, + dtype=torch.long, + ) + return input_tensor + + def _ttt_mask_token(self) -> int: + return SEQUENCE_MASK_TOKEN + + def _ttt_padding_token(self) -> int: + return SEQUENCE_PAD_TOKEN + + def _ttt_replacement_tokens(self, input_ids: torch.Tensor) -> torch.Tensor: + return torch.arange( + SEQUENCE_STANDARD_AA_MIN_TOKEN, + SEQUENCE_STANDARD_AA_MAX_TOKEN, + device=input_ids.device, + dtype=input_ids.dtype, + ) + + def _ttt_non_special_mask(self, input_ids: torch.Tensor) -> torch.Tensor: + return (input_ids >= SEQUENCE_STANDARD_AA_MIN_TOKEN) & ( + input_ids < SEQUENCE_STANDARD_AA_MAX_TOKEN + ) + + def _ttt_predict_logits( + self, + batch: torch.Tensor | dict[str, torch.Tensor], + **kwargs, + ) -> torch.Tensor: + del kwargs + if not isinstance(batch, torch.Tensor): + raise TypeError("ESMFold2 TTT expects input_ids tensors.") + self._ensure_ttt_bf16() + if self._esmc is None: + raise RuntimeError("ESMFold2 TTT requires load_esmc=True.") + self._ensure_ttt_lm_head() + if self._ttt_lm_head is None: + raise RuntimeError("ESMFold2 TTT MLM head initialization failed.") + attention_mask = batch.ne(SEQUENCE_PAD_TOKEN) + output = self._esmc( + input_ids=batch, + attention_mask=attention_mask, + return_dict=True, + compute_sae=False, + ) + return self._ttt_lm_head(output.last_hidden_state) + + @classmethod + def from_pretrained( + cls, pretrained_model_name_or_path, *args, load_esmc: bool = True, **kwargs + ): + if cls is ESMFold2Model and "config" not in kwargs: + config = ESMFold2Config.from_pretrained(pretrained_model_name_or_path, **kwargs) + if config.type == "experimental": + raise ValueError( + "FastPLMs ESMFold2 supports the released ESMFold2 and " + "ESMFold2-Fast checkpoints. Experimental ESMFold2 configs " + "are not part of the self-contained AutoModel package." + ) + kwargs["config"] = config + # Pop the precision knob before forwarding to the HF loader. + esmc_precision = kwargs.pop("esmc_precision", None) + local_files_only = bool(kwargs.get("local_files_only", False)) + output_loading_info = bool(kwargs.get("output_loading_info", False)) + loaded = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs) + if output_loading_info: + model, loading_info = loaded + else: + model = loaded + if load_esmc: + model.load_esmc( + model.config.esmc_id, + precision=esmc_precision or model.config.esmc_precision, + local_files_only=local_files_only, + ) + return (model, loading_info) if output_loading_info else model + + def set_kernel_backend(self, backend: str | None) -> None: + """Select kernel backend. + + Args: + backend: ``None`` (reference path), ``"fused"`` (requires the + unavailable source-built Triton bundle), or + ``"cuequivariance"`` (requires the ``structure,cueq`` extras + on a supported Linux CUDA 13 host). + """ + validate_kernel_backend(backend) + self.folding_trunk.set_kernel_backend(backend) + if self.lm_encoder is not None: + self.lm_encoder.set_kernel_backend(backend) + self.parcae_coda.set_kernel_backend(backend) + self.confidence_head.set_kernel_backend(backend) + self.structure_head.set_kernel_backend(backend) + self._kernel_backend = backend + + def apply_torch_compile(self, mode: str = "fixed_seqlen", dynamic: bool | None = None) -> None: + """Compile l^2-heavy blocks. + + ``mode='fixed_seqlen'`` recompiles per l; ``'dynamic_seqlen'`` compiles + once. + + Does NOT stack with our Triton kernels: call ``set_kernel_backend(None)`` + before compiling. + """ + if dynamic is None: + dynamic = mode == "dynamic_seqlen" + kwargs: dict = {"dynamic": dynamic} + + from .modeling_esmfold2_common import ( + DiffusionModule, + DiffusionTransformer, + PairUpdateBlock, + ) + + compile_targets = ( + PairUpdateBlock, + DiffusionTransformer, + DiffusionModule, + MSAEncoderBlock, + ) + + def _maybe_compile(module: nn.Module) -> None: + if isinstance(module, compile_targets): + module.forward = torch.compile(module.forward, **kwargs) # type: ignore[assignment] + + self.apply(_maybe_compile) + + def set_chunk_size(self, chunk_size: int | None) -> None: + self.folding_trunk.set_chunk_size(chunk_size) + if self.lm_encoder is not None: + self.lm_encoder.set_chunk_size(chunk_size) + self.parcae_coda.set_chunk_size(chunk_size) + self.confidence_head.set_chunk_size(chunk_size) + if self.msa_encoder is not None: + self.msa_encoder.set_chunk_size(chunk_size) + + def _compute_lm_hidden_states( + self, + input_ids: Tensor, + asym_id: Tensor, + residue_index: Tensor, + mol_type: Tensor, + tok_mask: Tensor, + lm_mask_pct: float = 0.0, + ) -> Tensor: + if self._esmc_fp8 and torch.is_grad_enabled(): + _reload_esmc_bf16_for_gradients( + self, + reason=( + "Gradient-enabled ESMC execution requires BF16; the persisted " + "serving policy is unchanged." + ), + ) + if self._esmc is None: + raise RuntimeError("ESMFold2 requires load_esmc=True for LM feature extraction.") + # Transformer Engine FP8 kernels require l to be a multiple of 16. + pad_to = 16 if self._esmc_fp8 else None + with _lm_precision_context(self._esmc_precision_status.resolved, self.device): + return compute_lm_hidden_states( + self._esmc, + input_ids, + asym_id, + residue_index, + mol_type, + tok_mask, + pad_to_multiple=pad_to, + lm_mask_pct=lm_mask_pct, + mask_token_id=SEQUENCE_MASK_TOKEN, + ) + + def _discretized_dynamics(self) -> tuple[Tensor, Tensor]: + delta = F.softplus(self.parcae_log_delta) + a = torch.exp(-delta * torch.exp(self.parcae_log_a)) + b = delta[:, None] * self.parcae_b_cont + return a, b + + def _init_pair_state(self, ref: Tensor) -> Tensor: + std = math.sqrt(2.0 / (5.0 * ref.shape[-1])) + state = torch.empty_like(ref, dtype=torch.float32) + nn.init.trunc_normal_(state, mean=0.0, std=std, a=-3 * std, b=3 * std) + return state.to(dtype=ref.dtype) + + def _run_one_loop( + self, + z: Tensor, + z_init: Tensor, + lm_z: Tensor | None, + _msa_inputs: dict | None, + pair_mask: Tensor, + a: Tensor, + b_mat: Tensor, + tok_mask: Tensor, + total_steps: int, + ) -> Tensor: + # Helper method (not inline) so per-iter locals free on return: + # otherwise leaks about 2 GB of l^2 * c_z data into distogram/sample scope. + # training=True forces dropout under eval(), matching the per-loop + # dropout strategy used at train time. + lm_cfg = self.config.lm_encoder + _per_loop_lm_dropout = ( + lm_z is not None + and getattr(lm_cfg, "per_loop_lm_dropout", False) + and getattr(lm_cfg, "lm_dropout", 0.0) > 0.0 + ) + _lm_dropout_p = getattr(lm_cfg, "lm_dropout", 0.0) + + for _ in range(total_steps): + if _per_loop_lm_dropout: + if lm_z is None: + raise RuntimeError("Per-loop LM dropout requires LM pair features.") + lm_z_i: Tensor | None = F.dropout(lm_z, p=_lm_dropout_p, training=True) + else: + lm_z_i = lm_z + + refined_lm_z: Tensor | None = None + if lm_z_i is not None and self.lm_encoder is not None: + refined_lm_z = self.lm_encoder( + lm_z_i.to(z_init.dtype), pair_attention_mask=pair_mask + ) + + z_inject_pair = z_init + if lm_z_i is not None and self.lm_encoder is None: + z_inject_pair = z_inject_pair + lm_z_i.to(z_inject_pair.dtype) + + if self.msa_encoder is not None and _msa_inputs is not None: + msa_i, mask_i, hd_i, dv_i = maybe_subsample_msa( + _msa_inputs["msa"], + _msa_inputs["msa_attention_mask"], + _msa_inputs["has_deletion"], + _msa_inputs["deletion_value"], + max_depth=_msa_inputs["max_depth"], + enabled=_msa_inputs["subsample_enabled"], + ) + b_msa, m, l_msa = msa_i.shape + msa_oh = F.one_hot(msa_i.permute(0, 2, 1).long(), num_classes=NUM_RES_TYPES).float() + msa_attn = ( + mask_i.permute(0, 2, 1).float() + if mask_i is not None + else tok_mask[:, :, None].expand(-1, -1, m).float() + ) + # Bias-free MSAEncoder.embed requires zeroed padding. + msa_oh = msa_oh * msa_attn.unsqueeze(-1) + hd = ( + hd_i.permute(0, 2, 1).float() + if hd_i is not None + else torch.zeros(b_msa, l_msa, m, device=msa_i.device) + ) + dv = ( + dv_i.permute(0, 2, 1).float() + if dv_i is not None + else torch.zeros(b_msa, l_msa, m, device=msa_i.device) + ) + msa_pair = self.msa_encoder( + x_pair=z_inject_pair, + x_inputs=_msa_inputs["x_inputs"], + msa_oh=msa_oh, + has_deletion=hd, + deletion_value=dv, + msa_attention_mask=msa_attn, + ).to(z_inject_pair.dtype) + z_inject_pair = ( + msa_pair if self.config.msa_encoder_overwrite else (z_inject_pair + msa_pair) + ) + + if refined_lm_z is not None: + z_inject_pair = z_inject_pair + refined_lm_z.to(z_inject_pair.dtype) + + injected_pair = self.parcae_input_norm(z_inject_pair) + z = a * z + F.linear(injected_pair.to(z.dtype), b_mat) + z = self.folding_trunk(z, pair_attention_mask=pair_mask) + + return z + + def forward( + self, + token_index: Tensor, + residue_index: Tensor, + asym_id: Tensor, + sym_id: Tensor, + entity_id: Tensor, + mol_type: Tensor, + res_type: Tensor, + token_bonds: Tensor, + token_attention_mask: Tensor, + ref_pos: Tensor, + ref_element: Tensor, + ref_charge: Tensor, + ref_atom_name_chars: Tensor, + ref_space_uid: Tensor, + atom_attention_mask: Tensor, + atom_to_token: Tensor, + distogram_atom_idx: Tensor, + deletion_mean: Tensor | None = None, + msa: Tensor | None = None, + has_deletion: Tensor | None = None, + deletion_value: Tensor | None = None, + msa_attention_mask: Tensor | None = None, + input_ids: Tensor | None = None, + lm_hidden_states: Tensor | None = None, + num_loops: int | None = None, + num_diffusion_samples: int | None = None, + num_sampling_steps: int | None = None, + lm_mask_pct: float | None = None, + msa_max_depth: int = 1024, + msa_column_mask_rate: float = 0.1, + msa_subsample_at_inference: bool = True, + early_exit: bool = False, + noise_scale: float | None = None, + step_scale: float | None = None, + max_inference_sigma: float | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + pocket_feature: Tensor | None = None, + gt_coords: Tensor | None = None, + is_resolved: Tensor | None = None, + frames_idx: Tensor | None = None, + disto_cond: Tensor | None = None, + disto_cond_mask: Tensor | None = None, + ) -> ESMFold2Output | tuple[Any, ...]: + output_hidden_states, return_dict = _resolve_structure_output_controls( + self.config, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + validate_msa_conditioning_inputs( + self.config, + msa=msa, + msa_attention_mask=msa_attention_mask, + has_deletion=has_deletion, + deletion_value=deletion_value, + deletion_mean=deletion_mean, + ) + validate_prepared_auxiliary_inputs( + pocket_feature=pocket_feature, + disto_cond=disto_cond, + disto_cond_mask=disto_cond_mask, + ) + del gt_coords, is_resolved, frames_idx + tok_mask = token_attention_mask + atm_mask = atom_attention_mask + disto_idx = distogram_atom_idx + + n_loops: int = num_loops if num_loops is not None else self.config.num_loops + n_samples: int = ( + num_diffusion_samples + if num_diffusion_samples is not None + else self.config.num_diffusion_samples + ) + total_steps = max(1, n_loops + 1) + + if res_type.dim() == 2: + res_type_oh = F.one_hot(res_type.long(), num_classes=NUM_RES_TYPES).float() + res_type_oh = res_type_oh * tok_mask.unsqueeze(-1).float() + else: + res_type_oh = res_type.float() + + if msa is not None: + msa_oh_profile = F.one_hot(msa.long(), num_classes=NUM_RES_TYPES).float() + if msa_attention_mask is not None: + mask_f = msa_attention_mask.float().unsqueeze(-1) + msa_oh_profile = msa_oh_profile * mask_f + valid_seq_count = msa_attention_mask.float().sum(dim=1).clamp(min=1) + profile = msa_oh_profile.sum(dim=1) / valid_seq_count.unsqueeze(-1) + else: + profile = msa_oh_profile.mean(dim=1) + else: + profile = res_type_oh + + if deletion_mean is None: + deletion_mean = torch.zeros( + res_type.shape[0], res_type.shape[1], device=res_type.device + ) + + ref_element_oh = F.one_hot(ref_element.long(), num_classes=MAX_ATOMIC_NUMBER).float() + ref_atom_name_chars_oh = F.one_hot( + ref_atom_name_chars.long(), num_classes=CHAR_VOCAB_SIZE + ).float() + # Bias-free downstream Linears require zeroed padding. + atm_mask_f = atm_mask.float() + ref_element_oh = ref_element_oh * atm_mask_f.unsqueeze(-1) + ref_atom_name_chars_oh = ref_atom_name_chars_oh * atm_mask_f.unsqueeze(-1).unsqueeze(-1) + atom_to_token = atom_to_token * atm_mask.long() + + use_amp = ref_pos.device.type == "cuda" + with torch.amp.autocast("cuda", enabled=use_amp, dtype=torch.bfloat16): + x_inputs = self.inputs_embedder( + aatype=res_type_oh, + profile=profile.float(), + deletion_mean=deletion_mean.float(), + ref_pos=ref_pos, + atom_attention_mask=atm_mask, + ref_space_uid=ref_space_uid, + ref_charge=ref_charge, + ref_element=ref_element_oh, + ref_atom_name_chars=ref_atom_name_chars_oh, + atom_to_token=atom_to_token, + ) + + z_init = self.z_init_1(x_inputs).unsqueeze(2) + self.z_init_2(x_inputs).unsqueeze(1) + + relative_position_encoding = self.rel_pos( + residue_index=residue_index, + asym_id=asym_id, + sym_id=sym_id, + entity_id=entity_id, + token_index=token_index, + ) + token_bonds_encoding = self.token_bonds(token_bonds.float()) + z_init = z_init + relative_position_encoding + token_bonds_encoding + + if lm_hidden_states is None and input_ids is not None and self._esmc is not None: + lm_hidden_states = self._compute_lm_hidden_states( + input_ids, + asym_id, + residue_index, + mol_type, + tok_mask, + lm_mask_pct=(self.config.lm_mask_pct if lm_mask_pct is None else lm_mask_pct), + ) + lm_z: Tensor | None = None + if lm_hidden_states is not None: + lm_z = self.language_model(lm_hidden_states.detach()) + del lm_hidden_states + + pair_mask = tok_mask[:, :, None].float() * tok_mask[:, None, :].float() + + z = self._init_pair_state(z_init) + + a, b = self._discretized_dynamics() + a = a.view(1, 1, 1, -1).to(device=z.device, dtype=z.dtype) + b_mat = b.to(device=z.device, dtype=z.dtype) + + _msa_inputs: dict | None = None + if self.msa_encoder is not None and msa is not None: + msa_attention_mask = maybe_apply_msa_column_masking( + msa_attention_mask, + msa_column_mask_rate, + ) + _msa_inputs = dict( + x_inputs=x_inputs, + msa=msa, + msa_attention_mask=msa_attention_mask, + has_deletion=has_deletion, + deletion_value=deletion_value, + max_depth=msa_max_depth, + subsample_enabled=msa_subsample_at_inference, + ) + + # Method call (not inline loop) frees per-iteration l^2 * c_z locals. + z = self._run_one_loop( + z=z, + z_init=z_init, + lm_z=lm_z, + _msa_inputs=_msa_inputs, + pair_mask=pair_mask, + a=a, + b_mat=b_mat, + tok_mask=tok_mask, + total_steps=total_steps, + ) + del z_init, lm_z, _msa_inputs, a, b_mat + + z = self.parcae_readout(z) + z = self.parcae_coda(z, pair_attention_mask=pair_mask) + + z = z.float() + distogram_logits = self.distogram_head(z + z.transpose(-2, -3)) + + structure_output = self.structure_head.sample( + z_trunk=z, + s_inputs=x_inputs, + s_trunk=None, + relative_position_encoding=relative_position_encoding, + ref_pos=ref_pos, + ref_charge=ref_charge, + ref_mask=atm_mask, + ref_element=ref_element_oh, + ref_atom_name_chars=ref_atom_name_chars_oh, + ref_space_uid=ref_space_uid, + tok_idx=atom_to_token, + asym_id=asym_id, + residue_index=residue_index, + entity_id=entity_id, + token_index=token_index, + sym_id=sym_id, + token_attention_mask=tok_mask, + num_diffusion_samples=n_samples, + num_sampling_steps=num_sampling_steps, + max_inference_sigma=max_inference_sigma, + noise_scale=noise_scale, + step_scale=step_scale, + return_atom_repr=False, + denoising_early_exit_rmsd=(0.10 if early_exit else None), + ) + + sample_coords = structure_output["sample_atom_coords"] + if sample_coords is None: + raise RuntimeError("ESMFold2 structure sampling did not return coordinates.") + output: dict[str, Tensor] = {"distogram_logits": distogram_logits} + output["sample_atom_coords"] = sample_coords + + confidence_output = self.confidence_head( + s_inputs=x_inputs.detach(), + z=z.detach().float(), + x_pred=sample_coords.detach(), + distogram_atom_idx=disto_idx, + token_attention_mask=tok_mask, + atom_to_token=atom_to_token, + atom_attention_mask=atm_mask, + asym_id=asym_id, + mol_type=mol_type, + num_diffusion_samples=n_samples, + relative_position_encoding=relative_position_encoding.detach(), + token_bonds_encoding=token_bonds_encoding.detach(), + ) + output.update(confidence_output) + output["atom_pad_mask"] = atm_mask.unsqueeze(0) if atm_mask.dim() == 1 else atm_mask + output["residue_index"] = residue_index + output["entity_id"] = entity_id + return _finalize_structure_output( + output, + token_input_state=x_inputs, + pair_state=z, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + @torch.no_grad() + def infer_protein(self, seq: str, **forward_kwargs) -> ESMFold2Output: + from .protein_utils import prepare_protein_features + + if forward_kwargs.pop("return_dict", True) is not True: + raise ValueError( + "infer_protein always returns a mapping; return_dict=False is invalid." + ) + features = prepare_protein_features(seq) + if not self.config.msa_conditioning: + for name in MSA_CONDITIONING_INPUT_NAMES: + features.pop(name, None) + features = {k: v.to(self.device) for k, v in features.items()} + return self(**features, **forward_kwargs, return_dict=True) + + @property + def input_builder(self): + if self._esmfold2_input_builder is None: + from .esmfold2_processor import ESMFold2InputBuilder + + self._esmfold2_input_builder = ESMFold2InputBuilder() + return self._esmfold2_input_builder + + @property + def input_types(self): + from . import esmfold2_types + + return esmfold2_types + + def prepare_structure_input(self, input, seed: int | None = None): + return self.input_builder.prepare_model_input( + self, + input, + seed=seed, + device=self.device, + ) + + def fold( + self, + input, + *, + num_loops: int = 3, + num_sampling_steps: int = 50, + num_diffusion_samples: int = 1, + seed: int | None = None, + noise_scale: float | None = None, + step_scale: float | None = None, + max_inference_sigma: int | None = None, + early_exit: bool = False, + complex_id: str = "pred", + ): + return self.input_builder.fold( + self, + input, + num_loops=num_loops, + num_sampling_steps=num_sampling_steps, + num_diffusion_samples=num_diffusion_samples, + seed=seed, + noise_scale=noise_scale, + step_scale=step_scale, + max_inference_sigma=max_inference_sigma, + early_exit=early_exit, + complex_id=complex_id, + ) + + def _fold_protein_no_ttt( + self, + sequence: str, + *, + chain_id: str = "A", + msa: Any | None = None, + msa_path: str | Path | None = None, + msa_max_sequences: int | None = None, + num_loops: int = 3, + num_sampling_steps: int = 50, + num_diffusion_samples: int = 1, + seed: int | None = None, + complex_id: str = "pred", + ): + from .esmfold2_types import MSA, ProteinInput, StructurePredictionInput + + if msa is not None and msa_path is not None: + raise ValueError("Pass at most one of msa or msa_path.") + if msa_path is not None: + msa = MSA.from_a3m(msa_path, max_sequences=msa_max_sequences) + if msa is not None: + query = str(msa.query).replace("-", "").upper() + if query != sequence.upper(): + raise ValueError( + "MSA query does not match sequence: " + f"expected {sequence.upper()!r}, got {query!r}" + ) + + input = StructurePredictionInput( + sequences=[ProteinInput(id=chain_id, sequence=sequence, msa=msa)] + ) + return self.fold( + input, + num_loops=num_loops, + num_sampling_steps=num_sampling_steps, + num_diffusion_samples=num_diffusion_samples, + seed=seed, + complex_id=complex_id, + ) + + @staticmethod + def _ttt_mean_plddt(result) -> float: + if result.plddt is None: + raise RuntimeError("ESMFold2 result has no pLDDT tensor.") + return float(result.plddt.float().mean().item()) + + def _ttt_select_result(self, result): + if isinstance(result, list): + if not result: + raise RuntimeError("ESMFold2 fold returned an empty result list.") + return max(result, key=self._ttt_mean_plddt) + return result + + def _ttt_eval_step( + self, + step: int, + loss: float, + seq: str | list[str] | None = None, + input_ids: torch.Tensor | None = None, + **kwargs, + ) -> tuple[dict[str, Any], float | None]: + del input_ids + if not isinstance(seq, str): + raise TypeError("ESMFold2 fold TTT is protein-only and sequence-string only.") + fold_kwargs = kwargs["fold_kwargs"] + was_training = self.training + self.eval() + try: + result = self._fold_protein_no_ttt(seq, **fold_kwargs) + finally: + self.train(was_training) + selected = self._ttt_select_result(result) + plddt = self._ttt_mean_plddt(selected) + return { + "step": step, + "loss": loss, + "plddt": plddt, + "result": selected, + }, plddt + + def fold_protein( + self, + sequence: str, + *, + chain_id: str = "A", + msa: Any | None = None, + msa_path: str | Path | None = None, + msa_max_sequences: int | None = None, + num_loops: int = 3, + num_sampling_steps: int = 50, + num_diffusion_samples: int = 1, + seed: int | None = None, + complex_id: str = "pred", + ttt: bool = False, + ttt_config: TTTConfig | dict[str, Any] | None = None, + ): + if ttt: + return self.fold_protein_ttt( + sequence=sequence, + chain_id=chain_id, + msa=msa, + msa_path=msa_path, + msa_max_sequences=msa_max_sequences, + num_loops=num_loops, + num_sampling_steps=num_sampling_steps, + num_diffusion_samples=num_diffusion_samples, + seed=seed, + complex_id=complex_id, + ttt_config=ttt_config, + ) + return self._fold_protein_no_ttt( + sequence=sequence, + chain_id=chain_id, + msa=msa, + msa_path=msa_path, + msa_max_sequences=msa_max_sequences, + num_loops=num_loops, + num_sampling_steps=num_sampling_steps, + num_diffusion_samples=num_diffusion_samples, + seed=seed, + complex_id=complex_id, + ) + + def fold_protein_ttt( + self, + sequence: str, + *, + chain_id: str = "A", + msa: Any | None = None, + msa_path: str | Path | None = None, + msa_max_sequences: int | None = None, + num_loops: int = 3, + num_sampling_steps: int = 50, + num_diffusion_samples: int = 1, + seed: int | None = None, + complex_id: str = "pred", + ttt_config: TTTConfig | dict[str, Any] | None = None, + ): + self._ensure_ttt_bf16() + if self._esmc is None: + raise RuntimeError("ESMFold2 TTT requires load_esmc=True.") + fold_kwargs = { + "chain_id": chain_id, + "msa": msa, + "msa_path": msa_path, + "msa_max_sequences": msa_max_sequences, + "num_loops": num_loops, + "num_sampling_steps": num_sampling_steps, + "num_diffusion_samples": num_diffusion_samples, + "seed": seed, + "complex_id": complex_id, + } + baseline = self._ttt_select_result(self._fold_protein_no_ttt(sequence, **fold_kwargs)) + baseline_plddt = self._ttt_mean_plddt(baseline) + best_result = baseline + best_plddt = baseline_plddt + best_step = 0 + step_plddts = [baseline_plddt] + + cfg = self.ttt_config.merged(ttt_config).merged( + {"eval_each_step": True, "automatic_best_state_reset": False} + ) + try: + metrics = self.ttt( + seq=sequence, + ttt_config=cfg, + fold_kwargs=fold_kwargs, + ) + for step_metric in metrics["step_metrics"]: + step_plddt = step_metric["plddt"] + step_plddts.append(step_plddt) + if step_plddt > best_plddt: + best_plddt = step_plddt + best_step = step_metric["step"] + best_result = step_metric["result"] + best_result.ttt_metrics = { + "losses": metrics["losses"], + "step_plddts": step_plddts, + "baseline_plddt": baseline_plddt, + "best_plddt": best_plddt, + "best_step": best_step, + } + return best_result + finally: + if "_ttt_initialized" in self.__dict__ and self._ttt_initialized: + self.ttt_reset() + + @staticmethod + def result_to_cif(result) -> str: + if isinstance(result, list): + raise TypeError("Pass one MolecularComplexResult at a time.") + return result.complex.to_mmcif() + + @staticmethod + def result_to_pdb(result) -> str: + if isinstance(result, list): + raise TypeError("Pass one MolecularComplexResult at a time.") + return result.complex.to_protein_complex().to_pdb_string() + + def save_as_cif(self, result, output_path: str | Path) -> None: + Path(output_path).write_text(self.result_to_cif(result)) + + def save_as_pdb(self, result, output_path: str | Path) -> None: + Path(output_path).write_text(self.result_to_pdb(result)) + + def infer_protein_as_cif(self, seq: str, **forward_kwargs) -> str: + return self.result_to_cif(self.fold_protein(seq, **forward_kwargs)) + + def infer_protein_as_pdb(self, seq: str, **forward_kwargs) -> str: + return self.result_to_pdb(self.fold_protein(seq, **forward_kwargs)) + + +class MSAEncoderBlock(nn.Module): + """One MSA encoder block: OPM into pair, MSA pair-weighted averaging, triangle update.""" + + def __init__( + self, + d_msa: int, + d_pair: int, + d_hidden: int, + n_heads_msa: int, + msa_head_width: int, + is_final_block: bool = False, + ) -> None: + super().__init__() + self.is_final_block = is_final_block + self.outer_product_mean = OuterProductMean(d_msa, d_hidden, d_pair) + if not is_final_block: + self.msa_pair_weighted_averaging = MSAPairWeightedAveraging( + d_msa, d_pair, n_heads_msa, msa_head_width + ) + self.msa_transition = PairTransition(d_msa, expansion_ratio=4) + self.tri_mul_out = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=True) + self.tri_mul_in = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=False) + self.pair_transition = PairTransition(d_pair, expansion_ratio=4) + + def set_chunk_size(self, chunk_size: int | None) -> None: + self.outer_product_mean.set_chunk_size(chunk_size) + self.tri_mul_out.set_chunk_size(chunk_size) + self.tri_mul_in.set_chunk_size(chunk_size) + if not self.is_final_block: + self.msa_transition.set_chunk_size(chunk_size) + self.pair_transition.set_chunk_size(chunk_size) + + def forward( + self, + m: Tensor, + pair: Tensor, + msa_attention_mask: Tensor, + pair_attention_mask: Tensor, + ) -> tuple[Tensor, Tensor]: + pair = pair + self.outer_product_mean(m, msa_attention_mask) + if not self.is_final_block: + m = m + self.msa_pair_weighted_averaging(m, pair, pair_attention_mask) + m = m + self.msa_transition(m) + pair = pair + self.tri_mul_out(pair, mask=pair_attention_mask) + pair = pair + self.tri_mul_in(pair, mask=pair_attention_mask) + pair = pair + self.pair_transition(pair) + return m, pair + + +class MSAEncoder(nn.Module): + """Stack of [`MSAEncoderBlock`] layers that conditions the pair on an MSA.""" + + def __init__( + self, + d_msa: int, + d_pair: int, + d_inputs: int, + d_hidden: int = 32, + n_layers: int = 4, + n_heads_msa: int = 8, + msa_head_width: int = 16, + ) -> None: + super().__init__() + self.embed = nn.Linear(35, d_msa, bias=False) + self.project_inputs = nn.Linear(d_inputs, d_msa, bias=False) + self.blocks = nn.ModuleList( + [ + MSAEncoderBlock( + d_msa=d_msa, + d_pair=d_pair, + d_hidden=d_hidden, + n_heads_msa=n_heads_msa, + msa_head_width=msa_head_width, + is_final_block=(i == n_layers - 1), + ) + for i in range(n_layers) + ] + ) + + def set_chunk_size(self, chunk_size: int | None) -> None: + for block in self.blocks: + cast(MSAEncoderBlock, block).set_chunk_size(chunk_size) + + def forward( + self, + x_pair: Tensor, + x_inputs: Tensor, + msa_oh: Tensor, + has_deletion: Tensor, + deletion_value: Tensor, + msa_attention_mask: Tensor, + ) -> Tensor: + # Every input tensor is pre-transposed to shape (b, l, m, ...) before this call. + m_feat = torch.cat( + [msa_oh, has_deletion.unsqueeze(-1), deletion_value.unsqueeze(-1)], dim=-1 + ) + m = self.embed(m_feat) + self.project_inputs(x_inputs).unsqueeze(2) + tok_mask = msa_attention_mask[:, :, 0].bool() + pair_attention_mask = tok_mask.unsqueeze(2) & tok_mask.unsqueeze(1) + for block in self.blocks: + m, x_pair = block(m, x_pair, msa_attention_mask, pair_attention_mask) + return x_pair diff --git a/fastplms/models/esmfold2/modeling_esmfold2_classification.py b/fastplms/models/esmfold2/modeling_esmfold2_classification.py new file mode 100644 index 0000000000000000000000000000000000000000..40ef916274039d645951ce86268e70381a10e588 --- /dev/null +++ b/fastplms/models/esmfold2/modeling_esmfold2_classification.py @@ -0,0 +1,230 @@ +"""Sequence and residue prediction heads for ESMFold2 checkpoints.""" + +from __future__ import annotations + +from typing import Any, Literal + +import torch +import torch.nn as nn +from torch import Tensor + +from ..classification_probe import SequenceClassificationProbe, TokenClassificationProbe +from .configuration_esmfold2 import ESMFold2Config +from .embedding import _TOKEN_TO_ID, _VALID_RESIDUES, _encode_single_chain +from .esmfold2_constants_esm3 import SEQUENCE_PAD_TOKEN +from .modeling_esmfold2 import ESMFold2Model +from .modeling_esmfold2_experimental import ESMFold2ExperimentalModel + + +ClassifierTrainScope = Literal["probe", "projection"] +_VALID_RESIDUE_IDS = frozenset(_TOKEN_TO_ID[residue] for residue in _VALID_RESIDUES) + + +class _ESMFold2ClassificationMixin: + """Run a task probe on the checkpoint-owned ESMC sequence projection.""" + + _classifier_config_type = "" + _keys_to_ignore_on_load_unexpected: list[str] = [r"\._extra_state$"] + + def _initialize_classifier(self, classifier: nn.Module) -> None: + self.requires_grad_(False) + self.classifier = classifier + self.set_classifier_train_scope(self.config.classifier_train_scope) + + def set_classifier_train_scope(self, scope: ClassifierTrainScope) -> None: + """Select whether fine-tuning updates only the probe or its input projection too.""" + + if scope not in {"probe", "projection"}: + raise ValueError( + "classifier_train_scope must be 'probe' or 'projection', " + f"got {scope!r}." + ) + self.config.classifier_train_scope = scope + self.requires_grad_(False) + self.classifier.requires_grad_(True) + if scope == "projection": + self.language_model.base_z_combine.requires_grad_(True) + self.language_model.base_z_linear.requires_grad_(True) + if self._esmc is not None: + self._esmc.requires_grad_(False) + + def load_esmc(self, *args: Any, **kwargs: Any) -> None: + super().load_esmc(*args, **kwargs) + if self._esmc is None: + raise RuntimeError("ESMFold2 ESMC loading completed without a backbone.") + self._esmc.requires_grad_(False) + + def train(self, mode: bool = True): + super().train(mode) + if self._esmc is not None: + self._esmc.eval() + return self + + @classmethod + def from_pretrained(cls, pretrained_model_name_or_path, *model_args: Any, **kwargs: Any): + if "config" not in kwargs: + kwargs["config"] = ESMFold2Config.from_pretrained( + pretrained_model_name_or_path, **kwargs + ) + config = kwargs["config"] + if not isinstance(config, ESMFold2Config): + raise TypeError("ESMFold2 classifiers require an ESMFold2Config.") + if config.type != cls._classifier_config_type: + raise ValueError( + f"{cls.__name__} requires config.type={cls._classifier_config_type!r}, " + f"got {config.type!r}." + ) + loaded = super().from_pretrained( + pretrained_model_name_or_path, *model_args, **kwargs + ) + model = loaded[0] if isinstance(loaded, tuple) else loaded + model.set_classifier_train_scope(model.config.classifier_train_scope) + return loaded + + def prepare_classifier_inputs( + self, sequence_or_sequences: str | list[str] | tuple[str, ...] + ) -> dict[str, Tensor]: + """Encode one or more ungapped single-chain proteins without special tokens.""" + + sequences = ( + [sequence_or_sequences] + if isinstance(sequence_or_sequences, str) + else list(sequence_or_sequences) + ) + if not sequences: + raise ValueError("prepare_classifier_inputs requires at least one sequence.") + encoded = [_encode_single_chain(sequence) for sequence in sequences] + sequence_length = max(map(len, encoded)) + input_ids = torch.full( + (len(encoded), sequence_length), + SEQUENCE_PAD_TOKEN, + dtype=torch.long, + device=self.device, + ) + attention_mask = torch.zeros_like(input_ids, dtype=torch.bool) + for batch_index, token_ids in enumerate(encoded): + length = len(token_ids) + input_ids[batch_index, :length] = torch.tensor( + token_ids, dtype=torch.long, device=self.device + ) + attention_mask[batch_index, :length] = True + return {"input_ids": input_ids, "attention_mask": attention_mask} + + def _classifier_embeddings( + self, input_ids: Tensor, attention_mask: Tensor | None + ) -> tuple[Tensor, Tensor]: + if input_ids.ndim != 2: + raise ValueError( + "ESMFold2 classifier input_ids must have shape (batch, residue), " + f"got {tuple(input_ids.shape)}." + ) + if attention_mask is None: + attention_mask = input_ids.ne(SEQUENCE_PAD_TOKEN) + elif attention_mask.shape != input_ids.shape: + raise ValueError( + "ESMFold2 classifier attention_mask must match input_ids, got " + f"{tuple(attention_mask.shape)} and {tuple(input_ids.shape)}." + ) + residue_mask = attention_mask.to(device=input_ids.device, dtype=torch.bool) + if not residue_mask.any(dim=1).all(): + raise ValueError("Every ESMFold2 classifier input must contain a protein residue.") + if input_ids.masked_select(residue_mask).eq(SEQUENCE_PAD_TOKEN).any(): + raise ValueError("ESMFold2 classifier padding tokens cannot be attended residues.") + residue_ids = input_ids.masked_select(residue_mask) + valid_residue_ids = torch.tensor( + sorted(_VALID_RESIDUE_IDS), dtype=input_ids.dtype, device=input_ids.device + ) + if not torch.isin(residue_ids, valid_residue_ids).all(): + raise ValueError( + "ESMFold2 classifiers accept residue-only single-chain protein inputs." + ) + + batch_size, sequence_length = input_ids.shape + residue_index = torch.arange(sequence_length, device=input_ids.device).expand( + batch_size, -1 + ) + asym_id = torch.zeros_like(input_ids) + mol_type = torch.zeros_like(input_ids) + with torch.no_grad(): + hidden_states = self._compute_lm_hidden_states( + input_ids, + asym_id, + residue_index, + mol_type, + residue_mask, + ) + embeddings = self.project_esmc_hidden_states(hidden_states, residue_mask) + return embeddings, residue_mask + + def _classifier_forward( + self, + input_ids: Tensor, + attention_mask: Tensor | None = None, + labels: Tensor | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + ): + embeddings, residue_mask = self._classifier_embeddings(input_ids, attention_mask) + return self.classifier( + embeddings, + attention_mask=residue_mask, + labels=labels, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + +class _ESMFold2SequenceClassificationMixin(_ESMFold2ClassificationMixin): + def __init__(self, config: ESMFold2Config) -> None: + super().__init__(config) + self._initialize_classifier(SequenceClassificationProbe(config, config.d_pair)) + + forward = _ESMFold2ClassificationMixin._classifier_forward + + +class _ESMFold2TokenClassificationMixin(_ESMFold2ClassificationMixin): + def __init__(self, config: ESMFold2Config) -> None: + super().__init__(config) + self._initialize_classifier(TokenClassificationProbe(config, config.d_pair)) + + forward = _ESMFold2ClassificationMixin._classifier_forward + + +class ESMFold2ForSequenceClassification( + _ESMFold2SequenceClassificationMixin, ESMFold2Model +): + """Released ESMFold2 with a sequence classification or regression probe.""" + + _classifier_config_type = "release" + + +class ESMFold2ForTokenClassification(_ESMFold2TokenClassificationMixin, ESMFold2Model): + """Released ESMFold2 with a residue classification or regression probe.""" + + _classifier_config_type = "release" + + +class ESMFold2ExperimentalForSequenceClassification( + _ESMFold2SequenceClassificationMixin, ESMFold2ExperimentalModel +): + """Experimental ESMFold2 with a sequence classification or regression probe.""" + + _classifier_config_type = "experimental" + + +class ESMFold2ExperimentalForTokenClassification( + _ESMFold2TokenClassificationMixin, ESMFold2ExperimentalModel +): + """Experimental ESMFold2 with a residue classification or regression probe.""" + + _classifier_config_type = "experimental" + + +__all__ = [ + "ESMFold2ExperimentalForSequenceClassification", + "ESMFold2ExperimentalForTokenClassification", + "ESMFold2ForSequenceClassification", + "ESMFold2ForTokenClassification", +] diff --git a/fastplms/models/esmfold2/modeling_esmfold2_common.py b/fastplms/models/esmfold2/modeling_esmfold2_common.py new file mode 100644 index 0000000000000000000000000000000000000000..40e1b9aa46d9ca3704d28e5f65310d5cc4f84754 --- /dev/null +++ b/fastplms/models/esmfold2/modeling_esmfold2_common.py @@ -0,0 +1,2747 @@ +# Copyright 2026 Biohub. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +"""Shared building blocks for ESMFold2 HuggingFace model variants.""" + +from __future__ import annotations + +import importlib +from functools import partial +from importlib.util import find_spec +from typing import ClassVar, cast + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from torch.utils.checkpoint import checkpoint + +from .configuration_esmfold2 import ESMFold2Config +from .reproducibility import seed_context + +_seed_context = seed_context + +try: + if find_spec("cuequivariance_ops_torch") is None: + raise ImportError("cuequivariance_ops_torch is unavailable") + cue_module = importlib.import_module("cuequivariance_torch") + _cue_attn_pair_bias = cue_module.attention_pair_bias + _cue_tri_mul = cue_module.triangle_multiplicative_update + + CUE_AVAILABLE = True +except (AttributeError, ImportError): + _cue_attn_pair_bias = None # type: ignore[assignment] + _cue_tri_mul = None # type: ignore[assignment] + CUE_AVAILABLE = False + +# Biohub ships optional source-built Triton helpers. FastPLMs does not bundle or +# compile them; these placeholders retain checkpoint-compatible control flow +# while the portable PyTorch path remains flat and self-contained. +_fused_pair_bias = None +_fused_trimul_with_residual = None +_FusedLNLinearSwiGLU = None +_FusedDropoutResidual = None +TRITON_KERNELS_AVAILABLE = False + +BACKEND_FUSED = "fused" +BACKEND_CUEQ = "cuequivariance" +_VALID_BACKENDS = (None, BACKEND_FUSED, BACKEND_CUEQ) +MSA_CONDITIONING_INPUT_NAMES = ( + "msa", + "msa_attention_mask", + "has_deletion", + "deletion_value", + "deletion_mean", +) +PREPARED_AUXILIARY_INPUT_NAMES = ( + "pocket_feature", + "gt_coords", + "is_resolved", + "frames_idx", + "disto_cond", + "disto_cond_mask", +) + + +def validate_kernel_backend(backend: str | None) -> None: + """Fail before mutating modules when a named kernel cannot execute.""" + + if backend not in _VALID_BACKENDS: + raise ValueError(f"backend must be one of {_VALID_BACKENDS}, got {backend!r}") + if backend == BACKEND_FUSED and not TRITON_KERNELS_AVAILABLE: + raise RuntimeError( + "backend='fused' is unavailable because FastPLMs does not bundle the " + "source-built ESMFold2 Triton kernels." + ) + if backend == BACKEND_CUEQ and not CUE_AVAILABLE: + raise RuntimeError( + "backend='cuequivariance' requires cuequivariance_torch and the CUDA 13 " + "cuequivariance_ops_torch runtime. Install FastPLMs with the " + "'structure,cueq' extras on a supported Linux CUDA 13 host." + ) + + +def validate_msa_conditioning_inputs( + config: ESMFold2Config, + *, + msa: Tensor | None, + msa_attention_mask: Tensor | None, + has_deletion: Tensor | None, + deletion_value: Tensor | None, + deletion_mean: Tensor | None, +) -> None: + """Reject MSA-derived tensors for checkpoints trained without MSA conditioning.""" + + if config.msa_conditioning: + return + values = { + "msa": msa, + "msa_attention_mask": msa_attention_mask, + "has_deletion": has_deletion, + "deletion_value": deletion_value, + "deletion_mean": deletion_mean, + } + provided = sorted(name for name, value in values.items() if value is not None) + if provided: + raise ValueError( + "This ESMFold2 checkpoint was trained without MSA conditioning and rejects " + f"MSA-derived inputs: {', '.join(provided)}." + ) + + +def validate_prepared_auxiliary_inputs( + *, + pocket_feature: Tensor | None, + disto_cond: Tensor | None, + disto_cond_mask: Tensor | None, +) -> None: + """Accept inert prepared fields and reject unsupported conditioning.""" + + if pocket_feature is not None and torch.any(pocket_feature != 0).item(): + raise NotImplementedError( + "The published ESMFold2 forward does not consume pocket conditioning; " + "nonzero pocket_feature values are unsupported." + ) + distogram_is_active = (disto_cond is not None and torch.any(disto_cond != 0).item()) or ( + disto_cond_mask is not None and torch.any(disto_cond_mask).item() + ) + if distogram_is_active: + raise NotImplementedError( + "The published ESMFold2 forward does not consume distogram conditioning; " + "nonzero disto_cond or disto_cond_mask values are unsupported." + ) + + +def _fused_active(module: nn.Module, tensor: Tensor) -> bool: + """Return whether an optional fused implementation can handle this call.""" + return ( + TRITON_KERNELS_AVAILABLE + and getattr(module, "_kernel_backend", None) == BACKEND_FUSED + and not torch.is_grad_enabled() + and tensor.is_cuda + ) + + +def _cueq_active(module: nn.Module) -> bool: + return CUE_AVAILABLE and getattr(module, "_kernel_backend", None) == BACKEND_CUEQ + + +class DropoutResidual(nn.Module): + """``residual + dropout(delta)`` with row/col-shared dropout. + + Same signature on both paths. ``use_fused_kernels=True`` + ``batch_dim=1`` + routes through ``FusedDropoutResidual`` (single-pass over pair tensor, + in-place residual add). Falls back to unfused otherwise. + """ + + def __init__(self, r: float, batch_dim: int, use_fused_kernels: bool = False) -> None: + super().__init__() + if isinstance(batch_dim, bool) or batch_dim not in (1, 2): + raise ValueError(f"batch_dim must be 1 or 2, got {batch_dim}") + self._use_fused_kernels = ( + use_fused_kernels and batch_dim == 1 and _FusedDropoutResidual is not None + ) + self._batch_dim = batch_dim + self._r = r + if self._use_fused_kernels: + assert _FusedDropoutResidual is not None + self._impl: nn.Module = _FusedDropoutResidual(r) + else: + self._impl = nn.Dropout(r) + + def forward(self, residual: Tensor, delta: Tensor) -> Tensor: + if self._use_fused_kernels: + return self._impl(residual, delta) + # The unfused path broadcasts a row/column-shared mask M with shape (1, ...). + if self._r == 0.0 or not self.training: + return residual + delta + shape = list(delta.shape) + shape[self._batch_dim] = 1 + mask = self._impl(delta.new_ones(shape)) + return residual + delta * mask + + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- +CHAR_VOCAB_SIZE: int = 64 +MAX_CHARS: int = 4 +XYZ_DIMS: int = 3 +MAX_ATOMIC_NUMBER: int = 128 + +# Input feature dim = 3 + 1 + 1 + 128 + 64*4 = 389 +ATOM_FEATURE_DIM: int = XYZ_DIMS + 1 + 1 + MAX_ATOMIC_NUMBER + CHAR_VOCAB_SIZE * MAX_CHARS + + +NUM_RES_TYPES: int = 33 + +_EPS = 1e-5 + +# Default for the quadratic triangle, OPM, and pair-transition operations. +# It caps peak memory so l around 2,000 fits on an 80 GiB GPU. At l=1,438, +# chunk=128 uses roughly 76 GiB, while chunk=64 leaves headroom for the +# largest foldbench targets. Pass None to disable chunking; this is faster for +# short sequences but prone to out-of-memory errors beyond l around 600. +_DEFAULT_CHUNK_SIZE = 64 + + +# =========================================================================== +# MSA inference-time diversity augmentations +# =========================================================================== + + +def maybe_subsample_msa( + msa: Tensor, + msa_attention_mask: Tensor | None, + has_deletion: Tensor | None, + deletion_value: Tensor | None, + *, + max_depth: int | None, + enabled: bool, +) -> tuple[Tensor, Tensor | None, Tensor | None, Tensor | None]: + if not enabled or max_depth is None: + return msa, msa_attention_mask, has_deletion, deletion_value + + depth = msa.size(1) + if depth <= 1 or depth <= max_depth: + return msa, msa_attention_mask, has_deletion, deletion_value + + indices = torch.zeros(max_depth, dtype=torch.long, device=msa.device) + indices[1:] = torch.randperm(depth - 1, device=msa.device)[: max_depth - 1] + 1 + indices = indices.sort().values + + msa = msa[:, indices] + if msa_attention_mask is not None: + msa_attention_mask = msa_attention_mask[:, indices] + if has_deletion is not None: + has_deletion = has_deletion[:, indices] + if deletion_value is not None: + deletion_value = deletion_value[:, indices] + return msa, msa_attention_mask, has_deletion, deletion_value + + +def maybe_apply_msa_column_masking( + msa_attention_mask: Tensor | None, + rate: float, +) -> Tensor | None: + if msa_attention_mask is None or rate <= 0.0 or msa_attention_mask.size(1) <= 1: + return msa_attention_mask + + batch_size, _, length = msa_attention_mask.shape + col_keep = torch.rand(batch_size, length, device=msa_attention_mask.device) >= rate + col_keep = col_keep.unsqueeze(1).expand_as(msa_attention_mask).clone() + col_keep[:, 0, :] = True + return msa_attention_mask.bool() & col_keep + + +# =========================================================================== +# Atom-token utilities +# =========================================================================== + + +def gather_token_to_atom(token_features: Tensor, atom_to_token_idx: Tensor) -> Tensor: + """Broadcast per-token features to per-atom features using gather. + + Args: + token_features: X with shape (b, l, d). + atom_to_token_idx: I with shape (b, a), int64. + + Returns: + X with shape (b, a, d). + """ + idx = atom_to_token_idx.unsqueeze(-1).expand(-1, -1, token_features.size(-1)) + return torch.gather(token_features, 1, idx) + + +def scatter_atom_to_token( + atom_features: Tensor, + atom_to_token_idx: Tensor, + n_tokens: int, + atom_mask: Tensor | None = None, +) -> Tensor: + """Aggregate per-atom features to per-token features (mean). + + Args: + atom_features: X with shape (b, a, d). + atom_to_token_idx: I with shape (b, a), int64. + n_tokens: Token count l. + atom_mask: M with shape (b, a), Boolean. + + Returns: + X with shape (b, l, d). + """ + batch_size, n_atoms, d_model = atom_features.shape + n_out = n_tokens + idx = atom_to_token_idx + if atom_mask is not None: + idx = torch.where(atom_mask, atom_to_token_idx, n_tokens) + n_out = n_tokens + 1 + idx_expanded = idx.unsqueeze(-1).expand(batch_size, n_atoms, d_model) + out = torch.zeros( + batch_size, + n_out, + d_model, + device=atom_features.device, + dtype=atom_features.dtype, + ) + out.scatter_reduce_(1, idx_expanded, atom_features, reduce="mean", include_self=False) + return out[:, :n_tokens, :] + + +def gather_rep_atom_coords(coords: Tensor, rep_atom_idx: Tensor) -> Tensor: + """Gather representative atom coordinates for each token. + + Args: + coords: X with shape (b, a, 3). + rep_atom_idx: I with shape (b, l), int64. + + Returns: + X with shape (b, l, 3). + """ + idx = rep_atom_idx.unsqueeze(-1).expand(-1, -1, coords.size(-1)) + return torch.gather(coords, 1, idx) + + +def _compute_intra_token_idx(atom_to_token: Tensor) -> Tensor: + """Compute local atom index within each token (vectorised). + + Atoms belonging to the same token are contiguous, so this computes a + running count that resets at each token boundary. + + Args: + atom_to_token: I with shape (b, a), mapping each atom to its token. + + Returns: + Index tensor I with shape (b, a) and values from zero through + ``max_atoms_per_token - 1``. + """ + same_as_prev = F.pad(atom_to_token[:, 1:] == atom_to_token[:, :-1], (1, 0), value=False) + ones = torch.ones_like(atom_to_token) + cumsum = torch.cumsum(ones, dim=-1) + group_start = cumsum.masked_fill(same_as_prev, 0) + group_start = torch.cummax(group_start, dim=-1).values + return cumsum - group_start + + +def _categorical_mean(logits: Tensor, start: float, end: float) -> Tensor: + """Expected value of a categorical distribution over evenly-spaced bins. + + Equivalent to ``CategoricalMixture(logits, bins=logits.shape[-1], start, end).mean()``. + + Args: + logits: Logit tensor X with shape (..., n_bins). + start: left boundary + end: right boundary + + Returns: + Expected value tensor Y with shape (...). + """ + n_bins = logits.shape[-1] + edges = torch.linspace(start, end, n_bins + 1, device=logits.device, dtype=torch.float32) + v_bins = (edges[:-1] + edges[1:]) / 2 # V_bin has shape (n_bins,). + return (logits.float().softmax(-1) @ v_bins.unsqueeze(1)).squeeze(-1) + + +# =========================================================================== +# Feature preparation and language-model projection +# =========================================================================== + + +class RowAttentionPooling(nn.Module): + """Row-wise attention pooling: attn_proj, out_proj.""" + + def __init__(self, d_pair: int, d_single: int) -> None: + super().__init__() + self.attn_proj = nn.Linear(d_pair, 1, bias=False) + self.out_proj = nn.Linear(d_pair, d_single, bias=False) + + def forward(self, z: Tensor, mask: Tensor) -> Tensor: + scores = self.attn_proj(z).squeeze(-1) + mask_bias = torch.where( + mask[:, None, :].bool(), + torch.zeros_like(scores), + torch.full_like(scores, -1e9), + ) + scores = scores + mask_bias + weights = F.softmax(scores, dim=-1) + pooled = torch.einsum("bnm,bnmd->bnd", weights, z) + return self.out_proj(pooled) + + +# =========================================================================== +# InputsEmbedder +# =========================================================================== + + +class InputsEmbedder(nn.Module): + """Embeds input features including atom-level encoding via SWA attention.""" + + def __init__(self, config: ESMFold2Config) -> None: + super().__init__() + swa_cfg = config.inputs.atom_encoder + + self.atom_attention_encoder = ESMFold2AtomEncoder( + d_atom=swa_cfg.d_atom, + d_token=swa_cfg.d_token, + n_blocks=swa_cfg.n_blocks, + n_heads=swa_cfg.n_heads, + swa_window_size=swa_cfg.swa_window_size, + expansion_ratio=swa_cfg.expansion_ratio, + structure_prediction=False, # no coords_linear + spatial_rope_base_frequency=swa_cfg.spatial_rope_base_frequency, + n_spatial_rope_pairs_per_axis=swa_cfg.n_spatial_rope_pairs_per_axis, + n_uid_rope_pairs=swa_cfg.n_uid_rope_pairs, + uid_rope_base_frequency=swa_cfg.uid_rope_base_frequency, + ) + + def forward( + self, + aatype: Tensor, + profile: Tensor, + deletion_mean: Tensor, + ref_pos: Tensor, + atom_attention_mask: Tensor, + ref_space_uid: Tensor, + ref_charge: Tensor, + ref_element: Tensor, + ref_atom_name_chars: Tensor, + atom_to_token: Tensor, + ) -> Tensor: + """Embed inputs into per-token features. + + Returns: + X with shape (b, l, d_inputs), concatenating atom encoding, + aatype, profile, and deletion mean. + """ + a, _q, _c, _attn_params, _intermediates = self.atom_attention_encoder( + ref_pos=ref_pos, + atom_attention_mask=atom_attention_mask, + ref_space_uid=ref_space_uid, + ref_charge=ref_charge, + ref_element=ref_element, + ref_atom_name_chars=ref_atom_name_chars, + atom_to_token=atom_to_token, + ) + return torch.cat([a, aatype, profile, deletion_mean.unsqueeze(-1)], dim=-1) + + +# =========================================================================== +# ResIdxAsymIdSymIdEntityIdEncoding (trunk relative position) +# =========================================================================== + + +class ResIdxAsymIdSymIdEntityIdEncoding(nn.Module): + """Embedding weight W has shape (d_pair, n_features). + + Here ``n_features = 2 * (2 * r_bins + 2) + 1 + (2 * c_bins + 2)``. + + For default r_bins=32, c_bins=2: 2*66 + 1 + 6 = 139. + """ + + def __init__( + self, + n_relative_residx_bins: int = 32, + n_relative_chain_bins: int = 2, + d_pair: int = 256, + ) -> None: + super().__init__() + self.n_relative_residx_bins = n_relative_residx_bins + self.n_relative_chain_bins = n_relative_chain_bins + self.d_pair = d_pair + + n_feats_residue = 2 * n_relative_residx_bins + 2 + n_feats_token = 2 * n_relative_residx_bins + 2 + n_feats_chain = 2 * n_relative_chain_bins + 2 + n_feats_same_entity = 1 + total_feats = n_feats_residue + n_feats_token + n_feats_chain + n_feats_same_entity + self.embed = nn.Linear(total_feats, d_pair, bias=False) + + def forward( + self, + residue_index: Tensor, + asym_id: Tensor, + sym_id: Tensor, + entity_id: Tensor, + token_index: Tensor, + ) -> Tensor: + bij_same_chain = asym_id.unsqueeze(2) == asym_id.unsqueeze(1) + bij_same_residue = residue_index.unsqueeze(2) == residue_index.unsqueeze(1) + bij_same_entity = entity_id.unsqueeze(2) == entity_id.unsqueeze(1) + + dij_residue = residue_index.unsqueeze(2) - residue_index.unsqueeze(1) + dij_residue = torch.clip( + dij_residue + self.n_relative_residx_bins, + 0, + 2 * self.n_relative_residx_bins, + ) + dij_residue = torch.where(bij_same_chain, dij_residue, 2 * self.n_relative_residx_bins + 1) + aij_rel_pos = F.one_hot(dij_residue, 2 * self.n_relative_residx_bins + 2) + + dij_token = torch.clip( + token_index.unsqueeze(2) - token_index.unsqueeze(1) + self.n_relative_residx_bins, + 0, + 2 * self.n_relative_residx_bins, + ) + dij_token = torch.where( + bij_same_chain & bij_same_residue, + dij_token, + 2 * self.n_relative_residx_bins + 1, + ) + aij_rel_token = F.one_hot(dij_token, 2 * self.n_relative_residx_bins + 2) + + dij_chain = torch.clip( + sym_id.unsqueeze(2) - sym_id.unsqueeze(1) + self.n_relative_chain_bins, + 0, + 2 * self.n_relative_chain_bins, + ) + dij_chain = torch.where(bij_same_chain, 2 * self.n_relative_chain_bins + 1, dij_chain) + aij_rel_chain = F.one_hot(dij_chain, 2 * self.n_relative_chain_bins + 2) + + feats = torch.cat( + [ + aij_rel_pos.float(), + aij_rel_token.float(), + bij_same_entity.float().unsqueeze(-1), + aij_rel_chain.float(), + ], + dim=-1, + ) + + return self.embed(feats) + + +# =========================================================================== +# SingleToPair (for LanguageModelShim) +# =========================================================================== + + +class SingleToPair(nn.Module): + """downproject, output_mlp (Sequential of Linear, GELU, Linear).""" + + def __init__(self, input_dim: int, downproject_dim: int, output_dim: int) -> None: + super().__init__() + self.downproject = nn.Linear(input_dim, downproject_dim) + self.output_mlp = nn.Sequential( + nn.Linear(2 * downproject_dim, output_dim), + nn.GELU(), + nn.Linear(output_dim, output_dim), + ) + + def forward(self, x: Tensor) -> Tensor: + x = self.downproject(x) + x = torch.cat( + [(x.unsqueeze(2) * x.unsqueeze(1)), (x.unsqueeze(2) - x.unsqueeze(1))], + dim=3, + ) + return self.output_mlp(x) + + +# =========================================================================== +# LanguageModelShim +# =========================================================================== + + +class LanguageModelShim(nn.Module): + """Shim holding the trainable projection weights for LM integration. + + Contains: + - base_z_combine: nn.Parameter with shape ``(n_layers + 1,)`` + - base_z_linear: Sequential(LayerNorm(d_model), Linear(d_model, d_z, bias=False)) + - base_z_mlp: Sequential(SingleToPair(d_z, d_z, d_z), LayerNorm(d_z)) + """ + + def __init__(self, d_z: int = 256, d_model: int = 2560, num_layers: int = 80) -> None: + super().__init__() + + self.base_z_mlp = nn.Sequential(SingleToPair(d_z, d_z, d_z), nn.LayerNorm(d_z)) + self.base_z_linear = nn.Sequential( + nn.LayerNorm(d_model), nn.Linear(d_model, d_z, bias=False) + ) + self.base_z_combine = nn.Parameter(torch.zeros(num_layers + 1)) + + def project_sequence( + self, + hidden_states: Tensor, + residue_mask: Tensor | None = None, + ) -> Tensor: + """Project all ESMC layer states into the learned sequence summary. + + Args: + hidden_states: H with shape ``(b, l, n_layers + 1, d_model)``. + residue_mask: Optional M with shape ``(b, l)``. Non-residue rows + are set to zero in the returned tensor. + + Returns: + Z with shape ``(b, l, d_z)``. + """ + + if hidden_states.ndim != 4: + raise ValueError( + "H must have shape (b, l, n_layers + 1, d_model), " + f"got {tuple(hidden_states.shape)}." + ) + expected_layers = self.base_z_combine.numel() + if hidden_states.shape[-2] != expected_layers: + raise ValueError( + f"H contains {hidden_states.shape[-2]} states; expected " + f"{expected_layers} in the official ESMC ordering." + ) + expected_width = cast(nn.LayerNorm, self.base_z_linear[0]).normalized_shape[0] + if hidden_states.shape[-1] != expected_width: + raise ValueError(f"H has width {hidden_states.shape[-1]}; expected {expected_width}.") + + # H can be FP32 even when the folding checkpoint is loaded in BF16. + # Match the learned projection parameters at this explicit boundary; + # this preserves the official BF16 path and leaves FP32 models exact. + projection_dtype = cast(nn.LayerNorm, self.base_z_linear[0]).weight.dtype + hidden_states = hidden_states.to(dtype=projection_dtype) + projected_states = self.base_z_linear(hidden_states) + layer_weights = self.base_z_combine.softmax(dim=0) + # Preserve Biohub's matmul path exactly so checkpoint inference does + # not change through a different reduction order. + projected = layer_weights @ projected_states + if residue_mask is not None: + if residue_mask.shape != hidden_states.shape[:2]: + raise ValueError( + "M must have shape (b, l), got " + f"{tuple(residue_mask.shape)} for H {tuple(hidden_states.shape)}." + ) + projected = projected * residue_mask.to( + device=projected.device, dtype=projected.dtype + ).unsqueeze(-1) + return projected + + def forward(self, hidden_states: Tensor, *, lm_dropout: float = 0.0) -> Tensor: + """Project pre-computed ESMC hidden states to pair representation. + + Args: + hidden_states: H with shape ``(b, l, n_layers + 1, d_model)``. + lm_dropout: Dropout probability applied to the pair + representation after ``base_z_mlp``. + + Returns: + Z_pair with shape ``(b, l, l, d_pair)``. + """ + lm_z = self.project_sequence(hidden_states) + lm_z = self.base_z_mlp(lm_z) + if lm_dropout > 0: + lm_z = F.dropout(lm_z, p=lm_dropout, training=True) + return lm_z + + +# =========================================================================== +# ESMFold2ExperimentalModel: the top-level PreTrainedModel +# =========================================================================== + + +def compute_lm_hidden_states( + esmc: nn.Module, + input_ids: Tensor, + asym_id: Tensor, + residue_index: Tensor, + mol_type: Tensor, + token_mask: Tensor, + pad_to_multiple: int | None = None, + lm_mask_pct: float = 0.0, + mask_token_id: int = 32, +) -> Tensor: + """Run ESMC and return H with shape ``(b, l, n_states, d_model)``. + + Atom-tokenized modified residues (HYP, MSE, ACE, NH2, ...) span multiple + structure tokens but share a single ``(asym_id, residue_index)`` key: + collapse them to one LM token per residue before running the LM (the LM + was trained on per-residue inputs, not per-atom), then scatter the + hidden states back to the per-token layout. + """ + b_size, l_size = input_ids.shape + device = input_ids.device + protein_mask = (mol_type == 0) & token_mask + + lm_input_list = [] + lm_lengths = [] + # Per-batch maps from (original protein-token index) to (LM input position). + expand_maps: list[Tensor] = [] + for batch_index in range(b_size): + mask_b = protein_mask[batch_index] + ids_b = input_ids[batch_index][mask_b] + asym_b = asym_id[batch_index][mask_b] + res_b = residue_index[batch_index][mask_b] + + # Collapse: keep first token per (asym_id, residue_index) key, in + # input order. ``inverse`` maps each original protein-token to its + # collapsed residue index. + keys = torch.stack((asym_b, res_b), dim=1) + unique_keys, inverse = torch.unique(keys, dim=0, return_inverse=True) + n_unique = unique_keys.size(0) + token_positions = torch.arange(keys.size(0), device=device, dtype=torch.long) + first_pos = torch.full((n_unique,), keys.size(0), device=device, dtype=torch.long) + first_pos.scatter_reduce_(0, inverse, token_positions, reduce="amin", include_self=True) + ordered = torch.argsort(first_pos) + first_pos_ordered = first_pos[ordered] + ids_collapsed = ids_b[first_pos_ordered] + asym_collapsed = asym_b[first_pos_ordered] + remap = torch.empty_like(ordered) + remap[ordered] = torch.arange(n_unique, device=device, dtype=torch.long) + inverse_ordered = remap[inverse] + + chain_ids = asym_collapsed.unique(sorted=True) + # [BOS] chain1 [EOS BOS] chain2 ... [EOS] + parts: list[Tensor] = [torch.tensor([0], device=device, dtype=ids_b.dtype)] + # Per-chain LM positions accumulate; track them for the expand map. + per_token_lm_pos = torch.empty(n_unique, device=device, dtype=torch.long) + cursor = 1 # position 0 is the leading BOS + for i, cid in enumerate(chain_ids): + in_chain = (asym_collapsed == cid).nonzero(as_tuple=True)[0] + parts.append(ids_collapsed[in_chain]) + per_token_lm_pos[in_chain] = torch.arange( + cursor, cursor + in_chain.shape[0], device=device, dtype=torch.long + ) + cursor += in_chain.shape[0] + if i < len(chain_ids) - 1: + parts.append(torch.tensor([2, 0], device=device, dtype=ids_b.dtype)) + cursor += 2 # EOS + BOS + parts.append(torch.tensor([2], device=device, dtype=ids_b.dtype)) + lm_seq = torch.cat(parts) + lm_input_list.append(lm_seq) + lm_lengths.append(lm_seq.shape[0]) + + # Map each original protein-token position to its LM input position. + prot_pos_b = mask_b.nonzero(as_tuple=True)[0] + expand_map = torch.full((l_size,), -1, device=device, dtype=torch.long) + expand_map[prot_pos_b] = per_token_lm_pos[inverse_ordered] + expand_maps.append(expand_map) + + # Pad the language-model input to its longest sequence. FP8 callers round + # l to a multiple of 16 for Transformer Engine kernels. + max_len = max(lm_lengths) + if pad_to_multiple is not None and pad_to_multiple > 1: + max_len = ((max_len + pad_to_multiple - 1) // pad_to_multiple) * pad_to_multiple + lm_input_ids = torch.full( + (b_size, max_len), + 1, + device=device, + dtype=input_ids.dtype, # PAD=1 + ) + for batch_index in range(b_size): + lm_input_ids[batch_index, : lm_lengths[batch_index]] = lm_input_list[batch_index] + + # sequence_id for chain-aware attention; PAD tokens get -1 (no attention). + sequence_id = (lm_input_ids == 0).cumsum(dim=1) - 1 # BOS=0 + sequence_id = sequence_id.masked_fill(lm_input_ids == 1, -1) # PAD=1 + + if lm_mask_pct > 0.0: + special = (lm_input_ids == 0) | (lm_input_ids == 1) | (lm_input_ids == 2) + do_mask = (torch.rand(lm_input_ids.shape, device=device) < lm_mask_pct) & ~special + lm_input_ids = lm_input_ids.masked_fill(do_mask, mask_token_id) + + with torch.inference_mode(): + esmc_out = esmc(input_ids=lm_input_ids, sequence_id=sequence_id, output_hidden_states=True) + + hidden_stack = esmc_out.hidden_states + n_states, _, _, d_model = hidden_stack.shape + result = torch.zeros(b_size, l_size, n_states, d_model, device=device, dtype=hidden_stack.dtype) + for batch_index in range(b_size): + M_i = protein_mask[batch_index] + positions = expand_maps[batch_index][M_i] + gathered = hidden_stack[:, batch_index, positions, :].permute(1, 0, 2) + result[batch_index, M_i.nonzero(as_tuple=True)[0]] = gathered + + return result.detach() + + +# =========================================================================== +# TriangleMultiplicativeUpdate +# =========================================================================== + + +class TriangleMultiplicativeBlock(nn.Module): + """Triangle multiplicative update block with gated signal routing.""" + + _FLOW_TO_EINSUM: ClassVar[dict[str, str]] = { + "outgoing": "bikd,bjkd->bijd", + "incoming": "bkid,bkjd->bijd", + } + _VALID_FLOWS = ("outgoing", "incoming") + + def __init__(self, input_channels: int, latent_channels: int, flow: str) -> None: + super().__init__() + if flow not in self._FLOW_TO_EINSUM: + raise ValueError(f"Invalid flow={flow!r}. Expected one of {self._VALID_FLOWS}.") + + self.input_channels = input_channels + self.latent_channels = latent_channels + self.flow = flow + self._einsum_equation = self._FLOW_TO_EINSUM[flow] + self.norm_start = nn.LayerNorm(self.input_channels, eps=_EPS) + self.norm_mix = nn.LayerNorm(self.latent_channels, eps=_EPS) + self.proj_bundle = nn.Linear(self.input_channels, 4 * self.latent_channels, bias=False) + self.proj_emit = nn.Linear(self.latent_channels, self.input_channels, bias=False) + self.proj_gate = nn.Linear(self.input_channels, self.input_channels, bias=False) + + self._use_kernels: bool = False + # Default chunked for memory on long sequences; tests override with + # ``set_chunk_size(None)`` for the unchunked path under bit-exact bf16 + # parity checks. + self._chunk_size: int | None = 64 + + def set_chunk_size(self, chunk_size: int | None) -> None: + self._chunk_size = chunk_size + + def split_kernel_weights(self) -> tuple[Tensor, Tensor]: + return ( + self.proj_bundle.weight[: 2 * self.latent_channels, :], + self.proj_bundle.weight[2 * self.latent_channels :, :], + ) + + def _kernel_flow_direction(self) -> str: + return self.flow + + def _triangular_contract(self, left_stream: Tensor, right_stream: Tensor) -> Tensor: + return torch.einsum(self._einsum_equation, left_stream, right_stream) + + def _triangular_contract_chunked( + self, left_stream: Tensor, right_stream: Tensor, chunk_size: int + ) -> Tensor: + """Compute the triangular einsum in chunks along the output i-dimension.""" + length = left_stream.shape[1] if self.flow == "outgoing" else left_stream.shape[2] + chunks = [] + for start in range(0, length, chunk_size): + end = min(start + chunk_size, length) + if self.flow == "outgoing": + chunk = torch.einsum(self._einsum_equation, left_stream[:, start:end], right_stream) + else: + chunk = torch.einsum( + self._einsum_equation, left_stream[:, :, start:end], right_stream + ) + chunks.append(chunk) + return torch.cat(chunks, dim=1) + + def forward(self, pair_grid: Tensor, visibility: Tensor | None = None) -> Tensor: + if visibility is None: + visibility = pair_grid.new_ones(pair_grid.shape[:-1]) + + if self._use_kernels: + p_in_weight, g_in_weight = self.split_kernel_weights() + return _cue_tri_mul( # type: ignore[misc] + pair_grid, + direction=self._kernel_flow_direction(), + mask=visibility, + norm_in_weight=self.norm_start.weight, + norm_in_bias=self.norm_start.bias, + p_in_weight=p_in_weight, + g_in_weight=g_in_weight, + norm_out_weight=self.norm_mix.weight, + norm_out_bias=self.norm_mix.bias, + p_out_weight=self.proj_emit.weight, + g_out_weight=self.proj_gate.weight, + eps=_EPS, + ) + + normalized_grid = self.norm_start(pair_grid) + bundled = self.proj_bundle(normalized_grid) + signal, gate_logits = bundled.split(2 * self.latent_channels, dim=-1) + routed = signal * torch.sigmoid(gate_logits) + routed = routed * visibility.unsqueeze(-1) + + left_stream, right_stream = routed.float().chunk(2, dim=-1) + if self._chunk_size is not None: + contracted = self._triangular_contract_chunked( + left_stream, right_stream, self._chunk_size + ) + else: + contracted = self._triangular_contract(left_stream, right_stream) + mixed = self.proj_emit(self.norm_mix(contracted)) + output_gate = torch.sigmoid(self.proj_gate(normalized_grid)) + return mixed * output_gate + + +class TriangleMultiplicativeUpdate(nn.Module): + """Thin wrapper exposing the triangular mixer with explicit orientation (v3).""" + + def __init__(self, dim: int = 128, _outgoing: bool = True) -> None: + super().__init__() + flow = "outgoing" if _outgoing else "incoming" + self._engine = TriangleMultiplicativeBlock( + input_channels=dim, latent_channels=dim, flow=flow + ) + + def set_kernel_backend(self, backend: str | None) -> None: + # Engine uses cueq when backend=="cuequivariance"; the "fused" backend + # routes through the parent PairUpdateBlock's fused path (bypassing this). + validate_kernel_backend(backend) + self._engine._use_kernels = backend == BACKEND_CUEQ + + def set_chunk_size(self, chunk_size: int | None) -> None: + self._engine.set_chunk_size(chunk_size) + + def forward(self, z: Tensor, mask: Tensor | None = None) -> Tensor: + return self._engine(z, visibility=mask) + + +# =========================================================================== +# FoldingTrunk: Transition, PairUpdateBlock, FoldingTrunk +# =========================================================================== + + +class Transition(nn.Module): + """LN + SwiGLU FFN with addmm-fused residual; optional Triton LN+w12+SwiGLU kernel.""" + + def __init__(self, d_model: int, expansion_ratio: int = 4) -> None: + super().__init__() + self.norm = nn.LayerNorm(d_model) + self.ffn = SwiGLUMLP(d_model, expansion_ratio=expansion_ratio, bias=False) + # Default chunked; set_chunk_size(None) disables for bit-exact parity tests. + self._chunk_size: int | None = 64 + self._fused_swiglu: nn.Module | None = None + self._kernel_backend: str | None = None + + def set_chunk_size(self, chunk_size: int | None) -> None: + self._chunk_size = chunk_size + + def set_kernel_backend(self, backend: str | None) -> None: + """Install / uninstall FusedLNLinearSwiGLU (no cueq equivalent).""" + validate_kernel_backend(backend) + self._kernel_backend = backend + if backend == BACKEND_FUSED and TRITON_KERNELS_AVAILABLE: + assert _FusedLNLinearSwiGLU is not None + d_model = self.norm.normalized_shape[0] + d_inner = self.ffn.hidden_features + has_ln_bias = self.norm.bias is not None + device = self.ffn.w12.weight.device + dtype = self.ffn.w12.weight.dtype + fused = _FusedLNLinearSwiGLU( + d_model=d_model, + d_inner=d_inner, + has_ln_bias=has_ln_bias, + device=device, + dtype=dtype, + ) + with torch.no_grad(): + fused.LN_W.copy_(self.norm.weight) + if has_ln_bias: + fused.LN_B.copy_(self.norm.bias) # type: ignore[union-attr] + # FusedLNLinearSwiGLU.W12 is (d_model, 2*d_inner); transpose nn.Linear once. + fused.W12.copy_(self.ffn.w12.weight.t().contiguous()) + self._fused_swiglu = fused.eval().requires_grad_(False) + else: + self._fused_swiglu = None + + def _can_use_fused_path(self, x: Tensor) -> bool: + return ( + _fused_active(self, x) and self._fused_swiglu is not None and x.dtype == torch.bfloat16 + ) + + def _swiglu_pre_w3(self, x_normed: Tensor) -> Tensor: + """SwiGLU through silu(x1)*x2, before the final w3.""" + ffn = self.ffn + x12 = ffn.w12(x_normed) + x1, x2 = x12.split(ffn.hidden_features, dim=-1) + return F.silu(x1) * x2 + + def _addmm_residual(self, x: Tensor, hidden: Tensor) -> Tensor: + """x + w3(hidden) via single cuBLAS addmm: avoids transition-output allocation.""" + ffn = self.ffn + x_shape = x.shape + out = torch.addmm( + x.contiguous().view(-1, x_shape[-1]), + hidden.view(-1, hidden.shape[-1]), + ffn.w3.weight.t(), + ) + return out.view(x_shape) + + def forward(self, x: Tensor) -> Tensor: + # Inference-only fast path (addmm-fused residual + pre-alloc out) + #: diverges bit-exactly from ``x + ffn(norm(x))`` so we only use + # it when grad is disabled (binder-design / bit-exact tests run + # with grad on and need the reference path). + if not torch.is_grad_enabled() and self._can_use_fused_path(x): + fused = self._fused_swiglu + assert fused is not None + pre_w3 = fused + if self._chunk_size is None or x.shape[1] <= self._chunk_size: + hidden = pre_w3(x) + return self._addmm_residual(x, hidden) + out = torch.empty_like(x) + for s in range(0, x.shape[1], self._chunk_size): + e = min(s + self._chunk_size, x.shape[1]) + sl = x[:, s:e] + hidden = pre_w3(sl) + out[:, s:e] = self._addmm_residual(sl, hidden) + return out + # Reference path: bit-exact with main: x + ffn(norm(x)). + if self._chunk_size is None or x.shape[1] <= self._chunk_size: + return x + self.ffn(self.norm(x)) + out_list: list[Tensor] = [] + for s in range(0, x.shape[1], self._chunk_size): + e = min(s + self._chunk_size, x.shape[1]) + sl = x[:, s:e] + out_list.append(sl + self.ffn(self.norm(sl))) + return torch.cat(out_list, dim=1) + + +class PairUpdateBlock(nn.Module): + """tri_mul_out, tri_mul_in, pair_transition.""" + + def __init__(self, d_pair: int = 256, expansion_ratio: int = 4) -> None: + super().__init__() + self.tri_mul_out = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=True) + self.tri_mul_in = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=False) + self.pair_transition = Transition(d_pair, expansion_ratio=expansion_ratio) + self._kernel_backend: str | None = None + # Row-shared dropout-residual; r=0 for inference (HF model is inference-only). + # backend='fused' swaps in the FusedDropoutResidual Triton kernel. + self.row_drop = DropoutResidual(0.0, batch_dim=1, use_fused_kernels=False) + + def set_kernel_backend(self, backend: str | None) -> None: + if backend not in _VALID_BACKENDS: + raise ValueError(f"backend must be one of {_VALID_BACKENDS}, got {backend!r}") + self.tri_mul_out.set_kernel_backend(backend) + self.tri_mul_in.set_kernel_backend(backend) + self.pair_transition.set_kernel_backend(backend) + self._kernel_backend = backend + self.row_drop = DropoutResidual( + 0.0, batch_dim=1, use_fused_kernels=(backend == BACKEND_FUSED) + ) + + def set_chunk_size(self, chunk_size: int | None) -> None: + self.tri_mul_out.set_chunk_size(chunk_size) + self.tri_mul_in.set_chunk_size(chunk_size) + self.pair_transition.set_chunk_size(chunk_size) + + def _can_use_fused_trimul_with_residual(self, pair: Tensor) -> bool: + return _fused_active(self, pair) and pair.dtype == torch.bfloat16 + + def _fused_trimul_with_residual( + self, pair: Tensor, direction: str, pair_attention_mask: Tensor | None + ) -> Tensor: + """Fused TriMul+residual call; weights from the corresponding engine.""" + tri = self.tri_mul_out if direction == "outgoing" else self.tri_mul_in + engine: TriangleMultiplicativeBlock = tri._engine # type: ignore[assignment] + p_in_weight, g_in_weight = engine.split_kernel_weights() + + def _bf16(t: Tensor) -> Tensor: + return t if t.dtype == torch.bfloat16 else t.to(torch.bfloat16) + + return _fused_trimul_with_residual( # type: ignore[misc] + pair, + direction, + residual=pair, + drop_mask=None, # inference: no dropout, matches internal's eval path + norm_in_weight=_bf16(engine.norm_start.weight), + norm_in_bias=_bf16(engine.norm_start.bias), + p_in_weight=_bf16(p_in_weight), + g_in_weight=_bf16(g_in_weight), + norm_out_weight=_bf16(engine.norm_mix.weight), + norm_out_bias=_bf16(engine.norm_mix.bias), + p_out_weight=_bf16(engine.proj_emit.weight), + g_out_weight=_bf16(engine.proj_gate.weight), + mask=pair_attention_mask, + eps=_EPS, + ) + + def forward(self, pair: Tensor, pair_attention_mask: Tensor | None = None) -> Tensor: + if self._can_use_fused_trimul_with_residual(pair): + pair = self._fused_trimul_with_residual(pair, "outgoing", pair_attention_mask) + pair = self._fused_trimul_with_residual(pair, "incoming", pair_attention_mask) + else: + pair = self.row_drop(pair, self.tri_mul_out(pair, mask=pair_attention_mask)) + pair = self.row_drop(pair, self.tri_mul_in(pair, mask=pair_attention_mask)) + pair = self.pair_transition(pair) + return pair + + +class FoldingTrunk(nn.Module): + """ModuleList of PairUpdateBlocks.""" + + def __init__(self, n_layers: int = 24, d_pair: int = 256, expansion_ratio: int = 4) -> None: + super().__init__() + self.blocks = nn.ModuleList( + [ + PairUpdateBlock(d_pair=d_pair, expansion_ratio=expansion_ratio) + for _ in range(n_layers) + ] + ) + + def set_kernel_backend(self, backend: str | None) -> None: + for block in self.blocks: + cast(PairUpdateBlock, block).set_kernel_backend(backend) + + def set_chunk_size(self, chunk_size: int | None) -> None: + for block in self.blocks: + cast(PairUpdateBlock, block).set_chunk_size(chunk_size) + + def forward(self, pair: Tensor, pair_attention_mask: Tensor | None = None) -> Tensor: + # Cast the pair tensor to BF16 when the fused triangle backend is enabled + # (its bwd kernel requires bf16). Other backends keep the input dtype. + orig_dtype = pair.dtype + fused_on = ( + len(self.blocks) > 0 + and getattr(self.blocks[0], "_kernel_backend", None) == BACKEND_FUSED + ) + if pair.is_cuda and fused_on and orig_dtype != torch.bfloat16: + pair = pair.to(torch.bfloat16) + for block in self.blocks: + fn = partial(block, pair_attention_mask=pair_attention_mask) + if torch.is_grad_enabled(): + pair = checkpoint(fn, pair, use_reentrant=False) # pyright: ignore + else: + pair = fn(pair) + if pair.dtype != orig_dtype: + pair = pair.to(orig_dtype) + return pair + + +# =========================================================================== +# MSA Encoder +# =========================================================================== + + +class OuterProductMean(nn.Module): + """Outer-product mean: maps an MSA representation into a pair update. + + The order of the ``/ n_valid`` divide vs. the ``Wout`` projection is + selectable via ``divide_outer_before_proj`` because different ESMFold2 + checkpoints were trained with different orderings: + + * ``False`` (default): ``Wout(outer) / n_valid``: the projection bias + is scaled by 1/n_valid alongside the outer product. + * ``True``: ``Wout(outer / n_valid)``: the projection bias is added + unscaled, post-divide. + """ + + def __init__( + self, + d_msa: int, + d_hidden: int, + d_pair: int, + divide_outer_before_proj: bool = False, + ) -> None: + super().__init__() + self.d_hidden = d_hidden + self.divide_outer_before_proj = divide_outer_before_proj + self.norm = nn.LayerNorm(d_msa) + self.W = nn.Linear(d_msa, 2 * d_hidden, bias=False) + self.Wout = nn.Linear(d_hidden * d_hidden, d_pair, bias=True) + # Off for bit-exact bf16; ``set_chunk_size(64)`` for long sequences. + self._chunk_size: int | None = None + + def set_chunk_size(self, chunk_size: int | None) -> None: + self._chunk_size = chunk_size + + def forward(self, m: Tensor, msa_attention_mask: Tensor) -> Tensor: + m_norm = self.norm(m) + x = self.W(m_norm) * msa_attention_mask.unsqueeze(-1).to(m_norm.dtype) + a, b = x.chunk(2, dim=-1) + mask_f = msa_attention_mask.to(a.dtype) + n_valid = (mask_f @ mask_f.transpose(-1, -2)).unsqueeze(-1).clamp(min=1.0) + if self._chunk_size is None: + outer = torch.einsum("bimc,bjmd->bijcd", a, b).flatten(-2) + if self.divide_outer_before_proj: + return self.Wout(outer / n_valid) + return self.Wout(outer) / n_valid + # Chunk along the left (i) axis so the peak einsum intermediate is + # X uses shape (b, chunk, l, c, d) instead of (b, l, l, c, d). + length = a.shape[1] + out_chunks: list[Tensor] = [] + for start in range(0, length, self._chunk_size): + end = min(start + self._chunk_size, length) + outer_chunk = torch.einsum("bimc,bjmd->bijcd", a[:, start:end], b).flatten(-2) + if self.divide_outer_before_proj: + out_chunks.append(self.Wout(outer_chunk / n_valid[:, start:end])) + else: + out_chunks.append(self.Wout(outer_chunk) / n_valid[:, start:end]) + return torch.cat(out_chunks, dim=1) + + +class MSAPairWeightedAveraging(nn.Module): + """Pair-biased MSA row update (AF3 Supplement Algorithm 10).""" + + def __init__(self, d_msa: int, d_pair: int, n_heads: int = 8, head_width: int = 32) -> None: + super().__init__() + self.n_heads = n_heads + self.head_width = head_width + self.norm_single = nn.LayerNorm(d_msa) + self.compute_bias = nn.Sequential( + nn.LayerNorm(d_pair), nn.Linear(d_pair, n_heads, bias=False) + ) + self.Wv = nn.Linear(d_msa, n_heads * head_width, bias=False) + self.Wgate = nn.Linear(d_msa, n_heads * head_width, bias=False) + self.Wout = nn.Linear(n_heads * head_width, d_msa, bias=False) + + def forward(self, msa_repr: Tensor, pair_repr: Tensor, pair_attention_mask: Tensor) -> Tensor: + """ + Args: + msa_repr: X with shape (b, l, m, d_msa). + pair_repr: Z with shape (b, l, l, d_pair). + pair_attention_mask: M with shape (b, l, l). + Returns: + X with shape (b, l, m, d_msa). + """ + batch_size, length, depth, _ = msa_repr.shape + n_heads, head_width = self.n_heads, self.head_width + + msa_normed = self.norm_single(msa_repr) + bias = self.compute_bias(pair_repr) # A has shape (b, l, l, n_heads). + bias.masked_fill_(~pair_attention_mask.unsqueeze(-1).bool(), -1e5) + attn = torch.softmax(bias, dim=-2) # softmax over j + + v = self.Wv(msa_normed).reshape(batch_size, length, depth, n_heads, head_width) + gate = torch.sigmoid(self.Wgate(msa_normed)).reshape( + batch_size, length, depth, n_heads, head_width + ) + + output = torch.einsum("bijh,bjmhd,bimhd->bimhd", attn, v, gate) + return self.Wout(output.reshape(batch_size, length, depth, n_heads * head_width)) + + +# =========================================================================== +# Atom and diffusion stack +# =========================================================================== + + +class TransitionLayer(nn.Module): + """SwiGLU transition: norm -> a_proj, b_proj -> silu(a)*b -> out_proj.""" + + def __init__(self, d_model: int, n: int, eps: float = 1e-5) -> None: + super().__init__() + hidden = n * d_model + self.norm = nn.LayerNorm(d_model, eps=eps) + self.a_proj = nn.Linear(d_model, hidden, bias=False) + self.b_proj = nn.Linear(d_model, hidden, bias=False) + self.out_proj = nn.Linear(hidden, d_model, bias=False) + + def forward(self, x: Tensor) -> Tensor: + x = self.norm(x) + a = self.a_proj(x) + b = self.b_proj(x) + return self.out_proj(F.silu(a) * b) + + +# =========================================================================== +# AdaptiveLayerNorm (used in DiffusionTransformer) +# =========================================================================== + + +class AdaptiveLayerNorm(nn.Module): + """Adaptive layer normalization (adaLN-Zero).""" + + def __init__(self, d_model: int, d_cond: int, eps: float = 1e-5) -> None: + super().__init__() + self.d_model = d_model + self.d_cond = d_cond + self.eps = eps + self.s_scale = nn.Parameter(torch.ones(d_cond)) + self.s_gate = nn.Linear(d_cond, d_model, bias=True) + self.s_shift = nn.Linear(d_cond, d_model, bias=False) + + def forward(self, a: Tensor, s: Tensor) -> Tensor: + a_norm = F.layer_norm(a, (self.d_model,), None, None, self.eps) + s_norm = F.layer_norm(s, (self.d_cond,), self.s_scale, None, self.eps) + return torch.sigmoid(self.s_gate(s_norm)) * a_norm + self.s_shift(s_norm) + + +# =========================================================================== +# FourierEmbedding +# =========================================================================== + + +class FourierEmbedding(nn.Module): + """Fourier embedding: cos(2*pi*(t*w + b)).""" + + w: Tensor + b: Tensor + + def __init__(self, c: int) -> None: + super().__init__() + self.c = c + self.register_buffer("w", torch.randn(c)) + self.register_buffer("b", torch.randn(c)) + + def forward(self, t_hat: Tensor) -> Tensor: + t = torch.as_tensor(t_hat, device=self.w.device, dtype=self.w.dtype).reshape(-1) + return torch.cos(2.0 * torch.pi * (t[:, None] * self.w[None, :] + self.b[None, :])) + + +# =========================================================================== +# SwiGLU / SwiGLUMLP +# =========================================================================== + + +def _compute_swiglu_hidden_size(d_model: int, expansion_ratio: int) -> int: + return expansion_ratio * d_model + + +class SwiGLU(nn.Module): + """SwiGLU with packed w12 and output w3.""" + + def __init__( + self, + in_features: int, + hidden_features: int, + out_features: int | None = None, + bias: bool = True, + ) -> None: + super().__init__() + out_features = out_features or in_features + self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias) + self.w3 = nn.Linear(hidden_features, out_features, bias=bias) + self.hidden_features = hidden_features + + def forward(self, x: Tensor) -> Tensor: + x12 = self.w12(x) + x1, x2 = x12.split(self.hidden_features, dim=-1) + hidden = F.silu(x1) * x2 + return self.w3(hidden) + + +class SwiGLUMLP(SwiGLU): + """SwiGLU MLP with packed weights, no bias.""" + + def __init__(self, d_model: int, expansion_ratio: int = 4, bias: bool = False) -> None: + hidden = _compute_swiglu_hidden_size(d_model, expansion_ratio) + super().__init__( + in_features=d_model, hidden_features=hidden, out_features=d_model, bias=bias + ) + + +# =========================================================================== +# SWA Atom Attention components +# =========================================================================== + + +def _rotate_half(x: Tensor) -> Tensor: + x1, x2 = x.chunk(2, dim=-1) + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_emb_3d(x: Tensor, cos: Tensor, sin: Tensor) -> Tensor: + """Apply RoPE with batch-dependent cos/sin. + + Args: + x: X with shape (b, l, h, d). + cos: C with shape (b, l, d / 2). + sin: S with shape (b, l, d / 2). + """ + ro_dim = cos.shape[-1] * 2 + cos = cos.unsqueeze(2).repeat(1, 1, 1, 2) + sin = sin.unsqueeze(2).repeat(1, 1, 1, 2) + return torch.cat( + [x[..., :ro_dim] * cos + _rotate_half(x[..., :ro_dim]) * sin, x[..., ro_dim:]], + dim=-1, + ) + + +@torch.compiler.disable +def build_3d_rope( + ref_pos: Tensor, + ref_space_uid: Tensor, + head_dim: int, + n_spatial_per_axis: int = 4, + n_uid_pairs: int = 2, + spatial_base_freq: float = 10000.0, + uid_base_freq: float = 10.0, +) -> tuple[Tensor, Tensor]: + """Build cos/sin for 3D RoPE + UID RoPE.""" + device = ref_pos.device + batch_size, n_atoms = ref_pos.shape[:2] + half_dim = head_dim // 2 + n_spatial_total = 3 * n_spatial_per_axis + + spatial_inv_freq = 1.0 / ( + spatial_base_freq + ** ( + torch.arange(0, n_spatial_per_axis, dtype=torch.float32, device=device) + / n_spatial_per_axis + ) + ) + uid_inv_freq = 1.0 / ( + uid_base_freq + ** (torch.arange(0, n_uid_pairs, dtype=torch.float32, device=device) / n_uid_pairs) + ) + + pos_f32 = ref_pos.float() + spatial_freqs = torch.einsum("bna,k->bnak", pos_f32, spatial_inv_freq) + spatial_freqs = spatial_freqs.reshape(batch_size, n_atoms, n_spatial_total) + + uid_f32 = ref_space_uid.float() + uid_freqs = torch.einsum("bn,k->bnk", uid_f32, uid_inv_freq) + + n_active = n_spatial_total + n_uid_pairs + freqs = torch.cat([spatial_freqs, uid_freqs], dim=-1) + + if n_active < half_dim: + padding = torch.zeros( + batch_size, + n_atoms, + half_dim - n_active, + device=device, + dtype=torch.float32, + ) + freqs = torch.cat([freqs, padding], dim=-1) + + cos = freqs.cos().to(torch.bfloat16) + sin = freqs.sin().to(torch.bfloat16) + return cos, sin + + +def qk_norm(x: Tensor) -> Tensor: + return F.rms_norm(x, (x.size(-1),)).to(x.dtype) + + +# =========================================================================== +# SwiGLUFFN (atom transformer blocks) +# =========================================================================== + + +class SwiGLUFFN(nn.Module): + """SwiGLU FFN with rounded hidden size for hardware alignment.""" + + def __init__(self, d_model: int, expansion_ratio: int = 2) -> None: + super().__init__() + hidden_size = ((expansion_ratio * (d_model // 3) * 2) + 255) // 256 * 256 + self.w_up = nn.Linear(d_model, 2 * hidden_size, bias=False) + self.w_down = nn.Linear(hidden_size, d_model, bias=False) + + def forward(self, x: Tensor) -> Tensor: + x = x.to(self.w_up.weight.dtype) + x1, x2 = self.w_up(x).chunk(2, dim=-1) + return self.w_down(F.silu(x1) * x2) + + +# =========================================================================== +# SWA3DRoPEAttention +# =========================================================================== + + +class SWA3DRoPEAttention(nn.Module): + """Sliding window attention with 3D RoPE. Has Wqkv, gate_proj, out_proj.""" + + def __init__(self, d_model: int, n_heads: int, half_window: int = 64) -> None: + super().__init__() + self.n_heads = n_heads + self.head_dim = d_model // n_heads + self.scale = self.head_dim**-0.5 + self.half_window = half_window + + self.Wqkv = nn.Linear(d_model, 3 * d_model, bias=False) + self.out_proj = nn.Linear(d_model, d_model, bias=False) + self.gate_proj = nn.Linear(d_model, d_model, bias=False) + + def forward(self, x: Tensor, attention_params: tuple) -> Tensor: + batch_size, n_atoms = x.shape[:2] + cos, sin = attention_params[0], attention_params[1] + + x_input = x + qkv = self.Wqkv(x) + qkv = qkv.view(batch_size, n_atoms, 3, self.n_heads, self.head_dim).permute(2, 0, 1, 3, 4) + q, k, v = qkv.unbind(0) + q, k = qk_norm(q), qk_norm(k) + + q = apply_rotary_emb_3d(q, cos, sin) + k = apply_rotary_emb_3d(k, cos, sin) + + input_dtype = q.dtype + if q.dtype not in (torch.float16, torch.bfloat16): + q, k, v = q.bfloat16(), k.bfloat16(), v.bfloat16() + + # ESMFold2 does not advertise FlashAttention. Keep this atom path on + # PyTorch. Models that advertise FlashAttention dispatch through the + # precompiled Hugging Face kernels interface in fastplms.attention. + q_t = q.transpose(1, 2) + k_t = k.transpose(1, 2) + v_t = v.transpose(1, 2) + attn = torch.matmul(q_t, k_t.transpose(-2, -1)) * self.scale + attn = F.softmax(attn, dim=-1) + out = torch.matmul(attn, v_t).transpose(1, 2) + + out = out.to(input_dtype).reshape( # type: ignore[union-attr] + batch_size, n_atoms, -1 + ) + out = out * torch.sigmoid(self.gate_proj(x_input)) + return self.out_proj(out) + + +# =========================================================================== +# SWAAtomBlock, SWAAtomTransformer +# =========================================================================== + + +def _rms_adaln_raw(x: Tensor, scale: Tensor, shift: Tensor) -> Tensor: + return F.rms_norm(x, (x.shape[-1],)) * (1 + scale) + shift + + +def _gated_residual_raw(x: Tensor, gate: Tensor, y: Tensor) -> Tensor: + return x + gate * y + + +class SWAAtomBlock(nn.Module): + """adaLN-Zero + SWA attention + SwiGLU FFN. + + Creates adaln_modulation = Sequential(SiLU(), Linear) -> keys like adaln_modulation.1.weight + """ + + def __init__( + self, + d_atom: int, + n_heads: int, + half_window: int = 64, + expansion_ratio: int = 2, + use_compile_fusions: bool = False, + ) -> None: + super().__init__() + self.attn_norm = nn.RMSNorm(d_atom, elementwise_affine=False) + self.ffn_norm = nn.RMSNorm(d_atom, elementwise_affine=False) + + adaln_linear = nn.Linear(d_atom, 6 * d_atom, bias=False) + nn.init.zeros_(adaln_linear.weight) + self.adaln_modulation = nn.Sequential(nn.SiLU(), adaln_linear) + + self.attn = SWA3DRoPEAttention(d_atom, n_heads, half_window=half_window) + self.ffn = SwiGLUFFN(d_atom, expansion_ratio) + + self._rms_adaln = torch.compile(_rms_adaln_raw) if use_compile_fusions else _rms_adaln_raw + self._gated_residual = ( + torch.compile(_gated_residual_raw) if use_compile_fusions else _gated_residual_raw + ) + + def forward(self, x: Tensor, c_l: Tensor, attention_params: tuple) -> Tensor: + mod = self.adaln_modulation(c_l) + if mod.dim() == 2: + mod = mod.unsqueeze(1) + shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = mod.chunk(6, dim=-1) + + attn_input = self._rms_adaln(x, scale_a, shift_a) + attn_out = self.attn(attn_input, attention_params) + x = self._gated_residual(x, gate_a, attn_out) + + ffn_input = self._rms_adaln(x, scale_f, shift_f) + ffn_out = self.ffn(ffn_input) + x = self._gated_residual(x, gate_f, ffn_out) + return x + + +class SWAAtomTransformer(nn.Module): + """Stack of SWAAtomBlocks.""" + + def __init__( + self, + d_atom: int = 128, + n_blocks: int = 3, + n_heads: int = 4, + swa_window_size: int = 128, + expansion_ratio: int = 2, + spatial_rope_base_frequency: float = 20.0, + n_spatial_rope_pairs_per_axis: int = 2, + n_uid_rope_pairs: int = 10, + uid_rope_base_frequency: float = 10000.0, + ) -> None: + super().__init__() + self.swa_window_size = swa_window_size + self.head_dim = d_atom // n_heads + self.spatial_rope_base_frequency = spatial_rope_base_frequency + self.n_spatial_rope_pairs_per_axis = n_spatial_rope_pairs_per_axis + self.n_uid_rope_pairs = n_uid_rope_pairs + self.uid_rope_base_frequency = uid_rope_base_frequency + + self.blocks = nn.ModuleList( + [ + SWAAtomBlock( + d_atom=d_atom, + n_heads=n_heads, + half_window=swa_window_size // 2, + expansion_ratio=expansion_ratio, + ) + for _ in range(n_blocks) + ] + ) + + def _build_3d_rope(self, ref_pos: Tensor, ref_space_uid: Tensor) -> tuple[Tensor, Tensor]: + return build_3d_rope( + ref_pos=ref_pos, + ref_space_uid=ref_space_uid, + head_dim=self.head_dim, + n_spatial_per_axis=self.n_spatial_rope_pairs_per_axis, + n_uid_pairs=self.n_uid_rope_pairs, + spatial_base_freq=self.spatial_rope_base_frequency, + uid_base_freq=self.uid_rope_base_frequency, + ) + + def forward( + self, + q_l: Tensor, + c_l: Tensor, + attention_params: tuple, + return_intermediates: bool = False, + ) -> Tensor | tuple[Tensor, list[Tensor]]: + intermediates: list[Tensor] = [] + for block in self.blocks: + q_l = block(q_l, c_l, attention_params) + if return_intermediates: + intermediates.append(q_l) + if return_intermediates: + return q_l, intermediates + return q_l + + +# =========================================================================== +# ESMFold2AtomEncoder (for both inputs_embedder and diffusion_module) +# =========================================================================== + + +@torch.compiler.disable +def _prepare_atom_encoder_metadata( + atom_attention_mask: Tensor, + atom_to_token: Tensor, + num_diffusion_samples: int, +) -> tuple[Tensor, Tensor, Tensor, int, int]: + """Prepare mask-derived atom metadata outside compiled diffusion graphs.""" + mask_exp = atom_attention_mask.repeat_interleave(num_diffusion_samples, 0) + seqlens = mask_exp.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(mask_exp.flatten(), as_tuple=False).flatten() + max_seqlen = int(seqlens.max().item()) + cu_seqlens = F.pad(torch.cumsum(seqlens, dim=0, dtype=torch.int32), (1, 0)) + n_tokens = int(atom_to_token.max().item()) + 1 + return mask_exp, indices, cu_seqlens, max_seqlen, n_tokens + + +class ESMFold2AtomEncoder(nn.Module): + """Encode atom inputs with normalization and sliding-window attention. + + Args: + d_atom: atom hidden dim + d_token: token dim for atom_to_token aggregation + n_blocks, n_heads, swa_window_size, expansion_ratio: transformer params + structure_prediction: if True, creates coords_linear and uses full d_token + spatial_rope_base_frequency, n_spatial_rope_pairs_per_axis, + n_uid_rope_pairs, uid_rope_base_frequency: 3D RoPE config + """ + + def __init__( + self, + d_atom: int = 128, + d_token: int = 768, + n_blocks: int = 3, + n_heads: int = 4, + swa_window_size: int = 128, + expansion_ratio: int = 2, + structure_prediction: bool = True, + spatial_rope_base_frequency: float = 20.0, + n_spatial_rope_pairs_per_axis: int = 2, + n_uid_rope_pairs: int = 10, + uid_rope_base_frequency: float = 10000.0, + ) -> None: + super().__init__() + self.d_atom = d_atom + self.d_token = d_token + self.structure_prediction = structure_prediction + + self.atom_linear = nn.Linear(ATOM_FEATURE_DIM, d_atom, bias=False) + self.atom_norm = nn.LayerNorm(d_atom) + + if structure_prediction: + self.coords_linear = nn.Linear(6, d_atom, bias=False) + + self.atom_transformer = SWAAtomTransformer( + d_atom=d_atom, + n_blocks=n_blocks, + n_heads=n_heads, + swa_window_size=swa_window_size, + expansion_ratio=expansion_ratio, + spatial_rope_base_frequency=spatial_rope_base_frequency, + n_spatial_rope_pairs_per_axis=n_spatial_rope_pairs_per_axis, + n_uid_rope_pairs=n_uid_rope_pairs, + uid_rope_base_frequency=uid_rope_base_frequency, + ) + + # Output aggregation: d_token for structure prediction, d_token//2 for inputs + out_dim = d_token if structure_prediction else d_token // 2 + self.atom_to_token_linear = nn.Linear(d_atom, out_dim, bias=False) + + def forward( + self, + ref_pos: Tensor, + atom_attention_mask: Tensor, + ref_space_uid: Tensor, + ref_charge: Tensor, + ref_element: Tensor, + ref_atom_name_chars: Tensor, + atom_to_token: Tensor, + r_l: Tensor | None = None, + pred_r1: Tensor | None = None, + s_i: Tensor | None = None, + z_ij: Tensor | None = None, + num_diffusion_samples: int = 1, + return_intermediates: bool = False, + inference_cache: dict | None = None, + ) -> tuple[Tensor, Tensor, Tensor, tuple, list[Tensor]]: + """Returns (a, q, c, attention_params, intermediates). + + ``inference_cache`` caches step-invariant tensors (c_base, 3D RoPE, + attention indices, n_tokens) across diffusion steps. + """ + batch_size, n_atoms = ref_pos.shape[:2] + + layer_cache = None + if inference_cache is not None: + layer_cache = inference_cache.setdefault("atomencoder", {}) + + if layer_cache is None or len(layer_cache) == 0: + atom_feats = torch.cat( + [ + ref_pos, + ref_charge.unsqueeze(-1), + atom_attention_mask.unsqueeze(-1), + ref_element, + ref_atom_name_chars.reshape(batch_size, n_atoms, MAX_CHARS * CHAR_VOCAB_SIZE), + ], + dim=-1, + ) + c_base = self.atom_norm(self.atom_linear(atom_feats)) + cos, sin = self.atom_transformer._build_3d_rope(ref_pos, ref_space_uid) + cos = cos.repeat_interleave(num_diffusion_samples, 0) + sin = sin.repeat_interleave(num_diffusion_samples, 0) + mask_exp, indices, cu_seqlens, max_seqlen, n_tokens = ( + _prepare_atom_encoder_metadata( + atom_attention_mask, + atom_to_token, + num_diffusion_samples, + ) + ) + attention_params = (cos, sin, indices, cu_seqlens, max_seqlen) + if layer_cache is not None: + layer_cache["c_base"] = c_base + layer_cache["attention_params"] = attention_params + layer_cache["mask_exp"] = mask_exp + layer_cache["n_tokens"] = n_tokens + layer_cache["atom_to_token_exp"] = atom_to_token.repeat_interleave( + num_diffusion_samples, 0 + ) + else: + c_base = layer_cache["c_base"] + attention_params = layer_cache["attention_params"] + mask_exp = layer_cache["mask_exp"] + n_tokens = layer_cache["n_tokens"] + + c = c_base + + q = c + + if self.structure_prediction and r_l is not None: + q = q.repeat_interleave(num_diffusion_samples, 0) + if pred_r1 is None: + pred_r1 = torch.zeros_like(r_l) + r_input = torch.cat([r_l, pred_r1], dim=-1) + r_to_q = self.coords_linear(r_input) + q = q + r_to_q + + c = c.repeat_interleave(num_diffusion_samples, 0) + + result = self.atom_transformer( + q_l=q, + c_l=c, + attention_params=attention_params, + return_intermediates=return_intermediates, + ) + if return_intermediates: + q, intermediates = result + else: + q = result + intermediates = [] + + q_to_a = F.relu(self.atom_to_token_linear(q)) + if layer_cache is not None and "atom_to_token_exp" in layer_cache: + atom_to_token_exp = layer_cache["atom_to_token_exp"] + else: + atom_to_token_exp = atom_to_token.repeat_interleave(num_diffusion_samples, 0) + a = scatter_atom_to_token(q_to_a, atom_to_token_exp, n_tokens, atom_mask=mask_exp.bool()) + + return a, q, c, attention_params, intermediates + + +# =========================================================================== +# ESMFold2AtomDecoder +# =========================================================================== + + +class ESMFold2AtomDecoder(nn.Module): + """SWA atom decoder with token_to_atom_linear, atom_transformer, norm, output_linear.""" + + def __init__( + self, + d_atom: int = 128, + d_token: int = 768, + n_blocks: int = 3, + n_heads: int = 4, + swa_window_size: int = 128, + expansion_ratio: int = 2, + spatial_rope_base_frequency: float = 20.0, + n_spatial_rope_pairs_per_axis: int = 2, + n_uid_rope_pairs: int = 10, + uid_rope_base_frequency: float = 10000.0, + ) -> None: + super().__init__() + self.token_to_atom_linear = nn.Linear(d_token, d_atom, bias=False) + + self.atom_transformer = SWAAtomTransformer( + d_atom=d_atom, + n_blocks=n_blocks, + n_heads=n_heads, + swa_window_size=swa_window_size, + expansion_ratio=expansion_ratio, + spatial_rope_base_frequency=spatial_rope_base_frequency, + n_spatial_rope_pairs_per_axis=n_spatial_rope_pairs_per_axis, + n_uid_rope_pairs=n_uid_rope_pairs, + uid_rope_base_frequency=uid_rope_base_frequency, + ) + + self.norm = nn.LayerNorm(d_atom) + self.output_linear = nn.Linear(d_atom, XYZ_DIMS, bias=False) + + def forward( + self, + a_i: Tensor, + q_l: Tensor, + c_l: Tensor, + p_lm: tuple, + atom_to_token: Tensor, + atom_attention_mask: Tensor, + num_diffusion_samples: int = 1, + return_intermediates: bool = False, + ) -> tuple[Tensor, list[Tensor]]: + """Returns (r_update, intermediates).""" + atom_to_token_exp = atom_to_token.repeat_interleave(num_diffusion_samples, 0) + a_to_q = self.token_to_atom_linear(a_i) + a_to_q = gather_token_to_atom(a_to_q, atom_to_token_exp) + q_l = q_l + a_to_q + + result = self.atom_transformer( + q_l=q_l, + c_l=c_l, + attention_params=p_lm, + return_intermediates=return_intermediates, + ) + if return_intermediates: + q_l, intermediates = result + else: + q_l = result + intermediates = [] + + r_l = self.output_linear(self.norm(q_l)) + return r_l, intermediates + + +# =========================================================================== +# AttentionPairBias (DiffusionTransformer attention block) +# =========================================================================== + + +class AttentionPairBias(nn.Module): + """Gated multi-head attention with pair bias conditioning.""" + + def __init__( + self, + d_model: int, + d_pair: int, + num_heads: int, + d_cond: int | None = None, + use_conditioning: bool = True, + ) -> None: + super().__init__() + self.d_model = d_model + self.num_heads = num_heads + self.head_dim = d_model // num_heads + self.scale = self.head_dim**-0.5 + d_cond = d_cond or d_model + + if use_conditioning: + self.adaln = AdaptiveLayerNorm(d_model, d_cond, eps=1e-5) + self.out_gate = nn.Linear(d_cond, d_model, bias=True) + # adaln init: weight=0, bias=-2 + nn.init.zeros_(self.out_gate.weight) + nn.init.constant_(self.out_gate.bias, -2.0) + else: + self.pre_norm = nn.LayerNorm(d_model, eps=1e-5) + + self.q_proj = nn.Linear(d_model, d_model, bias=True) + self.kv_proj = nn.Linear(d_model, 2 * d_model, bias=False) + self.g_proj = nn.Linear(d_model, d_model, bias=False) + self.out_proj = nn.Linear(d_model, d_model, bias=False) + + if d_pair > 0: + self.pair_norm = nn.LayerNorm(d_pair, eps=1e-5) + self.pair_bias_proj = nn.Linear(d_pair, num_heads, bias=False) + + self._kernel_backend: str | None = None + + def set_kernel_backend(self, backend: str | None) -> None: + if backend not in _VALID_BACKENDS: + raise ValueError(f"backend must be one of {_VALID_BACKENDS}, got {backend!r}") + self._kernel_backend = backend + + def _is_zero_beta(self, beta: Tensor | float) -> bool: + if isinstance(beta, (int, float)): + return beta == 0.0 + return bool((beta == 0).all()) + + def _can_use_fused_pair_bias(self, z: Tensor, n_queries: int, beta: Tensor | float) -> bool: + return ( + _fused_active(self, z) + and z.dim() == 4 + and self._is_zero_beta(beta) + and hasattr(self, "pair_bias_proj") + and hasattr(self, "pair_norm") + ) + + def _can_use_cueq_pair_bias(self, z: Tensor, n_queries: int, beta: Tensor | float) -> bool: + return ( + _cueq_active(self) + and n_queries > 750 + and z.dim() == 4 + and self._is_zero_beta(beta) + and hasattr(self, "pair_bias_proj") + ) + + def forward( + self, + a: Tensor, + s: Tensor | None, + z: Tensor, + beta: Tensor | float = 0.0, + attention_mask: Tensor | None = None, + num_diffusion_samples: int = 1, + ) -> Tensor: + bsz, n_queries, d_model = a.shape + + x = self.adaln(a, s) if s is not None else self.pre_norm(a) + + n_keys = x.shape[1] + q = self.q_proj(x).view(bsz, n_queries, self.num_heads, self.head_dim) + kv = self.kv_proj(x) + k, v = kv.chunk(2, dim=-1) + k = k.view(bsz, n_keys, self.num_heads, self.head_dim) + v = v.view(bsz, n_keys, self.num_heads, self.head_dim) + + # Expand z for num_diffusion_samples + if z.dim() == 4 and z.shape[0] != bsz and num_diffusion_samples > 1: + z = z.repeat_interleave(num_diffusion_samples, dim=0) + if ( + attention_mask is not None + and attention_mask.shape[0] != bsz + and num_diffusion_samples > 1 + ): + attention_mask = attention_mask.repeat_interleave(num_diffusion_samples, dim=0) + + if self._can_use_fused_pair_bias(z, n_queries, beta): + kernel_mask = ( + attention_mask + if attention_mask is not None + else torch.ones(bsz, n_queries, device=a.device, dtype=torch.bool) + ) + pair_norm_w = self.pair_norm.weight + pair_norm_b = ( + self.pair_norm.bias + if self.pair_norm.bias is not None + else torch.zeros_like(pair_norm_w) + ) + z_bf = z if z.dtype == torch.bfloat16 else z.to(torch.bfloat16) + bias = _fused_pair_bias( # type: ignore[misc] + z_bf, + kernel_mask, + self.pair_bias_proj.weight, + num_heads=self.num_heads, + pair_norm_w=pair_norm_w, + pair_norm_b=pair_norm_b, + ) # A has shape (b, h, q, k). + q_bhqd = q.transpose(1, 2) + k_bhqd = k.transpose(1, 2) + v_bhqd = v.transpose(1, 2) + attn_out = F.scaled_dot_product_attention( + q_bhqd, k_bhqd, v_bhqd, attn_mask=bias.to(q_bhqd.dtype) + ) + g = torch.sigmoid(self.g_proj(x)).view(bsz, n_queries, self.num_heads, self.head_dim) + ctx = g * attn_out.transpose(1, 2) + out = self.out_proj(ctx.reshape(bsz, n_queries, d_model)) + if s is not None: + out = torch.sigmoid(self.out_gate(s)) * out + return out + + if self._can_use_cueq_pair_bias(z, n_queries, beta): + kernel_mask = ( + attention_mask + if attention_mask is not None + else torch.ones(bsz, n_queries, device=a.device, dtype=torch.bool) + ) + out, _ = _cue_attn_pair_bias( # type: ignore[misc] + s=x, + q=q.transpose(1, 2), + k=k.transpose(1, 2), + v=v.transpose(1, 2), + z=z, + mask=kernel_mask, + num_heads=self.num_heads, + w_proj_z=self.pair_bias_proj.weight, + w_proj_g=self.g_proj.weight, + w_proj_o=self.out_proj.weight, + w_ln_z=self.pair_norm.weight, + b_ln_z=self.pair_norm.bias, + return_z_proj=False, + is_cached_z_proj=False, + ) + else: + # Standard attention with pair bias + g = torch.sigmoid(self.g_proj(x)).view(bsz, n_queries, self.num_heads, self.head_dim) + + logits = torch.einsum("... i h d, ... j h d -> ... i j h", q, k) * self.scale + + pair_bias = self.pair_bias_proj(self.pair_norm(z)) if z.dim() == 4 else z.unsqueeze(-1) + logits = logits + pair_bias.to(dtype=logits.dtype) + + if attention_mask is not None: + min_val = torch.finfo(logits.dtype).min + mask_bias = torch.where(attention_mask.bool()[:, None, :, None], 0.0, min_val) + logits = logits + mask_bias.to(dtype=logits.dtype) + + attn = torch.softmax(logits, dim=-2).to(dtype=v.dtype) + ctx = torch.einsum("... i j h, ... j h d -> ... i h d", attn, v) + ctx = g * ctx + out = self.out_proj(ctx.reshape(bsz, n_queries, d_model)) + + if s is not None: + out = torch.sigmoid(self.out_gate(s)) * out + return out + + +# =========================================================================== +# ConditionedTransitionBlock +# =========================================================================== + + +class ConditionedTransitionBlock(nn.Module): + """Conditioned SwiGLU transition with adaptive layer norm.""" + + def __init__( + self, + d_model: int, + d_cond: int | None = None, + transition_multiplier: int = 2, + use_conditioning: bool = True, + ) -> None: + super().__init__() + d_cond = d_cond or d_model + hidden = transition_multiplier * d_model + + if use_conditioning: + self.adaln = AdaptiveLayerNorm(d_model, d_cond, eps=1e-5) + self.output_gate = nn.Linear(d_cond, d_model, bias=True) + nn.init.zeros_(self.output_gate.weight) + nn.init.constant_(self.output_gate.bias, -2.0) + else: + self.pre_norm = nn.LayerNorm(d_model, eps=1e-5) + + self.lin_swish = nn.Linear(d_model, 2 * hidden, bias=False) + self.lin_out = nn.Linear(hidden, d_model, bias=False) + + def forward(self, a: Tensor, s: Tensor | None) -> Tensor: + x = self.adaln(a, s) if s is not None else self.pre_norm(a) + + swish_a, swish_b = self.lin_swish(x).chunk(2, dim=-1) + b = F.silu(swish_a) * swish_b + out = self.lin_out(b) + + if s is not None: + out = torch.sigmoid(self.output_gate(s)) * out + return out + + +# =========================================================================== +# DiffusionTransformer (token transformer) +# =========================================================================== + + +class DiffusionTransformer(nn.Module): + """Diffusion denoising transformer with attention pair bias.""" + + def __init__( + self, + d_model: int, + d_pair: int, + num_heads: int, + num_blocks: int, + d_cond: int | None = None, + transition_multiplier: int = 2, + use_conditioning: bool = True, + ) -> None: + super().__init__() + d_cond = d_cond or d_model + + self.attn_blocks = nn.ModuleList( + [ + AttentionPairBias( + d_model=d_model, + d_pair=d_pair, + num_heads=num_heads, + d_cond=d_cond, + use_conditioning=use_conditioning, + ) + for _ in range(num_blocks) + ] + ) + self.transition_blocks = nn.ModuleList( + [ + ConditionedTransitionBlock( + d_model=d_model, + d_cond=d_cond, + transition_multiplier=transition_multiplier, + use_conditioning=use_conditioning, + ) + for _ in range(num_blocks) + ] + ) + + def set_kernel_backend(self, backend: str | None) -> None: + for attn in self.attn_blocks: + cast(AttentionPairBias, attn).set_kernel_backend(backend) + + def forward( + self, + a: Tensor, + s: Tensor | None, + z: Tensor, + beta: Tensor | float = 0.0, + attention_mask: Tensor | None = None, + num_diffusion_samples: int = 1, + return_intermediates: bool = False, + ) -> tuple[Tensor, list[Tensor]]: + intermediates: list[Tensor] = [] + x = a + for attn, transition in zip(self.attn_blocks, self.transition_blocks, strict=True): + x = x + attn( + x, + s, + z, + beta, + attention_mask=attention_mask, + num_diffusion_samples=num_diffusion_samples, + ) + x = x + transition(x, s) + if return_intermediates: + intermediates.append(x) + return x, intermediates + + +# =========================================================================== +# DiffusionConditioning +# =========================================================================== + + +class DiffusionConditioning(nn.Module): + """Conditions pair and single representations on noise timestep.""" + + def __init__( + self, + c_z: int = 256, + c_s: int = 768, + c_s_inputs: int = 451, + sigma_data: float = 16.0, + fourier_dim: int = 256, + transition_multiplier: int = 2, + layer_norm_eps: float = 1e-5, + ) -> None: + super().__init__() + self.sigma_data = float(sigma_data) + self.c_z = c_z + self.c_s = c_s + self.c_s_inputs = c_s_inputs + + self.z_input_norm = nn.LayerNorm(2 * c_z, eps=layer_norm_eps) + self.z_proj = nn.Linear(2 * c_z, c_z, bias=False) + self.z_transitions = nn.ModuleList( + [TransitionLayer(c_z, n=transition_multiplier, eps=layer_norm_eps) for _ in range(2)] + ) + + self.s_input_norm = nn.LayerNorm(c_s_inputs, eps=layer_norm_eps) + self.s_proj = nn.Linear(c_s_inputs, c_s, bias=False) + self.fourier = FourierEmbedding(fourier_dim) + self.noise_norm = nn.LayerNorm(fourier_dim, eps=layer_norm_eps) + self.noise_proj = nn.Linear(fourier_dim, c_s, bias=False) + self.s_transitions = nn.ModuleList( + [TransitionLayer(c_s, n=transition_multiplier, eps=layer_norm_eps) for _ in range(2)] + ) + + def forward( + self, + t_hat: Tensor, + s_inputs: Tensor, + s_trunk: Tensor | None, + z_trunk: Tensor, + relative_position_encoding: Tensor, + sigma_data: float | None = None, + num_diffusion_samples: int = 1, + inference_cache: dict[str, Tensor] | None = None, + ) -> tuple[Tensor, Tensor]: + sigma = self.sigma_data if sigma_data is None else float(sigma_data) + base_batch = z_trunk.shape[0] + target_batch = base_batch * num_diffusion_samples + + # z conditioning (cached across diffusion steps: independent of t_hat) + if inference_cache is not None and "z" in inference_cache: + z = inference_cache["z"] + else: + z_rel = relative_position_encoding.to(dtype=torch.float32) + z = torch.cat([z_trunk.to(dtype=torch.float32), z_rel], dim=-1) + z = self.z_proj(self.z_input_norm(z)) + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + for block in self.z_transitions: + z = z + block(z) + if inference_cache is not None: + inference_cache["z"] = z + + # s conditioning + s_inputs_eff = s_inputs + if s_inputs_eff.shape[0] != target_batch: + s_inputs_eff = s_inputs_eff.repeat_interleave(num_diffusion_samples, 0) + + s = self.s_proj(self.s_input_norm(s_inputs_eff.to(dtype=torch.float32))) + + # Noise embedding + t = torch.as_tensor(t_hat, dtype=torch.float32, device=s.device).reshape(-1) + if t.numel() == 1: + t = t.expand(target_batch) + elif t.shape[0] != target_batch: + t = t.repeat_interleave(num_diffusion_samples, 0) + t_noise = 0.25 * torch.log((t / sigma).clamp(min=1e-20)) + n = self.fourier(t_noise) + n = self.noise_proj(self.noise_norm(n)) + s = s + n.unsqueeze(1) + + for block in self.s_transitions: + s = s + block(s) + + return s, z + + +# =========================================================================== +# DiffusionModule +# =========================================================================== + + +class DiffusionModule(nn.Module): + """Diffusion denoising module for structure prediction.""" + + def __init__( + self, + c_atom: int = 128, + c_token: int = 768, + c_z: int = 256, + c_s_inputs: int = 451, + sigma_data: float = 16.0, + fourier_dim: int = 256, + atom_num_blocks: int = 3, + atom_num_heads: int = 4, + token_num_blocks: int = 12, + token_num_heads: int = 16, + transition_multiplier: int = 2, + swa_window_size: int = 128, + spatial_rope_base_frequency: float = 20.0, + n_spatial_rope_pairs_per_axis: int = 2, + n_uid_rope_pairs: int = 10, + uid_rope_base_frequency: float = 10000.0, + ) -> None: + super().__init__() + self.sigma_data = float(sigma_data) + + self.conditioning = DiffusionConditioning( + c_z=c_z, + c_s=c_token, # conditioning s output is c_token + c_s_inputs=c_s_inputs, + sigma_data=sigma_data, + fourier_dim=fourier_dim, + transition_multiplier=transition_multiplier, + ) + + # Atom encoder (structure_prediction=True, with coords_linear) + self.atom_encoder = ESMFold2AtomEncoder( + d_atom=c_atom, + d_token=c_token, + n_blocks=atom_num_blocks, + n_heads=atom_num_heads, + swa_window_size=swa_window_size, + expansion_ratio=2, + structure_prediction=True, + spatial_rope_base_frequency=spatial_rope_base_frequency, + n_spatial_rope_pairs_per_axis=n_spatial_rope_pairs_per_axis, + n_uid_rope_pairs=n_uid_rope_pairs, + uid_rope_base_frequency=uid_rope_base_frequency, + ) + + # Atom decoder + self.atom_decoder = ESMFold2AtomDecoder( + d_atom=c_atom, + d_token=c_token, + n_blocks=atom_num_blocks, + n_heads=atom_num_heads, + swa_window_size=swa_window_size, + expansion_ratio=2, + spatial_rope_base_frequency=spatial_rope_base_frequency, + n_spatial_rope_pairs_per_axis=n_spatial_rope_pairs_per_axis, + n_uid_rope_pairs=n_uid_rope_pairs, + uid_rope_base_frequency=uid_rope_base_frequency, + ) + + self.s_to_token = nn.Linear(c_token, c_token, bias=False) + nn.init.zeros_(self.s_to_token.weight) + + # Token transformer (DiffusionTransformer with pair bias) + self.token_transformer = DiffusionTransformer( + d_model=c_token, + d_pair=c_z, + num_heads=token_num_heads, + num_blocks=token_num_blocks, + d_cond=c_token, + transition_multiplier=transition_multiplier, + use_conditioning=True, + ) + + self.s_step_norm = nn.LayerNorm(c_token) + self.token_norm = nn.LayerNorm(c_token) + + def set_kernel_backend(self, backend: str | None) -> None: + self.token_transformer.set_kernel_backend(backend) + + def forward( + self, + x_noisy: Tensor, + t_hat: Tensor, + ref_pos: Tensor, + ref_charge: Tensor, + ref_mask: Tensor, + ref_element: Tensor, + ref_atom_name_chars: Tensor, + ref_space_uid: Tensor, + tok_idx: Tensor, + s_inputs: Tensor, + s_trunk: Tensor | None, + z_trunk: Tensor, + relative_position_encoding: Tensor, + asym_id: Tensor, + residue_index: Tensor, + entity_id: Tensor, + token_index: Tensor, + sym_id: Tensor, + sigma_data: float | None = None, + token_attention_mask: Tensor | None = None, + num_diffusion_samples: int = 1, + return_token_repr: bool = False, + return_atom_repr: bool = False, + inference_cache: dict[str, Tensor] | None = None, + ) -> dict[str, Tensor | None]: + bsz = x_noisy.shape[0] + sigma = self.sigma_data if sigma_data is None else float(sigma_data) + t = torch.as_tensor(t_hat, dtype=torch.float32, device=x_noisy.device).reshape(-1) + if t.numel() == 1: + t = t.expand(bsz) + + # Step 1: conditioning (pair z is cached across diffusion steps) + s, z = self.conditioning( + t_hat=t, + s_inputs=s_inputs, + s_trunk=s_trunk, + z_trunk=z_trunk, + relative_position_encoding=relative_position_encoding, + sigma_data=sigma, + num_diffusion_samples=num_diffusion_samples, + inference_cache=inference_cache, + ) + + # Step 2: normalize noisy coords + denom = torch.sqrt(t * t + sigma * sigma) + r_noisy = x_noisy / denom[:, None, None] + + # Step 3: atom encoder + a, q_skip, c_skip, p_skip, enc_intermediates = self.atom_encoder( + ref_pos=ref_pos, + atom_attention_mask=ref_mask, + ref_space_uid=ref_space_uid, + ref_charge=ref_charge, + ref_element=ref_element, + ref_atom_name_chars=ref_atom_name_chars, + atom_to_token=tok_idx, + r_l=r_noisy, + s_i=s_trunk, + num_diffusion_samples=num_diffusion_samples, + return_intermediates=return_atom_repr, + inference_cache=inference_cache, + ) + + # Step 4: add conditioned s + a = a + self.s_to_token(self.s_step_norm(s)) + + # Step 5: token transformer + a, _ = self.token_transformer( + a, + s, + z, + beta=0.0, + attention_mask=token_attention_mask, + num_diffusion_samples=num_diffusion_samples, + ) + + # Step 6: token norm + a = self.token_norm(a) + + # Step 7: atom decoder + r_update, dec_intermediates = self.atom_decoder( + a_i=a, + q_l=q_skip, + c_l=c_skip, + p_lm=p_skip, + atom_to_token=tok_idx, + atom_attention_mask=ref_mask, + num_diffusion_samples=num_diffusion_samples, + return_intermediates=return_atom_repr, + ) + + # Step 8: compute denoised output + sigma2 = sigma * sigma + t2 = t * t + out = (sigma2 / (sigma2 + t2))[:, None, None] * x_noisy + out = out + ((sigma * t) / torch.sqrt(sigma2 + t2))[:, None, None] * r_update + + # Collect atom intermediates from encoder + decoder + atom_intermediates: Tensor | None = None + if return_atom_repr: + all_ints = enc_intermediates + dec_intermediates + if all_ints: + atom_intermediates = torch.stack(all_ints, dim=2) + + return { + "x_denoised": out, + "token_repr": a if return_token_repr else None, + "atom_intermediates": atom_intermediates, + } + + +# =========================================================================== +# DiffusionStructureHead +# =========================================================================== + + +class DiffusionStructureHead(nn.Module): + """Wrapper around DiffusionModule with diffusion sampling.""" + + def __init__(self, config: ESMFold2Config) -> None: + super().__init__() + dm = config.structure_head.diffusion_module + swa_cfg = config.inputs.atom_encoder + sh = config.structure_head + + self.diffusion_module = DiffusionModule( + c_atom=dm.c_atom, + c_token=dm.c_token, + c_z=dm.c_z, + c_s_inputs=dm.c_s_inputs, + sigma_data=dm.sigma_data, + fourier_dim=dm.fourier_dim, + atom_num_blocks=dm.atom_num_blocks, + atom_num_heads=dm.atom_num_heads, + token_num_blocks=dm.token_num_blocks, + token_num_heads=dm.token_num_heads, + transition_multiplier=dm.transition_multiplier, + swa_window_size=swa_cfg.swa_window_size, + spatial_rope_base_frequency=swa_cfg.spatial_rope_base_frequency, + n_spatial_rope_pairs_per_axis=swa_cfg.n_spatial_rope_pairs_per_axis, + n_uid_rope_pairs=swa_cfg.n_uid_rope_pairs, + uid_rope_base_frequency=swa_cfg.uid_rope_base_frequency, + ) + + # Sampling hyperparameters + self.sigma_data = dm.sigma_data + self.gamma_0 = sh.gamma_0 + self.gamma_min = sh.gamma_min + self.noise_scale = sh.noise_scale + self.step_scale = sh.step_scale + self.inference_s_max = sh.inference_s_max + self.inference_s_min = sh.inference_s_min + self.inference_p = sh.inference_p + self.inference_num_steps = sh.inference_num_steps + + def set_kernel_backend(self, backend: str | None) -> None: + self.diffusion_module.set_kernel_backend(backend) + + # ------------------------------------------------------------------ + # Helpers + # ------------------------------------------------------------------ + + def inference_noise_schedule( + self, num_steps: int | None = None, device: torch.device | None = None + ) -> Tensor: + """Karras power-law noise schedule.""" + steps = self.inference_num_steps if num_steps is None else int(num_steps) + if steps == 1: + return torch.tensor( + [self.inference_s_max * self.sigma_data, 0.0], + device=device, + dtype=torch.float32, + ) + p = float(self.inference_p) + inv_p = 1.0 / p + k = torch.arange(steps, device=device, dtype=torch.float32) + base = self.inference_s_max**inv_p + (k / (steps - 1)) * ( + self.inference_s_min**inv_p - self.inference_s_max**inv_p + ) + schedule = self.sigma_data * base.pow(p) + return F.pad(schedule, (0, 1), value=0.0) + + @staticmethod + def _random_rotations(n: int, dtype: torch.dtype, device: torch.device) -> Tensor: + q = torch.randn((n, 4), dtype=dtype, device=device) + scale = torch.sqrt((q * q).sum(dim=1)) + signs = torch.where(q[:, 0] < 0, -scale, scale) + q = q / signs[:, None] + r, i, j, k = torch.unbind(q, dim=-1) + two_s = 2.0 / (q * q).sum(dim=-1) + return torch.stack( + ( + 1 - two_s * (j * j + k * k), + two_s * (i * j - k * r), + two_s * (i * k + j * r), + two_s * (i * j + k * r), + 1 - two_s * (i * i + k * k), + two_s * (j * k - i * r), + two_s * (i * k - j * r), + two_s * (j * k + i * r), + 1 - two_s * (i * i + j * j), + ), + dim=-1, + ).reshape(n, 3, 3) + + def _center_random_augmentation( + self, x: Tensor, atom_mask: Tensor, second_coords: Tensor | None = None + ) -> tuple[Tensor, Tensor | None]: + """Algorithm 19: center + random rotation + translation.""" + bsz = x.shape[0] + mask = atom_mask.unsqueeze(-1) # M has shape (b, a, 1). + denom = mask.sum(dim=1, keepdim=True).clamp(min=1) + mean = (x * mask).sum(dim=1, keepdim=True) / denom + x = x - mean + if second_coords is not None: + second_coords = second_coords - mean + + r = self._random_rotations(bsz, x.dtype, x.device) + x = torch.einsum("bmd,bds->bms", x, r) + if second_coords is not None: + second_coords = torch.einsum("bmd,bds->bms", second_coords, r) + + t = torch.randn_like(x[:, 0:1, :]) + x = x + t + if second_coords is not None: + second_coords = second_coords + t + return x, second_coords + + @staticmethod + def _weighted_rigid_align(x: Tensor, x_gt: Tensor, w: Tensor, mask: Tensor) -> Tensor: + """Kabsch alignment: align x to x_gt with weights w.""" + w = (mask * w).unsqueeze(-1) # W has shape (b, n, 1). + denom = w.sum(dim=-2, keepdim=True).clamp(min=1e-8) + mu = (x * w).sum(dim=-2, keepdim=True) / denom + mu_gt = (x_gt * w).sum(dim=-2, keepdim=True) / denom + x_c = x - mu + xgt_c = x_gt - mu_gt + covariance = torch.einsum("bni,bnj->bij", w * xgt_c, x_c) + covariance_f32 = covariance.float() + u, _, vh = torch.linalg.svd( + covariance_f32, driver="gesvd" if covariance_f32.is_cuda else None + ) + det = torch.linalg.det(u @ vh) + ones = torch.ones_like(det) + rotation = (u @ torch.diag_embed(torch.stack([ones, ones, det], dim=-1)) @ vh).to( + covariance.dtype + ) + return x_c @ rotation.transpose(-1, -2) + mu_gt + + # ------------------------------------------------------------------ + # Sampling + # ------------------------------------------------------------------ + + @torch.inference_mode() + def sample( + self, + z_trunk: Tensor, + s_inputs: Tensor, + s_trunk: Tensor | None, + relative_position_encoding: Tensor, + ref_pos: Tensor, + ref_charge: Tensor, + ref_mask: Tensor, + ref_element: Tensor, + ref_atom_name_chars: Tensor, + ref_space_uid: Tensor, + tok_idx: Tensor, + asym_id: Tensor, + residue_index: Tensor, + entity_id: Tensor, + token_index: Tensor, + sym_id: Tensor, + token_attention_mask: Tensor | None = None, + num_diffusion_samples: int = 1, + num_sampling_steps: int | None = None, + max_inference_sigma: float | None = 256.0, + noise_scale: float | None = None, + step_scale: float | None = None, + return_atom_repr: bool = False, + use_inference_cache: bool = True, + denoising_early_exit_rmsd: float | None = None, + ) -> dict[str, Tensor | None]: + """Diffusion sampling (Algorithm 18). + + ``num_sampling_steps`` is the number of denoising steps actually run. + When ``max_inference_sigma`` is set, the Karras schedule built with + ``num_sampling_steps`` entries would lose its high-sigma tail to the cap, + so we inflate the underlying schedule length here to land back at the + requested step count post-truncation. + """ + n_atoms = tok_idx.shape[1] + device = s_inputs.device + target_batch = s_inputs.shape[0] * num_diffusion_samples + + inference_cache: dict[str, Tensor] | None = {} if use_inference_cache else None + + steps = self.inference_num_steps if num_sampling_steps is None else int(num_sampling_steps) + + schedule = self.inference_noise_schedule(steps, device) + if max_inference_sigma is not None: + schedule = schedule[schedule <= float(max_inference_sigma)] + schedule = F.pad(schedule, (1, 0), value=float(max_inference_sigma)) + + lam = self.noise_scale if noise_scale is None else float(noise_scale) + eta = self.step_scale if step_scale is None else float(step_scale) + + x = schedule[0] * torch.randn(target_batch, n_atoms, 3, device=device, dtype=torch.float32) + atom_mask = ref_mask.repeat_interleave(num_diffusion_samples, 0).float() + + gammas = torch.where( + schedule > self.gamma_min, + torch.full_like(schedule, self.gamma_0), + torch.zeros_like(schedule), + ) + + x_denoised_prev: Tensor | None = None + token_repr: Tensor | None = None + diff_atom_intermediates: Tensor | None = None + + step_pairs = list(zip(schedule[:-1], schedule[1:], gammas[1:], strict=True)) + num_steps = len(step_pairs) + + for step_idx, (sigma_tm, sigma_t, gamma) in enumerate(step_pairs): + x, x_denoised_prev = self._center_random_augmentation( + x, atom_mask, second_coords=x_denoised_prev + ) + + sigma_tm_val = float(sigma_tm.item()) + t_hat_val = sigma_tm_val * (1.0 + float(gamma.item())) + eps_std = lam * max(t_hat_val**2 - sigma_tm_val**2, 0.0) ** 0.5 + x_noisy = x + eps_std * torch.randn_like(x) + + is_last_step = step_idx == num_steps - 1 + request_atom_repr = return_atom_repr and ( + is_last_step or denoising_early_exit_rmsd is not None + ) + + dm_out = self.diffusion_module( + x_noisy=x_noisy, + t_hat=torch.full((target_batch,), t_hat_val, device=device, dtype=torch.float32), + ref_pos=ref_pos, + ref_charge=ref_charge, + ref_mask=ref_mask, + ref_element=ref_element, + ref_atom_name_chars=ref_atom_name_chars, + ref_space_uid=ref_space_uid, + tok_idx=tok_idx, + s_inputs=s_inputs, + s_trunk=s_trunk, + z_trunk=z_trunk, + relative_position_encoding=relative_position_encoding, + asym_id=asym_id, + residue_index=residue_index, + entity_id=entity_id, + token_index=token_index, + sym_id=sym_id, + token_attention_mask=token_attention_mask, + num_diffusion_samples=num_diffusion_samples, + return_token_repr=True, + return_atom_repr=request_atom_repr, + inference_cache=inference_cache, + ) + + x_denoised = dm_out["x_denoised"] + token_repr = dm_out["token_repr"] + if request_atom_repr: + diff_atom_intermediates = dm_out.get("atom_intermediates") + + # Reverse diffusion alignment (Kabsch) + with torch.autocast(device_type="cuda", enabled=False): + x_noisy = self._weighted_rigid_align( + x_noisy.float(), x_denoised.float(), atom_mask, atom_mask + ) + x_noisy = x_noisy.to(dtype=x_denoised.dtype) + + # ODE/SDE step + sigma_t_val = float(sigma_t.item()) + denoised_over_sigma = (x_noisy - x_denoised) / t_hat_val + x = x_noisy + eta * (sigma_t_val - t_hat_val) * denoised_over_sigma + + # Denoising early-exit: stop when consecutive predictions converge + if ( + denoising_early_exit_rmsd is not None + and x_denoised_prev is not None + and step_idx >= 1 + ): + with torch.autocast(device_type="cuda", enabled=False): + aligned = self._weighted_rigid_align( + x_denoised_prev.float(), + x_denoised.float(), + atom_mask, + atom_mask, + ) + diff = (x_denoised.float() - aligned) * atom_mask.unsqueeze(-1) + per_sample_rmsd = ( + diff.pow(2).sum(dim=(-1, -2)) / atom_mask.sum(dim=-1).clamp(min=1) + ).sqrt() + if per_sample_rmsd.max().item() < denoising_early_exit_rmsd: + x = x_denoised + x_denoised_prev = x_denoised + break + + x_denoised_prev = x_denoised + + result: dict[str, Tensor | None] = { + "sample_atom_coords": x, + "diff_token_repr": token_repr, + } + if return_atom_repr: + result["diff_atom_intermediates"] = diff_atom_intermediates + return result diff --git a/fastplms/models/esmfold2/modeling_esmfold2_experimental.py b/fastplms/models/esmfold2/modeling_esmfold2_experimental.py new file mode 100644 index 0000000000000000000000000000000000000000..ad90ba7c7d91ba6e87189afb7347e02edc105a07 --- /dev/null +++ b/fastplms/models/esmfold2/modeling_esmfold2_experimental.py @@ -0,0 +1,1083 @@ +"""FastPLMs ESMFold2 experimental architecture. + +This module supports Biohub's experimental binder-design checkpoints. The +released ESMFold2 architecture in ``modeling_esmfold2.py`` intentionally +rejects those configs because the experimental trunk uses explicit pair-loop +re-injection and a different confidence/MSA stack. +""" + +from __future__ import annotations + +import gc +from pathlib import Path +from typing import Any, ClassVar, cast + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from transformers.modeling_utils import PreTrainedModel + +from .attention import ESMFold2AttentionMixin +from .configuration_esmfold2 import ESMFold2Config +from .embedding import ESMFold2EmbeddingMixin +from .modeling_esmfold2 import ( + ESMCPrecision, + ESMCPrecisionStatus, + ESMFold2Output, + _drop_transient_esmc_state, + _finalize_structure_output, + _install_esmc_backbone, + _lm_precision_context, + _reload_esmc_bf16_for_gradients, + _resolve_structure_output_controls, + _transformer_engine_version, +) +from .modeling_esmfold2_common import ( + CHAR_VOCAB_SIZE, + MAX_ATOMIC_NUMBER, + MSA_CONDITIONING_INPUT_NAMES, + NUM_RES_TYPES, + DiffusionModule, + DiffusionStructureHead, + DiffusionTransformer, + FoldingTrunk, + InputsEmbedder, + LanguageModelShim, + MSAPairWeightedAveraging, + OuterProductMean, + PairUpdateBlock, + ResIdxAsymIdSymIdEntityIdEncoding, + RowAttentionPooling, + SwiGLUMLP, + TriangleMultiplicativeUpdate, + _categorical_mean, + _compute_intra_token_idx, + _seed_context, + compute_lm_hidden_states, + gather_rep_atom_coords, + gather_token_to_atom, + validate_kernel_backend, + validate_msa_conditioning_inputs, + validate_prepared_auxiliary_inputs, +) + +_EPS = 1e-5 +_NONPOLYMER_ID = 3 + + +class ConfidenceHead(nn.Module): + """Experimental confidence head predicting pLDDT, PAE, pTM, and ipTM.""" + + boundaries: Tensor + + def __init__(self, config: ESMFold2Config) -> None: + super().__init__() + ch = config.confidence_head + d_single = config.d_single + d_pair = config.d_pair + d_inputs = config.inputs.d_inputs + + boundaries = torch.linspace(ch.min_dist, ch.max_dist, ch.distogram_bins - 1) + self.register_buffer("boundaries", boundaries) + self.dist_bin_pairwise_embed = nn.Embedding(ch.distogram_bins, d_pair) + + self.s_norm = nn.LayerNorm(d_single) + self.s_inputs_to_single = nn.Linear(d_inputs, d_single, bias=False) + self.s_to_z = nn.Linear(d_inputs, d_pair, bias=False) + self.s_to_z_transpose = nn.Linear(d_inputs, d_pair, bias=False) + self.s_to_z_prod_in1 = nn.Linear(d_inputs, d_pair, bias=False) + self.s_to_z_prod_in2 = nn.Linear(d_inputs, d_pair, bias=False) + self.s_to_z_prod_out = nn.Linear(d_pair, d_pair, bias=False) + self.s_input_to_s = nn.Linear(d_inputs, d_single, bias=False) + self.s_inputs_norm = nn.LayerNorm(d_inputs) + self.z_norm = nn.LayerNorm(d_pair) + self.row_attention_pooling = RowAttentionPooling(d_pair=d_pair, d_single=d_single) + + pf = ch.folding_trunk + self.folding_trunk = FoldingTrunk(n_layers=pf.n_layers, d_pair=d_pair, expansion_ratio=4) + + self.plddt_ln = nn.LayerNorm(d_single) + max_atoms_per_token = 23 + self.plddt_weight = nn.Parameter( + torch.zeros(max_atoms_per_token, d_single, ch.num_plddt_bins) + ) + self.pae_head = nn.Linear(d_pair, ch.num_pae_bins, bias=False) + + def set_kernel_backend(self, backend: str | None) -> None: + validate_kernel_backend(backend) + self.folding_trunk.set_kernel_backend(backend) + + def set_chunk_size(self, chunk_size: int | None) -> None: + self.folding_trunk.set_chunk_size(chunk_size) + + @staticmethod + def _repeat_batch(x: Tensor, num_diffusion_samples: int) -> Tensor: + if num_diffusion_samples == 1: + return x + return x.repeat_interleave(num_diffusion_samples, 0) + + @staticmethod + def _flatten_sample_axis(x: Tensor) -> Tensor: + if x.ndim == 4: + b, mult, n, c = x.shape + return x.reshape(b * mult, n, c) + return x + + def forward( + self, + s_inputs: Tensor, + z: Tensor, + x_pred: Tensor, + distogram_atom_idx: Tensor, + token_attention_mask: Tensor, + atom_to_token: Tensor, + atom_attention_mask: Tensor, + asym_id: Tensor, + mol_type: Tensor, + num_diffusion_samples: int = 1, + relative_position_encoding: Tensor | None = None, + token_bonds_encoding: Tensor | None = None, + ) -> dict[str, Tensor]: + s_inputs_normed = self.s_inputs_norm(s_inputs) + z_base = self.z_norm(z) + if relative_position_encoding is not None: + z_base = z_base + relative_position_encoding + if token_bonds_encoding is not None: + z_base = z_base + token_bonds_encoding + z_base = z_base + self.s_to_z(s_inputs_normed).unsqueeze(2) + z_base = z_base + self.s_to_z_transpose(s_inputs_normed).unsqueeze(1) + z_base = z_base + self.s_to_z_prod_out( + self.s_to_z_prod_in1(s_inputs_normed)[:, :, None, :] + * self.s_to_z_prod_in2(s_inputs_normed)[:, None, :, :] + ) + + pair = self._repeat_batch(z_base, num_diffusion_samples) + x_pred_flat = self._flatten_sample_axis(x_pred) + atom_to_token_m = self._repeat_batch(atom_to_token, num_diffusion_samples) + atom_mask_m = self._repeat_batch(atom_attention_mask, num_diffusion_samples) + rep_idx_m = self._repeat_batch(distogram_atom_idx, num_diffusion_samples).long() + mask = self._repeat_batch(token_attention_mask, num_diffusion_samples) + batch_mult = pair.shape[0] + + rep_coords = gather_rep_atom_coords(x_pred_flat, rep_idx_m) + rep_distances = torch.cdist( + rep_coords, rep_coords, compute_mode="donot_use_mm_for_euclid_dist" + ) + distogram_bins = (rep_distances.unsqueeze(-1) > self.boundaries).sum(dim=-1).long() + pair = pair + self.dist_bin_pairwise_embed(distogram_bins) + + pair_mask = mask[:, :, None].float() * mask[:, None, :].float() + pair = pair + self.folding_trunk(pair, pair_attention_mask=pair_mask) + single = self.row_attention_pooling(pair, mask) + + atom_mask_f = atom_mask_m.float() + s_at_atoms = gather_token_to_atom(single, atom_to_token_m) + s_at_atoms = self.plddt_ln(s_at_atoms) + intra_idx = _compute_intra_token_idx(atom_to_token_m) + intra_idx = intra_idx.clamp(max=self.plddt_weight.shape[0] - 1) + plddt_weight = self.plddt_weight[intra_idx] + plddt_logits = torch.einsum("...c,...cb->...b", s_at_atoms, plddt_weight) + plddt_per_atom = _categorical_mean(plddt_logits, start=0.0, end=1.0) + + length = single.shape[1] + plddt_sum = torch.zeros( + batch_mult, length, device=single.device, dtype=plddt_per_atom.dtype + ) + atom_count = torch.zeros( + batch_mult, length, device=single.device, dtype=plddt_per_atom.dtype + ) + atom_mask_t = atom_mask_f.to(plddt_per_atom.dtype) + plddt_sum.scatter_add_(1, atom_to_token_m, plddt_per_atom * atom_mask_t) + atom_count.scatter_add_(1, atom_to_token_m, atom_mask_t) + plddt = plddt_sum / atom_count.clamp(min=1e-6) + + complex_plddt = (plddt_per_atom * atom_mask_f).sum(dim=-1) / ( + atom_mask_f.sum(dim=-1) + _EPS + ) + + expanded_type = self._repeat_batch(mol_type, num_diffusion_samples) + expanded_asym = self._repeat_batch(asym_id, num_diffusion_samples) + is_ligand = (expanded_type == _NONPOLYMER_ID).float() + inter_chain = (expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2)).float() + near_contact = (rep_distances < 8).float() + interface_per_token = (near_contact * inter_chain * (1.0 - is_ligand).unsqueeze(-1)).amax( + dim=-1 + ) + iplddt_weight = torch.where( + is_ligand.bool(), + torch.full_like(interface_per_token, 2.0), + interface_per_token, + ) + iplddt_weight_atoms = gather_token_to_atom( + iplddt_weight.unsqueeze(-1), atom_to_token_m + ).squeeze(-1) + atom_iplddt_w = atom_mask_f * iplddt_weight_atoms + complex_iplddt = (plddt_per_atom * atom_iplddt_w).sum(dim=-1) / ( + atom_iplddt_w.sum(dim=-1) + _EPS + ) + plddt_ca = plddt_per_atom.gather(1, rep_idx_m) + + pae_logits = self.pae_head(pair) + pae = _categorical_mean(pae_logits, start=0.0, end=32.0).detach() + + n_bins = pae_logits.shape[-1] + bin_width = 32.0 / n_bins + bin_centers = torch.arange(0.5 * bin_width, 32.0, bin_width, device=pae_logits.device) + mask_f = mask.float() + n_res = mask_f.sum(dim=-1, keepdim=True) + d0 = 1.24 * (n_res.clamp(min=19) - 15) ** (1 / 3) - 1.8 + tm_per_bin = 1 / (1 + (bin_centers / d0) ** 2) + pae_probs = F.softmax(pae_logits, dim=-1) + tm_expected = (pae_probs * tm_per_bin[:, None, None, :]).sum(dim=-1) + + pair_mask_2d = mask_f.unsqueeze(-1) * mask_f.unsqueeze(-2) + ptm_per_row = (tm_expected * pair_mask_2d).sum(dim=-1) / (pair_mask_2d.sum(dim=-1) + _EPS) + ptm = ptm_per_row.max(dim=-1).values + + inter_chain_mask = ( + expanded_asym.unsqueeze(-1) != expanded_asym.unsqueeze(-2) + ).float() * pair_mask_2d + iptm_per_row = (tm_expected * inter_chain_mask).sum(dim=-1) / ( + inter_chain_mask.sum(dim=-1) + _EPS + ) + iptm = iptm_per_row.max(dim=-1).values + + max_chain_id = int(expanded_asym.max().item()) if batch_mult > 0 else 0 + n_chains = max_chain_id + 1 + pair_chains_iptm = torch.zeros( + batch_mult, + n_chains, + n_chains, + device=tm_expected.device, + dtype=tm_expected.dtype, + ) + for c1 in range(n_chains): + chain_c1 = (expanded_asym == c1).float() * mask_f + if chain_c1.sum() == 0: + continue + for c2 in range(n_chains): + chain_c2 = (expanded_asym == c2).float() * mask_f + pair_m = chain_c1.unsqueeze(-1) * chain_c2.unsqueeze(-2) + denom = pair_m.sum(dim=(-1, -2)) + _EPS + pair_chains_iptm[:, c1, c2] = (tm_expected * pair_m).sum(dim=(-1, -2)) / denom + + return { + "plddt_logits": plddt_logits, + "plddt": plddt.detach(), + "plddt_per_atom": plddt_per_atom.detach(), + "plddt_ca": plddt_ca.detach(), + "complex_plddt": complex_plddt.detach(), + "complex_iplddt": complex_iplddt.detach(), + "pae_logits": pae_logits, + "pae": pae, + "ptm": ptm.detach(), + "iptm": iptm.detach(), + "pair_chains_iptm": pair_chains_iptm.detach(), + } + + +class _TransitionFFN(nn.Module): + def __init__(self, d_model: int, expansion_ratio: int = 4) -> None: + super().__init__() + self.norm = nn.LayerNorm(d_model) + self.ffn = SwiGLUMLP(d_model, expansion_ratio=expansion_ratio, bias=False) + + def forward(self, x: Tensor) -> Tensor: + return self.ffn(self.norm(x)) + + +class MSAEncoderBlock(nn.Module): + """One experimental MSA update block.""" + + def __init__( + self, + d_msa: int, + d_pair: int, + d_hidden: int = 32, + n_heads_msa: int = 8, + msa_head_width: int = 32, + ) -> None: + super().__init__() + self.outer_product_mean = OuterProductMean( + d_msa, d_hidden, d_pair, divide_outer_before_proj=True + ) + self.msa_pair_weighted_averaging = MSAPairWeightedAveraging( + d_msa, d_pair, n_heads_msa, msa_head_width + ) + self.msa_transition = _TransitionFFN(d_msa, expansion_ratio=4) + self.tri_mul_out = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=True) + self.tri_mul_in = TriangleMultiplicativeUpdate(dim=d_pair, _outgoing=False) + self.pair_transition = _TransitionFFN(d_pair, expansion_ratio=4) + + def set_chunk_size(self, chunk_size: int | None) -> None: + self.outer_product_mean.set_chunk_size(chunk_size) + self.tri_mul_out.set_chunk_size(chunk_size) + self.tri_mul_in.set_chunk_size(chunk_size) + + def forward( + self, + msa_repr: Tensor, + pair_repr: Tensor, + msa_attention_mask: Tensor, + pair_attention_mask: Tensor, + msa_track_mask: Tensor | None = None, + ) -> tuple[Tensor, Tensor]: + mask4d = ( + msa_track_mask[:, None, None, None].to(dtype=msa_repr.dtype) + if msa_track_mask is not None + else None + ) + + pair_mask4d = mask4d[:, :, :1] if mask4d is not None else None + + msa_update = self.msa_pair_weighted_averaging(msa_repr, pair_repr, pair_attention_mask) + if mask4d is not None: + msa_update = msa_update * mask4d + msa_repr = msa_repr + msa_update + + msa_transition = self.msa_transition(msa_repr) + if mask4d is not None: + msa_transition = msa_transition * mask4d + msa_repr = msa_repr + msa_transition + + pair_opm = self.outer_product_mean(msa_repr, msa_attention_mask) + if pair_mask4d is not None: + pair_opm = pair_opm * pair_mask4d + pair_repr = pair_repr + pair_opm + + pair_out = self.tri_mul_out(pair_repr, mask=pair_attention_mask) + if pair_mask4d is not None: + pair_out = pair_out * pair_mask4d + pair_repr = pair_repr + pair_out + + pair_in = self.tri_mul_in(pair_repr, mask=pair_attention_mask) + if pair_mask4d is not None: + pair_in = pair_in * pair_mask4d + pair_repr = pair_repr + pair_in + + pair_transition = self.pair_transition(pair_repr) + if pair_mask4d is not None: + pair_transition = pair_transition * pair_mask4d + pair_repr = pair_repr + pair_transition + return msa_repr, pair_repr + + +class MSAEncoder(nn.Module): + def __init__( + self, + d_msa: int, + d_pair: int, + d_inputs: int, + d_hidden: int = 32, + n_layers: int = 4, + n_heads_msa: int = 8, + msa_head_width: int = 32, + ) -> None: + super().__init__() + self.embed = nn.Linear(35, d_msa, bias=False) + self.project_inputs = nn.Linear(d_inputs, d_msa, bias=False) + self.blocks = nn.ModuleList( + [ + MSAEncoderBlock( + d_msa=d_msa, + d_pair=d_pair, + d_hidden=d_hidden, + n_heads_msa=n_heads_msa, + msa_head_width=msa_head_width, + ) + for _ in range(n_layers) + ] + ) + + def set_chunk_size(self, chunk_size: int | None) -> None: + for block in self.blocks: + cast(MSAEncoderBlock, block).set_chunk_size(chunk_size) + + def forward( + self, + x_pair: Tensor, + x_inputs: Tensor, + msa_oh: Tensor, + has_deletion: Tensor, + deletion_value: Tensor, + msa_attention_mask: Tensor, + ) -> Tensor: + batch_size, _, depth = msa_attention_mask.shape + m_feat = torch.cat( + [msa_oh, has_deletion.unsqueeze(-1), deletion_value.unsqueeze(-1)], + dim=-1, + ) + m = self.embed(m_feat) + self.project_inputs(x_inputs).unsqueeze(2) + if depth > 1: + msa_track_mask = msa_attention_mask[:, :, 1:].any(dim=(1, 2)) + else: + msa_track_mask = torch.zeros(batch_size, dtype=torch.bool, device=x_pair.device) + tok_mask = msa_attention_mask[:, :, 0] + pair_attention_mask = tok_mask.unsqueeze(2) * tok_mask.unsqueeze(1) + for block in self.blocks: + m, x_pair = cast(MSAEncoderBlock, block)( + m, + x_pair, + msa_attention_mask, + pair_attention_mask, + msa_track_mask, + ) + return x_pair * msa_track_mask[:, None, None, None].to(dtype=x_pair.dtype) + + +class ESMFold2ExperimentalModel(ESMFold2EmbeddingMixin, ESMFold2AttentionMixin, PreTrainedModel): + """Experimental ESMFold2 architecture used by binder-design checkpoints.""" + + config_class = ESMFold2Config + _keys_to_ignore_on_load_unexpected: ClassVar[list[str]] = [r"\._extra_state$"] + + def __init__(self, config: ESMFold2Config) -> None: + super().__init__(config) + d_inputs = config.inputs.d_inputs + d_pair = config.d_pair + + self.inputs_embedder = InputsEmbedder(config) + self.z_init_1 = nn.Linear(d_inputs, d_pair, bias=False) + self.z_init_2 = nn.Linear(d_inputs, d_pair, bias=False) + self.rel_pos = ResIdxAsymIdSymIdEntityIdEncoding( + n_relative_residx_bins=config.n_relative_residx_bins, + n_relative_chain_bins=config.n_relative_chain_bins, + d_pair=d_pair, + ) + self.token_bonds = nn.Linear(1, d_pair, bias=False) + self.language_model = LanguageModelShim( + d_z=d_pair, d_model=config.lm_d_model, num_layers=config.lm_num_layers + ) + self._esmc: nn.Module | None = None + self._esmc_fp8 = False + self._esmc_fp8_module_paths: tuple[str, ...] = () + self._esmc_source: str = config.esmc_id + self._esmc_source_revision: str | None = None + self._esmc_source_files: dict[str, str] = {} + self._esmc_local_files_only = False + self._esmc_precision_policy: str = str(getattr(config, "esmc_precision", "auto")) + self._esmc_precision_status = ESMCPrecisionStatus( + requested=self._esmc_precision_policy, + resolved="unloaded", + reason="ESMC has not been loaded.", + device=str(self.device), + transformer_engine_version=_transformer_engine_version(), + ) + self._ttt_lm_head: nn.Module | None = None + self._esmfold2_input_builder: Any | None = None + self._kernel_backend: str | None = None + + pf = config.folding_trunk + self.folding_trunk = FoldingTrunk(n_layers=pf.n_layers, d_pair=d_pair, expansion_ratio=4) + self.pair_loop_proj = nn.Sequential( + nn.LayerNorm(d_pair), nn.Linear(d_pair, d_pair, bias=False) + ) + nn.init.zeros_(cast(nn.Linear, self.pair_loop_proj[1]).weight) + + self.structure_head = DiffusionStructureHead(config) + self.distogram_head = nn.Linear(d_pair, config.structure_head.distogram_bins, bias=True) + self.confidence_head: ConfidenceHead | None = ( + ConfidenceHead(config) if config.confidence_head.enabled else None + ) + + msa_cfg = config.msa_encoder + self.msa_encoder: MSAEncoder | None = None + if msa_cfg.enabled: + self.msa_encoder = MSAEncoder( + d_msa=msa_cfg.d_msa, + d_pair=d_pair, + d_inputs=d_inputs, + d_hidden=msa_cfg.d_hidden, + n_layers=msa_cfg.n_layers, + n_heads_msa=msa_cfg.n_heads_msa, + msa_head_width=msa_cfg.msa_head_width, + ) + + self.post_init() + self._register_state_dict_hook(_drop_transient_esmc_state) + + @property + def device(self) -> torch.device: + return next(self.parameters()).device + + def set_kernel_backend(self, backend: str | None) -> None: + validate_kernel_backend(backend) + self.folding_trunk.set_kernel_backend(backend) + if self.confidence_head is not None: + self.confidence_head.set_kernel_backend(backend) + self.structure_head.set_kernel_backend(backend) + self._kernel_backend = backend + + def set_chunk_size(self, chunk_size: int | None) -> None: + self.folding_trunk.set_chunk_size(chunk_size) + if self.confidence_head is not None: + self.confidence_head.set_chunk_size(chunk_size) + if self.msa_encoder is not None: + self.msa_encoder.set_chunk_size(chunk_size) + + def configure_lm_dropout( + self, + lm_dropout: float, + *, + force_lm_dropout_during_inference: bool = True, + ) -> None: + self.config.lm_dropout = lm_dropout + self.config.force_lm_dropout_during_inference = force_lm_dropout_during_inference + + @property + def esmc_precision_status(self) -> ESMCPrecisionStatus: + return self._esmc_precision_status + + def load_esmc( + self, + esmc_model_path: str, + precision: ESMCPrecision = "auto", + device: str | torch.device | None = None, + local_files_only: bool = False, + ) -> None: + """Load ESMC with the same precision policy as released checkpoints.""" + + _install_esmc_backbone( + self, + esmc_model_path, + precision=precision, + device=device, + local_files_only=local_files_only, + ) + + def reload_esmc( + self, + precision: ESMCPrecision = "auto", + device: str | torch.device | None = None, + local_files_only: bool | None = None, + ) -> None: + """Reload canonical ESMC weights and discard runtime quantization.""" + + source = self._esmc_source or self.config.esmc_id + old_esmc = self._esmc + self._esmc = None + self._esmc_fp8 = False + self._esmc_fp8_module_paths = () + self._ttt_lm_head = None + del old_esmc + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + self.load_esmc( + source, + precision=precision, + device=device, + local_files_only=( + self._esmc_local_files_only + if local_files_only is None + else local_files_only + ), + ) + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path, + *model_args, + load_esmc: bool = True, + **kwargs, + ): + if "config" not in kwargs: + kwargs["config"] = ESMFold2Config.from_pretrained( + pretrained_model_name_or_path, **kwargs + ) + esmc_precision = kwargs.pop("esmc_precision", None) + local_files_only = bool(kwargs.get("local_files_only", False)) + output_loading_info = bool(kwargs.get("output_loading_info", False)) + loaded = super().from_pretrained(pretrained_model_name_or_path, *model_args, **kwargs) + if output_loading_info: + model, loading_info = loaded + else: + model = loaded + if load_esmc: + model.load_esmc( + model.config.esmc_id, + precision=esmc_precision or model.config.esmc_precision, + local_files_only=local_files_only, + ) + return (model, loading_info) if output_loading_info else model + + def apply_torch_compile(self, mode: str = "fixed_seqlen", dynamic: bool | None = None) -> None: + if dynamic is None: + dynamic = mode == "dynamic_seqlen" + compile_kwargs: dict[str, bool] = {"dynamic": dynamic} + compile_targets = ( + PairUpdateBlock, + DiffusionTransformer, + DiffusionModule, + MSAEncoderBlock, + ) + + def _maybe_compile(module: nn.Module) -> None: + if isinstance(module, compile_targets): + module.forward = torch.compile(module.forward, **compile_kwargs) + + self.apply(_maybe_compile) + + def _compute_lm_hidden_states( + self, + input_ids: Tensor, + asym_id: Tensor, + residue_index: Tensor, + mol_type: Tensor, + tok_mask: Tensor, + ) -> Tensor: + if self._esmc_fp8 and torch.is_grad_enabled(): + _reload_esmc_bf16_for_gradients( + self, + reason=( + "Gradient-enabled ESMC execution requires BF16; the persisted " + "serving policy is unchanged." + ), + ) + if self._esmc is None: + raise RuntimeError("ESMFold2 language-model features require load_esmc=True.") + pad_to = 16 if self._esmc_fp8 else None + with _lm_precision_context(self._esmc_precision_status.resolved, self.device): + return compute_lm_hidden_states( + self._esmc, + input_ids, + asym_id, + residue_index, + mol_type, + tok_mask, + pad_to_multiple=pad_to, + ) + + def forward( + self, + token_index: Tensor, + residue_index: Tensor, + asym_id: Tensor, + sym_id: Tensor, + entity_id: Tensor, + mol_type: Tensor, + res_type: Tensor, + token_bonds: Tensor, + token_attention_mask: Tensor, + ref_pos: Tensor, + ref_element: Tensor, + ref_charge: Tensor, + ref_atom_name_chars: Tensor, + ref_space_uid: Tensor, + atom_attention_mask: Tensor, + atom_to_token: Tensor, + distogram_atom_idx: Tensor, + deletion_mean: Tensor | None = None, + msa: Tensor | None = None, + has_deletion: Tensor | None = None, + deletion_value: Tensor | None = None, + msa_attention_mask: Tensor | None = None, + input_ids: Tensor | None = None, + lm_hidden_states: Tensor | None = None, + res_type_soft: Tensor | None = None, + num_loops: int | None = None, + num_diffusion_samples: int | None = None, + num_sampling_steps: int | None = None, + early_exit: bool = False, + seed: int | None = None, + calculate_confidence: bool = True, + provide_soft_sequence_to_msa_and_profile: bool = True, + noise_scale: float | None = None, + step_scale: float | None = None, + max_inference_sigma: float | None = None, + output_attentions: bool | None = None, + output_hidden_states: bool | None = None, + return_dict: bool | None = None, + pocket_feature: Tensor | None = None, + gt_coords: Tensor | None = None, + is_resolved: Tensor | None = None, + frames_idx: Tensor | None = None, + disto_cond: Tensor | None = None, + disto_cond_mask: Tensor | None = None, + ) -> ESMFold2Output | tuple[Any, ...]: + output_hidden_states, return_dict = _resolve_structure_output_controls( + self.config, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + validate_msa_conditioning_inputs( + self.config, + msa=msa, + msa_attention_mask=msa_attention_mask, + has_deletion=has_deletion, + deletion_value=deletion_value, + deletion_mean=deletion_mean, + ) + validate_prepared_auxiliary_inputs( + pocket_feature=pocket_feature, + disto_cond=disto_cond, + disto_cond_mask=disto_cond_mask, + ) + del gt_coords, is_resolved, frames_idx + tok_mask = token_attention_mask + atm_mask = atom_attention_mask + n_loops = num_loops if num_loops is not None else self.config.num_loops + n_samples = ( + num_diffusion_samples + if num_diffusion_samples is not None + else self.config.num_diffusion_samples + ) + + if res_type.dim() == 2: + res_type_oh = F.one_hot(res_type.long(), num_classes=NUM_RES_TYPES).float() + res_type_oh = res_type_oh * tok_mask.unsqueeze(-1).float() + else: + res_type_oh = res_type.float() + + if msa is not None: + msa_oh_profile = F.one_hot(msa.long(), num_classes=NUM_RES_TYPES).float() + if msa_attention_mask is not None: + mask_f = msa_attention_mask.float().unsqueeze(-1) + msa_oh_profile = msa_oh_profile * mask_f + valid_seq_count = msa_attention_mask.float().sum(dim=1).clamp(min=1) + profile = msa_oh_profile.sum(dim=1) / valid_seq_count.unsqueeze(-1) + else: + profile = msa_oh_profile.mean(dim=1) + else: + profile = res_type_oh + + if res_type_soft is not None: + res_type_oh = res_type_soft.float() + if not self.config.disable_msa_features and provide_soft_sequence_to_msa_and_profile: + profile = res_type_oh + msa = res_type_oh.unsqueeze(1) + msa_attention_mask = tok_mask.unsqueeze(1) + + if deletion_mean is None: + deletion_mean = torch.zeros( + res_type.shape[0], res_type.shape[1], device=res_type.device + ) + if self.config.disable_msa_features: + profile = torch.zeros_like(profile) + deletion_mean = torch.zeros_like(deletion_mean) + + ref_element_oh = F.one_hot(ref_element.long(), num_classes=MAX_ATOMIC_NUMBER).float() + ref_atom_name_chars_oh = F.one_hot( + ref_atom_name_chars.long(), num_classes=CHAR_VOCAB_SIZE + ).float() + atm_mask_f = atm_mask.float() + ref_element_oh = ref_element_oh * atm_mask_f.unsqueeze(-1) + ref_atom_name_chars_oh = ref_atom_name_chars_oh * atm_mask_f.unsqueeze(-1).unsqueeze(-1) + atom_to_token = atom_to_token * atm_mask.long() + + use_amp = ref_pos.device.type == "cuda" + with torch.amp.autocast("cuda", enabled=use_amp, dtype=torch.bfloat16): + x_inputs = self.inputs_embedder( + aatype=res_type_oh, + profile=profile.float(), + deletion_mean=deletion_mean.float(), + ref_pos=ref_pos, + atom_attention_mask=atm_mask, + ref_space_uid=ref_space_uid, + ref_charge=ref_charge, + ref_element=ref_element_oh, + ref_atom_name_chars=ref_atom_name_chars_oh, + atom_to_token=atom_to_token, + ) + + z_init = self.z_init_1(x_inputs).unsqueeze(2) + self.z_init_2(x_inputs).unsqueeze(1) + relative_position_encoding = self.rel_pos( + residue_index=residue_index, + asym_id=asym_id, + sym_id=sym_id, + entity_id=entity_id, + token_index=token_index, + ) + token_bonds_encoding = self.token_bonds(token_bonds.float()) + z_init = z_init + relative_position_encoding + token_bonds_encoding + + if lm_hidden_states is None and input_ids is not None and self._esmc is not None: + lm_hidden_states = self._compute_lm_hidden_states( + input_ids, asym_id, residue_index, mol_type, tok_mask + ) + if lm_hidden_states is not None: + lm_dropout = ( + self.config.lm_dropout + if self.config.force_lm_dropout_during_inference or self.training + else 0.0 + ) + lm_z = self.language_model(lm_hidden_states.detach(), lm_dropout=lm_dropout) + z_init = z_init + lm_z.to(z_init.dtype) + + msa_kwargs: dict[str, Tensor] | None = None + if self.msa_encoder is not None and msa is not None: + if msa.dim() == 4: + batch_msa, depth, length_msa, _ = msa.shape + msa_oh = msa.permute(0, 2, 1, 3).float() + else: + batch_msa, depth, length_msa = msa.shape + msa_oh = F.one_hot( + msa.permute(0, 2, 1).long(), num_classes=NUM_RES_TYPES + ).float() + msa_attn = ( + msa_attention_mask.permute(0, 2, 1).float() + if msa_attention_mask is not None + else tok_mask[:, :, None].expand(-1, -1, depth).float() + ) + msa_oh = msa_oh * msa_attn.unsqueeze(-1) + hd = ( + has_deletion.permute(0, 2, 1).float() + if has_deletion is not None + else torch.zeros(batch_msa, length_msa, depth, device=msa.device) + ) + dv = ( + deletion_value.permute(0, 2, 1).float() + if deletion_value is not None + else torch.zeros(batch_msa, length_msa, depth, device=msa.device) + ) + msa_kwargs = { + "x_inputs": x_inputs, + "msa_oh": msa_oh, + "has_deletion": hd, + "deletion_value": dv, + "msa_attention_mask": msa_attn, + } + + pair_mask = tok_mask[:, :, None].float() * tok_mask[:, None, :].float() + z = torch.zeros_like(z_init) + prev_pair: Tensor | None = None + prev_disto_probs: Tensor | None = None + for loop_num in range(n_loops + 1): + z = z_init + self.pair_loop_proj(z) + if msa_kwargs is not None and self.msa_encoder is not None: + z = z + self.msa_encoder(x_pair=z, **msa_kwargs).to(z.dtype) + z = self.folding_trunk(z, pair_attention_mask=pair_mask) + + if early_exit and loop_num < n_loops: + l2_converged = False + if prev_pair is not None and loop_num > 0: + rel_l2 = ( + z.float() - prev_pair.float() + ).norm() / prev_pair.float().norm().clamp(min=1e-8) + l2_converged = rel_l2.item() < 0.25 + prev_pair = z.detach().clone() + sym_z = z.float() + z.float().transpose(-2, -3) + cur_probs = F.softmax(self.distogram_head(sym_z).float(), dim=-1) + if prev_disto_probs is not None and loop_num > 0: + kl_per_pair = ( + cur_probs + * (cur_probs.clamp(min=1e-8) / prev_disto_probs.clamp(min=1e-8)).log() + ).sum(-1) + kl = (kl_per_pair + kl_per_pair.transpose(-1, -2)).mean() / 2 + if l2_converged or kl.item() < 0.05: + break + prev_disto_probs = cur_probs.detach() + + distogram_logits = self.distogram_head(z + z.transpose(-2, -3)) + + with torch.no_grad(), _seed_context(seed): + structure_output = self.structure_head.sample( + z_trunk=z.float(), + s_inputs=x_inputs, + s_trunk=None, + relative_position_encoding=relative_position_encoding, + ref_pos=ref_pos, + ref_charge=ref_charge, + ref_mask=atm_mask, + ref_element=ref_element_oh, + ref_atom_name_chars=ref_atom_name_chars_oh, + ref_space_uid=ref_space_uid, + tok_idx=atom_to_token, + asym_id=asym_id, + residue_index=residue_index, + entity_id=entity_id, + token_index=token_index, + sym_id=sym_id, + token_attention_mask=tok_mask, + num_diffusion_samples=n_samples, + num_sampling_steps=num_sampling_steps, + max_inference_sigma=max_inference_sigma, + noise_scale=noise_scale, + step_scale=step_scale, + return_atom_repr=False, + denoising_early_exit_rmsd=(0.10 if early_exit else None), + ) + sample_coords = structure_output["sample_atom_coords"] + if sample_coords is None: + raise RuntimeError("ESMFold2 structure sampling did not return coordinates.") + if sample_coords.ndim == 4: + batch, sample_count, atom_count, coord_dim = sample_coords.shape + sample_coords_for_gather = sample_coords.reshape( + batch * sample_count, + atom_count, + coord_dim, + ) + rep_idx = distogram_atom_idx.repeat_interleave(sample_count, 0).long() + else: + sample_coords_for_gather = sample_coords + rep_idx = distogram_atom_idx.long() + representative_atom_coords = gather_rep_atom_coords( + sample_coords_for_gather, + rep_idx, + ) + + output: dict[str, Tensor] = { + "distogram_logits": distogram_logits, + "sample_atom_coords": sample_coords, + "representative_atom_coords": representative_atom_coords, + } + if calculate_confidence and self.confidence_head is not None: + confidence_output = self.confidence_head( + s_inputs=x_inputs.detach(), + z=z.detach().float(), + x_pred=sample_coords.detach(), + distogram_atom_idx=distogram_atom_idx, + token_attention_mask=tok_mask, + atom_to_token=atom_to_token, + atom_attention_mask=atm_mask, + asym_id=asym_id, + mol_type=mol_type, + num_diffusion_samples=n_samples, + relative_position_encoding=relative_position_encoding.detach(), + token_bonds_encoding=token_bonds_encoding.detach(), + ) + output.update(confidence_output) + output["atom_pad_mask"] = atm_mask.unsqueeze(0) if atm_mask.dim() == 1 else atm_mask + output["residue_index"] = residue_index + output["entity_id"] = entity_id + return _finalize_structure_output( + output, + token_input_state=x_inputs, + pair_state=z, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + @property + def input_builder(self): + if self._esmfold2_input_builder is None: + from .esmfold2_processor import ESMFold2InputBuilder + + self._esmfold2_input_builder = ESMFold2InputBuilder() + return self._esmfold2_input_builder + + @property + def input_types(self): + from . import esmfold2_types + + return esmfold2_types + + def prepare_structure_input(self, input, seed: int | None = None): + return self.input_builder.prepare_model_input( + self, + input, + seed=seed, + device=self.device, + ) + + @torch.no_grad() + def infer_protein(self, seq: str, **forward_kwargs) -> ESMFold2Output: + from .protein_utils import prepare_protein_features + + if forward_kwargs.pop("return_dict", True) is not True: + raise ValueError( + "infer_protein always returns a mapping; return_dict=False is invalid." + ) + features = prepare_protein_features(seq) + if not self.config.msa_conditioning: + for name in MSA_CONDITIONING_INPUT_NAMES: + features.pop(name, None) + features = {name: tensor.to(self.device) for name, tensor in features.items()} + output = self(**features, **forward_kwargs, return_dict=True) + for name in ( + "res_type", + "atom_to_token", + "ref_atom_name_chars", + "atom_attention_mask", + "token_attention_mask", + "residue_index", + ): + output[name] = features[name] + return output + + def fold( + self, + input, + *, + num_loops: int = 3, + num_sampling_steps: int = 50, + num_diffusion_samples: int = 1, + seed: int | None = None, + noise_scale: float | None = None, + step_scale: float | None = None, + max_inference_sigma: int | None = None, + early_exit: bool = False, + complex_id: str = "pred", + ): + return self.input_builder.fold( + self, + input, + num_loops=num_loops, + num_sampling_steps=num_sampling_steps, + num_diffusion_samples=num_diffusion_samples, + seed=seed, + noise_scale=noise_scale, + step_scale=step_scale, + max_inference_sigma=max_inference_sigma, + early_exit=early_exit, + complex_id=complex_id, + ) + + def fold_protein( + self, + sequence: str, + *, + chain_id: str = "A", + num_loops: int = 3, + num_sampling_steps: int = 50, + num_diffusion_samples: int = 1, + seed: int | None = None, + complex_id: str = "pred", + ): + from .esmfold2_types import ProteinInput, StructurePredictionInput + + input = StructurePredictionInput(sequences=[ProteinInput(id=chain_id, sequence=sequence)]) + return self.fold( + input, + num_loops=num_loops, + num_sampling_steps=num_sampling_steps, + num_diffusion_samples=num_diffusion_samples, + seed=seed, + complex_id=complex_id, + ) + + @staticmethod + def result_to_cif(result) -> str: + if isinstance(result, list): + raise TypeError("Pass one MolecularComplexResult at a time.") + return result.complex.to_mmcif() + + @staticmethod + def result_to_pdb(result) -> str: + if isinstance(result, list): + raise TypeError("Pass one MolecularComplexResult at a time.") + return result.complex.to_protein_complex().to_pdb_string() + + def save_as_cif(self, result, output_path: str | Path) -> None: + Path(output_path).write_text(self.result_to_cif(result)) + + def save_as_pdb(self, result, output_path: str | Path) -> None: + Path(output_path).write_text(self.result_to_pdb(result)) + + def infer_protein_as_cif(self, seq: str, **forward_kwargs) -> str: + return self.result_to_cif(self.fold_protein(seq, **forward_kwargs)) + + def infer_protein_as_pdb(self, seq: str, **forward_kwargs) -> str: + return self.result_to_pdb(self.fold_protein(seq, **forward_kwargs)) + + +__all__ = [ + "ConfidenceHead", + "ESMFold2ExperimentalModel", + "ESMFold2Output", + "MSAEncoder", + "MSAEncoderBlock", +] diff --git a/fastplms/models/esmfold2/protein_reference_geometry.json b/fastplms/models/esmfold2/protein_reference_geometry.json new file mode 100644 index 0000000000000000000000000000000000000000..5cab64b9eaafe3fb4d4bcb37131ae2202587be5a --- /dev/null +++ b/fastplms/models/esmfold2/protein_reference_geometry.json @@ -0,0 +1 @@ +{"dtype":"float32","provenance":{"contract":"biohub_esmfold2_input_v1","manifest_family":"esmfold2"},"residues":{"ALA":{"C":[1.2127548456192017,0.4737588167190552,0.19521640241146088],"CA":[-0.04190138354897499,0.17447763681411743,-0.5729365348815918],"CB":[-1.276943325996399,0.4288230538368225,0.29937705397605896],"N":[-0.01003183238208294,-1.2073018550872803,-1.0555061101913452],"O":[1.9390329122543335,1.4484562873840332,-0.13759790360927582]},"ARG":{"C":[-3.469440460205078,-1.0612813234329224,-0.2755832374095917],"CA":[-2.0503084659576416,-0.5735036730766296,-0.4097220301628113],"CB":[-1.4193516969680786,-0.3735991418361664,0.9852858781814575],"CD":[0.6643245816230774,1.0068185329437256,0.3963329493999481],"CG":[0.11878877878189087,-0.3112654983997345,0.963895857334137],"CZ":[3.098905324935913,0.3215920031070709,-0.09047172218561172],"N":[-2.0170421600341797,0.6717798113822937,-1.1794233322143555],"NE":[2.1090238094329834,1.0977025032043457,0.6120952367782593],"NH1":[4.461230278015137,0.3844667971134186,0.34141138195991516],"NH2":[2.7856509685516357,-0.4166366159915924,-1.1148239374160767],"O":[-3.8218462467193604,-2.1369943618774414,-0.8294969797134399]},"ASN":{"C":[-1.9211044311523438,-0.6982439160346985,-0.42196929454803467],"CA":[-0.76087886095047,0.23876343667507172,-0.23573364317417145],"CB":[0.5504899024963379,-0.5078350305557251,-0.5390339493751526],"CG":[1.7250099182128906,0.4264017939567566,-0.5778228640556335],"N":[-0.7595629096031189,0.7503494620323181,1.1369825601577759],"ND2":[2.57365345954895,0.5730618834495544,0.5608599781990051],"O":[-2.677666187286377,-0.5753439664840698,-1.4223182201385498],"OD1":[1.9470350742340088,1.1086392402648926,-1.613560438156128]},"ASP":{"C":[-0.9431572556495667,1.0356197357177734,0.18555717170238495],"CA":[-0.6379959583282471,-0.41974392533302307,0.41681644320487976],"CB":[0.48594576120376587,-0.8970447778701782,-0.5209363698959351],"CG":[1.780342936515808,-0.19918935000896454,-0.2310730367898941],"N":[-1.8452696800231934,-1.2169504165649414,0.19437327980995178],"O":[-1.5183608531951904,1.4045922756195068,-0.8739855885505676],"OD1":[2.5202910900115967,-0.6044584512710571,0.7049641013145447],"OD2":[2.1454880237579346,0.9208861589431763,-0.9712985157966614]},"CYS":{"C":[-1.2652032375335693,-0.6832379698753357,-0.3594406247138977],"CA":[0.11344368755817413,-0.09400428831577301,-0.45952197909355164],"CB":[0.6919880509376526,0.09034398198127747,0.952482283115387],"N":[0.0469963513314724,1.190075159072876,-1.1607273817062378],"O":[-1.4631439447402954,-1.8851220607757568,-0.6826791763305664],"SG":[2.4619927406311035,0.5235707759857178,0.9020372629165649]},"GLN":{"C":[-1.7545503377914429,0.7091967463493347,0.8433493971824646],"CA":[-1.370002269744873,-0.6000258922576904,0.2103111445903778],"CB":[0.02040259726345539,-0.5004461407661438,-0.44764479994773865],"CD":[2.4745187759399414,-0.24800164997577667,-0.09364881366491318],"CG":[1.1377512216567993,-0.28680720925331116,0.582992434501648],"N":[-2.370004653930664,-0.9637529850006104,-0.7942749261856079],"NE2":[2.947425603866577,0.9601329565048218,-0.6888364553451538],"O":[-1.8520662784576416,0.7999289631843567,2.0964975357055664],"OE1":[3.1685523986816406,-1.2966246604919434,-0.1717153936624527]},"GLU":{"C":[-1.7741456031799316,0.9664392471313477,0.09259600937366486],"CA":[-1.0560977458953857,0.027459044009447098,1.0306966304779053],"CB":[0.4706551432609558,0.048803869634866714,0.8114414811134338],"CD":[2.398822069168091,-0.3097084164619446,-0.7210537791252136],"CG":[0.9133604764938354,-0.4219329059123993,-0.5830985307693481],"N":[-1.5850872993469238,-1.337684154510498,0.9490851163864136],"O":[-1.9012441635131836,2.181349992752075,0.402479350566864],"OE1":[3.1389315128326416,-1.274524450302124,-0.39029765129089355],"OE2":[2.9647817611694336,0.8781346082687378,-1.1732689142227173]},"GLY":{"C":[0.9440054893493652,-0.10314033925533295,0.19859643280506134],"CA":[-0.39974430203437805,0.5488945245742798,0.15242962539196014],"N":[-1.3942985534667969,-0.39875128865242004,-0.3370324671268463],"O":[1.3352899551391602,-0.669218122959137,1.2541258335113525]},"HIS":{"C":[-2.675257921218872,0.6571555733680725,-0.30441102385520935],"CA":[-1.3396095037460327,0.24797579646110535,0.24960045516490936],"CB":[-0.3041955828666687,0.21721023321151733,-0.8885309100151062],"CD2":[1.780855417251587,-1.1011489629745483,-0.3814258575439453],"CE1":[2.9566943645477295,0.4924798905849457,0.6477115750312805],"CG":[1.0887513160705566,0.028941065073013306,-0.36419469118118286],"N":[-1.4532867670059204,-1.0689626932144165,0.881072461605072],"ND1":[1.840459942817688,1.0411773920059204,0.29804590344429016],"NE2":[3.0280203819274902,-0.8751969337463379,0.26084381341934204],"O":[-3.1311378479003906,1.8079776763916016,-0.06785715371370316]},"ILE":{"C":[-1.3896740674972534,0.8142145276069641,-1.1164065599441528],"CA":[-1.0636085271835327,-0.35169270634651184,-0.21393552422523499],"CB":[0.061667006462812424,0.01599610224366188,0.8057394623756409],"CD1":[1.7929610013961792,0.899773120880127,-0.8863027691841125],"CG1":[1.502519965171814,-0.08899776637554169,0.24154816567897797],"CG2":[-0.053174979984760284,-0.8521055579185486,2.0702083110809326],"N":[-0.7167549729347229,-1.5426139831542969,-0.9983330368995667],"O":[-1.2377792596817017,0.7302915453910828,-2.3656840324401855]},"LEU":{"C":[1.9905058145523071,0.24182087182998657,0.7879968285560608],"CA":[1.3077669143676758,-0.6677430868148804,-0.19492436945438385],"CB":[-0.20306941866874695,-0.8093230128288269,0.11243502795696259],"CD1":[-2.4228057861328125,0.29949337244033813,0.573042094707489],"CD2":[-1.0282856225967407,1.1250264644622803,-1.346014380455017],"CG":[-0.9916267395019531,0.5234957337379456,0.06723011285066605],"N":[1.9657520055770874,-1.9763224124908447,-0.18391533195972443],"O":[2.06896710395813,-0.07880014181137085,2.0048046112060547]},"LYS":{"C":[2.7168593406677246,1.595757246017456,-0.20924785733222961],"CA":[2.0314927101135254,0.2786507308483124,-0.4298512041568756],"CB":[0.5018402934074402,0.4873858690261841,-0.49062973260879517],"CD":[-1.769762635231018,-0.5552700161933899,-1.040329933166504],"CE":[-2.576533555984497,-1.0221366882324219,0.18493641912937164],"CG":[-0.25062066316604614,-0.7894009947776794,-0.9055535793304443],"N":[2.4221372604370117,-0.6473312377929688,0.6370573043823242],"NZ":[-2.269151210784912,-0.24293844401836395,1.3849012851715088],"O":[3.397681713104248,2.116427421569824,-1.1332510709762573]},"MET":{"C":[2.30391001701355,0.8367712497711182,-0.7254616618156433],"CA":[1.2630571126937866,-0.24417810142040253,-0.7626462578773499],"CB":[0.10567972809076309,0.10861825942993164,0.19741646945476532],"CE":[-3.265165090560913,0.7033554911613464,-0.11588376015424728],"CG":[-1.0658042430877686,-0.8736631274223328,0.08811883628368378],"N":[1.8903918266296387,-1.5252995491027832,-0.42638593912124634],"O":[2.465414524078369,1.5928632020950317,-1.7207728624343872],"SD":[-2.4557132720947266,-0.3332225978374481,1.1461700201034546]},"PHE":{"C":[-1.8900631666183472,0.45833414793014526,1.0232222080230713],"CA":[-1.591969609260559,-0.8545162677764893,0.35214468836784363],"CB":[-0.760358452796936,-0.6342853307723999,-0.9257160425186157],"CD1":[0.8468314409255981,1.2480632066726685,-0.7146694660186768],"CD2":[1.6827683448791504,-0.9758077263832092,-0.1423054188489914],"CE1":[2.1801748275756836,1.7875733375549316,-0.3744623064994812],"CE2":[2.888307809829712,-0.48277512192726135,0.16804970800876617],"CG":[0.604112982749939,-0.07200468331575394,-0.6148118376731873],"CZ":[3.149812936782837,0.9656873941421509,0.04440271109342575],"N":[-2.8484435081481934,-1.525790810585022,0.01789816841483116],"O":[-1.3424992561340332,0.74432373046875,2.121629476547241]},"PRO":{"C":[1.6121541261672974,-1.1711241006851196,0.31082412600517273],"CA":[0.32722190022468567,-0.6164458394050598,-0.25072571635246277],"CB":[0.3248198926448822,0.9028244018554688,-0.33368146419525146],"CD":[-1.8495968580245972,0.026575811207294464,0.2681289613246918],"CG":[-1.1425083875656128,1.2730128765106201,-0.2590600252151489],"N":[-0.836250364780426,-0.9899801015853882,0.5561304688453674],"O":[1.6127740144729614,-2.2771971225738525,0.9156193733215332]},"SER":{"C":[0.9941009879112244,-0.5374617576599121,0.73505038022995],"CA":[0.00013792862591799349,0.4966467022895813,0.28510504961013794],"CB":[-1.1279288530349731,-0.1659376323223114,-0.5160963535308838],"N":[0.674650251865387,1.5018702745437622,-0.5367295145988464],"O":[1.0545241832733154,-0.8683545589447021,1.9495396614074707],"OG":[-1.8135979175567627,-1.085249662399292,0.28947514295578003]},"THR":{"C":[-1.294381856918335,0.7077372074127197,-0.5549946427345276],"CA":[-0.5433306097984314,-0.16364754736423492,0.41697052121162415],"CB":[0.853203296661377,-0.5363803505897522,-0.14109353721141815],"CG2":[1.7225933074951172,0.7054727077484131,-0.3651331067085266],"N":[-1.325830340385437,-1.3728225231170654,0.6882233023643494],"O":[-1.6939635276794434,0.23654410243034363,-1.6540418863296509],"OG1":[1.5220820903778076,-1.379003643989563,0.7635167837142944]},"TRP":{"C":[2.1113572120666504,-0.6121063232421875,-0.7733646035194397],"CA":[2.384092092514038,0.09079249948263168,0.5325262546539307],"CB":[1.281521201133728,1.1139036417007446,0.8559791445732117],"CD1":[-0.42329534888267517,-0.15470874309539795,2.2227554321289062],"CD2":[-1.1023900508880615,0.2158389836549759,0.11529432237148285],"CE2":[-2.045644998550415,-0.4881173074245453,0.710669219493866],"CE3":[-1.2173502445220947,0.6102271676063538,-1.300106406211853],"CG":[-0.04292375594377518,0.44645074009895325,1.0942792892456055],"CH2":[-3.3817875385284424,-0.5677337646484375,-1.3032053709030151],"CZ2":[-3.256009340286255,-0.9164394736289978,-0.00984987337142229],"CZ3":[-2.315925121307373,0.2306906282901764,-1.9776310920715332],"N":[3.686030864715576,0.7599999904632568,0.496155709028244],"NE1":[-1.7030320167541504,-0.7665823101997375,2.0595016479492188],"O":[1.796526312828064,-1.8323148488998413,-0.7775964140892029]},"TYR":{"C":[-3.347280740737915,0.3588399887084961,-0.09830684959888458],"CA":[-1.913882851600647,0.23552845418453217,0.330669641494751],"CB":[-1.0093992948532104,0.0004731413209810853,-0.8981552124023438],"CD1":[1.0992432832717896,1.1877919435501099,-0.3579142987728119],"CD2":[1.1803174018859863,-1.253401279449463,-0.31122180819511414],"CE1":[2.5253450870513916,1.1990256309509277,0.029804613441228867],"CE2":[2.471151113510132,-1.240687608718872,0.043534230440855026],"CG":[0.4520410895347595,0.021162061020731926,-0.5305932760238647],"CZ":[3.180687665939331,0.04672492295503616,0.2214856892824173],"N":[-1.7900604009628296,-0.8409399390220642,1.3180142641067505],"O":[-3.967811346054077,-0.6449354290962219,-0.5423302054405212],"OH":[4.523719787597656,0.0671030730009079,0.5877485871315002]},"UNK":{"C":[0.0,0.0,0.0],"CA":[0.0,0.0,0.0],"N":[0.0,0.0,0.0],"O":[0.0,0.0,0.0]},"VAL":{"C":[1.8391697406768799,0.4067850410938263,0.06351757049560547],"CA":[0.6014357209205627,-0.10503966361284256,-0.6336286664009094],"CB":[-0.694736897945404,0.4259096384048462,0.03581475466489792],"CG1":[-1.9276031255722046,0.09515828639268875,-0.8172357082366943],"CG2":[-0.8938426971435547,-0.08640842139720917,1.472349762916565],"N":[0.5987519025802612,-1.569443702697754,-0.7379124760627747],"O":[2.3952062129974365,-0.2666190266609192,0.9731166958808899]}},"schema":"fastplms.esmfold2.reference_geometry.v1"} diff --git a/fastplms/models/esmfold2/protein_utils.py b/fastplms/models/esmfold2/protein_utils.py new file mode 100644 index 0000000000000000000000000000000000000000..1d7d0eaa25d0f4285056c0a8a7eb7ea7f2bd54ec --- /dev/null +++ b/fastplms/models/esmfold2/protein_utils.py @@ -0,0 +1,178 @@ +"""Protein-only ESMFold2 featurization without the Biohub runtime package. + +Input is one amino-acid sequence. The transformation expands each residue into +the checkpoint atom schema, pads atoms to a multiple of 32, and emits batched +token, atom, and single-sequence MSA tensors. Reference coordinates are loaded +lazily from a provenance-bearing declarative package asset. +""" + +from __future__ import annotations + +import json +from functools import cache +from importlib.resources import files +from typing import Any + +import torch +from torch import Tensor + +from .esmfold2_constants import ( + CHARGED_ATOMS, + ELEMENT_TO_ATOMIC_NUM, + ESM_PROTEIN_VOCAB, + MOL_TYPE_PROTEIN, + PROTEIN_1TO3, + PROTEIN_HEAVY_ATOMS, + PROTEIN_RESIDUE_TO_RES_TYPE, + PROTEIN_UNK_RES_TYPE, +) + +_GEOMETRY_ASSET = "protein_reference_geometry.json" +_GEOMETRY_SCHEMA = "fastplms.esmfold2.reference_geometry.v1" + + +@cache +def _reference_geometry() -> dict[str, dict[str, tuple[float, float, float]]]: + resource = files(__package__).joinpath(_GEOMETRY_ASSET) + with resource.open(mode="r", encoding="utf-8") as handle: + payload = json.load(handle) + if ( + payload.get("schema") != _GEOMETRY_SCHEMA + or payload.get("dtype") != "float32" + or payload.get("provenance", {}).get("manifest_family") != "esmfold2" + ): + raise RuntimeError("The ESMFold2 reference-geometry asset has invalid provenance.") + + raw_residues = payload.get("residues") + if not isinstance(raw_residues, dict): + raise RuntimeError("The ESMFold2 reference-geometry asset has no residue table.") + geometry: dict[str, dict[str, tuple[float, float, float]]] = {} + for residue, atom_positions in raw_residues.items(): + if not isinstance(residue, str) or not isinstance(atom_positions, dict): + raise RuntimeError("The ESMFold2 reference-geometry residue table is malformed.") + geometry[residue] = {} + for atom_name, position in atom_positions.items(): + if ( + not isinstance(atom_name, str) + or not isinstance(position, list) + or len(position) != 3 + ): + raise RuntimeError("The ESMFold2 reference-geometry atom table is malformed.") + geometry[residue][atom_name] = tuple(float(value) for value in position) + + expected_residues = set(PROTEIN_HEAVY_ATOMS) - {"MSE"} + if set(geometry) != expected_residues: + raise RuntimeError("The ESMFold2 reference-geometry residue set is incomplete.") + for residue, atom_names in PROTEIN_HEAVY_ATOMS.items(): + if residue == "MSE": + continue + if set(geometry[residue]) != set(atom_names): + raise RuntimeError(f"Reference geometry differs from the atom schema for {residue}.") + return geometry + + +def _encode_atom_name(atom_name: str) -> tuple[int, int, int, int]: + padded = atom_name.ljust(4)[:4] + return tuple(ord(character) - 32 if character != " " else 0 for character in padded) + + +def _padded_atom_count(actual_count: int) -> int: + return max(32, ((actual_count + 31) // 32) * 32) + + +def _residue_records(sequence: str) -> tuple[list[dict[str, Any]], list[int], list[int], list[int]]: + geometry = _reference_geometry() + atoms: list[dict[str, Any]] = [] + residue_types: list[int] = [] + input_ids: list[int] = [] + representative_atoms: list[int] = [] + + for token_index, residue_letter in enumerate(sequence): + residue_name = PROTEIN_1TO3.get(residue_letter, "UNK") + atom_names = PROTEIN_HEAVY_ATOMS[residue_name] + atom_start = len(atoms) + for atom_name in atom_names: + atoms.append( + { + "token_index": token_index, + "name": atom_name, + "element": atom_name[0], + "charge": CHARGED_ATOMS.get((residue_name, atom_name), 0), + "position": geometry[residue_name][atom_name], + } + ) + + representative_name = "CB" if "CB" in atom_names else "CA" + representative_atoms.append(atom_start + atom_names.index(representative_name)) + residue_types.append(PROTEIN_RESIDUE_TO_RES_TYPE.get(residue_name, PROTEIN_UNK_RES_TYPE)) + input_ids.append(ESM_PROTEIN_VOCAB.get(residue_letter, ESM_PROTEIN_VOCAB["X"])) + + return atoms, residue_types, input_ids, representative_atoms + + +def prepare_protein_features(sequence: str) -> dict[str, Tensor]: + """Build the protein-only feature mapping consumed by ESMFold2. + + Every tensor includes a leading batch dimension. Biological tokens have + length ``l``; atom tensors have length ``n_atoms``, where ``n_atoms`` is the + smallest multiple of 32 covering all heavy atoms. + """ + + if not sequence: + raise ValueError("sequence must be non-empty") + + atoms, residue_types, input_ids, representative_atoms = _residue_records(sequence) + sequence_length = len(sequence) + n_atoms = _padded_atom_count(len(atoms)) + + ref_pos = torch.zeros((n_atoms, 3), dtype=torch.float32) + ref_element = torch.zeros(n_atoms, dtype=torch.int64) + ref_charge = torch.zeros(n_atoms, dtype=torch.int8) + ref_atom_name_chars = torch.zeros((n_atoms, 4), dtype=torch.int64) + ref_space_uid = torch.zeros(n_atoms, dtype=torch.int64) + atom_attention_mask = torch.zeros(n_atoms, dtype=torch.bool) + atom_to_token = torch.zeros(n_atoms, dtype=torch.int64) + + for atom_index, atom in enumerate(atoms): + token_index = atom["token_index"] + ref_pos[atom_index] = torch.tensor(atom["position"], dtype=torch.float32) + ref_element[atom_index] = ELEMENT_TO_ATOMIC_NUM[atom["element"]] + ref_charge[atom_index] = atom["charge"] + ref_atom_name_chars[atom_index] = torch.tensor( + _encode_atom_name(atom["name"]), dtype=torch.int64 + ) + ref_space_uid[atom_index] = token_index + atom_attention_mask[atom_index] = True + atom_to_token[atom_index] = token_index + + residue_type_tensor = torch.tensor(residue_types, dtype=torch.int64) + msa = residue_type_tensor.unsqueeze(0) + features = { + "token_index": torch.arange(sequence_length, dtype=torch.int64), + "residue_index": torch.arange(sequence_length, dtype=torch.int64), + "asym_id": torch.zeros(sequence_length, dtype=torch.int64), + "sym_id": torch.zeros(sequence_length, dtype=torch.int64), + "entity_id": torch.ones(sequence_length, dtype=torch.int64), + "mol_type": torch.full((sequence_length,), MOL_TYPE_PROTEIN, dtype=torch.int64), + "res_type": residue_type_tensor, + "input_ids": torch.tensor(input_ids, dtype=torch.int64), + "token_bonds": torch.zeros((sequence_length, sequence_length, 1), dtype=torch.float32), + "token_attention_mask": torch.ones(sequence_length, dtype=torch.bool), + "ref_pos": ref_pos, + "ref_element": ref_element, + "ref_charge": ref_charge, + "ref_atom_name_chars": ref_atom_name_chars, + "ref_space_uid": ref_space_uid, + "atom_attention_mask": atom_attention_mask, + "atom_to_token": atom_to_token, + "distogram_atom_idx": torch.tensor(representative_atoms, dtype=torch.int64), + "msa": msa, + "msa_attention_mask": torch.ones_like(msa, dtype=torch.bool), + "has_deletion": torch.zeros_like(msa, dtype=torch.bool), + "deletion_value": torch.zeros_like(msa, dtype=torch.float32), + "deletion_mean": torch.zeros(sequence_length, dtype=torch.float32), + } + return {name: tensor.unsqueeze(0) for name, tensor in features.items()} + + +__all__ = ["prepare_protein_features"] diff --git a/fastplms/models/esmfold2/reproducibility.py b/fastplms/models/esmfold2/reproducibility.py new file mode 100644 index 0000000000000000000000000000000000000000..92194cec1867567899ff7fd26646961ac44bd607 --- /dev/null +++ b/fastplms/models/esmfold2/reproducibility.py @@ -0,0 +1,64 @@ +"""Dependency-light RNG scoping for ESMFold2 workflows.""" + +from __future__ import annotations + +import random +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from typing import Any + +import numpy as np +import torch +from torch import Tensor + + +@dataclass(frozen=True) +class _RandomState: + python: object + numpy: tuple[Any, ...] + torch_cpu: Tensor + torch_cuda: list[Tensor] | None + + +def _capture_random_state() -> _RandomState: + cuda_state = torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None + return _RandomState( + python=random.getstate(), + numpy=np.random.get_state(), + torch_cpu=torch.random.get_rng_state(), + torch_cuda=cuda_state, + ) + + +def _restore_random_state(state: _RandomState) -> None: + random.setstate(state.python) + np.random.set_state(state.numpy) + torch.random.set_rng_state(state.torch_cpu) + if state.torch_cuda is not None: + torch.cuda.set_rng_state_all(state.torch_cuda) + + +@contextmanager +def seed_context(seed: int | None) -> Iterator[None]: + """Seed Python, NumPy, and Torch temporarily, then restore every stream.""" + + if seed is None: + yield + return + if isinstance(seed, bool) or not isinstance(seed, int): + raise TypeError("seed must be None or an integer (excluding bool).") + state = _capture_random_state() + normalized_seed = seed % (2**32) + random.seed(normalized_seed) + np.random.seed(normalized_seed) + torch.manual_seed(normalized_seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(normalized_seed) + try: + yield + finally: + _restore_random_state(state) + + +__all__ = ["seed_context"] diff --git a/fastplms/models/ttt.py b/fastplms/models/ttt.py new file mode 100644 index 0000000000000000000000000000000000000000..56b4dbae55b65f1d6a9d86111a22815c103d0162 --- /dev/null +++ b/fastplms/models/ttt.py @@ -0,0 +1,882 @@ +from __future__ import annotations + +import contextlib +import math +import numbers +import torch +import torch.nn as nn +import torch.nn.functional as F +from collections.abc import Iterator, Mapping +from dataclasses import asdict, dataclass, fields +from typing import Any + + +_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: 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: Mapping[str, 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, 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: + # x: (..., d_in) + base = self.linear(x) # (..., d_out) + delta = ( # (..., d_out) + self.lora_up(self.lora_down(x.to(dtype=torch.float32))) * self.scale + ) + return base + delta.to(dtype=base.dtype) # (..., d_out) + + 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 | Mapping[str, 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, 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, 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: Any, + ) -> torch.Tensor | dict[str, torch.Tensor]: + del kwargs + if input_ids is not None: + return input_ids # (b, l) + 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"] # (b, l) + + 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: + # input_ids: (b, l) + 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: Mapping[str, Any] = {} + get_vocab = getattr(tokenizer, "get_vocab", None) + if callable(get_vocab): + vocab = get_vocab() + elif isinstance(getattr(tokenizer, "vocab", None), Mapping): + vocab = tokenizer.vocab + elif isinstance(getattr(tokenizer, "_token_to_id", None), 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) # (c_aa,) + + def _ttt_predict_logits( + self, + batch: torch.Tensor | dict[str, torch.Tensor], + **kwargs: Any, + ) -> torch.Tensor: + del kwargs + if isinstance(batch, dict): + output = self(**batch) + return output.logits # (b, l, c) + attention_mask = batch.ne(self._ttt_padding_token()) # (b, l) + output = self(input_ids=batch, attention_mask=attention_mask) + return output.logits # (b, l, c) + + def _ttt_eval_step( + self, + step: int, + loss: float, + seq: str | list[str] | None = None, + input_ids: torch.Tensor | None = None, + **kwargs: Any, + ) -> tuple[dict[str, 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(), # (r, d_in) + "lora_up.weight": module.lora_up.weight.detach().clone(), # (d_out, r) + } + ) + 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, Any]: + return { + "version": _TTT_SERIALIZATION_VERSION, + "initialized": bool(self._ttt_initialized), + "config": self.ttt_config.to_dict(), + } + + def save_pretrained(self, save_directory: Any, *args: Any, **kwargs: Any) -> 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()} # unchanged shapes + return batch.to(device) # unchanged shape + + def _ttt_input_ids_from_batch( + self, + batch: torch.Tensor | dict[str, torch.Tensor], + ) -> torch.Tensor: + if isinstance(batch, dict): + return batch["input_ids"] # (b, l) + return batch # (b, l) + + 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 # (b, l) + return updated + return input_ids # (b, l) + + def _ttt_non_special_mask(self, input_ids: torch.Tensor) -> torch.Tensor: + # input_ids: (b, l) + residue_ids = self._ttt_replacement_tokens(input_ids) # (c_aa,) + return torch.isin(input_ids, residue_ids) # (b, l) + + 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( # (4,) + [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( # (b, l) + 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"] # (b, l) + attention_mask = batch.get( # (b, l) + "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 = ( # (l,) + self._ttt_non_special_mask(input_ids).any(dim=0).to(torch.int64) + ) + prefix = F.pad(position_has_residue.cumsum(dim=0), (1, 0)) # (l + 1,) + window_counts = prefix[cfg.crop_size :] - prefix[: -cfg.crop_size] # (l-crop+1,) + valid_starts = torch.where(window_counts > 0)[0] # (n_valid,) + if valid_starts.numel() == 0: + raise ValueError("TTT could not find a crop containing a biological residue token.") + selected = torch.randint( # (1,) + 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] # (b, crop_size, ...) + else: + cropped[name] = tensor + return cropped + return input_ids[:, start:end] # (b, crop_size) + + 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) # (b, l) + row_has_residue = self._ttt_non_special_mask(input_ids).any(dim=1) # (b,) + eligible_rows = torch.where(row_has_residue)[0] # (n_eligible,) + if eligible_rows.numel() == 0: + raise ValueError( + "TTT sampled batch contains no trainable biological residue tokens." + ) + sampled_row_indices = torch.randint( # (b_sample,) + eligible_rows.numel(), + (cfg.batch_size,), + generator=generator, + device=input_ids.device, + ) + rows = eligible_rows[sampled_row_indices] # (b_sample,) + 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) # (b_sample, ...) + else: + sampled[name] = tensor + else: + sampled = input_ids.index_select(0, rows) # (b_sample, l) + + sampled_ids = self._ttt_input_ids_from_batch(sampled) # (b_sample, l) + labels = sampled_ids.clone() # (b_sample, l) + non_special = self._ttt_non_special_mask(sampled_ids) # (b_sample, l) + label_mask = torch.zeros_like(non_special) # (b_sample, l) + for row_idx in range(sampled_ids.shape[0]): + candidate_positions = torch.where(non_special[row_idx])[0] # (n_candidates,) + if candidate_positions.numel() == 0: + continue + num_mask = max(1, round(candidate_positions.numel() * cfg.mask_ratio)) + order = torch.randperm( # (n_candidates,) + candidate_positions.numel(), + generator=generator, + device=sampled_ids.device, + ) + chosen = candidate_positions[order[:num_mask]] # (n_mask,) + label_mask[row_idx, chosen] = True + labels = labels.masked_fill(~label_mask, -100) # (b_sample, l) + + masked_ids = sampled_ids.clone() # (b_sample, l) + chosen_positions = torch.where(label_mask) # two (n_chosen,) tensors + if chosen_positions[0].numel() > 0: + random_values = torch.rand( # (n_chosen,) + chosen_positions[0].shape, + generator=generator, + device=sampled_ids.device, + ) + leave = random_values < cfg.bert_leave_prob # (n_chosen,) + replace = (random_values >= cfg.bert_leave_prob) & ( # (n_chosen,) + random_values < cfg.bert_leave_prob + cfg.bert_replace_prob + ) + mask = ~(leave | replace) # (n_chosen,) + 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) # (c_aa,) + replacement_idx = torch.randint( # (n_replace,) + 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 # batch, (b_sample, l) + + @contextlib.contextmanager + def _ttt_seed_scope(self, seed: int | None) -> 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 | Mapping[str, Any] | None = None, + **kwargs: Any, + ) -> dict[str, 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, 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( # batch, (b_sample, l) + 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) # (b_sample, l, c) + labels = labels.to(device=logits.device) # (b_sample, l) + 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, + } diff --git a/fastplms/registry.py b/fastplms/registry.py new file mode 100644 index 0000000000000000000000000000000000000000..94c1caaa3a2ffd46c281e2cad8c26d324f215ad8 --- /dev/null +++ b/fastplms/registry.py @@ -0,0 +1,1543 @@ +"""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", + "backbone", + "backbone_model", + "publication_status", + } +) +_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 '=:'.") 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 + backbone: CheckpointSource | None = None + backbone_model: str | None = None + publication_status: str = "published" + + @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: + if model.publication_status == "pending": + unresolved.append(f"{model.id}.fast:unpublished") + 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-' 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.") + strings = tuple(value) + if len(set(strings)) != len(strings): + raise RegistryError(f"{context}.{key} contains duplicate values.") + return strings + + +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.") + strings = tuple(value) + if len(set(strings)) != len(strings): + raise RegistryError(f"{context}.{key} contains duplicate values.") + return strings + + +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) + digests = tuple(FileDigest.parse(value) for value in encoded) + paths = [item.path for item in digests] + if len(paths) != len(set(paths)): + raise RegistryError(f"{context}.{key} contains duplicate paths.") + return digests + + +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) + if not ( + prefix == "fast" + and table.get("publication_status") == "pending" + and revision == "unpublished" + ): + _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.") + assets: 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.") + assets.append( + OracleAsset( + role=role, + path=path, + url=url, + sha256=sha256, + size=size, + ) + ) + roles = [asset.role for asset in assets] + paths = [asset.path for asset in assets] + urls = [asset.url for asset in assets] + 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(assets) + + +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.") + kernels: 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 kernels: + 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.") + kernels[implementation] = AttentionKernelSpec( + implementation=implementation, + repository=repository, + revision=revision, + version=kernel_version, + expected_variant=expected_variant, + dtypes=cast(tuple[DtypeName, ...], dtypes), + ) + if set(kernels) != set(expected_variants): + raise RegistryError("The manifest must pin both FlashAttention kernel versions.") + return kernels + + +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.") + upstreams: 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 upstreams: + 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}") + upstreams[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 upstreams + + +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.] tables.") + families: 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}." + ) + families[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 families + + +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.") + runtime_assets: 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 runtime_assets: + 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) + runtime_assets[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 runtime_assets + + +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.") + models: 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 models: + 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] + publication_status = value.get("publication_status", "published") + if publication_status not in {"pending", "published"}: + raise RegistryError(f"{context}.publication_status must be pending or published.") + if publication_status == "pending" and fast.revision != "unpublished": + raise RegistryError(f"{context}: pending publication must use unpublished revision.") + backbone = None + if "backbone" in value: + raw_backbone = value["backbone"] + if family_id != "esmfold2" or not isinstance(raw_backbone, dict): + raise RegistryError(f"{context}.backbone is only valid for ESMFold2.") + _reject_unknown_fields(raw_backbone, {"repo", "revision", "files"}, context) + backbone = _parse_checkpoint( + {f"official_{key}": item for key, item in raw_backbone.items()}, + "official", + f"{context}.backbone", + ) + 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)) + models[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, + backbone=backbone, + backbone_model=value.get("backbone_model"), + publication_status=publication_status, + msa_conditioning=msa_conditioning, + ) + return models + + +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." + ) + for spec in models.values(): + if spec.backbone is not None and spec.backbone_model is None: + raise RegistryError( + f"{spec.id}: backbone_model is required when a backbone source is declared." + ) + if spec.backbone_model is not None: + if ( + spec.family.id != "esmfold2" + or not isinstance(spec.backbone_model, str) + or spec.backbone_model not in models + ): + raise RegistryError(f"{spec.id}: unknown ESMFold2 backbone model.") + if models[spec.backbone_model].family.id != "esm_plusplus" or spec.backbone is None: + raise RegistryError( + f"{spec.id}: backbone must pin an ESM++ dependency and its source." + ) + expected_esmfold2 = { + "esmfold2_300": ( + "Synthyra/ESMFold2-300", + "biohub/ESMFold2-Experimental-Fast-base300M-step1500k", + ), + "esmfold2_600": ( + "Synthyra/ESMFold2-600", + "biohub/ESMFold2-Experimental-Fast-base600M-step1500k", + ), + "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 six 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: + manifest = 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(manifest, _ROOT_FIELDS, "manifest") + if manifest.get("schema_version") != 1: + raise RegistryError("Unsupported model manifest schema_version; expected 1.") + legal_files = _require_digest_list(manifest, "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(manifest.get("attention_kernels")) + upstreams = _parse_upstreams(manifest.get("upstreams")) + families = _parse_families(manifest.get("families"), upstreams) + runtime_assets = _parse_runtime_assets(manifest.get("runtime_assets"), families) + models = _parse_models(manifest.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", +] diff --git a/fastplms/runtime.py b/fastplms/runtime.py new file mode 100644 index 0000000000000000000000000000000000000000..49e60dbbe93b8b0f0d48a4814727dc587edc83b2 --- /dev/null +++ b/fastplms/runtime.py @@ -0,0 +1,69 @@ +"""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"] diff --git a/fastplms_bundle.py b/fastplms_bundle.py new file mode 100644 index 0000000000000000000000000000000000000000..c881bff49269cff54b4f789dc983fd4a202312bd --- /dev/null +++ b/fastplms_bundle.py @@ -0,0 +1,3700 @@ +"""Generated deterministic archive of unchanged FastPLMs runtime sources.""" + +RUNTIME_HASH = "ef1f313b70da02019bf0712085e332968c1cecf189633f721fce2554cff6b61f" +RUNTIME_DATA = ( + 'P)h>@6aWAK2mk;8ApjK1*rl8T002P;000yK003rTb98WQZF4VQUukY>bYEXCaCwcDT~FIE6o&8fE6#kCOkKrpD{WG>fiw*S0w!&%' + 'ka3ff)~XXnw!{T2SB<8Gb_BVutut-z!08p%#2b354F#%3QUNP8H7Yw!7-LCN8eIYB' + 'W$RjlocEQ0s41s#lFWLh3n)2{NcX|JwmQbG8{b0@OzU-$B%I$-1!P3#Ux0lcAB~SFFF&3B{lkxcUz}fFUs2#O0lL_?O)FQrH94{g' + 'q%2NjY;Yl|xiuT@A+u3uaa_%@D5#Dfb' + '4TP5wdRw}|2I8v*^8>b;ee=k6xG48K=iO!OgQHPDZr&U?i(QWNPkylorNUB~v8^RTf{%mcazjg~z>=06^&^}!oTs&@Z~|0J2A{tA' + 'MKuW(QE}akzq-Z7Jvsuflk?Vf-zn!H#mqmv2}-=rv2pMdyVnkc1@t%1zD7ukFpY)*OH0=OJv{P|f%6aM36{d|{^#WI=wlRN8H9x!' + 'tb_BNP_3`4?=c;0=X9%=Y`2*#=od=Qg|cbZ^+Ha;F~khgpF`b0h{G(TA7q%HqGO%WFmEQ?xTXA9r1bu#8}n@3)^ti~e|A}x{0C4=' + '0|XQR000O8001EXV%|RdDFOfhdbYEXCaCwbX+m6~W5Pi>AM0si@w13ctz_MM$' + 'Wmke$s;bHwoXKM0*vK{%_3Jw!aqeN^1vBT&aAxc?Q501u6;kDaSzG4w?8z8AmhwsNgEe_2}1g;=HKJlFE)Ff&zA3Ri_MP3yWM=VTfA-9G@dQzo9FH1#{3zt7SE$-?Db6xc|~+8H9xLNX7Q7uvS))Vm0%C7' + '6v(881j_{WfsSCjvZ`U0sxG$A%d)Jl>q`}dQ*Hi&Nd~bC7>TMwIja*SOsLb@6_LHxk~z&FNrV#@Z2T7orIp|y89yLDp0rS|A$H$v' + 'pGySOZHEIT+LB-p9QOKYXB4r#`MOBWE#r#OwMxtghdN29ges-NSR!48ZAPit9@3sw;|k7-7mUuJ3|cK|J|hDdh4TxbNl01nEH%Nk' + 'O2x-`FLYmOs1Oy&&y(P5`O4`+(4@f8wx9Y-yxqO8*4xG#e)+MP<%jYS4=~p8cIi>Dj!;`Y^|iQ)QrrckNb;+629%}dGBmcKw$_yE' + '3' + '(9ZJF{vPNcUH5pny9?U!_BQa?y?x+0z448j{!>D=6+H9)Uhu5P&F3>d*(F5#6Z{8IO9KQH000080000X046)9m}DOS00Vsh02=@R' + '0A^uxbZ~5Kb1z|ZbY*UIX>V>XUt@1_WiD`e?LBLA+qkvg^(#>2!&1F88#|p&M|;(&lQ`YXCL1@&ZhI$_p(qNnSW~3JhaIn*``hn%' + '03bj@vXb>aI&(XnjYYf;4$cb)2M6Qv_(gerl*=obi==5}(PU*IE@f57ycS8Z6q}^J6lq!1P1UAY8jVJ8*HUcTMV_T%k))RrszHsW' + 'O43H;*_EtCR*31QTu$e!q;9tPrjAVS^GLkf;vY#qG6PjoRI)DfE4dV9k>7}9)yPUDB2Drf{*WKHFk03?NhLqDvTkb8tmP&Ybt$sV' + 'wyc`0xEMWy(avAItVIDcs>HIqF7h&2)}pKc!p%0z0k*VU%4jqmk4K|bRc=HauiB=qWE`uWVBSU9BnV?Y8mVvBNmalQwf^0dRl4R5' + '(h^`K+|;^RyCfhAkJFm)*H*cXpb($6B*3^Nn' + 'swyciR3XhGDhge^YKxTDMkKX(rZ5!6yijot1M8}YSo$EBs(3USHPy{@B;bEe@7>}@t6YAJJ4#c1J;r`s0;Po_1qd4%FYrsep-*m`' + 'k^Gp-Z6lsj)3d57D;{#tQq076Wg!_JcU-H)jIBS8;+LoYjGw(Y`)B;*^vO48@!5B8Uq3&41Lemje;&n(ke0FA=hR;GcTKy^WhlzU' + 'zhKqhLsP3Tn2fk>2WstYZ9v=AL7Jp%$*VGBvVfw|Xen1BuH;5;7P1PclGA?(9DXwuYLu%aZ{>t>2wd^Js3q}$$QlXcSTrT@;i7Dd' + 'B`{|$fyIt0nZl&=o1?adZ!|!J3`NU;zuv(B2sHKGeU3Kem5iG*29Xg!l?fGHuK}IVbHSia#9zdYC14=C?zZI?fedm0YW6J2Yso!U' + '5*WHrqcc8IfV}E08RM@GCgSKXwDtxGPpfh(v@|&aF@n=n>n@%_CnT7Vb6HB#N|LOsK|W$9Z`WBZmNL&^34km?a4k!y41$H4iYc<%' + 'bj~cIx~?J+g0h*1h>~kcAZH`$>l8TrN{UrkflLG02*^r6' + 'UV}KzD&T5>4nD1bR*F>3E(K!1VR8w**U1WI6KPEGY_GFsE&4Rst=2A{iD+KIq^sEH3sXZ;^@VKui?}KRYXAF?{#nYaECqUJ%z*JIb0$LWz$J=Cy0)HHiI=+0(G@b0ALtIZ}TCFMJ' + '51F=171>3RpoUYV`NX1_rQ8B9f`TmU1~}m(`K{m>_uw1-MW+P>OwvWL-2<5=%@!EYF%7rh1UGWh0N!s(_g^' + '*$S+nnq`{m%aH=0z4DfDv`m`hi180_8I&oVFZFzJ{76xDKA^HrxdDp<>D;MuHA;Yy;9C{wwxJlaz&(BY7NqU-SKq}?pT9YM{NgNr' + '^6KUJbNKh^+2im3HVy>~+5z@F;{fOYQ($Tm{lnmjE7dv(5KR(EKj80^aM*zOb66igTh*Cx%}' + 'x=}Q4eFj)!QtHDZpV+@c(C-w%Rsn8ejFtY*VDFtX4{_+P_Z@8x5dno{h#1j06#dEy3qgv2lsQsZWn)O`14pzPTSRvHTQ>p{g&kW_' + 'gQf?w5a1!9vo1mOmLOp8LvmGSOBnXa_fJn*$K4d^x+;t8N7f!&*krNvt;*|2oEZ(mXa-sZAzH(bU{IDTElqjZWtxEX*wbd)93Ucl' + 'qu$CiOL8C!)J>?*t|ZFivb|VSM`|9_Swm=qTD8nqII~4rg56Nyz7HrwUVv%U*9pVz!cZ+F89?~D1p%bv>|b&HH-LhQ78Ug$O7im?SkVI4NOeXOb<;>~0s}l?^{zqVqR~qF6BK5x3|KH%' + 'p(CO|fLH5AWtZT@;tq`(6CL8l>9!poj5qZqM(Aqj&+b0D%H(yxE8W6m{yd3b`@X%Qb)p7AiQBK)_{o^Ypl{v2DDErD`%s2KCs{*B' + 'poRlEYO7G_$q-wwM+XJ>pk{k&1GVkI{zo-yaO5c4ervDJ9S$+zM3<*+$sZ0%ER+t)j)$KdtPe{2t1ztLaJXB%(=6y?FziSy?4x6^' + 'aWGg4+bZZSiERRQO~6EWAiilYE|A@yC7@L1%HuajJ}zruq&m+o)+i^I*@`xujW|7jUK4jegMIz_)R^rj&>LtblQq3#*1%!' + '2B!0vM}-!}$#{3~rc#W~gIyz8)%2iLd*7&!#-si|NSSx!=5vq}WN1m?hAm;g7hr}k#QR-#d4>4(y|G|kKrgIFa^3+NCD~tO9!eQ5' + '*o+*X=t@zB3Torr_?n=X)i2GKwM<54b6sb+n{S5_Ika_Y8aeZT)AD(=pS4uFm+_-eLcipT8+Oi$X5Y)' + 'M>vQQKy73zSfrzER-oTXK|1cWyzQX;^w(P-n4d^y+`5SSUG>Q}t#}0q>`@;%n(0vu?UM<1n@Q(k@O1P$ca_j|NrB{b`R5?+km<8X' + 'U5Z;DFo1}!2T0WriQ9c}_Hd}Dj1C&$roi&a@-`{xdBX|@3>jqbFOxnr*FQAzdd{X-PwybQ&44T^A2~c9^s4r_uEfej)S3*Ez$6CiJDBjz}j<@3%T}f)X`ko' + 'UCM^c?sAXPQ2SkLV{YwuaUN{x)}m~i8g?GE8rf#kHY7xppwRLIc9cL&mUcJJ=?Bi#LTe&n5A0wO3P3(s$jb@pfuK=8tX*cTfCA*&' + '-A2kw+I6nK*hBEQPnw1B%){w*f->W{qv~=|%j$|;M#v1&zsjr#pgM48(AoUNMQT@2xL{hWb^wPAWT)N?*5dD>_?CUZKcM)AeNiQ8' + '3rFL8E$~PvLr8Bm%tN73@z@L+OlgQejhGIyUK{N?TY>&)9-WwK2$<7dR>x_(Oq>eb6iSeXQ&8iv;^AP7bv;6*2QYVl8r<~23}m;b' + 'FEI)=B9flRM*%)X$wq&T`#zxs6&(Kws^~ulb-F5UjZZ^76C?nL1r}|poJu4y+)s*fzRcvQ+wGR9^otE1riWS+gl9@9rh(keP~EZ_' + 'wdQb*uJ;o;LiJN~oPx1lmNLx~v`?8JfRk@E7#_vjam#RGpUAmFPB-!R+3}x*rZwIZoj*D`0)`|{OkX#tW4;|4#u@4^y+B*e--Kdd' + '$8M(?NFX`#2#isbM1(PQw4jtht~L;f-dEhaQsD8B(*swWslzi{AE-~RLSQeCtr2SAPSx$?Fon>0i&zz?)X_F!`_G' + 'nrX1!J$$d{&0ivtvoK?pr-o?poIv}!i|LuHRW6W_tP%cNb)J2E&!Z#Jf{x-%T=>UKE^fdf#g(>Qudwq' + 'dsYuQ*=!e-PJA?WL?FXpL?F2=j=JM)g@I6%Dy|OnQh5cN+>B%&`MnrELGUT(%0Je!5OW7|?L(c1)QZ5;c6PGe<{1W?HROy*+MN9=' + '^JDP`QGckKU@c+&m)T}AkHmRW*S*eX$I-(%(ia^Vuw!nW)WF>s-^QMcZ6%MW@>2XA1TCK|T&a0du+id1oV|H@5{jqiFJAJ0;m6r=' + 'h@YR8`EtT8n`cRu)9g!)3n{K~#((EC1v8GG7omxe&|D+SQdkc_c3CU;e`SvRWpaVRiDm(?>YJzMr>?W;%QVcuEfMZ?QXDQO_PZ%&' + 'GuHuWX`8m$wm6pEF6wP~1DtAKNs>5XY&C1(g;iE5Kb-N6&m8nPWtTGxSx(8-%toDCwyer6Kpbyp52J@}kxPiqFSjDI_>BJvM?-JB' + 'l1PjhP+*AFtj0sGqyU47zpL%|gb!yF7JFFwj3s&WAG;*?hxdJGb`LHFE967`prW^y$1Fe}X1BZXI3a2zlw$5TQs!b=rT`)~N5l1cedAS<6acr-=gu|iJ*B%)VJ;;6' + 'X=yMORL3gyi*=B*t9js%M%W{SU?KiX9Ebm5y2^6IvV3fQW#MR}c;?`HSG7yK+NE2~;4UK7nYuKjQ2$q?)GXb|W?e3o-)E(jX^=J_Lofhn!LnjW-YLIgW<5kzLFE(_L_hgjH-b-_' + 'uv3~VDy0i$5FT3OV6A8^kP%!t1w' + '5)K$d`ZTd!Mnfjx`SEOAN8y)Ghq}}z2>SSr%q)`p^QOVj-8B(hyAA^MAb`mQG+M>n-Aq^r0Ui#W-^bJUUGSPr%Yu9m(8%fnqw!=O' + '3>aRW=OkE@CHjWRtS3X%ow-2K?RbHekZKnq!0O{kC>|Pmj(~MR(#s!Za6G}&bl|<<=s2{%>js$MogajeYW&^;dNSI>PKLyZpO6h>' + 'mR^=K(2YYzoB6%XSyrRzaN>IP!U@LADehXv>}suTBDDJ1%T~L+tA0u}M1@z;3Oc8tu`)dPAU(U1*Rn!wTX?8)?DF4E>T@a0p3IQG' + 'tWv1k(S`NhOC6fDwpF=o(}uoI++JCX=p)7-WmVQpk;SfN6o>BP9nT!zzwa4y99U@A^aUBxNs(KNP76IMxtl};8Xe$Sv@z#5cWew!L>#IApR`F$8%a$a)`Pu{97J>a8q{cv7&3Z<1gy#x3BQK@RhT=w5XTRW9' + '3pgg-bp{rC;dkS8a>4yf<~)1a1aHge$1W$>gZ>~vGZZ|H5@}}-ARH6F2iaTc`T^~4E_A6PGKJwoH`@%0C=J8xQvu-29f8gW$spBG~iG>{sLU+QpZ~SZDWU0cW' + 'ccqVr53e_b0!;82PY8h&^g*Es0i!XjyEe)}0p){%qPpD#&~^rP1k5AQOF!v3tkInu!RiO=YzxdX689!@k>' + 'cHYg-twHqOnC_q;00=PqfU?Q7D)a#jyL-;d3H6F`(e_jSHo`>ulcx{^|(WcDmcQx;lysE$W=;YghfSE_OXB_mFkp>E}III#3w;4`Uc|K3|df$~uwJ=>_O2' + 'mktNOdhD6iiHm+}l-FJI4J^NgdoVun3ir^rygwkjL#79L' + 'zWEK2p(56P*ALIO!Z1IZwh1B|;BLEUUCZMTz_BA}hw$P3T}y%E(Q#+~Y5C`#zFN6r@aX2(mw&#K?BB>)_&Miw>z{?gKftToGSj{{' + 'eI{HYiBC9%l2|31EWhcx' + '0zSk}=Er~N!K;!S&{+(2&k+Lvg-)?1f*o;e$RR40O)N{WLU1Y`oKmPjBwnGL>pD}HmF^1JWAreMUL+aV5I333IYt^*d9{+99SSC-' + '=Q7t}ad@h^tg;pKmZVh)_$70uzPGTq2fOtyN?4Pg=a8WFK99S`S^t~=h}C5%`(k_lV#op4tA1X9fJBG4X;6qekWRgj^eqP?Vcq*6' + 'JL9N2_~}~e7x!b&S-YLSq%WUZ^G7oetIhVC&X%GX-t+gK>+d?%-+87#fT{0o|B_gm1M%TxjEH&DyT_=c{m9RT2(>xtM;!#aBayXn' + 'F0amf!hOdvK}W2EB@ZpW8S?Ywuwd9|Lwk>SmF0QxzXRHPHx0$q@zA65%-gk_S40i6;|9l0*nS;){3TZy1A)G;Wp%Vl(u`6Xsq_{U' + 'D~j1r^1Qqz3-2wZYAUM@M*mnj4U0xr&{{^&)wfb^#e5E}=X0^m+ZujgLSgt-m#Wijm8B+Y=uK@>8|ol7`7x^-4k-dlc$HnWm8!0C' + 'IPtHFf+zS8K3U?f0KlKm-!v5t7{!3wDvmJ~86UKo4dP>-VnDaovF>Dw1W=S8lBxLq{NW$@a_`yc-_Bm+Hij-;u!)IxE$~yl+$Q|R' + '&0NU49y>Fau7wwmx&((OzU|S8zwRUdR!X~>o+J8kJ5+Ba8~dDAl{2ed;(Byz1zU-hGf@{Q819jH-k?5C3e_3qfd&*OU^3JFH$ZKQ' + 'X*Zy^Ufyk9%g7)Bca9D&)S=wADz~y~Zp?vPrmtwBi;cE}xbYmYDkoGctZU4aLbkVv(EzjI=hc7I)yIqlKb*dJ{?sadeEQ_uv+tf_' + '3TF1M3Z!7fC!~V_q7#)V1&t1X+j-kuOuqlHMIeczDC+ZqQnaP`*u4Q`3+hd1YI9`uV' + 'koiV9C#=C%&>yE$!;HH<(SU49IoU8xtICz8GW1bc1Ca}?5y8eRIC`L%u%nH30UX@4oF0duh6N|Ed#Q<&1?Fd+Vwydj0mDGIg5tgt' + '3RB|{n7CKHXk}e%ASq=rdUg_b_|c2Aznwn$pZLi)XHUL8fA##kxBlhr=P+mU@|F=k^%0Uy>`WD(M%oROCPi5we&cR^EcRU(5d_P+' + ')IX-' + 'j4r>r{fE8rGzDSW{~^W1m_`c!4%E;pnN^>J`s5ngiBXNS<*-&DziRh>yhD-y06^K-g>0xyDh~Zk={*CnE(;7(=R1`|RTTEid}J`+' + 'CFYvx_407cV8)ZEOZDD)k@j^ln7-WwW_;)XnRQ+E*Pqi2s<#;;{G;2DQpz7h+9#61U5B^kPrSzMyhdQLh{oI4>ad2AdUeD6ipV8N' + 'ZJwRLoSq~KZw_l3od!J#r5kcdjo*q^0KO{+F?c+Sr)RYj%$-`tLWjazqeNgrbfZHChk}$Rk!*JmoU6^;J!I1b;eL*txO>lR=-RL*)bQ' + 'M#k?z+Q~ps&bL2r0V^YCJ890k2)PJ?k_gf`22>SKj=E4' + 'w2{k_aLt&q54^(RA-;dEcM1MklnytGxgC4@6DQ=Ia' + '-uRgfhxWJ&oIHlT_Je)o4(78`7j}o0p4^3pMUJ}Hes*e(+W|11^n1M&d?=F2mAKut_ZdWw_H!9#D+@ofyG&)3rtXe+5$+`cHYzpE' + 'u8TIKV=a|t4;cRnGc#n(&ZI?ZN{THn{N`-HnoyhTmL^bHCU>D$x{=SxJ&^L&OHIY5x&9P}9QoJY^F^@8' + '^qQH)sg~zB*tPBP+koDOAymNLK@UH;qu(8YaoCGMcR$(Vk7(D1fbf3oULpIp9rF%s@|Yy18!1Lft|>h~*i^E>w`VGI7*7rP$%tAE' + 'qn*P3S<)Yu!@S`ABmF)qJ%{pxQ^vhpb~l{f+y0=4x?XNp7l1

h|R41IczNZasi~VlTOK+oM-)S6JB~ZZGfzOi%WO?O8EzJLnuP{{vy8f;69o%R@Y_*V(~Mb' + 'Cx(dzMUi&(4;axS8+w;t2w`4c^}u_l%{g`6O1k^~m4ok&{w(jXxM#1Z()7C!ieIzkM-?qje|_M8F0b(4UK#=n9qZL}${&cAL=h^p' + 'n93(M?TLNdc>qn+<~l5p6DJ9Xdc8eY74V=GBf!A$$khw4bRb0u1*o+$)GCftW0z5{vsG4Pjr^LPQ&s0Elz_CA3NC)rOe#x+odI>^' + '%fm1{O14nRn0r6k_ddwZ*QT7eJJdrZvprE0u7xq%^{|Kig2znX78vb&#lv4~&r2RGs(Ke|hnB|7_smNF#gf_Z>5|F8T;E6cKmZsc' + 'MW1>ik^TVyYLnbpCix8!7N!7X+P#wo(TCc-CDc(6olq$1v!4-C2kvy_gr^B;1RZ4s6r@1Iitq+Tj_+R1J12@E6jPHd?3n(a?~g_H' + 'no@n)Y4;Y|Pivs=!RJ&EO(aM=Vy2>@(Ynh_-51aM1oJ22d#)Tmi9ddQVRMzE4Q6O%zwdBQ1H0Bm6$' + '=sGtOa3V>%uX%Yv$8w;9gGtDP2VdbBUrki(ta(7HEHv;c&ugV@TlcGeP;8F9%mpS)u5oH}oU95Tev2Wj#-D~6jcWA<1;8Xu^J_d2' + 'gE^EbyW$urY5=-zM2RsAfzv^!MfcPz=VPH))hM8ygM2Y@7tk<(!H$ol;NS%p&*aEavmpedACzF~nc)^)1m@CYw(7ojai5*()3;sl' + 'Sk;^2)=Q9`6W2j^C@xBspMQ4S8~)S4g445gE;z?|b}2iW2yH>g=9-zJ?!cy|_%3o9BO#0@E(>c&O!xl>P)h>@6aWAK2mk;8Apk;!' + 'S7Vq6005R7001HY003rTb98WQZF4VSbaZ8IbZKvHFJEhAa&Bd8Uue62msHMKK#' + 'dM_Cd1CgMGm;zV;w4?Fpf4^OPkRW9%&fVNY{g6obSS)sb``cY8Ns=q3WpmF66&1Os4X>zSHo{(UOrnrUYv=S+QCwM);teHUVNvl{XH!W{iwQv%M>V)ANxTYqcGT}lS+Q7WM9*Qz*>ty1e*C&{Wd<&cYpvUCw' + 'FW(*6MeK9QT0`F0b*D;6MW`kWKMYb7<(iQzcd}iS_H$PIcY~hftf6{+Vhj^#&R_r?K5SrIzpfcY2fznhopX^Vy~K=Ii&`DL6}|b2' + 'MwWL4`oP%_&EMPaih)FDHg76i?PO6v%!*KeMNLcarXpH)3b0eIS$QWr(_7Dt+cU`09*_cB;9^CQ<;r0(&FoS@8z}*n4w77u!6=cG' + 'yrl}S{Nd<^ye5hFM2@;8^Mf0Hcy2=$`L`Za$ho$frqR6ho-<6pZZ`XpPe+F$(l%UcgtvqGdJ=#K?dmiRenf$+D&T%rMp++Myfh!+%N&F1=6E0CRtgG)7zdYO)Pml$(t~_(9IXn&KAV(@QdXe(GU-gc_*;@?aTuhR*VYtsC+1(*mlwYPvMXnx' + '>}8=*C%+52Y8Ww+;K5T!9+Ssjvt2SmyLYJ#Y;8kZlq`*JI42F)COfFW*nOK$8&XPPC_t1FypxrYCs{a;ifuYM(|Txbt)E$>+W|y*' + 'tj!nnIK5}EXQYi2V|7}wppIO9Vr8$;h8lnm`tpNY{p|6myPe2ynMo6bPr)lN>2eui=_XvWVIuXyGxE)n9KD?3>H{`qp8@;8+Cb)H' + 'Q223B;2X933C23eV2`m5yO{K5QK6m6Xkt`yg)s)|4Z=vSeD52ZZeBw&5j59V@O8nw3@2tRb{pC2^t4' + 'jy4o5c&h@EZD+mkm>eZ@l1PF78r3oUf6H{Ti*-f}^|z5aXX2hKh#~Zhpc@8zk}QWn=53mR6(5*_i=g>N=k+2`LELAFH90>k&VT=B' + 'ad~$3_RYm9D)9d5IU+(baPT`!!#IIW=&#hO7DL4(E^uZq~ql' + 'q_eDv?l^h(CbIE}z-;`l9iak^vSQu!1qNUY1JdY&t8<&6D)}Ig_`+Tr4z+A5rV6y=|CH~OOJ~c0$SBh2klrJe7O2vg*sSQ10;&ro' + '!9YH>PH>X^dU|zn`u4i``Q+zw@Fpgd-U7MDfSsG>;=!ju{s6}4*w=eTNj+&S2MI0@6;RQ!r4Rc6tjVdpzn6OoK;{m!ItKf7xS<64' + 'LYlo|wF|K6%^~zOuYfv?cGnzL7LQ@wE*Y_I7#?aBC>jEV_`=1Vluj_`c*GAzT;hiAq5S{=TM_y=Drj8VyP3}{=eTb@>-;j=p*PH)z}v>nj;MP?jG(V1w~@umu))e8cU+leGb^|171xOPzXXqW_iWyngME|b{aW6>)}7h`Pw-5#(6HFdu#cMM8i=IuK^wayS)g1H7Bp~' + 'Pm+stP%PWUVf##~3Pj*R3B-*zrHv1cCl}|P;lq~D#~<7gRFBUIW9~o%sb!3Q{Lj(SMsV#fxqRY&>3~uPA{*QT#vlFiV5$OG#lk$0' + '0vFSBT*=0T^}q|^9G^O%*~|F9LpkXBA|(68Pr!Jf%sDA+z8)hdofs;Mn}%WYv~LKmz>3tV@DiUy;$S8b!0E|X*zq*bQQUb%Il0h2z`oLgI4Ry*@u#fM7^qx6C!(R{po-9Ixrw5SgW74PP+@)CgW()3#+|@~)Ya_t>?`)~SXt0}$t!IOhNJxZA6{uMs_vR*5k=#sZ5Z$$wv_g52XQ>A_q@wT#O{h~EG?En)I8u|8QW`)_F6ATo*C6<' + '(5P5!iJojdM*0Hm*|*^TF44}zO_P1GAxb-2Ui*(lQ6' + 'Mb)7e=c!&`KC;S?Dm`a>9OLn0Pc=X~)X@G_WcY=Ko6+w&I_5IgxFqLle|n46fNOz_qH%F+V>2h00yxvfrIq8dqn8#=sO&E|UIP|d' + 'ie?MjBQhhMcDi^*UD%B3Mfu7B;$$W78L)FnJKu64gTg#gh{j^ga60rd8Uu@qK&4~CtOC5pCQKub_3ntqoLiO?Fbpe3Id}P#^ZYLg' + '2f22vWH$Q`P)h>@6aWAK2mk;8App)lObK9tt8si8Vzq02y1&`G3FN1waBMWjmhU%;h1LNMNzpAK!jhD2k#B#VeglwHHdCq;g;Lrr4B%' + '@um?~Q^<;GQHqp4!7%oQ>*m9+?=;)XOq4o~qG&eDmE5x=$=jw?B1u@Wucc}juPWJaoT_It`>c^F-I)ogEUjc6^G)hzz2kLVR9iD9' + 'Z>zMCvefQv%~ezI(!6aBIMTg3s}AOM05}N>)Y6%Ou^2^tzTs' + ')g+z$l8rxhE9g~~fZS5-VF=mByu22w0-Gjz$@R{YfJQJQl+d#LB92MdAU%~;UTmiqS-7p4m6B4X*KUVOd~OTyo>QU*!F&tWm}Meo' + '-D6KHscyK+jAbAi+ESbv_r_HfzoQ36tVz3>vTt' + 'JNY|nS{Qs8EJ_luPH74R`tNPU;6d+j_69xGMFq3YV9id+cDrN$YPVYq&^HiP&b#r9#QAyWB^$*%W(4>n6~!09{=H};BeII{M)9|nHc5cHk1^_53G&~x`&M;o2B;th@44l~%d-1bU' + 'Qq8+}w3{in4P9LTvQ)WHqDsYytjYts*@?>Yq}xOk%7%FmJbwx$WA+i$IRELx83AjvgIAs6CW~7|fgATS78Y?@4bit^8-T{ZJcUlo' + 'OE3^T0Cm*d7-h+Vmv4Bf#bROPX_+^Bew%2BX#rw=iTDgOip)qGSF3xmsy8P^L!v=z;G4gj7z$C{Y7IG`C+' + '3JADHnlM#jTL6|GAO*mhkO48sS+D;oDxrAURnpju6KcI?ZJnXn1KfMgQivXXa#Li2?M>NVf^9%ptEsRzYzHaEZe&{`y459SB;wIg' + '2;M=;m#8u3l=uP>2hxbEfU-2TRx%Y@FX4Mywi&KtxD0D6B>?C|h8uMjiDz%l&M)B8d2)7f@$>tOx4*ql-h6!jBYrrCOc*nZ`C@9$' + 'yR-9Og7E?Cn4Z(Ea~Ta2?PYWY!Ymy!8F#!$PnE;fgh-?a-=@Z^zg~@3qz3xWqD#g10=+k%(wI%mUl_1I=7dX6E^7a%agZw%tzBUB' + 'ek^W_sv`;Khbo2gUlo6fV@1Aa)Ks6-RyV-POpro_g~662yrV}s-xp;eFzp&{g1Atu=bC7J&=Bi0YN})Q8|JPu1nBYQ^A9UW!{1>%V>$+T?HX53hQXfjQrnTtV6ATb8w|%T26@05~mKEoNs{7Fp)N' + 'WZhsjW_=kh;K80(MJ{y1U_T0>48Wi%a4p`jQh*nU%*oseLk*f6IpjhX35VrBu#}ERd$%Epo2OEHbU^a*;UTZWAvtJRmwP|K2KvDG' + 'z6eJGc)W`G!V;m={O`_HMworGFs>%{55=;;knvW;P@6*}=gTRYO^g~#Oi#Ojo`nN}2qBfkbjagJZ0a;%QijES!YV$&;qYG`eT4&}' + 'QF0L{q;Znmu}lh0Yg<#Zui--9sk_%q+}4nzkRv8W1~kr!f2`(raX8&KE#+|nRV(4yAAkf?nb_}*4DR!g;ZU-Ea26ul>jZdZV-`Cl' + '+UP>clUA1UCV~GQfKLr1$HEf5Gw@(!j-?*NA+VW~!WVgsX&lylEn5Y=nUV?%@^u}^z@#wJMG6Q*>c+_?^=h+CdE;R$cOwuVCu0c)VH0}cVJ!PPgtnDyE=PXPAI^*Uz13!tC{96b*wpIb<^' + 'yuiRyOu|Cy2IE}ZniH1;rQT@(=N0lC#?w^PL|!^A3)`k)J1MVY-(xdBEys?(KhvVlR$VH=IW6mz?Bl>4dZR_K;5crv)P#89FGIjm)0&%g;N)nOnnlQ$JbEaHgme<52cgKo=?1EkeJ`6#CZ`|;yXXKV{<#Df@}Y@E#(Q;|I3' + '+1c!`i}s_f#cvM!2Ls-(VMa<~>4527EZK{`!xlE&Cdh%tSisKHipk9%P8+U?o7Z5em)r5a4VB90FIvl4?7GuP3ON0Uz5T&6q)+?)K)M*b|41wbS+A3' + 'nBlPot3HjEd0wQjh$TMQ*tp>umqUO=ZldHKyRy6)h%un`_&|IXA=`NLt%j78e+tr' + '*5No&LIgBmby=j;%~-GfoL{fWnYK{^t{m!CS#^4Y2Rs%iV+r9Va=f&{aT^ZFX*w3^L>!n_T4ApUunP@4J9C)QFG6ins0K>#3=Yv-' + 'CD58xHdIdo?BrBMWeeWV{Pr5eliUBuc@rqRisKl3M+I6!eyAcFlw1I-Uc>JkCXNy6{`;^Fd?#Y7tHG|c1x2l|XZ?&pXQPk23XB4v' + 'kQ!D$h%wvba1QgNxt?Xw5IC8O0Oog>C~IvzS;yO0HBBT>7@iWk7K-cUco&M9DM|NOY4J#E>je7aALWF8V#h0v^rfC~8n$=Jmv{k)' + 'kFK6_>J5t~w)&`#haT>TZq`)_5K0ak#oCPj4C%6YVu6MRx3C(437gxYy~@>czPK{SHWzg7*9Pg' + 'M1x-s7>bR=g$H-28>;(Ay#qS%_4)HtDv0;jbr(sdZ@jVr`#&es69>XEi5@*!PO!uWGKTguOdHzlVk)nX{$^l)WG+$wb7$Cf=ge0<' + 'ABUbBsA=FhFz7<9YDMsZNQP@ko6i1#|3)KiqmGli2NRUkoSm&)+g`MH!by@9DIodjWModNNiSNZKoPLxK+a{i(A8+t@u8!0$H-wJ' + '4%p?Fqo1BE4A&1_IFS?P*Bc#J`@f9w+m-GX3OpPl9U#;vpvU>aRGXm0K~0y~PVaxSfS8B(pWf}8!^5ecJy7@G-4PPW)jq$mml=3~' + '=y1cWWiF#FIlyYW$D0vnwi9z*Nw@!!;=mCed-->=E#(GAxZA@1dV|8#E|PTjoRfD%-L07U1q|K<-#>vamCRs6Mw8xjD|XqrFYqN=' + 'vbz8ae@ifE;~#J3dvC`#e)kS|6-@yY9u`q0n_(k+(ZH)|4le7IR`g#}q>g2yrO9KQH000080000X02=`jzW4zE0O17y03QGV0A^uxbZ~5Kb1!9WVr67xX>Mn8FJE72ZfSI1' + 'UoLQYjg(Dq+b|4<@BS5nPaBZ+FkrjwkfB)@AZ?1I+hM~XuomM8QzT8I+oivLie>pDc82=q=cS~V6sbJV*G`~AF=GzB6cxUMXK=`@' + 'GlWtzJUpThQkSeEsU%7CS9zXi+0K~*yRycplPC|?IAZtf^=f%{!|&#|i|vfv2Nm6nSt!wg{>`6JmV5HeVpnO6F0Fl&vG6U@dAIVf' + ')U%q^9=HGkj}(z4bbJS%ApTo$9{a$f!h-l-3W1sjOvu|@{49Ihc<^Lga5pq@z4$W!wtC>(uk+0{-!6YF;-+9NYQg=VPzdA>3*&^3' + 'QepHHI^qgTD2`1?luB@q;Rs$BReSQVLqhG1^Ss1R0U&8~vyarF#-QX2w#QUW8I%E$_P4a}!Ekkzq8{iUF&2~kj(d2T0GD=&;&_)B' + '%l#hSfA|?03{*6kwrr_wMQ3bN-^B=a-d6B(e5n0xa2Fb07tQTN_#IS*CreqzIVi?;~2fAJMwoM>J+NflqlQJds)Jub#j3*UL1a$%aniWJ@P;Fp~+sveQ%eWUsekJK>3(_5T37K}jymz)o1|{`**^sKIz8czM7=#INihP)h>@6aWAK2mk;8ApmR!j0lJe005IF0015U003rTb98WQZF4VWZDM6)WNB_^b1!gj' + 'Z)|C9XD)Dgty)`e<2Dw4_pe|*50$dA>SU%xk&cUMx3h}{)>C6TyYL_g6k4KfZe&WOqHCXf5*{iR1OY+DjJ4xGa|yp+e*pt#~fV(a1d8(P~?NuaJUX' + 'oa9-?5;ZlXtE8szEBkHFM8fnWHNPJE$V8dT(P*^z<(Gf`^Zh^K_tE^_uVhAG>Sq{ZmG@=5V^oX-GTzY#{1Ca^sh=$6dZ4UQ{J5hM' + 'Ka#A(kAljIkDq)Xp0YKG?`XzTTCtdmVqcZs^3re*$h`jQlgW?f>*a{R>lJ}%-YOtjKvMq)oT+~KYLu@D*oN?ui?X7ip|=bgO1bO~' + '&YM9^9AA?1l97*(1$!rDF1_&*44#rX+3m}UtQe`(iZUEGESI53`Hp-wBd@!3{zJ39A+!ccipG}SF;Z@6!K@+y*??$Ysy#S{w+z*&' + '|K#vM$(JE7<8_w9-iu9KT' + '$ck3UHZJ*Rc1a+XdN~S`_nbaVkORr+Ba`238BIyUfJBgUO`zP|m>iJ#H#gx(t-rZZsq*HAYMCPTMqUg{&$}_rx{j{ui6lkp2I)`$!pd0pqirbKUU`0j~%m(jSb)U=Ir7lOR7%5uHBXEC{Ed|!V1KdI~Fh7zFYrRLc' + 'D>zx!8a8O`>1?|38^`?3(6CP^;{fl3@jK8Qf}ZV)>d{fMG+b&g(1qh2g)$e#Tvv>(v3l6dh#>5=--6)Dt*0}Tl?#%zj;7j~gr%a0' + 'u&iX4Epm2oy(SF9dP!}nNS6=DDc`oN5}B+yy~bQ@ole!Gdhq-j7djrr5^4&I57*2?7v?*EYIq?w0Z1T>dFqFN7Kc@MiO{eA6rn$EB|FXbEs@*)+9w-zi1dL=1~K6pBfz40V2@2Ij@mZS' + '*|>mA1WtMz-!s110v1#&q6QGEpJ)0bg%Sms>QW;bu=q6m{#!j%K}wgHbL`+bn|!ASVOL@(5R#f4k0=Pnr_*3Gl)?W70Ayd;BuR{D' + '1X3s94XikVZGhNaJH1wtKfp%n=#Ic;(@M**Z|0QdmxqA{Yjn?Z`v0Z<&h5EMs8K(XD&E-3C;&G;ZVKP(G08I2f+CMl$L' + 'L?koOF|xrj+KB;zU8%&~kq3J|3$fJPnx{ptQ}T9JhYWd3&X2|WAC$vUxe37ImT{&%XPaiKBNdb~uHw!W?*!vCw6s8$G?-|5%rJjc' + '2uJ01_6am@Y+I!}Q{l}4}qtkOlr}P)m;r`4fgLZd@@Hkfa' + 'p}~u4@4tJav!NO4j$yHNAHBX%)E3HCDFej0^?cRzsUFXA@xvs8z>IfXcvFoiY(j@Z5SM#^5@Zp=r-B_5z)sjJ;P;Pm&zd#odW|U-' + 'R{DLG#Tmb4-U8-?Q}T`3t1h!ev*douB=b7Mt~@`IY0x}e07Wz4=wg_2vCchLTgbkf&5o|3qp9}1(EsTna0zwKS7p%6+okP&PXz~gGpZ>Xm0y$v=WmMB1V`#>)SF&s!' + ')RZof$yPKIu+Lgf2cF-+9(L8YXR{?H*~dlNboKvbtGZl?M}h~G*5wGC}T' + '7IddfOq=GHpdrm$FR*za}n5R(oqWNAoXZ!$UFd|@Vv1#dLSgU!@bxV<3PkP=6W~fU~' + '>T=kNTan)jupOR(Yw02KTF2zBIA;Dc!mP*tHgx-8aY?=a(ymzQ8Cn+ms(^s4fOx6YewvA2ly!S%R_0NzkmP%zrS~C6pFg!MXemA`' + '_0RRxp>WoX4dC~EQ7E}U6HB|L1O{?gJEi)x(XLPO;u6-e(u90!PuxIfd4`Cmi7SKFxEc-MAh5TfoBIm`-uci>Q){CJt2n@lp;!pq6B3Fe(3PzLE}11ouI>(d3KllCpR18xg=!c)^6I$' + 'y==GdS_qZeV?6ZkkDl4DCmqJZBh2dHIruTKI9}1@mWi~k5WU&G#((IFsV%f~+Fqa4_iwU-)}rnyy;>Kuhk}X}yblliTL8!nCu^2CeKu=W$vLWs2=+C#ihx*^O;s`2yeA853s4(pFm->zuZ+xHnLf{@B' + '*w1d0Vs_UF1;>+nQmBZ)8+556$Xh}!MrQuYeY<{vCR{c7K&~czc6Q;P?t-=1Yf54Ho2)#)<@IPNjinlXr))Fcp&!`et4q-B`XH5(' + '_16LAUXsM)x-r9p*a1d^wP$daD!cU+P1C_wB5wZlpisi>YsKnoKrRHkI5Iv=t#iDZFby%9>#cBjt$$SQ5`kX3DLC!g$>_ApQ=+kcX#}M-azDD))0U6cyh*UWh)>a;hTzqCS_h=a_V5xatnFy1?Fpa3raCn{rRDdRX' + 'h@lKZqQ~~$@bKFM=XStQCvuqnJvOJ^Tat#Ub8xf_ouiFI((0WFX;DDnd-vogwB13OHLW!8kJL4NWvb~ODqKen>pt}l>t2?2*5gI5' + 'lT%ptEBSQfe*#4P18O*bj9170d@zu7+S(e(R2%t|Kk8y}Z3AeEJTq4peis2U;eu}k(?_FG94p|Bu{3Tqpje6Z0gV2`Roh?2*Q5Uf' + 'P)h>@6aWAK2mk;8Apm?pnJ1q)000^Y0RSBU003rTb98WQZF4VWZDM6)WNB_^b1!mrZf<3AE^vA6eQR?YN0Q)o{)(P<+__r=!jkNr' + 'n7|CYD@*iPJEEi@=`kG)M!N|VL|Aya8z4pNxWE1KQSV0spyj#Qi*@Xdi0x=(X+%|J)A?%Ith0KtSgxsAdvGA1UDxgPYhIO8UQd?IT>Wl;nN8Q?' + 'OZ9Ut=BvqcCiHJagu5OsXEQOPj{5c0NcI1`p3Um3ndoIN*P^lhf=_$dalKl>czW5X_+=v&BcWg|*5bn&a8zLAlX<t?Ko)' + ')#NSwaX9zchJ3hR~%&U-)TJ&YXZ@pt{YL$e`tVWO^;r}@L&!r' + 'ref8I(Xtu0{A#+GuBY{E`m4~dHw%7+h(hCbGh15^G!XN=HGRwHmH04i*R6VeBJ-Y6>&t5K8o3;1sgvK5`G`PL0f4$?WI1v-kATuBa2M5*hv-9fhvy-!z)ydK6%NO4too2)AY4!N4KOa;l-@mTDefj!t' + 'M<;JjUcNr7PM?2wbo{LP$I;2@%Qvr~-IwtFn-@o~s?)P)XGhg{&rZLqo`3iK>%Uju{CGx?p8P3$ls$R;^s7gYo*q;$j=p{N{j0O;' + 'n`dXwzpH+D`TE71AF5+`^78Ggmq#bm{-8Sj{_Wd0Cuc`50PLHSXMa1YzI}6Y{Ok5M>|Il@uCprvD_CT6lq0}ql9((d5X@&8bWmLNvRO&L' + '7sw;i@doC22<_L!aeuum#yIOkS}A>a_@}SRe!H0$ps%yYwDBnLzowqHBKt=@+lZs4SvJKaCkkZk' + 'W`$hV!k}OP?K=AygVVduuBOY`^4)Y)&ob3%KUZ{Z1&X5TdO8OAS+#4Vuoj*zbMUNg-4=@I%#QY(c95+%tC_fv4YGc}f0_Lsni+2W' + 'DEz2G3I1V#1O>7lK9Nfj2r4R&iIvyE6$lz=_ia6Ek?(Qhxi>xar!=Sx`jq`-^eyND{A2V!6TlI7^5N}teZAa(+`O(A?@*rIAgxP>Ow=G49n>=5t>#WLOfMVj*-i~bPcm}Ea-AjKM%f@w0EIdJWiy3|mRxNUfaYC6{4Dzp)msK4ZY2okg+#lT-Cn~=$QIcI' + 'G}ZF9&DPg)k$m1}Kr4`0K+Cac`i4yfDANV2i#`##aL@v{v>dQsyDvqekRP%k^1{IHA)E{e(E-;%&5)9jzIP#_M9AspB7}A6BFM`Q' + '^P#w2SU*|2+%zM>lcPBsR4ti8IE!eK;6htwb1`2wcShDt(SVvxrlL{U' + ';G%_vhX3%2l$HmHU|5JV(4QC3^b#6fTvC(Ca`tS1PyjW8q>v-!KI>7zg>YwExi~D{X$4M5k' + 'JLt4_UN1c{M_=N)|A?aG-qIPW&=I_u+{(>m>WcG-t*d(f!f7Z$&aDG@C%PWpiQBM1QrH(fWth{ysNaDGp>EE)Ayu3enw5m+p3R?&B-xMwbuZ17pJ' + ')V%m|KHsd-0{~W2`~L6>%qiHtXkJ0?Cu~Wp!VYJ;7&T%JtO+*$c=}GXYr}@bW)<$CtICWQ?EXnl|Ki~cwBzAoHy*oW2;SaHf$vx7' + '16ejT+QiBo(mzwrUVwcy0u$$sO$cpnbhLB@E-cH*^usX6Wwn~k+e0SGp%I&8Eo*6pPl@e{r%|2UMA(rI6FH!Kg(Xff&c&h' + ')holY?(;H6J!UBwmvo2q*)3w3P9u=X*x<2#f;tE|M_P%@I4H<6!#%Ew#;wzVYXCJ`2(rL%jXh@Ch@qvmR$RKeFhg~U{Ss{Ku3%t&' + 'T;%dkz1HJa0@dUSGzeEmi-|399+Lbq1&@4XEFuY?_h>MGiS8t`Z)>LxVy6`LOX(-4DBtVCZd%>sBV2CEC`^KOoK8(a_oqbcc;)wo4' + 'i5wyv8wt_(Aiip7UOHDCfn`}TfteM}s$s8{1avJv@LxqJ=~+V3zm{f}o(7Ayq%BCNhAPH!83mIP-MuhFoj|YCxne2S=Pg;Aj2T_xjW*DjA853Zf6IHn@!U$aaYO^BQz@$`t$VY)zR}atpqK#_wC7>W3_hQoV)aZw4Cl' + 'hYf}jrT#>$N7p11NXV6;3X60N3mCzKi$A`z3ILQ$&HvCxHHNnD#2U4lBrj@zoI6Yac^>e?IUk>14OWs4UxEVK+fmO6t1#W-N@&*Y~' + 'uVF*a@*}qvI1CNHbMSYyV%*cZzs;jv;e13aH;)nI^t%{t=EEdvW@CzW>teHt;$~w&s96GcEM~Tjjcm`=nB$%(dG^PD53*uk*8n_{' + 'Aq(aqihARbbIt?N&rz|TgHQOR|0?Axk)EuUoobcb+B>ox5$|zVF-Mbp@A|HSO{m`kIS=^TvX3vpCMW{K^vx-)jc!opaHPFnZf4_(' + '?XjUv3zRPIckqBb?BPfXqO&5Hju?^F>qA{U9i$`D4m>9k&Vq4vN=hBD|9zJk=zzj#*u)_5e8I0TJ3CRth;_ssA2d(eLZNCO6?iT6' + 'sjf9^f$3}e`=q1d!1yNA;hD$O_Bq;~z7Gjk?dx~)d^rAih1kI-+=xnANYFtf2dKzqvp}!;SjKZmp`(eETK$gbnAacV=gOWm(FUpm' + 'N2WMoBHz~2nTi9^5g>Q~4r8-ICk|!k1k4+3p|e|@diabMycpx=NuQ%mER2*6#bv8;@pvUN;nIWUOt~bAGv|bo-ddvD#dqB}I^Aji' + 'wN1E#{m!{2at8^suWgKv>aTltOHaqyN9|MkLvv5T26@a&;NTUFSLsPQ;}ML(dic@&eIK)2dB>*cJ|uzbovLZ|M+g0&HE+PUFX+br' + 'J^`zAfD`KvkObaP%*j53' + 'g>^iczOyVCY7Zt1zi`YKG3xR%@4&eY7`4#J9SOEECR=3wyH&1ck{=+Zx`(Q@3T}e@d>OiC4hLPP-B#CZYEl-EH2W+Y5>I' + 'rA?-f<{+#Qw2+rzyAB9)-#qEr&2B)=!rbac{M%|mVRxp5{Ao21AK*Dk|C_p*){C`j=|8tVp;)^%XZ=}FYknln$$K@~mV2(kuwAry' + 'm_N*7fPZ5rqSPv8GALpptGyR@Ov8_|5zsrt(S^?eYAk0WtEpVL}7sXyZmOUKy5V)o(' + ')17<9V|)=uzmx%=;k39LWnznbU_dgXl)$O;<&EIe!y;#4ntT3;oTB0soAKxLDAMlfzx^p6xBm|rLFom0nV&j2@' + 'xINknICbS' + '=Tt|_HLZ78~b38FP&4rK;~#8(^%ptoNdXIHaQ0p13ElKGU*S=k!XGxMajqu!pMyPzjq-|i632$Ljd|ksP_X8T3u(zwTv!-6Om{tYGq8VR5`(_H9?!sFv' + '7l7Rea8O~-tL1Dux`WSfR2$r=P`KXuPKQFPS>w*W+gw@js5ex6(ESJ-%<5n7Y#OUHBu*uBVv|SpVoXjA+8art`*T45D*F59O=7*q' + 'j0<}wq+p^$xkA3ED){>2ec6X5umMy11dG|ErU?$DEM2a8iW#`t4T=COwKY>;aJTd>62}s?!W`WMN+4eV@ot3Xs>mr3D@K}7_#ho)' + '+mgFIz>LTsKTm|p=;QP5kL^X3ghPuh*`iDg&)z&aSm+n2FGuz0S^zUNd(y^yw+~UQuYN~r%s{El$dley%f$@$ooH8jBM4gc*n?iC' + 'z%moFOz5OmKAyIyS;l!6W$ZNpv@qWLBq8!uv%CQjv=`jnXq?RSU(!bB|99$Uh0skB)$clRCuDZMXIV6BbJe|TxVY-FhOze)iA_^v)3R)fQmmc(rh+M)rS' + 'jruKg`l(lrk$|w#in(5nBcI?~r#Zn$x2M;>5n|S+ceZ8YD|ef^sXPhO7O912x{-|?x3bh9h2O-y^A_l$3?0DcJ1|(Gr#t&V(e&qx' + 'Mq7&QaS0xNI-*|^aQb19DFRtbtQc-qYOg`hPRKTJtfwg?Eu02hpmk!k' + 'wk*v9d%pBj3zGeKFT@h?G-jcgRn$dc>Ji%Bq2B>Vq5YuB5cK(pYFapQ%*w4E#As!2+|g}>{w#7ISwr>PRHMzf)(vH|K0b8Z0A~2u' + 'HbM>0Sp$2e6`uSm(>=C-+UA?5>~Es$7|&hiu;AMDbdBrf)W&>M_s&U8Po3{fwJ+D|N4cWeUpuSkSL52`80kD3J3L}Z`;f?SO)g(WV{x?b)mYCsodFlwm@6t77ETRAe_kp*bIf@M>$et@&FXiWe' + '!ln^f!y(Limw#QOUSDwHK0HQo(ylR23Xjay`7s#D_EJ`7u@#>hp~~TPp~v^}z~4#fkGf9~!sHe89WK8>Z0J*An08)!12h>e_P;+c' + '0m$J5vS%MBFg&_0sSt@ZXV(+?Uiz5hlkVf_ex)22j-S{iC+Z{bk*$WLrGx!PtR9bXOc;Qwl>u0^@2#*jw6on(6y~rKo4W5@ms`AW' + '#0gDoCVktKf2cn@91Ga4G{9fd%;#4vpV=PUMME@py#OPtn8-1%R_C=}OoZ%RT|S`HXT5Tdfie@%2n(!4qV_0!Ag6xWoM(H7dU5B%IHQd2}' + 'odihB*2Z)QyVNJS=|@&D8FHgOPjX=;`!Zit?&9RLh3wHIo6e3j~Zk4oc_{2bwHwD*Db+O4dl@nEgkKe9{5STD-x`^Hk?W`V' + 'rMLdoe^*#h0Np6gr%hi!*A5_*&{E*3m%Y%Rni6htfp(Ct?(i8NvfMWHs^xM+O?~SmHhlBuw5JYS4#8H$+2ajQJ(IsZcsNZ*Z%(D|' + '>Su2jv%9Q@50;1Z2=oPIl)&JlRihkOBy-~$D=CO=LYI-rz(wQ}T0eKba4H`UyWtN9*>pGfgCVfTfZ~v(4B(O;oHYR|Xgjbc!e%um|~kQ`>*H(=lQRa$eNqY&prs&5}0r7C2qr4K^uQ^hJDt' + '9hL5+?^{(;!h~t`u$QU4;1Ad*1ECYg{>k>88}sB>{wZY?TQonD5McK-ZHDrsCQ$W~?=EsH90j?~;)XxQdUfkjbr}x&Wu=>wZb@#o' + 'L#P%JOmctN%vMwS;n13#M~}Q=$C3?94^hAg7_4B>D-qcY9{9RZ_Ti{T5>KpCJ;CgnLFRKhmWv4S-2K3k5JvGW#BBxCZ;f(r!s|*4' + '8@*Rb?g*3HEo&b9L3L@G&rrnjX}(Sk5K|ehASjmQO2x>5Ecsaa`lMh}&Pr#BnvEXHk3TI!(DB3x!BQ0I85rXL)Z8H##F!!aEQqFd3aH3di41HRSn5WQ*#W7l{9H1T+Ju_w^snG7' + 'J(@2N#HiY-0=#S<>ajfG^IJLRXu^!H#pwNNi937+v?q-aiZsI5h74zYf@Z=5bP-A!%W=8SF5Prisdy6hS4U=1QzqES>%*rAxbUnz' + '_`Kp`BNc|txy{Rt_yVLkyC=JF?Z#Kb36P8<+8sd<2H0A`I}gD#O&ERZ5E49&kJE6qWZLaGk=#6td5F#Spc^Io8yeMKgwc%b0oiT#' + '9`;p?B=OH}(%g#ayX$pp(%bd_npD(*Sz9qqhW-X6~DpFt(LMRAQe#<8H2&1?ozc6Dbv@=?R~`+aRPp7-%=X;ze8@g^%&gjYS1G+w4MYpI$AkkFqFz6m1Eh476#Y%f)mOFYYntuV_{^W=9RSLiT)Y|h47G78j*B^U))o5=(pZzv2}' + 'B9uLU`#rMc)UGP}S@s&MmuD23$vz%go{#Tq;2;Ja+yH<7^b;?WzB+6B>8JjI@}j&oXPY3WfBK0Kvv1Up0AZojt!U}M6>C$%-Uxm' + 'karbw58YYi(1cWI(^j4w6Ic~XQ4Tsz1X2lXG4KC`Z#R)cpvv;uM??Tqsg=J|&W`QQ)Q0o4CH1?1YWH)x2qm-Yn<1b>>0_8D(fSK8' + 'cpEJtWUm+-F4r$5--~Dtm$wVzCkmsN3=yPjtG{IB3D@f-Q@M2ZIg$E+Qk$@=L`VM0$`{(i@oNM})eBy7g4T_y42&&AJvY)A_Vnv!Gg`B)qpMYC|9Sf6wTgJN' + 'qGS+_IlvLNn_1n-BWvW^3)y$3+CgsyR-#a})nW%Z?7GNFx$9yXK@F2gnbHSVcipF|>u;#$+EK-IpHh`unGSC+vgR)y*RWq17I__%' + ')!M{CpmA+oE8|2h>=^(-GI?NUO}p_Fok0*zl!4xz#@i)ymR%!44xEjdik*r!t|od!p3x56npmbYKJG>lX_b5Dhqqu1EpPpINH%%S' + 'f;KTVbqUCoOPp#>b}TuUcN$n{18zans9f?71WXEQs`7{(H=>Qa;T8v*mLn;M^ztCjROl' + 'y;v*D7}AZcXWln4BT0_jf9!oJjURSxVe)WJ7qL*S$K2cD=zWn@^Rbj^x_rr+y=A`-1h1|iHy$}VNgpHOhKysIYMSYxE0TccK0e@fkkEBkQK9*cGah9*o%;a*B5cZdW3?%L$pxO65D)n=Q2N4y4WdHxf*>ekwr' + 'k7Xpf^n$w4v>aQbU%GA{?PD-+flVW(>DUT(x4YSz1f;vC3=JEJGKGaMOU3@-#4))pus$#aLoHKuOva7!YJ{@e%c0OdOZwM+&L@40' + 'q<3&9y-irmhl)U1RnJ<_1K@hBwmaFAkt4ghU>@jx+1d4i5M}+{hdW>bt^X-pFpfFOf{0{_^qdH)yi=)?-p(vJ9u23Zon;km6;HbA' + '@XADdAeOVj^TpUM4C^xC%R%(!Sh)k#6**~jh1b39c0T!Q;tt@+pmwtMUt5-G!B^B$jqyY=~E%_zRy' + 'AFGvJ5UEpCIv#@?o3hnfH_6;PO^}G!8}FJBi&XB' + 'XR#Mxn`W6*tfc#es>Vs`rz|7)0{HDu<0l)zdh#%{gaXX(J369ODM(Aa>ZYL@O_l;qt}ks4xsL+}CSC&*4-zGhFIxzXt&r@p+{VnY' + '(W>J1SYeJFov~5Md+yRbV~U?C$^I}*h~mheu?9%ILhATXJ0`kpSG5cCeZ=fX)bZli%jNr6GN0KttEQ3q1m;{n*lETW' + 'wa)!>`nQ&&3m11icEjS)Bjvj9Sobm1Zwy>ijRYF;{Hie!<==?upx&F(n1`^R2JPFLm0YBpYq*6vHKR!g6^n-!Lzi|$b*yGYLHr7dRVp<2!A1sMB0' + 'PO;z3OiC0$Lka+Qyj0Oxm+kQ^54!3}*TAI0Ebz&tNfvlTfDW%VqxXXEuTpW>e4?Z-e5D%KMBX*k7LXZ`dWF(`i8)K|MOh}kt$Jy}HUXF!%zaQ15Rj~`MVs~W4@Z#&Q{HOQHh-qC9MN%X=nhPzeY7^013RvCoF=aBao{K8RdpRT+A@DBDTqlmW0$nBN}fD^' + '`qiUHU+x^3>L3>nBTbF8O9?xTQ-@8N<)PmF;K>2K?ffa)xbqmw`NB6os|l#H=h9TY==b|QFWdfdCnzc^XzR8^eh(ckZ(&Jq7Se8m' + 'zvDsFD0`nNo(<89W~7Ycq$QXRBtp78oI-N{O&9NSQRKWL3pOsgn-Lr!NsQHcliD%BY3H4(91ua>5FLDRRT((8l~w>ESzM' + 'P(_ZwD4tnG2}}2ug01qZGk;NBORA3}z7WE$9qRMS_>r8f8FA@VquBpQ(^~5fJly2Xk6QM!K2^I@RB;>JBbJfCf}R|!>5WXITqF>q1yjPtt>gWZQlkI6V=cE~WGuR#' + 'xB?V4$j!W0s$qGS9=m$&1OP1I%<$Avuyf-EfRwKprk%_!pR<8MjX9u4C8m@IPAyv>70TB%T}P531Il?IC`G{HUf@F@82RkM#?rf0' + 'M}ft#pet8wqITh6p{&dFZ5A`pwpk24d(x;x#asHLaQZ8)X^_)uZqIX^g6zk1cNHOyR#PQU)K0=8hm|$3f{i3D-XdWGQ_6v?R^D$n' + 'lgadhs^w1K_Ypv+0r^8SrGf*i>FE7T9A4Ee&<#J!tjUnwY8Oa|tkf9Q>GyBnzBxHNdQqL8y*YXIx1;LYHz&u>&Q3eo$K3*aZg~dT' + 'M&#s)&uz0jU5VIwwQWAErv7{+&E%cVDx!^c!6x`8t#w+A4EWnu*Gf*#r|SvIz$|Y' + '(XFFC_KtkeMVuIXg`9%^_;i=%I!egEpL`sUf$^Y5x3UcP?u=7;JSp1gef>gCZ%' + 'aGdK)`_McMnz=S-2bO=ePc9#%BEXcGi5mabfJ~icirv5fD;--SkGu*Mf{#hWT$Nxv49zZvW2%2|7Ruu~mZpSN;!I^7J4L#<)i9RY' + 'tw0Ukyzc*>wikm03cJRq|CX%=)qVMAY%0=TKHARi@9HzRekC9)Et#ZtIJN{Rrl^khPYsg*Tat{gpfA!zA|g`HzR_&6c)wWQF4U`P' + 'c_W$zv@Ttjm=>J<@pK|{jtlFxrdJZ|rxDxp`i|wcf-ALYYc3+OY2+GFll@WTeTt1' + 'yg-w$a%a~Di%#vP<@zVyj@d}~2$)@ozFUl?1goXUawX*TmI0v^9iv!zG?Yb6pWEg5aF(Q)LyfiI;vP**zL+4bcy' + 'v6_UajBKG#dUP1EkQ0_IN25(6?{j*%mm*r?tk;?TmCAe8aV@>Hf{W!8c82J?y$`*7xI`KW-d+(NAOwxQn=BH3VK0E3#q!Yp(L30G' + '-(sc2`_#P1*|1YZZaMuvgu$8Spu9QX=0c=W_?%Lo8$%I3Vn>SU)JM_zm%14^nh{AO~*rT=|yi6|GT(LTqN#y<@Len' + '4{h4UFFJo193YVTujwZn4nwEH790&B0R^E&bP5DXLaX9B&QTH}n}Ut(p%N7u_GL0D2vdC5dU' + '0(sNURPi2u5*pls4TgTd@n*!gJ=K{qrPn?&em^QuSsaNS2Ho$GZM@qsF++*`e5N!t*Gl3l{lZxBIt@92Eq_+cIS+~VY%UVT@`^5M' + 'OdgGS**3`)I' + 'RfW?(TXsrTlW8Ut%U}u?_6y$*CvOo<)>MieRz$7h{-`gy~i#gmc0*iy434a%2g|dY2tA(zx!2SkH>jT@Oy2jg`@JeV}X_x?R+zJtJ' + 'nCzx?9UbT4y3EuhWO9MB?b{CkY+ctVZ2aly{>Zy#ke4N05346Vl)Q2-Quz6cyBE2vkkm5&(~Iq3`NoFIicd&e?6$nIGMye+oL%y)' + 'PVJ>6O|PQkrE@57aY_DSXO-AE(gje#a5&`@%h+n;7Q;^*r+!dzQsLuDXG=UYE{0HU0b1?hs@mCfB%CV&I0=s8r71BOMY6t*TBE2!' + 'mfY*^2hN5YL>b=+J=EV_JihGHqUO;TrHsTpd~)fZQs}S$7wJFG|MrI+;xvVah(V(B%v4pg4<4Fb*DaU(N?V' + 'E*xyX_#%R?VrHwp=xqaB(hRaw^X!zE;|F>XSF|DfxLGI8bjQeZ!ne>V0$FUxpy>o2gp(j=*jV3`f@Ty02OflpB=4zBrgc*<+Il3D' + 'H_(?w1JsMv-*&6h6JGj8#RL=r2niAu2is&(>*)J2Y4a;7*lr)7~_6Kb;vjJ@I)*SX|m(|6<(Csq%!EUu+qMwrR' + '@;Zz*bZ=mD=o7S0(bGNVS__p(*p*jz+Q!{MThJWtI*`anDxoU{!q7Y-qGkp$3LleL3oy@iy7' + 'pJcurGx$QcL8Jggu`q&trp8bvVZgGKI~i7=avMgs24D5cPFi{6?sYbnS4>9I-RQ-ZkLj$GF~q_!>DJ}L*#og54pF1ug?`b%h(4vD' + 'R;`!`SfP)5ipu4lXg{EuY@9#*njim7n8Zt``&Zx;7p43%sP>d%4nH+8!6AFFUifCT0rPLJs|v;6OuMpYI63E;tWwuM*P^kPjVX@@f0)t=^fzwaP7mAC?6wE^oH+ZQrV#)^#uAK+_RT7s(7D=+b7tH$m5pfF*A3*EE;MCLKKfNq>AiP{th0A2iCNU$s' + 'uBM~{+$P)4Wo1EGB}OVRN(XEhz9RM8>;A2$w^BwiHEwPUK_2RAowjDm%TQuDW5I2@jxkicYil!Q*' + '(sH(ZHyzb8CBb5%iA9*)V^RJL?O4yfD-JtKv7Do@{jcNxX3>JZCVmx#s%{n4z1?Do)AXVpTnu{hu6g;#<-vXw_lCMB(Z31EX$Q{f' + 'uoQIXOSQmft0~Rpnuu{MK0^#BiHE{ntQmdR)T?XG=&C1p-pdXTtl`zj1P?*wIH&D33^fYf`' + '&+cvK9wP;$qGSAjGP3ljrV}elrrUiNIeY6IMgErKT8#9R5gHi3VWL~D2o4j~rDAvb`JQfrA9W0j_L_2vBqm@!>LUDkdhd9ger0V%' + 'WjXO3_M$;p*D!dp>wa{{2q4R}iNnyGO@UrNdYlUy2Mc=?k4s`=QDc4S!Xa$L?P{M(s5^(Px~It12fMgv-P0)vp>-BRsGs6|?4Ivh' + 'XDCDefVVK0@%ra!5`)r}tP5eEiC-M-+_isv!^ge$(Pz&-%6I7;oUMnEc-^WbY~L3j));CS*JR<{(%vEMax!U!zGJ#f9zMG5aNuIG' + 'nPaM~C@flTGo{nf-N`t|kO+JSfC=QKp5p{ag_H*I$Gt8FA60r)tX17vl&U0MvMg4nGn&_sx?8;KcAL2n0xaF5`Aqo)wF' + 'FvPBN^1}A&PKP8L#r2^ZP^*@K1B9a_my`8v4O?6pb4Ll%!9y8^1m^0`oE0C5dNTf7)66;U&<6vv`i|n46ogb_sRfcg$>X`4-34>)' + 'q(t)ef-nX>a<^aq_LmkKUi!oacoXj;la_uSWH5|on=yvtUF6|5tIaRH_TqzX9|re*^X4><#QYt_AFpb@dn*29bV};L%SPG%6ucISTkJTA>1-Y-U)aq91wr%ch>m5*8M3' + 'sH-A$1vBdwKZ_Y0M@|HRN!ZZNdrih^2k|64Lmt#HL&F3z>=l5' + 'z-i|rGk4)(vh}*>htBfHjkuXgHt+h}v`fI5OlNc}WG4`A6hd&M<-ifUek%$?9b{cu_KzNQWW1{qMmgtbc&PwhJGp~E>YHD-)bMUB' + '7-Y%X1ol-JgDiG1Xl^YSWZ_w$ddpLV1?VXUz;2uD!;0MIGGRP^PZ?RapbWB9{AA?TjX}0cQ2LlC}Eao0usgWA5LbQ_S(y}4K#J&sx~XvGE)wfnL>B+dDIc#e~ktR1PMBnX58AZEUlIS`-YO' + '5h>w{+c~=>NFK*6@xnMT>#Y0Z(znN1*r!1nXeA7;q66njdoFW#oKo!O4mkA?r@z8T%33{H)lDnhM|>9O' + 'Lrr_*Es}b*sPd~KnJlVgFk6fV?zF6g}axt%WcGR=$yuO~6' + 'm-^3gp?~{py_oCY>%UFQRq?I)b5*TJEB)7eJzL%;d7aFc`VmZZbj`E4EbC&l(lg+{vZ}ngUh(_UVmgIKIH!JoInoV}Rz;P=vU~g+' + 'bl6Kyi@&Xl`KaIzm-z}Ou0MDKfAWjf?GiwtpC8O`d&vkO)k_vPMKxXIV{W!ymt)-sfB#N>jZ@dp&x(1ys1!W=&~Z`gSDl>%{vB$_' + 'pA@4-HSYNj>-BWS4`1egzvb>uib+x7I(j?Zot>TR`PnaT4zuHfSBIyAWL%C`=QM}@7{}`&urHEfa!0-Jn+dF9wg3GfN!4#@PrSXf' + 'uP;OI>DjkW_1E7%ZNB~fJN@?i@0_>goZrz;{|)s?Km9i}8T#qJ`GH^j;5=We7wg8G%f(`f>w+Kf{C+>Cwq19E;~7V_WI=cY4+yuBzt~%`t0QB&DrY{{s`Ld' + 'XV7>jJ3Bl+g{MCsoc^3W`}ythpRyl+IitDle{Xtx3fn#|CP{{846(d(lTX($aX{8?vik@8IiOXSlccQ6dA-U3U$|9|aHPxc@n2QI' + 'DW4~IfGC|0#cg*W`o%SN2@>$8m*AfW2kOWIAUaZKTH?W3t(ViH+Z^&dG>|jKMZ9l3u+1)U5!9d?|56-}whsDLF@qz&tcpqbsgv?o' + '>E?cmn{rfaGc)=N`bmx#^CH2eHvtG=0odmNL^@8r1W`8tN@~9&9w6=4<0))e7F7-VQOsA_tXSpae3dg5(tj!PL1O|6UGT|y!av4E' + 'J*vtjEHW$SbC|3%1=+U0Tr4~3V!c|fS9RJW;C7o0R>USUg43_9D6>#IR2b7Spe6N`1-)551J' + 'V@I%>Zl7cr_;|J8C73gJ2@4;?UaAF-uGjMqVuihAmVe@9)4`z?Yl-EY+t{P>OIb|ENwxU6JIntH^y1l@w+Z){d@SKKtXI(8N?{)2' + 'B%Z;^dOBs$0FJ@KC^RP1d;#Cw^1!yEu9|q5yWGdMmPBlaBMq}9E@9zrm#6c' + 'jb@ALHiMPG#tj*zECQl7gbv)T+rm9nNC+C3=VBT`Tp=2Ql{-a=CxXDeBqpi|@QMjO;LNCz0Lh;I?t5>-0AF(&euiaf0gW)S' + 'M@)jr0(Zj1!(x1s+`~dt_bo*J9nA)YhV1B}b70^l~(N@8E_seIdCFM(?o9SOH{1%it{s5y(VIq>cm$kO{`' + 'K%%-s0%qr3pU|M&^*(bEI&>(+=DTB11Bd#n_t(WI{>_ql9#r-3;Vn&&nLPbZp0|Gn4iR8rBcd>ac@@i?h(SHv2DD^_1V<<8hR@JT8DOxJXc0z|cxBau~P<$E{%T@BH;+BAhW`+WPX*ozPT=C4n4+U4-{ckD+' + 'E(@vDqXTCKEXrAA~@bhIK' + 'Cw94>ObUQGK(xmz7{VJxenHBN5*8@cHd=E&gzr;T=U1CCn#xAfpJXnJo<>9XcQr7Ged(+2meQi04fA&c!&b{' + '3vU>{hUvZ0O${N#Kl(MSxBP@MN7_eFrwug;n?*mV^GTuTS@_F764H!6>8K%`71Q*7ixvsgM(Dj>W3`@VvwU7ofM$GwrIJMgXDx4$+$bTmot_(M?D@B3Ty`cTMvU3$IwdVpQJ5#aWSU' + 'GsgtLyWcTqM;y=~X}I?r@aI)o*T`DLF~z-2F-ShIT}+M(^K#C3J96Z$q#ajOko0M;>x-^Xz8~>7x2NUix2(q%a)>0FF{k&uz=9Dzx@=R{`y=)B^42_dD)BD+1mab;n?f' + 'K>nvvJYtyvzV(i!B2Bm&2km{*sX8E?v-@c1Go(p`C~FOW32>LAg@myt#97Xxf603Ws#ZDH-Br>rrPKB?81GDN6Z#wPsD~T+IUvW~eGKTj(F4U%MRXm%>?RFBw+=w9' + '=YZ_dFf{Xd5d1t5nLHiMT{h69*@Xf9NcBn2decOWtww10OnrLJgq0#M$Q()z>**>hKEb)K!{#xH-|>9GAFm70#48jI-IxH@!$UlR' + 'hp?)-V^@yb7OcXp%*vhDrv$j30Ai124n(i#BxTUkX2iye9JrhW8QMJ@^M;+Z0tQ%YNm&VoV@8jNM_%B@Dtqbfci^Vs(yRTIZiuRvJ>3)nZz=)6S{8' + 'u}d9`R;)V4F$jwr*lD4rYTi-^q>T!q#)!%eXYm^-LC+4V9D}4m4jcV(C!$T%lq1@NR9b154{D(e^F}})b;eLCshu%@kvbi@HWr3J' + 'KOc`fW<+u-mIXsD4P$^E9xHIfGzC(7srNku=O}X=N?55F0d%oqI=vj1DHWr-bsAVhvSbJ@!1yg_biI1Xz$#&QMPFVPTuFdWLPx!{iUcWY0=bxH9vK!pzjfs7d$n@M4%W;l-wsCa}<;VxVQ0?yZd*@dAB(Evbdb+)J0m!L`Dj&C8gtUU{iol4_q*@^*AM@eUyk5VrYhD(w=3tW%kJlb>-6V?U7Tb9!I(DC' + 'xl?G)(HCZ`x;-@u%iKwAqwoeP2nq@!-N#6NgHBHG_QoLQRDtgw>PcbUP{3ic_NrD@jvbCsYS&1uZofEh^}}?7y`vnvM9n;Y#Ez=9' + '9X$_hW^)31hVJD=suEt`Q&{3-Ip' + 'n;Y_e?FiWm9`f$Vhp>x7#Z+7aNl^{On)#(ekUiq)P!e;d%`#CjAAxFFCz;p={#4)F>*!z3y}mEHHe3T>m6WsDnmkSkhN{uMy+p+^' + 'ORPp3`cg}sMmf#2F+MZhWP*0Gl}2xS1{hkdvJA-7bka*C*J!9k19Jh<3qfIK;wowZ5zg-~e!c1uuBq$X+RG)eRKT&a04}soQBNGA' + '*5@1=rKl7=ZxetHQb=M04qW5$kH|' + 'unjDceXy_ZT>{_J(dh1bWa?GF<%Ho5P;%YBbv!7`>O' + 'Z1mY^>1h+Mq86VoD8qQu-VB70p(q50?WB>&lxc3g##XBFD^d{ouZ4azXByLIA}epe?>9q9' + 'lpeU|*489jZvm8mZ+`Pwy0Q4mcv}Vy9Pf0cAXO>th?JbtiIdRIDAR-Z6826h4NA9SFRcn{YFw4}H9~PsY&+4$^+aeT' + 'hv-K~Xd35Y8W-nGpf0+e1hCWGlQuq@_L4?8e2lSf!&J9olzcaL8&jwNhR;&FjwXydCiZ%Zy)pV61JvH*CVRDoy06^w;u|Zt<#@{5ELZMZeI4`>UrVnmM?~sQE(ly=ivIKe*4(6yaTbh=@;>dhV${J$D>lfyV=|SmOVSr4GF9!~daq' + '-4Q3(SAjXd3pO}kF6t8L*9C32{>$!lW1Oxv0hrcm=;RxR;OIC`)I}+3G_i>*v@2ZqfKBB1e;CTHLdKIL?2AdcLMO*lsc?4aB&F*3' + '#It8c6yeELaZ_T(ieZD+O8S!QCi@q-2KkIuX=U{30`DUTFqvO>%`nb!TTCW(vC^?@xC;g5NvKfvZaqughLMH4LKm+)ivVlXIcadr' + 'olQ{ImO#NMPghlu&;C;3+>EX7M>gs0auNkxQoWv$*>AbNoR&2Pg%RDg_CT0OUe29Z6=C{wEP4)y>xrNY|7w|j)1E~~WYV!1gxM_k' + 'SPzqC7O9>H#4hUu4C!(Lc@t72Ujnf(rpI?#lI&ctFEPPF7QpPFW9!H3-pMt9*PqZVKZzFoH%nl+3)|;=-VnSsr?GFffXa(wKc1ml%17fAFo$9wTD5_`sJm-5x}UZ#>nM' + 'e^`1kGZf8dn2KH6wP$vh<*c5^bZjv0aIY5-!mLEWTQd~f>TO9A<&m~FpL%kzL1FB>fzECe*OuX3?p|?}B7V~v!Dv7DDmg$tI|ihd' + 'Mt{nr1?mOz2ui?{C-dTC@{uSkp8+Hxc!F}#mu4x`TSH^VWKITsrK;9TFANaxRrQlMrX@n77Tp1E3F1sK&Bcry+&X9w(`QZ5r-S!=' + '0I3q#+;WOnH9%r1SIJM$`ra0Z;fI`' + 'eo*{~GtlT9{^Vz)Og6+v1!t+Ip-hJJwk2BUVkb7hjFSesx28l?dP' + 'ye|ht2IS~!S;6*rv56GaBZ;#LIrKj2J{-7tBdDH4tsIsyA)oFl5hM5j)>p06F+?-L6)~&k3Q-l*7}NE%7yaDkcI~o&3{&jvlf0Xe' + 'h^xPEiQJ6gpMME&*tmq(a=7ts5e(gF={XWU+l>t25dna0Tj$;dmOZcA=B?b|ZLUpAa3WZ@oOtZYm2b#!nuJXm5!m;OPb;);@oVC2' + 'P9q@V{tOxz5y4Y^shDM3uz8&z?)oN' + '5v9uJg9@1KS^_76s+BZ?+diWyNcj`g9P;W`YviHO=SkY9{J{79)ohu{i2`ZS0!Q>~{8B>O4B1!roj5Spc4nT3bd49VxI=m$L99u7' + 'QcPHW0||=wu8_3qa3}~V40mD*uwAeY&o>lb!*9-a_b<-(E#ohYaMa%SVdP04fq3EbZlpz)E+f=B)k)MjqZ3Ca@ar#Gdde5-cakjZ=16$k(Su6aC2Yib#)m+K' + 'ohzwC2ZtIXN96FukKJ8Zs4mBp^n~)dmE&Hxk-1Z0nwog$@CS;K3?iep^LM8eUQ8P#PPBKU4PO)+NbtoQ^DQmj9+kzjy|olI!Y%Kc' + 'kMCnbrlm_ky8Td^9itQ;z3y%@*#*(2jZ%9HX(8J9x*MH|qeFi)Y)EQhQ(TMrZz}_yVFv!PT_w&4b5S2vP+e0xrSd3wSEeJ5r8G_7' + 'u!lj9ntZ%2jen1<)}Spf;63)np*0JgT3mSiwiE(|e>Jobi3wNGa18~OnraAr=d#H5^#>fDLP2s{WQ~4I_GsCABbs!slW?h4I@A9X' + '{lHDt-!O*q=YQ~>F^`X@vL!g~^gi8NU*8=Mc+VM0$;U8aL*}F;u;md8iY6|0iub!*MnSODSa@4w+3tbFS;S;L=lFdSg;z`=aMAFK' + 'CU)x{iVZAQC|1||Qq)|2S8=4;y%YY58>dJOh{7~iKmQP`iuw-*M(cQ)oKTs%H9T)?*=T30bNbKX|W?VFeJi#1V_=GF-zGTpb=#{ULiJz{~Xv1D8(m0tb#cW%xM=WxN?5mv`JPBahZ=uITbFg' + '^Emf<+{OOcZuW@R|2)-Tr@BA-k67bIp`e!29Pq0;`l5BDdqy$sv>7B%JsXzFzBkyIde8gLIQa(dg1lzV>IQ+edF`7%pg~Ob4CNWM' + 'JEmd>wv(2r?FLy{Lfs0>Az8UG37z*erY8qy-Bcnoza8`$97`mq_txFN3x5#M@w-9Iy)g$zl@E9R&sXXhI^ng(^?PN9Hn80?&2bBC' + 'SW@8coZh$kTH6eQ%a9%`dT4F}dVUCZT&G-%+kw1j9k!_l@ib;+m$M2TRWDommmye~gxqV-lzZSW^6_-)Uc}^5v`=PYh=i$b8wwII' + '#z~NbOZEC-oWBq3k0yIQo7#kL4L0%^L~GK^FqXN{?8dutkSs_xmC4^r' + 'L!pyzMm|JIT#M@eMda@n({;_E5L&5bD-WvHc|$hO6CsqxWbrmV*SV_{*IM;2J+s)MEjX$(p57E8$gA8Lw9s$~=y2${0{nIX<#DhA' + 'Bq+R|JMwa-hO0{=@Ts~a&)Nc0MUf6;5&Rt|nqFEEsh_4$Z5KDRzedB@htVi^W0R0gx?;}OgGpf5pGs?l!UWa6p}-=_cf5t@lnzQi' + '>QU5u*CZp2oyB0D&e-Nrotj(;fqu$cYHD_Is*<-Vqn{G(gjF<{ym#B&obmMaZsZ(B^@LvuwV|@)04vib&4@)0i04{xutsX_9jSVd%79__`_Bi76-fWW}kiO~iZgX#}jVyg|u(bX|--sGEbQ' + 'jRRjRcZkp&uS=|Fh`iAm<+2#%sAA?UIObDa+wCsKY$XG_j5T=gZP5Ef-J-2c4batfwOC&vQ%06c)d{`I;l!PpPvbV!{p2}Frqz0+' + '1X|iYIt_G52Sm6<#p*8AM`S{cNrj2;1A_fuQ-m1H+&BhcSB|8U' + 'uO=X(@^hL!kN0u0o{j;63Yl?F$A;97uJieo=9P`2YZ?<`#c9=?R7#$flB}wsNH2k3#FOf6gUp*p-KLlJDd~CZES*~&vYqW5M5#or' + 'h{Z%;oN*YXeo)-zDG}22rKqn;eEwM}P*zkDXW^>Km)8Iq!f{>No?Msn+toGwQllOT?c#bR|Td04C_(9VOe' + 'IW4Qjs#&KdR8~u+9qt9qG0{dbXolK831kj$%)~G?_2K;&OtvkukA*XBx2D)`%Q|JiOc|OVksr#!OC1oC<+(Kh`4VfhS9D?Ao^UKd' + 'D#F5^>^|v-&GxdPhz@z8fk^e=hJfjp`-E{xi?+44FJEl' + 'MZfEH1^WEorerk$foz86fcc!+3g9)p`L_SeJoDmNY~wyCZlm=|o!XAHc9`GaoE-f0>Hrs2lvi^KYKL8Wef+;oi)SZ?2WN-L*};!5' + '50j%8$?@y6Deh#yV`RKA$FAG8T5Vh>fq#;S!?QQB1!~wxGY*nz>BrXkZfixp+#g@Q{xRMQdNx{m5c`#MczGh*gxmJT>yyKy' + 'pN?hQZgO(?;t;mu*&(BbdH(tsHtgjgfcM$K>9d39hY#EUt=0cZ8(+4}1NP_&hxROr`i)bR3c1!sNwt+WMF>v}fFbALzo' + '5f&dQ_MhU|79WWWH)u*HRg^=sl*s}oSa*Rf1KR{WWjgJ2?=>Y&y_1+~x-G^R94|23!OJt)K*kQDI|t97C(mBLeEaG+cFa?GYyCMK' + 'kITFUMxuy{mXv|~Ucq0~$B4yvHu8>>;-;u>F_BDFP<5=n=V&Kppm>!P)7n*)&+B~T`l-H34(D9odsbd4txz;HW|!$L`pSMh{ORa8IePW#@c9u!' + '0;Jo;N70uuE%5EkWMMBJ(+mimgGj`v)^S|~EK*e392RyABo8tYpV^g>cZ7}#T9*jsgz01{2Ds1B4L2?S4i8WN&r5U+5an>J#G$k$' + 'kH{cm5R-3L1x)Auef1tJrrRiezM&ccufm$fNZ&WI&H2~fAYWhjwF{Z059R$rn-hIwNLJLj5sAp6?cX(5tK+Fo(1^{Q3zj8*^=>~(' + 'p9fW6Qv#|&M^%GuRnCgVdNurZkM3&r%Wwtwdx!40hydu6JK#8{C;jd6cO^;t3MZg+e7' + 'esp|#cyfk}>$Mb}dM>W?d{W*`!gTwOJ^b(4>w}kvr_T;MuMXZd=p|m>Prfk}9>(-=B0Pm|C=ZUz!xqj2rtrztlHPr>2=>O+9R)5O' + '-^kHBdVqJ4bjEIA>N`K3ar+iFwoz*he?EBm7Noci?}`1dn>GU}8!}$6b!qy7Qn%M~H}Js|*WA82@hzupFEsfVmZB34@vcU6fy2OT' + '1<-OlRKka1cgk%GyvBU9p0_z`+B2;wR}SYX(J=^a8Asq@&q%6Fa*=sKa|o5o1nu{iYfNS$oTZ9={Muab*JFQ31Rucr{%5fWj>)9(' + '$odnNGH7}}+w&~Vcce=MhRR%L2hx>E@zP*R63`*p+0+k}|5oNZr%EJoRK*CCayL{g@b@BrqgmFPZ}Sn^nk7P4F2' + '(oOOg)`kuCjWZm#x%W!Q#0Vyrv*vaE8@It4b0r>`3UYit~LGEQ)' + '^?yqqMrrii;~Bp>q*~H_Oi#^ryVxjbG>4a+{$|)XesPpMvESRvPqeo-zm<^1j-EHFAou;$r%Cdm9O|KvOgOP4{l|Jm!e$7h`OWle7^5?;99?8zvFBYUOq~ymSA)52fD|Q%wUss)ed37L_#c&Y5QXLe{1%' + '?K)O9o}$17&^M~Lf7x>WWy=XBfm(cg0L^lZ;oH&ajiiX&mcI`D7BYgUUV=W;y^C&3xlJ!G@0O&Ef-oJCrnO|VoM7lqw11G<)bC|W' + 'wH^4v3w7cPFpqa)`hJi#y30aHo!w3+=zF#<#<2Hu)t;njsz3tR4WuW%m|JN' + 'A!=J~s`xL;IT>wHs;(B}^{Biorx=xP4gd$QBfp&Ha1?5~0)(D?`(~uzKpJD(3z=g^f#b#M1(w9&MTmfF' + '6}RY}FbsHgqL}-6ovo{~)4l&A+9`(B0!{!tP(BNB3iZ-V?LEiWvt1aZ4qY(ox5R}b4ZLXXs=>6|1D=8hpAuuxp;t4W0bv3f#rQ~p' + 'eZ%By8vpCmKl2nKB?7jU8->EoQ})OvA{Tn9gOTqYPmTHrwz6JdVx%*#w!7`UhDmyimbc^u%>2#ESG7P+g{?Z)WRTo(>IdKJ{I*uN' + '$Kf2Q4CoAw4qNDQm3Cbg&!JhmcN`c8q_0m2`F3MH=HghTY16qv#lZ4`_T%z}=syYe0d<8iLn{wq?rqhqnCcyhAZ&1rn9#e45-B&N' + 'CP4$6@=Q}Bb=tFWA@!~&Fk1B33bp~w)J^2YN0FA9y{qCr8&v|aN=0!s)qqxuzh~=r%FAsq%QqHm>*6&k%f_tnO%I5E{GgGQjp69|Rzq8;%8~Wc2MugObhjE;?3;TxWs8w)-Rxy(G^b7(' + 'H*&N^J2tmhWyz=oYIV=)sKx4&eb)T^a%64gZg;}jYgJa*U?#kKKZnMq?64*yEN9%^*)IR$<%%4#1|j}vRut4x;~1vkC*@;}S5$BQ_ZlxsjzC?zUWf|0faMviqPp$4p;%(a!m?V`7bmeVcCoT~wUbp}iFB*JJd7)q4$K#T' + '9ksto{^Rx0vAxBfygfZS{wZ|HBN(k;^@SSiUjhEiBG`|qsk}joHoF}arIpu6!Hd>+e;kT{Ttt0%wz@R5_oY=#^N=E;yWc066GWv@g%k8vWC)uZSNKu2e=QE(gk{;`viO6-hM|yJN' + '@X->@qmHadqvQk~Z%+6k*kmXTnO?ojY7a*$n6k{UCfOL(;ZS4Ex;C}lobySD6lk18l-fm<>LrV%x#&}h1Tp|W%0I$^d5VY}Vzw4&' + 'sFE0SgIRuy;Z~R?ep(jeesWSQ`7WMrQd0?6T~)iE9B={|&?Qlw(PE3~knY}*7rnq}uIZ{=VmecGVGlWhO|r763vGy?dOv)VYf+U~' + '=y%}B6M5iUe~S^j*I_tJ^p6ASVZA&njfib)0p+bx%gyW}hg48f%09dvY1N<5iuDVSW(@8K4WhQfXYE~H4-ESUX{Qbp}u?TI;>`j+b#V)+PsldD4|Mv2J=YOw_UHt{mHr>@2$0pGPgW>dMU8bZ8;SY3KF(cz0w@2&`3pp#!#KF^!%jNH=dW2VoP*' + 'JNAMXaNNswc5>mUNV)%Z9ge3t6xeCD97krQl|v=ntQscYSxFXbcS8jS=$IlmUNgHjjkFlB2YSc|)CDN1r7|($l<9ArD|Zoe$IIFl' + 'T)xSL3y1i6!J&^dH*f9vV5Q7uk*qF+xuko&yIG21g^Y8uFc-e@kI%XBy%Vp)6TFohyOH~F3B}RzLNSRVxFk~JG=5!L=bearXp|NI' + 'TP%&yQ+KoR_vGc#tE01IKTzQ#Bmj|dFj!Ofp9&YVJduy%GgsDJ)%Gjeb~Ri1-@i;4(ec9b=}ypbirB4OW@*>s_5ejn#pe66D{%ol' + '9BbqT+WuDvIYOPlWZ3$jWpxRKHvBh;+I4VjQ|zuT$;pfBK059t=IB&*8!@2`^ZDxnrT@>K&#!?Rp9AMc30Dbsy7CaKJp_zd!emq1' + 'juSuh@JMXua}2~58}{g!>~bzrW|NyQ5TRX%A{mWU*t=Zfq|%K;IYJfnSuTCck?WK7o0BN+xA$fm+n#NKXUMN|r3Ul;YAv@3H{lrC' + 'V0X2IyUpM#FkIN5$lB`qpvdsr7=42c)=N{UOtD_aO#x_=JIwvl4fqczz(roUZ8dsmcBVN?>u5Pj&WU)IP1*' + '#$F$M5D?4+w*!SK5bK^RJaxr5hQm&wPKw8zm(Y!!58Fmia_pnk6yqa1-mEz3Y;_nu(y?6j4WxL{iKy@aI2@?mb)-)_8qIZa>DE&reC<_YBi#&UDy0ZB_>r*doph*y1~ft)xqB|Vl;re' + 'YBq2dS1QMMM~rYTKzBh^?{&U=(a2k#iqnI0%El_kY~_;=Z4A>{&eU`+dkpVg%DH7oH5f0oiJ|l6!C_Y~ZeEvPvW^jF1W>7!?wS^w}!`Z(u6w>+$EQD~;7jZ|xAY@^{' + 'W7M(5Wtz8v_;aK@7tQ)h9{{)0L4YV&t8%)5EeFX(;~OFj(yY9pSC6nn`m^R|m*VvNKINl7HJv=aP+VRGqGC5qg|' + 'NbRCt?Xu<54`$m`w!$y_L*7XhY`s*TC-)l8lO9^ZffKFgn94d^mZJ~T!mTOM>JB`@T2Lz+y*&e(g+;T_oSnRX-(OPs1H3<=Zsr)5' + ';9|MjE$5Ad8=>iJn|=Q!&jAND1!qQ7$AZ{#!rj^8GXR9Tk$gWe0lJJp)gV-;`l?#jT+z&nPZt^SQ9N2yqieFrIaiASM%i@1c|?cl' + 'Xt@R*>SIw}U9W0#1IWrhf#_hJaNjf7H8sTpy(t6`OQmTyJgc4BHyhBseFYb5;1Yql@4D^;>I^tK|YH)6?WH@amKYVkMChLTGOn{Q$bThJ359B2V6griKpr@+e' + 'M2N78#*neN>W%{udUWxAPaq=xyLUmw*TdUBa6@IY5$APMoM4|;Y<5ro1|~&C|Z!*!B+ExDmD8LCxU+*f;$de#' + 'E@x!j<+3XM^?ZpCMbW2|{=Q80z4)gt@qn*%AQk=ohXl|{6{}~s{T{&fP4Zpxgc@w38L#pw9=_%pEETz?3N4NgM+Nn>l8j&`8sa&Z' + '0}Gp^zrVv--P`qj-*pyU)w1$a{);B<<@x1clFS}8Jc0XWERcs)-+1x+3@YII9&<2ysC06TH8VDM{iy8^OEc*ocXb%' + '{BgWi32~nPkbhvFjOA^IB8jj=hEX@@mmzSvw>h3{SSQGxxD2`t8L$AQZ>q&j31rlpA4M(aAAzZU`&98K4#`@D' + 'I(dlPt8aVy`DFO!klnc!yLaTEyGsb3`0sp!g+1cYlfB)}$lWo3Yd`7keEWgOhA2!VU1%WyRz&d+JRVVZ{H6&d{yG>69ib7Ck(zQk' + 'AHIMr9wn`{D(>neYt%4zloT}iu{YGacFKc#QngJS#8O=N?(vSjchn5m<=T1CFy8ItJ0}8nqpsbaYxqHdy3csQlnUc!4|H*Rv0A88' + '=3y7qbbYK9(NST#Sl*jKgft2(9FR*}Adj->`}dNRLdR!_Ry)?6%^|$`2+2JP^|!j{VLy3Gz)=PBrdnkm7aC89V(Yqx' + 'RQNy^s#Hb=H~5w(pUD_pe$*A)?7ZejHcZxbAgW4%q%A(RFSbB>WjJM=tesMTIqdm_E{Lsk@kZ+2UQ%sxAt$oYk~K2Fx{awE=55E<' + 'uCcPQxojL1EZu&(n*&9?fj35rT7k26@<6v-U$DCpjhT$aHGWnN8+(u3IBBZHdhw&n{j?Sc#64JX|T%c>s`' + 'JQL`S<7h|88&xViGs=sS3@F{^o+RJ9>mwY90g9Jnm{gm0CUUgiXzq&p?+-M$=eYvJ^mqZLb(a)8W0>WwFSty5)zOx?A`p18{`CW`|Zg#5A2?aapdNtYoAw^@6' + '?A5%^>rq*zzVFXPF2dP^vAv~KJjlITR}i%?`@tRQN>cNU)8i|Mp{^CSi#d{hSS+hjFw!O6ZmAhE@+I=8C@GZNoD4(L*I4xG%1T5>kqFYx(' + 'Nam07pm&~Hk;FVCS007gA`A4~fBGtMvwD;nCTh`cjxT&=LbTEs?M0y~IYz8Z)#i(MynXx{lNL6LxpMJel!5qG>rfY!%uatkIC-9(' + '9{uNGn;4GL8jXR%qSJr#GB!c5S!uZvHMe?VLPPs&H#8RQMQoTrnJ;;1#O(^b3~l`RFI7q9#uy)WkvWe&my>a?8e_0g;hKUYkF0sw' + '*8(PABd`3LXa6;Edh6yCC0`mY&X#@tB73*FO7?b}6DeWyWgwMhKs~a|L|UXxR90A@p4yI1DSzT$sOGqS9jq6jUk3GT{qicWF7qp-' + 'fK&tBftf9lZu$xq>CrDe=eF~hf#*DGZjbv~LS4UUev@%|6WidX`AzT~yRSooIdX}97|*E*EiQKcA5cpJ1QY-O00;m803iTb-v`t&' + '2LJ#+82|tr0001HVRLkFY;AKdWo=?*WMpY>XLB!fd2nTOE^v9xS#59IHW2>qUvbqw#3~dQ7_fct;GtdDVFkLaXtTk9VaSX{$3$gP' + 'C8@Nom;b&yl6tiqr^Qe-2pWmJH{Lz>LN$t_*X4CC)^N-6l7rnB+`vX@$e3k@Tj0BEo@GMbOrj_njW$~CAWb)=Ej3S55W7NY3rtF7' + 'nH5T!K(khP&ev{b!mig9;d7Sf>^kQOytG_1t8@Z?u%bW(3B2X+OD@;kOCSr@IWq<=3f%Ql0vo~e%oC!~$X?Mrlluf-36h>Cpl-R&' + '6{{$fUT-VHD=v-FqtWPTEq07lpSYY|=@Kn8xw6JR$0N55FDhreltpQ$BfuYY%U_a%WILWgp_S!Af(a)9Qx-)oaK|c`Zc4eHuAE2H' + ';NO*ZD+vEl~kgwUUV;oT+GLFZ(0nGW-Uw4luv4EIjz*{TEj4EVnf=ye?2A6z>+S+#NUY-nf5gbso0EfXpT' + '?1oqE#T$|O3>|goJ@U~<9h&hx?og3aJYg&BlryfSoI#cHMA#%7(hvUuB`V>3!rWc3mC-T0r9T+x4T;E`M;>Sy{VEy' + '^BbnKoFj16hPu4nVu`m~VQ`DAL|25F;PRTS-_bmPodC`I+6@mioQld&FnSvOd0Si)HyC)uGylhR&@s*=N5!MX*U*fP6P|K;TuuLr' + 'Uc8raQws*_0Ff9sZeR$|<(5&&0m*sh$8MT-R8}ICkzm)j-{*cv``U7ciq<)t_$iz;uA!||Ppb>f0@Ea+`^XZ@Bb5^b7ya=p14YKU' + '>%2-7GwwynI{`y|gGN~FykRMUDBUn)i+pENvD=lFhQhQ;?kE92_W{R$!?Y1tn_$LW18W!0us^NsxPGTba~HRrn{(MNsW`h%`w&5r' + 'ar@jIe!9iA1ypy&BbOYBCC(waPwNX|qXm72AVlYjEpB7I3' + 'vax`f)P?y(hCFT_-5n;}ycdzk5M5D?G(xRgYX(I`8_gZuw757GHe%ceZ7e-tt+1Xu>@^98yjrkaeBzzr-O4Yr+^1oy^)EXm{kf7D' + 'LR%*hRt$!sT?RubX>w*v4uo6h0i6bM%5Q=@JfoC)}Ou!o*xv^9zw(P)=P`>=*r!xprWq@Pz-)DsIWD^*dP?30yc?_&~;x' + 'xQqKce=9BzYhx{j`+FH)w&}aeiLhe$4|g>!UAL-M+=cjbpm?U_`B|@gCKSx0p}fU_f`ds%GkG^Yvl_m~ln=YyG?$!TYXlz9hC4Vk' + '2H|-B6RfvNp|D2#ECdjulTsu4;@4tuQLaIs?g{*t>F3Vzr%?5wP(WMj{XGJ5xh8xSg(%^V^m}_U0rIDL5ixJRdn^`WsuqjayAhE' + '?hJa&eGJaiC80yQ?8@p4JxOaS`l8$IFU+pD^+5YlA^;`JVW%>lUs`pkODVDCP-bYhwqX=lt>&v$>m$u-HNH+DPaunNord4l3gITQ' + 'lBYXn-VG5W0v)#^65v>g2hP)h>@6aWAK2mk;8AprG;kqGHh007iO0RR*L003rTb98WQZF4Ve' + 'Z)9a`b1rmmZEWnlX>(i2wI%vpzXGBD@?Nw^!hsvjjAeJQGAw0n%ujTAtqU~?4tA>Rntv4' + 'gO7AmRF*=FF~*^yPni|EF&*{`lfgFaP(&@Bi=;P5gL%Rab3PZJQM?dliBU*(w!<^Qu~tLM$$&+C;Q?SdX#h5|XL-' + 'Y>h3Jq>U43<1bI1{O5nB#iCs=@NO4x+U<3_od4x7Po`Hlm+cBaPK)uSpZ3c%e|@p}ta|$7w%yF;)5Yxe9Y3)A@>!W(UEN(z7w?`A' + 'Upd29u7?l*G+hqh)0$nj(i?-(s9rd!gEK}I5o-@o^x|sl#wTm>Lx5e;!m0xc2t3JD|>Y|kT@NM#R_WCYA|MtfpzT-#k' + '_g}W$Mbjq8WMwcdJLzRLm?~R*Y&w{1f{*U$lbf`7&2ONE^P4x9PoIp|$sfhujma`9UtEqbaeS|L7^MoBbQpXRF=@Dov{sAUT1y%2' + 'rVs=^SS+43{aeG6jPzPK8_bh3yPn@&wcCrHuBMmow%_sTlf~`bd~xw+dOf`RYmVE^^!gg(c{432f)$q9k~JrQ=Wg6n~(' + '=dh*c=@kOZt7*ZMKYxMGp5u4-r$eZ0nzE%?9b%3RMpfV|u36)XQ>jZOT@Ux)@awP7zWU#1zyJE|#}{YjT&#b$30$kFg`fg6TK&3C' + 'U$yz-;RjYxr&M8N^*wfk7VTOH%A}&Zw+3eIyph5?t%U7}I9@1`lQ;NtcC8i>0vf7Q49qm&4PLgZGoP@nozpd%{uwN1|DLY^^B8sGOB{6|qZNm5LHA^0)_QKOv6Bei!th70qCc4!rS3Jl|Q)Cd?$6~WXIB~*6Ut_2ZvzfGuCCBWN(gQl}wv>PFcY4es7q{tp-eGpbRX{BG0=DFS&sbmFFhriIM}y1bk!3Zs-uG&}6t(r(iY}r%rECzQ38ih9HW_6x-rzP_z^kHE' + 'zJojfRHedr9}Fv>u%X>EIJ;c@?U6Yff4zT}{`&Lm_2tj_*7(u)XDAbsi*c=V>$?T7B%4UCM*-@Xvs9{Boy!T~BiWa+xJ;gdVUd{Rw{`TTe-+ujvKYgI~fQoP=aQo7j1udZF2wo^LP7a|Btt`et8E)aQ5Q+=ika5ZrMTel}Jt~WrB&Z>xy`y' + 'Yp}_;QY(Z2c)!{+{gW5J|7Nmj%gvn!;0TS@y<;y>P|J9#ge*cHx^RFXV9V7v(u@0>RD6MrON^77DwJ%tb?B{NbDLzx)2{7aZs9#Sx6M>*yC~n(;P48$o&^WORvc25_|&AgGcENmvY11$^YLwL|>&P^4*p' + 'Ea_=_`?z%2d}4oG=?GeRHhTl1+0t!!y<0-44E7i$?8;&UU?V9?W0lMiFI50B&bkav`^hoCCRbrOx>YGlZkY9|q$*z7Xng^-1)D8g' + '*NQO|=&xk5uF*RX5*u37r+~^wxAiEfL@QFocvXmGFbm@lk)aA|D!r)$j_cvpAQ?XDW^_z3X%j^^ZN~9cGIr8^040!HlKJz' + '{QuWw%CpIo^`&@MNgW8NuM2^}8DAgjrM' + '=zfaDeRNxo+C(kRz=plHx&YEE@CBkEYw1cdDPjCyn(kr|SxbVE*B>uNwwzr!F{qZLc' + 'cQKyrTKNCV_DcJ((5{D1ZY{d!F5N|Y+sfT~|KLWe3-ml)HFohL-7RJ%&FQVTtqcw|KCrf)?k*R2$WQq2#bP?X`f8+~K0(kfZ*FGO' + '>&3-Gr-BD>|LltH`g5DUdcBxmjI%iGXI!Mq%h}J3fn#xd*PhH5Y0)m$OF+wYb9M2$Uo2)99leCMonQO}ErdB+?nzuAnjpYX`|=Cl' + '(>i`LS-ftOaanyb>z%9CmlN9FWN2ejr-dCzx9`p;FIVjiefrHL-P~Nh8y-+OL@+Wq=(J3BcV+UXy_;O6o5|0w;kuI>Ql;CUrq{0~' + 'n5g;ecJfn#Xqr#bb;ZZ?JBog`TNxu=lDQvz}$+NUp{qyKKaA&`Fn5rvu5n@6f+L{_cnQ' + 'WHU$Slm7<|*X29jw#$|z)9%L+DVO+Qom?X_<9TU5yvb?=pRL{RcCy;7cs4v~dcAax^U3d~u!q%H' + 'U%#7NbE^avafx32jPA6_YI9?fCfDuf&EqFi4gvQ$Zik1Cz1^6?gnoHGY$0(PhF!Ct+S-f)xNe=5yLL@@A%={_KM;E-RiI7=a+w5c55-ay4?Kr`GY6i' + 'iyUj~nQwmqFLeR%OzUQrzZhP2^H;zze~+Mzm~cUx{!yQF#(%(*pNOyVh;KSo5{?4dgJ_gdq%_0K&*Y_aq1So&l{$xG61DfM;v*C+tC>@_ziyW|11kMJ!CRjS+J6`M4E6M?o&EA~t3M6g>JNom{W#p>A;b5UQ`|OOaTLxQTU#RVQvcoXKVN+(' + '-B!O=5}DNUyFY45{@E>T;UH-150uN}lESD>@!K70nXkBp>?N?R1@e95zPsLx!OH(9!>dnWLvuXp-GJA5iT!+hznxv7!RZ_wh3D0M' + 'c&$&sYyB^X)B5A_S>Fr$dK6Ugzyq6plV0P+B+W*Fw%dq2InMo%6-Ve&P8`)0)`ez*Mo?#JT4*Kdx?eR-ZQXZiURrtSG^Bj-a=Q+^C)-h0(8' + 'a3Tz>4g0L)Z8uOydjsRzi>&YUX!1~`;xTZQadmXP?!5r#21iSk$`V+}CELM(HmlTJV;*qyagf~eQU}12Lv{p6l14qdkNbam@$X9q' + '9-jR`0PoSh_XcEM#`lGf<3;HqKuO}~8X(bfBAwXUCy>Kv?M6{~A+0!OAzA**0jT^~_>^0~A|g-ebg%NnEv`vpgTU|qt|u+LAPDElV~2`aoYOThan+{?vw1m(h%hl5^_g_' + '&n7d3_{D1D^zQm@&d8nTv&*|hI~&f74NDd&U$#L{jyqmkt^LBAuzcCirYH|=xxQ*}p%rF6gv!w}a+*))_htUOMf+x$#i8q4>Ay=R' + 'HDdC&FTVNVTWW7u&~$;>0><97;Xzwx_f{{sWGJw5ba~2VMkC;(_$ReoZl3>kGQGOG8>g@R#CJzW){=IZp;bSA9Fp9hKg?u#ums+=' + '6Kn8SazYS-yZl!~M)qs#1qW8P8HxGq?zXgL)4v5@nB?18BC)z1`Z}zdujKbmo-kPOaZF(hO#fnA&UM=R-4Gl$rg(qzdGD_DFdiSr' + 'G=?(}Yxo@gMq4|MZ|wQOFn|gu-O6NWnJTU#9CfcO' + '{f}LKaKxysN59-~jx9RF+8FP%nHRtRpa1cY%nVgDKJBs!VPSyRar}6h6Mr=4Z7YqAIs10px^3onbZg#i2!o>9hw|5G5n$)mi@HqJxI`v`P12Of6LYt$un@+5Iy4<-qu;M=-cL^FsfV)j~' + 'vjs4&Kb}1xaqn#iZ0`UJFSd6NU>IRgnRXayNB{VAK)bH`!*9YfE=J4THcT|p)8)MVW1#Xd03Qdee8dI@hQD!cS%c^}WPRiTr$Xx^' + 'kJ|d?w=3IvA7mB&`H)BBP3xc)(&fLe+06%I_~UJTN)Ug%rNaveLz$5sD+yl>KYuSpZm3_fGjm|IE^tEXz;?FV5VNnRLu7pM$M2sF' + 'TO6y@f-#~j{dapoVXM`{xp2!j`Y1eJ7I4`G$bX-Oh0Or`BAIUEAYQC-;(}D@gF$yRvPUs@Or&0JEZjWEVOxHG-J3J^W|R50T%6k1' + '=f|;kN8e7~gkNRdt>r%C=DaPQty9$Q?Y+-i8zu3wpU)mFxQ%-!^M}y9meZ3uYa`U6wdvJroQ7W9U9bDFPI?=(JC{xfcC+iowUAwK' + 'ymFjFSiJh&12o>4=-*f2E`PcTVJ86OSc+*KWXJbB9}VOXpt02Lr4%2)n&1IR5d_P35cAu9!' + 'cOMv|m+$@Y$;#tX{t5Dxf9U|e9GeStH+_t7TBiB^fX7UhM!B7ZHJr_vF6OI`4@ZM@DsM6Z&EE8@{e2sAA8yv!g?t=k$b55Vkf=;2ZsfsZnhFITV65X#^!gEZQbUbpQE' + 'BR*+T|*X=Km7UCi4Hj52UNX;TR8)`q!ZdF9dU8&7%9_e#_FA' + 'k&M8*-mloLi1+8bt>V=YY2)Juj}yH2koTYdFm<^nV;G=vd+)gQ68KKbIUj26+&cz*vHI!yEbwJLIs^Tf$K5WAt$a5|jgQGoKdQM;' + 'oS}YHn@63^8j3dW&ybAXA}L?!dCo#^s~8llcue!0P4aznuV`&o0|f6tlD2pA$HH' + 'D%{E1N~2#%gHJ-F)Td!m>K}_r@fYJs?A258#1B51y`~Mch+l+*t@Pnx!qrlF{r(}b!80D_ypwEgtVl4kU9={$k3x~m`e#cR^_oiS?c-K06=-TE)_{|I!^6r>Lgk{y0^}sCG@bS@Gc4PB$&)$m~%Q0ILq+vvtGx*!dwqIb9' + '0$JFyC!6I9@TSG&?mGRHrkC8cKdbF0*!gu=VNC{^elln{}h%EdMZJTNXO9@@1u*ihq|SN9McGE(U(N!=#p@H)SH!' + 'V}x(-n#$_;tmpo(%=Hday2ncEj6?Xj3I&W4a8Hk1tbe?+zmqtjx_7Ma{)0aHbmE={F76NUy1mas2~X?5UViwW<5RcqFBywV-Mj8X' + 'jTeJl;G6B$@-~y@@!L&)%Eju)7k_&_TU-n=pfI!z_kIj=?&1}C)UF4+;MxO%Vp_AL=i@KOVovDA`CsR=>vi4T=xqzGvd*sZCWPEo' + 'a-IbP;e*M9Qe;%tlvtI?E|n~J1Bz7mFe?jVv{&cD_O>@qy)8eqE>viQlHNL3v@RYG43)S4#M-pgxmlBR)~pk+WUsxdC@indDBZUh' + 'yD1#zfwVKXYp7JPgf`v-qt3z)PglR*x(g=QR&qtRovp%(6l!WLml-9Ru3C#?d#qyl)q8aoM=u&j?$va0hS!k4F|LZ_c>_>qd3{}6' + '%LPLzq_n|$ElW;W**I2J986-@Sd^elbAUVBvzNuITx=@BY9U;V(fA;vVCh)pN=@B`=4;qS&ana9TuL>-neD1teo#Dm`EBznwN' + '*NiFlQJe0qwar#7oA$Z36gA#tS@wA=O=j1t>(Yy@Srq4^2sUE6ib|%k+bKw$N_AF6J3j0uaHl-fOkO|0swPjz!ZaEbbRucjhU}zdsTIKG_i%d2tQo>`A=4S^K|>Oo2f6D3}HksHB_%_fn^JmVFAr)6`gJD>+G$`!9mD)t6&i_;Mp1drV;c6x5HaRl~wDS#bQJJyG(!Q>$YHIgOQ`BmUXT>YZ+S?gbEl4' + '_;>p>)4z_V=Q9Mt%dOj?ZlPB>ze~QosVBQtzHZw*&Trl?CtEMidrs3)01dOgENne)NIZY+QJXuW3|>H}6oVHfVl+z?Ers9_FlC2(' + 'gygbyvW1Qai@(GuGJJ>Cz-5jqtnY||*Td(dKH{yH%dr=m<=D*tY<~7$&(Gdp#{H0S`G6KS!xg$u##&#k_fqHRVV=zd1Dh7YHv~eZ' + 'G5rW3v4~WRPa?@I6`U>S%u6Vi!zHbh`v6zle+a}yH|{}h1Atat5b_mlAVs69LI_b>=^#yl<42Q(RI!#=Ol`^+%QDf?Id2gP4ni&)' + '+=VkShn$^srH5cJ#NMz5ds4k1khQ9`4NjvsQYtvME3tJ6h%v1oMo51EbG6+Cy4t54yLA~))Xrv|Yi!C^*Ae2Pw4wukMLz*_t5Kmu' + '1TgDW*LYYGqGsuoGBJADwJ#wwb13S{8g0rB#*|W&xwjHvcU{Vw94{L1*=yBRWv~Qms5v#45Ffk9ImPk?bR`Y#aoBY|qX7Ys%mN-A' + 'WrOz?Ffb`)0JVc9b{}gv33hdNMj~j!l0~bsAZ(~&vgL&;c&|`=u*!#izyLrr9jR4lf>&5-E732^N;Mn-Ua27PF`rFVIqFn(iVB3!' + '<1;YYBEG=xA=uqEkdFr5_a3&v65kbY6;>vzgo{d#xpca9jB_QLrVf|k^$P>VYKW^{1P5;Ibt#yO4VGxbfe`jnWAAq7KNtEQ)5d1Z' + 'q=OSRVQT=)iUI|wE3cJL2wZewA#Q5AZp>ENtPfH' + 'ta6L0HH0Jtg95Kazy-^PI(1$ls<4*5f-KGlA$VwU7X}Lr*o)EuPbB;w01RS?W)0e_#Yj^jjb%)5%t*8LAcxf;*rkw+>IX0w&F%u>' + '=wlOfY)FSDiok>h4K|=DTnHa%#0Qs-24P?~gieAwQQ21w2-4OY0ZP}Ug*ru`2N?B3AY5hl#UVfjFpSi-3SdAMZ||*$(7`St&JXl3' + 'C0|NX&C6^PG_UTVTm1!uwV8Lpf;UIU5O%f$=YfRf77Czgu+%eVI=Gd7@w+h4)|zVvZe4gv}^;AM)Pu0' + 'TR^a!D#<~4@D^FAic@P$v4}?Ac>l>Y@V(xDE+9Utkl5>S9OimV{`5cOj4yt5k2v*yMNV0q(NG0=<2Zx#-W' + '-m8S&aY#IzjryEG>>dI!5JYP@3p8lfDR*0|GKf%>E}lh|1N9~W2S64G)&{yV+~Wfqfr%UhVz&>7HHZzIH`OY%Q=vK4FrJ0C??dgz7Ag2jURA%NR0&z>`%hS{YM&P!Vbvb``)Sx-IFD' + '7)n5013QD%@4Wc4)!Jg#-bcr!?dC*mxk(pINk+;RwKEXyjTfj#*y%E0Wu7)TFr@f^JaG%5W`{|' + '=CC4Uc8$;wyr|kc0ko9TAr!Y%f%)x_1&9F0%@~phdf?%7MM6{q$%X0F!AAn9f)GAKQQ4*;v>CRAD5it6ev!k6fq2_OJ|2kQd)fw4' + 'Fo;AFl`%^xcn>5_4s?nCLOhT{34(J|p;Q~3HDK`~%MR)`@UAtaD6&Stw~3z|iMPA|*^u~{Mz+_advmSk3SO&g3Dtn7V1x=_xDNNM' + 'drwM}A=O-<>01z2*`sq2@@J!o;VXXcnshz_hu?m7AlCUF3|nDzR$A)<7lNF#2H_3DK$4bJnMyX<@xnJ9p%qh3Jq2Nnwwc!!n{aqZ' + 'I`7nfcw_?R3iyBH+Ba7kf-MI6l=su|-g|^K;YCNoV2r&$q8o%aAArz}AdJBnBIMJTo7Dd0NrlOe-~4gR-x`v?c?6ypMs0GQ&Q{rt' + 's~e^!C0$PQ?IyWtOnCX7nk+aYi|4PI?Q_W_pNTjZ?m7PDpKU()zhXR(_>HvqjSSBweR?}zoZYnB!bFnGwmxN2W^*Rmh;zW_bN*Y6' + '={h$vxKrK!br^o_&nMsREc~`COZMt=meb|ra)#hGKf6~cWNkWQ$&gh>5HqRU!%Sl7n`?us)FA2RaI{^JSGR-Y?3h>lexXHO6RzFx0v0Q' + '*B3iXUY+OZI&x!{JTH<5cE_qf;#EqKt8%bJ8H*I!$_Pu0s(5X?DY^j5dJ_>-VWisEn%dbKUO=Fmz38Bei?U}A?J8{sK<3iA0r#xX' + 'bkZ2u2E4#Pl!RbF7zNe;7nXRCC2oxK<2mAkXKihn)E2T4uoEc15V?q-GAX|I1)44({`%Bq>sdM|f!JWXa9>OSUL7i6vs42lhRo10' + 'd~$}k>HKHohliWkO7K7=7ow#MVQqrs-h3!vL&axVh3b%T8d6@uHAEJO?J~FuM&qmKC78`F6Odrfj>-D6ta}7;>&m%beCuS{qkCL2' + 'XQ2?SOV*nxzHCkbTrE_vm^97EMfa-$=H' + 'NhW0~uRu2U=;EK2nd__Z`pw7|ukPkN`nM>rcU*5X<(ZgGvU5g+=(F21{CASd?QA}uw4WMg=Jm94cvorfJU8&*M*~fJHd#i(vAWHY' + '0H(35@T&43cps0d48Qx1m8M^;&kCp^XMCB0fA%UZ#&UE6;Z3i2q;IkP63c>?@6ImU^{e6JBz`uXu`f*0ze(4y$K|2QHKjcR_iaN`' + 'cpIwZk34uX>*T*GKAUVY;_-)i!-vmCK78*xChlj&XG_|<{RL3pabW%j&#hda{rEavUk?Qlw;n2lIiI}T(BP3puO2sEsBm*8hrIsz' + 'YZ%*nMTYOL3m*6iY#5{4-cDZJyAJQcfa}$_#+nX8QNT3k0gp5MFkLR%?RdNwQw#FV!GE{8K5sahaJD)r@p?8bZT_3pQG&}^VF~pA' + '06QHgaOu)`i`Bf1GxTi8@;`H5J?t+l3br~Qp2lgoH-%fXwDY|>Ty8$e?Q-_>@&Pg0htslo$T-+MS&>#X40*_yM;ope%bFDlFLYfu' + '5^H#?HyELIIpG+*M&y~6-%ik_Reh!9xtX)A`+8ZjZ#?2KR5pXLpaJ%l&p)6186G_N?N9=35J*d3U5+z;{GQhgRW$E;4bQM$t_lJ^' + 'B8yg`coV1}2qfc*?SqDB2Arr!Bc07gnc`&vOUgn7ZkZh|kZM=m80ZpdiNHmBP(6raP*4x#@j-)AD}NI#E%gMb|^bhE%wcT5j|$Zdro6MeUI7|&-uA6!lguIcaIVzt~&$86Zn5s0oippIi>MuFsmY<6(hc9LL-$;6f*gB!J7' + '><|nAPzSJ6MnL_A2*M4{abW5NBzg_KIRmV;4I>8wRQNicQDVX{NLvA*AoK+go~yRdAw@KS&x7>-qEpW;#OOO{CCI2!wH3i7%j6+)' + 'gakirMzmGsG' + '>0UFN&1)DzH1KTc;NA5gAsd_Al|v`B7OFy3b_guyZdk6b7|q8DSbWz_R4cMd*@28f?K&t^BM^U>`S*p8eIaCD2-z1x_6Z2tI!-@a' + 'u@21Ylun0?pjbgCqBkM6A`|c*aNtuiN`#{cnf2O%ghfdH2HZ3VtpNX~MJk|oJk6mz{;>y8*`$sr9g3~Zg<`61n#aJPW(`!WNd^pv' + '%Pqk~2yJgbhe}AcHRenQE58sjxZERz40r@qtO{6}ayh~gv#_?e6nhFLyMp->(VRrcK)o?X7)*tk=-Fv7i=doOA}SCjj4zMml|x*-' + '8KqUxqJwGmhIjYE$_7fGywEYaQ~=IO*dVe2t8^j4Kz-8H*cvhUzivY2@(FD0?PsTB?*2wE3udeizo!qkbG{OoCe`nM2f>XpM;2`vJ}0jRy>(M*V!Xp?+X$8Ld3ohu`fjI6A&@CD{lZI' + '&__H;X93i0otxMh?%Z7|)T$*8=Ms-+Va*#+a_LISuB-$|iK1+hy*nWh+Q3SZPcH7w8*n>U$+W040EG<`rD9iAu~0y8PKy{q@CvPj' + ')^dqS0tbDScn9rcY=PNTd-Dc%ww84b;GZF7D-`&TbO_Y0abGD{uprm*JK$^zx`d(pK{Y{=ELoe2)O+&=cD81fjn;w-*r-h`wz#If' + 'ZjR}phU6l&6DOhtjZ)*aVly(@Vfjp0C%ln5!W>(y' + 'BI_^Iok4nE>IEw?BwKyY!u8a8u69`K%(b+RF}BX&CDKX69j?r@YRpQ9C--bKH6TVt!nG6+mAA4h6X)@0k!utZ@~nf_Hpko$Q4*})' + '_+(QPwJT9UZe{IgmZwz^o@6W);|)^`>9lCky=_XuvY|L4WJ7~E&`rSd@MajEXu9-{wr->a)3iC@MZJc)Rrp;xEn0LJLC>K|N_223' + 'II-~{D9SrX7FMXpwyK5*={+J5!XGq>H(K{zdlOlm7A=ee7RU_FiZGjVuTV!GQo%sz=A?!7U1qIYQkwY+*;$da7l^+pSE=i1(PGF0' + 'agcwAx7surI?K(-gm4#h>xlC?6c|KoUTev72wGNUA-^zzVI_5rjux(P!Br>j3V7yMIH>p*hZK6tW2mi1ot@zb!@yOCP1+84@mWF<' + '{hHOD-kA`gx658GFG_omK=Dg2%j+F^5*G6bH;k@i7Pn9Vj%yHTKwB)(0}aT0t5~f%)*B<)P;OZ3hMoN5;cNtDHxIz' + 'fYZ#34A0y{3jpjcX}};NC}gJvOzR+|N=};ukTOB&EFV4FqWA87`dYJ1AZ9%_3!9GcSlBi5FdDEpm|UE0MIKG^oUqC@#cYoRmSB2L' + '+}mJAr!_XT4u^zG?k9O_r&cy#n_fF8avLkSV-PH@fvAzBGhEEWeArU8YjF2gm^=>v$rVF{Pwr0+-Pu~kw?#88%8a4SLV$R*S``Tq' + 'oIeL^G>@Z4rqUS%qV6+aS~8RaRe$AVfrfrIzdGZU0B0=lz9(O^3s?jgpsT4Xmq2lKOuAGj1<)5pQ{j%M)tL~{`PMO(@Jol02oa$#' + 'VHk2`kQTjKm`GgGIfjPk%hbNnJv#*Scz0YI?$NsVNu1v-UT^T+RS?70!5IqhTxti8>`(v@HL3y-9a;@dHWOhA@MGUQ0Fp<2lrM2#' + 'v-7J94i;|**=o%;X0b{Q+HZqEl*tZ_$xI*hFcT@6AW#e-Fen9pM8K*AfvV|{O5NYus%Yb$x767yP<>7}&+A2%q1&2o0CnFaGYXR+' + '6l#Z$^qy4o0$~!&3aY@ZRL5iOE^jG&(5`W;yp&plm;-|-@M6b`T1KXbaHk~`IuWHP2DG)%2}M?|3)7@Q&%{1VlJ6HakiT?IA^Cni' + '121H9CPCJ)oKNGDqGHM{X1PbFv_c>)(iTWgn@ufX*9uTv0yt_0R6L3`a&I*QXFcn1r=TePtRngdt1aMP5FtS' + 'B)tdQfptM5?FX7a-p=~vUAtbKeJl4>HSnNG2&tf|%*lm@6p@vHJOC6Bqy*2{CB((!w|1OrtR;fSmE;Bd#Pq_*0rm&Pkkw$sq(7Kz' + '<9%#k{-!A|4FGdX)GV_QON=bA6si;n4R{+P^f)*rXs(WhS?5rsy9458M82BNo{yt)wXfWrskyhHK-VxG4Pq)H001uqCIXs)NA|q0' + 'r{%8$8lkMCM3;@@NvXyPkh$gV0TDo(YT$4w9~hmSl(SB2%oUzzV{jM{4f)g&x?vOutl&rp4#t*-&xDqNinUv>3iyfH@0nU0;hCQ|' + 'V|MJsJ=X+n@S*Dt^Y`Enp!FWFjL2VW6QFt(+DUM+iA2$9YXWebp&nmqMW+L>=00sK^3x' + '#A+%-ys$}RJd;&_KrTU*(ANyFWKqWyRRwwOD)XR(GMKg_U30X=ovvw<3g`pI9*CSFuN7CpTyTwaoWU94p=W(srkLa*x(J_)4znl`' + 'eh_x8>^c+&h-ZlvhT;RGqY)c;mdmKlWf;@I1L1lJ-GSsD$}FjB9e5mqZs_14&SlnnfrVv|>iw?ij&aS;s?o+i-1AHrRBjUFSOsZh' + '0ogYe1PRoLXQmN(U?YNjPeOaR!c_mxCCa3$^E$Qzp6T{@ChI%|*kCfw$_j>TEhtduaTMqyV5Sik+{6pge*qC$CBk2AIsqcA)JJ-z' + '+wGaIgba5tL4)4|B7^h5Q5CWjvi07pzz@7pM|&n{lPw}^m@YyG;<>l1*9*lbr1K^q=vA*0=4M%XPnR^6K!t<+*@8o`l^<' + '3o7jm%K*dN;6d' + 'Mmmu&T0WMv+EzKZb^)%?MfcTbICJEIBGbT`JGirFX;i3)Ad<94tvX=Lgg_{1Om*Np+6<@GaJ7bUgzgJg1YPQ9I%=Axw{ND`uX269' + 'Xvs4jyMEn@mS^)xp{gNja=)L9_Y>^SwgeasVF1WGPOtD#5)^aHG&+kzlj' + '38Wng85)9XnYv=~Bpw?ArYe*dV_c-PWGF2UWKLHhq_zMYXAq=' + '0Xz7N(g+HbnLnU1lKDvrtfm@iisp7v2ow%P!g?^4?yIoG!z6-B5^y}ox(+i3!YbK&&qC!Fp2YCPhyC6#gEW)6pcgNQo)~wcJA&>${fYiJj0g)LT4SE1!PqBtG#LOYt+Y|_&&!H7rgl95p!3#TZZ3Qi)L>D@U' + '0t*BE&#EN&E@lV})Cv=F4PRoZ+aF6kEX;AMkX08di>D#90Fqz|g{PM^YmD*e=}G7HLCm4j@t*{?9uVK6gTMj{?X6jdunOAi^r2V^' + '3E8aCJQEQc43DKVAnOyifubfq69vI3s0JNUcwa<^Gg6zjOveAKjivW1XV$x)F#aJ643l*zG' + 'JBXkH(@4-n%6TDG;Ca1vIFRnHk>jDnlL(eekx@#md6ru6ZWe^3?zk^tT}WUG+S`c9N}ygAq|9ijf7uR%(Vb;pJj-|1sxe`$bdYS(' + '9OhFqwI&IHe)#M8>W2*Yw%y>-OFmO{}E(hfr20b5z_MJ3b0MrJ9b(dUn' + 'a7jvF4vSxc#W>6nP{Zi%xo@C_X>fryq>Z#<19=WBb{F0!bg;V1z1C{NWGSjo3=F`kRv;DEjTMWhb6W}eG|LdFF(MjBnbLd(SC)s==@*-NWYdHQ61ZRW*z@#P@^q5|O0GcCu0XV_wn<^sZV+Xb_rmFLAbh>N8P9YvQs7|cBw$AytY$J&$k+%^Rz' + '5{fl=VgWzr`Sql#Jf{z8?jU>ao=+Z7fDKmAJuq|NN=*OvJ%OYZUk)g$+(0KwkBv9JfM66*U9g~`bf5=E7BPbRnPggIj*H0Tnk>zbz|xgIT1sc|RQR%q$yx~)0Ch#zqf+qf' + 'qQH(>5P4M^bjmU1e7ZwSgWKqydBI51T{|#Mhv3)%Ztxu{uWGH}4NuU%u|t0AFQNRg#W83I7qc^n#i0Me9{Kwm?a%>Hp84!STS' + 'lX)Md=U;Cqe}8rCu|m@SZL(Dzc`T4PnJ1pvy`A(l9S%75*=?I_LOy;{(jAWhU(Uyoy+HfZ`Q-Pr>voPO{x}}kU60my7Il3i`RbVA' + 'gQCO35yQo69*bO;VLso7G;4$0y3hT?Q0Ku~9?h-^#q$W6m*7YhW#oyF#vFYcV05y?mb>A~tdX(+bSOgu;Rmn>PmHRGJ3pVH+pdxd9-7x>}?FNblT&R_4D' + 'ybm;l^QRvygNkvjK*tZ)p}4{U16egG+ZKg`j9ITU4hmYk7MV~1yLT5$SF4-N#ZRfN^4g701_eO)!u^wy)5?l' + 'T_qc}0`s*(_EuHM!8pmgO;q8TD}zW1;tDH0Z9!D^uA}Y_p}DiTEyA+r>QikF{K7PDiSRNr{Gk}0~FFR9CTx^`*6UtUDcr1;9JQxR;@8)ahf-4$*IQC(W8vg5sBUr56K' + '6;Y5}Z&^ps^%$zmG5VIG)H3m~vj8cK88+@uQ+o=7Uz#zcx(e17WaQ;N!jgIToFOnhVy!-HMYImd5+H{TJ&Ir$@E6Am3z*U$c*-W~' + 'rg&`)-X|3gwaJtyP+XUx7v*UyqP!tARX9wsojc~TDxRM(QNsZ7^I;vR>ONektC`5ns|!HWbMDb(d*X^{<@L}3bjl#!SUii1$spyr' + 'HPS3NvBXx-P`_bLMN' + 'ho@(nf2fLHcj12BnM^*;WOBs?JDxAKo_U5@tsUTC>#aTf>(2-1`=@RDcFhPmw>B9N-0FaQ*g' + '8Kyy2_^Xv*JrTn!N&%V5&fa4*FMj_&|6@fL%8Z3BhLRg=cC$R|Gw_Q!#*g&kY%x1q{dw}br8+;KytrG;e&Eb~H@p2I&EK^8yB{X)' + '?ahoAfbfa)hE_0Fn853ViJE)|*arO1Cv&oa#>uRctKq9>!*}`GtJ){4q8!(oo((-`>mN>QO09tfC<3@3h^OqvG}#*dt51%sF~6>Wt`HTN0+i#t^buf5>Yg17i1H' + 'UL$wkTn?(DY)OS@MBu1RTL}X;#){06tKCbRd3PBAF)Jo9@$yzC_ZOa2-?RR3KE1B(?ZcClXi=s-WF!eE!f^0h;CU`|n``0Sso9ug' + '7agq42zqcyp(2lfd({wlOlGxx40KH_u;Qb9{9qvpc@;beGm2Q8R@(>x$M$_8KD*#I+14JAzo~j)*5hK9mh>D10@mB$R2>G;CWCm}Ow9D4G*kbY!QJ@w^fWq3uY6L)Lqtt<=' + 'hnQ6gTwF9=fE#lY5CG>vUSn-cdCV{*8%l^VfaoisKy_R6=sIwt{95LJ4+g$h7IzN_ex8-&O&9K;xR0O}2n<;$S)rYQZ4}fPa|943' + ';Xtl6ZVteGGnbVRPp#p$2DEi6he2>}Aza-vuN)`{8j=S&QUrg-5Vy{QZ7ek7GuN|VJ=yXaJz+Hni_(-Au-h*vxW2m%`i(c9CZ>$Q' + 'xPvJWevM@mT-M!kPls!b5wu);@HMT37m@dMdM3a;ih}!!guVWlfp{PRJH9(rfZ' + 'KM5p3jFI46x9AY@ST|mioMr4wLiyA8M@}QU{RIf{wm)k>hdChV#5yH7!`F>HQ$14{oAsnPS#7-~pCMa2rPa3QIu!(B<' + '+h{AOyyH1wgm(;^;(oi~6<*~O&~@1hUwHeM2QX9~)8<9%tt1z?dz~Bi&~^Xk4Z78Lu!x)-D=>x{3!;;RKm|`Ku|wH6uEdCVBHA_s+Nn9MBQMebb#t)4&koa54_6u' + 'CasH|X$6|K>OxN@GAoz#4p?uYJ3Lc$hO-COJ5W8FZRjuuTfYGB3V3bwJbR=%Tut&uIs8IP=IU^d6+ey&e{R3-3F~u7lOxN7Vp&pa&JGEWa;5k(%wdQrBB)DWz>g&l' + ')ptbq2G>-^D4nh2ndABe$M5AhTTDUC9##YxlP%WDW7!#hLx12N(6hP?<=>CMadigQR{Lt0^CTr#VpIslKst0MLqR({L);X?mw`Q|' + ')2bSQDFVH@^va8X{E$6IH>=X|h^>s=k@V)X_{UyWplY|U3K*bd)*{HWLeS>IvGCo%RlxemXkDetFz5spe~3juUOC}@CFZo773hBe' + 'R-kvV0t{~E5dcVSCIKZMV85MXq@cZ>gQeAENmIMauokIOMwF7Cm5GK!q>r-#{m8vEAd}G~5RCu=i~!^j|EtunAjl#UB50)zu=gem' + 'Q(#yX7Lk!v_lBJMFRWn23V7#j_N@nyc3Tm`t?WfMUilN6i%FOupb){=vjb0$Jm~_jNUW3$S5A+y0=<_N6o+tBK-?s(NBIP8QNS?0' + 'ZM?Eb*a!!%V3kL+f*>_!*cC0KH8JtZNNokna*9|qtDh+WRb-EgdX^{R3yZzB%z_yq@Zt`in1331Uup^pv_qZ93OeX%vE11*nX9Bs' + '3IHRE+^8XHvt(dkVT}sGZ@Cq1WUeeu0}f@)X{?~g#AE%4&fekOH5gtOVjy#>7GO&+K9wrFQdM(96XCgW5kB%l(|aei4zhv-W9Y1}' + '0`=Ee|4*taf)fU+M;BJvg3$B&*f0)$2a<@6cUlO*p)zB%6;_5oAaC3zV0p3C&bxV?Yh62?qwZPBka5S~{dO3|auUo7<_Ep;E2d3$*Tk9qq!iOk#kReYVxU?vS3v>Aza05i4K0Ww5L#_K~5F7Ilq' + 'zZtrg@1~c0N4)*CPuuE3n%WAi37uN>_6jjynf}XMRNgfjbgP^D' + 'i7U@F>ylfO&P88XvME+{Qh44qGStW8rE^f>Ui5^B44IgC7b34d42RQS@3p3xLls!Ro_V6M5U!hXQJR;jcu^@JL6EJn;2qU5xiVxI' + 'giA0~zp6I>Nc!u&)`Uf4RzvC{CTp;8o1Dhm+5{-#{UR-ByfN=7ioLKhCNn*p;u;x1#(we0Pfp77>~irpwN1%aKY9Q7|EePU(G^_R' + '^z^IYndEObx3kg7!LO#*8}=yA#rfph;S#g)HSbo(SQf)=YfBEdK67$bc-&=qr64Sy{lddXb~=1-q`D3C5InJ!yWs%L@`wsqdaL6_' + 'b>ppdXYKl@>Fw;AHyn=FwebQsFywacq)WNG#UB-C%B3%BE8ojPkHAI8f~uQt?o)7kZyz?EI_`{r>jHAO8D)9|j74`VN{!t?q`K@NT&I-s|ALT3gKn(g|j1Lg(sv' + 'OOh8;LjEJDWW<=Dp>#k$y<|a)qu2FZP=&|05M6zPZ5N-km7RG(shePOhjdDZ%Bw=QIhM{5B0)>(B*4B{X``uB0V7#M-H}iAsU1fz)I65ggxIXug6l3nBs@;>6xF7^7njQQ#Z~uvkY`kguRp2_-6~PZ{00' + 'tQsK}rAUUGjjYfI-d4!p&}MQT(V`b=6k^(H=q~5yi*&mwlnTUgH9GPzmrY}EQJLGUuIt2Wev(5N%e6?)nsy;{KwlDtcimLK|IVKd' + 'Lwj)H&-kNH5OGev*M}|NuS^>RFgrTwAS7TRd>gzu%W+}Cl*`fy2licZh42ZksFsA|wY@EL`S^0cs=Zz32K`VPnWgA70Fd`Y;r}pO' + 'h%StD?9EBoi3zYCE|bURl45~#?s4Rj_ZIhD?5&oMXw?Il9bTx1A' + '$>0DQ&08H>v})HabNaxUKFR8X6+wwtf5Q%%gdURQ5)f1xQ3-DZe#xE45YkYsx_|^T$H;Vb?B1xxz*|;`RiwHN1dKJiN~lSJSPODG' + 'C0D0@4zLcqq}1?>Y$b20X=D7rGBTD~D**' + 'eY4(gQHfMk9=n?&dG@1P)5)t{son9qh5#VQVB3N>5;QLA9G6<2>ToPAc9|2?$W5z0Ql!*cDxlJ3(Sua`4?L#aA?u8XtYG&F5`FG`XE!4vTNgJAXsZEPa2ZSEy9f1kC(opQz|;bH77`$N~yN*n>EYWL-|%wS~!INp5Fh}s%rJw+|FWx6zI' + 'Xu}MF%tV;J@S*agGiIwN7V>}ee!TBiv+;0GyW#~91BW7gmU65{*fp;jsTt8$wK!zF9)CL?^swwhYWlEQmz@@DL@CV^E5Wb|yU>S<' + 'yz`9Lcp8#WoLME=vBFZ1xmyegi)oswS$E7JihH9&(mi&;`&iSkw)W^b6vrc*21uOg$BU+3@##}GI-QmSI2{;Q8oh`x7}6S4V+sE' + '9U~WQc@y6et~<86>)Nr#Dvi87Pe{xv5DMTOBOpvAOi{9#2To-Rh%Kyi1Fw?2P|<)CN?65V|2_L|IXe%W*dP{on|aj1C~H9&c=rKc' + 'o3}6W=2*mg>5s1i;UHgGuQwU#cuiV{3dYk)gAV0Q6U7hsd+6pKNUg;C%0mq3MQ)kdy2fZ{J18}8TxX(gRJ?&v$dJ0XPzxTE4?I4Y' + 'bd(b>2W6)xvqUI!z99)LFAwKwOi86I3QH#_pdND!)k&VreSWJP6f6jP588kd8Rs0=^?irJHyPq2t6ySp9`WS%_NLuVxio~gDt8OK' + 'p%TjelOXcG?%#8FOd{emFWBT7X+n%M)NefU)|2YIsLgcVU7sLBG;}H=fQZmmDneqFa+JI8>&UKRAYi_RTl0+<);jKb1=TdzGMH0g' + '?N!z>sDam2GabaUMkOmlay=z>cSHpEcB?mMU*DV6hgT|T)rV4T&?(&e3Jp}j!FE{xJA?OT%R|-2U*AH4AHA^^v8|xTIw8a=9$?VC' + '!&0IV2t0eJteTM8hN7kCje~f9u3(W|K`0EjFhVsK(ZOl509%0SZ4O+^LrMUQW!A?9g}vjnyY@0~5&^%Jy!X1)p`NMY9pg;A6lhl)' + 'QB^pt(g?3xgh@hKV7ws|fzSyCSzX6lMk=B_)Is3;_(S^p>_JjqfZ#?ztTVQ+RV;LGvKY+O>dSZjihpwM~7`wIeRE>;=b0*b5-E0w(>sU2q(^`' + ';sKe^s5P%F2HDhvEM#_N)Uz*s3LH{R5t' + '$}cf!Uy$KfhYVVXUmgI)HOuh<+yTEE0xfrU@u{bW!rZdU52D%K_3%{IfP5IR_H4`+o(<^`epxgNz03j_rm-ACBH=U9v~DVo{yyS%' + 'pm3fyd9~g^=~*nUTEQ#iqZLgsam`DR$I@DormH&CgNOwRuQ=c9(+`SNtY^b(jnC$bb|bA2Z}z3jORJ%DBbZ=`N>~+)m%y69yH6Fb' + '%{5v$y_3uSV~^T8Gh=&nh8%_M^n#zKi`SDh>9@1LfoE{YmMH^x(3?EFZY!ieYeVM7^mpy?1(VHxcpVH;Dj#1}^8xVEOwrJ?n)eYbKkGB|3y!^jU@aXp2$(Dl-IWMDUGJE(7Li|f2#L6GWJ-#Ipl#AtsDXZTOv;KaV`g*uy' + 'edQl`vHNm*SpU!-deeURDB`V1;>L@=u4>%VY+9x6Fj7i1_Pat)vQlfBOB2HUwupk+0a>`$41^x^x' + '>}ldLCyGQR2z%l-;OVLj5uZ`^G+t3UX#YGb{FJBm6;|NC&kB#{GS!0+;SQI5iyQLS!`+2$40hcA(%zs' + '70o~1&o-uXzEMyrnvnr2fs#r6p$Yt=vM(3-RO}Y9Mr!c9yS)6<{RR>EvC)*27B<&ogtxYdI`<`9Hi+647rLQJafe`%NOy6lf)kVr' + 'm30&DZYL;8=#*Sm>U@pHkfd!~30BfhHH28eot?1;5nlk7LX{Lf9NoKfqcIJ`kkYKuAm@4aT`6bazt' + '7;%De9KYcS?(WksP)h>@6aWAK2mk;8Apq5;RC;{`001uy0015U003rTb98WQZF4VeZ)9a`b1z?Ib8TO8Z**aDc`k5yl~zk{<2De!' + '>sL&@7b&GxCuj@Qg$wksy|gKc-8~ct0)>`FF%g+mNXkKg{`d|_Jxr~<*);;i=FEIE^UZ@=tyUi`-!};FKi+~?mg`r*;~uL@$VUT0' + 'R@fpvqy<&OC(Wg)l|G_&GK*UB#^gy734Vg|z~v)Wuvhi~_JAPy5rNAJO5QY3Drq$@J*2P(S}ji8`WTWvh66S&>Ky*y#{ToytszJ(' + 'pISvu4Ri)QDstTj)L>uK+M+Z{kI)Q^N}P^4I*&(A-{~R9zjJpKR-6zTdKmJhk+LxsdF4Yown|&dANUYjoFFMQcqC2;P{}-5tyW1=' + 'YjuR8s5{$fEQ%h5l9npynI!4YtkUHn82GP)u9QiVR9HjdSg=?ecvGh#vu#k#Klxu7tbdCRwh7RO0{P%AJw3Ic>OdtZQbYb-YX8Fi' + '6@J-3g^wE1)eu@PG?mJ9d(-j~P42UCXzyTfo#uJYpvwGlEwpj`B5_^8u2;&-Lnoio8|EDD*4LS{@9p3RxXD=b%Ohzdu4JD_$!(gh' + '11bwnvL&#r=ZLp$^IC*~_N`)c2@=<)h)TS8;ogW8smKfEnY`8J-!4vj=*f6P-' + '`K_wD2D2H#J+o9;;fCjYDG|%h9g@&3lqY=c!#e9+Ax}lE@j0$)kooy4gcP?O5pojPCC!%s%tzkZ+0PaU^&qrY_Z(ROq)O2&Ifx0B3)qRM*IlFJ@' + 'MWu%Cc0XZAt!jmY`S@fBOgN42p^?733g{)>#EC(^H4>zIF*!9Z9yN_TO#gpZ35({p^I5)FV9%C3LBa+6thakiRX%8S7=!ZJAq@#tDtU*_Zoc*JI^$mS%UZ2N;Cs_TuGmX+|eBI{PZL5st13q!*&4{-xKcxHlu+xQ9Q}?eDQUF8kbi' + '78+X!)NL4gC4Aw!y0lXKJ#}fL(ABjh`u1hwy|926-MZtB=-(>*l(&+' + 'CAt4j@y%Z|-UN1WFTv9Jo1j_pKTt~p1QY-O00;m803iS%`Y)6O5dZ)(P5=NW0001HVRLkFY;AKdZEs{{Y;!MTY+-YAX=Z6-VRUJ4' + 'ZeMV6Z(?OGaCz+VJV?B$lM}a2C;1W`uvjeiy$jGm' + '5S%ab78j^%^O`K0c7!$5sW!-c+1tvQ?' + 'Vr-RnOAg;{)@6M;z#i6{5vFU}`2a-!9AN5N8BFSW7+>@?BPp*?PS7&}Bu_<>Ra4*!cFUT+TlXCqR)#wW{(GDg{E|YPu-TXmo{J+t' + 'sd7TfMTuK=_M&-*N;mf0Qlq$>&Ol}79=m95J9Wkg;r`*g}A@oUg#uk7&>lHzu%Aj2E(@%;E(jij4kUXdV1!8&idGI$RjX9q-rvUyW4%FFC3uX-Hv9|zQ{#!Lkc' + '5b4^O3NHsT*hGYO|4RSQMbFm3PN2DVy!c9' + 'DHUG9TERwzP9RJo5X31;K82g19xW2&TdebdV-=&Y%4@^nh>8o@lB#BmkseGjHA0)dt-RM7{qS&OOKF70`DhPQ?lOtbLyZBW6ui1e{QZfp2slkX;5@0OlmAI?D!_sgMH^?k_<9' + 'o0iYq=l`qwRW<;XH4VR1n?TVyJ0UeWye-g$k|m(I%?coqCMuZcITWi(xJ!G~c~c{`wU$x>fIJaOQ}p0K' + 'la;Zv>W!rORDgwJND979jczyxtEL75`QN~<_tV&?yqdBd!P96ayxzEu^DZ%#sz03!w9QHLE;sq8*pp%8+#%jLWKq5PpWLxhy-UiQH;D>V7;rOxTo)Wi+lQkQDSC7#E}nmzOjve3sAgF3lByPCzcp23jj1O18A4' + 'LCptk?V6CECMob$IrkV%s_Xss#v{1=ELjoE87OumKqCGbhr17=c%VfYQ|Y1GIktES?W5{l(8My8_rwUErYnzHtWr>2k>N>' + 'uPM>7!U}mPIUg}s5KhoDVYLo&W@QaYbf!_2U6k=9Jix&=qK+V7?siq#NbBfRZj~pQAaS0NXQI<*ML&KYtRorP`>`~' + '!$?*mIK1g&tmMirdufxX*4#lvx=vn(G4QhzWK{U?&Qh%$OvSI3$QESG_MWq?`' + 'g#XrX+~Ul2Sf69p!OF|JX>nH81%8)O>SH4V`QQj3bIlLY)8HlDgPT705UN@(bJk0MB!YG@IR2ar_yegs+*g88Z4rnZtah^0bOW&=' + 'XI?nadi67Z?}2P#j&;N~y{>xRs8W1|KmM' + 'LVQ}p5DjQ&RgWm1;J|v6d52@(DWJmRtCRV%uiu9IF?A}b>JvxXA!XHvbTPFpzI;6RDJ@>lftsRiacvdS*Yq;fHSjK7v;Q9z6)<>1c' + '4j8G|6757P1ijP*rNV6(B_r)x@KGp6e^jWv#hmA3JDgSbI8*U3n}*m>lvgET#Q`WC7#L!Lm^XdhsYXp%n5gShl`0)E3a8}Omdph!{dA+T`!?zjPOYu__!' + 'DT1!Ckgafi*)3y~(Hw{^;E}>Kb1`>S;%nV!D_FI0a|dHm)e^_tczVT?C+q@g4Tc{=4Bn4yw|9%Zci^L(3AU1C)Bp5)0(}Ac7c*Q5' + '%+*159c@GJ`gMh0YrDEozPW>>v=bQu1Mqeeqqleyk3HnU1VU|LB%J$0+qp9&!@^PR*r)`8Ei4cilK|TGBM2Lf9P|#5+%Q5pq4`3K' + ';o?42@d1$axNj})nx`)UYq#{G<<0fJ&aZNsELFVud!v#t_F%tuX6mebME7uhiQ9W)Y^yfjbUcdq=hTJ=-5m49Cmz-%gZIJ^Qfh%+N8&h?evDsA$zCDxRrcERa9|KyO1n_' + 'EuqiZ)p|mHYCFrNYt_KQmMv8Vc_R`jxnS' + '$DupJGsybN)sR2l(-+@E*HG`to4OgbLd_AX3FG)2Vo`%o+;{i3Ft`s%nM|Axg#l4lF~I0A&OYEl+@c_86YGO{6=oCsd!K-8S1' + '!8$h|9_Uh0`MUUTGgNB1*#D!bvkWPz5n>#SrnZ(oNcN`uor|2d##Pp=c~HFxPS1|cpB%n8%D(?DI6ySW#wd9F;^fJ9*)a?}+<)&k' + '6XI-a96~+cQuMNPUqY$?oykt?s^' + '@p)T(VjW=HOCBirhBA|uz_dZC>5-1yGDoqSovJFk5BLKc%tN|dfu7}HrIF8sz0D5D)36}T$lcAo&IbALn>>q>2EJ`R{R6oNa;2f7#7=8r-@nNKQmSAZAi!bRQ{h!VU+S|=UC*zJRRG0Eq{H<(P=7t;a' + '^sM%QBI`(|xOMJ`mbNhp_@9b$K=9rM7?sT0Z$j5xLg38;7LxKRlkU_(1P?mjfWAY}=r&^rQeT#_Fr+_NGNbFLIeMnBbtg*3O*h2lTC-v4}12S*!rk$7h1m*V#Dn1Nh*U^;fEy@dZ1>#A_C5+LT)TL0$dd_NX&5VbIliq6XuTYNwD*@UBPZIs' + 'm(pqxz5oAKM9b8}Q=8Cy-|%aZr#4Y&XAe#wORP%Ly&fI+8$kcs%JIjl#Z)S9JzxUHw%$xvB!KD+PnIOM{I{WI*xTUhM?ah@Nwy1TuW^T4@EITh4Zryck=' + ')4t(k?Q7}|Eify<+K80!t^u?nB{I6e#0=mGu7cxF}z{xj(Lz2_J@x?gw_4cX*j)SiP=@pN_D{>uGMb)1i_scSHHJiKLP' + 'K`L7X{E*|LNaK5MwJdPe3fgr9dlpH$jrvWCV&lyf%|pNqwsl@X9!8UFrk2IHgNN!Yo(QLZrVME~OklIzn0AMCVO@`3xEleNO`' + 'HaP$p?|uIuz_#5domK9eTmvQ*JOVgk$5%zESVQXd!Ty`rcg3v$Pv3L%9p$4kpYI#nL(?(qS92Y=JBrwS!DG~ZHq)K=x52q-=c4{n' + 'u4TBn_4EsPQK2Y}KQ>YfOyct*7?8y3(wGmf)aUK=9r}PrI-QVjpudB6JUJ451aH-Ae0AykK39t`Ljt6H*kLC3#xr|D)V6{5W2M{P' + '=1bpUkE`US>JJ;#Pj~G7I)%`0-BYg%2(1IB>^mpN=LV_kWUhEel8Jz4MO5V<7;)MGO;pw9nKM;U0IUo}UHqYmi9rP9BrF_Fq$Gm#' + 'w3A{SfbGSnUEM$J-)H&IP(MU(HPW}kXH<;{`(v#CWMVu|^nMTW9}?u9|E|wg5&pLn6}>UV8=rnNVdj<(Vn6*Drjm4q?uf$%+z+|#uM9dW2d+G(P(ID' + 'mn*T}CD$ZZmg4LGe$NZu0diNeoiwd8t;HpQ1K{A`;Jo2rGMRk8IXkQBv*=+lEB9y3YF+fz>9UNDo}' + 'dId$h{@LRvT{@Xe_VyNSvx@S3vFSH$ndexzY5S!t@-P2KJ7iFf>GS)F<>(pA|P^P(^0msQc9%jc%!MrO@&SgOV7*K' + 'o3fsjvJiCh5@E{k;?1h4i?g!jQrOgNS#(|5sSbDZYSt(E6PH|U>RI13OI2=Nw0%`9`D=f*M&yZ4_v))8dVqERUbJe|ebdg)tsiM!' + 'LvvAGyKhr!4F*^&vA{!aq^#;@t=ezPqHT-%OpLd0UX^wAOWCRdFXS7yK&qz`%iVu=1lu-WsBR)T`7pV&&|KbZ`t_z)jl_Eq{P(2jUX}C5PkyA23BSV{<_k#~G$qxdDqHg_b-@2L6(p@Y' + 'PWv*Jx7s2nR2s~d|(Cv!)Jyz22>^iT?wgzGYphb6Xwh9!kOSNU!EH6sq`>8wH!6G?ONLM1I|_Ex^@6Gucfr_B^Q%(@S^Pc7uAZ%TsE{{VEV`({P^VAkIx^yILd!I`rFTsemeRte|qo9(eo%Xh-2bX$3!D+qUghXX%F1|5~iI%iQvf0a*Ry#O_=V&B{b!k-|fog7oYmf' + 'uVeIQmK!9moeC)$G)s@u(!?a!YQb}#*2Suf{*b|V1%XW>B#rt?&p8!=O*N|O=xo`X7RxT4M$iHN0&fRthg&?ubaXHEi`NL;l=Z_yfJ-YW({`}{M4v' + 'u14b7B;|_xRCYpFTNy`T~j_d<9VcqBh_dR@^UToxNx`C9Ke8(|7bd-J{PY{{=|T0neL`#|VPSPdsJNVZ?vbETFS6l*{$9' + 'npLo^>8&O@s?TtkB=Gy-}CaKnw9#uv5n;=!Oj7LT_x&7{~JARYDPcECh!U5D#D8p' + 'ZVq8C`#k}A0sJ9H$s`7{T1=z;ZzKG541f{=4NAkm*gXOI#F*(4#IjllHw*%Bj9~NOZpPG%?&_f!)|Cg>7vNF}V`9MHXZRB)dcg|=' + 'b)Uhi0wK-fk6xREJTV5^ZT*4p0QwDxt;Zyk!9OR8G#*2x5CJTUU#cY#dIL)jmTCb+4>a-FXQHPalG0U#k4O=kLM+4pt!2j0_$glz' + '-?wlgx|*)9Nct1F;c?EY^>x9(vRyTG%5`l-$B(EsVKa$-d@de-+5jtCmGo=s>_Uo3Sfw;CqWy3Zz~QXfEa$|fPNkfLYv(Np%4JLq' + 'Ow&9^hLz{<;61&1H&KA6yv@>U$srNidC>t2X=5H=0(%lnegE}?P1^!Tehw@cc`kwkCKtJ!%3jBjEC' + 'Y(d(rT2_&e*IYn(plW#Ga27q-J%g-!7!{9BEzmMwtiOisWnG+B%c{SMC8uG1fqxmKjx2zf!im8K3HY1y61vnNRG1{3%VSak{=' + 'm`qc*CklL!?SgdOR;Qc3ye=CKO=e*2o})pbi8dW}g;z3;(4i=MDqWnwTHP#rKnhijxz(J-9U*x+V|)SJ;3(8z3y{5-cHe-uk6Mj$' + 'Z~;|aRl~kr&u~20tAv@Qu~gIyLWhxsih{>?PV{)^K!ep13Ussy(f5OsDYmdIYw>0peH$Innp7ns&_POI#Eb!W-r8)S&56YhRJ|cg' + 'iF&Z&LMUrX7G@B2KI;^aSmfSTT~sE90`r1IiVHJdyN)v?%&L#a%jx!`)>hCAjdMqFk>)D*R3ysqk&6(RX80(820Pc|Uq<_WzwFDeYZj^AY%lf~56?=VEFrV29DN' + '>PW3R!5L?a1}0YYP9w2-Q$zb)iWKpTXRZaswCdwY3Om%n)GiKXL_dyo^%G+`69uEi9VM1hkcB-t)T6Yj^F>*p>Agz=pHbrue?P5?' + 'F5^P@h&v`S7u+@Ng6pC!R-l5S;3q1)DXX*dKBZT7VHTd8V)QO?q+a|F30vB}oBXUT=BBvkU_yeMHS4Q9Hd~llRbUl_hh;-{p<`(a' + 'HolmKrX(#vvX7^!6mO~Hc}&b#SLlS$Li5X%k(f|)NkxKbv0lS^Cmv&~2FU4RRSq@yzc^kE;=I6dm7SV22QCv^CX6HsY7vTk9zmVf' + 'QDE*rJoswAgYwW;GNd_()-0uglqu1=8t_$BF6YkPB`a{BDuGl{H6GXNCahmzA7HwxdQrA$_1|ym<<&PJ5tS0E{DMpXZ3JjKlXRB~' + 'oms14+xZOQSe506G^8`?g$td@*UcLMbT{wdW7%cAY3scNU}9TEpGk=HEB*d|?*H`YJJWb*ET3Z2H?yMah4zxuVl^~2m~?u@rB4^w' + '$%C&9v@)zx0}`kg#y9RtrU&e-Q=Za2-!R9HpxE^cX>#M=o0eRK5!Xky%X3Or_LRx#UVCP+-3#dEY2H=803knzf2Dp>' + 'LoAj34#=jajSMRAmq>VNl%aLAGUS%?vY3P1@pQ8~1$u`{dLyGxRyENS12iicG!$wMq$m&$%?e_ZRV`q`J^@A{)R@(Ywoc}GRabqU#|GzHE|TCpRkWXd^<@&C8vguWU)@Q@r-I8J{3#hj' + 'Js0}IfhFv-;u;3c_IC>_rkDV#G40IK(@6u&4' + 'I;J(UJM@d@)m$MuF>EQ_EARUk#$ayWe)R?fBq3`UOak`;g@0c?7yBS_0ua(fu`FU}XsL-O)%-5KpV<4o0X8G~*Mmb^nh%{cv%8vX' + 'H=qc_Xq|Kk+C@xe{;=gA3qhInD;el@rXTGhHZE<+3D;Pvmllj%|F8=gD-22*zwBaxg$Bg~Z~bc4!V{IOcXkzXeF+PJ{Am{ys{y5W' + 'ezHG%3k;1p?mN4Rzto_P|8lUBwd^c^EdGb9Fdiz+G&lsxkW!f}m4hfS?R{0|Z<=;)_!OZpBu3Z;%~cO2%%8AT-@0`x`l;#3%b^|J' + '$kN+}T$bm>auHujG4TcLnrO7CP&{0g#f6k7W!tQi8)daZPVi`VUQ`s?atQ;5P3j+LD*vxLu#dnNg=hfnvY3G!@}}9J%Y7QkgV0>F' + 'zeVcKMkfs3Mb`H$x=c4Ux-)@p?n#I}1Aa>{T4^bM@miO+jBfNh0(2os%g@P}M5iKbWQp^j)Eq-jYV;^G*FK!YMwh' + 'J5b@HW{#Gi705kv6sa45g0@3{224NsbCB+?KJ<`$1YJUnOKRe%F04`JB' + 'zpMGWVJ014rrmk5E|2#QPNL61A&ox&TrfWi8z2S%E-rc4vi>2UL%cea?+lCx!w+aL%d5O4l0w?d2_ml%myCq&-GyF*kP>=5NO=4i' + 'xHjol?zdsGyoddbWdJMdq' + 'mywLVV|yCU4P*07HHZ1lqPx*;#G0N;xge}^-$uq{yPA{j%7`ao;9L_;wB?-X(H3(fwH_zZ7DM~mtl1EK$reSIpp*Gxvs`M23(^(5' + 'Jw7z*0ar?h%V9j2s' + 'HpL)J*oXFNWi5-DzvBI|3PNOC?9}4lEjqrtis(V_J3=$NUv#h_8YcC1fd*Bm^pLcj!|13!M=Pi8awmOL0DdL`tG-M^8XxrEQQs({' + 'L)KB{PN6r4)F)nSu&OpauirCN45LNMae-CY8RELZUH)5oOB=Et7~goBDX3L#wW!ADPNN+w{#GH}?ia-j(0qu|w(!?1s?e(;!yAwr' + 'FLNZC^%r-oe(|VEJt?rqfU-|}C>GwcEpms^1FFDUHovoVy2' + 'I2~^6iVI#vr<=tBRX5-aFJE#CFJCIL#tQ!!wx=_YToI1hUn{199iaemQq>nZ_F{rei_s>#g8nbVL1ukC))ir^H&Hts@6sTY)S+_Q' + 'Pdf<8sU_j5mIRN^<%igUGQG~CI>Sk8Mc-rH^+dIrM0##B#}`pZ{a6@_>1TW=F<91JNokqYN!xBrZ0w7>Q=55vif}C2`z*SXeu+|#' + 'Xa#hRC_Zpr0U@}2?^@chhZ_Pk5$w|9KJIjkE--g#`6`4mVDnB~k%zQ2*A-B4C&Z!(wYYM&G4@|;g$aS~08wNiQ|%>=Z77O&nar_S' + 'gr?LqxQNPy3dOnJ=ZI%1RL2}8FKhKe5g2=bVx82@Cp' + '-dK^9A}&R3<2*)VbeZ_Oh0JU}xn7Tc*TQtC2PL?H<8*b_@}LcHfNcnJ;_;em0L$fo#~C}Qh9o^hcXA}Yk>z3(V&l0eLzce(b>%7~' + 'B0ZC8r+4>vcSw#CoU0ie!}^TdNKm5qHZzCl)!|Y~3i8Atw;bCtGak2BiemWKbSF+un3Z8Nlw&fUhF7A{<=8s$5$0@!!$}1r7lVu}' + 'CZg7;D{m6tb_m@qUnGMD-6>1fwocN9$UM4b{DY' + 'IxeERMOrbW1GkJdQN~NC6Yv2x~F~7%#p>h9Ux;KK-*qHO%LT(KpDuCG7nE0(VO-r@eH)=yMV6fl){P-C(V8H8|Rmqx=adM_wtzevORV-2U$LKQ9u=}W(%{HqIo;L5IRatbK7B+z^8{WJu' + '*W8rqV(JiUa9ew=92ANjNZDKR4Y4;*16zn*u7Xcu?3lM8gFJLwkS)!qu!M`PO?t6~wYFjra@~SbHA6R1haq<)S>VG39zW?)I}CS(' + '&4=wbVeb$Iuxam@vA~M9jDW%MSrPPxamdz#RJw<}y=+jhw>B>qr{m)PfXnf9Vl|' + '4jq@5>w=sbIfcOq+cuCT1u8B{_#)3Gz?uHVAu%XlfdL7Tx{RpV&M6r?21x{Ht9?3^NY+' + 'i1-Vv;J&(GvX!YU(y)tR$o}iw*9Uf!%h2W#RADX;p!8jGQj&{4zN?*K?ggwNGAb?bIeD2lyH)<}xG|&*q4LnP3d6q6Er3qSv1u{qE7Dc8K1OHOT{$5`()ZNi^x^>jM7bnI0!2LV5!~' + 'o=L%eXRWt>=wK#$b7LPFH;>*WesW^Ubof~);=1?bS8mhs;{hFP0L(vhUU6-mD9#Rh7GvL!D71_*(&T;LiG#yZ;+S8W;46g-nql9?^w)ctcD9$tY9o69vQ202VE8G' + 'IOr`SMhl$aI#Xn9GC1cHs*4kU&cuSO!WJ8tdK3P^XZ4WXATjHq!#BH&=E!VoOBsA^jM!1*TrQA#Z$>M8=~ejBso<>qY%qgbaM5Sp' + 'n_vaPS7!c5Tqlt(E%7S%{n(fYkJ*mQkGT_CGh}mMW3GmHv6{7^0>=D(;^!ZnpI3`O??TBXcsZ_$dQ&U|t8z!4llVwatWi5wD$h)1' + 'OWv861~EOl_t&GJ4&8ODnN99;m|uHH0o+<0o;ph}N(^D3Pc~kJtSsem6w3_lYbz{+N=6|O3i4A**Z~qZcQQq^$Rif)kpNxb)zG(TYN1gC' + 'muEW6*8rkTIPRn#8(ld*KOwXCIeL_&S6FMg(-T9Y3GN903WCk963Pb=BX-pXcVH!lt{!0VFwM3?F8o!|uQtnp6`s<$c}&J94F=ON' + 'MOd_Ik!!W^u{`Vsb2@DPOjne1I$u=FW$b}Wm(5uPT>^^l$iuX27Co*7n2hPR2S^UtUMN*tWl_m}|md(mG8m$0)wc?1+RhlF7a$XGjc{x!v^t_;8!&9|G>u#hjk6hI@5V2c*n+C@+' + 'His+(l-Re+Oa%D$0#QhhGGsdth0Ni+rzK%lY`S9U#_ke3l6NVPOE2@Ke(FUs9Pp=?sEPn1S)WkVu5)!6=lO=I|CHo^TB49$c' + 'S8eA|j>$u)7g_bJ9DRHM7~YokMb$PnC6F>A(ii9_31>-_tQW|i>?Hbsjp^THdz;ZLNi36' + '?k=_XSeCtAO>QJ37-4^ItH^NjX3f^Px07c-9X)?=@9|Oow;v{lG6&-Q$3H&!A%6-lzPS5tr**m87V53cfHwDgkx6%5XD)#}@or' + 'leW!J)QfrAp7tukK7(gcSWOwu$hFrG&' + 'mA(tk7@0mo;|3FWx#PklCmOM=&s~X8*2C>3k+}{6gsYKnoX2I0G^27N}_Z*&>*7KJr-y0OZ8JspaKGvd6EgEO&WUSv>{6+18M&5H6w#Rg+{5_1f!' + 'U;lC!Z2=aGCFs(-&M9qmTKju~yU#bv>RaifH$^8d=#^>6`CxQW%S%vX8WZQV_ONJjE4WbcG=D4d^j2hs?;DQCXIui4UN&#ac>tWa' + 'oRSl9%j2c;n%p#cfxgN!SkRpkl`&m3f9?YCVBII`^!NT(_bKOB^!0(j4s^4`9ISN*Q%hH4Lb->yhol}bd%8Zuq}7R^WyfQJrg0y+' + 'kU?7{BbdC=pL`GAu~8rQsmoxVp8@errTdYYp8N(c$InBA0SAI0Y*8ntTd4CqEFE8FmxE&S^$deIHW1#(P9F?n3HQ6;-d$vdh&6hJmr)KGhBZ4L?=VEOQJ_XS1j?N@`si8Mtz`N09LqyC}s|g_)M?vPj@y=-dMJDrFt*d3(rt_*(x5ko{Ypc9z>!mxWpS$awW=3x*bUG$7' + 'A!hE7T5yLvke;4fE*2<6U6k6+=2|u0%YirGoKx<&;7iwccyegpge0zeoEX@hu!Nv~`=JLDd*UWGNO8)YQ_E=}TVUEC&E;9pR2{>ty)O$*nDtaT$Cb(+FKj=yKkaN00QVjWrYUEw}TyhHMW=oarYW+H>2wQqIE<0W0dG2j' + 'L4t?_dq1U_qgzQSFCL}}A+nCz=8ctjY=>dmosnwxA3GXN2^u@!mAKDVCX`a5v$ol+p+MD$vmSvvI~Z8sobgQ^e2+6*!gMZ6`8kcA' + 'iElDlA>S6|KwF?ap(Mi>#_1HO+l$+us#a*`WB{O>TZPX>Q${)-+2Ik4OEkxa^31AuI6+K?ONeqOJ0B#xes}D=nr>IjuTF~Iz|}wY' + '{ujG<*#E2UhCtrlE+_elpGbJv4dVd`+UuaT@e|46)%J(s)wbK_K&CcvwJ$b$PhTKjFRvnT($rY629uR*jwI+b3R~MIKSK{F!xn6b' + '^L0V#wVZPRd|el;$^GID<~&48G*qw}8-UJjPqr8u^y#J@4A(kh+@6>nk6Q1<76a0asJ3e6LX4>9mk9+B^V4ho^XdZsPQwOP%(NtH' + '4ETx(5f)!DTBA06&SR9|=1=wlcAVd8Ob8^(wu~uQliH284|W8T+<3-C_r{|2tRpD;^*A5S%pW(FtvRY6gweD9d|1H75Py#6-&S&q' + 'D1IF&dw?Sec9WY*U$Ep3t1VcMwM4;a#+#WA3<9V+lMl(>2Vy=Bb?oiE2MZHnLeJXr1z!tp1WQryV=RN^6kB^zT~>8sRH1S@oJd>}' + 'S@n@TUh7;V-r0!~Im=_t9m)JV%lDXyJH9P`iRPI_h1Rx7*{weNjN6l&P(B#Wfv1Q#RC2k@`wEm9(Y+j#QXfVX|MmBxJ=OtMc$4!9' + '`F8ncOVx-z&uRh=eb4W1E%Aj_Vgcx0FISx*CD`rF{OTCghSD9CrhRE#XjG8ckJUA6Lkb4m_jxY^LXe_U=rQN>yr%pVMobw+ZKD' + '=(uVNdRJ3t6R1fyEs(xfmdc^|O*AX&rY3F3xYks41I5x!t?p+?LT4K3k{V7iSg3r4G0Xq13r|jW6e*OzM1jZ-gl19uZ>1{7@t}I$fDl$t%xv1f-XuDcCo?Ja}S|gs#!wBdWw#)%GorZKUp`*(Yh+JWX^&' + 'NT=~2bsE|2A1GSTS23|3Uo3%==2t*FobH{Z_%hPzVmPfl=bBgWjxLKC-g-=~3N8L~46L40vIKZhq)2*gN!RhtztIMh!x+Io&D*nJ1R@ODB@q24h6ryv&}Gqn1;b(7cH9ojDeCnR$T8qIqzG^!<9M7~fH~7*BSAS2*Y%' + '34%As{2MfLwsTC}&O=xg>n>>Gbdv33yaguTdE3#?7?V8G&fpJd#Pq4|GOWwXKCbBgAg_L6-gIbK9p&-%L~IbKLG(0t`K(wMe1&x!' + 'u!LHNWgo-10kvQfjnqEP%UYPRMgD?^ZWsknm*^pc#BWx34gD#ClD9WDscz3eljQ56GEqeu_LDV+a*L_-fA&D<)tgt?5-+F{mzM0^^c`J0>H>@xbgMlL>k#eyJivF;L24y22&qi$>CCOP' + 'oG8vIZoflZzKC{7n=V`P!Pn$Ftn>Wa~+fo' + 'fv?5E8xlXJgs9Ola^!y}VGI`g+)xgNeEA@zpmU&UuMWk{zD<`tgC0DtUXj#qgmlV?YI#r!gQ|_sw@s)Q)T=}Pw&29Rcrl+>?NG}@' + '6)8`StYk-x+XX=uE0&lwtL$>z{D;EDcMBFEc~&3cGY_5JJE8*^_6<>wVFpBGhzLnUaum|nqHR`D5q-ZoHKqWCHtI7p(^7uCD{+x2' + 'EQx}ns8@#2ix)5ncuOSrMr3@UDgR?9H(Cmm+|+=D4zF3p@x6Q*b;VW3^5`4j+KEgTqk{qj5WP)_E6utF`*9^hE29b)m&l(WJVMG}' + '#8skI)veGHHz3XGh;t0yLAoW{)~{Xnh{RXoMD5u8(j>%8J|zAno=Ae3dUDn}l3GUpI*mD`mANm%JVlE%q2xLInNWuJUh_j|CXcYq' + 'hwdqsxwYc)yqyuvFVY&u)2Yu;R3Xhj9AM%~M)l&5kC7{b!%6sF#d}P*L#v=~ObwJEo0p}{X63r?tR?gErNh{Ol8N*e3s7}Yg5aR#' + 'moHH{fu_h-$C=}yto8@iJ-R*R&0r`Gmr~h-9;5L|AB+u?ElLuqA2M=HV8@l)C^p-;Ra}-#BKLTrFHgOIIUjXYA%+pDr0$;Lde=}<' + ';=o(G$Y<-%gB&7E?FOI>_>&6IQ{4?kToc4G<}UUFxF?KXZFNd9)C}WoU?y4d?!ucj`sQ8*6F3q' + 'q#KxD{urEKeQ}Sh$c3Sn_%uWpvgRMMA(P4bVcJG9%lDmZHG3-_{?NWl=%PV`-3_7HkYUm8Y4kLEJ1~Uph!+Hi#t)R(;}HJQ5vNaH' + 'ON74XBC&zSBC?kQRmH{8=%JaqD@N+c$XMe^HD90tx+*Rw)8K&{XN?Dj$4WO1$qu@lj-AWVgItWV0(leXTWkhqJL22g-BDriPzG~M' + '5CqXo+=;?M(j<@%Z7SEH(|a{_8&A#&W2n4TzQl{kqPNf@^!XdM?%^(o#nH6MBOTZlK<%RKh5|XZQlLAWV`vn;3F^(BACQdSwj7c*' + 'v<{W_wJ2HA*PrO9QeL^K=e!u?j}Y2y#{b$!QG>C+c5X;ovwB!^HybAnd10l@^=9sp5)VqQBm=a2LQ;R)8b5d|QukUhbTsF*pv%+^' + 'FK>BYHc4BoI}jU^D~OJJ^3gE+L?uK2sU+oCMAQJiilnbq$C{A_hq?+LPpjt?>}DO~Y&IYOUo%8>!U7T^U)(E+0w|9$r;GbKZNkEX' + 'cGBJAptA`rATx3sc7_tuw(_>jNTZstz}V|W$L=m@xl$+ewSi;trPgFy2QS1G*7a4a%VhQjNp5rI&((Zu041>(o_(5BrrWxg=VZmK' + 'eS?ez^g|d))Uz)g6jPp9&v_7xIzRN)GZ{!?*7!Dy?$~!B_%yx)NW0Q_JggsKXo1p*b|jHCh}9C$j;RIHZBt~iFt?+X7x$O6EmuGe' + '92XD^xp+FY_Dq*W_*$RPmyK>39N6Br2hB&YY8GxFZi=7A=$vLbf=rnnMv' + '3NcHJ)+dg}rm%)mev;%BytG_wsXOhT?%o+Z(@n<|^g7=;c{7gteD}lLCUe}H7)izNNI4oLt;8`^I0#6qzFc*&!{?dh3oYcV@YCwgei9rxxoS7M$hsaI|KX' + 'eqDSAFlRt88(4)yy0D`pNx)mxAy3~!?Aw^(ZC|V$p$qvP8!(FmxVP^%(pXUZf>~-DfN5?OW1ticFtehXbWk(tvJdqF)XzXxNTtFk' + 'AR`2zI^8!pasM(NxD~H66(JA~(22*g)zB{d8)`SIc_|X_rU@C6zzE=v^ik)o=6AlzwL|2J!g(^57bvlj<0xbus3jB#VMyvjVmlF$' + '1hRmU$9EFh*QvOFCmYf&eAQ5CvEz-aVpa$Aps*ACKesD)(mU7J>C8g1MdO{Tguzyd4TdmbUSs5I' + 'XF;8COR9%((k5p)@1ozi&WjKo42S6R$e~C#oGQp#n1^>ylEIR$N;^_gcrh9D&`F1II-%*_C`Ngn-L({&)jQs{dBer%u73{5MG>u<' + 'f&Ovc%=JFpx7ErD?r`Nn{(b0YmPm$}=b^U`ZOcbzt{}n1#uz?#X*eUDZL45A1AfFam4{A4e8U}UDwPDCS|F$4WLfE83K0Lip!Ie-' + 'kPa`FEzL' + 'myG!BCTz9QFSzw#u#1B0g-*fj)?0#Ys7Z8WMg;na=2vb|1h|GmtsElYQTH`s#>e4EDXp%xud=kab89Yuy)R2d+3;xZKjqn#|j%}m}anfTh2D1T#C%Uhr3p' + ';is&g=c0Vd(*QG87Bz~DcnocBiLts@ZB8Ky9=8>sFz4}jf-Ttx0M&G=3s}#&J2^Zq*@9KR|E48r>qv!W7hg8Nvn;M^ul5)ZXVxyD)I&yZ&1+ZvBT;|M^T=?i{?$qb0SW8L~lihC?Y@x' + 'tW#7#vj|4tzumRV!$(5+Ruj&#H!B2=k' + 'uFBG{eH?Xa1aQ-yo2k>W?qF{;&Onkxvv*E@B#s_(P{AGZ3#NmV89e?T0Pc)^7l6{A%YEA0yLIc92r+xYr_a4xLHMBi@EQ0-2X)Zl' + 'cjv4-1GkPn$`gQ0s6R@_B$w3ds#8$`mb^_j0+tsbWs|5aroBIJo6Q*;=3c(Mx9J<|+{>3~6#anNnPkfDjlAu&q+|0Dk5|T;*x__i' + '&noeNCu_9(C|oyvVif5oaOnSVu0DAY9bf!eyOyk*Y_hSeKi;ARpTB%5tN=|$ZTcaMnD' + '!7E6tt49J7%iu`?WdWdtm$~JL=orG+j8{VW%v5QLS>fKOvpPggeZ-m@j{a2LZLApcZR@6+nvP2892NSG8T+yv;(r+3T+XqW5n}lAr9n3s' + '{JlF}bcmDtZ^C!#|Dr`g`hCedy%`v)UlRx)LT(y&MB%?f$HWb3Bfc>{POX`d>8EA!I' + 'w$QtRa$rzDtvuHvno(cw{qm(1nrrSPFJBHu5rzZlJ-7N$FrXPb{Xflya}Psjx^+22^ow9e-AdEiz$ZIKdhh7y5BEyMpUY02Hl)zV' + 'rhM>)aHC)dG6wW{L5ul_E6*4mjR7C`2t-Hw=fN{J<5}=e?bqdpJRC0bY*T7-t)J$>$%Q>yPWaO' + '3*q=4agCK7ymq6urs3qz*E_MINp(Ed=6&T&eU6&LZJpx@=G9)olTFU{l^M@TI4CI7NWG~rx%E;oBO&jgKpl)J8j4%nvs@JzW+!b8' + 'oW`etUFlL>e651;?0i6ngBdw?6u*Na_%$)N6ZYOyX1`I+W>*l~atBUW5X-`^li)p8!n9#5fFVP_K2tKlH3mV?RD=Ov@vX1m`D8t?fe2zj*lSmm;)&o;$b$zM>qr_zt4Sl(#^H$x~uwW<1YhEeUz$%{c)' + 'WZ?f~+x}$R{yo~ZF;K4QM0(u^=o>n!d+74jV|a|m4dMQ`zVG9=lV#HuIVb7JMW#@ZsBLl>O`ZU|sIa#x1@+|JC!_i&qxvVK`X{6M' + '@5-o-+>Qi#3eq6KSVqjK5G*H;n8!Pg-6=o$Pczyti&M(D_Q}9+Sx7iyV!+1#>w5tVvGN*TfKR66Ynzhkh&iqZuZL!~96=phx+VaL' + 'G%ikM%4(Z_GGhOW8nHjwus_+b|9x!OpX}EE?sn_{EL-)C_US>J^yroMLu!!`JdGQ;Y~3;v_;=jX%D1=vz|e#Z$UVaEiz41p00DCF' + 'c=`K^C0a+o$aH$ZQ}QCB*u4tF7>jA%lEp8o@=d&d;2Xsf30jnxPWODrfhPpVUb`9o<0jLlaq*uqH^H@B13S3`hFt)@lVgh({+sv=' + 'o+%CN-UAE}i;j;*;+22nF{)g(9FlQXNO8MeL>-8^x#s;PRB?7%{?}sTq&PiDWA?K{5!Zhe$CSfuD8Ke0j_ll;p3WS' + 'CC?A0(+)}C!*1(c)5isx9mv*ZLbHH}?n{?r)`jXaHWSQ+Z6};CW4qNh(Ccg3ZNo|KhSL3od+>3bN(1BDmc-5iw=Iv|ZaM5k{yH1-' + 'P#QZ!nAKaROhXW|djF~hZ4X2CsFHwjD6=`yjV!e9x3h0E>KZI0aHFJS1q)mmzLVc!C*$e++8>P+!-g>In;sKuRZIu$ms)*M`HIaB' + '-pnw0J&U1p<1{EGr^|Acn^)8YIgh8RdDk2VG(H=A7EM~EviW*2u~@g6Z=(Zpn1xx37_r0VF_S3Wz!>}rI)Kk2u0ImIifvhPH*;He' + '3k|I+uiwN9qX!3C*m+Jf8|&{mx2C-rG*L-tizV@ySlgD>o9Q)>5r!^d1~%l;QIE}#-@a{04;#@nib6*8>-BDAa>qM&I`a`TYQG;6' + 'D}TQhebYd1IOBd5J3Zw-{1`@hzAosd7Wz=_k7}p?x3LQ*0khwLsqfdd^x' + '#a#a=8{^BSex-gqmJj#j!(ZjY-{r$|`EVp3ev%I_HXK' + '07&y(RF~zPvs?iTOt1>d7KIgVmYAant#qf&rk=x(eO5=;&9a(Zsnp9x;f(G)jDM-tu}yplUtJ$#sD}MzI;CJD$+LyQMdsYjbl>Cq' + '&yHf<7A9>64Xw+LPZEvdOg-6v(qazth0c*;$=|#9SWQBrai16{Pc)_zAHgSf1M3DjGiZo2P*EE?>0PeYOW-Td_#GY|*AgVzurgkW' + '(f_*_nG26;p=!nF!#hv_{-GlH2M;D&k3F@KAtGa38VUK0cmm)u#^nw5SpAUCTK0(%wnU~`pqn+YA@qVa}#A$Y&1xR3_POC6jve!' + 'Bg`)X?|Ng`%V$Sc<&RuFIWm$Ozg|6dAg02CUQ>1+ZjB3Lso7>$g%P~~?Ou3z@!kyPADEqT0kl53m2^oVsr!^dd0rF}`R3fr%X#oS' + '!X@V5;{@2b%R_vgXIHi9&k%fX@rI&89We~@{*TY;eqs^!XyzBkD|Xvn7&}E)aVqfrL#zgSRJojs+}r-m5vOuoW62E^hWEBXLFTp#' + 'AiOSa0G5ZhQGlo89rJ)~>tIt~001$G%&j{Spx}ZUzKvtKRE-TKjiC5a-qzi-2!n7f3LAk<0I?IY0u(4NUt;$HIF1{;EzrP_xe|Amw_0gWr0ndu{_mX^Qlhl7lN2?OSQ5z@a^{&=dODr{s`FfxN5)F4gxSg3C?OkL' + 'sWPi_RU5IaYH@n?;l&FP!LPTm*uwY1B$LURH1_iRgAsbWE7Vq%)(E>+V%IDRZPqFmzvyb+EW|qGadmPkvZ}PT%&f?DtukBeXTnrM' + 'S}U`)E>}X9IBilYl^ao&#a>+Qud6y+&%|{tO9Ok_!ro>uF(;PUT4kGErAsTcHL6%9p0}~OC`5H%^4^g9z0#|-HOCXX)&`H$6iU=e' + '<}^ju!2xQutBkf)y%&YtSB;&)c0f*RS;Oe-wSj54Iu|=ti-pWK3nWO>' + '`8cJzgApLFjji>fu?n|Ufm$2|ko*ZR5W5;SpTSPns)k)jtIbj?KB0~Yq%C}lBb*5$EtSSmc-%TPHncu>d`n>_jE1r}~5L+RQQN~YbKkiHJ#zITrH|2xG*Wv<|e^tpJU0MYYwIo;oW9q#Kc|rECfd47PIPuo^k~M!bZtUt;l-(D_o7VCJ&We-Jbfv>M!q' + ')pnoI;XYnozyEk~a-LpYpIo1&7bhQ1`Dm~fj{%h6ATW6G?&QFMd^)%)|07clXaI{Gun@lC7R2&DaoD(BaAqX6Vf~ItmPKGh=3nJH;7}WIxDM|68y{9$;P4kE{}NlPC>pmbk1-6OH#*PdX%Cd-es^-h)?hde3*1nrzfL)ZWdwsk|MEL2@r;6E`|q0|@Bg28|a' + 'V6n~{Ws>$xj%K=*t5Vx0SLWE86b|?4rbT7iaq%B>XnT7f{zj`T|r`RByg{0SN_' + '#1n=YlAnx%L?a6HX~xtr!=U>cJVIo!2z1C>#638s);Jigcv_dcfWn+XZomKw' + 'S+A45?o_)jLp;_o2C1)iT$hpreIEHedmIRg!~&9Uhto' + '#_|GWH}DPMV|#`J&bEDL=YiIdYM_#)5vF7yIrf;New54t4U4sjKc<1u6U9kOCGLhW#ba^+Wq#-n{woFL!x+GX7(k${BhC%fc4zky' + 'FG`$`^P5XT7Z)=Y^O?xvA$8E!H&73$wJaP8q&~3vs?lN7Z;fj#^UJdGiP5VhSBvqh!M3n' + 'U^}2wyKfOf0fXQ{006jl`Uk9q_UeKt1x!JyJ0$T4UsmIZw' + 'dkjR;;lbvIpW@I`T4p5e5y89(-QSDQTODHIxvlli@J;)i4UT;vmwD9R|4|Rwak8uKqd&!AS`*LQuDdc@SwHXu' + '0G17;*js5fy(I5>%SP`;=?U3qQ|_RIrj7$!v@al@%vkW}8ogjig+8W^5XZSplQmNvEB~=xYYBW`O%M3hIc1B;4iZ&t0|n=q1a7M@' + 'sb~SRLv+0E35!0*Pk35DU`XV1xA3MzoAZOTGbY{H7iWL@@u(tzI;wQ&(}w7EG<>$-^@FH9k6Z%mKOTdIuuiBogYid$^l(ctzV!XIj*juZ+rcP!9J!xUuCU4o^CTsouH5)#X+d!p+*Fi9Yij^do;' + '5Orr}K!8_{hErdPK?5dEv|&`lot0ys2!VXC*n?r{`EkJ*A%cn2P}rQo$VFw(K#ut-76H+6$}NIcQ4a%tk7A8IvPwL7wvWm50PDzl' + '5|(r%i+0wag%u4P@~%pCJ_qQ9((8L|&wTD5hEAQWVD2CK!>uXIsSF2swO=szp^xCFut`-%?a*KgWFd3=VRIpE=' + 'W*z!T&ShzkTD&O=^q8dBv-ao8pqYz5caS<9yDWs?+=$)px&eG?O?)KgP^HWm;^cYWwH6~3Q@2kDhRlal$d56f)K}4S7k?l!MpH`H' + 'rmcDSIZhEBZW*)x#LUNo^`YbXhzk(*k`jftl+x4iYjOvS+8>G8kKBt!;5wWCrO^e}$z<42B}sP#^cLQSTR2FqUGFQa' + 'c6*GsvQ3^H1L|#J;Z53!V^?L|op#o?jVZ0F9HPf#Ij+P_E53~W4!EDSA-LwhfISAhM' + '6+mtVwRmvJ%cAw^A|_@*h61=eNFiNe^U?K8`zc_11<-MH-b2=Wd+UD3wTRPzH6L' + 'jKd5|lU!a7X3hHb;pa7kdlO#oP&MV&P}=LScre%jMXX^PbR@taEc;A*Sw85*c02VR0+pDc+|ea~zpxRVfr}+RGI$@>Pu#i;wD6YJ' + 'Gp>S^k2W1B8P{pNzXO3ow17y+*DyUXQ|GP?->lZuCCKBIRo{}k3~Z2jhrPNdQ^+oTY@fQp(X{XzQ^Ghn4x9upUn#1XC^%w*2u-c}' + 'Vk70Q8s^i+E|2~)jeGG0$SbUY=Hl}~V0q?7cMw^a^UWX{INHxAXThdNvojv2^bD_fjdYRc#as6r#9rsd9McE-^*D^TOt0Ie' + 'vrhDo9Aw*}UX1Nf+^=1pwace$llS2M{)M+0{{74Jf2{dNjQ0Oq&-G#Rx$x*1t?v$H)#uI$`+TgzJ4ArH#+jpeuqk|Q|2pGlP`?gu' + ';d_gEf5$Pla?_i>8J$C?obw+f^%2s4btOa3BnLC7IHCV%U&-eTBy?#iA!kp~^KpP1F;|D4=~zc^e2j(e$Ho#y7`vJL7f?$B1QY-O' + '00;m803iUjZormJ0ssK*2LJ#h0001HVRLkFY;AKdZEs{{Y;!MVb8TjCY-BPoUtei%X>?y-E^v9RRZ(x#Fc5yvuQ>daNTClrRvOY&' + 'GC^%avpT' + '@Xl&xfEy*4!%YYwL6B=z0Ht|lD~*)ySb#~XEUN>}4%0XUU$>CvhZZWU1U4AEExFvZfYtY_gw8%Bvriu{-^Xw+w?V-3zIhse`+V>u' + 'RWj$BO8c}FO_BR^?%3LM&pZm;LHd>sk&}eN2|36yuK`^OZ+H&FqOxUWyA+LOjDz}C@tmjrge_(0W9qM^O72S3yg+H0xKMhDKP!|e' + '&idh<+$ROZ$MUSyEy|OVJb@l&I7$686dzYPrYsum4tCCWT5kr|Xy6>?)mUmbF)3ej!R@w9WDqmz2Ee)@tr@%jgMcRAt`=WbbiR1^' + '^>ea>Gr0GZq(1^=3go~Pk+>f)Q0imHi!#P>u!jeD`=fA#`H6=97wyZ6d_iF&9oUC9Yi1nVy*wVdcj+ANAqavDbGP|7==QwTl!lTO' + 'I0a*M1h3!L>$+La&4Q=3*Im55t2lAtKCrptZoPNm+`y%hIDMQV&A350clr3bvY2SCbeNMS2JWd3qx6-2fFEuOqypAr;lB1c4{1{U' + 'Vxp=A?evGe=P2x`h@qv8V)!R@Bh;D+6Gqpy?Zz>AM(@+;UPHMvw1vueT_j%{!HvCkg}Rw`;l@9|p>Y*?B#+<^)P}vXMUkVWlnLRU' + '+%wNNfCPU5P)h>@6aWAK2mk;8AplRa@u^e-002V<001Qb003rTb98WQZF4VeZ)9a`b1!9cZDwz5WHK*dbaZ8IbZKvHE^v8$R84E#' + 'Fc7`_R}8*%H+X$WZ(AA&B@2aa2~FuG1R>kwjaXSS8o5cKKfWV>yZhn8uB>_UX6EV1aUB12(okt#g^o74t!n92rV*s~XuPrp2sP&I' + 'ied2ZwBK1>+EC&@p4+td9J*sFJC' + '5;&k23^sys8Y_;q^`a=swImw*UgoXN)2+7=@WUhhu#HFXuLxn`U@q58qx;9Npp3UIQd?PI(aBTU0(GZNU=-BUVGI&ut_TXEDd01c' + 'H~1fnidvqJlp5dLi7$mEy>^P+GGZq+U4eXDLP}Pp*44brrL4sGRf}zPLU}+JC(}k~d;UH&+mZE=b*D*zC0j-r<%L+W*~*07VM=vK' + ';6}Ow@!!qyEv(kCxr6U)w(ZEw!=PKO(~;m(>X^|1Qlu*HW}o!G;Gv##)-GKDPHlxT#SV3W6eda(SO?N1@Z-TTt91+' + '_8Vnp*Q)QDK0$m+iSK+N4bL{=`|g1GfWJ7B&&bpHR}W;q=SFv?ntlraNr#i0;ns(NY`JIScL63X^($B' + 'Kdu(8sFh&pje+dqyhCX=dnDQu7hL!`KL7cQU!s3dO9KQH000080000X09m_hH1rJs0K_o>05Sjo0A^uxbZ~5Kb1!XgWMyn~FJ*IW' + 'W^Zg{GB0CqZf0p`b#h^JX>V>{Wpiz2Z){{TE^v9JT5VGsITHS!s{99K_5*gxI6w%SBgNHu3Av>XND7nP+M*~$o{?=G^L8US4&n0O' + 'uUpb+G~*fD;X)E)O{>-2diqV8U&v{eKNMm%SL9&t;E4PtviV{%B%?ScdRUTzOJ02C(eU}RU*Ipf5Ft+`kH{j8ctO;hlTpq>_^?j~' + 'sT%*3Hi(N$R>&bIfDr~e4F}&gefaD' + '{+TP4AHRCFTrP)93p~t<*{j%)%2yZXr)O8=vt1C{PWqh2T%rN~vk(O+Jb5522cbeX0a0VNBw0b&tl;oWWhnHr5K5%80g>5MEm^^l' + 'O(di$#AKmrHY-I0&3yu}8cRtq8k6&JKzBs5vG|DN#;BhQxSq{>1@GfoXoPXyhuSKlJg=F5_eUCfDy8ZMIsbadT2#2Sm^n)AP7#gbSh?xLemimR!^8-' + 'u;4LgxJ7A%a`HkXJXI_vXXDEcSsWb@pbhWyOr%N&RpzLN%x!4FGgJj_eHe' + 'kfWxSc$7X2$jw8}|6s*{giI>OS`{pnphE&MABMhybfHA-*85a&AQdT(j1?S%DR9fP6=6>r2w#)R@zfXMJFu0=V*QW' + 'p%K)FO}hEx;aJMHP3)&!3@(!Y*HP*xFPhDPgIYa;bJPk-}e7JtQ<' + 'P;Iu}UlAc%!C4C(Xdx@P{jbnP5`os)e*e2Syj-24N(0S7{|c7d0B*OxR#S3vQs?Gxt+B7;{xe{M2#Vnvla9mN{Wp4d+UKz' + 'ILB=^q773pB2GJ1?P-O`lE!1z1*aAcclBAGX+ZYJ+JqKhH8J<2Aw(Na>tOUK?82P@tpY8kORf_Ve(l*^tqp1nlm8OTDODtc0W=ll' + '>=NE~q(b>0up*=&lY=$9ZCKZPpp#2nK8C^wqZ%Xx^5qLg?k`_R$`bB_94soNNTtGJsaF?rpv7x|KnLFjVV}G>A-y0N^cw~?@dspH' + 'a1(r1?H{<5J!epCq*kjcxs^NfPFH&jt%8{rF#Nc)bKJ1&k6}<>qVEY`qXgj92X2Ehz`8sNG~4CWrahc(D^^W+sh;34pm{?^_AHAP' + 'Pl)y236zQj{LU~93k4;5phA7l(FoOTz`b_S$)A=LEZ9bwI)_C<-!?Gr-YgN(rUcS`N`|X^y(-9uSKKYY;6~l}dAul!6b*SA)t8UF`ewYC}' + '+-)?3jXy1>Fhx1k^@>1aCv^CWh>=6D0-#`6NRj0nTQX1{YFPAFL=`p>&~yD85KybM{vQ)rOXJT!Ty6gRnf@Rnl_jGRyG=#L5J^F@' + '-K9vQYzY{LbYB5ohtoNq@DbrP--Q}>2@PW;;ZoQXnGXyqvTCgl2sp;Z!*H?2df' + 'awSK+os9c?@JEZ;4yyNJI$fZRF0*J6^Cqm`mvGU&E|k@|nwLilVU`tuIZV(|$pySeEwkd@N0St@gi&0%Q^+G{$Z#v@h1)?dq~AMi' + '|K`Y%NGj5Np}L^=X5ST>3y1{;1rOT96gqn7!e~Jg_N@c<2X#`}OBA)TygYn=#|iWUO{x7L8w+pN?fq3MYoj3?16VV+QH#A5y{0saK&t)x`4gT^UgKu&XNg_B3`Vhg-^+pGd$lR' + 'OA&$Jp$F=~4i2B;DN}pw-VV+hA(qi`TW)F2g^ntO3jE=ug|_cj$AnG+(a;2@LQ&e2Jf4~aT<@E3eCP{-9c5c7utGQ>T@tJfaAUwk' + 'h}QfawzX;Qo5E&xw6RF1yx@4=il=lK7Nu@8{eZ{Q3^VzxX8Lj|Z1-T+aeH*_G;Aw%2{defW(5N$h5BJreAI_+6h-7T&So^>EbYSR' + '?*8zgW13V^XUHDDv6E&jf$+Z9?E*)8b<#wn-CXzgel7h^iy@(O=Ih-{qe=06J?*}ub=zA;-L9mx*FL3N{6oHbxH2Q}qJG=L9>Xf?' + '(y_xjvpvbvn%G?HZB60<>4Ijv&*{ezD%9p{)ze+TGAe=!xEjNh7D*0HVTx!?Y9hMtZ9K)x2LR|K%ikie#S97Sr4>%7_X' + 'i6g#-^ytthWm}N~W+-}aXC>*N0DxRe`KBWMxsWh5Ih1NU9=D`aH@8AgGeE^GnMAB@a<5FPuBvz{P-}9oMQM$4&>~bPJ}<|kGm}f2' + '*!-9)(;L}hh3`sJVDnvQ4`(SJ7OYLb>rwfEL6_&+r$n;NJ=R=6d5~_XBG;ERjf$pw-EvhR(4hjjj5<R~KhlPW@Tymz-U=;f`H=0uD#B?74' + 'zeo{%(?J#I9<&mLb1aoCjZXqe0XGYWp=Ym+S7}u;Jb6T#x1L}^;tryFQd{^Vg^3xCszMz+!wSVcff5zsWc8dz3p`w=uDWPtmN0ow' + '^KgY5@-f#I##Hr3T1~xssDplnx29KT$c1v^BD^t$_of!xb1b6;X9mhrEmH)rMt)o`u<%3=5bv@{QN5QivWL3y$GW5' + 'cLcMNzTls5vST^#pxJi(Wp@T-2I4$A>0T5+R?EdEVbRxwioPf' + '@7y(e%gY8UZXz)5t*qH^<>|ZBdZ4)}RggKsE6&yR9RVUlzfIF`0ak-o?QXGpH|Gi0v<=1%Q78$3X9wTw!B2dz2R;#B)BG*`7PNu1' + '!djg#pl-Mw?7Yh&tzkWnZh^eUJz)cg{rU}G-O}2s)uoiy!^?R?TUQQrUDeVmq=&b&;sfU0xoM-ls~j&6rg{z#t%;8|Qf' + '5H%k+BGKw|w4`C66JJBm&Sif>(VDa?T9W0Qr;u;6uTWe>HJzQ28~XN|JKeq)EANrCC~uv`7@$5uUp02>-Pw>lj' + 'qb5v+bJDWHI|xQ0*L`EFV?MUZ3~s3b@z^DJyWe(g3dJ`=@pV_soVfM9CTdp*A?WFTY&8Mv@=w3sMa^NX3^SdQn^Cl#5ThvoZo*r`' + 'X$Nzcxy?+&0v*r)s)KLcE{A)6J6gwfOJRL657wFh98Win5TVhyv5?#Dsy`DJx()DPB)F?@Yt`E*=zv71s#DDR`zW5jfZ(}S5#O2a' + 'HD6mP%kq++uU%Cc(4Q&|%!;=-6?+2nns;xuQZUpTBmKULnQG5mD2roCfy1qBH0yPEV7uC&(Z{;+W|!<9_`R_k+_DYpqQ_eH`Lny{' + '&;AEcO9KQH000080000X01FV{W)B7c05}l<03`qb0A^uxbZ~5Kb1!XgWMyn~FJ*IWW^Zg{GB0IqVr67xX>MmOaCx0qS#R4o5PtWs' + ';A$Ug6|%5tyNkeu7f2Hq4Yo(NNrPeygP|qbWy`@NP' + 'iKIDtO_V-)`%WQRtZ1Gyu|WqX$7vKrqtRNH1;Th;>snHb5i6=vYD9!6wKb+jBmYd7GTXWX`nob!yx>4wC+Lt9rOt`;3P!h)pHZPo' + '>F-L@w1uG>jd}Cr717x?>3)Bts^&Uz$Z<|IQB-?FR26IEI?e#4CPFI=kDfO+Pv?Jsn!i7sk5`TL}McK+e;;MHg}asv3q' + '$46*}ZgqJ{#T0^>(<>WB`(J|WbNGR%s0#{VrsL37d^f`92XBv$@W=V-@zJOGDVPJCH{d(f1Lehw>GQp1j3($dN}~it6Z7>Wx&uTx' + 'T_a3IR^}8dhybTJ+Y%OJ;IFwP`Zh-kzqunU!fqDq#ULWn@X@|vcRzf^pK%Pb~H3(dzS9vuV|g+b(6u>R;HC|im7JO-I$8LhOkKx?rf5S6?|YyE>=Wp7Tv{z^PF5$2`3`LIo{th^_F3%Jxx_KP=(ljp_y$*px@RicaGr32kd;f9v-&^2N*fa' + 'uBUBtS>!C!3*di(%9Sy{T$-wFlM&T|>1c`~AMGgVL`~FfS2H@cGBV&vUyF&1*Et3GB68(7vFsSKTSd#>DSFkO3K{w8!)edDBkBD|' + '6|qhaAG?bom#ko#N;Q;=QcN0HHeP)V5AV)Vz;qR#5tp0Nbj>2#_1VY0p-UZxB;aVo{w1sRKYKXfF!pB-@R@|HCz4gn}' + '0r8UL*4*q&c9Yw7WV?x&>AiS#D~W9tw3CU6F1pKJckxZ1Q?UU8%}_zE#s#Uy@E~dKiftl!duj!GARjkkBVD!9oPJ{&wHpt_#Qo|*' + 'C8}CumMf6kl{8(~oQKphj2Qrn!$!MflMLiJtiyEX?DJA=22u-Y=Kdr**c$I_s4e0fDogbfkOLSi816}U^`keN*cO&}tE6KP(aNL$WjEF3y3d!Ja&rqh#k)=+CHD!4Dq)' + ')&#!%%%1eWXVPF2hNC64|9RFz3$4!{zUuR|f7|X|r<;qxTeTmGHbrjuJvs~Bn|ZL$SsB{3H7it`$H6WIo6`iADB6ug|0)pfbUsiI' + 'L+wz;=7k=cR}*aJRkZt^Dq4>I1yD-^1QY-O00;m803iU9yAY675dZ+OPXGWd0001HVRLkFY;AKdZEs{{Y;!MVb8TjCY-BPoWpiz2' + 'Z){{TUtwluX>MgRWG--dy&P+g+s5&`e#Q0*6y%vwPLh!}LS-avV)RRjHa3a^2Z5lZ)gABnkV(q-NZjAv*_q|OxTK`>wT3NI5ip!TGUNaY~+lHr9~4mkQ9D}=j`ug!D+nsQ*n@0!Q|+0' + 'WZW|lhb>9cqT!}_L5wkfPv@Hxp2yk1BleB?!4Rz#t3gV4beefV@pDp=pLY}iBE+_Ho{^QvS9L|oRQEXMn?PMh>I2g>({rlu7p' + '^6Ep@u%$#H0Mv1o-g8L|F`=d}8Qo^@*p)H=x*f1mkOXrQhJI>^gr>H_0A?Ckq|Y3tiECaJu8I?{L10mw0}DaMca2SHB(~jcbLvw5Ey>fY0bKj4+#WWxZWT|m@h+|E#-xV0GFC8oh*$!;nSlBx&x-p)' + 'u?B5J(gfec+#1HNVeA^lt~+CqWIedmTv|@p(D#ATbbUo6fRNS?Yz%~mt4DG5B(9#s)w8&IjyzKy*&8|KL0sXKXEo)?H|0rTc~(emOk^XUmH$V13kr1hF(P=8uZ9kE0?Dz' + '!}ujS@Pd^;bG?tOvZ3(C#+HmnedZ|)ufbzPY`R#rdIF~RHp%09bDO8zhD7riCy#0EQO9oGk57ThmK}9>%ZoBkiv&ovJ%a>xc{vLw^&txc*ojcnIqVLI@Q!nj`Sg#WyB=H3}B(j@W{c5)Y^VpAKNM*?^BRQxi436WHR@^H^%}`B~vJ1%3mLxt82SFa!J;!Vlf(' + '!`!0R=6F6^=}Dsz#%JsoY0(Tl%5#t$7yw@Fk_|^!lqXGHrWg9' + '#}EvVQV15Y%S&=}@%=+m-PEY}B4%8O=L66oqTQDHSHs4w9N%ytl;?3A)bJR>zpW89y3~N0?7=!bsrfSg5L3I`c4#8f?auoEUcezR' + 'sbg9gJ+@uWpYeWoyjV8LjPflDwl2#I6zzSI)!YK}PZ;h1*_04$*IWWAn_CXdd}#4|#0b_ySk3tCiOGJSJ;~udXF4QQR5Y^Ibq|WJ' + '+_%D{xpK#woe*CPk^v@-=!J>KvI7+F4nxHuePAp8=>W#N7Vc)Np^Jqaglf$j*haDN5wX-@+)RxWac2cktS(NN5;SrQ`#nP$GKiha' + 'ieGP%EE~hdTj`~0O1aLLm+qeQa$}Ka@&|wvh;U@A0+_?w7KeEH^{6CFA3AAXOa&HstOt|Iz`=kn*|iuI;6^F9NHvA+RbN4&t+kZF' + 'nzZPvLf0#Z{U5YnLggcVP{tmR?LnP@I$!|8t;8RvYL`HUH7(LcSyeI24UUY+*D+68c95YLkop=~IzozU6auJ%yaOk7}VX@_sU5PpX-9meP4Y?)z-IBfaA^exfg9)ScRQ%bQG~6biu`H8jBK!wbK!_0Y' + '+lbu&uZW5LO2q^NdH7{=VF3;Yq@+a~p=ghZ$Vfu8oWw+)6sCws;f&EsV-;yUKrr|w$C|Hwr3vzGNjp%RVdY3V}nJF*c~K@r&Mcjqd@2+YOu4(6(o!_1sp|^-5?@so@*ar2Y@qa>b@dQw~!BV@4NP(17zb9z2Hpe2yHf_XsC@%1I>q4h6guw^{g|bv(K$X&=bdCIJ%>(|t25a7+7_y8-|lP`QpccO' + 'X$bboWdxT-B&PjA#OClfGy_*}j%Mzafz2~eIKr*qFB|pkc%uksZGNOlezPeb$TC#~d(zgQv@i%)Yg9~Z=clGFBB!3*H5M(y1U~Ff' + 'YX|x9$kdo`$k7S>r7E2l`a=)_k>mS+%xAJ-G#K^+99L>XvAr06UT807oEF)Okrdd!gScLlxA2aR>mlC%&cRlqDaCPzVwWWiKn_H?' + 'FQ8rJhyQ0b6M2IPezw58fu^0CJQz&csRLyS1XN&@6v9$Hq+8gi-%&zw;YmlVBtug=kYkW8nIidL2FZv(NxUEQx$&~GD2wszRl{XB' + 'P@wIhTrq#*Vx@WmWHnrFtIsKa2rh#u)(2v&3*ZV#i|&g`MT;^9(@iBvssrWh-U#H^hOnjY*' + 'ud|sOqkHNEu3@r6IN1*JgmItZWQ{A=IS^t2=ndQbp*x0~kSM$nT#C=>RD8~6(UGBJ96BBzIyRu5G;};ffG1LR#bt{uG`oGGMSEsm' + 'gCeJT(A5e|5c3cPE00wB9$n3l>nqQrDRDeK3!p*Iw3nVKt}Kr?B=F8=y>UL(HGRHJ%' + 'y6Egp$&x&{^u0M1yG+OdVi0Ysh8ypV-Q%9($_!?g1s' + '%^Vre+^QZMEd_qfD0h}QCVf#IKI7;F4z`ziB=b(KX=c<2G&g2~jbgBmr*@NSY9|U*=cU#-<+R+|hY^^Fm{VIAO5+Jrn(wKy#bCdv' + '!^#SdKCfm~Xp84R`iBEG#4st(fWd%NX*h>N;s)p%JKLI2U*Nnd@N-2@Ss>aC3+v*~9nir1>rp' + 'uXY#&;lB1+5CZzjM?ny%?yY7d=v12J>!dqjLBgyBD|od{G`M$x>^nj{O(ppwai5fi~Go#27`wB?n*iq' + 'MW+ztDrGBvUn-4pweHc*L*r0T)X&pIDW|GnbtVlex2%8Js@*-OsGm1mqa5}B)o#-hwE2kB+n#%R+mrl|1weP$A(kq{hBSw(fxYXA' + 'ge?v7u%obSVIMCYwuhqn&b4`Wx)aT`=IyN52}hk4Jjt6XY&qh{E!Y@|fXZYXB7X*fppHgNEYuEYcbblbOfuw{&fzA9LfPD~AmxgWFTVF-w8wP5X`efaI!pUB&zW%uOZPfd}p5qE{Mf^W!g$}J-%I>' + 'NXmG=@DpO7$4veGWDikxla1N@%Xo{8J@y$-H}xe7Uejq9&9e+$S%nEgG95mZmwI1^lr;SQ~`j^t6>-3f;i$%Iz-)pg0TK!C%zlerBIhf)$0)$g3#u{' + '%QIh#7*Jp89BP#2ixC4Y?-SkE(k!-`n$uVLKe@Mk=E8QUaibxNGNeIa_pfQ&$gmCs8P1dXBi-^0_9dvLVzV{q9woNe@M?-+r$K|J' + 'SR=sz%^Zw_d~VXkA$v4d*}`L5P-dvvG@dG6U`MsIzcQgU#O=I|-YfUk<-d~LUr*d8?Gmd=X6_F>Mk{v*8Uj1|V!u(ub+*}Mmjq5*' + 'TmxyI)b|pnxOCrI5`CUnl$B1}rI^386Rb@rVB-rkod&RB=|cH#OSg{>vyBF`=eV*4(s#h(n8hRtjqpO Hm>7H}0}s`T{N*OsdC' + 'j>303`_y?i8ko*=gg}q&a%uPc{Q!F{`BKELL`@04?zv#cjL1f(Z9#gwo*p&#*ho3+SZqFS=T0HnI^j2sKT^F*cal%4i;1@a' + '8NSJh_lD2?4p+l0{8;g7*`+od#PK)Ht>oL+l~Wx#lNhVYCDl2Hye~iGkmix3FmwG@uVF0qgtAacmKimH?k#;gt#_hJ-4YP8KQMU-' + 'xrvMP*s&gcX+APGr1))uHEda_mI3_|QCZX{6@0-3ZxZ+?&hfMpt@v^l4qFo^v)&NA)|X4Rm#t!KGN3dF>CCj;pY-VZKU(IZvQNLg' + '1j}C~H(UcCqa5#25WO)&TP|h4uIMK!rlB~OWO(>?3OWc+H1q(VvPe-Go-sqh9R>N8f`|C' + '*u|udo3E*y0OHwPSLhsPkv(<_G8a|$b6|EH&$YHJ+il1Jr@hve&|EL?AAHE@?^_%`b4x47|wvD~6w;al#gZ=1QY-O00;m803iSi2^P+}0{{Ss3IG5s' + '0001HVRLkFY;AKdZEs{{Y;!MVb8TjCY-BPoWpiz2Z){{TUtw%%XKrP3E^v93R!xiKHW0o0S13L#Hh9gla|kqa$dF_%OF~$NT$)C!' + 'w#(CEWXUVZB%NQMl4MK%m}ZC%-L_Qq{OZ+*oagyh{=!?r1b>k|NJ~1UiLrWY?5KgX%(6x>W5CpTo@ZI7RZl4GMz>8V;r*brB}__X' + 'ndM5F$!6g>8$EFlZACN;f;AA#>}}xkWwv;fw~GAYy#8XkB7bUSRii{jo^J!-QkG4Ia_om2VTQei4|VwFu|R(d+By(Q' + 'qj0oydKpR&QIb>38+$_ZDu{VP!OrJ%aQzuW8bI5de`=v4%&Pw3M>;6OeQF8=qASVApHQ}1woGewJCM4r-vQ{Ed40#%SyL)NQkpC=' + '_{k)Tsy&nAoXsLnt9eB(YLx9UE|uuzQ5&Un%O%5{T$pXHMO62YlHuPMQXR6P0WD-Gg_;>^{a*5|x8hW!#bShG@MzWOQ0q~e?<2s!' + 'P<)Dn%#~?qsRIbo^yUKEy;k^-f>D6E_<>i~@sp~%d5ktg(ek&BsXfcUU9_9*TO1Xf;gFL{Tx(*N>qLowP%' + 'R^g^@)Nm`x&GzX~k%z+>t)7@Luup<(KBbED~-TTU=K8dsu0!jRxujb!L-G;{(QS9Nb$`J(v4K}07^4Yt-)2_^aA=pO@}GMRbF88<*4aqepOg5R{ps?R6&aG1@4^`mjK' + 'PFlt>L043k^?(U8yd42IMrrdh+o-bWMT4aGtvNB=DE6+~6v;B)jZUi=-$ebmUyV+&Ro&a~wC~)r)jpTys#NOfXn6P)&oVq25cAHN' + 'QLI?+C9K|wYQ66X!#~_u<9Q!r1j%m8bmHg!YHpz~%6p^U^No|c6N=d+`qak>Gqn>N266L_)xLF}Jy?27h}e=}11L+M3d9YT(8^2*' + 'GGH~jPZAd=Sti?eHqN2g*Y29F|Ih49&X{RG;kM_+8yVc49H;5k=XJi8$BEbT;qxL<>)97`93pe0c2#M0kU7vS>*QHC!&sDc7^~%~' + ';FlDGD%~gq75o*7JCRiWaA)#9P>dB3C$ifrnZEC0TWukZ+l4aw4^T@31QY-O00;m803iT2_hkvd0RRA)0{{Rq0001HVRLkFY;AKd' + 'ZEs{{Y;!MVb8TjCY-BPoWpiz2Z){{TUtx4_ZC`0_WMz0|axQRrZBk8Z+%OQm>sJgu*|6(Fdnhcp5PIq*1ln6ETG<{aYV~2Hak7T~' + '_>Szg97A-lrTKVz?`evncvKe8t%}>H!qo+K~t4K#{j+b<+{X;71HJ*F{kXF$8Zx%3-1@pp>BP=mT+w^Q1_77wCIxqn<*PE1;ij>iKiu@uCl!auc<)))JLD6fQe(ytBSW&GIJq`?`ohB*Eh-NtL;xFliPT87E%AA;sy(Hn)>T6o5uijs$' + '*&&W4+2kgJ9t#+obNoGq56ZY;AFFbU9^`ON6WXjOk4IcIz&|iL(l6fruHgKYXRda=3PH`grdMSqeuL%%2IWz2m3sp6*?lmU<%;&6' + '$zo))sjkF-P)h>@6aWAK2mk;8Apql2SzMbC008+u001ul003rTb98WQZF4VeZ)9a`b1!9cZDwz5WHK*hb8TjCY-BQDV{dL|Z*pyA' + 'a&s^b`AYwbm?lQ%#$g&)%;z%w@$*xPKLO@_Bq67jg9I|Lf|9!iA26Nz`sKi?vJ=i2L)6>V-J>5M+$8p|7' + 'AGR!tV_sFPF4*CUuhTfn*kQ5W6gkgp_9~6*w8*1!%ij&(hs(1~rsv>KUJ(gCC-*2)sPHR|t^4o7&5hY<9C%^s1>U7P&' + 'XDLWj*03I{S1@{ABuxhE@DgcOEH7A6yw7nOD@wLZ;iXVaFmW7bGFg@&P#7*76gdnfJzzL5>IiwLCKLIrsM=4PG``Jv`>le3?dO`W' + 'H@JwP8y8u|NzqjhE#j8=NzF^pPT3C1Ykps6>7sovpRA)iy5VwE64gn8@XC@a|5%MQOv_tqo}Ty_$_R8J~%c2`9NoRtt2^q9UlMkGCVvuJU)^pSw>a8$<|d+@|(1(0jF~8' + 'd6T2%2`WCj;dKZ$=ULc}lk)>!t(QfXya;=%ogZCPJcZ)KRf;T==3uD_2B&E;N_bmHfPS!#m|SY&)&SE' + 'L2nN(j~y}o;Ns%wGCX;OWAs{a0+^copiJWeGaq<9zbqT>P3RE|&k^8vJRuoo69#{1oSCo*eyz)rZ' + 'SY27VyCxF+KIWU6{lvFNWeH9SDC77+FHlBl#anS`kgML|-_pUjwTa?ea2O>Pcs|Dn03SGL9wM6LMKb)+;^njQA#fB`j}eN94|m3RG#R;86fd_&xBGZAn9|AQG=^v8{PEf?5(d5Z)jP^#lO9%Z&*jE9P--CP1*hjtWgVTc?ng%V6dNVm}@rZl|7vagpo710MElp4fFZ48TFnHK1xWS?&W)>`kD8EKo' + '!}~H2B2k6~pq5qJHKYi)T~QP{TZXU{kuNm4X>A%JA=^JRP2YOnqR@UiETZ3Wik|5W7A@O@(ojsL<@2K-cWe1^wUNpoHZ~BAXoK$Q' + 'jfhdwgq&5uS$O{H?DWk$a2$xPQC*jA0r`}37QQ?=Jve{o`0Sb2e|h-k?C3v7hxFDU`TM*(@bnD!wl`=0)*};M#bvqyf1I%tGDUm0' + '>3NYJE*M5_x`s-9v2+g+O8Zr?B&BOYo{8I@3C#HeM)*{QFj(kVJ>hpgvmd2-;o)vzqqoZ`PTV;F$QWpSf=rNt=Qj}b-mfsfg!ywtncir9Fe-~W#3-7Ve3L;!zs713@Q~(H`rZQSeVKxP' + '$5cnLs##Q!SvEEAP-^nQ9h*f<$O+gT@S707f$a-?E%*M8xvf#c8PwQ56yjBr-vSB2F%%k6;;z%&eg5plH|z_B|M_~p=k*B4@I-WG' + 'gSc1#(c9jeR=EKw$))*y>RQU_9E$3EXk1)uFm03GqUhE<%UpXTUn($iv@t{POoorW6|MR#jUqrl0cfa' + '_lf*;M<5i-e*y=NYHV@9Ci2me|8ZYd=2>P7elOp8Be9YQS+SDleu8LuaeM%T' + 'RJ{Bi`#eoW?U{nbWq}PTh&IDUaAXEIZDRpu(GcDC*X)q;wp2s3Lqs+Ng+`WbtpwTF+ftFF_Y-5-Bjc>7xRtEP0({TZ_UUC+w}z5c' + '6Mi&`dRm+nxh#ZqBtLEzx@p+iK7h{5s5FC*r~k}4kO+RpZnMq>S0HlNO`<;' + 'p`O0Kfu@~ZGtwk4mBTpsi*9TIvs0T' + 'g^W?#kF+I4T@UUAdFo#Zjh2}ELUCAj6NJ*sMT6OHY>kHPMOG|Yz`8?-#Q~MGBk)^#SL8-iYjS7}Q4=i_XjOg+Z4(SzZKri*Mg?47' + '=zA|!8zq1e2jV-XXCk(nPo)RA-d;gnv0LNW+K#40yM&JK)bec~LvkX|9NIHLbp&_;BFprK@|O@&^*RjQif2noFzoaP9|!9q<8hPm' + 'YQ_=(z!l)Y#{;A5fi97nq|X=4S^87H2{6kmqwV!I5cv3DPsd|KYi5>K^;O+$GA_Q;KTsfF4^F2;$w#IS&y?qwVUGd8#TCsOn1m-y' + 'APfXb1Qj0q$(dqBU{kDUiY=8l$U}Mq^hSs&_Q#1S^J3?KOzX?6E=fz(k_Q^L0m=Thg{@rK2E%y1EYWBWpLDyp*ck2}U-be7Zh}Zg' + '<=J9JR>agGaFtZPqU-ezsFXvhHhM33iN3_e(cv%WCztPL6p(wT4*f=z^2zY_xl}>!M^!;?|iNl_E>LS1V@O^P=F`dK(ByWS7=3fl1P9QgSDN;j#~y=q1`!0^dDyP{Y_#l*cJ^Gh9Z1N30Ybf' + 'vQXuVE#Lv7EYTY{p=nbnA)Da_WKpVVz~C4YAafR@MC2Q6?p|&;ZJ}~$@j}sGo_N6)5UX1piMquR7+uU`I0I%dl_K>mRV+' + 'n#OBhuZl$P5SHh;yTVAPKz)L<6Hg;v0`hI1-=!rKv#{N}qZbmaa|umh);~2f' + 's!}<~vx<1_KN*s0WgU;@Hq~~#P1LMstPqP0z**<8ZaW(eTkRTyEuzl-lHoyZB^sv5e4v)~M<-!Yg_oP=mK6SuMMTXu?^SWbbG%3T' + 'SbN5@UB}HM4LT##iB1tlyQ7PrU*okE9clu00R|Xrk$xP-PG^ErDqn3+*5llO?s' + 'z$9kk)&beFYmzqU8VBlST9kNo2(nHGqjm$TcV_IF&zzHlM;Uy6ZsoS3osi}UzlTY9?9$m0^BdSQMKe(@=ygnW%eV6^S}&4_5uKT0' + 'Lu#*`ctx+Tk^Q%2v8fI!y#6Sa?4ZLtp3h!*ez%=Kut$WJ$!%J<_s4K|((hMLCTv5)iugJ>=j-B*A6H6ukf}#dU_l}{e;q+IwCriq' + '_@fr=8PHO%t#f~mHK*)(&l~JnLsS$Ax5J0wFOU0&U4R%RqTGOI_C~0ux1Ho}>1pXs^RGsz`eM|ThrN-|-zm>}JX?j`fDc<~qh@Jz' + 'FD6`KKLl^%g0QCk53pQyz2WmnoDd@5u`>jpce=zN1=;X0@>)gH7nUnTH(x!Tjc%W~Yo>X_2M25lynQk^t_Mru-b!t|tpF##tvQPMvIsU6}yERn+2UG23Ky3$?MSU7wS)K(9JbqZg4#C69|TsdhTj' + 'c*o>P-Xt0!BaRe)5>JxS%yzLnXIIzzIMZLrWp=m=y8Ir#A~h348@!9}nq*q8F1G7Mk-46J@|tewAJPSVaFGG{l1+Qam5Ytwc3{rF' + 'ki~A?ABF4_(JC5$|2MS$DWK91EO6HtMWl**v&TXC818$EJPdjcDzr=CB`SH$bCzjYA!tY%VhBmlG!C@^uD&C5SD@TH8JvEUCyz|F' + 'U5*$r+)Xe>O*cJAtY?f85BB%Sb`_d7)%LDx6zpe;37R`_iRFJ7>{GF;2FFo_?&LbA' + 'h>i}{t-tjKK%VYAcg8_8u#qS??JX{!)9oq88KsFS+0MN~={AUx#Klho=(LR4Ln%!k>>ciyHIlwKg!WRk)zs*enw~11dFZy4c-{TV|$lo' + 'JjI$E4<*QKc)sLDISzq$Em`-XT+$eY8+}05%1BY0_}PZJP>2%c_`w_Y*qXQ*PV31F6=gWN*SmvX9OUjcAn9BrT9Is`1$WtEUmJkk&tB7s$)^fkd5*l_m9>FS_GnlZFGHThOA6!' + 'AKpZevW;;jNBMYaX-pm&ak{NZG;%Xk>VRu_g?n+R5V+nD3qFD!cI%YMW#-|oce' + '$UrjIz`OC4+N=S-Ur+uEP)h>@6aWAK2mk;8ApojkS~ffh0068P001rk003rTb98WQZF4VeZ)9a`b1!9cZDwz5WHK*hb8TjCY-BQD' + 'V{dMAbYX6Eb1rasrC3pO+DH<9*RSZ)Jp}d$wFvJeJ5_sCLXv?pNUS8B5SL3O!`M~@BuWxavMK-f`(`AB5ZKw=xr7;dx_iFq?&)rg' + 's;VYmuER9UqbE+AWX+d5{G@F8$TuO~Cn;Ip^fFnmG`iBXfJj*Hh5@(cSX1C4ql!qHi`BRiZ' + 'g2H%3>uAa2OlV4;cpSz{-X;Gr%9KsA+O0Y9uUI3?c%fsNr1-T>;#Cws5P&T|Z<8p_-`6rb8MEcl#(%o37*1rv=T%LbBZ?GB+cnd1nFH4=}-c_kHFob6>PCi3Q~wY<)9_0' + 'vf1UZAZW0`5__c#zj;?tRaH?WmLRy_33~^@euXfOlU&+8Qxwk_SyMfjxz4p^>$YLhAM{-z{Hc341oz=4T7Od+sXW`<3n!|bf@PV+' + 'c?uJWxw|NN+}#CxmLQ6^yF7U6f%jQaloh|H;EOOw3)|jt2B>M9-u?$}%<>!jxMfml$^kbf$<)oZ>nLxvAwP8{fxnnp0;+F~ZQ%U%' + 'Ks0R~gx-PRN`iJE*p5AOCX1=%g5Fn@e@eF$J|gL1iZX;C43yj#0aJwt$-U5iw_S317Oi$1^c7+-gEH7%m^P>UOG0;i8SW4PfAKVt' + 'k($I@IvfGW(EO3|RQJLCq%Smz2$}L4P6#~*x`b1|sv&qoOAcu@W>A_;s9Wu8*' + 'Nyd;2lWdEOPh^jJ_-IE0p3Us0fMbHsZ1ByPd5+IH<}(%^GmiMYVSRo$vDu$z06G?yw6tT^IkR41b(*XTHmq3dGo!J`MolF)a2uvZ' + '#_6KtiN(eZelQQgZ1(Dm6Kl>Ur|6{OBR&}9>dY>1Jnc_)#;N7AX_L_tjCHQr$YPf-EGze8Sol6tj1XkAQ;|B0XSvKfVsVbV=8s(F' + 'H|C{Pj+opT`|uJpdk@X(1D=40@8eMP7cN`;cx&X0>QdMGgk7I=Sjp5@kaG5?(6zj=Iky7e5kr|xsw9~fiGYbrz`x4#Cr2G>l~WxG' + ');!O7ib0qso2)I4CXaU;Sr1x`!`>mSjePma+BfRdv(#Hac%8}e)!RzX+`g)V*EB^zhJ0ab$Z2`&uV~v+^`S9WJCug;18-@4fP2^O' + 'h{O$!Km2sZk#~firuuh^`D5)Z)zqO}HHY%QA^-boyISvkhx*#ddgkeReN7qavpS{8bns^1cY4BomFEiYl&Q0H*|Eoj!y0aSKGJt^' + 'i8@uQnV#?T#Y%-Tw%7hYuO#%{s2_`UusKzlqFklsGv5#ft*q$2Da|RVVuJ8X3VKo3pDw' + '?aRfKMjs5dizaU!E_q8O0YGglLU_NUYsb(BN6yHT&pZyIm0WxOqf?pg_~e}6=+H&``E-T7(fUPH;xic!b=wTa_H6D4V^f6eg);Tf' + '2k9ol^~$mf>aUV0Vyi~%;p~)QB|_UU%qkuH^2_%s$;dqLg4c^|f!(}vgu7ABH(3km;=WSy4-sJqHMz)hk*hR~72(h+AylDIJGSrJ' + 'A|hl!$kQgnK^$$>J_(7J^oZlwZ`UI0_}q&B=ah+vF-GVytV2IjHiAl3?CejhN1_iwr-6|uz^Jzv&72G7XzXn{>Q|!+F)9q9dexLy' + 'O-+Xkk4B$th6go3XLvNav@~{UVKBAXR7{jQ42R)C%_#;bG;@Xr1XT7hzOPV`QTfqD=}eRth6idV49@|IZurdk=nIOin}BZbhj!1F' + '^0NbRRuackO!wJgKl+q;CVm4bo<8Zl!Kg1;R8*$Inr5K;&eW@7uO(hr&}hbaheR@Ii=vdbmXBfjzz^+VgR4UYgX|Yiy2?gFdF;ow' + 'bnVP?$BsER0&g)LI1{hw-4Tr}nv7{+5vnZND7TlQvqXPB9g;pFZ%i}tU`Y6%u0%=p4D!cxO@<36!3JbF@B$nHXYNfj)M;pu?vV$*' + 'kaNjJ*XYuMKDlHMXzGw}kUJ-LL^G3g**@W|!NW5k-=Z;ngkp~dI!$!ZSe1%r5I1QyC2vaBoQ5X(12PvhatMjepVDGLlR1qp$=B%v' + 'xfVGdja{1gq=P&pe?qeX4X`VRbf4S`@+C0pq+@e{q|;ou(wbOP%l5^&7v6*!?fJBFY4&Y%m#mp|RJ06``R;gHdGd(bDVdwiMZ-HC' + '>J_1Yq*4lqoH9a6d9J_!!oYF`rb4R1bF?DcD-c+UVdt1bU^%zFNIgpoEu#w-%cu)2_Xe8s`WaB~VLrz3`yim(Zfr|2?gi}d0g^aq8WN}IsvM&nhig*4puhZo9f5`i;WPg^vITv@Lz%e@Gxqs1#' + ';cj5C+j4Dl_!DBogP8FXV#Yzt{|PZ)h&{_7YV5c=WbWTnRr24~4`9(BikF4zs%L+<6%@A`EqT>rqw=sS|0HS7SFMcat=HNhc@czU' + '(Qb=WB`%Bh4GH=ht=(o-$CZXr#>w5qd9>kHnkFd@Mw;Xt=+fgG8`|tLaiQ~sEOm0%gG7QwN$YmCURSO1<(kK5dQOQ-`;}y`9LwFD' + 'VJf*^$0{zQAX;I4N|xcBxJUiAha2Z1_JY-xVVuO#5-3aNtOm5vo!;~=Zi~(liA>=98!wk&?_Io(*5lZ3vf=q-1Sl~tFlLY9o=e_s' + '9cF6A=U;0;&liUWm=H%^n(TK?0M_c2q*`I*5@+wY3%^?Y(4{wEFS;O9zHEa`ZYlww=8(Bl*{C>VCe*Q5yadRnF' + '+@p)cOQ2}dP|>(p89mm?SpqEx*Xtm_0ltaDP4Rjm2e8z`I=Kti<*S1rcB5>!-SV{6zE%DUP)h>@6aWAK2mk;8App+Sw=(Yp003bN' + '001)p003rTb98WQZF4VeZ)9a`b1!9cZDwz5WHK*hb8TjCY-BQDV{dMAbYX6Eb6;h1Z8I)#d5u<4bDKC2e&<*0@r36R#&+y9=`@+a' + 'pq!Z`-~!I|n#mAiY`HI!LV$RoA;`xTy)bF*fv{QlkqJm49Uiozo;XqD7x' + 'wyf9+ScQb}N5`5oP@9&FXLZI3NQzYBC0%4m-kw*{JP|7Xl$Bg=O23wQw#aH>jNYZRr1(+PIXX{yojl}JP{_E3w3-~jmlbthE*7sY(-5?*L6*kKE&~|shf(%ac7Yf1*;Xu)~0?gRcwX&Iw>(!kF$z>eU4C}X0aXfqFq47;>-^$DR3K)B}I%sb@D=;~mmCOud)0f_l-ivn=' + 'Ka3;)l6bq>1EqVAhawjxBPM$hlP$!0rte0G' + '=&d7msw{n0opbv9okTc?x{{i>Vp2EqeImtMMpee+FK&pyWh90Te>*vS>pV=aLX!4|bmUB~rvlyg}7+;W-dJVa$_c}6r8u%gXc^i=Uu)m+;' + 'A}Lsb6ml{DWN9PXaZM`>7OWaxcT{n&*A#1HPmNuT!&D=s(USpGc2!b5QbRl=2QV7+2mdF4nAXigOjxiv2>#@liE&M_%pM!oNP0M%' + 'yt-_Z#Yvvyq=v~apB?cIb??t(&K{B+`3kp$(9C#K;;d-&?zDdaP)h>@6aWAK2mk;8Apn=AzMjtp007Gx001%o003rTb98WQZF4Ve' + 'Z)9a`b1!9cZDwz5WHK*hb8TjCY-BQDX>M?JbYEh1X>4R=axQRr)mTZ7+cp%w>sPQkS%G69IT}zAcP4F9G}9oH7AOospe1^uR$E1C' + '8prj2@8eEN)Xe14C5MUi`1tnalj1o3aOz39MnByCbn~3zBSBT$kBp*YH=x(IZ@=y8@+CSF%*KJBejq&_Fsr(D6UT8B9fz(#d43%E' + 'ah{{9>AQg;+_oJPY*g_SU0stxP&W8bXvSOe$4J_ONTw1qT-2CSLN$XSNe*`EtM*PG?%LBDaR+vFg+Y@dw`0?v5T>Z@mG%vx&5?VO' + 'Hx%pJZ*O-|1b^OehqqP^O8F7`3*B@riK6J0wL$zvZ@aQOR)z3y8zJ}yvDl%{Wayq(ZAm_n5)Fh_CIiW(^c}5)OvnKfl%goNh_XRE(59*h70=2zuyo1BcSFZW)k@#Q^{UiY>x!~H{LG|Sl>)1%ar^^)g5C8wo2VCeDBdC+%m?gx' + 'f3y2XYznBp{`XY>Sl!{ae2VtP9Gh}57oOlTuy%I?Zmznv6r~0ACk-e*;;PLnr@TtdyoJ8mv?U$}d&&$2XwGH~58>3q1yePcLl0VAeGglD_wHcJJlOpytY4`gooU_Wy4tTk@;>Q' + 'WhKU6q;1g_}5;kzk#Ky1Y}k4)d}k*(;TX8j|djqM65WGbpA^e`4!hR=k9M{64*q&Y@' + 'r<*w!6q8e0n>)f*oGe{R678&Wp`uPEg@d@F1pR{Rk-Q#;Zde`TpIbWi938+4*KH^f5r-bI6wZ#4OPr{FRSAKa=f~VsybPTZXoIH+' + '*3G5R>A+Zb6HrFrpt{Xg?lt7>DHHh?D53h8IQ7DPi}uZvUNH(0)$b~N@7?*s>Q=p5@%d>Qj?ue$yy{Jc?1fOPd&^>KmvTR-19To{~+$GL2rI#8ecN}J40?%aIJesW0r&h=i#?N==`^kIRJ9PDDXj|Kzh@3' + 'EgvEdn%MK-*(Csr($pgGSWcy+kA*&1Sz}Jpjhgw' + 'RWzM#P*6)NcJlHeOUhZ%PEE(lsDoXO_%d7R(J5sa0>$d57' + ';3Jw|k;&D{jESLcp9qyiw05X(y5cT)$%M%EgD7w*t&sYymZ4@X`zt17(|Hpat=m!d&^HB@Jt4m({Zq@fgdarHF|*ms(lJ@@o2op^' + 'J0^KKu)KiX_Ehu|#m)xJ>Ty-gWno#}a?#S_J|%-Nk5lTlYkI~-?ks7R#w!A!qN&@VXE|8+AZjN1yjLp9G?+e=Fd=i%IKq&it7SQv' + 'PFAZ`$@U~KIVw`Hyr$sU$GB}IS036Ld!~f7@UpKO+9zcFy&2vr39!1YXu*&v*^zPab(Uo*`CB7cO13v-ePjkF6`ziQ_QH5g{?zp`' + '2qiDW6!(e$94!4wdAGkz|$op`r01)l>#u(%mGY5)q&a+cZ;*5pEY`rn#' + 'E$b|#w&b>HDww_k+2D$)Q6Yrh6#z|QeM6LQ3MR?>_m@hM_wQ9u;wvfuuOE1SN8}xC8HtsNEBHxtI3iQYei|71xnkk~q&ug6CizDe' + 'gDn*yFTujRe*y|^efGrS$!f4Zdw>n?91lI9SxDS{nWG^ku!%=1^2(SnAt#}o;XU4Dr+^eIU$dgOvZ>Iotm~AeO<>ATj%WxaD!lEe?^kkg6f5eN7<(L9r7Zu8*QKBZNa(;lL{Rh$|b2CAE$rL%B;9^)e{7il2qF?WGkR0{icSs{p7ULKRblsqo9' + 'wNwcws^Sz))0JFM7&H|`r4Q7Opmx&b3HI2+6xdWn=yTE@xhwY|T@Q8AP>f)ps|ZN&$%B3AVYAz&ew}jjFUritFfJto&32WXL8uh_u()D^{4uJL;<$5r1p*kFI' + '4UI_f8$+yxRi>}tT8#k?{AQsyu}P9ST$JIEmB;w|ov}nhz9X^q7K}QGX~Qh%7tQj!oxj&tW`!ccMjIE-Tt~Ec4cq-Ai8fALw0LTouamt03e?Kwr7W3Z=WhW&h!B2^5Hwar<0Uf=ym8)5qaN_mieVo$nPW)=cDoK*e6o#t9K;xL)j5gSwjHIz}_L1!}HXrO=u4t;E7UQK>v{~Qr0pGE0KF?7mLnt;wi$}U}H}frC' + 'Z_{G}Z#*`$8u#ov9gR(U@ArI6W&yxV;>LnT-=M*wbG;@jjGd@JWsEWKz%UAJCo#=HlUSk3#x-4%nA>;o2lfy*T+YHW!aX6fk>(f+wXnln7jX7lxa3+Lu@xaC;g58yH$=G#srcW|h' + 'XjnkCt;P$D}K-7aETVT?`W}j6gbTQIvf0>o{pxKg-UMZI7r1#%E@Ho<|`1F&G5T}TQEje' + 'RB$oNA-FAD&nr+V)b)Svr%bbtVawq*#xAw4s*~l4;Y$)vw*G=_lF`u<@}TzrJo+Vi^Y@Er>j}zgMukO$OTrNpzJLzH@{czm!TF40' + 'V<2Z&06;*ZeMvffWRlp#V|TIdUBL7C{Ov(jmFQKin=TPr' + '{c})tt`Z^Z0{*L;UiM{Ei=i*8vM-eugI0yu9>uG-Uw+Y4#br91y}MT8T2-y;v?y!QUzb`OnqsJwkew2Rs>-eEWUs2D*mX*YW{2~t' + '5*om(RJ=a4P1l$89`>r4*$TgQYS#b&01Q75TGR~yyDf7CgjOm?$`X+`hqkOFe>9!w)S>ACbDk@$)7gAJpUrk%a}Zg!8~UMBStd$C' + '=p9g-HF;ARz}hY<)J_gnUzB-&wkDO^+^+khs_~G>#@9Xkwyl7TZw*83&bn9kJ+Q-Wda>1gC-eR#&wMFsxmN(89;(W$=B4A}*Xs<1o?GBA|McSZSJ`K8UcLBB_V(SI(@m&_I)mekO!#7`Nfa-fC}*=TU;Hh5`}Nh;U*EiY^?COB' + '>$mS-eD(5`Scwbu5ooKG2@IISM{_qgPsFwv>SARF;@NEW={awR>jK1{OQXEfp!nPG^iou%?$;wxn^-J9@rYcCGzEK|;pScY{4$dRP$+eu8C{DEY_u`r^Ab$7' + 'OJ)RTT36SgPn)F>zZJZ`+}B{UmIrSXSX_-qwHlR2b!86z=LqbnHTaiXCGMK;2EE7^SI;lU^OrlHDzRFL`JvTw<~9D)' + 'F(S`yqWGbmc{`l%B)v`Lp-nAQB5z+JKqA1wV)ccrw2CLsJB4Z}}08x_(oQPdewDJ7wzU>*LrBtG{8+8' + 'XUW^X_fY&ekBu4}6-OBdb5*0RW&E' + 'j5!cA&q53tL3+HyLA$clC3%rtS<-!?;){Zy1r&%PBs)-OPm*c-B;E;+-Xv' + 'wFAbjYE139!{`|WtW0nma3sPbBOuaOgHS<`HD~G}YoZ#Zbgmw_n2e41#(>|NX+hELOXEPslL*yva' + '8az4qZ;ETaPN959D$_ZGr;oKIbOV$)u^Ex@awig2$xgME1W#7hy9hI&G337C8L-A(S(lgv(h2E&8naeNDlu~vFRor=%z>N}LM|Ua' + 'h&2H9!>VQ3=b0#Oe8@s=h_c2>Qo%E@UKJV9_Bd6*+sJVrmqoBC-jGcr^OzJGBvr-pp2o?w' + '?Dh*B*bt<6HeuHxRe?QGI0)y{!iINrdZT)9O#&`!CDtH}b=zP{)6^RZ6!r@v1!;!t!3n)rU2}py-b|YMDUHOfcoI7z50`Hytoav^' + 'dje~Mg8(4GRFV1RO;{`@x}F)+e4c*Sl(k!TmglT_;ew2G;~Ns-OtXv<<1Bjw0A`q=p8F((g2Wi1K^Qk(S;;Mo9#U;1VaD#1Cv#Oo!T`wv<&dgRo!ClF@*FhmR*f#T!Yf&gN;' + 'dAf8mBXz6!8TP^6-KOxUB$EA8CMR)dv+*m`XMn@&9=b4)x}4nXL(iU' + 'O`&Xnk$Z{!73x-Dd8S*aj|#vE$E?IwYG6>V@Z_GCH3P1YtpEb<8yU|;_@6+`Sb7ps<2Tns#+TEFH%DD5H%VW~Y-H?6izTyD`dR{u' + 'qKhOJPi%;&i}E0zif93&ywmocF>yt$F$;%UgS%Gm6)c6|Wx>Yc4)cweSQu3@?PRKhA9gn' + 'kHU1XJglJ1tDB4#W$>+6&qne6k+wj7(tw{r7eU1+til#4WP^9tWrZDNbm{2Iz@-=Xlb{KC?}tmq0C-Eb9Fv<+LYH;2#^wO#Ge}uV' + 'TpdHEyf1-T4JX66si<3vmO%v_tz9pJLu~wN1m}_IEF*+vg@ZqNH1g7' + 'osxLVrMW#RK_|NA4l`WHy%6YvxPI1nHME~7?kG^-Zl1@Z{Suwy=i<_H1A&So@Ab&|3v1-ei<`R6_ML368D-w!{m2PV#u@4tP0hBl' + 'Kx24t-j-5O$WBc^Nlcd_1A3_1g9O)c`Lm&<3{wJk`$etF{q+`kOdJq{ro_TtE)^vQH4?+2)HJ(S%|Z3uab!DiW#3<;ccT*=dq%Lfxh4_A0`8zya0~h_9' + 'hvqI#Gsm);as$IDPg^g^%byKDkK+_O8xcC>#f5e8x-154EuxLs_-DKg%PO%U%YfvsI8-kueV{0Sc@};ioP}VL%B>z7(a6Kdopt?_' + 'FZ+%7SWI7g0+6bT;zuflAUA@#gxe2LmLTA}QY(+aATvJ?2PyKTMpe2AZ?@C~HpxQQxEbE=o5n`I#nb-kzBWGZ7&L}AAorYuu4bR*{61=;zAO`{oOT>;VF`JTxGsa{t5KT|V^{c$n*!Q&?>&f{EKt_p~?X6ME9FFsh7Zsd1^LB2$PIY}IqvDk;gzUCu-(fS&' + 'Ttv7;N@J56f*g1sQ0MXh30jXOA{}sa(9O`dL!a4%X6@9L;+Z|E7Z9I7xdQ;PT;NIe>A@=i%&6R65CP^skOR7=KSSGi1UO%`8$iPx' + 'GM~DUgpjkRdALIimA+vqW;XnbUzSd1>f!dYd>xMoirS+Kp(i=I89F}k`?8Ig;PIuArn0`w`_;QHxIOh?m|TjJIZFAMnI<)51SYe<' + 'e2fQzRxx3?U#@>}%er)yn^RNeb5?v>+sTh2C' + 'l)M?R-5CZ{R8|uMQ&TiP6&D}YJs$?E_*0uzau##B+>F|To{>4Vay&cLOtX7h@9gB6#<93v$bRP3&?Vmr3T2y$V+lrVh{LkRp+hb8' + 'j+UHqFH~5jV`RRua+FIgIBj5TVdKiZUju?2Vc1WmyUXti^$|E9GNCRZ!}YTNlg*FZd_tEpLo0&35-auengkLCdgJsE}V9+RkUuQDRacDmWn4UC5GQw=45f)qQmodDTFCW8BCld)p~1ynIt$Sxl+Ha6m%Px1c%&l!Xy9n3_#Dv^|K}66($E-1I!&' + 'FvpL*zzAq@J#W5c*(-im1LWC8HijX;H2Auyq59oRyr)UT9*+0D2lPjai#8X5Xh*X`EzO6nGjFFj2aQ-FiMUb6RV5GG' + 'LgIzZVM+hTH&k7Jep$eH@TC0PdU3IVwM`4A7#`LjWI&-g@K$ROW?>vG!E2PV9>!jB+d(YCTbh-Jc@{#%AV{C(`Nu&)Uqj~JUr1a%sbqilArBUQ<*szxK*cg@hEs^BB*HNH6Yl?M6eJ;J6tKd((WK2pakE{_3D' + '5(fNwoJ@&;jTJXb$D20t@kU@7dUhug4b+6%Odu0%u+zCM))ZA&2L(87wl}naik&k$BCsn99O+si3QHn@u-~BYLHpDNgO>b3=zNLH' + 'B-Dac@O65_VbqYRH)^<7eU@L#vfg$|-emN08X|(y8Aj^N#;vd{^r}J-^~NaLc)PVcc|3eAtKFh$@qQRl@l`E^6^Qt4RrZ)4)5L;^' + 'fYjp#f{@n(cSNDQ%;hE&wX$+4qhHEVZ_)|X1VhuL3FEcf5qKoK^f#;1*X+y3y$N20F+D4vyr~P%hnhFUH@yBCkVW7VnPDDZ@D3PN' + 'x!HMt>p_380sjuQx^L+bO&WG`;tvstJPQfMKNI*h#DSPWy`UP${*#Hy(-bCW&?zOMSM;mP+F%VZGBP%~u(NS*rEwF1nwS@i' + '8sC$U0*&iq@jG!zk3&x0{n=N7-com)g$ry#MyH&-sD!drzUD0`0x@Hu`xTFU1N{|W{`Y4@ahL9a32VIB^7*b%a9kUEx{' + 'FR#8HjiYcI>5k0(y+-;E-4`QNP_#H5*Yc~&^E)v_G78=N4NBzW@cJYNa5Ghtlb2i3oBVO^4twjp&hQMB<=T4hUr$ki3_MJ{gM{CIUuUS3|jzB+&N9$r7p0{9nRWNlxO&*41yOJa^3oRMFRkr$-sYErj^{=6?f!Hchx' + '+2f2{-(0`FdGq1w@*TOnzI}cBf!uz0dx_LQB)`(CqkZ%c0@RrR3zH!H>HL=ec=7K1`r?Pnm*HbPo6VRV1lRT3w(ef^}JMJ~Cs>iV8Gb2fI3_UwecJFXc{NU8}eWlh^Tk^?ndBa#LC' + 'k3j(e+9G?P&YaGf>K4ExdwEa$ZJFK8gEjaK48O|Cj`|2+?KZ{s2b$$TNDwb+)=`4?Ki6L2(g!ZVD3O7Aur2G%$D=Z7vZMWdohvSX' + 'y#B}a%}>`Ti=1wQw%%7c*m8Z`AT!rxmk@wv41I5_);q1NiaX&`QhDn$jHf+Fy-+7>yjHZ%~ZKvY3fW-zd=Ti_sDmo%WC(P4x_' + 'm^`;Zh&m!2IAaL9>v5jsB4w)Ol=8w-+jUW_i2{o_S@NckYY6;@=2ujEss`(}' + 'xC8m19jGZwM6^;MpR#ffp>d&L)2t&|+h&KTYEq`&zx_2<;2nzx2oGxUB' + 'Li8(5eXt1NLPbcM2{iVF2$>OMav?ZRY$~zM!7l4QE|etj$z+3@EZKmcfbZRdr<1bi`iL(VWy35L$CmK0l758-6D-*TBw934%%s6a' + 'lr++=DT_V|lQ2eoTi`26o9>$aAY**Q$PFS0BPPZoJD@Dkhmt}3ScDK2MbucZXrUH*&!PjtjfOC@=0ava6A^7#{4j%w9So8go9~mwCZ_$4@)C<5rUgKi$3BoG1^_mfvh*uZ' + 'S(wJyVEST&IT06RY=qsfFF@H70ZYO_f>UhzB_vphj@K)1?D7H5Eo7MIKB7d_fEB&us7f~2z`>_ok7hZnR&9}YmaYn+C%zTq9;85V' + 'qqBi=WE32?BukhlqlwnTL%D_kR1iMt6K2?XP1VJnj+K?KMJR}D8|C$z9' + '(8FBZBm9#zlx7g`s$6m%n?X&53lY}P@^_Gibx;LkannUi89Fw>BH2m>X$3bQyk=p_?}(zV@Rm9n;p|io{UQaU0u(8ai_)J7gj7<}' + '9%^uwx{^&+c~SiaUq=%oiN%c1;XSsjD+OyMi7x^rRzW2#l#d#ADv_}AKK)l+RFR2WR1C50Ae;<0*vT+pTdHDsDG{K-$kCbC9)yK<' + 'n|r#Wx_%e(76iP}K@56TVAm%&FLFW>(f@`yg>W%yFUq9ETt1TfdMp*CUpIgu4u^P?o&AtQlJ3)bgZJ6IqP8RCy8I{ceooHa+G#' + '+H?Av@X4y1YAj$b5rey1o1-3f_D7o)dy8cF`5w_w0b_0>Iu{llhLy3aIJ>hHC4b`L6TuYZ#eWG=IdS!WuA#|5EQ|_$QpqLlXb*;$' + '?aMw=%YJw?((hrhB^fs|4mkw7dGOoVy43=yqOS`;((JX^B5|19I_QE16O8vyM`VT$efZ^M``' + 'LcVi!u#_@G)02&^3DGU=$7F)Bh5cxQ${v|vwX$pt7OXqN=i(WZ=;{L*a' + 'q_nHse$Z+`-s!MYOvkD(T~<-$^z+>3VaP3eJLcaQ6lg9ZTjym{Y&4u5bxc+78Yq^jY9I*P7ki3)I?47=vuk_bk!56}' + '%UIV;T(cB>ffyx+W`Ng7%}i@t<&seBMr<`fG(QA6H;RfER#+UGp%i+S=R`NTw^>=@U70Ttxxz9J)?Iqu*SmA9og7i)(J(9^G80)o' + 'QG0ro2lrW*^?e(uw~%W|VH#&Q2&%86=vIQ>lc%w!T0tSfvxAlpfj;)0p$|J2czMIcW*>!{@JgoBS*?t63iZ;nO330=@54HJx8%z<' + 'BC&3f)`fV)WXVwqh{s<3+AGDOB+Ut~3&*IyyYUo2a451ZG;giL96fcF!O}3(Oh68;f58=j%@#7=z>^>(e`UMITEo`leD9!' + 'Aq+kDQF?1tS28&2KAzED1*usF`50bVQ%t65=4sAk4^@Qb6V7-}?4v*7Y-ez;BYqDk;' + 'tW_7e-exV{tJ9a|jQwdx`z&X79(6OF#Y$OZssuq@5d?oo7mJ)4B&{C82_w&HKVmhV>bvwwi40gqW51@6G9a!L^7z~zMg^xkJzBkUK-H(u_fek^!Hc}=6n;iy9@Rl^yl8*fE@4dz#UP_x7d9BX}V`I%jx;VqL9;a8SFFT7B' + 'LK!wdgX>SY>n@%<2d5f=ybOVD77Vam!7d?hsF{qK?&9fWw0Qq^cuzOHt_?C6|P2K)1=S6pKqiy&q1A=+b*X-&&H7kV&hT$=c19m%{Jvqi8jBLabP<1lyxoc9f%X2wUPD{leZA@@2QW^6&8)z(kR*d?=' + 'Z@gv*W!0MJYLJ&*2^f#+$*CjVblBFu^V%w$EKq+|QiQP(l{8p%Q&oBK5kvn0h!vH}y;v9c$zcYJ@aq!2ZNdjxju}^#!3yYnJ_LtrsIXh`q9isTeQ%Ckunv4TRS5zW%?vfj6JDUr3X7#15' + 'p!B}aaQ|ef2fOBkGveSIU$$*%W;v$XmH2|G8y0!S9A*Hwt@%Q{fYTp?8{%wt`aZ109)4P#JF*SX_fN)ZY~^zd*XMj~dlceGvaK8D%rj5TvF&M3O0B)FYFy$F' + 'jBugabo`;#nKv{Kf6(>v;lw`9z-Y)Dh;Om5xbc2cGKghNtVZDR^E#e-9*t0*9CJQ)(*c{Fk-36Fd}0s206J)gvi6z!vOV2okYeYg' + 'X=})%`-2K%5WR>_tMK#rC&jNY46m{#pnx4woe5M%dstaGb`lEwPF#Pd?5n_Zz!Qy}sk+Pzn@)MI?h)(el_&E(wFIBSCL@15`wvh{' + '0|XQR000O8001EXzawacfg}I`ou2>zHUIzsW?^%5aBOXJFKusRWo&aVWpiz2Z){{TFJ*IWW^Zg{GGA?PY-MA0Y+-U=V{dJ6Y-M;Z' + 'aCz-LYm?i?k>B|%5Txo3P#{KYCC(KuOl~d7i7K|NljN%;3=6`N;4UH(paD?QdcFVsx_e$d4}hecy7=lYRgssN>7JgRo}Qkbo}M_4' + 'zuXpmv}?BI>aZ=^hpT3{-bq#)Hho#u51V>>8(p>4`nt@b-R{ZrFQezb' + '{3eUKzCEn^Lt93>vTv(ZmtnILq#S!ytorCf)grO&uIQoZ=Y0g^R#!wV;JhiisBYlfzO2_}y}G3i0Qg}KxXWTU`J!rW4p$N2dRJVR' + '|4xPbs;+@Y-Qi01qLN$_blyhmYO^U@=u1^>BVbG4O($_2PbQld_>kwD0}`F*QMKDQZ4X4(jp%nbnaFQdqyB0-^;cU84C|sVR@Z)(rTUqaRkD{Pz{pIDe{K+4m' + 'J^9n~7k|j2%lcOL>8fh_sxPNz+<^aA-C*&wYNq@3)kmy`|J4;iF1D5Wwd<~Nc+}7Ou-o4vM7>v^fYg=j_i!7sZZe7Be;{H{Kt!&g' + 'kD30^Y!ACy|9lFZ)<3=kHi;iy+2M+u_{#EC*??)`v1*!Kn8|4dM%xjfO_sGnd5Sr;5&}B^<{SduqfP^{yN7U+fvuLwz' + 'Ky~?r)oAIyD379FS6$y+x5e%_Aw|{ZtDB;#yIiz;1VC^oFLzhv8g+QymERpe`IYBGQrr6mxVSzFtVFv4P5bd^ZB(B9ZC)M6+ZAQL' + 'DGQu#-3%uEv|bl&Tin`Jqv+?+pGn_$H^sh;l5G~%8ZMm@s=n{l1e-;uPZuzgNrf$)hEI@SCGA>qYUHvl8m&7eV@`y^L' + '6oralZP`?7G2w-Zf|BVY7~%ZIw@E8cIO%e-Bfy?326rGR~s-4gT{7{O1Y&^DX}KluDjb(NijW`Xs(jC;97V' + 'uby9i_582Tp5~XYUS0mh#z^Gc(Z3`?_AtkpQ=$gZE1*I0z6wtB&qA{LMsSpuh2|zStg0fB8wHhrc5s&@<`@8f6g~$V9v_t8t1kn=vZ3)pafO2ds!yD@^TmX~j+Q)4@j=E9EZl@M(jaWht2M~qZ5i7|ljmR~}Z$T6g6zXBJC$osk' + 'I@HG;%@lUMyDnjZ;S05Z9FX6T41+G9M(@MoE7ZpwO?81go#wP!=lQ*oZApqK7wT{h1_`P|`AsOkN2K?Z?hnz$WZV!Y1)8H$vMVs%' + 'OOMsiRI`QAnri|ui#|JD`P$e#Eh&ab+VeOCDD%iBA' + '^gv;^edjF*_AFe?vgpy0!8H3*KrWWn60x7+UO+bZ+vpK(5A;_=`CCbjU-CO;X&sYESAW1BwMUPZkO80+rZ*sO4*kPs^Y98Cl$~@n' + '?2G#3xy^TrA&ZKNtodr~Btsyg?3|hV6Iyh~$r#czo{9PCOR!}=>b>)3TlCt+Ol4va<)@kJXP(f;xj(>&+Mu#0aqYkuf;#awmjeruIfW+t1P+Wa<*6GuqwKDP^uaf%+B%a!lu+{C*}Pz`#uj|7Y!moR?tMI@7Q}7D-^)G' + 'slNw4CrE-L>snVpgWidaPzyNyI%Ixgkf%!vi`nTGdpv(ky#6c<>@u}r^=3?hYKqWIhLJH8mRp#X@srqwC+YHl?K(ovRHZKqKx8VA' + 'ps6ykI%Lk!rs+gFDf>gq^)N0G9ZiOWxn5||n-lY!qJR7uQ^O3vY_diD=IZa|st?ztPM%Ux^i@=|(B3mj-s' + '9K&i=N;P)G3uvh`n#nkQ$=~*gDov3jsX#?wrZGyFH5;Tp&|W>G4aPieQYC08e_p^BWL*BZZ!QRA$xSPV#u39Ue$i;XD;fZ_D>pbLetqnauF~%{uVJ6Oi_OscUotjuCEoO&sag4x>obt)uU&z2U4wxq;{uNJmq$TJ`fc?ZNYME8=jta2YVWM_i2)x#3H3' + 'fC3$&;jqB?h^@r-MTi_;b65NvN9b6YS7MT7thZ_AIXa^EwBX#?)@{FY%LXnklU' + 'tR@WhAS3*=^~%%TkzTcqTjp|!s#QFUa>_e6Q`@YDN@yt@w%a5zpjq_zH{#d=@fUT#sF*NYQz=guMQT|BBMx{nFqA4~E0z>qW7}YZ' + '#`3lzAu@H|t;|o)_|`0Nw*MLJ=#Opa8FuChDCxkx)XEo6V^iNI#tvro>AR;WJqcmFSrDN!q;r`agJ8D$5ERq5rob`IYYg_=R^P+i' + 'nI;)Vr`OQ^cQr&Hh0Mm(V);vVvU+9&W%MufW!aP$7kphGS3cD%elElysj7tqFM=p4S*95S6<dJiqYM^P!i)CC5cA%_YL~3grj2hIo&jssgQrra()RaE0`FZ%m3wyis`LEJogv2L-OJDfRC}InE7|;^NPk^t8>re|O;JOqD' + ')E;jU$IIwJ6u{)|AY&ss$BFQFwXy4?U|h7399`9Xq$rJ?C>LQG*zMk1F-5|E&7y{cm@5_b=sZW#`3|H3ItKNyJ_efms$*eM?|YDB' + '8^4%G7tD?l6J}v1$U`190qcH~{D?uOTW8Ms*TAf==qW' + 'B9X|2ZMUGs+I;Z#u6eAx#*&FB&pGJ+@?+l?@Uh&kyZIu%$~Ohx8;F-`O`@$z@~85ab}AX$wVbf|8KJ@34n3p)V6f|%m{f1j#M1l|' + 'QSrH0F;?=T068i$a;=&%H;evEs!@bQxsy-}tgk&`H?R!#V6%#9u-`2@(Ms&gg{U)CH4Y{Pt{9p$!z7lj0@`s2#LT%7jAP->F6q_DYZn58QZ}CQgxM=gQ@>CC6hf)*4c(rX-xMfq6P3eQ;nmDBcl~Bgwbgh;E' + 'usX$7-|X@Q16+Fi)ovqQ2?+4vsEXY2&mGaw)iRe3F;ef!axXTb&{mhA)g~~7x1ce~Y5wx*7dgo!D(0AsL)DDw17ZG`F?s^HQd)>F' + 'v{ypWyzPl~<1NwIcvUAyIUuk>4;S$iLxq2j|KR^Gm0l2|PqwOH#{wH!a^Y50V7p7gO^!6=pDLifcv34IYcil+Y)Nb%(10!^E1sq|' + 'UIv&cLkX{_IAv&G-}`D)O?OX*m6;G*r|vNmRSk;Lj6&g+1a+Dq&)U;#Q$xGT5S`l5YmJ1BFEgW|M59aX?i5Esp0x(8Ep!XW2iC@9' + 'Top@_@9rml2jnM86su;=y^p0FY-%2Fxtt4b5ug+bp!ax_mz*ymi-&R{tRCOvALzanO_n8yy' + 'W(ZnH+)IAgX)TkI6Fl~yFS+8N3R2WQSq!X0J+Wjc>CDVe<|4=jD>*+Y0)YK?kNsH@Z5R0mI#`tV%@n=jn-qH)cN9FR`N@VjcwQaW' + 'c~e6;M#aTB=CNidX!?lF55+IjhnYA*^3IZ+ncTE}xlW|Gr0Ko!{_5IJotdSo)!qwdpjUT6IF_CYmNU?erS7qhPGDP~CT2!=vFPcEAPfOO!5{t{UW;QqIcA=?3C0IHde|3|7Qia6l8CaRgL`U~pLR=@{T?' + ';Tar4?aferz!vJIJ)2X&vly(c=@3agPc1)D1sN$v@jWgBeTUjxzY!t' + '{1&S?4iRWj)wap!zm9&PnqfjL6$Z09WGLujnQAo>KA+}iWHMkxJBnx+t5Px1VYCdRez--qL*X->EfI8DtVgFZ1h>L&mO2!Vs!Uem=$y2aYq!0k^i*^k?oX$()>?p2cIS91}xhFORl8fV&7e$L)IQ;7Eu~IyqHO84h0k' + ')5k=4X#J}EJe)(IW=n!gQfk)pNM{6D9vugz{<(Qaih9VrlfE5GpV|cZ=7ZG;x2u>Nk=s_r#FJBu;ybLJI;|kqKYY)w6z`c}|JN4^' + '5kH;yIzK&fBV$PItKlQ(@I0cGWsfe~b(wsX+E-Lp1uQAmsz7UA1!sB_W1LZBYs#Dwaq?%->)|e2XvcKFLw^#`iBO}ht{P%ggB^>#' + 'O6Z0EHXbBE2Tl9vbrb8d`FOM)71qUt;BFwJ$}w5+h&!pL>16nqfy=jL-%C#zk~fQDMdM6ng#kimGB*rRJHu-y_eZ70Z-EO0{U<%v' + 'Z_o&I%Yvfg<4gq!TY5n1F~n^vu$1mC#=nmzt^lDTVonrpg=LnMWsHK%^fFDoNE1RS{W%tEJpA}u-XWE20Jhf7J)RBCcB~YHw~-C(eq9dpSH`!vmrRyrvcWo~B2Pr5r&|x0porlb7hBb(bA`qaQ0h(Uh{lF2Z1BRp~(=L-5%cg7Ui@_E1vrDH;v-' + 'boCgz#~wQG^uq*JoiL`n<|m9PcHzu%C6ND-vBfj(^M`iV)Sot_`e^D?2TvVA72EAqv3i%AKA=2;p^J+dk+_@$UErJe&v!DCD!RCj%}EMTT-4Vk1wzQX*rKxQ47RxxE' + 'W7izUoqIGESKagyFkYH~pL0_U8GyLX' + 'GT|`lDB(ySta~5q(#C1@Sl9nf<5Ue&u5Od?*bw9Y4g1^^cBdm9^SRsiJ5ncUkbe~}yRUIKgy)Q=5wXsA{*Xcp8Be' + 'Ip=Xsm7`aSe5eyxC{@cS55)|U!Br<*V78g5HImDmz`a;Lp`r)D!_bMBZNNormOxuqg$Bny*^zVOLT&=sDk7qMECJ}emBvKcvttV`' + '9!fRF-(Gwkr=d1v_5=oRNsU49Z$7+UCBz?YG0)Dx@EbN7_{yb*KBXa+z3i9T-n{ZxpB)Z$kz#wHdss$wGkb!R*;hOqoQ>%H#BlH0+7@H)ku%|I<|K' + 'yr+cZ^5AvXT+vX6k5)k)Lt8_o4>BcSNZqdLFhl<=LQ~g@3E_FrvuOC2*$3j(6O^Tx2mpTyQqD*ZDoHJQu0p)>-l2k(_oh^$UJ&U!' + '*5-Zmz7rp0tM_hFLQ#=;TDGHsgaX4UH7QwGICRi4y6yd@B$}~CQYQ^_{fq8gi(e;Qww9zzH2N2z5x-7CBimQINlAcnbDEPpV-$PT' + 'C^V7)2Ms5A#wg}TB6J5;ZfiMhY4p43k=FdYIVy3_EfgY^8O@mnC!Amqyj{_jD;pLnJ4>cHNo}d?8' + '5{i3){Y+vrF-@n3x`RD!`F)u@O0yFX*!(<#py;G#7yMGzgeWBULjXLro<~SThI`>f9e5o=Uy{01InLFQB<|tNg9#+^*N5iN@kf+|' + '+5R%33?#CI{6AGq&oEIc_`rh*QS?SM2Ige+_HFg{EtqSB>+Rc@3r7o8m_2glf4o>O-@cs!gxUun7Im>pV8aEXIP^E>a&rqc=CD+T=f^Clf$)o<9XyYMEqiWEy@@R-kq;LWm|4zCG09?Z2x|p>xvu05kSy7J~UopQDYX*0DG8*aS_(JAl5&e6wiWl94D@G>y$k8}4W2' + 'pp9cxy2cR&H`5@;QJi@pZPf|wMG+|PEE68X*33^J{^H1JwO+xaYPKbl`csb7f#_3#6lm;6Cbtl?uBss4nyY@' + 'j=!Tbf0Q{k9D7N}Ei^X<1e=6+&mJf#(f_e;@tMT_R?9ch#+Gd}Q`}qQ_J!(~87;7QQ;l#uG&qEPFu4A!>(osA0)ipwCBp^u|4p;j' + 'B$IjNJpPuf(BZdVxoLU%Bz5c_`~F76Ud|ByIhBx;?q}e(K9CW3ze-0l!7jXeTI4KpQf#e^;?XIaajzWz!0s@qMG&wKt%6A?P+gpB' + '%F>cyX8Kk~^w^mWC;s=gdmn1Sm>zH0O@9Qk(L+__}LjW-_l51rKKZ64vvk#t7duye6#-87ERPa?!+ey;u)RP@drD_2{)cdjh6' + '#Ua_`K|PxT!oqRz$rPq%A!7yz456#*w2HQMol$G{Y;Y4;1}8;WHo5_pb)kNDL(38?NKuiFW+xT~oSXHADQcc92eEniB88b2%-u+G' + 'fdfMdeJ~T;KZ!Zv4hM^|^C69RsH^W#k5QbVIIa+%LV0;*Wxb!})lF~G&=M4WRCRO#@`GDz6$P2jj^#rQMwgP=FtB8y%4Q8@ytVX5' + '^r;=tjY{0%)KK)4wnUx' + 'VzaPWP7AuO>m)9-LG09i8pg@^7g+jv?lBGKCnXC4qDF9GM8dJb-Gb&1OQCl10-PSR(CO`FVV76+UWN@yKbXDoKVqb-^FvhGI>8@jCFld_>M!Kw*JHOq%e^W++}vpHCO_Ms@`Cxc8aXCDGlUeI4i}1>8X445={SK%pcLZclp)wY6$yK' + '`!e*i76W&@U1Z@a-?sO_y2s$$w2%xsa>-j;6!50DI|lw}xZW^bpUinJVR(2acrRnrx`2t?9Qyr%BF;)Oz*k)R)g&U^j@bX!#V' + '@1Tswy~)@&ikn57`#omq_fBK`#ytc^CHxc^{GYSvjvBqUuiK`yKeNNnhteC%2rn^SGJ@{Bw8k_vxd^@~b3z1DwLTQvL@=6bwP3aX' + '7n7fin9!Ntn-Z4iwrFpkR&9wv=C?_Tl&{s#VEh@Hfao|x^HA4JqhCg`IJ8_}#i0PSVKB*m4^T%%61Q3Fp>b5zSZhu-w>x4$KIB?O' + 'u5QZJJNYwN#Kh8AT%G2Pi!Z#LrEjilq^5=Gx9(a-=k6)zRz' + 'TuP+4h-Yr&*$GS@6v}blh=v7~vp8K`EY$-<=}}B!URk>dnMB{{sEmcRq{oR3ru961MEv&|pEQT927s-F0iZ@(E5l)XaK23DcDBp|Tz%neNe*8{mqgiqf->Jt*WhP?;wj^is-c}^k&#g&l*8le+pUM$UXbV0zb;wQv*5xqRbb)?tk' + 'ItabkqaYz7|Dch8>5En{#?06AWS9rEIP?E>+MRNFVJs@y-Pm}tyJPB;9BH-60cF9`!u{}DQTA12Pxbr' + '3A=|iNxT6ZK&ZY}a3o0>MG|>m0bhC^5d?Pec{tZKP@Fm%}n^wn)|y|FByu3(q4t)@7~VP1|MJ=&M`%M0h?~Ln' + 'M-oWZ#d)(JXvhI7*&p%D0z1scM_^Eo$M>Th3(BW|L?j!+^zpW5K0XV3Ux{Pm2)#U!qk<(Ed$Xnc%7Q{+jB6TEs2Bemh@Ttn4h2C&{Y{4~L%`i<}jYxl{lm7-#O9KQH000080000X' + '0Lgz-o)kK4xayMD!n1qw3J5huF=8X*kg' + 'G#9{t;{>)7pj8EeC$3J6`A|vfeDeAF+dH%G2Px^!m$WHjAdy7w&dz&gXJ;=6g70@_Uu;Wy*3p0MXuYCmSy^1w6|H-6os~OEn#*74' + 's_#fv=cH`9uB2T@U}TdQ^;H)KK`@!D+omFEy59A>mZm8ws%_KuFr{w#tS_3nn@q&BJnOSnnROlQF4FI}fR%jntlldU>Rq+nldL24mJj`=XbeQ3|G=Mit75y4yDcyv' + 'E91QA`V1&g2;l?4xtMmAEO^hExhq5pva1n(yH`GBYje96)tbvf1znE~stRb_TKEWq|8N-M9NG' + 'X=#>g_*3$(1+M7_E+qZ}DNdD0Y$4*iZ*6HA)AwCZD>>oEtl#|5ezRMz0p&^h?CJOE^OvtLUcGt#!^`yBix(GfF48y8zIky$5>%qN' + '%HGnx&FU_k5cmf~_c1Q;mhQWdYgwFuScI(+b$zfgXX%hO151W-DOu9Px?B*;nKe^lQ%00ZVyB!9Z&AZbSPGZAnJd1b_wRLV!e~z>z9z-_n+w7(!_IbUi6a5PQ17~^cB`PI^{$aj?cgS_cIbUVUiDu`*DcOL$7#hA&e%z7FgkgcQxIh}((<5KCW+0}hF' + 'fP1}->pW}QY_CZ88;;dzc4XVq6-DnwD0l^iW|j5u4V=%esx`jg``6446s<(Uj!iGiU;%cSYU*HDWL{=MoZf' + 'fL|pbu^i?FQAGYMpMEKxzK5r5rouR_vhFQR*PuSgNr{*g%Sk}ny0{Vi;i#{gdIdbLkU_E)30!^*XdLSvVj833$Jyb)K9{0W&W|A&8VEeBGB#mWSv==HznU(~gO^-$ATeXdzEY32t-&(qAsQz^_3AQ5h59' + 'AGiLLO=qAM^BJrhlttf{G!SfFt2v(}WY`OE(2U5e0RI1sk@2E!n>JhrKh@oCi~fq{64IR9if?yf6;PUzj$t*Mf9|1`(`~;Q8UOkG' + '3vc{olkZ_Td;i6J$rs(i5cG|F@lQEW_KI=OQ*8vYn6qh1sm{TN*Rvps!R|NNmWB+(RDk45Zy-dV-6v$hzmhe0B=NPV8EOj#biRyp' + '%8Vw8*@H0nXB|W+13r8yMFNR~g3CM&W5-wIT#Le3HVBp$bwU+|iim->nn^1jRW;X>9F*YP3K5=gv}qX3W)rP>^D(BM%{0da6_%NB' + 'G>K^!g6O$QotTmEkc4>8wU$jaZESMP7B$=nytoX4_?M=rL&bjm4gU?>P24wvOEZHdQ)?z)TaUCjGCGSPexFHP5SBOuai0mQ`-a1c' + 'Fed=Ddnc{KtoTH9T33I*d`>;Z16*$wJI{^+VUTqox!=aEfG7~`fS8&*6' + 'wxY>*S=muUYvw3-2Tm}A&uojHJv^VvH6ld$Yn2d&U?*Nm+=NINCrx(Z0AdA`pE#=KX9' + 'r>kAplS`mTkS8dM%w_dDutrf~FnfZTq4AvS-1!pIhV#L9GyZM}IV0v-Z}H!H2;WTxi6dkTXSCY(d%_kVGJb%`U-cWvnXZ6aw=N!c' + 'abT_C6DplQq|L~0CNQ=`rFCcyh>#EL0k@39gayFfO9)!cT16K|3QrHY&jGG8AYZG!q*M_0UMmLJosl=l)YlWn7&Pt|?+I?UH1!qc' + 'o0>WbfiEQxVJ26>2BparC?njhO4tOtu~JB1L$`wvn9^8s7r6)l{Lr2q9Yaic;`?A|EO#g)h$54rSFrnl{V6bYqQh6Y)=@^dVdiwf' + 'k1(GCU(u(q?L2%ML7Io8LG<`mM7nj%am3dpntYW`+#pdgtECX)vY$+i!fztEUFk}R;tctz+' + '8Q%aIzMSm=aCuou;C6ipLv0Vsb2otKYh`R0m`d3(B~xfZ@WtQfk&{}l8s=Mij%Al9gxzG?x6&c&WX~D?)|=1du1(={Uw>ssiR6T6' + '!7Cw}R~RYOkE+W|UnR8BFk9RXfC|5*H7hsSX`&vMl_gNbWRB%rmO9K?84!lD)cKIt6Zs=U=18t06x}HAV9|{yTNrpSivif3t=gIj' + 'o~YI&WcWjRJ(!ceZ2ISz!{dnn&Bf9N5aEAy2?QOhAj?>p@q;;IAdOoD8wtM<2c&mtqRCUAo~Mi8NHC3?+4u=o4f5N|P<{4&%W#l0' + '5DCpmtPMJ><8`r{qS{rdtj)l0iJ(0WsF9>ve$N6BjWdAja>`>cf(l*u+QQcMPa_{@MklOnlGFX_HsAqG7?f(%iL|sCkkhP+uz6a1q|n-' + 'UV!mT?s#=FwnTxuIyk&BUa%}QtM}oU?l}?>nW>e@sSIo?M^Z=}Xp6FCb8Mo4wiZL{j3vM1m-)%y=!l+4j8F' + '#ZInhpW*Vl&Lt2NhO3#j6w>EMXq{OJ)0S-}UF(9DunVT^Y*+RP&k|Nm1-oTg1ALs~m_V4HjJZf-' + '#baI-*J^?FMT?y8At}jl87wsk0jaYI65LZVU`CqcoQRfvxN;D%~M`=!=pr5CKD?z&X=*I' + '#iGQUdb~Qw*rqjBBBv^WCT1d4m})-g7+~Qy!?O|Q10kw=c_5O*AZQ{mVRVo`AuIGoFcW2Z0G>neKxg&A;P`cbgU`tZqR*{^!yQp+' + '5DvP9|J|{HW1wyaXmiV>cKW6)oT&Y)cwldFL{)nc%Sk$a!jugM6ChvVJU*BV(UcPRGjUeOgTuGvcq{kvg^Hr;|TQYi+upac3~W?=5&l`mN{' + 'lVKV-pA0FdXv5=WDQkQb16JUaL)N9&Wc|}8gJlA>&r%>0_3ii}Yu~KiZkqCU<7(K}g05g^ZFmsYo*lFEUgka^JSf6~BfiA2a4djw' + '*k~?)7~gMhe6TMP{4er5BERRfTeStNj@$)6oMr_uwHI1%jXt6y<`KJLEFLaAO?zNf^pwiz' + '!J?>L7=4r`9iM5MSHs;or$~K-8|E<;KmRx(27=9|c&$X>zv}1gmI+$U6z>QveLBqe4HM!k+J}L97@Qz-ho14h8h~E}N`Hgpaypzz' + 'B%wU4x2iz;&=7Zr9p4`snqCFf2(;m+&p' + '8mvpY%{REx!Z4zxSL%%Wh7~qgT=mE@`ceGisNf-A#n|-=fv8PkH&odW3G4hYT|+dnMF(5mhRU2~yOmKbh>k-wD;t`4O{U0gjIA%L' + 'R|^dh?OqS2T6G;eL}BN#?9}C?JdhIXT~zISzjDvwsGqlbp+XrZeTyHTw0A?dWJFhSqLKc%e(1s5_V&deu@GmB*%vLZ%Sx5R@RCpk' + 'E#V*2@e;zYu%%@@yifmswntz3gaR)w)8CL1t}yH_BHu(w6=DjYP#s+AOEr6FH3PC5HYP8B9eI=|X5QyE!|(Yf!wI(-=uO{ZHPkIm' + '1ZM&H)Y_4-;Fs!_KYsUMIo_@^u+fT+cCOe&;U6iGM~sk%;Lk4PE8E*8S=7+G5%xPebOK}jJq4pYq85!*>*XT`bu()qpK<=*KUM$O' + 'PUnyBa=ynIG4EsbPj&t~O~djqLbT^%&TpyJ~XOVuhlmk6N&wr0d(Bsbg+!8wRhEVV!c;E%0-|%wKA6Inle?{lPmt?%{p?P7(x>e!cq&VKB-qCFa4zXd0d*EF9' + 'vT9KJxql|;QKrD^)0AP%HHiD;e#l)JbRl8$fjb>iq(km@T$alJmZ5N6&P^8O(1cqYcgq95%W_$>h>$g~X}-@P)A{{T=+0|XQR000O8001EX-=$OvzXbpQ&JF+oIRF3vW?^%5aBOXJFKusRWo&aVWpiz2Z){{T' + 'FJ*IWW^Zg{GGA?TVP9rxY;1r$5n$JzRXSLYKhD&R*D2isYQmY2?yzFeJd7gu4TBR*8DV1fU$INCuO_}h^' + '@}@0C&AoKXY*mX@5dDhBK}L2>y8~uGw!N*4o2FEC@h&$zgNlNf?qqJdRjbvS8~0t35M^Fm#X<;w{K2fC08-@iLVa6HpUL+Z5iILn1Y2PP~' + 'TYIx!DAU|jtmSckbcg%%WeW3)Bt%zvP<#eEPYlll7G' + 's4e`%wE|%f=v%fklm;dwFEY5cz}dP&e|YPpeLPt6b)}Ty#t}7dkRQ$YB%(q>BKgBCH+JF#3O+nghzl3n!vi3t=InsU=@}S>M8*s}' + 'A`x6OPJ=_3w&MxEnK7!G$>&p)Y$;>a=JJV5VOI0fKK*L=msBx;ld1MZ00L<2u9;2j0z' + 'Ndc{}Sit!q`NESAjsL~!uD3qi(5IaGQn4qFamhJO6uSeb5aNQjwn{OXa$~C`g$)M&GBNoY4H^gECy5C#Okt%|osho=sT(;Eh;OAW' + 'Fpim%B*Hgbk7*m2KD?cVV{IN*8j**33Aq!c`yLa1QtG9*yd|bdjgcFHJgRipPOpR)2ig=IBwycfx{Nw$x|U`%;&ogvzPVrM>$NxG' + '%7(qZkC|ZI-(re}DKQ%Lvx+7|teMzMY>OS58NzX%$p@$0yHrW#sHCxr(4{RlIX!XeQlg^=IpVeU;KC;?yA5!1RxSATW' + 'elVd2&)g&*PF|@WEB&_68_f4mghf9Pfuo^tFQgB5Tj6RY@W&yj1U{U|+L~df7BrT;jk0DITZv#44*?~|?73Mh?S!V$-H+;)aJFxL' + '&>{J=TL7GvAof8#3f;S-3pqYCl_bkiN6?zFw&k*jN1Y?jSdZ~`zc5F|`T4EtOtc)8sRy5N0H4I%?__%GD9!XP*Gk#gtuR_?|H2nU' + 'l&=OtUGpMuc%yJ1E!-0*j?N1HbY`k1+OwlfRvLOT#d}f-VPlVy36mE>qeTZ2vex@lRrvVAu@Ucg$E9Db=-e_5H==O67=8NvJna!(' + 'mVwrj*6QhiGLZ`uv%UoPE?$OqqGP-V^eLF{vIkgTu;F3e9NhR$mN*NAQH0eY*ymxXVGM+8GL;ovZ~IW13ETYbq^Y_tYaHyuC^5A6MTjaj<=^ibbi)f!R!Tab;' + 'awzrrKq_LtG>4?Q`l0Gcu8&lU;y?_c8oSmwVx`_V!A#^MUTm5ZtVliLd_3vFDf`%Dr#Q`e!nqtl^X5OMxG!6ZgTHlu?mYAscepA}Yks2sr}Gi@zwEohr2ha2220lK@txLjuR{%EPy2)jfjzy^aj?udj;Oj|?F$I(@kENA}#P)h>@6aWAK2mk;8Apj#(' + ';Y!#8008m{0024w003rTb98WQZF4VeZ)9a`b1!9cZDwz5WHK*hb8TjCY-BQDZf|mJVQgu7WnW`&Z*pX5Zees~b1rasja9*J<2Dez' + '>njE-3dks}ZsNlh2E0I;_L>51Q4}x&hPFmF5s6euDv1}xetU(' + ')koNCZ7QjRLtEx~o@KSsEpT4VL#%&Ip{T4j18O<0qHSy4' + 'RJYuql~s>?PpO?y&U(~$i$ipZ5_K&VuBs94E`I-%Wp~Di(=R4wh0Vyk$Y&clBj%;5gkNL>$3>Q971mKZ?lj+t{qs&M~_4<' + 'RtA~YV-PGqrS$6z=tn}|d*Y_wld=%6C=~F(;C?BZ?kE^ZUco&iJdg^Wp8kA#iudGkRLaI+FQy;qoeG?iDton)swyhkZkD%;fMSrY' + 'S|y65C&CX#97W6R{@BXO75lGcI#U?MA$WY_yA$^+aYdySAx$*g8fG%Hv=vU0Xrv1K4UM#JORtad0MT6&u%l|0w=9d2PH~d1zhQ$U' + 'D%}at?Y7#G_X)TPImkMVCeI2r+F!g$d)CvEp&JauDU+m6Z>~KSXM)TVjJkzyI=4OVK7y=gDA-t>lo#SiAi0l`?`&L#v(%Dkwm#x9' + '^s0aJ*B^64$%r^1kTCG=PzI@uYR>S+cZUPt~FaE#9*60&=AY-@4_>GWH4{~=samTl`AK=W{JFXSd3W~rD&QsFm7h=0p}Iv;pn^-' + 'FA=w8Su%Jag!qiUJr~Pu@4IYFQ?|lBW5sJp8NLj#3_dJMIu}!}76U{dCc$OG^TC_N_2exI#V6`rkog{pi2WVh;PQu=y}ml~XP9N(' + 'XyJ-Q`lIN&<`hr**myM;mSNEp>tvhYp^LHh^IJ#d+CbxqVQ4nEqaP73XvpBYruBe2{kL#NC`;p{4eb@S-^o-42&$kA)Ca0jiCD(O' + 'Fkzt#u%&N}QGECWAL4`@`eiZ{UXN(7m_uE%7`U)0V+K@6aWAK2mk;8AplIce3Q}#001i;001ih003rTb98WQZF4VeZ)9a`b1!9cZDwz5WHK*hb8TjCY-BQD' + 'Z*_EVb#yLpd8JrwZ=5&~{?4zk_ykvi()M&8R^oJhNl%rsNtNcLR;!ie1#gxHFbAgB{r5M1#RfJ_lYEF4`y' + 'bn$upfmOvTlrz>9WkWbcoHjk{dQ>)?p_)}RA1czF<*aTi`U7<|2WX#$VHgDaj@1a`{lJHgVvNeVWuTTc4daBDtm!o-XB8-tBPrSC' + '2KbTFj_fL`Kn39>uSnn1-T>@Sg7zh?3I!}dj5AhQ*j>klc6U^>{MeSwotb*m9D_jrZic!&BGRL#HA9?r`9T|*(!Smcci#ZU7Lw_YCM^ZtSKkL2SE@NbdT;R$F!m~Z8+|a' + '^^R2$;jAum3|rgLZY81?qq9Fz&%0HC;1A-oCM_a}ccg5#K' + 'rl43La6@Vu(MAju@?#|e%aRFE)6wCbrR?QSCs({HS7qWsWB;H^k57o!*G^{W`|S*O+q' + 'M;b$5!4VgyU}xDzVt0R8b_Vi4!a8a<>WrxE^z`2WY=NB$>oy?1g}q5ON8leaR2=U~4n#a=0uf88`~cCyB`8s1V8LU=IJ>6}=1h9q' + '$XW%DT9`HavTzLESByMr?UwbWJmK!LO8PP4n%DiW`N^{CDRKu2iG*pv9VC>;N2X%zerQ*yD0AuC1U});5jg3trEB4};nw()kZVV?-n<(u1DyUPKJ=S*n' + 'BFqz%pFRw7GsY)b3eVPI=58o~zLJbcS3AitpoI+lw7{(#p33w*1w@_~%OV@KCKCrMpJkp0eT>|QH9TwD6e>yAk>g}b2`ZqApaLMt' + '4W%K47#v7TqqCO@dL1w5HKduTH4{cF+kJyxXrjV*{Y-~H?bruLcWT|MiE2ZUnNrIgU*8|UPZUV1%%v@*7TmSqA9+e5?wH5!LIax3a%i' + '3U#WVPe_;6q`zO9?C9Q*{nWILDu(@ci)(B0HG5`CH*|u!a!fKu5jNK)a(Uu%Nb}KT>pA-U+<$z}S*0A#n>LZ4YABwuUulEvt?-bm2x(3%{jz' + 'QnNv2)EcUz4NNL19>r6(0J?sy%4bl5tGKWfc{(nj)?)Fl41`SqVJP{oOI8TQ&pTRL?0=KI<1J{O}z4&' + '1c?#dfChV~NdV+_WyA!kBh4LEic`1PO-JCC;0)z#fSzDAJJfJtSZ~G!wsVB|JBmgRMT>JXjIIiM{Yg&Z7w(4Q+G3*SYKtFGGrz8l' + 'RPl5ytI$CcE|wG)-gQs96)JlS+&9$EG~pqiVUfU5mMz|_$7VsjLb%y=qg3erXKQtuDmN4Ma=Dh)V@f`BPG>0Zfu7Owni-n;AaLUE' + '0auecL>lZJkvRv{ae+D~z|m$zZxsDhH=qLH^Ih#iu^!^Q#7!46FL`25=jpj$e!MT^zJ*K3(w!r7*V~+LQ+IM|kM`;1$G^j+Vksk=' + 'd2yb@RJSs5^MW?nXL*^$wPbXz0bGDnF>FL)j9U' + 'enqz8Jz{%AP}|WKTAgAv+xYJr?6iE?%m>cMb3l71{dF3^m+$Z7d{R^f6=7VN<+E6$9&' + '3PVTe3X8CgZM73vJ=kwNeJgjdNxS8}Es0e%7UD0P8Q*G0LYnQh' + 'esk%G*{WSt&(7`x?1OJ*u*zn;*q{1amkOPyn3c72yV!y@j7u$>5CRRyvIlyECwfs^vZfQ>Zp?eoqKeavtV&m#>FIt(Pph`2C&F!7' + 'a#~06uXnZfjGgLbe?BtHSuZm5GI&-M6)C{d?4*IFc0Q_V%HNx58nuplxE|{0pDci09{EXQ051Rl0A^uxbZ~5Kb1!XgWMyn~FJ*IWW^Zg{GB0IwZDwz5WHMiHVQF$@' + 'WM6G_VJ>iay;^&38#faFU!Q^};DFuhO0eRhXctiKjK(+wMeHDTffTyMnwDsL&-+k!xk*%?zk6rM<#Q=ZX^k_`#v*5i^Z3n!l&90_' + 'k4@e2wqJ9?>6fN%s@-f$iMzv|pHG4E6@B=#M|2})4UtRumev!X@lZe&#A1<%X-(KXWSCOR#yqZmMljvZ$Rtu-!76{qsMYZ{le@nM^k9mTcRGv!c#R' + '#yM+q-sIr%obR?QiplIX*%WKOf*>WNsCkwO{lM`b82Gy;{GO4fW;4kRr**}qXXI0}rZ;^_+a2jUh~~{sP-0M7y>2$FmBB+{7VVNe' + 'iY5F3se?}Ga}Hba`<5xnOmib*?487t2ZlI7nHQW@T@-8P;@07{=>r-FwkOY&dA|invFLF+o2Fx6km<*18VvK3RHK#dH3EldEZ2$H;|YjphImSj4gLN0H_33_8*bY=4y439;i4Fg09lkHlJmtI*KID(x7XKswG}jQ!g7J' + 'C(#JrZ%ekSxAS^K+m`MUVt&t*c_J}7hz0E3_GKA)2POfOY%mR%zyODR_pL<6h5;(EtVvP7v$pBRvNY#y*-(D^O7gCovaf1jJP{v@' + 'x??TMf`Zo9T(9r-X=|ZdXs=zS-2-7dW@l{f=>i' + 'TS)#bnw|tR+|(JT;&`5MP18ZQtPzooJhu2*Z>s=XSkM6&5>UA+3i0+vkGclkT2w8HuWdl(A_$9+a?1desW7nZp5X0U$cb=BBV%@WPVPrKBu6ZW(~#S!GVT9wfx1' + '!_#({aMsc32+@gC1A_HU!xRN-%@b$OuNltfICe<@wzd=I6hL2ZBBiCFa$KOuvXk^Dy-7|#?|+*+JkFmp@mXlLlU7!DaiM6YTvVkpfuAl+0T{}YW-dJi1l?zBgpV~%V!L4yjBoY;p&55{2`@UvlCgf-;64I$xX' + 'Pl}5f<({Rd?Mm1|(kuy=rlsyi(IW9por@%dG>QYg$rWgPy{u5U<2rl{`k' + 'ewy%EE!{X(_)`W(K4s*$CB1u=Za)l;53K-=f>;_1axz>AzvQS`H{5UpCa|)SKB#1~-h2Z5mqT8YNDZpgTzB*uRF#THw2dsBErILxCGtTBVc8Q8>m(T3NRLT|mO@#o0)j(AyEU*sJiKmgC;xB;`;EwEsrL' + ';j>SIbU}-C@>#WN+iar~(Gvo@`mMyrnL?2|#w~MpQBjoQ!>AfseQ0T={m3Nu-h3pZWO7SnUk}RtKTG7%MUOX}mf@m@+byf$h@!qD' + '!sZC`5^mQZ6L2bEH+(=G=|fGb;(RAc!t%rvtlUcXqK>p0WM-V4xX0TmZ==ES{^KZ!rq{j<5Dzy<_45?C7=}RlXjDAeFZwBd7Kt*5' + '`;mwTgA?;@N$Z!KvHY$X0+m{Wi?!2SWw#DuJBWB^9LF2Je>Ko1C6*1bn+?e4jhsQYWH=|*>58|ZBV#qZTws(vuEjyF@LYxeOQghv' + 'SPDxe6JuXuLJ~Xj+#a-L8_l_^%FfAMehi?St=gfw#H}E_(n6' + 'tH9x}t`hOSS8;r%@*R3m8=7QM^;P5`a%CiDR_U56^bBuHiwwHz!l)0OxzJ8_o&MUM>_BZp8dACG`L^d)xWIM1ZfXQh' + 'WLV=6lVhFDXKadtuWy)c#g{y?4o^Ixab`nJyJNL?LKjk2Mq_vS&Q@unqLJM)3-T)G>2da$p*N040%KZUQO9KQH000080000X0HjYs!|Mb90NV}#04x9i0A^uxbZ~5Kb1!Xg' + 'WMyn~FJ*IWW^Zg{GB0IwZDwz5WHMiHVRCb6Zf7oVd97DbZ{s!$e$TJaycZjAjcrdufE4Hr*Wq5S!&YndLNU|!x&nq@pgoWzd;pH$Yl7N+' + '8~C#aRe@hwt!ijoqM&T($drP1dl|d9oIu-74(059tHA%=jUN3jK-R$mSr;w!HbRadHd~phkk3(rE{sdiK3FgbtD4vGvaw>gLfUMI;D@yXXsySyoSYp' + '!yx23Has1txiG+f<+6vjTDLk)u1O9x#z+j+l2Bc5lPok;*Uq_}%Z&g>JX78F2AYB@W$HEa@Xlbm6=ZH>LVNqXB{e4qFOf}T6X=>?5N1rHEGQ;q;GL-b|*' + '^D|nUQKDg|gOI?mNdF?bsW^@WeL=IP4ui)qa(Ch#H=Gu7moaYG0mqn(3egu^7!MPFda;7^(`*)QBebz&&@ab4MH&U?5Vi+k&dQwo' + 'oe$jSAIZuQt`BeuZKxe8nnmzOUn{I{>w}4e$oN3WQ`a-3(r9+&^I@B_MGz-?LpZLSO-2*XCl}d38PM+T#L*r?N$o($_-ddE=heNz-9{)bF~uhr' + '$Ovu?2G!lg=^Wu`vq0ORY;nA=g--p4*$}~E9~1Otd+$EIrUIxojuvGX_Q^{Y6tF(mNj=i<@6q' + '8!t?R(<>9f0T8p&@U|YW>sw2g;H~d0yzlLD_A*{)t@RrnFwF)k5^A%(c@4-&{e=UBO^RmQwj6d%1&h9fbjlCrID~kj>Y%Y$eVKW-Y!_vYC' + 'o^74kcvg_Oyj%SPP)h>@6aWAK2mk;8ApqDlyQ@_N006}g002Ay003rTb98WQZF4VeZ)9a`b1!9cZDwz5WHK*hb8TjCY-BQDaB^j2' + 'X=8L{WM5%yX=iR_WM5@+a&K}jaCxm)ZExE)5dQ98LHNTUr7{~Q-LM&Ou_8@@edrJj*)X687+Ruiq7tc*RMHf~e*2DZk|;YZhSf-7' + 'i98V3oiDkky_*Ba=BR4O16lQx;MRIgdpB_QW=mI(on;t(2GS_Gg6g@t2*rjy8yD^v)}EyHVFI6I!Zp6=u4vWm8=cTGw2#%PNY!@$YXJ3vzYyJ9&Hk$1hjcXd7;-*aI(Fwk*4TnJ*TLiq(jevh8}qNJklI' + 'ss5-l`*#y=%+cy~c-$-yd>{b)o^wqJNBWs~J$8)6_5DclSFSI^jImaWz~`tI4yi@_#q%Z>}3Ih`&2WZyQZZX`8u>E*g`8=`NW' + '*amT7*ME5B*S(*c#B11|grry;IR=^h%^PnH3~cPcOeSpp9vG^)HdKI?8ox(LQ^{g)5wBd^TC1>2oNFHtr*bK!+bFqSEM;pFxg;p|Iv;e_M^QxL^R;-x`29Q&|ZZ|1Hxg!Wx0NQb&I|8j#d~f^z7I{$&bE(iD9FsVjLL2zGn;H}Zxg(@3A7$?_?YibE3pyjQLd@)p`#gbK#tkEAx?!l=KfiHIx^UBU' + '@f6{Hb6!^9+RC;_**&>g-pzJ^+Y-XXKeGzOe*rqFxt-3D>nC5M}saDxdkhuCroG0u)`61#i|4KrUkH3@7S4WWz7XmT@*J^*y8' + 'Vj~2dRIYx*R3)V;6$MJ702H64*&0q1O{n@~O|b' + 'kzx-DJ*!*rkqEvmdWVyh(c?Jz!cYf#vdaP`lkI5KxCmzT?w`r+Qh-Uz>4f{0mNiso9!efevizl?FSyq63>hZIdoIjuI|F|8ns^;c' + '9C7U;wVBxK$Z2*Nc6Xeb%iTuMR?Ls(i+=%7O9KQH000080000X0P&t651JzY03V+K05Sjo0A^uxbZ~5Kb1!XgWMyn~FJ*IWW^Zg{' + 'GB0IwZDwz5WHMiHa%FH~a%Ep>Zg6#UE^v9}J!^B@$dTXqD-d2)7Vtuh9=^RCE8IF+mc6Q7NiIp7bHZYwAQF-gg8&17T5EO6Z@=!I' + '_X8j$JISSPQ;|uar>Cc+Gaufu*w!ell^s_mGdmf%hkGV;Geb%+Bvah29M;$`{?zT)7n^Qx$B^Evd8my7H(vhA^jx>B!Mxc^H#u=' + 'ZT7h>@@oOotXC`G9>=jr+jL$4XETLEJc@&xJS!FqFmJQEt*SzT&8y{gUUFgy1mpt8%ikZC4{{9Udbxs`Y=Uwn9<^0HzvGM;WzF&i' + '#UlZ7Z>o9;vLVofg8=@!&Dvx>UnE7A-a{V=&S4YFSNS4Ki~QGYK`)Y3rM)ZiTbOlfO|24~p9CbThAH-fU03D8edFsR^3V@bP_`EN' + 'M?V~1ygz=M9A2HDUNX#Y&kmEltMdc(?akrQk8jS;j%}zPjt~F%C;Q>W@#V?epO2HPbNrZG{prK8_7eX8)5-h8v$x5~*@vI6l9RXU' + '<PNcz$&FhCMz#|B2csfY5q!c7FEZ{HH&i9$#25I2!Beg+y?9n7lvy07U$F' + 'd?pDNeJf6hCkI#Od*%~|PvS)iGwoPz$1gCcL*' + 'DcFeF46r>uK@%=@O61%|2+zT<{`)-n>G)A#h3J|CUq-?J0?{4@T&B*-@x!zUQ!;J>88' + '8_?h25Z3nz=%F#O2=w)NnFV$A3y?!LU|pi!$e1MwNa>GAt0Q=sF0(NjBKnpU8QLrLfa#qvj*Pz7*$t|aF*3MZq;;J>1i`-r7i1Tk' + 'yL6QW!18EBz#hPlUO=*n2^lEr3OKd1M{h`X`Cc@ube<*ad|^G$o8%t!9z4DV?GMkdYB0MVu3cfD2HV~!TPf`>bW`eP4s3rNRKQ;O+BXn&XX(N(^+j9e^SNDd1m}S=151K=y' + 'SVnD#Fu{U!3*v1Mmu)l^Abzxt?-GP5ll(hfGe2y;uE7t*' + 'wNW!CluJm!eMWFHHb|PaSUl_G3t2-B32!L0&PHZnLwFS=>1vgg%pGBNg6PYNn^=C6i$OioRg!wD=_U2hr=0#LID_TNV4`5>d-^&y' + 'x*yT7MvOc~8U7~*X+9N{PVi6cyy6s3_`fayBYwjEabNRcOhW!)G_O_K-xYI9Jy!_r48mnX@WH&=4cOG' + 'XeDtI$;mEQDTxeJel;PZf&0lq9eg_p_FTClLc0X(l`bW=K6c-E0OQQX_<4' + 'C>#c(C&ssKU|LuD7Z+5_#6*1XVq+lUg9{>NXd*s%q?=95g#YUiY4*PBwVc9<{M508Xg(+EV+R<`?nHgua_MNZWUU?W3lqWY-u1&s' + 'TV->z0=R?fl8XG}*`6iTx|jF2IsbmfkzL6B00a0l%>xI!6Y5P1xgY@?uz;)ib_c)L70' + 'ISHzp8wd)?y(0%9o!@~EI0-re)zK|Pkcf|s0vrS00{jAVdL*XlE`b3_!DJE);eX^eRP`dnhcKYvV4r4_U|=E<|ET!Ox@p7L(RBQJ' + 'Hj`L8$eLmk^aok5JbZz0y%bLvUolI?og-pyK2OT^GSodiVlePldCBPtf(hYUq2|4=F#a-_$&l>;vTNY_eKk+7*F{=C_!!#}&J$~S' + 'FWYn1o08H|jn*KlH8@72RaDU}q?0RdEsiWfYfHy<{mtKq~=9@IGIKa;)fAD#)?s=uqdwqoINnP}AMnC~p?|' + 'ZQh2F(?o;Yn_)xLlImd;d>`yFX&n9_k=rmJwEGlIYNG%l_2pT52pIUu+MNk*&Z{MO16d;c2DRj4(*raG6bKe?ky1N}=up&?>XU0b' + '_U8Md-u~X69Xfq=q+h9@fgRGsUHL$bEITV}`yE}xU3iGQcphjXgelG3fa>9' + '9F1aCA40jZ2rpM=TAU3F(8-@uf2>99zMT#~4rh##NN^8qa`BKXwbX7Cau-v0B$z`%HfkXJjh+#nanFP)u}n9$V4}*~40A^OG5HlO' + 'S7K2xJAUjitSG^#)ygnZECmG' + 'Y4M;d#p7m(oChq7VIhO937n@gurXhej4<)A(!jBm-c`{JPhn>P=YrImF|Scz?|BUYj2T?CqocRD^S**Hmu(~caM2P26MfIrFgHUm' + '&cldiRoJXjGJU4;n|btohR3otOk^5N!(&S)Hy`ROj&;t)4nui{8>uu?T*mSz(So=n!}<(KR-}g#WmV$d69b13?ezF+$a20K8jRge' + 'FOPZjroU2{s^zj0VJ_}Q8mn+<5IGzNJ$bpI@MQqZzb`H!x!3}rH(556uf1>OLKuz%Ppa%mU$F+gV%17Hpz0fwm3isgayOF`x=BEZ' + 'aKZo+XY8;evU446HtI>PDZA)(0h|p3vv)6yVb#149xhBx%rm!Pv>{cyi8MVE+e4qs#`?YQjqG>^kX74pa~Z=Te6SJ0g7dRy}oaC(=Wq' + 'G%s0_iV|hQjN?Ew4l-j6vAS)lX9IWLc@w3y;(U%;mvsBYnGx&R3qWqUkCnGGVO(SFDWS1v31vKZ!H$SrVBOV0&LZnAy4tvajF5)0' + 'oN+o96F^>;%aJ$W|d{#6|+m=BCevAB#&7&=kc>HO0*I?ZeW;IdekpUUc&5?(c#@}?GLN-?TGA*3Y}l4ZIANKwm&' + 'YYN-a;uR%i1bzoA_(AF>tE$gwroE%qf^Y+YS&-)7FkcfIK3HLt?X;7F!pP;=UZ!He|6?EX>v@KHzZ$5ko^&75zhC?J^&JUDC_Vt$' + 'sQRAiY}yl`t(<4Nm}aCac=gpBOB21B1JQ_{>Tt8@#SYy|+^?529>XGB$>W!CYZ_{4gYH~aup9G?q>~mhRo<+JZMs(b*Y{e##7VVy' + '_J779eSsHgE-uNTSdJxnvD5aI-7qB+Y(z*aZd$wPOY88$!w8sdDJ3O}+S(x8i`!(OtH*UXAdXJjGTN1R^5jY`!jVocM5bS{3i0*(IAIHaq&3Jp*<$gGVd8OQ9`}DOja7ykz8$' + 'gGYw+iR%TMfWyV&hb9zHqk);*tm*~rck?lfj>MabwET2PIz%?bvm`wDzb&geo2G4B?|?v}-kuHg-Jn?Of;kMJWq-|A64~f|)-v=_' + 'Y?DLto6p+`?wV?Pm?!DUKtS_LDn|+-gZ*#%K{RDobL;$V^HjQ@k^`;AE-IWgjlYvT7+7he|6m$r{5ByyyRMW}ga~MUQ@?dZAf9NedM?9Lwh~iEtp$NMn-ci}z5qYRSIABrQ5R1a880ZN6WBK_(`~2(7bG);^D*kJ_' + '@f(~}#p{mKCzDm{$})-Dle%emfJ;7R0$!zvUpLp;i{Wc!qxYn!w~$;cTm`gRc7cODgTPiBh>qRK4oM>M^u@p$|;TA!>dMh);~GYk>NE_' + '*w`gxb|CciA|RHlGWHfWW$63z|J}L&q0TCJZ|H6fx8fY)*HpUJhs8(i&3eS3pPsvZh0HFS)Sqw>7o?2vgKK+lhG!MtrhpY!_cAr$K+JL8OuH9!_)V@Vwtgw>{V_ckA>Ug<7X&uj!+6nBTiGW|*T_ytq9kz+anHoXMF?haGVrbeLMu()^*S#WeEbO}m{mmgHnp-(g)`==}@m&)`0Z;ZB' + '97jGD4o1PE31-QsVPb_Gd!1b3z{n7MegeV(*VjTwp0#+g#~fYxov+L%uM+}v?W$@86kbJx;Dxp`4t_GmVqV`lasvy1_3#!DxBT*n' + '5E{pnSM$JWVXznr#UtFF&OG*Ggh>oKNk`U$Kk8(hL_79m;_+G{eXiO0K%#DT8|?XZ0|K@m%H+hx`AZ@c6@KQ*' + 'O26_1r^Dy^+WOUdtpyGKr3c}^@U*?|0aC3U8)!!=LI)%+L_IfF2w=ZkXl}|It8~(9RfH?fGj@QEhwUNBwmC;Asp9dyt1@XIt{z2fw!P9x|1i~+>NZqO%jcp4og^pUb>>X5pgDwL?6' + '?6YPlaHd}(59#sq_c^Ezl0JE8PU}3A~2G1K6x^cgPN#k*5!8k%&4oLaAsIVL@s<=V{X{ZXWqlE' + 'DmOXquk*rbQsoRfM<)Sezg<%oM@{_W^)_QuBsek&LhWv=!2fweS>pw}' + 'ms0L-uzE`-lOu7n7kSNMm)hXlOrc-o%tdV&#^$ee8;0S@rIMd&^0xZ=0${#mw=-s6OY-' + 'MBA+{D!dnwpyySwUX}^P>S4$cVc&bRIMBZCnqT*4k%dmq%A8HlaOP6=G1Bcq%gUrx_E_xz%nY~Iqo' + 'HVuEOF=48(&P;i<<&YGGId)VVA=f;n0EsRQ&wybRrjNds_@U>)CBvL5o{@#6`}#Vy7RgfZ2#X8WrNUX&ni3Y6KX*f&n|`Xv|Ad' + 'Bne2vLDYqS{iRq%I$J%|iMXbRL^aYai1YAA%sZbdl>&fwXk+^Rx9hK5HB2lx*ZtDzIJ~@YTAeug1(?9bfUxjQOaj>%oA`yv@3G|3m9)QRcJuUa*J9xX}DFp%SR$;2p2?WA~;P' + 'sD@|@ga(55Y@b>Hk=*PsgY2HnPvp9A*uYGUSi4*r@<=lhW}qXb%nUsp&y)Vfj)C=Yx-$f3)_XI~&bJy{&B4uhO(E={*}wO-+V+;f`tCWE@ubZEUzlg7b~i)X)QnM|%eH)Mv`xD(;4IFh2Y2&aat&51+ad((XN>4lOwe`J05t_Gkk97Y69borOM4q?)_#H_f($MLZLBZuywdD4_vC;2wtTo' + 'rbS(x9N15SZ}mrX8>}ULW=#KwKy3-h-+^fK_TvNw{&6>6+CplV%#g5pE0{' + '*=Nl>5V8$=-#pZzzLEEZg8{CHb9wP!Txo=zZQKl0{hrG60WZy*b{JOwjjt8xU#' + 'M3okgY*ndBV7Bl)OYq6RCDZGosY<9$-)w__=G#h@X3a1-f9Uqlp5>BxNBKMl3Y$b*7c-6B7?`nsVoVmTT(hvrj2f7DZr60*$y;ju' + 'as_}sr;hV^y9I<#glz^d*~O@;KgCoc5nZ%=@p&)8B1GQ^dys16`myeNuyQ@cH`<|RE=pAm' + 'oG=^jZ)nWb1TTk9)Xl~{t5~^UUv0m=cUQwl;0yMx6=9)lWx~gvJ!!k+Zl9oh2)7H#`a*sk*?(=Gm^6pSZBJUp>-XRj{DyS%`' + 'U$9R3nP|Q>(7?zAJr3D^Ik!VX(VSoF*HyV@H#2#LKQGwP*q~L8`h~z3+4^{oqMjH!{?=qx?{1Uy9$3X7|j#z_E)!cx@5is-h9Fq84LZ1K#_nv`C9o!b8t=H' + 'r0J7H8M-YrWhv5+TO|?I`tI6k_U$phP$Qco=QN5v#5epHzP_Pr?cXtH`3p8@{Tr;;{r;BdUDOU*xuNQZX#7=aCXksbGMbBg;j&2+^`n~2mIpODGB|cYBAz?rht$W~X-DYODDBX|w8IL_P3EOF' + 'yW5`!cwt>W4n;gS1BaxiEZ=bTo2f?h%w9Cxp7=AzfxMa;-t^45)J*&bd>ql48_AV$Y8*40MZWWd+PMUHx}j{@zbN`>HGU$*o>}^-' + 'EW6~G5Lp%t<&^)x{gy)7pK9Eg^N-KddaqVi-jbiXt+v?oq#zy#l0Q' + 'kZ9ucpUUR`R6~5wpyGk;$ny(K+bS`C(=ZxXixaJ|tr&8B&ON+u' + 'rB$>XCXQO&#=6LWWgt*2=a0L1?7j?#*Wtij3956APqDX}!yKZA4|>T7^JDrDneoFj(YZo0>6Rb;_6@OMCgOv>8@eerRczCIEzXe0' + 'pZ3tF3is~4_7Ij{^8WBca&`XW@tOWBMh=GE4h*J{%9~`grXyOtQ&w!C-65mHu=slGbxDg%$z;IWYD%wHp9{JiPu_S8GGFCwOf>&>->Lu6U`#v2MTDfAvRU' + 'U1n+NARXb;aFH)3db!?^v?Z*LlC?~mPjsOG4iyBqNs9bahAcKi5~{Nj7Xz5~PDWNhBZNYT76RFSh~=ZMlR1nftH*|gn+zWUA_Qz*Plrb?f>zSQTq{0n6hJ@E4{l!1)MH)n(^dVGNRyClXLbkpa7P' + 'h8>pGqs*{WhaeMgxQUmk@hj}G_vbcOC3xcKKKZ+Tv2eD`uj$+tcC4-q!928ZSN&TKh6VC&J@C^ahzenKC;m$h1X-v}o7i_6hHZ_l' + 'u@oxTvH=AUjfdSrA+E8M&REJKyG~CMr}I?B7@K&m&di#8CQGI_k^D9*GI?|VP07y1);yasaTq38XV^7|KVhkEdW!Wli}_iwUy%9T#Vy@D|+dN1alB(U0d6XLZ$&Uf}tPvwCLSip9mHGb_D0&i4EU#ZY4q_+Q3X' + '%~`X5rNVmPy;n+p^xZ5qbl>SWR%u^!mr?zyh5Bq#Dc{S_+Iw*cUxNVh=PC4SrW`N2TDJ%$idrmhe$aJ?B1|ye;3=wW@pFC+PHMDW' + 'Y}p}H8ZjD9M8sxR?D`q*F|Q)xva{j1d6(pra<4&2@<3asMp#vf8!oz!ZD2F;tKwQy+0w@kY{P$4a#NGvs}{U0$wS@isJkCp{aaH;' + '9Ogc2PYZNV6Rib!)1tm(yMDwCoLjS7G-L%R+Y)D;0yWQn+C>AR%tZ)7=k-(^U@FrDdb5EJ#U|MTWDgBEnj{NM-PRxoQ%}C11iQ?l' + '*>yDL<)LKAO|(MkOF;Uz2}sVD1*i{r(~S#qohTew&BhbtG4`(^iOeXDb1q' + 'w}fRXL_xOWfrog$F=54SeB;|0(lUQx3hGNL%(vR1Ic%D2d0k{wQjR&dA_L!JU0!3b2#fys^7I|FC!OH;*0&cyH8' + 'Y9>+#HT)P)7!>J;p~S|6??9P#en=#qwZ6PtXvTJqIE56c{2xyT$kYedKii}yXkV&&W*2B' + '@(BX(JN;mtPBVAC985ha-|j*bZ4nWwRfTB|N67_?scaS-@$KU!gkOuaO|8o(W9L%S' + 'uFE%bH8Nw#_*J(TJ6{QrSPFyx1yD-^1QY-O00;m803iS&*kQz63;+OfE&u>70001HVRLkFY;AKdZEs{{Y;!MVb8TjCY-BPoWpiz2' + 'Z){{TUvP47V`X!5Z*nehdEFXabKAJ_UB3clA4(+*$9H|x96L?YT*otw&vSPBqS4S03E50Yf+ax5PA>nwyZ9nNQdaJIxqf6Oi3IE}' + '7W=&mMNxF#c7r5c&pO&uS@D+8wj_LzxLmS=msNX94!kFCuHJsn>+)MtwFlO-wqR)#MXS}J=MBm8!yt#A^xN#=-@9~ec6izx}MOO{uyN9pw)t=e~K1#hojCUAkWI#h*HX+kc0E*X*p+=N`oekhQe%bt~$(($=Rwc#bJK{>f6(53-f8|mk-f4KO|`Stk^7v~p$&Cf3`|9+icynOrS' + '%8&r2dtMGjwXf<*K3Yt{Sc%$^>_M(pzbCI~@%{tt%NHmrV0B+JxNvn(B_p)1zy{EigNzI?&Ltn5>' + 'U}<`T{d2(ZvmtPhUQxGNxG9`mvIEHrH2)mHdMRS{1$cEgUXvGp(%(P}yNx9Yu$yri#J<%x?YO9vell{V7SN;=QP3hW)YZmFD1{&g' + 'cZ@CW+U+}P&off9R&%!|&&f9mq-Wq(t)XJCBk-PiprTmoRI(wqj?$A`e{%q0_`' + 'fb8@IzRSq+`QPC^9nQ{uGIC4>lVCU$5>9MHx!tW(h#qK@Rs!qQf(Am110VlNhUSHfqXy=|EEtxczkxZyY1-4FxIs+IWuPy!xK)tI' + 's$~%LfE7?~^(NxBOQG&+Z&_>%Xl|+N(|YN&9^$!Hl$7NFX=STqKZ&%vWFcr81!#c)_t7Rw-*H1ewb3_KN5|1m(Ryt^Z=$uQZa`fM' + 'yc%Rr1(YbI6;4%d&vr_e0<$B917fqG1IwB~I%O94=>#?!2P(U2OZJeA0h&QM$e`9`mY-31np>cYY?Wf!ZRXH6xF4cV=31ZgPwLpu' + 'qd=el$rgkikET_M4E`o=AqWWBV$er%n5#L{U3(P3Uo+o#pMcklcG^6GU^uZEpGcx_Hp;~Fu2GmMcaRd|iFiHpC4^axamyIG;l!L^' + '?H%e}Uf;9a?8fH*&6G)H_r4=+T-buU8LU*77(r3t@0z7RgKkrC=2l%0)eH54_Y?WBQ-)xVpgt$@?#?)p9' + 'JOEN<@t#S{5E!{9S%OyTp7Pw*->iJo_q>mTupYS%*bsvXyi{>q74QN(D*?!8@7NR!9YiBW8_fX`#T{!X^mLL$Azjx&h&#BKJH<43' + '(5ZCTJuy=DyaR8NDfu4SjLW~h6(EhnfW5N3gYze3O#?VK8=``&B5RNmNL3=!6x#i|e({m@++NC@nvTTQp(IO)v;4NFO+LN_t+-C4tWxK-+rB^=|(rFU}eZeos;C`5x>=EL-7diT|MWz!nTpDfG1=MKmc2u7|WxfzbVf20z@fmJ9+ls{+&>4Dh{3' + 'Cx)SE>Efw@!kX;C2=sUjXsABFfr_F6hY0jT`2HOfhawHhd7Yxoh`B2#=H;JPo*1E;D&Jrmi!*sN?0RO|RS{gM%944yb=6r=3RuH6' + 'yv(!wv4jxMtWIQ>=fvxsF=S' + 'bRyBk80(JjHuNEPGwNJaw+-E(-<_7KVj7G8QLd54O0@AA8;vCz`H40ja~{M3PXM~Gni%=`O4hp75?zq6=Gw{^b!WClRinK1sLduo' + 'C&upMVseppIKKY%s?SH&PK$kg)yBWl!1q_Pe%mR?G~a6Jf6IaHeUK}|>Uf=4%@' + '?O@LB#|t=VG~}m?Qoj(#{5x8SVaf|9r@S@5@cQnF>+2lv=J9r(=0$8EbiBYG-S0yM67(m*F@Bk3dFy?y9fKp$z?UBZV#7hA=w^yu' + 'm_@ZkPJ`WeiU~NuadVX^i+Ngc3S@&K`yW-N=T25UP(VP52t5ioabfF2FzVuc2ep*^gHRJL)Mif$IeL=m2E%$gBhbN7Df?K7D#a(N' + '3LG<HVTLr@S>@#qh^AxLBW3iQ' + '#^FRIt)iblb+mCdAhnwgai{$nU59;AEKVAVDDH(Z2e|t+ifL0K|W0{^Sl0Y' + 'I;ZMQhK@ucbo(@n<&}O^bTP<=mCQt!pJz5UW1{AtdT5+$!<8-_?_!;FLr*f{aPew>62u2ZNP_$t^$VkCN)@vc$()eUELVzDK*QK4#}-y6GuzOIyh(Jt=a{;XA_b=JLF8ydpulk^$3olj*f=!2ECZzPIeb)czVvf&l4c0(QN#fOh+=-' + 'R*f8WbsOtS9PB_@fVMN^T4DIm6}VNt8nzI}9YJOWuF-7cX)$6A_kt?oe' + 'XNCJyd~n&S956AIpoYoe!R8cG&+pMUCs=O7$MABBPAUI2x#ZdSu^pPctPY1kV8qmt6grJ(A%qK}%ODtKgZ4wg}6OZZnuN*WCmRE8Fg1S45AnS!+F&Ni3?3;(;MX{0$M5+p0d_Ok`YrU#m$t;a_gWc~}=|LwJP;=fn' + 't^YCHQ_SPf<8XB3fWAFD`}sz6vP3*OS(`x~8?wy+Pt4n9@Qn8LW6mBbkR$e)0PsBp<@}be@I~Fo;+Vo2#w56h+PAvUm2eo9bycjp' + 'cKsz;&HE)c5mj^k*_hNHMEWDgC2dS*M!Gv_c@7z%gk*Peh-7}0n(2Gbd(vmmoH%ED&h|6+X&ESVf^5FDj5c$MY;kptJX&Ekew>5p' + '*prv`w?N6l?VVPVBQH08(=<&~qkxm;$=$IEH{AHevN&B1uku{=ub4NYrTKA`kZ5UQiR->?MZ49%08mQ<1QY-O00;m803iTiF_OVa' + 'Gynj};{X6M0001HVRLkFY;AKdZEs{{Y;!MVb8TjCY-BPoWpiz2Z){{TUvP47bY*F7Ut?%tX>KlXdF_2`liS9z==c2Km*|J(z?oTzx|jQ%nSw~cO^T?t*A;80nDRkrl()i-7^>rUYA9kr1|0eDoS%Nj_PRSEt6uE' + ')a7>MMS1L1Nts4j`b$)&MeZ$%();G}#kWNkp9F)!;NW0U7Ar3d7n^!hCSmBMt94PxE^BNxcL%l^f' + '&=>Z$b*YZ_elyd=y><-5{i>gF|RL-lJSwaC6q(!ifn?XH$' + '`fZve{e|yj>5J#zUPN`WED#y+p~yC?oPItdRDmD&Xs$}?(3mOXMTkYYrD~q8w?PF90;s2eDe{}-?oXR68+qS9ds`-nqi_(XRUHAY' + 'sFlVSBX7bSqv}cQLA=ZsV!G!=^PefTVEAKcT^UA|KvBEiYMw5_byQZ`<`^D$@aOfPpFMjUJ`KM;' + 'zj*uVO?Y{6{^xH-^yLLs`dbM*J1P0i8=hKzD2q)Vi2?2DFZNo(wt&*@VxLn13ECK1(O!ibu' + 'M5ws9qd!@?!Cgtva6Q7=mKQ9C+IH1Von}=C1f&5F>B_6}>1MHjl{~ok)A{q4;qzzV#jBT>Z@sZMI0vG75GM-{`wOp=Ee;^fYy3I9' + 'iLy;noq0^}0C@`7IPd{pWq{!ms515bulKUZ6KD(n^S#6WrQ&BadZ5`b<3Lkj5IJmZKrs9W!;i|Xhb?*0qK3Es=knFdLqr{zn^~5a' + 'e1J$lHpdXe&=(B+mn4TBJ=CD`fL#oIpVyYV#j@`|&>LhXP-^7Sr~n@RQ3XInR@KmFh^d=}mjasc*5*kV5!^oUC}_k3c7rVlfSRiK' + 'hua@_&!{n#W9_%E+L7SnYiJ6p58;<@fkwlDUKCtG3cP8V)#0N-*|Qep;dNWfjdU!aDPVpYQ``5AuW' + ';2NSzjQvdV(g}b)l3Q*{r~z-z^scJSYA9L^B-w-RW=M%Dnq@*RTs!d@r@gc5fR_}JfMHK_!257M<9&L;d-5g!cMdaeYMO4G290y;WoM_*pe3>Qs=ii?^;m>gj%)N}V=HuZ>aQx*dY;Cm(o+B9E#o;WeBTpwIw1F>101_?X#bK1KuZ$ur_fSVL9z~3X' + 'Kb`^C4)QoE%VEgzQWo`SpweG$uP}fn$vhZ' + 'VALIGTPG*_7a=w94q=9^5Ar=z%{y4)u{I=8nRfxdHovcfd0AAIC^(`P#C`_>g<2cpyOGxvrP`wCNHoS(ip2s52vt&GYPjG%qE?D(' + '*dTyE>MjY}QZLz5mk6R&KHRT1zO-ul`tG@iV5M*gZ-8~e7I`GbMEix`jpQrXHov6nA{BkcuWTcPIBpx?pnJ(6y$gkfwEd&F<' + 'qfZ*c$QVOkcT~LTDK{#1_aY}nkZo8dETE>!>hG^e%|wT?j5a!;$%n~mU2ka>^QG3UDlHflfg7lGWde%by9+u9Y!7S;m<^E4Ng?v$' + 'B4NS}(rXE?d@sBIukwS0KFMLpwg?i76pDHC2Z2~s11!ByI~chd9~9?`CbQOix1Olpr0v>Y7k#{;JQkO@$-?X`ks_Ta|2W`>51X%}Pjc1_4Dq)H$k_75#8fqdF;cD;^' + 'xct{&3S0ystt5GOUc`yJ?qDG64*u}Nmjl6e0Ln!giy4JlfC`mfOZ7%<_w+ODfCs8%SW@XWDU8(f>HrUlxGv09BegfHBC=3c=ao#r4b>L4|Up>+kY^>@&KV&F%A' + '6J}W*2v4l$#DJGqFrYSR#rt9MjmrNmt3$5>cXvV5(ZE##8cyM=4qp~ov~224^(C^dOoGWcI-zzTVVceU-4-3~X&dPdL~--K-RZ^l}3N_NMe$-*2_0!#ucr' + 'HaO@8m51yir~pJ#p1Ipr1iwM`QExjF)#G$3Cz8hLZXb#N1Di(Cztt!yoQ>qB7N%L0i*v`Cm;~J1)-;4!xQIw*eeAT~9OyJzaaJKN' + 'MZiAfJHJn6x-3==?i0Yb5ep~PIswv|!hq&QS;C&k6g?I@qR|OtSB>~?uB7!IS779(kfRQ=*Sv4$ta+zI6sZ9m9Q;w)o`-A+VZH5u' + 'UU|aIenHB(R6(D~8m!MDsHacRa19Vb4{UIUofCOzVapOv*rAxPHd?3<=navld1DCuxHY3qeWk=bpgD&izt+(#OH5hs3@rn~l++eE' + 'bXC_ORkS|y;F=ZcDu7GX`qnH#>ytyzG0k6kP1DMeD>Hf~I=<#1D9?KSL`t^K@Js+tf=fQGLuosSJ?hBj9wQ7nq4v2QRxf!8@Q=4R@J!l1yhIJFc}S!98li!_5;VYEu~B8=u~T(RijgAnl#U5J?4>pa}Ld@(tic4|iozxDCoUnjeVZni^5(^_guI0aWHcW`#K7LhF`jnPA=5x%{MJ*@o`vUsD-Es+=?i0x%b<1wSM37&5*WzKr)rpK%{1(fsl(B`8EOjKvy^93m{' + 'sO?0c|IIw1TL2|swNP1-ut2e?1MgX~h&EY$#uFL1+2fEb*M)!`I=UqN89-#pPlTmf7kP|;1Irk)xa%uo`k`w5)2?w1!c?njq!2(`v_lDDfJsU4C&{K6I||8ugs0_ErT930o?(' + '9+cPu508UU4wHEtUV~f5KVDvFOH~on6SX(w0W|Hf!`_Mf%G}Fq&Ga29{N71De3&WHAru^^`1X77IKgEGD;ce' + '4)+*mhxpcvzI6~GKz{dLb^hU9;0-z&`YRrBs%bm(-aDt}!!CJG2*zV;fD;A7)RahIqRRwX6A><=1*EUzBVn_UN@K~utN-z^L60<|' + '62w)(LCL}sbcy||=x*p1SsW4t40W7Et640l(^;2Gzz%!#)VI00LPH=GEgZ3TO&3y_CeOC~uGB-GPGtiPtiABBpwH~qWu!A>fIr%6' + 'sEt1cm2qBc?Gg8HSU(V9pfAxD4-;2WhJI=Z5JZQxTSxP2vH*`yobs4?I@}yg$mKt9Yd>z`xlg|BkY{_O*FI6TOE%Rs9H(}pA5ZV8' + ';_Oy$KBa0SOqeabrcBewfjzz}!!H*{Qhjs5>;5K6@r87UZ)VyiM||&~%tbd+Q*E' + '8Mmv@Csc_#ZM)W3I(4Q_>u)FUUAEX=-F12e+DkGrn)1(j@n0N*=&yqLdNcF`@_ilpu2UzYH0T8@wD^7LWJImvWL=&dO}+mbd&h33' + 'w4Z=QEarmT?Y%3IVzb1XMK!v&!>X(|=uy-fY-8RYbhfz|ymtfr!ClASX(w^aFAc*mUT-}{I6$(u' + 'h$dc5XrfCXovRE>cIt_9R>z_qpV6ugPoIpu(_K8JEiSkLzluu|3$Q;PYhB$RiF@>+-I@^EWDvW~MAXKwuNfIa$iJIz' + 'v{TcWMf2+!Y+epy(jpy9<;UwR*Qfn>j~-p$M&+_{oQzzQ2~@XFM)*hsg9Fh>EIVbNF+B{U!KSxYK%Xec6TLF3bhXa5UUIjFEmsN0' + 'GvdQ{Ah+RPfuC06nr_mBgOH)Wss6;Q528&Au2EwiTt{)7G;U#yNQYyCmeZVsrft9@fYyrJ%G*${8`y1{9*VE#JYCRoBZaLohP6!b' + 'z3X^Mfua}GaJF+2*BoGr1LWmH05o{q=Mg4h@5Gm0I^S6fC|c*YelE~mtzT>FuJ@bri0&CdmeD6133jq$9D8$@2QAGG9j^yT_U1)x' + 'C8`_JJ$Y4!hO`@*q;j|mcMea4nrQTNI_i47zaqI~4xAe04ux_TBDLo2pvM>>;|@%3oZWj)DZWNU3UXMO??lu@XN&Y*XOvygPR>T&' + '8BEUEv=7r=;!U925_iiMid*hv+Ut0>w^v1%%@7n`K9U%GGVchAdu+SA-p3UJ(-n{-g!l5*+i%YB91V}(#4!oRX3VP@W-_2sRnWn~' + ';5Q!6fh0hjwob}L1hOh|#3GBE2H1^LxiF|RlFB=5hK;Ngcp(z61^y_6gS&D?YYhLp{Wq{{=~7WAceUHEVwaP~KhF*0%sV%oh!sfQ' + 'yLS?$-n|nCdzE%b&CZR6?9es~xzDI^t=LMF#y!WBD`cKL&u#4pb7tVi82W(UFsH6#L$fa`jjj0sVoq4WN~L$h!2pB}J8nXC5$vUs=XcJ#tGfdZ%U05%HxfrN#tk' + 'fnug)==96ghGs?k{6p)My>;zypdE2&9brHZn9xOgMsWy-I0T;Do*2vn*D>`mYQZn@=y9`NXDRz#LnlPi4h~%A(LwN2kyRT7((5rej8PhTY6&S<)OMAUtfsJc%bS+_g7?&Blu98q@Zr;4sKBkaA-Vj=Mdg#YsNB;LWre{?PAT^aXIyD7gQ+1<>aGX' + '%v=pFYrL~7SgV3ZTNNK}iIExK2ezAI3$?hG{o3wTM>3-Zz8Oq=&`m=PF|+E0-N}Pc1U2F=Yd}0xDPMjK*P&gZ5Z(B>fOd=Ek`kpn' + 'e}&8%MiDF$(BFm(HApTn83Vubz{j?;SivUL%g9E4(e@`LR@y9#wcn#o6@6?)lrltN-hcewcfZWU5W5LY0oml})AU_&`q8(}>2n2?iU;8pE?zBkWGFcU9<=m8%Sr8*Ppuk93(7f9s|6JMK>#MXf_6>ol' + '6!IbA;d)!9%d5JDin4rQB(*g3o+eY(27@*<3TGj#5peGhJatNb@T;zhMrsGUbeyLg?{s' + '(&V;Mco-sStV2Aq3PX@cvW4l4XIb{KKp+zAqKr?wf@2KFH@{ZO}!+3izFM}Ss)@__5pyN?La(m6xBA3UbJW&@1i9S`uHBv2aWjE=A4}JduaLNG8BG07oVVv%`_>FZ=Z}G{IgHKvYunmf=W2J1f' + 'D`~uPWb;bD(?6nNH{Ddn|7h0Ia-?XJ7)f6Y_Opv&bnG3?E%d&-*!I@kv42|`lf1D(?wSvda0%0I*brLX0Ui%?KyCNt_U++Z24!&r' + 'rDAInT9&bWAv4WFLR&J5H)!W(EFa<5(;heXMWg5AZzG#lUua0n(HJ3q2rt}w>J%E7`0X;Z^`QISdP)b9bfCYm!C1G0Nd1@2X<6?@' + '8;5oyL(Ax@eWMB^PqsF0D>0bLelWk|O=xIgbm+_O_TRE7GLTD#be0}oZaFa3|6db}$!fVJVoDKG#j^O?gQ1&WJqDgcpB{(8;fAHJJfInLW$' + 'SRI&+JvRe(v!aQHVCl5$c+Xu|=Tz<#0T@DzQ>eg0?M>dxqXpoV#dxLW$&!9dyCn7' + 'b=>Vaa$FauR7sDmE`rV%807Y`Vt+-p' + '2-oW*kBO(YY`%bBP5Oh%Z&FlBk5sQNqr74}JE?BSd+`7Y;}4=ZTEoOQ=3R6$b6kj{F&lBIk#QpMU3!z`O}K{Pz%6bbYBHEogSx|}' + 'c!dZ(+oYLr0;SX{s0F_q$qP8dZ5&2y1{!TtTi-!DhnMrKRT={zWqnmF(Q{H>SmQmfF}r@cOvy7;^&s!}P#K^bD?K}_T23Eo4uH7{' + 'XK@_1e9gmg#zqK-jj@P*AY(npD>$jWri_*BpRwcOq9J$CKG0{78aCcsjU>bZ`hBLx(0>9X*PAa(;?~%~lfV>tmjs-jR~->MZO>S)' + 'K2$*TPjLSt8NhpQGlt=NK@=hMQ*pou~muUOC' + 'IH!LHQ-G9No*!)T84PrYnt~R~{!%qq7BD2}=I~VeIthY+y%#BPf=lL6t|xf=E~z#00jrzh$%IHXZB%s}NK4zU!vY{mcZZv_uxp2N8)6mCaC?lV2X(96#=*HoSu}tj3=i)DSn*-6|INeN' + 'Z&Z>1ZK;PEUgpVdD04dBlQfQ6zGUKei{`3McIeGSt?+iYIbJr_&@0#NVCrTvDyv0|;&q>~S1`50-E2@JZGS_n=<|TorFUzOJlK|+' + 'AHiqdWs-QCDq51oxh|i<0xU7_Wt?E=fmZUWuGiJs&k+*voHp14a=LPyar(&phE$_OrlDF0_{-pRE*QAV4mhz(p8zGgZRqcmSS2lR25!JoFeVNMW@yrbF&WRMEf|?4c+6^' + 'ft0Ae)GruqWf#1Q$eoFz7ouyK?KUlM(T=DtlLh^$Oz$YKZd_2I+DCd=;!+S7i}WGPqY4)pJY=4p?YXM|QaBu01a?Je?^s@u;eIun}q-+N|y8wFa^' + 'KIIIUAzg~`&W?E`uIyMuH>_!+g?LUfzjbmnP+vF}7THF<0j5UzN;_JUKy&w?4-EElc@;hFF%cs?%V8{>0js%I+?(#-Y56C*vcPsY' + '>s97iIdGs1SQ%%pHvC(6p6iRt7xYTLd07c}xbqO#0TvzM8{yQl+sgF6TZeUVoe;%--|z9(WQsaGwq57j-fntcsa9!+!=8#yI~IOM' + 'Q{n%y@O$0%pv*5xSyZCp)QG0*^RH_FJksI&&~ORFI1G*C)2(XA|7jT4&;&fZNY-cF$JBrU7DBcK4@}W4y_zRer(nXL5$zZjw;zIP' + 'NZE@{p^85)K=Uc0*)9yvA9FkO5XVuA!c5Jp>`7Ix=k+tgQgvIG)c*QMN+Si-p-d9B-F' + 'u6(2cpV&1=ZY=4E>B2BRo*{hD(#$oiov(p^`sSPS@58sRzW?TB8;|&~r)@@a#rH7$6#3W{w537BVZm&tQ@)Na^3)MCF*qM|g-f(A' + '3~|M`;Gaj2nze0(MNkw3kSEM=zAm!u3Z+RTfb+U?IHD6v7KwnXeXb{HJ6Uw~2DAS2CdHsvc`-y*L@{2pfoh)~27>&j(_=zwIf9Pl' + 'RH9{A?!9IZ0p@jY06{yjap5Q38PPgzxBbIqhuxXAM+R7u0PO;`#{t=FjKerT{yr$EeJC3C-Z!0~8Dh;r^CaWsg{}|M4R^s_cLRjV' + 'X75y_Cbq6*l0~I6j=ed_hwf<@8GTuoAf=VWH)usQ_y_MRw~x0hOiF;i?~P+z2W@WBT|=GRTTLGMzI_duz=9@Zoh#Z?)2`R_sw^Q?' + 'f0L&`c!pZL6C1**HA{^7ErL(mH%pg*tjPQu{AhcaXhku2(I_pj$x*>E2B_+It6VhAMm{)3p$>Za(Q&>*r5b9;sZ~NvbVIYLenlo*' + 'xajqnhj2_diVrv@FgA`}3h&VHNW)dX64aL}fLx*X~S;=+8`%q+$`XA4gQ-sO%YkA!E' + 'UM%{Dk@wBrI^np(_Y5Mudy~wQ^oHI=(T9P(APp!7-E`Zb0_&i*&%|Mgo4uspixR@3gI-gHoP5vPP|dceH=xf>%uJWvPg67PCC5%$' + 'S7_FSmgPx6AyNSDa!Mdz&b*=Fe6h~rxVCaC`ySEhHewulMmo84Y<&6QRHy%uPpMC@Z6a}q}MjD+_vI%jfF+dXYxmx1O`Ri6E7Y7#}6{-|F;tIvO91g5#|JkZsAiW' + 'bDghaU4|$M0|coGppTY`7R(B2Ad0Rcl$8jMjC7ama*F&i0|bKq' + 'J5$_t-bm5G2avEaznu0}>`9TzEx^B4+EZ87CjviQZ|f^iA~@4t7*+ZLYOD_UK$mC9f$e{JuVf%*hKDxZu%`e{NpGSDhx2S#&6b(l' + 'EzHZ>BM(zN_9T3;MO(e*v@_d-wHJ2)YgcT0xayrAe4Jxz_nkPrJ+dTkM|lgRaK!F5rdv8MF|_lMQewjT3fBFq$bj7wSpp3GU2xQC' + 'c-OQk>QC9xEro(;v8E}t`}3u9bUTHXY-6GhXJmCB?Ma8n2U?O*%>Ee0X+++ILs=EQT-q`D;-evHV0^yt#&7581bm-on;18R!11cf' + 'o{bmCSJm3PEy`;&=+~Hw;r4OSY?;b%5m@Mn4BF>&t!_=LmzC8(p&F~SVP0re`!12&(U9i!DDodr{zH2LLk?L}$N2qohPVY4&h|>b' + 'Kmdtl;xr1a*2nEvd!=Z=Pn6Tq91rI#=|Tx@#S!R>=igc@Xv`;qeK%Z;JZoM#cThOXin*X<#%-UU;V2_ORuSM?WPtu92+2uY@x6UH' + '3ZrXucJ?0Y;u>Ihp*2HU$CMDuLDl46yMP^dh}k4GgbWCr$zTt7b-cG$Azj?MBm<;6^A)!1PXApIU&rXrLl' + 'Xbg|f;D5F>kDGG>kV;Obek@JO_{YvZKD0Q1SWUE=!Z|ww>oQkR!{O{FCT%#fbbu=9?1Q#;m^_+qzk=3lQ*y|GBttQf2rxPmQu@' + 'mSi9tZC9a%Ev+SSEWm_q_|Kn(7q4DkzMXoH8a8Q##ahc@Dz+mRP{JB@51`zE6b$Jep&N4_Y9Na^N9>U$5' + ')nydw->w5zDY9RsaT^&y?SSg8;FOf(d{<$kPG)EX%!YExiGT3pJi#S%;?stp(W2yt9R6XPBxrLpX^Z;mXN00#80PC`6N' + '79yC+?vMOK1f_XDRiJHxu$NL)0?b9r3Jjt)rJtwv=x7=lCtfGNN`C?DLn<6|@zi+&g*nsA-`>UmZ<o5LxZ^FwyuH#uf@5qj{P@GiYhGU!Dt|*lEUUJ9o{J1#mYn+%HQytfXRb0&y{r~z)Hp2;KaL-<-*74Up{e;~u;&n)_)W|Ku;BG<#' + 'NS+DW#0iQR6cRn?yGarU-*0REfMy+S(H8+P`>vL2s1wAS)mlUWXmrkK4tWl{bdN?~r>9x>' + '-#tElAzuOIIXTi^fFMVr;x$LlkWpq?6fsB6G=lLJrZlWQNuN1ZCH9aL&`GOmJx;6ZfTY7}!@5wJeBKz6qpfK^W)Xnr>!8IscipK&' + '(Kz#ff5gj1ePoX;pYt6xSB##KI2FsWSBHDhK{li5o6Ad_@&5' + 'm=+vFCzxf5K~$`%H7OgasFJ~Kv#oRn>N+t1Mxgt&xlH%1TW-p7zx#qN&o9q4*GSz}9$pmLO(O2O&|T62_;8JGaVlL3jfXnz%J_|{' + 'Lo-rK+rQkbJZSKc8YGZ_BNs|18q886dt{deq6v?4P+;qz3>P;E}S|q{7Ba$T+gK(ah9sS(Tk5^)v5{DD?bU' + 'MK%cN`jBh`Nla=Z?Ake+G|+Uh$u}4iicay~Oi!~dgdC((QJj<|i1gmy5|xdEv__GSi9PUWLX0zafwM28mtKLS0mI?U4)Bhc8YUsL' + 'eEG5_Bww4yU=`h2+fyySoS3og0N$GOiBq`NR&v<%2+N#cqV8&J;F*t}NTaTHIKaR;G>2V&+QSCz(M|r*PS(UBTUywfg@!bhRy#hr' + 'cd0qv4=p89(T*rKHDG}=|E(iygXLx=5;s^d9dv!CxP6UStwVC}k%4H-WtlAL%t>8Gh&2X%(a{$IiPbi*FvliPNv$=`bCy?lG2ig=' + '{IzfRPw`yn#X-VHyG&)0Ds}}bR%uU$Zg6{j' + '6&(FW1pW2C^t_cdl(K~mcpIWzDbqkn8B;wF$r-eD92uYi(8}Ffp3J' + 'F*hXaI6oq-dZx2n4iRcAq$DBO*p)Fk?&jy?=e?`FwLHq3_rnQoUfPD^!<{%$CHa=$HYwxpGvvT42N;Y_LV5zmbB+^zPSt#j@TGTpl' + ')D@HAs4nX5Fa35e2fHUEFBnK*B@cMKwtIWR9f6wz9n2y+xw%N#s0VK`K4}teWqytB(wODjq?=kosNAvl}`&`V^' + '=alBn(Vl{_' + 'rk&QX^N|-yxWd&9h(Zi2Mh|3TAc0y+Ec_fYrlGU+m%uxLuB}s|)2IAY{an^~oRFcT)G~;ZdU)<@p<${z5(YhN_6{uzyn-rG%#-ref(lLzQV>6*L74)PR&>)fjZK>$blm#xwu$^n2VThcn@zYnDWM4K{tIrR%@eG)l;{riSSVTP31*JHQwij' + 'qe=5boX^4pCJRtsQ++z{zzMaiH+TaRV8ffq@f29-n+bWr(i6QsGTxq;Z;z+G?)fd}R(KbQabPd7Ds5kR#)L=hisCQ_E#>LlI?Z8}' + 'f30WZAj}_;tcbTykB^Vwk8$|>ZJNhu49L^fW+gEX&@wBCA~#bBD5BCD^S07ypa{gnEjSk>vq4G10-4J)s`T#R&+zVW?gwwZKYHiz' + 'm!5EXRG}+Sy3uzFgVXN*gc|tkX1-qIG0Fss1Q@zzM$Xd5;H|!O(ECh)-yAW5SVbkxt0?}W=`F)yj6xI?)K>K`le1VP+Ir1D^QrbR' + 'Zjx}(J>o-AU4!O1yW5G>u9X&o0m6_m2?-gujx}@FMnAF65s5~;H+)`J%piwr9$9&uKl8r7iso187C6Iswk&`hT&)Z(FQkjDF=lKi' + '5yz6;euTdN62t{jM~wa$ZkoKUz7Naz7npXv#tVj)%Kn8GyV>Yzelz@$D*4uZ1$T&)Bkw1X$PsqS3%0sM<&vLF<)+4b6PSb-D7vWa' + '9eOeRJtn%&tB87%_W9U1TfsD5IJv`P_^6u1N}G1M$xpD|$4xQI{~nj|3w<4%wQRcAV&87Ly+L' + 'skV6&#VZiE_$Cs*%XFFO7QndKA4uQ`!R|0f8E+2JKA`cuXX6jnvB~matfa%PsFVTiLmfWx7dbFIy8f95yL6`O#1ba>Zj42ZQ~Q)TSFz^L16s=mK*f' + '{AWgNrxDZXyxy1_eo$4Eszc1-BW+~P`;pI_HvjXC(nV5gZH{PAT3?f4os|WKus4Y)7Xh3woU91>#3!wBRyOOxJnl25@C7~av4(Oq' + '#hqk%p9T&nL=$hh-+MOEno`b#IwD6~49L?nqLNl4Mgt>NRr-=?(w&-!yXX9?9rj8Rc{3U9t2idwSIf#GCL3DN=E7`&KTNdmRvWD;' + 'a0c&)O(n)7@>|m;)wc1fiXXP0(D^u=UnTSFuuK-1h=k7Y+SvV1-h4u(RZ|t>kngbz3Yya72KmVe}498x2G' + 'J3WNZeOi07V>vElaZ18V@5$Yp_O71Bpf3J(H+#ntzDZb3W_!q-T-b}fwz%m+$zYC*d#vt%Bip5L!AAye+=R~ekH;gH8uXlz`kn6V' + 'jbWyg>@bx7yAUYnvW*oUP3f7Dv#JoAM5' + 'CMVgT4a;IDT+qo*mK2@k)#awX>bUskpG)lflNvb|i-twh%p|{ApM(0;huc{Az96|c-+1v2#<%1+p7v1O)L#c)RrO$t9`oSX-bGHxCaPMqJHf?X)*w;+N*w{(!Gh=k-m1n5r%(Vv<;EP)vC-06+' + 'pR9X3`?fqy^!?!#28Ui!^XouNFz|6CCuRDmxo*)lc|VT44%cq)zjx0Dwb`BscblJtjxr_^CFP$ETH1VYQPa-55sct;(|Xr`U4nkW' + 'o0(GX3v&M0tjMvToy-e;CkBK4FCy*&>LJ$>k38+ZQtdJfK9$Ajq+jK_;Q!q6`@ma%|NZay{p1;JRKGGayP1+^pSywS(C' + 'ZE$R5crI{x?S1`|+sKXR@A@kkx#~>X8hU2DStq(uUTlxO=gZpOTRXS!o}ywXu{A@FL~2QD##+bu+aG}b{-C7sI=NS`&Q-fZk=zZA(tefb?yVo!4RdJpUhr_{O+0^SO%a+@2+psK)' + '%JrshI;d6EUETqNb}$g{%33{jtm*1{)yfwQv{J9tFD9RE^QJ1R%a%8@s8=htAfRbJU&yvz0knL+Vx#CC`+3W%1>?oBPm2}ysa0sT' + 'MY-rk`qKy++pKa~Wx1^ugqWyoR@V66inN$*)EXm(5+~$XgP@X*SPWg?O`Q=XMFR' + 'Td%5{d{q`N;nj<#sT;ebLOm_(bW>vMY-9fEpTF)t8#z3c;7H4=42}0vUa`1`IQeC~y%Bo-=jiT`G6@TRxN2yaT^lSoYwe47TSpZpqlJfVefESD0t6jEANVKbR0b{u9uIehQ=O19I' + 'WqB*+fm?6AD3{qLZ(3sp#)Am{y~a;JHTed_!HC|xsf8>hHs|k}dRr9&%z~J*9SsthWnrtx-po1LOW*~O7xJ1eJwt2Q)LZ!jr-K#g>eg7i+{>A(EFMiCP|J&17Z?abx@1jXGeERJ9' + '#fz8U{oAX5|I_!c-@N_dKYo1o{?GsU)Bpaj{{wQmV9U#^^27CNUDcbPo3`8D+#5z)BRY*B}TY1vb*KfF2QyCGZCSQa1%__`-%z5Gh&;i|i7{bQ9D6' + '9R%QqXcYfxk>ev=}eWvbre>uHMUPMVJff6E*_4CiH1c3wv2`+!w6PWEQ@vQ?_QfJ@@7@USTym9!Zh#7YRjw-%Mu1<6v2Af' + '0h%$eR}8{?%M!wq9%bVa&y!(d1GjNtA3MYWIv#`s67iD7X0s2YGXr45Kn7MeOgI1+J' + 'P;M%q%mP!uM9o@TW8<4K9%&%UhXx{Pmp5whUs2^j5PLNOa^6bLrMIAdJ4BlpEhZMT`XGV#{-`wF0sUXCEL`t@&AJ>*aPl&c' + '4c#tV)%7(>IU*wngAt5G7yS*T-pp8X{9>vAt$br^v4r|>>WTry@eh@{7rnnKVXXl@gn&TzKn;$%s~lCHO{=#fF#SOzg)fYzB(nLo' + 'DnZv53@X5(jKa#YUJ6>ss=jTBtF~=^$*|^M+uO2RTxGJc`ON7`fXEm1;`%?{f%33o?-q5#IC){r${Qxe' + 'W*iv{mm(Zh97oGlop)lr2I;c|e#n}2YkmOr3MG*9b_s3Rzw63*g&1BjtJcdi^B3We>34no`;oQVRmZ#5!30@l1-$;*1SNW`IsvW1' + 'v%y0!%(%-j>8xmB)}qf2!A2gLXMzuYDK{}!M~jrVgEBl$Mgk%0WfEP&t2WxzTVk+HUN!(yfRbQ<&PB{?CsBJ0?mWeqq(*eFMvBa5Ag);X>AlfYYyp2*Hgm;0AMCs5(njjd04tW=Aa(9Fo?6%7U~V' + 'dTzz;TE7tj0}EA~R4AZGHu``I5)3fTTrpcEgw|jIzw#Q`ssuE|URPkhr-f81TNDsI_o' + 'Ny>MK`NafyOGP6&B1pMq`H}G{TSae1(etP9A9b)zqgMz&Dm&B+QPGCka04AKl)i-j-r{x=H-yL}u$M&T0jL4X{OuKMSp0YRQMRx*' + '2C+g&lD^YUhi`|oB+XZ=IR99-6Hs&07G?qa1$9jW`VAR>P?n5I_rr(}_b2^Gn&eGFw&mNDW)AaQ0^f5@k0=H@tUYf+k' + '<8!T_s4hx)$TEDlYa%xtzCr!?`BVD;a0ZwXs5nWt)f_q+7v&mo)^HWTriLy-GbfLYuT$89A+DonJX0mBywV-L7k$S%;@xF~a}dV2' + '1hE&*cl7f4)05{XZ)jk&(uoC;+|u;GJ9F9OSNS??7gy`D=wj;bD0-u@;S7`6`W^?Ck~$tss_Q!LT!tEoT>EraV~RUf@Rcmx+sxGC' + 'u90tx7p=tw*amkM|0Ri0(LL+Rm%QTF8c2K|4KExJ#8+X{aFu|QH)q7' + '}BQN&QzcAzU=UU9u#e1av3wblI&mJ3vpDpuZ;s!S8*V_`Y' + 'dk>PLVRu|N?%01r-FmBn1xPjjV3WW))7SA?HlrXSb-U8!6XJbj|1Szd&7d3~Lf}QNl$b(BoKIS;6^t{fm5xbl;%iFSr' + '{+Gz$3nn(7u*q+LCRC`3etiAz!rx}*tHqXYBXJ;Z8B+GP1{Q2VYvdc8T&!-=hMceHtO07lFy21$U`iHBO~(6Y3qoU^0?DziZy2tU' + 'dzYwD1k1~5Vyr0d7!%t&9QNjdh!UZ$hUXQjQSz$mHtqP)qf3}f+c{vfesuYvuF6NU)}#5VoiC{sdOg^YrW+V__AFNuO--' + 'r*@#-DQPR9xy5AGA`(l59CksINi_U!C>Cy1^UPJ-li%Y*^7{ERz995r6I^wUhI;HWa~#0`$JUiU(%ub((p6okg@#m~dIM{HTW=e*' + '53Qhu&B9gS4f8SVAKqTPjsN%yTZ3!?v02~_5cTb1B=@id167>42yGU7FyT5oZ@->@kVbABjoIE6Zd1(Bq>DyJfRL62xBJaMkj)QZ' + '&P>iW|wB1vwqW#GbqxJ{Ow}x&3uEy|4h}i>8KMb5Y`f2NcBfNHRXbTCuI#?FLUpSTVFnk%b(' + 'D{CX=SeLnxluDQ4L_0?6n8pg1PKM1q(Mn0y<`wSC`QfwbUIrWJjBl?9IOpRU*2!iM7pVcjZP#xm^n>aop&|HBH1YERkUj?HW_;U<' + 'RbTDF0o-MGHB8+0w2ZQhH3|j#PIx#4z6w^L8qZ-%2u#8UZfH=KOhKv@(yYB4MgTMB0eTQ$^$*YzAW_NXAsIaXF3' + 'S|9s1;VvFuRJE>z4HOSn=lS9q6h0=bQeix0VLlEZFkl^B!|YMzNt2}F$vCKNJm(OI{;2kK2}XL@|(_i4DC1X1iVHYtzt' + 'BWGk+5}?J3&l4zFr}h#P9>M3bIBBNz9<+2mg4Q7c%yk=}HOjAf1zap5wK2k#U#iaHWg6(J0P&u8?T#S!jrxzQ#Z!2`Yj4reiUg8{' + '_9~IwrYRLcvJb}`TQrIAkL9wLgCYiuMs)`cOf-}T7XwEH&`=G^q5vmuCxD{N)F?0;g_s9`7__y4@DEety91-2q@#hqcH}pg^bPtY' + 'dPp|sfvW*JO0g_IvT7Pjfbci@(pBKUN3h>POyj=?53X;au~rCfROiSvXtiaD' + 'b5n|ce{u1?M~Vv>+bhSnh}2z#Il;J#4FB&Fxe5k1@Ho8;(3d($#N>ikla!GGz*-y8=29XtJl!{sS~?~ZO@I=)Va<6KhiYhGfsz@;I5qt(u8%K@Obk*@gBB' + 'Afn3MQ9efS{bMjD9N0$?04=wb{aSK+a#q@p@i5c0?#X}ua+n|l)2Brk=2wyHoFBQnB0UCVyZ+E8-T74ErXKFUvfpBe!i}OJ&`7q$' + '&4o_J!6Y}i=hD>9qxd>Av}8w(De<-H9V^V-fGLa^6wZyt@~~C^R@)Y?W(>-aCp+hmfIBx9&{kVZ3B6Wp#^|j' + 'mC$jrHzTPr28sR|}4-Y{{HKZsC9Mknu8Jg%c2z}y4' + '26PzSUA@&3AAKmu1LkDKU@#l(Qo}Y3&9=e>^6CP;6OkIXCD3!3m#Z}TVZ|^wV*wO4pua(<$W{3Z(^WwLFHsKj0ld9p{4H-IB2PJ{' + 'eQ-Durm*jtoh|4$=olt{i1Wv1fXfSd^tY$Jt#-(9fe!Igj<$ndzZWXJlIrO&;jS>srBLcG>5(Q-;nAf#hf~I`' + 'Dx8$sc(1LU`%1O8m$y6(EJxa-mf<9B(ARXcuq~Iqiz~238U2CGojtzyfl6Vk=li;bDOv5LPZhBIZG#gI9?-~*E=0p5gaNkBP1BMT' + '0sDJ@!#^?F%}2s`-pY-LU~bKoNGM{15E|O8-9cp#z!uL6_&I^F72^*dUk}&i@M6nD0nFI~f1H^-jy?gz&kmkKb}LVHV#P2gkuCjD' + 'A&)i50l<$IgY^I;Mpu-L-9>l-+=vkzF@Cg>r8&~o7}FGE!AH@MUcqmtTf$?~E-PO0(zfM>BHhU4TL-?yQ$E3IwX_>kZw!z#jlrs{' + '$ZLngNehODy0P-+|4U`Ie-HFfJVB=A>5AnA@>jhk1zb}b*6@z|`QbSy`v|-uCn>gO&rxchZ?51erON1Ga1{?<112jo3PR#af2T$}' + '9KRjMB68Uep7B@Z_c31U3tgIRPwC&#Si*BisNrfs' + 'tLl1#=YS-6U=?2o7eLYlUTZ>5yV@hIod~wM?&i%7a?YlVmR(j}$|G+TRcHtOugy7C$1IYUt;TqIHk0y8#mYs)9uAaOBuxMr4&aA`qf#5gK1D}dwyqNXOCP=!=XW&xXq6?&s9' + 'l<+H9UPt8U!rzO+d@Ck)&Wn6gcDO0y{(REgB;cZm;XE1`{tC7%qI}Gf^deYT#Q*5R;zh9V%QIc>R`pngxoRtXSR1;2)g*' + 'H{Rals=Q5t!#j^4ZYGv}w_a&sAHtK+urO#L580CykBS)adNg-oMOd$smZ!9v>oA' + 'EZK9gs|rc0PZDfK9))sucmuMpuo^L6?qQ5tq9=})F$!!YjA*dQ9%aI=YDJ4pJmwhn2(ig0k6q)^74Fm18K|A!+p}3RHl' + 'm_oK-!T&r8ou^OqScE>)`SQd(lO7EM6GjM1-gOyO^Oo`$;AG015`7y-l;;3wB4TNXg>OoB+iK@9wC$oiFS{(`ThX!SXWEhtrV`GM' + '%1xA%qo9Y?dl0!|$1e+Wzba5;a#E?P%hXsOWuc7we?=4JWTt&m1J-|gJcq8a8q|iMlkzB74NbaE4SeMkwnD>XaqFd1!HyP_HSD)e' + '6&+zEi|aQ|Q5AA0OR5(xOd(f+XQcsBHacE;{uoAC0GLGDOtumUOmcvij>}P|U-fY)D>>{4*' + 'BD#b&KH-8dz$a-&=KC!BoDRLME-wSkmus}5pGGn(L7P(C8pr}P8ado~gzYG(9V^^HFAvVFTK^k~ik#xg5q`cVzoweD`4Q_Pn>(6n?O`!(ZH63U>JNj)r=vvAtUd=^zd}i09z@1@*EA' + '{bx{<qCt`S~E<{cXG^cw{t2GhX+nTMma!!wue@fFjL_%MIMo<4&z-Brio8TVA(?$JF0)5' + '))Ic?NO&Kvl<^~5#(M>YlpmQ=-YDdYZc|hR=S%Xdg`zuQ0+U0|3`=kx5_LUXs40GD{?k>D;yf>G+6jcxA566UTH-j29AHqJbCy%D' + '!k(42Z%Vm6bJtgXJgDQLEkfJldo*UY5{!N`ZgI_aVUgI!x5rCtqu;zO>OD6PThaOoy_GULmi?yY*()Jz6L`t$I(y724ZU=l1k7DgX;w8%-n9TPWI{K9c>-rJV-E8UES{#' + 's*=LY^*MV{M`#MAECy&X#XJ^`H2V<3@cNp$-a{s{8te(=rIa?_@hzSRP2taMUqu{|V}ned%j}vW;8h)h(%PbE' + 'i6AL6h48EzCVNruo#kbl85H5^30fZVH*G1v?{7m|7h7c&ks^1N1ZV;O2LL2ODJ&Rx&uU^|`QxqVM2(-mZ>u(6vW$2Tk@muSL7N_?5FOrVrpJI!IL7rwfu%$@_o%+x^H$Zz6j=uS`-nQ|7u3m(si|^@kcnL7Lvbk*>Vi$A#Mc^p+DPOBMt2(;6(zUL!Qs;u&V4@f$@Xpj' + '`@4(yU>~cINVe*KxlmTyw#nZ7ldh4)nlQ@n!-G$~WBifxh69->Xf`>eDmk$=OP;p_MPJD7CX35oeWkfxIWct1zLJ)XJ@M7J{H9zU$|CQCTL@6wKa' + '02F7oIc8GLo07E>q??ghLZ&b@ogK(XD(kgjd=ryfLYw#|Md9IjeoXts311BdvQKIU)BVkhoP4B18Dt2#x67#WO^QcPVQP;ek*&z}' + 'DLsIIhM;jU(zMEKf-@sy>SB&+7p!vw$7im!eQ|6C^VXiWf2EnrbD1(3_%Bs-vnO+>@nVka(nHKRfsgn!OM+|TTrTK3u7Y*q3~0XN' + 'QcY!up(JCdfto?`SajU)Wl78R2vkfAHpC;Ial^?g(+2&vX`' + '&+MU#8%k|th2xq2D3OZvi(6)ebkVayJ#<`0n3|M}+(_It9e5t)zSkz%_CdO4^-?&DNv&CCPfN{An9>aHGNuV&GEx20iI`mDg@tTr' + '6#m3^_87dnh(1ZUp9#hi2f*I0_fyg*2Ii9%tTPeG`Ju0M$*uSryQGiV(g626kPX)JlsvIZu-gi|od(=EUpol|X_K`GP|Bs!5uLD5' + 'nx4*#HPSdBRm6nU=7*fyw=FY(!z+dQ4D6T(bX7QbmeYma~3(SPkZ2^y03q}vk;' + 'd`8-cwK^)hy++dh{%TY6Kdr+uRp>vZ%l0@?(L{CA8rL4kJju>y{a#FbP&StsbC+`M!bKsg%SFE{H}|v' + 'YZ&vk)AzXQ(J}77E}fg_+Y(d*eo+k1REXCF=tRimm0A!)?h&VR^gK0>vDzs{Kh6ffQjwvr1o`>XK5fP*Bp;_{u3h=kNV0q4T2@7u' + '3U)S{gtV9J@xyy?W`2vP?C)B6KwyT#*k**sM{9St5yY0TG(i|x$A^{{3&pD{k8R(S7MABcs(MOI7B3Z6~VNKn}u`W0=z$7wt' + 'Fxm920gPeS&3g`o;t%r2c|uNaI#XG_o|imlvlB1v?ACoY1-R#Z$9K>PK%Y06g6L{?*&K$HD#ij6NpklT5#&@WM+XN2^Z22DljE-8' + 'ROH7*m!;xE;x@`DH-)}`nDMG6u%6ODoN`}jnS@7H=;fq->%iUE%i@zgR' + '?4rkfb}X;-k7r(>fM&~*_O^DEA6oVBi@q{wZ_;0MjGz!{F_0pQt<<5#!7ESm$^q})SMefIOir7Wh(Wv&0o8O15xmt&hF@K@Y4k!Z' + '2|T`d&H~=}cd&4p^1JpaDvIgZ%E_b!P;g5zwcmY;7Sv}{O1tkA!M2d' + 'yb?+qTrTof`zQxFjP$j=whnQDADQ)^sv8_+UAnfCCN@vnO;%Ao&kaZCRmMwvF(ZrB1rH;*(}j8aQkGk&*M|8usnw5H0#7F)VFr_Kv>x25+' + 'RCy{v10LFa6%{?EUX@uiiGnDG>GBW>;|p3fPHteA#?nGQAWP&xw~kmF2&^UWU^XAZbSorqbQD59#(mz2+etWY7UuJz%g+2x*3cat|uE&+nt-(x0{?u9jtHke$gVDb0KFi_s=ovK&Wp|TofQLbNl&JIpB8m8mt8!vwJq%*_ceg)S$qi|=Az56uPgLu>=HHYsj' + 'bk*X9FY7Id4d93jQ!(@GO)|jo*J7vm?5j(TfTE|>e(s{^OnrLWCr{{O!MeO$g`~ovs5&yeZVl' + '*oy%Jn}%ToJ##LQ@O`#io2a!hhmPQGFNK+1#E3vb?m24b8&i6v$r)_-K6XtGtjtzLJDCpW*)qpWX+vx6I%F@Lx8?*k>=*%ai`?=|' + 'xu_W>syJ@pI8C!x=UuLi1!R)@2C#{X;5jnS%twS|$<8AY^itXEq@YcZb0=l^&+NXG@)Lm%Zj1Gb13B3qhawuM!hB^p0J}ryjfE' + 'Cz`cJe=3=mR3Dhojb@V9tds0_C7KIln#z_3jh>x&DI`&}6}1s)Mdn)2k#jO-VtfL~H;mbId9QB<>08Z0#=DVnhETj7pW!ec_zaXEeEy%>V#;*}gOBh9Ui#u@;9@$v#jbc&K' + 'Zs>S%xbER)T^GTlBaCjvMKT(@UO}O2;uABd!gFnm3>5WHv@#wooIOMN4oA}Rfpx-3*oKaY%^G-SA^L@X(A-7rsuQaIVREk8(d+fj!YetuiLoUKmafHJHV=r%l?b;(n``4q847BBaWYWKf-`}TtnjY#Z>ok*N4(mAN' + 'mt->8L>(d!RXPwHJcW@s0F?LXpz|zGb2y4Vv~>mhu%@KEWb+jh`49a7V}i#Mk;1Qf1404hWg#4ZNS_=*A93{Q1TLan+$15Eo3n2z' + 'q8x8fHHSZ!FKQn06+E@-^*Fk*Qjy){jqtA0' + 'ImirFvdHNkD!**%Usy$Z3{tNtojQq%b$(6p&t)Ux?bi%b5r9}&ueV)J^oV|{n`@kCxdV37ls7mEl)0m=DaSc$JATcqRPNOUU7LKj' + 's$m^la7JCMSWf5u4p7V-&YJs}p}Alu;<7b5L@TQkT29hFr#?)rx#*@=;L^;!8P=(TcXK|b%SQEXz!(Y#7`kcb--A#=9kMlZ@G-tj' + 'ebAQFhtK5d=_d*p?ve5)P@{&6hxutQBEF&wn$W@b8&l!+%sxZ#O2!i&%X^R%N|Ue$g3&A2!3Dfapr>2f0@' + '%XuC1?)nf|@__pG8y;vYZkV`ZD^%<3+db`M{Sd7j_XZI?xfupwAk@`#b!8w`s68O|y82p(`@rBR*m+A*{wA@fhgwo+@QjNE@uju@' + '!h0WB%T-{|aDmbAvS|>V!$ZF1;tU;`Dh(Lm3wO%yJ!<0z!Q' + 'qFj4Gi)_n^2e?r4tURx%GJUX' + 'kMS$8-K;(#vE9Gg9I;OEm*x@_5b=UO<38_mqg(x*;vF&!ib9MgeNd0+tGS{byCNEIUERy}*&au}w(SY@(82*kyhf5uL($xJGOmLi' + 'QK9iSgZ_w=9zG7FQN;0eeR9LYi^)idieB;AQP51m-vilP8TPc|l=ZUH%C#we@W3hRy9cgXUh1GYBSB4qtgJTMj$0iu-v*hO%Q}6A' + 'V*1rv_(1~O-l<^LqWF)$tja1GM+ho<7_~gnI>=uXYY&CYg$LS}G_2U-Sud#Jb)#M6)D6&lEBo8??;f4KHjEzUr^XA|i_mW8T3l6x' + 'ZiTp0jh1E8>P<;|g|2K}B&9v6?P$7!6#}0%^rCF9Q`+FPTYkB4!|r2dBr<&6<{+8P9Sm2~IK2j*c0*TzB%FjHa|4l4UKv6_5;EIm' + '^ubYB2%;W^eeejB>F9Ng-epxbFYoA>u@d7N=|muil@xCZ5o;-D59fkeo`B`5uCJX1nqM(8a@|H@Sw%-4T-heh_17ULPu(S%9uG2o' + 'vuFj=^Skav4E!}?n~a+c_QdD9rzgv?Y3fCzAV5!-40t5w;87TYI+*n4Om{m(>SIhsm5BslAU83d?K23WULEQ`hi{!V8G4v^_F$M=' + 'LuXIMv8}R=>@C-SR;A1Ijz#E#^B+affa@~M*-tpTRy^E{J07aq&ZT2=>chq!G(t3#seoaX*bUu55eP!2f1%oB<8ZpDQ`#l&K3pV2a(2eeCcUV2QV5+JTJY=q-5-$;2&HG0M4?xf}C{wAHZ4Q7U!S~ufYOrg<(xRJtRm5fpT<%FP-g~ycFu4?' + 'z%1!UM*o6gA1LmsG^Noakju9vqQlfLbpO{>|yeWx`uGM#KC@1%Fr*b0Ilk*kqD(F=t^ifliC2*A^T' + 'c-2rocy**tDRwac5A#U*eG8n6DsJW-*8Uv_dmaD-s2HP5ZVGr!tSK1{`QXh6P8n0YOoJ2^gs@#MI45wSm0K-o}Brg>>Khx8wi&dAMDcnLK1n!->G-3ZXc8qGg^k6m$H)b#;S?rU6-yy@vIdn8?IPUWyh|OUdPpnb%Fb^m8blTMsB4IijEHE|d&F+ef*Od^^a45t+' + '%4?7Z$jx0F_kk)8KyW6UT&lfZehTMuF)jC{*ofIfR#{-|CFN~^FIbXq!Qh4jSFNn3FI&_bLLF^d$^SFP-ox4bTPg4#YF0<9XcSG9NH`FyB-Pc' + 'i*BYIF5p3A_ovH|BA(8z@E)aBXBA*(0_i+d_+|qt;3MCn_zHa~ZrIYO7y^d(pEJGsBdNe08GSW7x+{USah0' + 'KxR_3w}peAeHn*s?iTfqiu=v#-$qEPewnX2crzK2aX|QYZ(Ek}Ew6O8rG5B7tfdq%K)1V2I7(@g-wuS7*B`yLMP{0$Onc}-sbhM%' + 'Re{(ZO@B=0?o%DRWT4{k2%Y5_iyu$-j' + 'oVw2II5c`jgWs{vFp=epMcqitQ(hFMFrwkj_G|P%ZwD$^r8_%NLa1' + 'MHWn=e~JFCehKIQAO+MsrSv?yz8wv~ne$C8mUUuLE^jHsJTu~!jO>l7eQTJ)rhWW)E?tGi#);nHXzEwoca`AaY;&47vfOXy*)as|' + '!4bPJ9%sEL5mr`Ugle{lYrEQg^+AY(OTVAcZC?(%Lka8lWvA=QMAwsk&(H=`;$gJ>%M+)+!8)IRFM`>YzfcR(c9=Ql_jQJsS4IKFNY5u' + 'NaRwU=@UoC&GrXcj}T9D6va@9lw9a$yVhW0uc%e;jLyem2$Ep%X*97$M%`jQbT-}hqL;Mz`Fw>8N5_X7Gebi=' + 'lJ_&AasoPvlbyqoTx1yx$tGv*QAm>_Xuak&XD-rD`Md~dV;Gbm2c5IAU1&}a_?6sDUq>T4p>E>XUu%G4PXZ9I!EXxFPmE!>e#K5hn' + '7GC(3pGek)U`O)MEmptr;%WJND{GJ4q7{f{MP0baHiL' + 'bairNE^v9xSzB-1HWq%LU%}2kL`pPTdz_-Fu~XEWG>bm81zK;Bg<;4XiJBoIQB#s~;&j{JzK0iy7tPocCs@=k9_bvO8_(qz-*%xKZ4%t!zadMH41E?h$g_?ndK#*_BvGg#M;1BIgL=' + '(Q#ru$%$y|uBnyh3K47Xs+1?TwgCvIFpt07#;' + 'JRr<72v^Q3tX^i;@pLwuRis9qs-oKwZgwo0A^7vGI7flVY?h+MH%RuohFlFBcy$%s@dd9T=E8z{yi0|#c77#60)mkD^{1hC#lnISC8tojV14dn0e&axyN8qP4mRb' + 'b!GS$x47xawcQ`eYQl?hO(ZTKhW#&PIJFL`c4Zdx1@Z#>nM{?1d<$6rn=3BbS;tmE+0HGR$;o-@>V*}uD@qK4y}&GdXCu}LXUNzyDhp5OSj6_7?^i(A{#mXp5R;lUYb>Y=z?Dz4aY>AvXkBK-' + 'U)#ngNO9B|4G;' + '1)hZDw3N0PD#@we?)>C`W=$D}O|R!(F~W-i24OGlxM{mJRqe=cV-nf`k>NXC3@sLKhu7%8=tpQQhAPN!J#-9G&}u^WhL#nG!vB`j' + '%_2`4%q2}ckK4boJF~WC&OV#v%4Jd&h7C4lod3zN($kWm9j*#J0#+8yV3vlGjG{o&p!+yU;Mi#7)GcSRBCqVoq5_Z&g$@y!k>)Ve' + '6WxhUjdaq8d!#v=M-vX3eG(*A6v-?ONEYfEa;|Q$gpp%NuTw!B%@Q$$Y;uw{+NE^&58FgB$<1{dH~Po@#QW' + 'p^|_sk^qGUX>q>1&H`$sK%y<-b=Nf!mE6*jtcLW^PmUw~xMoG(?_eigv5wTdX97x;!1Rr_%6Vc%ZZ3^|_aj48VoW)AU&kO8ER$r}PE1S$5V4+QU+nKZIJm~4msUytznCHMmPTa}{r' + '&4b_;a|;8yHOtM-e@6KIy`I8oi*Esjt0Y9KCYkCO83J3?6CftjLxp;sYshX=0MIqJM|~5m0pI@&s2#hj{(qH(rD*|dF=M3(1p(q{rrB_x3~SwaVjOfm8rQp>^|hM>t6`Z9Doz8Q6h' + '+bVIY?J%?rFg7F<8OV!`IxtluSF(WM?LanN5!YH2V@DXQcWs+-o=o3xTy(2L0%uLho@;txRbIu_o' + '=%M3E-S^S~' + 'Q1|f&*Lu|P%_^U4kSUVoj!0' + 'LDb%YKqoynq$D-x1fQVpfdC%rbIsS#EkIDR{{VWlj-h$ma@^ChClDCY!?lIqxJ;(X(rG*($A7;k+J8)L7?d9U&KP6(R71MQK=h5B' + '6P}E)sRqJSQ`HerekTFRzJ7*I$>POlz}@;2Ndf&_m)Bnfbq`-|rnvewl@>R5Om+r1YY1y2tL~M}XEWW=>yEdYtV>4Mt1H64W?zSR' + '96l92gaCuv8M5Q^!9zG2qroZ0jbMe?plW>RtJnGH=RSzx`Qf*Q2L{Ph' + '%@NzZcMpSsvKlAJM*|G?XW%5Bs`g{gfXvOg`QG!8;S^=ZNpY&McH$KO1J2m!>11{4FNZl`)Od216FJLA$$37J^ZZeAUQFb?c#Iq+WeC)>(bZ?8' + 'vBNJd{3Ks~=~t8{832+_agRwXbH%DTXa56GO9KQH000080000X02~I%0OKA20GEIO05$*s0A^uxbZ~5Kb1!XgWMyn~FJ*IWW^Zg{' + 'GB0IwZDwz5WHMiJWpim{b!A^;Z*FsRVQzGDE^v9xJ!^N{xN_h1D^T4Jm9nx{q^M_`o;tQRb~mweZD)5A9UX?2Xq$~pYDp@NZ|eWv' + 'nE^lmeCUz5dz-UslLRoB7iI>)0R9_0idK(taC?_ByJPp*Cxou' + 'a&iC|Q=9kU(oYf=`+u&37-+qIWZnu8ntRuP=F)p$QOvyC*oU$-LWCdUAPvIX7E7X=^udc!1dAX^D9^k`z!YQ;l;(_+12UflAT|$qw}-l)2q|-GkE)j9iDw)zn-2Qx0nySf^Ph;s~GVIBm?B5zd)v4`o4j{O~esO' + 'R{lJ=3Fbgoc)Rv)eRdmt@#7H4VJknr4-%AS0_cGTuqEn(m(rI4QSDl-Uauc{VH5(dmdv{k!sx)82Mc!VNB4djKe7~gn4si|uanES' + 'zeLN04Nb2A&=M3ahL#|rd*sd|yZ60h9sBowm;w&5&pxA$d`!54sf?}HxK|y5#2+D' + ')e5NN;I7xfGG)sM@LfK#YkwI%@S0EtSj<0z;`?8{c?zZ8EvQ?P63GA|gw9?C(cSu*_)ItwYud<$TJ0u=&Rq9qodN-_iOf;KYVjoAG7JKC0;LH<7)wRweH' + '8U8%Z)iBH8=h-Az!z{z+kFrLF+T|IdR)tY1Cg;aieWn1s`!N7VA1*5ZT)tMdUcWi}D_D3wvGN>x{pS5Yt>+KA7M+}4?#<~loztgp' + 'CY1oJG6)t&PXqC0@_ug|OaNTk|JU3++AtK3XegW&hQiy))sLfgxdFAy9BQ)TcPEn{W9qw$^X~wl+FwpCD%ziGL-6XPx|W8_)y2CX' + 'qvqG@XWDZ!G90}MBpL~tXXhF{X;@PcNHZ8x$H4RuReuvrv(8}l1}LdIO5rCR`8^`LD>OZ4?69h-D@$c-R*Zu2AnV5Aq}ini?MYj2SK1sw+fdLps-~@9Mw@1B9bV6(' + 'dficJ-)InfRIk@3dL7Ftfk(khpVvc30IxSF19*0VE^9zF?GCiC+ue>;2Jyre6=C*+)OOk1i(Ws24$F>qD@yKdhL*i|xW97O&wz9hYZkIamjmn1!4=}e_fM*#d15s()DCy@$$#9oZa#p!*!t3>E=0mR2qwARNBwTS*Nm+&R9#DR<1m@wUfdYvAa}^9oRHx`@L?da+OajbU`Bvtmj}seBXze3`odjTS7J%@SWe&n%XorUtqh6`4' + '(B(872^-+#<&>LaR`7qoZ9w-Ke3mt+=Kpr3&dq_V)P*Z>{~5HUx8G>p!4n9B)!m{iM8TjuT2j`1C?wml`Or|tUHW3>F$N9|K&xL0' + 'bRsK4AVjzgdjh0gW&n6i%i)AOea`ESlNXd&YbY0#BfQ<9`%GSYw_r#N?2+^?-CY9rGSg})zs-%CQz5a}wyG&Y>tT-%$*Z^Y+frDk87i}iW6' + 'GJAWbrcJ#bpEL{yrP^G+?3QuC;*58mX&Sr>D2Nt}gJJgOLodoNSm&*fYhNrdSkLNIHw)I`8rZi)' + 'usqt}Q~AWT-`^Ynb%(@S$LVc0QqF+n@7PYMRS`?tRaRvr1I52?ZOs~rF(B4rEbIHmz`3?+z|BS1?v?UDEZA3>fk9?jpUtCfV?jOw' + '>wV|cN(-xmWwFFCv-4h33)UPlovBz05E~a;9MQAJReL%jd(V~T5~DfVG`~QPs7FTu*y4y37yRqmL3@aK<>vKj?z}8s#hJHnz5}w>^hj-Y;Mo#&9lj)lY0nvSa7Hg%awhkpy;F#F0{by|i(k`*w?<~*uAt1(fGwna0*K-?w#d=0|!pE#7&)GSrN7~q;ZL8eANENiRtS#R?pv{5YPSwV?d^<;h' + '$n6?bQY#X7%=O*PG;6p{IIa`Bo?D{d0Rms4Qh;)6oH~S@WW=&!bssj*+^{t6SVot*b+0l=Je%gVE%OzMXo!WKYBVejn|#|3mpdre' + 'L=o`OY%iwamF*`v5odWWa<;h})`ExJ4C~6SFcGtyom3_l?r&1cbET9bFfQxC&NotOl#dUwx%P5syuC8sUK?-75dGCQUMIQNuJQWO' + 'Ed6Mfo*1QP#>>R$^}JphoBXg4@yAI$=jOJJhMX{ksW-=K!2r9lA7ZNQt(!b}E8QIPBpNbWV{Z4J3~Tg_tiza_nI}VwzY(QlgVHgl' + '^vU4V=88|XFJ9;_{Mdc)yk+9X-Xd5Nia^vDmu`qAe8LG#{hUAAII!q(_*8opg$rn*l3p8i{8YyTD5T$_v>WUj-|Xyd>4&%J9X=)L' + 'V&Q+O&zc155EFAufSjPLa4Tqh>t-+51S!+Isx@mh_ws6TF*$dSP9|?pj}G6sFreREy6+CJPM~|tz-`ohEqOn|M$m)`yL0z{(KZezy(s?h{RIjC0v' + '%fAct`cCLVN*RVN4IT*KKeNp9_8bpJkVKY|5{RjM4NV#ZK9{*hNo%b=zSfld-axNP*Rlfxn2f)m)PD9Qh|{&VBw`X$PJH47C2cM~' + 'JVn51_7~J64J|yz!VeET@^c62UF`b@OFsoFW%eE!gf{z)_mT;gaXUgMAQ*7VkMYz=8t_vL{0s%A{Y5K3rnG)Om@s1=$_NY{Bv?hz' + '>R<6#$qjH8&$it79*%Ri3RXT!kR_4O${P>?9w<4WMvP}mA_W|)AbkiDq!Wa)SkLh&L|d}4W@=>1)SvRfw!u)fsR53IMy65R^^g?C' + 'rBX%v{-ly4qX*21%upI<&X?ywdI2;x^DjBe(ORK31#r-gkHbdOY@Hb%_ELcpRdEz#wGFh}iJyuS6}imM>8#oIR%iqpSyTLnz9DDo' + 'T3k#Gj%-uDa-5s0=MuI}$+aV&rn6e1Sfq~b!CNKc9COhiEd4xtdG&>lN0<`HW%<4nx5e+d`{UZoT_^q->XtlOPlw!@KOD0P)%a' + 'aE&}XPTD47HUqUKPy<96!*&Cj`L#-EH57YZ2*C@-!d|>!cA@@HyakXiFxsiMBB`=^POy9kN+Zp^;)a3IE%WDQXlmGiR{79wmDwu@' + 'U$6c2!S_RmY7$<;J96i4h{h9l1o=LXW6dfn3+pY~;D!a?3AFlV8<49!)B$z*Imjukx$=CmPPf;sZsExea9!IOh!z&$#U1G5id-)b' + 'tFsjgwVIMGvZy7MZ0$m=tMK&%`DkYqU#q9A{8BYKR!^#Vwj9IP_~fGks=l79ua^qeL_H~}*{Y7VqC*S3wcP3qUs4qC)z@AFyk-oOv_3b>NPOJSh6X+$eJmP4;4' + 'E9?3}u56u&uOw~#zv{S+xjdFzzAIZQ+hJkPTyQg2>Ws1z%X~G~?Y2DU7x6bRM;*6S3;LUxBsLXxg?YXMOu1pO0VXBv?|`H>EmK~(PqLBWB|9b?xg_^bc95dJoV^+5cI_=aT)YXU+-uni^HPJk5n?{$' + '@BhNM?V%s}O^lD+5qo-cQYinkTyco^>;hO{PxmcEW)$RBK*F6sH)VL|IhlFg-_cwRPJu=Y{RzdfZX}YL)^z-' + 'JjvKfi>tF}`XY<`l-pYqx<&ny' + '(>LT}*$!yO03}41Z@GB+mWx}==4=wX9=7|LEkwGyZ6`B{h=tR(GRw%ZvN8h+@10?$b$=JUA7++pgiN9WJ~$_8zl6(@8jw7C!T`u5j1!4}NK0a_6ml_UTh&*5X)^QJ7^vLelNy_d#+`Z~mTa53dn)-%Rr2uZ~r$U{;' + 'DaGC#_Z~>Gd~Fnd*' + '{a+1SGQ%soG$=FK9LH#|kU=|s$p!7SD_IEIsb@e;3&vdyp>Lc8LQ+CNg6~2coQqGjlap6(POmQA!$bG#+&w)uRZqI!SLmwsAVgC)' + 'TT`X3CT2Yd>1pcJ*HlSH3jal6m8L^OO}|EZ_|Y*!rEEJ!Yfe#!tw%^*J=}6MO|yD>Xd)4NuMtQw$UsxRY*zAz-X{-T#LGvG5eF22' + 'qw<{|BZC#YGe&pDYIbam)hyT^tEq8!tmeYbSWRquV>LDIkJapWFjf=f;aE+LM`JZ#wyn5-K_TZ@SVoks}Juc$m;X=7vc>XYT2e0(AzWzZ{)=^4QBjK' + 'y+B<+r_bYaP$0g^(kt)1&%T?a9(tGsPZ-9V{TfvvEqZDT=HnMG{~z0#-YOA!#eX`5*Hslu`8i7dn)wI59J' + 'Z~3~%xR8TH*Ymw(rk&IP&CF~k6-rLkZs3%%QwSw_+n;e7#>`?-5Rog>91(gi1)q!!K!%rBhiAu!7su}5)%jcZ?cwFGn8sHrR3Y{3' + '0zxWl$c}uNSO8zQh#ndr&_#t1JK?>;QCyr<;3-f%J?R*t&+3nPLXzr3Br{R;hxx^P(FP;YM{UEDbf6$|DPt*FC0KF8Zu' + 'E0L7Tmtam6$7VBinI&Su`z%8F3HSCb<2FoDq~!u;43ZfXk{MH}2F+S&JvD-y{FpJ*De&&*hT?mRQcO~B{<)FCY_?OdAPK~$exuW3' + 'c-*pyibG69(3UF#LY=c*ldPtyW>;6UoB2zj_0hCNTC=jyWdl5lcY#)JSf>)d1f!j4ht9L%e=|ciSuJr{scKcczcU-7U{zf#<2v9%' + 'J1QUR*K<#E+4dqgkAH>1$EW<1r>C`QV5MeFsJ+OeF`3Hyi~aH&RgL65Nf5toG-#<8(^gh=pqWlemm?N%wJWwKu(>%}8Iz`pE_|Iy' + '&2AWo^1hl}eR2?xW5}v9=%%~lm#vbPLxp=v%XXcw4%ZO0g8B&YAI*1ax1>gTLFFpUQ55n+8|*H+k8b_YkJbsxdIwj^FIs>`Qm_XsI2HZ3{qM__rJ2AvBPB$ddu5(2*4WtS25X()4S+4f&v#Uss9Li4OJIW+VFt}4g(X=_>%&RwT;GwF7(o-YWNR6' + ')aGD?k_Lw=`VqBaaj!i^A)F|Y;zsr-*GByQhrw?M0KBh3^0L!xUlB}=qnO_;Mtj1?d&4}*N@sDKt?l-39))vI-w;)}QG#_V->pKc' + 'sm+lo!EJzZC+%|I%E*U%%q+q|H=KvGF>=bahXP7I#uzCFM!#xUyJdJIT=;)0Xv^7$R?DfT6(DV;RcpTiq5Wya?7q#k>M(_gcImH3' + 'Y=k?cBi!Xk7i|-VK1a*pg{e4La|Zr2vnubR)fxtcPbTpojSdI}<|RpR8;ZY&L1$wKK}KMe(a}yceOL4>x8i;@h}KfsHf6e1cxO1|z&v9KcBKv|mr>`bu^QKvWqp0zHVm&i0?tUE6aNV9N2LQ{J8IkdjyG{shHm*>LRo~bFD)BJuw1ZS5(t-xB$1(-pgqkT9!' + ')c?)|?HXwTf@R!v(QIu1+o}~)-uEW0`vasMZ6_^PXjVVlRNJI)w+VEK`UE$kiD^b}f;dUV1Ow)Gs^j8f32#2_%x0y-a%WCE`*J#m' + '=hQ~_=^w0FYh#1Xz76bI)yXg8PfD~{$;?c!$r{ga7U{+5ODi|W!hqEE>1e#-x9s>wnF~|r3_oOak(K+hUPrlATB?uRgN*}sJG}hi' + 'q*`rb@FfB?sv?X%{&hAyC$bI~38x1CEZWjrfg@XjjA%!Z0z7mvG8aNWSCZ9pZ#D!{m1QyDt97tk@E!CP->7fp' + 'jv+rXn8zeInJli)VWexB3v>iZP_;EmrU+aQV`(agSo{zKe%|Ig8}D{nNSHD@d{=#z{hztE{2W2CXRIT8*z$(;*&7h%d}M)FS^NWRWU1}9Ol6ve}s+BhuT' + '74UK_Qn9DPZGWiOWqoXes*m{naJ77^vp={>Bnez?&vV3h(6!7+Y?~D-Rni8mnBwJ*JvWPuuXku{moz2`LMy729W&kJRoDw&nOoIb' + 'fzvr)KYcGSaoYgSK$h$K5V>A*YPDJ)N?drOTHLwXB}QPAu4{PC$Ok<5`>Ayf097pj5Pob++dMzTvZR_J)-_7%V}mX&<6grSD(Sdu' + 'mBi&kV~H9sdde+Xf_w4quY=n)E^bn={|xJJIf!%>l}IUpR%x*G8!9acy}Mc&#H%7Lxd)}z$qc|j`sm_yDP;>~Z@u_4=8VyOF$efZ' + 'd&M0xv}eE_Ee1{eY8C(M{0@a@%MN<4v@M@$+))fUF&z9G;%J' + 'uP!Esze2-bCuf+7`akvB{{m1;0|XQR000O8001EX!r6}l4Fv!ISPuXIIRF3vW?^%5aBOXJFKusRWo&aVWpiz2Z){{TFJ*IWW^Zg{' + 'GGB9Ladl;GbZKF1Uu0o)VPkAzb8{|md9_zfZ{s!)z3W%7I>`$~NP9HkV$*gPNQ>^FTl8WWGA)f`ZYWYq(%CrK{P)g~`nF=U=r)JM' + 'mWFTM%)A*+;yC_=J6=d`EwjhE*2b~DHjGPgr=Vm8tSFCgCE)=i+Y6AT&EhzYqP@`-%k%xw9S!6;6O}*AmC}wop_OfH_^xPQ-u(1)' + 'ir=j>yl@q`gD!(jp{0bvpUMXIw+;;7NeD(|%M@&D@SX^z{hPF*m?H67@}fEEPPI^X9Se0z*+-E3C~BUlV^yCRw@lRy=?rYO*RuQ}' + 'uY~RR6j~MB<$CuC$)Xc0MFBZir{+-<-Kr8E*@_9}*jFZnb!+^%!Oz|n|EfAuB{znjq9`h1&vIdN3!jgmoZvF?0W3pgYoxkau&W>0' + 'PHVZ082%W*ge(6b*k+hWP8rDVz$I~jDj-hTdoC?37OkK5P0k;MeZ>!w4Q|LL#BtF}UWy7mU85H|CHCPUY5FF9M+F&V>c`' + 'LA%?kmQZ2DP`)!pnZR$iKyB0W*%83G2d5EP)OrY62;5;F8r' + '%mqYj-*%0f;E5+tt3M6B3FV%VK!Ova(GL!t-vH{q*PB-FMirMDdz>?S;HIBk3?&1~;$+Dpk8^{v+691$oBRHJ1NTL^1KVy}xv!r4' + 'j#@IJ`?B3;9%;KB>T3v7VXU=Bm;gc3BclsU_;$wYk8h>b@nL7Y}EAIInG;;|Q@4yZq' + 'YcTA~0TfoiI*#5kR0E62aNZD(es_m`1mlT%IA{9NPGg$QJ8?DiNn45bG)P9M7aP*UF~yxTotelpPTh_^e@5><8#)^6ODft' + 'w&Jch5YbMDfnr~5D)Y)#PBkYun2N{JJ|tx2_mtf_`2(;QR^RJ9+LK(HlKu)b&|e#-waicLc#h~' + 'job0f;7QP`Fk*32*Ye~Wx%2K1QCs4kWKQFmGUlYTsqpm^l=;vTT@wELhJ+w&+q8)YH@@83V*uHtwcx=_`Z;Odk!^lkviEALq2P@#WJrc`?Ff}W!1(_GpuQ-dR_YbW!' + 'NQ>herPWnmh<5P_8!dLgUKs1V^RdF3G6q)B#kn!Jc4Z+nB=Px(*B3H!*jyZ4k>Z29MfBeVh^hPb`tv!-rz-TRSYK~0Dsg!9-`xgH' + 'VX%XT4oA*>2Iwgzz|+}`Ao_bAVdfU@-XK0>J}dIcDtEb=#{bxac0ytIPRBmSD`5K*wiZgm1g_b22WaE{@JZCko7N{P2!#|D!_7aW' + 'KPcT`OAdU_rJ9AmN!OvT&Vybz!9Fv?$=qSK!jprZ-CBDg(NBdrS;we`zGM993hvzFJ=lthn(f&fhP7-~771&PJ#|c9aBaU0Hm!' + 'Gx|_eff4zf34zTzt}Z+rqNn1;BO{wC6{yR&970PVH={g9QBYqbM{c+@mdP4tk+Yz#vS^CayaK#zuvQJmcukpW>;MhpKvp}Brsz=H6lUapE3dWBlrJOIBFc>{X@=gJie^=1(+TK$<3m?>D(w+SLE4rC23h^QT6>DeNMo|04mObIl' + '=g%*rfTKVo302n%w0jlr$?Gu|mPYf3lfa->UFRp<-(1$hH8>Ui>_OJ@PCK01Nt2T9Bjdi%<>7Pc5=1J)^9TrSRvu00i4P@Q#@7@Y1Uy)aq^8;^jYGaAL<$d`eALQHVPoa3v4~j5u5U6o1?=S>0OGg$Ee#e24{RdD>0|XQR000O8001EXvK8d?GywnrYXblPDgXcgW?^%5' + 'aBOXJFKusRWo&aVWpiz2Z){{TFJ*IWW^Zg{GGBCgaAk8YaCv=_%Wi`(5JmU%6(g@J7Tsr&<{?!o6$;C$vL;|sE7(T%P$_?3NaEBD' + 'q^!)n$Mf+8apco-9m-lW*K4^Q0#Fhfn#*m&qIj^YaQ4f' + 'JeifVwH^x{8=Zumqo8orQSrYQrwFXvH?_+0%OFnp)P$0KUyUzccvb6a3^3o&BM?AJYTIFP7We@;d}4TA@pL(+xuO>5cl)1fJet3I' + '5S3)IQA&3pYSXFQoMaE8w!5q8QMZHt2|rLv0|XQR000O8001EXVQ+#&8Ug?SA_M>cFaQ7mW?^%5aBOXJFKusRWo&aVWpiz2Z){{T' + 'FJ*IWW^Zg{GGBFcX>4;}ba`-Pb1rasb(7zV+b|5r-~CsJJY88E*d7BN>%k5N8#i`+JBrUql&e?A4z@hjV}JZ)=huyGc}OiqmOj!~' + 'EsElySH{2)Yoi;W>wp2mS$S-sz5<;QK0bW8bEbUOPQw8`Y;SGHOlkYAR-b=ItbUzsEQy%*aRbx~L)sqY5SC09tFZcwSt2GS{wBq{6@jXJP4T>|G0e#pqL1SW@@{GlEnI*p(!L<#<' + 'ykjn9K@9^BC4VeMQ4~US-t{157tBL0N`qF)6_@e;M`vSSGZ6n*CKd*p=KODvxgP9ZVZa-B%hF28OU=u=hsWY?C5tT-^Ipv7d3%bJ' + '_zh4?0|XQR000O8001EX3~BOsK1cunFFFALEdT%jW?^%5aBOXJFKusRWo&aVWpiz2Z){{TFKusRWo&6~XJ2J=ZDwz5WHK&rdF_2`' + 'cihNv==VKm{{uX}CkC`2R^+aAH+0|=u0*YMElRXVId*gz2pnKYL^Ff&03(u;_1|wj`rQr8P+G^%J=w$}2I#JSRaaM4*JCo7oZg+c' + '>*;lLe0K7*omU6ZqOI!rVbovOQP-Ets$5smx}Mi%S65L9MVr2!_M3H`ef`zfU;SazOy5Ddb$@t>za#i>wr&?uzb=>EtX(habr&^@' + 'RlDw08zWjxtj{v}&gmy9X0m2EtJeUx' + '?dxV)lwASk;>pP$o<05J@6Vr~|MdIwEXyWIf}MNXuA{|f-Zy*GYv@E8&6_Kjk2HdRqLZ_us9CNyeV0a?u8!Wml?}kVqQ6_!-P^ZO' + '-(J-)k?XI&x@r1rRcHnyh@qhbia;asH4JxEuIr-f*BhGQf}n|#Fx`{M1Up9aTNJYml_})xmdj<^mwnSNJ1iyMUQN{_|Cu+J{+n!3' + '_hnV~rTVle`)kqEv|aY~Ef$kai8qUKSzh7js2EnC&daW=)qdJYS@F6U`9' + 'ExXfi&c$E>*ueMgx6`tH!!nKPKY)qEj3)n$Dz5g~dCpH82joY!4{-YmFt(Cd@t7V-c5{QNOB^z~Qu' + '?X+Ies5kTaMcY4ZH_PgHy#{QCIj!+QG}+AfXSOUCb@Yoon#lGiq6Gd24Ry8A$LzCuC8y7uJV`Z<wJO`e#X~~dcL?x1Svls}I6MBs>*E)XkBc8)o)zaWe}DWUbv_&qeKCuz$' + 'hqLpe7f+5}Jt>Zkp!%P7DD><_xX>S8K0f*pUM6z5Y6lqOUM%zyiSa)E-O($qT%0}o_hSJT`*L*t^5og$;>GKeACF(j_h(1N@!Jee^-~4_=rme{XAc9mhZ%?XGG|E8SAxp`Mg|SZGimI);+s!' + '7P?QTWwZXby1BaU>+0w|kgO{JC_V#?sMn`!*f`Vvq%N1@^{cvjR^1+TcZ+A$8UAw&c-r6L-)URv>aW@xg~X?AJJ(;%ZknH-zdm_>' + 'Dt@1@8|dM@KEWewHE*VX-tX(zE1<=-C{nQ;P~U+0VBx`7#4!ifcdI`R3c#3QeO!wd)#' + 'oZ{2AU02RW2HCe%LVQ}3cb9chu2%EAV$qdQZ@yV9i$&SJ)1%}<-R80@k90Ssxy?*-i*`I)snIQdHW^h0tqw${;qv-LUPLCg-A3rHhU%mW4@B~gDSfPoT`%?~VXv;WKQY4{#vWef%;0=SL6I=(~qW8pZqjzdnHHN6DuIs-I%X{d(`wSH+8$FHT=R|MSW5tK!)cc>6u}_m|O=' + 'dR781L~te%ch?I9Nu!siC(PBX0CD%|;G*e?=KpbkJQY*ME{h(Qt1epAs}9iNx$D>w_Nsmpo^$99iz`Yy?$m$2Q|O;tzl' + 'n-YOZ-il611YC4Lx&=%18yIlcMK9VV%tjkkO;=vdYwiHj@feS)dM*3dT_d5IN0%^?moHEDV0ED)fB;5k1wU>ZV4CXrUG(VT!vxs1' + ';>q#Tqu0;Ri^sov{o?mT&tNqk(QxE*HhTy=wYytV@FU8svSfw7GQecdkXnR%Jy8z&QA>' + 'sWy06RrT#y!J3X#ILO8dt>#tL?=S^aZaW~@b~NDaC8yMbkS&fmjfy91-7LG!3I&$B%8bg1' + 'Gq=|2jZLprP#4luZ9@{XtqYpd9)RWc%Rn0FmFUW*tE~^sEV6!@D{Oo)(U34K4i7&`ov#Lr#Mo+d*rQzS?zPz57>5zodjn-sv3D<=' + 'T_(g8Qp8z@YiiJd(1+ojH?-yCUb8tj*oeIONPUIp?V+LS6oWD}o-fT&2HAsL_c_(Ytd8cfH}2FJYBGpYw5$eQp{abpMNG0Vx##y?OGY4MaekHJVS&7E&9Q' + 'J3fG_m)$v2)mr81DYZ-3E0)_oowCMq<;Js9Dzs2}=w-N`aXfo(oEpZ2<<9KWDJpf6c2}H|$=37v&<5*r0fHnBocRK7@W6KfYi8-y' + '2FAOnnM*>x=$migfUJlo5(G-g=mK=aptDKhwcT_L>7u7~Y|p>AJ>iX)1bP!18jN^ksCZQD4v5=_pVN0SpC>MlQ8kPFVd9bJR@})b' + 'He=p`OeyamP#Z1Jp91Ka(oT!rK`MSe&+yKQsB@|$?4?vPn(wRneKUn^%su8`@)O<*sxSCMe3-YWft@w;x+~h{{LUi+bRvu#nq&7n' + 'For@Fx76_Jt)qj-|BNOqQO5a3Ix5X2@PlT;Y0_XA-knM8OsrV9uHCGs$c(;kI=b8{pfb>xWitZ|=E~FSdirkFHcKIcmD4_UT2%ZB' + 'gDKv@F?H3MoFi~%7XbRBalX21Fgf&=MGW+l5@($L;%>!_=j>RgA8#h|#RU7{z%cNR*$HM8?Zxn}$DIjCfrE2el`}kgTNoSdrWS^a' + '>c3>Y8?-eiy(nFkM0!cOXw@tuRAcXvMi__=Yp)}ObBQP9`*IHCD3kDB)uex?M1g;hv_HjBz=%8eMxHbKd{=;`Tq2k@buZb%s@' + 'hK9W=R68JaePz&hIV?q4V6xI+f#pWsYs||#;HkJ6aL~(WP)*ND0F9coyA|EKoB`G8_a70?*xe+KnMzAi0OV;Wgyrr6nwpDde5aeCQ?ZQ4lta^-sx+@}@p1u6M2m7o_xE(CiBcROM7b^(b&qBsD2oI7RIr91>k(&)mjgnFYJzKLXn;_!8$@aZi1i-*' + 'sIdsv;=1j6c?mIV)?J?&5|p0BHxuqXkdYN_xQ~+w(d(GsC!h3|v-euq5+$(JjMkFioc%)18Jc@|SOQ=9Wcb>`N9YXx0SQXiBF>S2fS^JTmY{Ok2BdfJee}' + '{SvS0YNC|X>avY~x4BdYP!=6U@0RV2IAr+Ks_mK{_YfRj<@~0+>*R_7MaFKc(+w|@@Gz9U##Wo9y!esUB<*63_ec6(gl~3mS<6UH' + 'z$J*RKo`A>k>z!0V{Ri|!lTCe4zrmC-V$' + 'WS)WypbT=gqFcZa562jRtE0_;ci_9*rNSg!2mr+8piEKa_z5p`74yfwsyj<4~' + 'OeLIjKp>26zOjp4eC?dOoptPEk3sGDx01Hm&t8gccWw5uR{*d9eflZ8DgtIrpJJPBurGk(CS=03OB4avDk2RXjAomHq#20q|GUDlND4N@ud^^G7!~J9Q' + '{v$H7p|fuWM1RprO&!zl=%WMf7wb=Q{Y!No?1|k_c^4Dv#zU@I@bbnE}`@`hI' + 'HBpF}i)W2U)Xm$TwIhjPyP%sn18HK@nP(tt{&_-HsyBF`;0Xf@TcL&sEom#X#F$>Uc$?Gorq)vFG3Y3^FE*@KWheBgwm`047?znu' + 'y}Sahu6Pd&k<`A5Gu_m5mK~{%SOts>fOXu93Kza^nC)M=n!$n$cOWL{8s=f>a1+3N`4lj#yt1u1{Z^g(' + 'QDS}l@x%t3M7JO)zi5Pb4z*0Rb#zR%qIvlN=qMUuOxtxWI*IpnaJ*La3gl2gle@jyx~@&ku&1X-j#rP5pOD0D=Cw65qOaNNjzp!2' + 'r97H8tNX%XgT{>OI`bVrqv|xNsbncq>xfKT)7hOE?b#$SkJENDXC2_BoJHItA5mh6sR_ImL5tj?_yq6i)hCmpFdIVe6H56OcN7RbdEn-n90Q!*1!h9_HJX;db8MC--VaHc{;p*`BCHlb6O+?Pp-TyAEO`O>Eu+V;pBzcp6AJ-&sMt)_bVd2TZ05+(' + '$H{P`1dQ&^RdET_-nonNKXFqm8aOMJ?u42`j#h~A-t;>g?WpVOHR=`uTP!sW?eKQ9dOqW>>lPz8fl%83qoDm{?He&cF^wYnZ65tI' + 'iUHy$@H;=ydB=mpuSs%`hCdRRPuWMR|4DBUVUB=!5}s5Cb3`jk' + 'xRx(S*h4Yle+GfikhMuNw=;&;B}A)_U=+>ermyc!yF=?5*x}dcr_uuNLy|G=cO0XOUe2m|69YRMt$TB>9g7<;=7uD%<_z>4sK#-y' + '(-tUagWcos+CF$S+~G}STuO1rQYNqF_EC!ootDNXQ1Z>gi^RM+?V%K^-01+L0{a(a88NSy;!P6$Hrl71s!Fhqb}|X7+6GzM#Jw}x!PiqZBA*zyMUIgjX?B!&5ZerY' + '8V#Pd!@@@?jEYg~J0|daGAz{OZ7Hr1h=BscnD{rFXqkqT8i3BIq;6-?$BuZfSU^cWrO_3@{ivD&IZ!D(tmdKVfj)n-5JrerxhP}<' + 'L1jI`kk(t`g9!i$gl(cm>v{y$7G>W|4S}s<7lICDG`s})4cZeR4`Hud{wtW%UwK@aagzh}+Z9bOJZY32iaDBS>GuX`' + '6nD6;jC-X38)$5x0CL4HxfK1#U2}MMm2&yueJ3i-f1GbZ_itQSI0c5K$$V7O^er0{4M{ax40656c4Fpp{N~+z6@#!XuqfK2Jq*8bl$&k&Q0P+hueJ4#7}B>V#8qYklbLm+9OEwM>V%0qXgD0m;^LvxU`!+g2Ag13#XJ!ITKYnMU4eAjpVZV' + 'qQHw*zF|et<%frnRIeIkDxoTc0tdf-^w2kl@`$lkmEmayn57JMBOHqIB0^-^Gp-aB6=r85eR*cimP7?DdYMXG!~D^@T7$wCQ0F3*' + '^b3' + 'vvn>$aY24G`|>G^PTKcIt;>3bS}`T6qf?apft?AQV;vDmB;x}-mU>F0%IQWg2>9f?RvO!;&?YX_ABD6vEDa>2MUeXvbJEaQXw8K(sfgxvXU8{Dd~n0c9{UL=k_#spF@*D71%%{lUm8*' + 'Y7JW;CPqUPL-v!kJA~89qmhI{a%TkTMbB?|`~6^QMO}!v^@MC&NJ1d?hlulmzUqC=HYJ8tx+Jgn5fW7Yw#W+3IKm=B(Cp@;s!CjR' + '^q7YRbMX-s9q#CpPo{wW8+Yy5*sNT(%f>N3xuF;b9rR^U3bQ}YTxUp{p$>XfWeaaDstuR$Wg5nEURD1Hljq7%S#8=mV5k|WBF{JT' + 'r(f75j8MUSitEP6G__wj?X(M7_}h*0tv44@yN`Gj+OS+;Y=gaQE2{' + '5r_l-{6tffO;9hZUKum2#37WI-qm-ct0BXCJ!6S&<^iQL3I*0-w5!Z6#D@fcB%f++(N5lf*>3}f1*pQ7u_BF5l' + '19Sfx&PpUtQD38Ozg*)rGxidrQUH^^#vozR&UA&eg!D>Xy?&2xdAFIfulC>A_B65A7>#DSbE()9WHBIJZQK4{LxX$GD5>bd7`fPq' + 'pzNxW;nsut7@mTAC}hG%$IozVST2ef#OGXvLh-#4(To8(n->5u`bLD`iDq?O?Fq*?5t#tC@+EB0chdYh6wboh*PgedWPD){_GC1R' + 'y@FMsxGeM*qr7~U4rbaJPuMaMP#(7yZ' + '&^REtTh@f+c#?IfD!hGj2n`(87oP84R<6*LCk~9Qn5Gg!lx@h(M$&>RmIz2juf1Vtt' + 'iW2Q%oDRB-l$Y+Pn^Iv>0E?H_0ep#8QYG46AF6|htt5rFzV|WGxDsUEYe`bsFR{zi;XK-$#y=dR157HpMxZKQUFf9|^)+99=cIzQz1WLm~BkY2r^D(7U@%Un|52=`$#B-JI' + 'iZH3%;@~fgv3!$6cqRLH^1RckP#*^Ryv>+&~fQC7dd@!*@W+{KKwF5(%ofFvmzYR' + 'a)*tFqyrHF95fRUG$gUjxf6%wzH-LhFI!_>>dPZv^k' + 'k5;Q4TCH|zRmXPPWetBB9gw$&P(W}zs3?B3St>P{`g)OMLv4ju_a5tckR&d*X@`{*$GYkNy^Y-2^=T(>3fGgzDy}VILrX739`J*5C<7~90+}lFQW%WHJ{tPy;qDCdF}*4tb7Q(yG0?s3VhIf5Cz_As@p&z4lxuC-C|ue0R&W-' + 'WAUNy+o5!K3nSc1YTgCU3l62=eHz9^43BDQ8Ri_0klgqv@Uw`tyEVQH&%$N9th(JQIW8kYQakh)k=OxZQ+=Ky?lco&tTAjyV(AQi' + 'BK}}UDj6Q?K)?bEuz2`p1dP>MXm;+^P6$9?z^c5#3&zewCfR1${bf_vKN0' + 'AUhp|ApqDgY#xy1L6-6@O3#BDAoGb!T8PLhKtYO2<&pt%Q7po}unKMO3>8PhG78U$VYC%0#A`lfm&oY}!CKvQk43Yg+6?=n<2ms-y@E5vH!A7h7Kun|d|AcG)2$(=7h*&>rEghkh|;K)3ojv008' + 'lqDk~zG+#`=XX)HS#xyaS-mEA%dRzn)#db>HPyfBq8Ynlt{dP~qI!8%UcvN5&z5_uc{#1MX1jdfHh?tc(t>mgBs4;k7oIINex8;_' + 's22Q>GNLew0QP2RzKb1J#$5FhZNMrDeTZgZ_*#a(uuW)OttCxyRnPlU?kI}`j#-MKQcg@WwK&eQXJu6tv4$&GBeBD>=>o>lF10_y' + 'BE0C3Xg5T|W@L<{D4IqF%|Mhr%QJfaS~2!xDcLy0HVA`dHB0`E^;fHpoFM=x0h>?+(+0%DZJO29b2gm=Mxcn6dlI{%UXHf#hGlQN' + '+g#<10&JXpW|(J1s%qG*KmaCLmQ7Rq=W_42@CP#qS(B5><`Ng64EId4;Urgx%{I|>s$BQ^!|WlVdX?{I4;|rIM9rJS?xRQ;yn;&S' + '?luiQR|8=g*U{zsC&EzBK5lZra)WyjwA|p$Y0pIFnuxiF^c2W%|KGCbv>AJAYtArQPY6sBSgx#_A}xSoNL}{5y_vd5^}uX5IFH*w' + 'h9St!qfzNw6ROyd%`(Rf6BeFYfywd29qIGTI$)uR*o`o$=g}C*)GbIs|2|YbJHkr6uANH1BJvKD*#E2yp?^lOUas|#i$G=0g$LOkN9xy_=>Yh79Kw' + '0CP>d$P{xItn~;=VSbVQ8a9rCkrIg1d?9!hqZ5`jm{b*zM1T(g+Fpo+D4i%`9j0mYuC7=3?6`zi4^csy9ej^uTf;Dq{(nx8>Hal{' + 'r%2gh<2?Nx}Yt3NoN)ofpXtOu{kB*LQ@3dz}{%6z`DdfKg~pr9&4KkdSNxCEZdPhDF_dYUT6_o3epy#3JT`ML8WKTK4>2V' + 'UR%QuErXbceKNcFz-quFm*&R+eYD#^EY&TcO@&fgaCB8iLX#)5)H@9E*Y{>E%o9lKF3ho*N8;GzxV_wRp18%g9VuofP}JOg`tYKV' + '+rg-1*s#NjVf7@#7-Mllet{v?+i#Y4r3U+|z??J=MgAtT05MUBV+98QOq)&6cEoV2Em8#7O#481Qz-5T+sN' + 'oI<7P!Nm|QO9J4(i>N?qL435{v{W^wq_P`-?1+jivUg5#Q(PI~ePnjqYI{|FMVZZc&d4vhWMWA{!%&>7hNr1b=&O}EYHV%%xV`)m=m' + 'l8A@A+unAI9So=AAgc~hy2DXrI=-t){E&rWs|JWM+F`8P8vuP!bWo}mCcgn{hOJ;X%08EE(68)wfBx~+vnOdVqk~QNYCVC{5P&1u' + 'F~D38Bt#UQB?@eHQ9`>ANynTMBR84|es^AX{W)yXbIbq&lXKGCHd61e@=r*glF&zGoo{|TT{FPCl1Y>dq?R=g(RPB>vN~z$id~tA' + 't-;AEIWkrRUDeTyB%YsNH=WGpjpupAj)@{SxC#Ya`5r_?jGY$FMWdZ#)^~l2eJyhUwkh4Z;tU)HsK{0SseqERQJl8=7+M' + 'rA2w{p>uoCNlb+^B}ZckhbvJIa*@)8_n26N`CS&BBAUK^yN|iZIre_J(p1hF9>GJCm#(h|=1!+sC4X-Pi2{L$YI>k*7NVC=jw2$p' + '>xfEa(Q_mV>}t5G(M10m7!L6UxN!8PQXEVW#Kso5q5|Qobv*?MKnMdD)eAx6#l|PSwdsw%(&4QDv9FPky>^6|ge)gO~b3zmgi)7lRN?=QT!AB$pV}h7FzKUD~x_3vCZF9A3F+=Z?' + 'kDkp^awdn-W32zja{We0MI;ujC;vAq$TW-tGyJc~#oxwZ(hBmlqc`I=$4JMeLRi?6wO>_Z`DgUzSiJE2AndcLi=Xd0CjflD&)gaV' + 'L$x>)uj=kub$is^EuK|p_|GwABEr9Ha~-=nVr2kMmqWL7?UhSH8T^_Kf?###00jNfRWw^3WkK~ajnXh7w`OX8mnoV9J=p-0s&eE!' + '7cc0XKD%z*JNv5mU^v=wY$-3O6x)jkH@Q&IrQy+#e#y`tgc94^g)u&h9kX$$)XssGoT?n$M(Qujm=jGkQb9$}F+HZn24V)$IdrmY' + 'VN~Bw!CHafI>WQzM%9pm!pNU8in&aQ=!7rDkh~OkHc7UC{=7(0!<`%o8a3H0IZD?g950H6N)9=Waf&+4f&' + 'qR`w2>p6ojJM&&{K!8V;+tKoFyrO;F2GP=WlocsB|I3|aO_hxzAe`uEoe3%FFyLgOOC8!4<7%8%|9o+3>C-LQ?Ct{?6#4?vnbYmS' + 'N(ieR6f?wz<6=zKwpJJq*>D_ob|{r|uU#cq3hv2JKrNmeH&&md?6K8n_Gb' + 'I+fTe2@%UO8Q!A;!X+tf!gE;P)qxo_YXeEDKFe<(ies;_i6Ra_^fjJlAwa2e9Oqqp=ZYCtEItR}Lyqpz-WUDnMjxq7&!E+E42!4*Cs+wI;e>2a9WyL1y0Vzgt~ABP7lI-<;EgXHnaA+N3~7u;+7M38j*j_9fEofr*Oj@1Nx#xi' + 'sW7YR$|Sks-a}V6)buNq<4IN5_4g|bxjaX;Uz|98_}hc(' + 'Yo?LP)r(bsS4_+4b?tRBgr0ht{=w2KKMhYv8oG98' + 'u^;2is>FYsx@RoZDQ9GvUZd@!uNoK}vq0j(VT&<4lavh1OCTTIz{R-^c5aPRO(F55*(B7Ex%Y;xAIwW$y5Y2vbc@vI&{84e_47N_' + '0hn)P7H?<$jO6vK=1vE{KW*11W%sVGo}Z*cA*Mc`a0Y^24{Gr2xpXRqSI)}_QS6om!pi-?%fS!|zE5=HFQa16v88Y~H1L5VB*Zo<' + '93ST(##mi^Gm$SQ7elec6?4MN!^%P%F)~CvwGmMUXnK(`6KtX(Skm>W(tlt4L!P5TynUlvyYPldw{dulKODrY{)?1Uv?cA3n8NDbrs||g*p_fz4<1abasJ' + 'Rf}6W1YOiwi@%I6f|a22tCbz1OB2*kk0>Lp-=p^VB+XJrpek}X$n!WNo*_{CZP9W7edH17l8ybGV0ZiNt+Vif=' + '&07h(D+F3rPi3?24Wk;->CqD&QP}&8<4(8Bt6{gZb$gS#otKy;`*G+x8tSpUq56lBJtG%th+HZjL4lBtcewY?2rLTI`vJ$1)N)`C' + 'Y)&5|z)Ji2NHwb`M`yq1Y4w1u%8G9>?*mV*k_XN(*=CTPF!dQA=wpQ13Jfots|V2L?EL7(lcQHpild|AHs3NKN%T*?kFBizFIy&PY*fNnRrreyvjF6>>YeCXt}{Pe{H&)t<$2bpBz44YFA1K8' + 'Dto$f7_uQ;XUrdO*3LPs|4~>tuP0KAvxB7V+x!q_bR#Z4A' + '_F{~+!D2+4fl7~(3<(Bi4@uFqQ' + 'oVaAmoStrXCM=szWL|Zlx(Uhhp6#ip7T7ib(Zj^|31@gJM=fsFbwe7A`6T45HlAitiP6GUKo4A|9qz=nhi7n)Q!)Td7TQCoR{pQ(' + '*yyhY#-r*60ha5UrX6JH9<$PuGXUV)9tCb9@gCE?^J1;a>AUi(&L&&k)TC-ltn7uls<(pPv@2FMtNR`_`j*LTMSN>Yy~UZ}fA^Fv' + '($<-y6*H4ram(4NUB$jh9$AVHUuz*5!JtqA1_6CIfSCMXUHg@x9mk)+Z4inmtD_XlLILt3=zM7Co?+yBf8H!Nx6$L*PmZGf@1pCr>&>WEbl)t0' + 'NVoRkoD83v%fV%#470{zZIV-Rw-#I@w|mrb?6dAs)3FHKy`oc$m^*qhw@8W3odxmY6kWEN%N!WEo0R)JW|ypruKo*%=ICB|2eh)8' + 'hU^yx_R(X8jc!GIFpl?k90<9ywi0{k61Y;eBc}{|%J@JkNeQU^7e9m&?fU-jy0R3UP%L9D9{K`oqiFvZeXaF;>*#3?=a(*^?zO8{EIp3*Lh??&@9MceX%Xo~UJ7g6S5p!4^o-ZtEc0?~^sK4%!Ax4BR1n#SSw=8nNZD75G9>|' + 'I5sLd@dW*YOx>vtpL@HC_ZY$h>i|zx?FNVW$0KA|4XAX-}?^d1fMT+@fc9)A^@^Vq1LhMG9iGOT5q!~V3H}8#qe1e' + 'O6bNO;o-yTgOioa_=0S{P#x->E2Hssl;r@OHYTx!S$rMF`lbeqnr06&`M*)kK-BHYnWDJmzVB8E0VBKKo1=8hoAXU_L' + 'M28nHgaxQY>k7=*gVqY~x~RNFKePjk>Hq{AV@o`Urp1Q{nEDd?)#}Vb-*TjmsGGK{Ix~UxB0P*bS}54ux?Lgc%@P?ZHax;<+pgRJy7BMl0#(}YAMW`UBv%Rp!|GEg-+1MXBqmM$nP>|h)h_Ad_7oEMm@8t>4_7cgVh6M;*$J2BQ(RtCtL{x8_HKJm$1DiGJ;&mb=${W9th$BM1>d_p@TG1N%AG4@0lZ@F4T9NxFki~KbS{(!_O{1ya@EK+doO7o-V1Mj1dyI5xA&MwUT5JWx|+' + 'hX?m-URT4DiW~C&Q;M~Ao=_b_K_K2-;h(UcVs!FyTa3(U^?rCrNL}tfB%(=w&mmzAGw-TFe+LG7%UuZCsBJqW+W!UyQW(~N@@#ui-ajAE=vI1IoDr04vwJvXN(Wq-I=&UB%~8PE-H(;W|Mod-YE6k6@x$;{ECWLQfrBbFo_}(cqZReW>&oRdIsm(i|eW$GE%u5V$BCKg_!%2f4U_{rVzD31dAbtjozr}RY6a;qL' + 'cyM-hD~Y4labCOTYEf>hBFvc-%5@`AA&l}48jV&lOlY^1T(#48wTke+6+5qbky(8E7NPWUbYrr$>GzpX}@+{6tY2H^Wwow>3B=r0o7-mIlZ)$cx)I5c#dAz}i6+agqVBT?E@I_5ZZ})hy(%54S5r2!Ot=sn4AO|i~2HT;GzTD6#Mri;BZaIi#od83&omQE#6q6<%i#zy-7pePd|3YcS' + '^fsp7BOw9}nbY6V>QfN6(G&6L{N>5B#|5DIkH@b9_#Kc--DbQSEw@6j@yEYAdd2?e#o4oeKMv$XHE(vcVz_E=mcVS)o7#|Q' + 'P%|tKC)#lGM|iB(R5~vn7=VWGc5uE1J`BSgYQd7;<#DTiYj|STOlOicdx{=Y' + '#VYk!PF8@x)+tM21)_FqI!G3BNe#sEhMdQTdA*TZV#;#yH)O5lGd{UojRAm4nUh<;21>K`DA$hz=HvSG2c91d;BkIs_CF14Q' + 'ud$e>+WbZ$fXwd@OUy8F{{VBZT3H-|#r6~bYF*?-fW;!2h1Jq%KayKn9pWw4b@?`1JE7XlhpqKQfx+)e)0wNMbTA+S^~rc1PYh+K' + 'JiQ4pzQx9~+yu$MqKSuZ#-{k2V8a2WzNKmq@TGtAJ?uBPnQ4xjYVKmLxy592i`k}Ykr_J{P4a`7Z~e0~' + '1SjJ|vi*Aeeg-Sv3Xdw^=&>qGFakM5Sxe{83kj)nbz9pn%SY_*dam)pXMB4rrE?rx6mIh>9!3~A(w)zxP2W;kM+{p2)3@s{i6+u?ZT' + '!^*QSQKH>fnf5tKHQ$`ax)b_y?qHBVqWqPRuOT?>E96#f_&smG89+2@jT9~AmJ9X60rrQIYgqB*rh?DYbUgbKP2Pj&_C>c7v~bn2VZ*B+Ds1G' + 'BipjmM8tUY{~@Eu{J=GVY?Svi$tWS+|Ip?S>b7(B6ew+9Sm*vL_s+<_kIe(`+?Mp!aBe)iv!HH;bNv{;KcmH6*Q<3fqc3otZA0N(' + '^|f`eUUXF+Kg{+YB3Gn8Xcx9*Fg`UOetdJodDbi%@n)iqN9Yz8o7kamZ4jT~5a?uUMLDl(IP$G>ifb_}4Ke2vFA9tqaawq5v7fs`' + 'I?>_XdG*PK6&^26j5$i%C0+UV+GN#n{EfE|P9=3XEB=AlXZ|#xN4llJZn@RPA#ncfSgtP5zF>RquC=*`I%}MkNx?QBkZ|uuF+S@w' + '8k*#gk|Gbi3<2=)JKIn|gV@8g#K&HmUG%SHb9+-Ut+iSYDK?J!pziecnPA8dK6LBAAn4W`zk)hDsIdCwnG6;)*xhpwP(6!>IAokK' + 'a1x;^g{0ykZbZw?9wYNNwR#ZI;aGvaY{_X{GI?kuNk>c}Li7r+Oc{n5BQ*#k%0fnL@GMRa-S`~>RFR)@mWCcp!~>ib`gdbDa{zb$9IR(d(@|x?&BQPK8@aE>k1smnu{(rV=)`i3}e6~QP9V#' + ')Qg@tmWIp9?8&av17Q{dnLr1a1jbxd1sa{!9V!9^`bw0i?9M+jC|8OaZ8i^rh?C(2VF4&xCk2n9&YKb*oE(Y*O@#wC)Zqfyh{R@1' + 'D5+U9KrIpvw4L)zG?Rf`Xx0ET*#bwpW9VBGd2z#TG0K};R|DB>jhp}jQEPDywuB=w!_N*I!{*U|pAO_ZRc}mdxWtKrdu2W!9nHFvwZ9d!US3wp=~^45HltD;RpaBMSd~V&CV9oq8p6{TK*$%S)N-$HA!AxMu2z1Re#GH;(0q' + ';1K&x!b5Wha^iRckB%m;>Asf`+A-u-ITfMQDLikrBG)Mew%W7*0}+RN>^bL*q|-qjD!_uI(J8OsGd>4Y!r4=f{0&{DFXQ%Y^6rG@' + 'e=N;LI9~=(tVJu*PP7ahV}c5JZsxBY+F>VI7#_!EEx%jDvKD75P|Kcsh>+aFkPoV???keBwA`cLNhV_D$CEuwF+WK_9$VRLgMiNU' + 'f064X8=~YF#4?6a3%C`=BgCFlzzSmv4#a$#O*r^@T}2<=4*g>NDUGfG)<=R3KPOwRJcwAs+zerqNABf{yfHB{Mp`83LaGYo@*$Ns' + 'Qpv^E$lQT5=D+iSvNN&%R}@2=sT*XRGfvdyvRKWlitKJT^PWUSVj>P*dtUgOQPIzf;f3|_mJcv8T0MXAcuF0OrF~&|2m|A?U6KZQ9i<6TmLZwsMIXeWcCZ<3D^cQT;9z`T)J0K=VdjJyk(fGkyT;i5^bSeV+' + 'n~Cw(RFdD6T_HofC?hHH+ldzGW<%@OcfR2|qMN)a4zX=C87SYY-k3c{P22n~6ho5w6q~)ZgIV%nVkYXf_ovvIhKm@{7qfyvp(H>b' + '@!foph|ZWJ68@xKD#;G@`uNBFM6JL#`s;q=u5#i@>V=tQGUYn1AgUB}|HBslBe2D^J^Eou7SWCMc!#*pIeCSZ#)i_927UVM!v$Lj' + '(sjj%s2ylaO)g$&M#nEB@j~-QL2IP}dmV`uxnhjNWwvn3fm+)csLy4hzKoF?W2SatsC+cZ_2brz1H;5O%m)vx{I*pec5fV8?fb9v' + '|95P%OSJz-pYe}IxPK1f{Ew5k%Vr2&4o2b*FdPFLox-XeM*d!(NW|f{UnG?qUN@(KS2Z|' + '&dl5W6OCColOY1|$;@Kc)nt%4%(0I0&_X3)Gt!9c!2C=W&ZtN@a^P' + '&0WMRSkc)x6XQKdfdM6hULw*ps?u+S$Z|1OMJ1QjuaaOCRHK#O8W=SY5Y<9CU3H|eKs~Gfga{WSm9+3`RujUQQ_-l9@t%^{@G)Rc' + 'fD$0nJ-H6zO(?e7#F$SZSc+ev!uB45^7e~h6*Zk$T>TgVSA0|YYRGLR1v79RzNfzDt7lq*h<&-@fh(6D$' + 'WMON^+w_y%Pc#afHZxN@35We>sq#F{4&6N@@c$gFc`~GH%fykAVYK=Mc69e)t*Xoa4#SeioOl%{8TY+l' + 'O;>Cub1C1~McL7knabBxnoGG1Z78J+zK@OaNk+l@3Y|6~h#SF<;3WT-bi`Oqt0R`xl07wB47O~9FvUoF_?}KdtMe?F^y*xePVDA)FlS<&z-c8+HU85QjFk)R^p-TX84o$hdn%5BPW=Fdsnq7#22eetiL`fo3r5c=@cgThd5bYOEY1kHwK-Fg2%g7Ti@#QJMJE15T' + '9A;I{K`Q1b325f6K@H?9B6Cc+poLmczU#$IUjHV(9m!7t{!jIZ5ye=C)S3BFj!t!od}zqv5Bh1z' + 'J1ID#|G7ZV6V89qOv##$e4695R8C0K&-Av2uNn19}^Up@-%o@C3I@cUP4mb)h9Qud*{;kTaG`{Wq2' + '7G%d(OYM|ON5s4p;d2MjAYP<_PJPbGEYP`EHZ($kuti4uz&MjFF!kcai)3`D2G+*hoc0kMY;yz-' + '9Liq;#zDQUY&Z0MO9fsiTMKmw(>NVz;1wK}pT63;h>KkT}CF|r@_(mrdCXoB->V;z}#Oa$UbtF0xypaE4Z{%?R)ugd69$e&e6sE{8ba9gfZC+=-zQ6pOtY3&JiJ' + 'GgtgexSX+exyeopd&%pi%k5YIFh}3Bf6OF>pQdg@-ma&b+16BQO#;G?yYfxp`B$V)ApHZnwq2' + 'S$3=Tb?}&LuLB=K0i(Buh`!JHIh0ZB_&u8Wv7sFsQ?YGIy{G)TF;`_oVcnwmGL7bx@*vByM5Y$*ub~P?Brrk?EXds' + 'HfBSdzYO5(ul^rUO9KQH000080000X0JxNCMB)el0Ei<106YKy0A^uxbZ~5Kb1!XgWMyn~FJ*IWW^Zg{GB0g!WMyn=Zf9R*b8TjC' + 'Y-BQDV{Bn_b7^L2V_|e@Z*DGddCggEkJ~m9{_bDF3I!qsTH)IJX26FF5(fy{Sw2&h5Ja4r+)MZX6Jv6eZIi)1ctFo4$' + 'u_z@&Euxi?4^>(myYH`xhmic3N?x;UwQ66=vW`#ILsS%m2~rfT4t&?CxZ#B;>sEUd#TgSKJ*2Ukl~#4R=bdQh=xa;nClm?EO?k%)' + '$Cp-dTo#A)*wk7r7x_T}(Q591d99?{Aw1vnBr&|DH~;xz`vLx?@7{!j{`Ts}cW>w~+v|64e%)Tf3$U$BI2AA?<1{|8wCGf|^HDHE' + 'NCvQ@3IFUJ+y4CP_QUHf{pr;kJxMz*2Ynf+n>F}73d{*peqZrA&3PeNHXgY(zG*G1RR2(zyOy4}ijEyu(_#!}0LCaL5ZZ^v3hn#!J}uTO0)N2YUmD;SkxbkXO&DOkJaL2OSLLUIC+PfZBD#DTywFL2' + '+GyNkKs9G5P##eL2_XoDC=0H=$3K*gYVaxrMCq{xYy7t?6Fjg*27O@H@L&r;$V?0*5`hY2g0Gb*P$IWTFnD%GPe|17HF0sy;iHFj' + 'E(3hYrxSbuPj03$CBZ$5F|1Z#IEPYcqhHww&W@Eu;qo-76tT;C8$f>Gb(VhuTA0hAZJp^}o' + 'mNPn_5OQ{bp{M~1W=v9{qHoi#pGC>cpwx3-6#^7&DgKhxN3qtcealu+U|JH3AOS5*+_' + '!PEh;tBgb0BA-%uDj`uA*|Yd=4>}`nM1hM?P+V&fz?iJbzAQ7KokD-=v6e^4zY-e$k!Aj%pytrY6iPr__HRYs9BN0I^rqs%2|Wf^' + 'Ch)|MGXQ*e)|8?J%c0KKv|{oE@z5`4L%TMe@N_DAU14bjEhqH5T_=i*hTV@uV>2h!;hf2>fipEjPYlzJEipY&4nS*z?MuvkEfS6_' + '4O}d|f%S7t!J(E2($Wc<&9UE}MToYlspa?^T);#$vs!;{oj!O^N}Gm_K6fSwbI#4v!i-Ro7pk_QWvO;!N8' + 'iNTzD^4&jf5o#mQ`h7@bQ(^PLA=g`3Tg6i2pasr!4x`L%=jc-v+^tw`c1cvN>%zB6RM@w72u1fBQWM%S@~6?6{i!{|hY_E~L-iLq^nj!(j;yL6bh`DwQmsEW9FhZ$ePI>QA40DO?FVdCgx!JkmO?1(IPaVe-C8ns$eaU0M8IA?bu0l9' + '_lUs3mbMTKxGra!HygDf#h4@+XV4HKqMc$h_4T$2^J3@PV)e!cYE#j7_h~^KX!V|Eyg14eD9zcupR>w`7hyM7V0et?(lo)q=*{e*' + '$uc`cIM(lmhT}Pc&=Q_*_U!M8DmHqmtSpYBul!>g^Nm{g^zYCSCMG3p>~N=>H8Mz{@H2i!=tEd4h7290gGTBtYe>SC&nSn=QkA2;6x' + '5reuLoLh0WA9?*iW{IeQgH}Z163x0rx|?&YY3)9@rCcRW$' + 'WO+Xp_D*R+Z4L8<7hb@A(Vw3C1zm8BpyP92ipOM_?=bZ2*wpxjS>e>u%V=>?jzKRH$GqwZc%_b$Ts@tJS4@(lf7Hu|NFRCVXgNk%' + 'HVq*zJADXP&5EPZ)oC@em;xW~D@bYEXVajo?_4c|c*sG^Vy|HC8nW(O^b&FGfV-=JQXLzNBX&xY1Xgl&LW=CuJ9|8wXe79tQ%hhP' + '&!GrTKDC}rKUi>2KX*RP55y!2XX@jyhKpx^{b>UAX7a{4zyf>YdT9?c5Z@TT&0Wdl' + 'F2{Y5r|xw8r3;@naEjyV=7iGjszjYFtjEVu^T(|Hcs+7_&gs6ES-61g55(G2j)Yb$t%G(+XtNMlwd7vpNZcuSj*HH}KXP-ovtJ#y' + 'Zrt0lc1KTVZU?1cx*U~u&=d=y8JvsUJHCuw?G$NLDqUA&*h#2V*e~@Mm}BOQg>|R>D0GXYY^}FF' + 'Ju$RIjTIxwI#IS~M!|rKqjMxCGvZ}JOvNE%>GQWzREtc0+IXP>h{oVbi(GTU~Y<1F)4yMzH' + 'Ue{4o)J1hsjQd+(ZG9EJDo4ePnJ}#=LIV+-rWj=yG;!7I+QdD)hm`D5l-6+Na+v3x9@;yMfUaX>7ejUvhO|+^D08!S_X*n&T' + ';%ZbZo2Z;eqs44FE%W)Ph%U?KjG7TZpc~OY1&GB-lSAno$}i#R)GQa}4Rsj*J8PQd;PKS*(tRKI6`Rwq`(P0ld' + 'ytTc(eRP&rFu;>lIUSetNpvz@jLxBpiz+%idi`QC9q&g!t|k-s_9BNK&KBch8eQZSbhxSe0AXwEv|7xfEIVB_tE$MdsGKbqRTJfN' + '7*~ojqUyw+#8IAV{Kea&Rq=UM@??&RY_Y5}8qanT0hzv8%nPoC|5bSjboP?kKCG%m)m|MdujcthUQY9qDXatn' + '0;Myco)ysk(3o2OCBAqtMRy25h7j_mnP6va@2k1g_SHx^?sR)$2MMFN&I2*aouRN1Oo#Lw(?-#IW=l=wTH=FwjJV' + 'q**Pd(dl$?3BWGT1gHq{hg1{2y?#%`1U1e8)(l_)Uz+B~SjIr#)6*UTf!D~cf%^cX_HZoU0~P^^8S>W*cziitEZ?Ar~zdf%O%f+gB=Tvt4BPbGM' + 'i)my5)Q<=!*?N6+kUjh9&GVPbZZALKjvoL%J8RpGsz6*(47' + 'w%)!we0%Wj@OgId^FLm`dU^2fpW%)yVW%mY?6k;n-;nSo4X~|6HLm50vc?%(Okr1+uTHCcR@7NJzEZs?>t=!cZtTC21538Hw#LON' + 'Y+ln6mv?qvRP$n*o#fDu`8XDT2T|Qr(f^?-PNLpl3_b|E-3wr9(MbXOBkYAM+}>dGVQx|b6&JWc&!KnYh@rzigxer&_r(a<1aFvv' + '6N?6EWx+ymCf#MdzaZ9`^PXlwiq6pT-I0vyr)#(~8tv)$gi*ep%i1r%(8dw%Q0_8@bEsdJR(LMY5`9ZYz#2sbp)>?5-`q9gI-Q*y+' + 'D%>gA(E;@}kn?JUODsjGUjr#N6)?yJl4`wLBEto81@5=HQr$UQ)QyEjXXzm4z_nH4(uXCGmGkAQsTnsc;|v^WhS%i||AG1gAxDX&' + 'rrtXoDOk*8Pr(8jLR$Gul#q#r()DZz+`){y@8Yf6G(MfTGrcK9RYrESE4C_eFN)6Cbldtz}V`$7&S<+fXSQFBaH_4w%^oO' + ';^zZITxz{@8T%lBTcd`>@CO!xjZ)4QFmUi3DoNYIz{WP0`X#xmrXIN;{;JBJ~7K6el3wI' + 'p>{A)dzHLL#Z0jiLz2JD*Uu-d=8EW}c9!4iwn;F&N(RwyU#HR6{ZG(9EPI;#-nHJApah_zmgTD}&;s)6+FEg9RoOdgnCjKGxaM9{uV0(@Wdn%1>o55SyRW|*XG<@4*B>|~D(gSxUdLwvLQo`G^cGp0MAATNG' + 'qnk~IffzuAi+~MghS#LfD0q-Di++P>KJUM#pQ+dtt&S0BRE`-yAB71Kj_2=4Mvfp' + ';KjNJYSun}RC?p1Xw2nf*wvam`S_8AuWE7mc(j=IYPP9~LVf0LN4Tq@uFi}A=JjOJuyzZo_Q#Jf^MbO35IaOe1Au-6A5ZdTbe4_F' + '*>LaUN5%y}26AoA;K$?)ii8wZs2$fO8aR4O9OmKz=nYk5F(u6N%XyFV_DBsdACD7M|EJTMW1u+O@be{~X*J8-g$CPQ|h{(TX0*R}pS!e!bW3*NS1xOaMn$0B_q0Q+vGZ=WaEQomf{eIH-' + 'J6a*x?Dlt&N^uV8ElIDIMqzZ?WNO$|O&-+NFoNA!@VP#{O6(FJ98x=mDtoF3b$)1Kh9S+1%M8e)ju~XKwWe3m4uRT2LuT*({sU;B' + 'nU)O#?8k$5*)KmmJNO|xdik$I;(LIJr{8S7KKMt5UytOwZ?^vN&wpjlU%pmvo=~BK_dmUU`7C?$^XnfD-^ovV``;lFUZU_H' + 'NvV>T_!Gd^p8Ws*cku7iZ+5=H58wU!76Qv&93H&?`Q0J5CE>}Yjk;dj@Xi?y%!e8VObQgs-W?ug@BjJsP!HnC)>d}-77c!TMekee' + '_4DF1UrifQ?$A>4d6kdxBxDpxLwp+j^!9a1d=K}j9?X9lkZYAWS_0(=R?F8EXq2f4Rv}Q{o7NlkP6c@7kc_lG%fNAD0Nmq*B5w%#bW_D)4ZA#HI8mN' + 'J}Ks-v#7y04Q%gkarY(Jx*<7~tX~9@1XdF3@L&+Q0PN7GWM8k(_>|}K3SdEtCo9yi1Powyy~Sxy47|Yo3Q82$DBLO>Jl);xZ)ML9' + 'UmX1W>OD->&u{+DvqyA@6EeKN{{R}Q#pV38sIbF5G*iJ?FYti0xsLMHM2WQZJvYs=m9J!yR0V>e1r' + '0!L&XW32?WzB*P-K_xso8}i5LVm?7I7f@x09p|5k3F*V$;8;NePBaDdr8DQjs=2FeWsFGEe3NDc$HN(`;X5I}pL!XTY8L+6DD_RA>6;1&6DIlUIPY?QgEK5*ptHe3QWr9vt?TZSAd3Jtnvk2VD;' + 'j8Jf42;t2yEEg87rO>Kz$}>H&T!vNCC`)1Unos0P_28o@ChYf95Oj;BMf$M`is6{_??MyHA^OXZhUUVRKl`hB4I5lRL+ck;OK2&}' + 'VbOXJll};nY$5bMSs!U*H;rhS(!qQuPWs4V<0SfDs{RA)f(OlF)@v5$#XLf{5D;F9huQBYIjZ6rH3k2|SQoRHpM+poYi$QVVj=!s' + 'j;~rZAJFK3kSYm?SCtrx1c}lIzD0og~$xvwm{e9_H$rEG@&C_p0i)>gk8uc2?z8GpJ{~uw@1TE6N)7!`SH?WlsMCp;@$y`NGn)9LOQK7s~60K=#@8}TP-LuKhOC?%_bLJBf!XnLxaV>w1Iyzsfr1*Cx6-n' + 'EjSbs;~7s9Bm-wST@!Q%Dv;Im(BD(fMj#G@$nU^qM7$2c`vLd_jA3KCn)-BQvdcNmEKK`2WB<9KjR=GH&NDn<8mb|5(Wu;6)M&!n' + 'E7xXnyaZw_VpSo)7gE)9V&K#6l$L46R~au2I>2jPH(ki!X1D=UZd0urmF*tsM*+%s!5fo}WqxTCB5}|ugYQz~a04y4a(!0Ccr_}r' + 'm{+eFd1}ul8k?}(aGPCXfI3Ih)fiNhVtP7c1EPKbpC#?`!z_-v}HQM?Y6m#x4r#d_*A2d5%V9' + 'RB4MG^;hgqa2?AYgck4qx~Ab71Bj-J5gOYO{UF~KS2Wgg' + 'Zs0zOFNzWLsjQ2!VyL)!K-WhF1~*;c!DvgV@qB?R;TaaX#+8%RVg>!L7vkWE+pQzE@>Q!j9bPkcY|iqA+@H}7*_eGeL1RK*UAHrI' + 'J?J2Jo`IZ=4G~NlF7z2KGbVj`Q~#1s!Uq*0xMU;}r*`BhONV^oKtw_L_z~4H;yBq-N=vAief(%s5#kRdPdsVAz;lTHG9TMhii?yK' + 'LWdR8;_dAnr(~JkO@zyWxzf?rsn{DnWYhAzu)#|&*;&1ssj~bSW0e%me8ZklttT+dWmY$N)sRao7V9HRFUHwvIi1Glc(7N2TB->^' + 'QEYrtElJhXSRkT*Jrk5#7{EOyK&(L3G0-XhqD}f@OX)v$LE9BXtrzzDM@&j+mA829)TA;^CpGKA-`2?1O9PumsD~Z=;V*P#}H$iO4V3Yzi)8e#I3r-(^Ag1`@' + 'g8CL-G!E882ma}Hz~P|W8#MMiBWDBEcw$iR6iOn7V~rCz925(bPAEW$J=Ze5(40TnXOg^-ZO7Ek2V6-2h#z3NCDG$(9}iA`!GfsU' + '%W-g{sThkq2p8>A8*$RF7pD!9H&Ao*SJ9ZIbV3omRXl(?_(Is|*)@hw^pSH-^Z8_zPl_H(gv)AyZZh!o-gQMj;(oWdJdn!pTM%#2' + '%n>-|hGIQD-Hnvy94)!@AYw0e=xB;E3;pVMo&T~swm9+~XM7S(kGDJErQ;*jIG^~jvv$l)&=!qe-zLvhHDzCe6=k>V`YnfuH=321' + 'T2%1xupTWcS~hm?;~Sgc%#$+dxNeH-)|s8O+Ogqxa4hT@sa433p?R0#o&-ysuBKD75DG}ezZ*r9zJ%a2$|yt6b|KhF' + '=-QhqC11G3@YK&NMbCC2{CElC{Nj>SR#n(}Z&xA>1@94Nv_q&$&U>dN>DP@i-ngXAgBe(g`{jHgG#yyr$nxyEcn=4}v}&;|vXdPB' + '8!A3&yw*K#FXQ6HDvFGih3)*RG_loD$_2P8$3`&&hW){=roOcdsO7uK?L5YBk*Y44=X?vxG`I?qL47@YIX=RF4qxs4<-4p>rtz;' + 'LEU8uEc(?{y`=Q4k2Iup>PHY-LB-Y|U;*q$J5h}PR8mdcPlT@#y&mb;O09lD-U93{#MOpR_ES-GxIc*Ycb+~KK%SzD(UX7ga!74q' + 'RitR|H;pFzp9R0tJ_ItKhzAiB>ufe)#sO!3l4t@PzwDR`}3|*2=oOc3u76uF508Q+WEZWuO4(v=;Yg3o)@R-Gs^XDl&t=' + 'bv5d!U(Fa-s21733Zit&C9zpFAc9aC!iPPkM^0xSIlXukf^Xxmpv}$HHH{{vQoKgpY#~?C%Z#kWft?tZcoq@>+TL^bgP(+1+ixhi' + 'r=OIc7$I`5L{|*HyHDQt{(D0uR8Ta9StItKwS5i2RHUE=?STt@*-N&JJ2Rg)>soJ@9WIe#E7-QQhwD' + 'u}Q+7?9?oou4@T7wsARxYEDtuA)5j4$Lws;#J7aF@63_pGi^9EXxp8)@Ta}xkH@_|h=4K3S#$r(TRy2rNCDM%gZT>eP|lp}r9F4t' + '&K+h>vX8cz3vrcR6yPa=`O^Evt*0wgHR+XS;o!3ri?cw_f10|6M$63K42IyLdJ_%IkDa1Bj<0>g8e1s^p' + '2VooS-3UnCr;RtJ+AV#;jCZ?fG+)icxaIrrb~ocZHs{a-@;FE5pm6C2#^u6g^si@1fhjSMI9q0=?ZT){J>0q?pV;)x(`Y1t-r!P_' + 'G0J;OZK~{9IUa+gT{mbX4OyO9f9i5p?Xobx725(v$T>WE{Vby1M%)~moN?7n0Z$m%L-rOXem2vtPQHIfpzA+63yV$B*78!Wj?(EU' + 'Jj|bD7=YISkXRQ7-jhmMwYaQ(F{y$v6y)bZ?oc2x;y9Vls)#k~Z0r9D4hhNC9DrJk^6rRBSOCS;zJnfU=0k7)7;)ghye&pf0VN{TT^;?_($4RKxW8Ae1@ZEm)0f{R' + 'i7ZY}%Mprvv|JVyhOEKpEbGv8S&r#O30@b7O`ennINFEZ-Z' + 'e+3I2BXOJ*(TlfF_Q~6f1HK{lPQ_Mkh5)T2&{D>9&=vjgV()3+fb|-Enk^P&Z2L>ebIB8vTsNU-QCI&' + 'nUm8IUtXmF(CR_R`Kale2biK88M$oTYGzqpVIS%+r^g#()uIadu|Jwb%EDkN$bqX#VP<##~3{-V(kT5GcqC1J_z;y$tq`' + '2OdvZ&%cPPzYcTtgkx!J2P*-)VixXqG5MXO;va9rVXkXj;fC1=V-jQ)&ljwHg&kVkgD3I`GqB?Xa|rv%S6l>' + ';a3~ceZqAz5^#>5a19?b_va?0w1mg=XuUA64JZxd!9^rDXNi&xOLTN(<&mh>^C&+>bqJ&CNFiPwJGP5Nd@2?i' + '$5bBbc+!zEh8xLyJOR0laT#@VM`L_#i5(h0(O;t7fs4BvvKY#EMff6(mP1`BjbwNbiZ|quBi48TWn8*;p}%rfQec4TYgX4t8(u7h' + 'D;u!j_Z)qXK0nAsN_n9J#%fIlFZ-gNjm-F7;xaedD$22YB``kj^1lBR&`;960nq_DPQi&mLr=H$WR3E' + 'N#0s9kGex%`0rM8W(ge0k4U0zH8axem8@>8uy+p%eJFU8r?_MVV(KXV@t<$g==ITI8XY`?f8PAK&mnHmPi-E_0dHAeZ;`l*hH)_{' + 'Rt|eYBWP^JUEutSP7l;PKUKaqDGC7o}L$67hcJg&7LG|C+=7zj{ru%qWoys;u-hc1cQU|CKi=Vb^O=' + 'bQBlMD1eCOo09R$&|AdUl!Fv7O^PaJT{e{Wvgt2{7!dSnTrK1+%KFH!WE{?>C{?(CSdu116v8Xhnku&>l`$ZU5Ck*@C3?%OB#EYt&l&J5Dw#5K>rf}df8NJZoRr$#9xWz?rgf*ZDMIu(lv2OEJ#r@*DRUm{{D160?DkD@hj*6-f}z3|#fGaGV&yYI#e;~h)}NMD-54a?wh(X`przh`#F%F36P8dp$IvYpfsUq(E^oEW7+>cyK*^=0' + ');Le9-W{QJeF6pKK01%Nrza`nAz>daLq1x~OV}|ptce|`&M6T5G3JjX$ESi_WumCFz0EVO0uT(qd@0b-0hs6|_319>OiN}Rb@L^3' + 'WSExr)3U+eX*}liE61O(ocV#RvX6eXl4AobKw}4&6TrMsuLYYRRk}Np|n*K1PNvx*H2_y+GI41cwN`GJ6PDUNfxBXji@$h!2-fLV4aX0j-=uwb$QU2DnB+51>@#htIWcZ&tB;eF5@LTV2z1&jZNx};F3UkU61SIaqJgL7|_' + 'u_LgBUBd`~Qn&9%FW!Df@tP?g4w`NSv?*4}m1cu}7JE;L)4a!shwuZKOmv7h{0o)ULPe6BlFt>IEcNW`(D$~M=tcH+zyHNtTpD^D' + '%ioW@DzF}t=<#FsJ^BhbuWU_$w=`qcRlBj&U__&a3&frm)gv6o!_aE@DgaK}!Smss5zJ$ArNHd7DTAlcz*usYbdFSKEm>QK5Hl}n' + '0|1M>5y0kp`6Vf^@+!yo*mJ4iPKpNBL=?{#x`5DwOkgl&Hliu`y9gk_vNTQ{_5m2$a5n(Jc(P&%xNs?uMxqhCrEO8t@PNjszG05S' + 'gkvS8P73rg`ak~Lz2Mt@!|iUNbX6ltnOT8ZpuTf`g~OxEGXb)Um?;X-xFmw>;LljK' + 'Gc+4b&6yb*zm?aJWw#P?;g)TQ9i59tu)V%jK?DJi2Aqukr(+eFFbsPF=xKO6!;acxTOtBxp%&h5wP73pZN6e(XXQpZys69?tlw6&' + '0hq4A;b2H)H@k^9u`SME=u?B^KE{e%fta{Q5$g6$%>waeGZkpBee*=dBKxqlWrs-yNd0jE-8qXN!fEQgh|Bxh%8&Omr^kC3|EX}3' + 'Kj3}bkSCQa9@#1UB2MAQ`7C)YVhpn}hv1)5{Cq3>NcQ5@PrqjGf65MDzB&5&^+2VP{D4<$KhzDyFR5rUjO;rAX0iZ4+k}E#N@rhG0?(LyMAe%qfC*!E&4k#SOH8' + 'T7Zt*{e1+xU&FaKs&dRZehpY3Kh_0y4It5uMu{HroI)B#R&&`5c^z;n;*+xJk^ez-a=Q1_00$~0o{Ew~Lfv=9rdx5&jxZd!u1ye4j1VcpuBdLmfW<&zS{yJ}i=l|9vY7u-T$l9!i>;{tVerAGh9T(8eJMNKPr^}~|s>j~cvNsVAhw^MAusX_V' + 'eUg(D$<9&P8{-LdEyJTIEkFiKHaaYC!;)j#2g486&%*C1BT5+Tw+|v%2gSVBI#RMtowJ86v+Kgy80kck' + 'rAED80>-6r+Wp6-%a)J`Vh5ffmy%wYl9riq5JF75#gKRz#O;!Z8y$_(yHN*NUEPsK<>h3yD95qUN}!VXv12TM$LY}6!Vvpf9ZscmpwYXT?69`Osn>}%*L82p*{n}-hajYwM|tJY4U&9pEtyz1XB-C!L!MA-?Af*#gGYnjhQz(tF=4JTudqcY' + 'XkWpI#FrIrfEC)y7PY>zri3hPP&h{_PxVDrDv!(f;z{?FQF#lLGhC%nCS?+(gOM@pR+b8(y&HlV5UK*qQku4@xC%ZWy$n0eQO>ci' + 'sbdc5PN%cZ`6XFBESo}NYhO-YddATD?DVVviZdfhWrfK#`Hq6jQa*Gti*J9=DZ?or-L@=dfOGOpvWqP_tmFMQCL2P=*L7qhwTn-#' + 'G2=qjxNi;z4|$N@!}2hoqRl=&3a6NGJW9pI_x2x&%1lseV%X|iq^TRHQDyMEOL!+>;;!}1%7~z?kf&&0' + '' + 'olWO5*0OGfUbOa0`cA7RfQ_j8ZrY35a!F__?p+CDt~CszDscCRaWwRjo;U~@iobpV0|i6lciMK%$#uA}9WRWODPbFmaS}dXP@?cS' + '@oheYg#N2H*{}W4VtJiup;;i#_tp{#>qPel;X7m!><=DTtdw|`23U)Wc@Oi6`}T&fLhG~t>)t*Rj%LjJI|B6yxjN2kkeUm9P!TQA' + 'H>d>x0L<*DLp!90sni~^Mb|Bi!HXG%0BSSKZf58|fI*)s+4PSr|>EHzrPV^Dzi0jyN' + ')gQ|X3ruV$q`4Z$Bc{Qcfnez5?DC0hD?_>{#+{aPjNZXaObW$yIbFq9d&$lfD97Ti2we&$2fNEB7O&>rT+XeI$O1GpjKplmvLi;s' + 'mObbm7J-7onh;=g0Yd`g`9)uL663k-8+ahODU%YDJ9F~7;W~_5fkXyqgivU#UCJaKH(BG2e$0`4Avp*ZKGnCmF^^_J{Z?!9&V=5nUD$wYsv7NXv#5PvkSWQiV7<2F(y%UaH=>K+_h;K3cZqgHli{3^?q$' + 'fj~swjCoOvMNA8M(He(pXG2H6AY88ms}rekW0RW%7}u4vG0C0TXFnlPhfBA~5X@>MM-CpzI2JUS9P~;bJNqDeFI' + 'Q&7`%c)*pUa=V(jw2Ye*4sbfgvW~EH?|lv%Cb>>aMH))&tOm=-!pYDmJJ@A#BxTq2#{WcC+q_XZdb3+%dZg$49=NhDQlAo^+hvNih2SPM9Mc6G+G=Ar7uGaiy9cFft<@' + 'FX;Vl(}X!Piq$<@9E>iu+BJn87{@OI6)TIWbKS^-Nw;0@05a&@V`Gc}D|3(pvBvh1+}e0xtl7n7k1iCCBXLQ-rwvNgaF;IxsW2b$' + 'k1rx|7ogU!X(vw1kT22%`!%uuuQ5@`vX;?zNUadz4c^IpTG*_FmC7Dbg~yms9=h*Ti}9wna7-mSP9wd^2@|Rpc=`*C8HT|?_yr2D' + 'Jz+aITI&)ht5cAWlR~)wI&gmHa3Q~m;(w}s-)^V!eqdB#lZYnc&Fi=ed*h1h?Jw(Q*p6Sj!F0!M=A5MM4B`vTXYMSt7|^>WVUsGM' + 'Iap__cf>wu6&uCv*wuoV8Lv}eLrN<}0};1v%Boc_)qU#di{m|Ib}H3g4zxhUZqu2xzSD<_*bwp^VBwFHzJM2)zSqGhQ5isNL;18S)#;sU+QAd4$sD7N7}o*X9a!W6P}qgVxRnGGh;Ep1)2XMKG6y1|' + 'X`>59LxX{b><@kj<1WJ4z=wwR(ygG|Od1GOm`{g+8o-=59Sx_Yqeuc`7aU*~DivvFhD?gXF9I_2>DP6HW+@Vlw>ms*7`b9LJWm@p' + 'xUT7D(^iu#e`%W|lc;9Aamr4-vXj=buE&kvCGC0BR<&g;(*4F8Nn3NS4~{pi&WeU3j#H6lKfrpNBxjzlPN#>QW^(J9+8Qspkh&jS' + 'GQ_tB-V$*P0IgNi$h@x+=yviYpp~n_8qmHK*KA0N6#}ZuoyF_68BLNq0S0AnW4I<0OQ9N~}A+E?U2Yh!NsAZ((>!&q;wwCY?h0J;}TG)XUqsnA|6W' + '{3?2e24fjLPbQ6S792cO$2cY7u#xxeq)LI&N!hq+K#bA-@#JzWl-Ed}JJEHOq#yl+w;3gfnnM?18+@If91`qqpf)wmx+aa9oRG`q+p#lbex`q$^EbDBv8t2rG>LorHL;8~|{snnAfeQ$Px1&NjO~EJ=7)|_y' + 'X)r)M!?jx%imze&`XK0}O+FMdMrwR2>|(5{Rz-pSkn=`ZCFmN{HThRc(RG6>1EQ5vcaxA)p^&Xn--&2ltZl=&Zx^iBM+cEebNfJU' + 'RX?pj@P1n@Fw*R6Ok*F|AgN4m$%P`$R!RZD@;UX;iKR)Q)N@82U%3@Q|K8&T8w|C9`z$W)j~^ey7D}!vA3x$b6JCb9sQcpMuONKD' + 'dm{{TS&I+_Q+PqC5%5^?<43NF!b(wPeB4SYprQ6jG0K6YjIg>g&?HCvQzTe}I$ISgWFHyhbtM`78deY0+ksYzHlc}FWF<@nqSG;T' + '1k=>3(Qyzp_>isc9WKdvB`WqR{VK2=MKLkZ5UBI)Ad)j)|=JEu*~XKQb9K=*iWaYjuRhl!r%Kh' + 'PoI7x@15FLSU2|b`w#HKv&ACg>L4JwMI;IeL9qSV2kZ9O32H7P{cp)?H2*(vl!u194&P>N*--cZd6${6t5qNcrE!{kd@>3M(CjVi^' + 'GZ!p)CVHRrJsnZh`d6Tq&IAys=Jxz`2!K&Bn1*1(zM4blqhXUtV+gJ#{>g-Zfg^JKC>3p5OUYb&sk_wl;S|DT2Uh#OEv>a!wrD14' + '*1q`WGK&?y$tZAOrl40C7u_*Nk5$b$ZCU((12f=AL8x=Y6n(Pz){xYiMyGX^d1A;{o<+8O>X`UIFEpTLUo#Y`d$6*2kGTWhvV1>5w' + 'h$6r}`0Y2^MuPbxg$KGSIiS-*WP>@$KfrcbnZ;N$==z-}J16*nas+<3#jUrwy#0pJ_GLbay`uN+?KWIaKQIc+PbVsIjQ7dOxXRzf?SLG$+LH%aK;wLgP8w=jqnKT#=h1Vyq>R9~50w$w' + '>qcm4$wDG!q`ZA5F0$w2{MDP@zlv&+++qcdF^ItUewL9iZpnLc!G*-_(O3Loe}wkYxn=3d%ymX)!EoA73Kqk_+?1tt7LY2jpt}kJ' + '?xWD=S$W!QP`P^vC`yiZJ!=W)GsfF5haF(`bZRm{MnA=Nmfr=qKvM?wMJjl)P' + 'RaF+%A?8WPqxA<`n{G3qr76lpg)B6PMvFS$-&vMBakFy?NI6N2b$Kaifs^E@r?y2qx@#u#5n3l5_swEb)?}1e;ju!zeYu@RI#BUE' + '9=Qw2t8o&j(Xs%V>@4s0J~OUG)tLxcPqkD)Uh2f9lgUoLBy*z@$^tjcgrx$fqra=-)htW+8PlB##OSdMfOGkQm%`v!67q?9^B^Le' + '5SSk8d-M;IN-UG=i@07W;{~Gue0?YZnmG77FoXO$1r6CrE*)`ExLr&nopFiB*yHGOZ=X+~S$-F{6E{BeDfMMT=Qzi?is*brS;P7i' + 'N?azDH!JANT@HZE)&Pw4NY4BI)o>XKK`*)KmX_#`vrg#{*W*X{ihjO)?W`;' + 'KR%Gw70Raw7@Ifs!|>Q@fg$B#@0y;*?S!sLIeg6KS=i$q-roko)#bQBOfh;*tyb3=$R62~acta>hCR;^Kg;&J(x_Y@sOa6|?I9C0' + 'SvPv)0?!K2<_@qP0|=Qo(Xv`zg-w-b_E8b_=2@WJ7&gCsqagM4XiY(BNURokxdf%0g_)CkClZbol$!~^X`|XciVI7iN{oQd{}X38' + 'HQg4&VQx~|nWjO~L>nioaykZrfF`Alttw72' + ';f$LE^far1brspl%@0b(D;WSSmAM#}IUX!zn8pSV9W|2@1po}0nzcJBe~}d_XY6TJd^U8=UHDI?0mb^kGFV_;6lNfcA8_E3Ug>!8' + '$#W*`N72tO`KP)KqQa*L+?XL#OatFM!|a%~QBWnGW>mq13PO}v!#p(m4$>N!e&pa!V6J^sZdrIL=NE(m96JcEj|~?>50&c6&W_`0' + 'Xyw-0O#`jkNsI{HPxjs1Bo@rap}yP4Ly4Noj$v3cZ4JTiwW6wxx=USDRb-d8M3#;0^vS+4lOl|yg?#K$?c0#&d3xUa3;a3XP9>z&' + 'o3St`>p5se7DSf11S5KmiB{&~!Wf-k+I0KKL@Fmiv6@02_`w=xsylRSY+f=Lf^6!~YW)lhZ+lXt' + 'oh&YGX9e~#r@*5VT}iCao@zJU+AyS0h~fh0p#f|NEQrClIN;rh2?7<=q|K3*HzqC!zuOCnJqdu(jxXed&*y9gYL`h=G?r>si?V4H' + 'UvZvK=x>@(lULHc=b`jqjG7nY?C}YLNH-?Dq#ww@qgw7zR8=%T%9RFEzD$UW&PTY%oPmnNsj{ZxCf}FK{rh)0#*jW~aqR2oQuxQM' + '=?Tgk(88l=|J!dVY1;m`Pw~~aPaTCZTP?#U5r*E$46OA?;xZdAF6SY|lPliusl*kMs_MHH-7KI`s)C?yw)?$VS|W`mb`;77(qW!F' + 'M}bA@Fb}l`yoQ2{VOkPTzAWc}0iE=Xmxd6t`q7Vh9sT8awABd(s3K8W*Y+vv+8$N-#%Lmq&eKRXvzlWtDl^|Vey5MTJAY16_50j#en2rhWEJ?FDuij6=PiZL&B9xD7-@Oi' + '>9cdy3B)<^N(S194S8wrY3kf1avV90;j1F(^m1-JT^LVh@>qG7A{s9Wj&YTbFN&(cU>+}kznv+r){p+qH@fiZ0g4$Ejcj2C_Ii8$' + 'o_ZR+MqR^&Ff}4oz;CsyxJ=AIco?z+G~2xO5S%)-_KjWU(ol9' + '*6)@>(JXH;*#)$jB4E>1d!JHlCi1+5c%Td9h4M&Z-x_x_H3N2uri2qt)+AmJooX1l1NL`=9QEdEMG^DTG(8eo7>{@qe0$|Xn)6bkw|Eba7^Mz`~-_V{4?rINAO$I~bcSn49RIe2l-IFKhZk6IH+Lpo;6Wa_YL8&7^P#mVlF6%U5GYs#3w$rN5' + 'Ue+}2S6Clf#5}mkrgsycfe4iIokJBd0JZK=!KXfK%rV^?-*Eh(PyXf%c=NsDxk4?{nRmoNy$M8' + '^5_>eR_Tr=dz@=Jw9H_RaEW-PUDQPYsx=L|YfEv#&dBdk(gc' + 'cI-DD<_*$-TCE%H2sBPJ0$*KoO9xV|Lx#Y27nF@}s%^irE3_aV>a~2?x{ok7mZ%sE{m+;miv2KHL`=z0Zn?oi<~Hoo0UC6}wbo^4' + 'S(A7&7F}ZSEi&*14!?B?YHKL|rrt-@?$CO2PLsqk8d{6WJ%(}*Nr&rKgbPBQw+(|!z)2JAo(m7xtz^5t' + 'bs#(JT0)(n(pvV3&xTZ^&3K`gEh10qOmk*lZf;j2RJkg;D6^j' + 'iG^^XSxqjRY-rc^be^n$=@3n>QtiK&E_ARqGCw~i=f' + '1s@9b+(W^jtDw{s2(CCo2Xq-DDd|Gb_GLvy^R+0({l1fRzX+$QnSS0{_W|15|uAFXP(FYe<$rw=y+1ub3gKW8v3M4Vq*ZEsQCq6' + 'dzM64t^pJ1!TXO|*?8QxkA(#3||YEa9n8g8n~i0EV@|Ciou26CNsXM;c!kQVZ3x9a$Dqp2Q9fiVd;r' + 'v_Z7E#%H5t3kW)lZzSHX@hq30M`w98DZ<~xSxfjMvxR(C5Qw%OhUII5q1Dc|`(8L1S+%#WSe=#Yiri%7r*&o8g?`)xZm;>Wlq{m%' + 'XQO;{R$vabQSf@(23iFx;B87X2o@qLOoVrqKU17ASIKph%cQSY+sBViUq60~=x+_HqgeLJIp+m$L=0ePW<;wvm10i7c%X%2DFi1`' + 'KB^XVZH3Dymi7AhG+qAxIv5&VxTb!)$$Ma#J0mo+vDtf+4aTX4A(=&>&v+aAB=pbiH2Up#dp9)U7`eGH_OS7Zj(>KY1i83)@ZOf~' + '&FtA9J_PNHLCc!I>g_N$;nE=ubAxII`)DuVu#j$#{MQHn$e#Up@a`zuiSTds%TLb^e#nkq{_8N%hvPs-TI1?Y*sjV1txvIp+!MPJ' + 'Rjg;n3H@fc#03SFzGHlq6J>L152m<4bz9BexXvXmRoxZB9n@X#M3=I-ZeY0CfC4*6Rc<&fHvQme*O=QAXJy(v6~`&3eJy%CE2G09' + '!tBUlNguYEP`A+>%0Ih{yIrP=?#pfvk^rcxeCVzs34y98BgmSLnlhDa&Nj>VBz%jRc4EBSv+KugR16`A{dHuqW)6Z2%&n&J?ug@R' + 'k+tAl??WR^C6<%wXv>S|8kB?4qyoZPJ9&O4ugJTy77W79a$Y!utRV>6H<38SXq7psIxFQEU|a&K(6BC`9Q)poTn|X>vko<9D5!De' + 'ItZv?M~_5B2j90shapP=pJq{I-)^3UNC{e)zjswEG6!1cU%' + '$UF*xmV#bC{*Re-O>|5)&%5{2G|vc{EC2D&f8oKv' + '(cSEoxwggK#qDI7O=q5~$r?lAMwZa~Htq%NbsHGws*K~Qdv-X}==DI`yd~a-xJ(Rl7bri;F&|fE)rk3PKr~3UY_}Wy_mKeh1#&Kk' + '@4lq#zVu{Wq^3VyzV%#qHj;1&S+|sNE2`u#Ya^>YdLnrq%9ZbB^bm28#bjVsfgdo4cN_}mYS@Z7;oU=pcXhiH63YID&fd`sa|icu' + '6LYUuQzk14NRA;qDYkMspWIgRGVWJ*CqcqR9l8NzM2u}a%iwzvx{qHRHE)y1vZ}whugKpUUF{AR-sEz(6U`Q_>na4Iv2Hzl_(#(6' + 'x{DtO*tmVo;b$@2GzI7Nq2OWPonqedi6?ggu|E(|?9h#giud;2L=xs5C984ol9E{gx>7g!ym9Mrq;Cau$Q+^Ku6p*%' + '!K;_gGm$;>$Z?O&@!t0@`E*qrR@I`4Pq$_D*{W`$lOm#{wbSS~2hQ(lG=T3' + '|7f7dXsv;kbxA93RK!>*jbic-k91&j&G+7hGV1DH^#Ipp0AUyZTVLS`BEZYo1`17wPjT%A0ul|jdO!7&ojDhuEVbwx!sg5_gN-#5TQ|=KH8$08x%zcuTY^^' + '#6D|AUlD;p&M(@dGo4fD3$q8<^^KdqpvV}vu5LykK`Rlg0&@Ba&alap&BY>C4`X=u7!v_QA9&4$;K2IZv)T3y^!{dZwVb@uExj^s' + 'QDl*tO@i^3tK*sBjKSw2ZR64FdeGICVRYBa24c*C6p}65l}*LOtn99deOvU)KLxgX>fhWs-4xGK5{`^~7ak|z(1p1ZB40$@0si0R' + 'QlK`Z>{2d5^~ww>;g>%Nj@*G@#d&t|Zx`DX2gU)c(fO`1-ef0u1^q@WK8|7JrZ?v%Zn~108#%!|lJR2C%*uLn?2QAv;`@9UTfuF3QOG7|hQ>Y(!' + '3~46#x}AqR67zo}7QVnR7a0ZqU!iwfGrv7ws)twp(*Hd4)@Xot4bRsXx)@%pFLX1!@n6yk>*SA{^SnzU?d)cF8e96a532GBiY|l%4lO&$7FGLK>hdJDR@#B$h6?Yis3j9sQV&0f5`3!khQo3h_uG4aUx`-|KeJB%LloAzO^YW~yir5j*' + '*%?PBSxO_BZ7C%MTDI=Tm+-Nm8r%pA!R)A9UgbyCfG$EiHt(PzBKm}KO@N5tL1zcV8_TJCtp2?GtVyyFt' + 'Qm~WpEsLr(%HBPmz;&E-+2O~L)m}9Q{y9Xcb~QBfp>M++hL`$p{Q;E5OGEuHut>;nB%a_K$)-vy{lonx@V()``u{!K4+LMmTevjeIi@9M8nvsk$h8#085xR$=hmw^!dJ9iPt;={iVY%I#Krf}' + 'eQ}pqYKl(Uqp8_GyP7!eHUS;-(U;7`TI>a>dKi5Bp9}r><%htLZJN0XAZDy0=PFnZ)95V}RyLA;pc~@vfAw?HhjE(SI0PI@C%HGB' + 'si`o))ZUcK*qz?AW$mU&G3BUX3DhP6N^OJwtQg?dqvL?}F2L9V_BC<~S!4VhuDbynoUGvn7haR#5H(&XpzeK!$Gh@l=03bXs*+Nk' + '8QkfiB`343PEEIFv4}i1pzO2&%zz+W!;I(r9S9kPt0HR3S%D!FZYMsCvKu9>eEZaVUk7jczHI' + 'l=HExjU(>DptoB|Dyzy7@;K*t>{5q&iqB01ru`+3oK5?OV{(h=5' + '{nM&T-eem6>x308kWe3Ef;?{#5<;DJ_xJtexb+@O-GLa5%?)r9@P#z%!|x9C2-?NXi+>$b=u2G023fmFt#EY)J8sn3pyv$7>%9e}' + 'yZiMQ&2D|~v%7)p(`09PAkL8HLH?JfTFuYf5oHDs28;17btJ^npRok5a&x^|`@B!*|UCF-5' + 's)BFsv{98lCK2JA1Zdyk@i=IxtQ{~tm9BgZT9Kn?>IK`YTD=t#%2@1@a}SN646;w44B>i{o|9<0g8{+%@Tg;6Gc&' + 'vvQ91|NUKKg{fFjVh4`}#6StX^~a7#5zh^XH2Ihc+D;AA5f;5}Uw{%AFTC21+5qUdc}RVXMU#it#JH7B&DR{$yfwz;ZB?(8PW(DKfT8|r7cq!>XsSwEUu^ES_DGDK5fw5+LMU9;kfW26Ol1X;3Pn|xkcPpS@S(sQqd;H|5fMa9K{#rMv#i6xJM5Kx' + 'k1;%)L^$o3`=8(*0g4krf&vsI!v_t6frtJpQ_%m!ePO|Fp+d-;rY^yZ)sGF*Tsi=iKaD|*?_J-6mY@!{|7keFkrkFG=xeYQI`RyX' + 's5o@?HYR=2>_{7aF+t<5-eIs3q>993$`nsU!?UB^KEFAip' + 'SnPnctQ=qcr?Y|P_4O%Bxj?t^tg^YzT83K8+v>G7x3Mjjp+D$BH#dPNbZv3pML{Bpu5w>s5Ozz7g2&cvWjxH~h2z_A' + '@6}{{M-hA?Y$Z8+v_D)ir))aHRr7Pz#0^^IFdsb;^o}se8DX**kXj1SV}_zHMbN_WI`c#%$r4(cW{GL^IjhghB^|`k' + '-=+8qMKd>Km#6OCA{%BPw%t%}Qg_UosrSF3{^=JHo87QAkbhv)8wTE`?ICQ3V$XBnX49dVGH)e?m&u(8wJRE&V%_&f+ur~=jmMf5' + '0pF>uD7-n+uq}kguE>Wz$R6ZxWwJIy=rXA_bHe_$$=D79aIodb*66e}wCrU8r=`$ld=Hc7qu{AT0V2tcZEVZHmRq{0{BHv*7O%lM' + 'WVlB7(y^|`A$2wW*g={F1J_EjcnV|zf?tnJTQ7KhKdF>?KZWmj9`r@DS$w+Dd%=vu!&1' + 'p|@rg<<(*}AG-%>Y~eBFW#kX-u`t<&z-cfC%2N4^o^N>_Vn(@eu%eNz0&O~dHzl!kJTy9;=!z_eTQZq0eAd>0V~25@Gp$(Sw9#Sm' + '99etT4_RP^DTY4>tV?+ALf`zO7pY*-=d}kYP->Gi6e!kujN;gVv23Sl*ft4U#RJAT*8+#-yRCALZO^U{Y9VN;!&>W&bs-Zr&{n_G' + 'l-+5omH_U@W-Y;eiS1g$xQh*2*1p{_$B|$l(b+ZhlHxZPz=G?J!%d5loj9NQ4Bqde*l;EvLM3LL80$$Fah;Ubwa@7VF@FlpV2V5PIg=dK=CU=S^n6t' + 'C}Q~dp999omfJ8(?LgtIDr2zo3tD7*{aw%xmR2-GVGd=tXhKQq^e*aj3^j8RI_TKh;f6s&IVTl|5Z;US=z?P~fSc{Q>!ZE5b3vn4Ho{Na9V~(mM)(Jj@7f*1dd{%RW' + '@aJ(nPor-VIVl!2F<~Ymn9~~?Vf-0z@Hy#Ym=ctr_Dqkxax!=0w0uUbarZd-a|Fof5x|t;CRP_evyrqQR=Ah*^0' + '1f-Dov+o<+TrTL!%07{*+gr1wwM`}pLUBPpJ_-T`ZvguGktDji7+@znMsQF%0' + 'Zr0D<`hI>A11fLRPr|EH*m-nQoazG=pzSB{KP${}fl97iT={A;lef*>-AW(R>2zO46o~>&ni<>aI}8)8*BJbhE7+k9rjtblO6Dxu' + '`}aY_eZ^#Z)Z0k10(nt!nk&Cd=k&`z_fmmO#YpT>P$&_+c7iD8$gTR`5d=qEmCGGgVNopcCtVyhJ%aVpniV-pa4~LZSUqX2FApiU' + '?2DN_0%1qsGtl4e7jG9EI?n>EW&x7J`@EOJAzevsd)0SYqaU;BWHwGu#&z#6C$oAx#mk;mat|ZuY{V)|4F--XiSKh1(<|b}1DJ)u' + 'akAxIhP`)$n=thySgTCumdMUIn9s7DfTER8VTXu~b$XRew4*|FsehVlw_SidfxS8+pzQHLKLHiZf*`Z`hx<@Rm)5$&5VkR`v7P8L' + '@fX*x&f=Pf7uTiUn)W-ER?+*;SY0c*wl0(Q>fXwlt#C+K6aU__@>MpH3v;DkP8xoV0DIhkcHmxcCiPL_6Y;!EPv)P1h|5nv#FuEE' + 'Mo>7Pqa*+!JAJZG_A2>;#Yn?owo0Q61g$gG7HHY|bkeUc#*Tf;0t)g&RbCX;aC=ff(QPDryF?!&7p}&+wyp*Q#&OZO4a2K=75x>u' + 'q>EuZPE`&ccu_)8Ln@IhA0`Ux3;C|hCmG#YiA{Os0|H4S{ujXLyHW`^jaOn?hR>`mFc`T`VQT-Xx?_aYgq=6)?I##2k0$)S)mCXE' + 'xo5rguZ%dY1fZjxpdeEB>FlS?FRRtb=F6QS|EtH-TYw)YFC+|nxd}h`F`zxYu#-z>~m|}#I%8KO~FlF' + 'xzLT%bLaAcKglhQ?ipo8UQMsF;;L-2YF3Xsda|L*mvx^>qB@EVVfs7RYgO@&AN?u$_z`8A=B&W!Isu`3acbO*VoO6lYF7DldL6+!' + ')?4?lXT=;q2BwW+)&;2S2v``Xlr`aGRZdx^Sl2TkO3Hw8xmZod(R6`kD3CzTKoIHiQEii#QxsFMol$O_o7W2vSkOQ>#Q=o(c?By2' + '@K3{%U6~g135YhRp(B(j-W^6eB&Z8m$Dj(3YHCauhh2l+V+C{!jHu}$ZyWK}U`~s1O9)Acw>#nmergY4@zWPYg^=BL^;#8Emm$~Q' + 'z`l*+;cvgoIFarxT)p+*-k@krZHr~mDq}iAJ9em6v#Vuxvy_eim3GMTMtkDt2lejH(o7ftnH;yl@l0ZScAO`PO65>Z5srlHNG-c+s!vQON`JgW3}x_*Cp%XB2WiIS8RL^=d=&2-`U`-^jmZT@XW+)8Y5L9ml4`ogJ3' + 'oXgJi0X?9EvrpE{QgP3R{Poq>UJ9C|^N`C-R1_!=`%G;%~UIhw+`g' + 'jBn)+ZD?g-Ps$Cv_}6;!ul44?Cy#zZlN^e_p>_<7a%esIg_V9eRKJ4wlbYtx_#HqJcde~wA?0%eAL1|aVh!Hw3EV{v>4HO#7sC=NMNIwScidRJT&6^M=6X;9ghqnItucMD5Y%=43JF&4L5' + '2W^~H?c+hyc`zuEyqzQ@MN>4;?~Hw3{ah3D_R77Hq6738Erhaad=>rl{P6M7^F!JZL$X3hR0v86N(foNykzn+SSnSfNr5^eS-??%j)%V*CI)fOEK}f{Qv?KxJW1yTky=+&$oxfNt*@JMJP9!Y?oYK_VtmGJO;o;quVq5EW#(l' + '`pBNVNs>^qGpAd9l*|*cxT8OJ*cHRvjn=M2*wzMwu)e}cRpjSwdsqO5Y2{u`n?b9I6U|Vz^S(d?p_nzR?SUJDd7HK`>xzJWFx~n$' + 'zBs}C{IK1!FH)SjmT&Gk5`)_Me^5&U1QY-O00;m803iT9=10L~BLD!_qyPXq0001HVRLkFY;AKdZEs{{Y;!MVb8TjCY-BPoZEs{{' + 'Y-w(1UuAP`W^Zg{GGAqQaAk67ZDnqBVQemNdDT5@bKEwP-{)8GrmpHpnPN1uoyyG~Q>!?3R#ld)vXs5s^QjnWh>^JBkQ0(4OWMnS' + 'zizyUCugL1lbkCRafn8v(P%UpjYe}ioxaPv{_Oop7rZ?`dDm3MVL@UlSRL2NivyS+>~9gYKl!of^M^3H*Mbq' + '-<8eH=IY-%z2;R}7o^=UNLOCh!QzH2Zr4p&_gxZP+>l92Dw21k7}`~v3d%b8=_hE2RAqggk#4m_m}LF%(@#*=Llb3F=T-Fpi2e?B' + 'x}d*lIub0JdRbm~!4+BL8~EMd5S^^P-PE@M6j5@ka#8ldIxpM(s%h4MY`?4#A7Bsix(M>1D3?ppLS=@#AoYTLb#i_bbbY?KO(xUn' + 'bTV1C%__*UlgSu~^Ba5lLfS$51W3y>DjxiXq$CLxm7~%u$4syg3%LRzAS;`dC#R2D7Kh3tnqE`#OWM%C6hIKq@*3)HT(e-89A3HkE)=20J75H4JHX2jfmT68ZQDNLH(1lo>nU' + '{N?Cl_SX+@j=syzkN^8E!+&!0U)j;chm+$s+3Alb-@X0FU!NakZ$6xUe|&NL;q>_QhwS+D?8l4j^yuX6IfDS`?Bm<>?BZ`{;??&s' + '88=AT3DrTy%X87uzmU8z$}WaZvJxCi=+Fh~1NP>)2Ih23j0DvCyuRM#*M!RW`AxZ!#GirE_)k(^-}Iz70)FMdxNBAl5<=Ru7D!z5' + 'CnV>Bg%y5W7a$woRn3C8`H^(T#r;wDusSZz@qce&j`R=schM9Q{1A#(rM3B@kSZlTFo^sNG;<-4WI*' + 'A8mL=uA3H$s%#~R!{LP96A(xchpcaINu8C&JumGDDRdDO@F1i&CGe|eV$Wsj8nhWmQ%lxa-Zv|#+O&mU#+v9G3cyP4V5%Uc*)3^n' + 'QqlQ9>cS{obvcq!l*kMe$&!wn0R=N>owq=BzPT@}GH)LwTr`w>OYzyJRGi#RwtdK(Api<6ixE+srAbp1eVSIwp_^0p-1T#Qid9(oTy^84WjJ~W{GRvFkFU9ca_hF&7Sl9pTp2*}*k2G{a%Iy5znm3ED4h!;`P9(^u5lF_vXr~^Tet4(N85ObzP6CGcd)u4Y^' + '_4oM$X;0y2C^#CKz>GgQL*yV}MNrRq8w&Qta*UuOW!|Ol@~X3O0g^{=eWZ13ZMI6R(Z>39c-Ej7q0;QjpgH^!Xkb}1(6GuoAffHi' + 'G5R&b!-v`}9<0!)^XLI)a;Hmw^Es1NS+-`Y2h?<|DI7sMFn|%C4nit%ua`KdZ;~ZjM>A?z>n&?V0Ie>$VVzY-cbBf0iFg()mjr-`' + 'j;0mRO-;n~l{u%@RZ;X=Rd1hFs4{S=>9Vz0eW23eOBcS+bX{jS&vKa81hhveA@LtmP>-Z-y3m1I9e6arH(<+A)Tkf^qKu(h=S%@R' + 'rBR>(7^Y=Zfs!jlM|xxVQq@-=bx3A#WF!>eFVsTEm$9RvYv@y~f%gxc7~!!NG)4SRjO178g`v-xO&nc*ODtDMe_hHvrbUW3dVuV)h' + 'C(t7?2fUnH<3`sfYAo4nT-aHQ+9S+FYFa3(Tlg5Mg6i*MPuV@fQ-tk6C^}uSv;KIrdSvJisy`(jr!u8*ky|TdD>n@q&{k3?H0qC0' + 'bh@h9iWK+Zp!|F&Z8hn`m{bi|iKWqfGyk1~rd' + 'lV#Q9eHh_7&x^$D7DZ#UXtrC(3|rcqE-9%bL9sZ+F4(`T;Rsnh$CSb?i!#A1f;<EcrpUG(S}+Y`J{VEj9rga?UO`)Y1FBj>#`rF5CK5~*_pQ~xCFecX8GpWEW6oOHP7{1C-Z8XLc' + 'n6^(3k^|T%){P`YHAVTh~m93%3^@lYwnKwxTM#tSYb3GY$MOXoT6Q`N<j2^99k|J)wqaPosG+2tk8@Vw7x%bh4u)2c@cyJsWGsBr6S_$|m_|SW$?MLAC}fd{c!{>~J$JH&q2twt#%#VjJjx~cj7Lfw4}uidMlL}0CfFHsg?' + 'SpfOI36(Z$VQLO_`HHx&EQx3&el82TY$7J$604{|EC@>1u;}No+q@>>LGl_{E>Xl3L#)2wYnMU-d!z3ssFj8fwi>`AZCr46~`DSpAtmv3tQH_QW0QMpXRVI7|jZh4S+PL9<_6lizmvqgtN4=qr8ka4#' + 'A&gR4^weV#4hZ%X8U|ZiSlpUOJM8!Duo&2-tKdDyYqbl!SunGaAeFv74Zk(3)I!HR(wabd92&w?J-My9%c~7hP6riDEw(251b+T{' + 'qiCpMbE9@c(U)UvF-aU(I41ONI;Uh@D|clZ9!;<|%7SeYjD;C1M@iX}RTxF+`cn4b-v$Q(sX7uIC{vJPW3x~}{UVs@;>W-;PVX)z' + 'uJ&HEv&&nqR+PaIcCijdZDb9K8dwf4!yZ0Z%rFRzsW#DWr2G9??-w)6+(C0`F^i17a<ykl~BOsGc9yLYEnP>&ZhpKEtPtnJ2AFG(lvvmw>WQ$c^8gQ(|%W#`-sYM#v|LAS0BNhxbt=xq5PbL_^j{OQYLdu~HvuBNi8<' + 'Xg4%AEx5u;LoAdw2%Zzd0wclLm>AJW)=L(lCPwPzp>n`R?e0#Zg5qzKgSjrJ<;;LJ>px?_8;lI3XC0eGPmOuNXNC(JJ4~c4mYo=)' + 'Qi}4fENDz8nui32lp2kHr#7%J!T@_jXUk_Ex|ro6b`dD?g*AF8Vst1hj17*BQPIm;hDAFH2Sr>wUIY3vXXfHDKS>DE3Zr|fkXVWn-y@{Z#NHzF9p#jRHEJs|Hlu#|oh' + 'F(uX^IpTO#SX&sq5p%o8j$isFWQMHZhv}l_rlpQcJ-gO-O=#Y0Z(aMyS7ICas^Diq^VubBg%v8;6^t75xSmh9s-aMHEFK2dc}AoK' + '2iSznE!IjBlplM%rX^+^j6S|lYD`R#7$&ttwaLyrnOF_8@!|Q`O)t<@%ly^l`Gdr=7' + '$5o1g3Kh8-4impjCpNw8#HN2%CpNv&iLy5PP}cv%F3{%0(=$3ynv!JC0ps=1nP-x$wP3z^Mi-PxVHmj%-Q%!u$M+1s-dq)tZ%?pd' + 'S^TWri~Mc|t`Z^7!&hIvj)leOjeth86}IogT>*C_^*9o`h2WLh=)AY$IX+pg1#{nG0m!jS837&uHEh`f!p1L^mb?>88Lw2etewF6' + '@TK}syNd0QK{rySdXbnM>3~+u)vq^4v0XY|*tRlyhJP)G@iPV}X7!>kVN%}n>1T(N`eIL!!_AHT{wD7-*p3s_DUL8yUXoGo?9;7O' + '=eC@APO-zV{lK%8H9fYmaT#H|t87W~;Y1R)G{zKReT=m}nP+@@kMz>Z*l|MY^0G<;goPt9!XgnVr0axGjEN^664n#Ofp5*QYr_zB' + 'afsWrv-xF`*AL8TIE#bBNI7K(eX|9x>f=@ihx?srDIW34b1)`;))`peY$x^La%}rTiD5`IZsBQ?{eBibaayioK6B6`JTgUXUAnS$' + '(!EDz+Fg&`ZsBb}`Z($3c6>VP*WS~%VqzfKMoicX@?Nsy@}k#BL-&p4*mG4nHm_QHZXdcY(r@sN{nf+Bb$sb`V;A%?Mty4DzNYs;' + 'Z^;9V*99EtZUUy$d&rww_}}O9uKK46R$mMfzeFGXr*`^3Nk;B5jEr6s{nzQ`KjC^Ks~J7v`r(#)jJk=zBa;cgA`H;_onfne9?edv' + 'd-jbtUjNxw-C9y%v^!q;+IHJGG>+S3w0sM4Q{2;7fRuB{T^jo^up=T5f}zk@JluAb>1Z0F?KK|Fc3@7$FL*P5I#joAe4EbCN9B4P' + 'tt9$WtupDB$4D!FzilWTUix}MZ}iUx3s!e9IbgHp`WuWgz@B-EGnTI-!-NxH(lFN)y}Xem2|Yh^5s`Jxrd<$rT};jjT2dB%co@<<' + 'ny^5?oQ1vjzfFXmzx~a;=gXj?aKfETOiFz?><6x@Sn8KTU(>K#78pqtxED=m)nzbNbp~RWkaPmW8J@}9d;b<7j??Yhp_&~P!KZLQ87Fx0nr+Z{*$>%1K^G$%tNjDkurVP*{UKNK&1?!A?R;(-OP1jS4&01Gw' + '9tvs!XSmvBH%)UJ`qK!B3%@Y!z|`#@q{+;*Hr3$l#EI@N*<0)E9?bJTISg6>tV$-tRLNxgs8p7zip)mm);&+crma**sLk6{mA5Pe<' + '%>#<9Fg{h`v{nbr;V@d=!9Q-ZcNi6!lFT<0v_9P+(ai|eCdxUyKs?t;g03?&Ty>~YHOIA)s``A_+|4b698PK$dB<})oxX=&QG$Zc' + 'Wq(5x3U;7E2XX^uwPRX6nbgqfU-2h=vu{@GON_xdy-*U8uJufIhQaTOi8a}hnlECP?_bG2=NszZV=#Te7}k(}r1S@iyl!f|Rl$ds' + '8nhkF%moJJB5#WTbZcM2VE??y;o*^{Dbxp)ZD8fSD`is~G_5wr&Al_|Zwju}b!;@h+~(i4*H^Y>W<%%?5k$(D*9+)RRgr~u*=i|H' + '6E;PjlwFqJkp_f2N*UBk2pKXGERcGY9d|gGJS8' + '(~$#kHBy5L%bM0uCwr_gZ?8Lbw(-GmEP=Gr-tFh1Mx-T|nKId*(#5~50|u#!1bg>MfL>az=fvB`-MS|!^!WbuWeu8V&|t`~L)#@p' + 'Z8EPaLSqlF@o^N_qz|Vi3?N~f@c|7_%!#b#OM2tNT!coDe' + 'E(QQy1II4;;Y!4i1yJEO29Kh7;(F>le+nIpi#!}>I%LI??+SI0uh-Q>Mzt(WCjbre-30;)=jU`;-oxN`ctHP0u_l}CGx2*jdgB~0E-Q%s*lW=e%rQ98&2hMwK##>XWAcwKsMMD-lakHQ;%*kRqxprpz82y?$?e1w4JYX1LsVu9gb{p' + '6LA8=ec}AP#F^Z0`*g4kOIsA=B(aM`DhG%7D3iqD*A(d_tP(LvT;UbSUV1sRTvvE+8-F{UpF0Fw@pLxsIlN_xdsTz)J)!5?PfWH4' + 'MEDhnIADHm0>PFn@etl!Kq`Vk9PUDxlA!F}Wi)F6wcZfQgPo>O&NeP)?3TB1Jon))Nfu?i7)FCRc9y81!unm~T7BtPeQ+>xyC{I4gks1^Z%}!WNwlqNU7VJ-QU-m8!8`BCHB{YHp4CckA!efNN_w3nvuG@y*fkOtN' + 'ZMmU-YSdXcigkgfm1ui=s@DxBL+hYH!hOR&ny}+L0pg{2>5$F3@_LnTslhFKDRZ6Qo27sof;3ixElu>aX&i73W-RG5z5s6%xa*#0' + '!`nJPI`JinVjJiZ52rd+g?J}5DjoV6fZZW?ySk-+<^w&FG+>5)hXn&>=A%OqtL(;$ep^4n5qi#CuJwi15|h+=9qL|^H!niO_yK_KAC&la0573J#)A1PUAg#Z_=d4--fcSrkRY{!k!c*Py2Cb?k4!k-`vm?' + 'LeP44)AXUN$dZ1tFhlAUhPkbMF2_!YX5&?VIU~dOXEqw!Y3D{ND;jYrqnbI6Sj;l8HfUnJc(KW6$Y)D|+23s(af$={@MTqc)|%~!KadFf@)eJu89+@~}$AU^R?NM$J#c~P`$2U#r*Aey+g>$Y~5' + 'MNVyHeF)9gSdKlt*@_|T7QC}*G0?V<5HUH66=OF#4vS>rl+Su)KdKe0KV8!{Kszk3}3-u_l7)oAH_vUaXrWhHZu)C?>973?Xplf?qE*0xSP4d' + 'HubuzdAYAx`fpSDU?y9~SX_55zP>O#h%*n%TDkg_lv=7Ga5g_~JKfE%YMpX!k=4{cu7a' + 'Mi_?=!(oFoMz&l<_jG==FGFmln@)Eqrn^haCkA7R$SxX1p?=!M+9Dn#2RmHDhqfNYaOfA(LX$5;YCDAq>F{VHw?!43!`rLug}k?F' + 'nwIvVJWpigxp=Tc>@OAA(OuxMXB&3MujWn$^W32{7a89D7D@8itBmbL93QRHcWE4KP1+Sqg76>?4&z`J2QS_GN&g0D8%ev<<1zLO' + 'Sd-*v`xe=W&fibIzo@-)%&|GOi1H!dX>;zMSVFUC-zH({Ocu>?z|l}azZ?{t9tc6L}v#jh=EA%APQl0P3O{Pcjs!EGMNH6o?ycbrH(kj-V`1L7pfI&0IMn_80)*Dq@{j65>bM)JezeX|b6k|1_S;T(Wi%Q@' + '^GeN-C89_lF|KQ9F`{d|D^H+6=J+(z|S!' + 'KQuY2!`NC%6XT*k)vCmc^iR@!' + 'rOMsG>-i{Ou3D1c`bAS6%5-407@4' + 'yiEP}L&ve`kqz6OI;p?t8UDK`IN5ga?fn*hxF`7#(R?Y3`!RRjPOhydIm*7AqiYARn(Nd)^~ZlYkF?X_bJWfXdaPJ#zjDN;+A&k9' + '' + '7;M>MAO;H*Z$Xz?8QBb|%bL2NrjfU8@5jV(cNPz@8~l>O@SVPx;m>Gg1EcLQ(aKr0J^kR8RaJbIeB$)e;1`HhdfaY&Pno25%-zht' + '{snZ|L9&jozU$mMu)ghT`<-^>#O_oXP2w~l9VVEg%ZMU?{;%)uL}BbA$$5!}%>kDi(%h6!O=AJK#l|b+Vl2D)<=!75#Hos$-CKKl!Dp#AS6G0exQ=2%Dj57o0' + 'a%ZxF^-A+2ke_8!BJ(KzN|ViETa3OP{fpJVRCv?EE|9_XQnZ|O%jP!8GpDe?!Tgcp4r=$+MR0M4f5)9{?PGZ>w+A1M(Z^w' + '-9KsbmTpS0)v~5Bo5o#DC7ZJdT}-_l1`hwM8S`c+q#4>jw;brW1xM%k$@x+C=ELdt#}~&RPLEH2$c|6Xe!R#|k51m6+x{X!0A(wp' + 'FcT{%#Qy`!=Rr@cKJ;a)*DxgwF$X|eBzE*#bz$_2q2#%F6oNc)NL#kFn;b?mLABqZX&Go1B}V;lp-WfR#g' + 'sf(d9FXxWsX!?hodOJ(CVLegarYBn(0R3Qu_{;ke}}' + 'hj4hBt29SfR!7H)8RUX&zTcF&9AQ$Q!LuWy4ZaQ)@F@A*f^C)!xi_5h;GhDT`&&?<45C7XX0Pi^*|JN~q(IdMddH=LQ1(DSJ5Y|%' + 'G_dH`(1f6#xKaEC2vF0001H' + 'VRLkFY;AKdZEs{{Y;!MVb8TjCY-BPoaB^>SWod3-a%E;^a%FB~WnX7yZ*66Ca(OOlb8l{Cnp=-uMUsWz=V$oYl=6P*HyT@}H4Gpw' + 'EkO$*Hx$!I*haQNLM{ILd|%|LI%j5ZgS&R^ojWsP#flXX``=G~eE-L<_fP)v-TS8' + '{_vh#{`vEpfBp8)@9*Ee`svLtKbG%*e*Noj@4x>&^USN4uYdmO{@wfUe|q`q=U@Kd=7+nU{Q30B+xvGv|M=VeJJFJ3$t^n-fWI7&-csidCfinoqjR%UmNSCrCRiKn|8o-J*xF&`~j*kr0{' + 'wb6HkFzH^stiI;TII7R?3eQ*}-SPWcS_cJqVa|p4WDh?(%1>w|86F=U>G0IkW1TGuGaCAa;;*8NKbr?Gy96gHhM+i{sUXab4hB7N~b=^B6Bl=N0_g-L!NP&~CHR;ZR0GC^VcICoSO&A)&AIBpSzrM7' + 'G~Vv2#tdL_XMl9?fc2N3KEj#+DMU?cxUh@kIg8J2Cecfr{gw5x-M9Py^Kk$;J%2ADjL8a#E~bs' + 'e84qPz;kQkmL~9KK)@g~HVkcGnd6$#T*c+=>~}oiGT(7WN|+AA6aY4YYG5{CP_1qfn+u(o?|#w8Cf98BzR*wpibLkkcl{&gLjgALDKyFMdg;VE7#B3f}LkORmv;9THY' + 'K1|nU!d?^9w*-~D+r(mEDPf7V%W#E}E4#qaxyxP$5T8E+g2=@|8oUp-jROsjY8ruNj57keD-oc3gdK1n' + '4Bk2^J|~57lk2ULu@s-Kx_izqi?>vVG1ZdA^})}haFFiNOcgyimN9ZbQ9&!' + '^W=c}NE9XyVu8GGij)NEp?ct|ED&~P_JB)QnV6Ar7bMxP&=@;!&I5mnjN*APkC8%DEo?G_f_T0e>Jl4lbn&XPW++A=ZcH_cF2OTN' + 'KG})rn-L}QZ9yOaZxZFLOWWW%D5#O(%kBXR9EZS2JcEt<@hR$tgrId{un`CN@u#i3G@d<6T(Z*w(W`qNQP=Gk;0tg;xaRpo;mym8rc$X|8@&@d~yVlhc><}BEn+osiaPy&U' + '&Ip;W_H_TO-gTOUk-@}a3XG^o6y3cbdgkreGB4SgG%g{C+v' + '0tO=>eTZ>vB3AkZQ=dK0A*-@t03M&IB%%v4nR3uJ3F{wg6gYqy5K7F' + '2mnBY5X`a%c$qSXwc9kMfrs_K5drx^spNR97}bXL7(OSGu{=5{&V-^Wq+af0AChe>aY6NHNENbhjxtUtsQas>Vu8l%DT6ZEZ!Act' + 'kk|O(VCtI(OnFB9t}K#?;Oqb*0JHo8N8MVZM~s1hb&w5xE}b~XN#r@THOGQWAB#yc91#g1bBNRj3L{Hf{oobv' + 'XXRB7mD6vj8(pZ{rZgZ@OK5IzV2T&sf;vj;W{kGTpr+DCg>{zC#uiM;KP#5V#sgIoH334V2EgGg(2WMM+94Z8jFkt8#oiF}MU8dJ' + 'u4K`I1Mn8^i4{nrkRwFUO`IprpC5Un!n9`L-Z>8lnnV*tBC_&Qi}6c#q#n;=<4XPLm}z)G$3clfHL5uHBiaOH-%}(L0!aCdwC8_i' + 'jVmtL5kFI2fl-uEJC(g>Afk%|Gu#)s2hU*KVf`Nk=|y?3&~%{7SP9{s?2aJn%EOX&IaZ1k+RWAI)PoUc)h&x_^o#lB;L39W(yE7}' + 'Bw`7|PT^kd0)Wkt@tR`l&tx&=>y}SxYL@998#*iN' + 'mi|X*UZ?=x6q}(5p^yTr6MTflI(eQjf^Z1DV(q#2jkK_nfkOZ@R^<-&kt~i$#-&(U5<2b;0bH%oNZ0Hrmbkcu0P}7zz+e2}-H^xx' + 'L%;@dJw*j2QnF9QrZRCG*g6Yz#0^r74Nv?(Pu0FUqBZ)h%}svK!|>8Rz6!Cr4e&TX|M$2f%DgMM%D?' + '<)#TXGkl>paQcL2G=HbKlJ=`Og8S6E!ZAUg`TYbL=s?8}6$pp9lAv4OHA4}lKIeAz&NnujehGDP&(1ji@OzrHpVDOyQ1wv@Uh' + 'Tt0`lHK)4d{UgGa`V6bFC=rP$e#8dgO4^Q0QH{5LL=mMgysmID#UGWjWdW)>c!BJ@`VDWDkPUWOAT1&~xH4=;xKnVF_ronX6WxSs' + 'F@_1@3Js15%6(NIMXM4o5HpS&gv=G(s}cEofFp' + 'E+*xAJ4GeLNM2G@g3x}v@$8~{se*6^@2~75fJuD9STHAS3oAz4CBp+n2X!j!BaAi_$ULCMqEKJKLwc' + ')U_8tfSzF$;MyAe)Evw5I0`2sW}n(b>Zn%pN`!Fer>B_mg;Zk+g@iFrv!G$|ECWb~!I_X(3S`%_GoSzVoPy9%GLq@oXg+N_n<$(~' + 'PN*akTj+Lt3^tjupRvJbk;&}g!!e{I;-)PnD~BYj6`W>QF>;3!0h_&=ydS;zU(dg~A_@`(eVFZdVI%C3Kl%=i%Jza=iNfgger6OD5|scbkBF%DOR{a=D2IYhNV%K`!)C?p' + 'DKybF6ou1{A{LwFv966w+Lv}D4DkNx_Al!#96D_ULS9cJDP{YNq)N#DbSn-=prWch(HJe%qtJ*j_U|riVW05_*`yl7xEw2w+8#NA' + 'bZt3Un%!7<)!4R0SyBo8nPW`~?P=%I;+ofJg)g67nASESDj-<29A`vZ-o^o&ZS6Ex@g>1l5_E(uLzsO2@vSi7EGc3GCzuEga#c3l' + 'BNCFElzNb_tW8RaE@Q4Y!lYXgsGErp;`Ah9Dluso`&ysaxbJGdp*;zLI_gdC!A;f9Td{tqaTbh498@!#HQ~v`Knl%;=*_m-=40lu' + '9fOBh-})J<@aVQXG${0meNMNX@>XFC5oD$5x)I$ZRxVUqHpJ@5mk9sm2Jm8)XaiWyIU7{IY^~>Zqq_5wmsuyeRaLVQ0#b9xnyC6?' + 'uw*79+qM`29c-V<&_)J`pzW%|x)TV-ar0VV4>dB3uHxONb!5LLX&*a*2$MM4i1OYDLAU|_>m-S!v7}Ml' + '1n^JK#mrFcBf&Qjv=SsZM^-7B&nWED^}#Mpg`n%;gU!Yq' + 'u>Rt+tE3v-V7&;FO>3r8Ppe(Bb2+r~j%ku2&D)$B*qOXGRgJA;3zNjp1D6)DeRZOPUS5@m)iKf$7|>s`AI=~0NL`PO' + 'a@)vJ0inc99XVvyj7}c<0*QOhC2sKpAcNbk#)Nvh#R2d%o=@2`?4)()4k+p(`JuVHGcFs(s%)RXz=1P5lgDVzM9Ps=_&`HQwIxbA' + 't9@vPoJ^vDXF^SU_3;b%=U&>NUcg' + 'uPYr*Uk~#=a?5gLhdog87t)rf01GRq!uw+_X_9pg4eeqwN8)h~X~omPW%)SSffp91aL`!erqNODN+r2!>A@&Z$l6hTXD1yYz16EF' + 'hIFW^w%e8ONO>Zqgb;1f8sMgmvPxI|>pG)$<_jM=~=V}dDsE|%p3}RCiY}U(QyWAF`HtQTD6S-)J9Azv_N42ta46r(ta5n_->hyzxQ%y%A_j3yBGYS?PqgGZY' + 'lf&75kwSj6L!d`6t&L7gBhSiciBMh-l5AemqhekxQ!oxW#meBAf3uoEV=++DIkwyzn' + '!%1BM=MwQj1zOulog+xUdui>|i?D{Jv;{d)DL}YCN66@zfmQmVBW9&Q&IwV1B8U;#' + 'w9S$q%QstQPE>2OQPSC~171*Lr+u%>$TmKG0O>8^4C?NP&u_lC`B+MUzVx4+jKF2J%wdV;42Bb&nI3#JLOCq6@og88ifM^eKlB-w' + 'k4n1URnpqBjW(9F96iYvS>PuB=}L1%lX-KrU>)+22{{BnMdPhg1uE~l!HnlEUv656DdvSvP@Kwcg|frS7^m6KNk{Kni_28YA_&EH' + 'qt=nyHUP<-)K}~9qit2&k1!hvr#^f^g(s4q3(5SqYQ1~Eqoj1^4l%LD`Q5s-`@Pzku12v8GBvnT@6LM7USIRF7R=D_thoPEzdYGyL9KMSM5`!NFOm|' + 'HW~9@(s4(n8XI4K`kz;In4Vt#y-|EOmp=V}S3dhcS6J}xAAWKzA?)ut!%r59^L%9)js6s3!fX~>#rf#Z$9sHzbSQfa%Tphx' + 'Ev~LqmsN$hqwQ1cN`$0by^wC?re9^qk_Z_)ArWm&oo?5GAN^L%Jye^=aFY{bZY__^_V^IZ>N*#{_5y6@yAJ0{B|c&r>yp!9`uJJ`' + 'PK(*mZMS{dsq%w{_~!5QKc!PB$6Ki$8*fBnP#_y2wW=GFcCw}0IIo}c{rAO8bTO9KQH00008' + '0000X06|#{@A3!$0O1_~04V?f0A^uxbZ~5Kb1!XgWMyn~FJ*IWW^Zg{GB0p)Z**m8ZeMkDX>4;YaCxm*TW{Mq7Jm1y;Oa@*aWrjH' + 'ECv~5v1!ITEq2@#iCfITFce0jY&8<86{)mtkpI4iS5X(IO}7pZThzJns7q#L0hKLQ@i>6|jp2Q=YS?(vr#N8q_jaR5$F7Pz>}(Hv~5z8D%6baRYn2a}XiSIT7(J2!h#6aVAOL0?H&wP*K&afn8F{ghGbTW@hapXH?VY' + 'ElmYuCAZ`$phomNy;2r8F`&rWCMAw4FG|8SmDtzPI`iU;?wuCFnsjG~2mGQ+tX>Q3{>o~BL#02PD{y~8DB$Ig^E00-29KP4s=a)7UAxm^z&h>$QK=@u#+Qxf{j+FgE_$&|Sa' + 'E$~gL43<(Y>B*D*6n1z2pkzQru)=kgPHM&rl?VZ<9wc!gNX0{MTN5Z)QD8zOkALvKuVIAK_hQw9N0!10m$KH8%wpx^wKn7)3>FGy' + 'AH@~yz*;F`=$kzTViZ^QGqrs!UIH{C9|LBU&QV!#G09Ovj<94ZbdGye12ezd9c3TzIKGn@uic&{+{!6K<%S`|VM`(fyZR}CcKgr~' + '30(?F1Lwq038RH#6bnt#11Co%j!6Vj3_tPKTPNbL_u5PbNIV<' + 'r=Yek-*i@dJ?!$J>&M)tvm%Ezu3IqKD|!v0+VH*6`Cw6K2=rRiNr2L;YOX3~Byk3Hrk&{Q0rgksFla;PsQ1qlT4+b1DFmNzta16V' + '#hVqpSjuS%ppNM@Y8_lPWW`XA5uiB_dh{XdVp_jA=lQ*&cG5&MIA*UN0y|>rr$QC' + '1`&h`)y5~j&W_Lpn3^z>f@b9N+^L2rG)G2gTM-xthy%^9%Z6OK4v6(zT~+q`;^ruT#??!BUPzx#MfBR2zv;6j*^kY`>#khtpF&o#O`_YR01Rh?&Ne{hgdZ`x(yV5mHMoLn_k@=MCXPX1WGN`svYHw~)kcU2IUe=7}=vj0?oUIxlaG3ct^j;Af' + 'i+Jdl%bDjLJwG%~5}dyZWYW;z{^97f5uBd|heY%a%yWaEJSnj{jBrAE6pb8Lg_sri$oSRsW{vFBEaqu&A!?S`6O&67s87g!{~0uLw#XFrgCDJp!6_p5?v4kuiB!_RQ0O;TNon?@^h`mbvbN9_Q=ppcS95FohZ$uPL_nvAbuCM3l(+(1}za%>l*pL#)p*@4SQw;8B' + 'U_W5Bv<_aQ9l>{dlR_*mGk>^ofDW3Z53*S^Et0#&si5ql8v=~6)iq)IKE!9Gyt}G_(^ouGC%Lb-tO)7GE~0' + '@*XDsLK?=y&x2a_q(IMR}HvUjW_Z1eCQGJqs}xj$$J36AW' + '!7fb7he3mGyFS&VwkaFBCdOV&{RmN*?*Y8b3DJT0G#NiYggkh$_oXmOo+_kbr8*WTBX7$x90>s*9^HC>6>lqe;`lv6SK^!r#>2eY' + '_oycAA~%eBB1~o7It$_(jXb-?4YkD+$YzVDB-jvZzCt3w`}dMvel' + '_UJO0T-mn4mA13UJ`Y`Dst5X-gL4d0;D3MZvGYvI;+BMx<wS}m3HNkTDuU@#5=vjzgFkrO>+7Y7' + 'RIS9NRmMFoFwuRA~-aQK>7V2P)h>@6aWAK2mk;8ApmtT29CD^007Sh001ih003rTb98WQZF4VeZ)9a`b1!9c' + 'ZDwz5WHK*uWpHwDWOZX{Vrgt?ba^gtd3{w)Z`&{oz57=X1_sze>O*%6;GqN74LhtT+T4OdP;4gACQF7Sw{@`}KT5J4=cCObh{?x~' + '_eeTvntsO?RgJ3ZucX-R9X#Cq1Y7A=s2yyzf!oL5KeVi`;YFLjTd7}ck)~;qY>jS!v2Eu%gN%V_T5TL~rL^NtXl0XRuoeaD|$zYP;oFac1pmz62?(TSwm${DV8Kkh-dI3#{U87(-N=' + 'S>HsK!`0UdZTVyo0cFfzcvG=M$Bf!V7voY=0Xjmp^7' + '3_Epx0Z+Y`2iDP@kJ1^m1Ur!~d@qm41T{U|sNzuCj#R>T5qaiu))N;yafztFH9yREjyv#(*okt)(?E{w#0ms$v39fG4vE9q?Ac&X' + 'Q+P3MPEO2=V-agM*famI1fkrpdqI%#Zr9%O6wA}=;gPU#AH-O|UDw>xrV~e4g~M{_FN+%?$?W#1U>E}6Gn$?l8strU&5WddLLz!h' + 'vKJ`p=_zjD7$K}s#8XumI$yv>YncPR(Y#n^UgDJ>5*sdPw^x0Ox5j9brJ;V)SqGcA^F9$*J`71M$ndq2UG3jCp-z#;d5w4W&V|NB' + '(Wc>2yph(SOc~jJg6#V9=P%d!@xEc5odle(+k2R}6g1(EOZGnOqpQETVY7^3zb+Ry{iVrmg`$v|>w}B*d&+zT`oAe*ozC|&eM@6aWAK2mk;8Apk-}Qz-Z#004cH000&M003rTb98WQZF4VeZ)9a`b1!ssbS`jt#XW0t+{TgL`74l3-FabwYf@G#ysh@q*uCJXRi@HvVE~#t(>0*1R*Em#BVS}e9Tf}MIRF$B|-J-Z$E5x4lqAj3jmi$!g_ho&>tv5wqtgE8y' + 'M5m|FZOV0@nO9k|Ek(8IxJQ43y;bAK_07r2N&f2f(?9;lv;X?>r{^#J`{k?GfBEbG' + '%ue#xuV3e{et7xp@sH2`_wnmze|nz(<%gHA;Q!F)4B-7ik9|V_O`yZ~O}#CzmM01PtJ-CtMY45Q' + 'n;B!(>_pukc6KM8*?Uo5UHAD$tc#oB+fIlL0;XBNRWtc_^SG$?*M))gyGIMFPhYfG!pvmXYz~!Rg!K^kzd;Xf' + '3CqwQ_LVq?4q38TEG`1w%DU{!qRKmvAOS;lqF*MLO;bUm*X@BD6^Fh7`6<`=rRe%#+owg_Zr)r3<$*-j2n064#u)?0hI-sZ' + '*^6B__ppQMkg$OffUOxPx5?ck0|WyB+Hznfd(Jai@lSv4H6-8wZf|iw-w;iVlEPnR-IqQvHVU}8fg;sdMALy~1Mk;46Bk#;nZvAV' + 'XL5n6HgjK@-B>j-PrRmP>(bEl6R)YaS*!Nm<4`Ai19!UFFT9T7twZ$=zxBx=w|`9p`P*w|+XAYo{me8U0`ia+68>OnHeh`unwD}wX!3Pr=!=3WyV+??MZ+_5Nxr?)vbO-wvHz6FCtW4a?=Ptqq9o;34vDE~z){xSW' + '$zyOQ)Cn1dIAKG^MQ{Q*z9{M_cLhvjwW@(vKIcVEXU&k&79J?quvXpOr(^mss0Jspqf4VGJEEZ`GXBP' + '0w%}+rLEEKgzZe|0pREW>-SFB%1yH~s*&W|Y8hI+{cW21cKQ=EPpEh-MlO6l;n|Sx)J~2JeFjt1rtV?MAfCJGtI(88' + 'YiWCboE4j5_kNbKJLPuD6o*58K|sDhjr|3H`64v5vdUoXCX!P3P{Y(seOimFBBlUrYi{~5POk|=X`xyL4daQT<@d5=>pUb|!BK6y' + '`Yz)-{f>OcdZq=P8>{-2G=;Sj&(_?kj;fsE}3&YK9SP1}0WOU0BE+!DxY<_wy-7hoeZ5' + '?I9Ae+xIueNGzQ>pF~>;BV(V@UCF927}ArSZp~H_wwp>>|lqCb2yO=bA+pi+{;AM+q}~8C-caTJ(;IhqUNaOQU*Q$s2;e{BkmCjwjB1Nou>xf`Xu++xi^cgIYwE9oeL<7-cq)GL#_&4u>*WDS!2#-C3R`t!LF=1jUBRk1+pgR-Ut#AjI)FlQgACtZ9njof' + 'pvZ|%wC^xZRP@TzAdi~El2zp?&Jvpm_5zJ@*n0~QDD{px_9JacGtyToS$9j' + 'R-i?q%^qG@2)}s)0>-KhRTBGv;0V(t9IZ#4(>80-b*EQVb6Hf$_b>iz1Z$ybVp9~UL%yHn{k5P-0!k8UGhBo&!9axuK}XNxmm8%4' + '6V_#=R3{7)))#|upa|BaT;G@db%QGl*72deR8csfSS4CmzL(Gch&z!ec7~|fM~WIS)TY@jU{&A*Uti5P;#~=IT%BKJi68~KSs4T^' + 'Q_7Q0voCFMPOC%{b8f9uMhu)Kmu1ndh|RNPgYKmjV?$c&caPEt=)>Lt%&_bQz5`?cgKNB~Ddd^Y-3cs|LhAa>3QjA@R}yng3qzG5' + 'OgQ&8oelD^mjHd`pvUVQN&1!KQ2rHa55T6iZpBsE0Uza;hwWB)X3Z?`dycT9t5YK~%SM)Up-1bB^)>7Uq%P*E+XpzY!*T0=g}Dai' + 'Br6+%`G*wStHvOZru9ZteF4qq(H8zndXhN*K-?Y>WiX1rwDyMpn@L&`#t#@30}P)o=+7bcntJ51XZABhj)iG0wOTi{m(NGF' + 'q~h6iv%itrA}o~YQMMliao8I>UT(UtMc2PBcj7gg$*^bqRDLMyrJgp~cmPLo4MOA%dg){_gt8ngV=_P@B)SBpcu3_Gu&qegueVpO' + 'rDk^vW<H?O3JdN_+&VoPJXy;KTmTTxL{Fr4F' + '60;zb?d+p$@SFDb8$907`8)X-NBH)ryTMv_OAuY6Nlr>9FR_-RyR7sp!&fee$Y2Im+m&5MDY9TGOI#!C?}(I^zJ;Co_vM!|@y5eB' + '^=4^gA?HBpkUUl!5G(JvW(Q3Nz$OJyr_jh2*=TO%0jDO-wrG@bzRRW48=!pAr~(dC0euM%&l}Qrc_8' + '{4xsHZ0+I1JhMNk=$NyI1<=)7X&N-)6g2W6_LlSHYIiSKoiCZ;8tkC$0edGUy!BiWk+A%B6EnE>k%&ivR7PudzdI!u5W' + 'ZE;q_-zXbOJwQ|7=S(7)vlpye&Gxwa?~Ba_G$$Q=2E3>tp3n5ivkPlt?&?#(pZXXUD(im0Ht=ljBDgSJVGJLx#POb>p?4Ir3?`xb' + 'iraM3YKk!5=*h_Hb8oRODHA++mB~e3@gG|jILh3>NS#HM?^7TkC{LO-*sd*vYdemS$&m$tn;GC2AdI>xg{ByS8xD{' + '89J0@x14Bv+|7tNYO2`248S6zVLiT)QB>ZJ;1EPD)8I`IYPsO)>a9OZ)}}Qm2M%MB1wh`59=%<{ay_|wc$x;RBNr#VA}cuyca)X;' + 'bHD*3XuIo4IklK<8MRsIW}6dDt(vYI1XKSRwk=)xbFY?U^6ngi@(f!_8z%I#gls~MNak~rq~ph1*6AMA$}YFk24bX>cZ7QD`pNbH' + '#!TS3y6_kXXBTTso#M-b=IONSj~21qv$fQV5_X932WL@BuvAjjF#jQ)2Fq2q@zkUfe}AyAKe^ooO*N^Hv`_CneI0ZQgV1+M)iGHH' + 'JQO}M5zO+TPryXy6%2b@5tTtaUmdyr;F7!NyzmV8@?u&}$ln9?)*o(pk1MPFL-wqupaw*w+zoItrB)N#p' + '1p4``f!PduI#50wCXLuvi(Cj_Ez1Nm;zwwjLFXcJ5zSZO(QZ;gfiQMQZe5V(iw{|h81#5<(nS84F6QV$QX' + 'tVD4B4VY;4A;6I=qjI7BrI%}N$u$a_@A))-*0^ufzqZ7~`t9I6JIe+lGVmp6Z3yac%KZS!4Aa>S' + 'rFPs6h1Nans~z%Zz-cR~ha~j~si`*I+_?xL7Yd3iU-Ap}&cck^D^G;e#&t;qP##jmTXKfI5FdIxAWP+Jg6$TnUdSij@vx6{lSZYS' + 'PJlyRa=8PGM`XDLVihYxa>xq|)T4G+7yItIkrbc^nMd)+k<%i+g7N93R*;l8Vj%S&M^dP*)Q8zJ8LU$*k%x|Ut)GoD1h11KJm4ca' + 'LIj#pKeD79!oH2TAXJjRB*xv@fCHWA6jAgXa)Wv`@Mm|h%))1fjA1UvSt4Kr8A~z70ViI*y4s#~W;za>P9uVu{a2Fp)(temfbWhx' + '3KkuXXVn-E&gB1d8A7g;U&{SlU*e#rbddoF!-0^YaS>xkPO-7!pAY%@A{~Y2C+q~E8TP#xI56rCOJrGb%~&SV=ZaLiwp=VM%ba>L' + 'HA;ItLaSo3lzlZ?{wGG4^pn(^T(WTTm%l1*LILGGPc8=H-APxuD$9|Klq`6Sol+6_gfkD)H}z{CaE8ayKa4I^J(JYFL#}7s@qIi&' + 'c9kxCLr*!OylHRPca?l;{8)jCOYna~joEDW3P3nX4KfEBuU=h?l6}NIWlLewVa6`>Yhk5fEsPQQ*iWDaYqO}Y1m81}?9jJD&~;z0' + '){)9IIdr|XUb>zUCv{`z)4c7HW>@xd%G!8M_^K;w=xtF}' + 'n5W-ar&P4xM6tmOWimCqzQV!?McbF#LgBoontF-cItWLia$zo}IG0nVYXB~9@K)TYkfywBk58GKT)ik0dvdwRs^|Gh=#Vx%P3!Y`' + 'u4S4J;f|OF4yH3!!{pVWbfy}Jk6ZMoLh8t|G<+fPVTGa+s#8_>2W^GgsUj|MHQ)OxQfz9' + 'NRQD5MZWv%T;gfrH9kN(NLI?(eC=E!#TkxPYo7^b&_tQ5yy_HJe@RBrr??4cafPOzBHEKO+j4dNvjvY4UZ|al;4{zu6%D%Yt>;9O?~b@T_d~C|Sz~Gf3I;K?L(>?{Ev_$%`LsM9mlD5xhAf?F#sP>ny9zZAPWC#' + 's}3=b>>v)+r;h)xtK{!tEgW8s0bi0|#Rj>+`Xg!jIT-WMwy>jMu@k;9)LE%^RFaJ8(k+q~#riGm%CNZ6W0}j&+%b{m1^7S)J4%0J' + '&SdMRZ4Y~Lku;a&EWz4E0k$8hQE3|zbfTUicaTw_P0t~vR2!8=c5ZrIj4Xa61r629NH4D1=yaWhNZ^QT>I57%LelXG4Vs%W4vd<' + 'q_*s-K(`4x(*jwoQq5<2*T|?4BD*GaN-M()8w})hlqILmrprI+}l7L_k}p-qc#S=4@l&&Nx-53LEg#bojojTxRpElBaknWgbFju|Dj&!%o7?lKEMdJe1k86>Q~anH>VFH_iK;Vq{2j27m5QC0Sl1' + 'r|QizIdxvLLak2m!2=sIe~SS0c-tEzM!ejCyP3&MbEoSZ*ltTRrTpn{)DX33Ww6;-d`4quR6S=CKY' + '@eGpiI5Cjm1Y+2Woe>u4P`++hJ-pEGm~*evZ{P5<`&Qq7|47i4Vw?fhh$GmagA?Pph6&*wK95r9chh!E=Rbq6r~2xWdk=%' + '0QiW~jHV?+7e$?2@c|9*7Ef^VG4NlKwPgz3dkGCqyAiFU75Adu&5z9~Lb0%toSfk;wV}OQA-|qiK?PG%(2p6w=bWZ@zEp&}P>V(9' + 'T^EjFKDD};gneOue9jRT{)4Cs_-@Opn*ZIv$&%BvhY!buSay^$d>6@M%%Z$%kRTZS`-Ya2I%TQUP|lr9A0RNc=ArL~+d39s%=HJi' + '4EJcxkwXQpgAB(Zk^D#U43%f+s-)IDz9mud>$i{0M+VmLfbK&2ZWRKB-m;_gnS#;-^VBXYKQdvdHU2*5ng30po{ph|=|a$teOQ5f' + 'Q_Z~RksCULl|CmS%%(Lzi#ES-HXEy^#{hph&Da!%k8anWY%u1VFgF!+GTVdPe!B!u!vR?Q@e|mYJNYGuEzeypP`TTQ' + 'O2A7~RU4u$iyi(mN|F^T=E>I0jDTu3OQPV7dSCj8XkNz3UA+dFdeVRn)*G' + 'LUV1;BD=QUa!q3S2fz-SqEraO&a0n$#X=^O7x9Np?^(d)SSb>Dixaau=|i)s9*Mk0*vlI0qxbM{tuvGLF&G3#Jh)`T-2GIHIR*re' + '0JlyVE1|C4Q`IQ%zCD(U%AsFduDikIVYKVK_U5CuN3I&X)}-5?$6GpOtbz>@JE~0M9x?1=r9S5A_3=&>?aeZ*E@*iI=cRGlhIh2U' + '8Vx?@6NewZ{BwtqxU}n_y>(s4@%K2lnbLSSiXQOIYdpV$Myg?nyh5@@i}e_-DA&-1zii7+JpP;A=*3kRB!}r`Vh#g%5rRlxfeFC)' + '5ocC|QMW?O^e|fY!s7+OcOK(WTnO;d0sr17c~~~EWq&GWGy}?' + 'C)zEhD3?HQg1(2|sX90aFIX`H^CcXHys&O79baJ7i?!JHV1^fXg>(X&ZROXERo$ujl!r^uv_%0vCDHz9iTQ)%?Do#=Md&AADOHj#E{pYB`p8X`5Oa6TH2D|R' + '8&59}rjm4}f-+36' + 'K%6v07GK#?j(|' + 'c8G6Kj>rJ=T=XSViUMJ?ipG}{$}^zlqr`s9<`_-OXoZmvM2Cr1jru)disrQ-2wve;~mm07HxUQ5k2So1~w' + '9cshfRV)5w7k2r7P)h>@6aWAK2mk;8ApkyXf&%d~008p(000yK003rTb98WQZF4VjWoKz~baHtvaCzl@YjfkqvFLYK`9Huor%Gf@' + 'aUXVE=Q4&*Xo>QwtXD}~KO*a(AQF_2B|!!NcUMvL-*5N4dma#}m6P0ioXQFW=Fv0L)6?&X<9K|x%BN8_ndEgHHOr{E&7*^?ZqAO5' + '>u9l@=JO~kr_rifzR%07oaE6WD~nlPH~r5)`~0&p6tD7|qHe0YsHmf&Z1S=xmSr}d!;5m7uX6aOY@+2X8ZWEKEo*LEWo12ER*Sr<' + '`|y4_FUp%}l9iu-1^`q!e9EDAy~-y|RIgX7Wz_(fliPgqZncC?>RvQ2CV5%s@PsWw9jJ%!)ys7SqlrEg&FymCeEwNc7EO`Oi@#$Z' + 'nHW^H1~3faLyo+YVJg;GNaUA0iXVr2MrRi+ltgAdtqhf*3K*@61WH{6sONqBtt{$7^Vm>dfMFW%Nd=7)b;{EJ;' + 'B3n6Z@+xbVRWCZuRx5x&6rW~IHkoHNj87I-uSAjAx|}pHbE?3+TBnl?=13IauK|Qe0Vn206j)`=t?~H`e)gg>80FcrErZk{caHdX_b+m4)nnR+>3tFY#Ap6|bTB#}o$QTL7~pt#' + 'c;a6mft-`PC>m}ZTjw){<2-|AQCYv0SSVzIEY+x_vcXBPI|L#<`je8n2LX3wv_i66{b' + 'XXm56!wUe`t~t$T*?Qh!&1~H)83-%c612j8{}*lS`-4~Cq@!O(dvC{{7HcH3lN8n}TjaPTvQ6VvihUtXVGWx*S`!7k-9^EnTp`8W' + 'ecgYFzb2F2>$}}@vRh}nU-MT>Jdh~BbZ>*hxw' + 'hi0?4xPjm#EAuJ*nY`oGz6WF;eat6ogoX@nD|CgH_VFw|fiGq&z72R~;LPZk@%hk36?b_l){ZX5>Dbu`lEOU8np?!>r}N>lr_1{%' + 'BLNxE4Eiq?P!HeAH5m=ba?8_j)1Y|+}t#MpI61KfCuDE' + 'X~V;}CpQG%gwU;Vjdf9q@p+qA5guzAh7J_QUrl52Q=dgRdR4|#EO+thp$mB6hxDab~!kTpWa' + '*6m<;c%F_fjzuePPA|sLNIaeQXV+P=T3-XIa`^+cVT@eO>FGFy^Xq8;vGuJ1R~8u}OKnfrnlE^lS7kn@U)L+xdwI5?' + 'Z)e$}m={((Oh{fK$Cu{vl%&8``Gl=eZGpN1TDO4a0bc`U' + '-or7PHiZS0Rt==SmA+kHKlVN?;SdgfoEPOg@(~aEJZ+>Rg95WG#WqZ};>4O^7+sCyY>M~v({wpmFZdU0(ok$+Blqt#+!xjtb&%kh' + '4n!^GBr4^&;`SgYnE5>g{&V>ur&|yim~@e?0=UW46ciP*$-C=inKQJVTl)pM=%Xi*0!EDP=(+#|aBw#H&9b6jZWQ~-(7s~a3u&xI' + 'g1uV$-g+ns`{I%g_Ih2SvSt2W9NM&m;yawqj4z%J)z_9L+dGCa&ux-S>)*g3p3`n6tCw=^5ykYZT$X?-^OWz0`Zjz0$8Q2?;i>OM' + 'B;eZY^QF#k&87vep3P%O+}Z%)3_{5f0ij;o^qX>PPr*xE*t%$J4D_z5bDVkR5RuxZFf~w#D)sDu2E~~49mp|{9)PgyK-lv1I={`{' + '7t4xH1pbe0<>*|fct=&WtU5ns^L5T1y8~W6j^p#JsPk#`;WjUs5}B)6rLvgoh9wz_nAQo-6{#l1;XnsK_4-2V`uS(md={msmqE51' + 'HpHBfN=D^6?-7DQ1R(Vy+{708mz{22rO078h91MfQ>XN?vs@UtaasD>zv@T~zb6' + 'sUb;|6b$}y+L_#DNIG+v>~Ev5USmP^iZmZkkJ1)rc6B5SO*R=7zIE!$N#x6FuPaPw*DcpC' + 'K$5=(V!Lkg4*f3HbAuAo3b2iQiGTYo#xc5(0~p@HtIzG5bOS' + 'T4N~c*%4xi2fAE^#KxjZv-n<+f(N8Fe3?1*7)1AM@LyC9av>{D4H!V2-uzh^SUQ~2C*x|J_o8~fY-;x0mE`aMPS$-U8EEn{`iRCV' + 'q$o>f1B(hYdlk(cn*gLbTtC3fON5jll|tghe`a%FT#Dv)VZ7r~h!Q2h(wML^4' + 'u>L^3r$YJxg2Wn=)Fq+Rg*KqT??tLMZB8~b2sM*`oaCz}(x*(6)lrTQ7O*}nh;g8z>pEJ%Nf@o6(9X9wvhR{_6@1?fzU8I9+d-1b' + 'Y*DNeIm7s0==|Z+6z?gagRLr?y|`;OkPECx=XrSpCrSdl2R6?xobj*XLG<;@o;vRb(KlZ|^lyN@P%s*k0+h{e3y}N-@M?oenZ<8Q' + 'ZAG2w88Wf$UOgMzM05>3c6beT_}GxZ5(y(+f3}{_7g;m8Rpqw<^RNU=E2ai44-!(`MfW(t0386!-&&B0O%3?ZRw&`iq}oax7rcXqcxxJ?aqEJy&6?7jmS`wltQlfPZ`LHYTaT' + '&uK9=D2wDsjaN);7(~q)ShY({HhNLN-@lSv3l%bCX~vIXF;uj*O)AzxwqMVxFe39Nbgx(LRNLm}&VAA3i#{i#_{fc#_y8vkTLcE=' + 'LwD;nIA!t@vkmC!f0Na>yMmU(VZ)eMzNe&aIl8`!vZ$1dBpMkbE_#v0Sk7}}+O1aT>t?lv`Go1T@CijKrsIlG0c1`z@BkTtGcdM|m;|oGrXRKh|E*9%+C=Egg1kVQXarHxENh~s)XN7Di)!di6gTwq#%rdXS44{0?B(n`N_w(3mWtklT%te#' + '#7EWzi%`}5AR_i|F&?F!4|Jmx^@B({0tmA0=;{sN4W+wD_^Gy?Vy>Y{H?hyuqnk&A(xOYar3-`!JW;9#1}K%6DKZnt@Y9!c&$C~yp3*sOIl0agoy' + 'Y6&7XQzK8nR7sNzUe)vUP2IDb`*4fQG}9qV?J=wh5Wt$f1mZi*tKC^P!9L+sUd^*R);c=}?bbjB5Qd+exiQyypB9mmbQu#vAnp#I' + 'b7;R_Ut05ay8i|%;wT-A$KOO!Pbfq)zs>~EFN' + 'bHo~Xs@vFZl07z5v0($`(A66-FgqF$rnywVq%f&lz_roeaWRULK+$~~lS5g*J!v%2GCwz@ald!L`8ceOK4gO&j+B}V&PskjZKJJ$' + 'P*M@&_(R``{nhBqV|^V&UXMQ8YfJQN+Aif$W(erck!G1R2bLt}JBut!r#gjjJ_HzkgJ!0LipkZa1jpeX*KA|LF_{PoCs=#R8$KJ4AW5>dV3vaO35d(Y(iIv4@D{g>4+~;9kEX<3)!l)|cFiChuW^zw4' + 'UMfuu9tXiSn3X>m&(!IVsRfL$hOA5mL6u5ZMh)p*XjL#yD6=$oJM@)hGh~#c-iTHWQ$dKN)m_gAlUZ}TCavZ?ikuBvHGDK%^-7`&' + 'LwAe!V`xQJ1ooK1pzB~cUWgpqn@u$M7hJC3oyi@Uo;HpqTrqX~bwkpGQzi^^Y8WYJ5&@WC(BbWc&Ny6N_!j&kZ%U*wB!*9DHdX`x$5Y3dFfQ@^{RHW=DH1N>pY95`&}5@T6t<5Sf7;nrimW`6~EI=XCk6MvF(' + 'fuKhE0zeuQ5AJy<4{_JOXiQtWg;!wz9{PFU' + 'K4I56dT1=m-F&fX?yz+jN#Ak~6dst)9OUJCu^Bd~`S31(JH7;F&0PsL?P|vLMDiNq5SIP4&z9x68Js($x*%fS3pAAo`S2&0x=9W~' + '^;COYnrCbK{4(#w)MNK-!MF|Z%r|5H5GZ#J14#-DH#hFh?K$7Kn_mw+2Bu0k#ArunL--bL3E%Ai3RbX(o<7UMe*C-6u?iz=onf7~' + ';rEhO=r7^Sdsm@X6u^S>RPI(nZK@5uRib+pYnzP;FgGDYkXGX7<4*=VPqxr4>CG' + 'i@RL$yk!5p)$M^~x$z4B6BqUWtEGIhG$OPdn{7}2*+ZTc*MOe)9o7SIWO$NdQNzW=' + '0N0Lyq}@MFC!lr4Z`RkAmzmZPh%O$~%5zoUjHi6^vZHib9cI*MKnymJ&FjgGj{$&55{>a1&TO)v2Ft_rsO8Y5f}Rn=r3g4aOoCWW' + 'nc0>p@pQAD`1o+lQhpUpY?=r=@sMOIGjf;FVtaeL-~YOMvbTFNlxfWD%x1L0S+$(OW-XiD(W~9zO_eikwoZRWOTSh6Gq*&S8gtxD' + 'YvMPbiLRyg^UNpAp_(mP_`1rocM5%cRXG' + '_#^RDW}Y+JPemySV4jVe)aLnES_ue>R2(@ePTKygOq^@?sC<0>tu=t9ze9Tqo{o+$Es}G`opV{v?|kN}m|0Ye;>ivh)J}`dN*a9C' + 'T=c}zNB9&u5B3~<%5LHKw$;El8J`mJfrFT)0!d_+*be>R~Nox@;h3gCq>f^e`xIl_$!NRMo' + 'u0JfRsatp(R?V4Mb-I<2)Nw>_#^bY#=ek{S46)6t^p;`YQijou^Zde6' + 'nK)BiHF8tDD?^Wj4P2_AQX+ea9JS)|9&XGjqq9bjEyAiUKFk+;dK%ckj)$-J?~pQ7U;{TfVbDD?Yq-@q3>k(-Q3<|cB-g+$3D' + 'tE`Il-ZWg}fWT`B+LvjmOKjw{rUn*@n1*nd7`EVhFAB_^7+;rA80EYbq%+1GiHus-1j!b*r#ZaA%_t8m`a94aGPcJFkiZ7iRw+39' + 'u$r8qw-au5-64DeaLVIH999@*&qL!xQkAYhJ%PUNq31BABj)`5vvtoZnaQlx_8KoB3^e11pp6H{6v;j&31&T' + 'd}|c)Nk=FxTmx$sDL8S@nyx%vepH-THp%Un+$y7tiTG=J?TS5{sJ(4?rr}H{?BAZx#fL9NV!!AqZHaZP5O?ft>~Midl!3Bn5v`(~' + 'gRf0R40VNqc#Azl+9y09s(aCD&w$Z8TDn9tk0{e;o_R)T#Anw8^h_FtsTY7kAe<(I+%I1Mgn9Eq&|Qja2(Dr(pLvz0=#JanD2dEH' + '0l?H=@@0s-p8Pz9<`3QWQL(`j^NIAbIxU~b#%BjjHj>br!n~t`+Y)L)j-f$dkI(~BXcDasx({SfH~{;=7qch_M>?bQh1C|yx0c+;' + 'Tc;f|!z@JdlY9Kb6bW5wmkmKLmLvqdo6kzyF;m%YdrGHM6H1%EAPv)oremu`l;G0U' + '1o0kq^wr#6OHRT<%c%%l3vpnKCM(AX-8WzRg+y88DY?JMP%hpKcM)VW`xhM&>;~l+!VXwQohj7DX@XySEQ(Os)t#?J$xX7X?QfIl' + 'cFpVcX2E%pNhC6~jhNbr`?a(-tJvS*5wGNOxilzxGYc#U4rsBrdIPM3' + 'e6W>qT64izbTdo4k6&hdBHO*Vq`;Ztdfdk{g8}ky@^XiVCtFx7;Qh?nj_h-XNlH7S+puDLz<_sB;2~4`p!smNhtOi^9g~S=(x(At' + ';Sy_IZNO%uRI?em3r$9!adjr!Z6ur|Bu^y>4d_e=xT3yFu*moVnL6`RMc=j|VERy^u#{(#98r)4N-mwrK+K2q1#iE8aWD52h}}DX' + 't(WDFV!e~Z#2gmevSVCHvhxu8=0)-_vd?E+L!1ECys{y>EIXnv$1xb})W|_6wk;(MDxXLqYHCyl30n+QeUqqt<&rZe91C~+AX;8y' + '{05^*dPn)fatDu#-jYkehTF+SJT_C-5R!0t=|FmQ1rl^qnG{@w6So6x&;vX98WUj6y9k^+;ReRb@-@Ak~;XlwH88$' + 'h>E~WUodDZC@Q>iv(EPw&$eyXhD9f&%rQ{QKaG!BqB!1wyu9hToDr3V!-L~UjFvCTc2z8a|!45?A>v{vPWL5*lz' + 'LZWm;eT#(Il~B-h>R3>s{JE`qKu?2g#UqE}1spf%(B>nCVt<9!$fZ4g^mE!Xm7{VGU(#)G^~v0v)2r(Q73#Y^7F+6->~-psguy3Jv+Qz5`jcNI%UE(9wXyyV(p!P%b-}-#C~Cw=W~+a9h@C+' + 'MunazVmdyWF)mFvP#@`NVb+>01LHgzNfZb2zSQQZeQcksScoU+%~WFqnd$dhL@fozAs$w*pN-F(>iq`gmdtk#c(2$M<#7s#}n@!q((=38`S-5I_((IyBirJgIpIgD4LNqXi*%zcyR;2*4KF3?u%F*ZCF_OH1XT6w&lnWQp5Rz``!LfG#E}Znd7Nv' + 'I^jCE7PRu1*0r^sL0h*Qn5~)qyLN->wp;%^a?^s}JwW!)ZwwY^48OqKkFIpQs;=Uh&L}^5&|0-7q;-G#9Vd6$Jz;OxV)yG*N|n?p1j{)4JcR' + '{+O8SUTBkek2yDXGSCoreltlzx=QBgZFabNX92~gDBJ3^|=DcDhVVkxQi90~A|B(j!hYbK2yY-Th5bVH9EPXF6t`rVc7' + 'YpK)Q`oXkJy(-_KA+Q2?eU>n~6OvXH2s&S)Wit#4b2`7JN%cQfyw%=HERCbQlGDuQXzo)_z-55v^IiJxIg5?RsGbv))^j7$l5ng#' + 'F|AuW#DyktU)gmyu!;EFdxBvmJs2Jz9{swB*x$}ps*UEIMA?$bq0;Nw;Cs?MK5Oxp%pC^' + 'Jl9O8*~s|c' + 'SBd}DR@H|gr8TY?<7VhkPw0C%c' + ';BCKn)&(|q;dKmLNaGJvdi>mp1MS8Bv&KKsl4^!wbCQ_we7)TA;k#eQJPEwl?)VRd(CuZSsWQ2n&3gcQ){fQISNNzAy{Grs=$G;N@WMQ8)g^p>av?9~Yo!#jU)7ek@1M?x$LSbzvjq0CzIZR?-k{Ws*N=zBJ%OK!dIp?c6Ro(7@ZRmZ4R6Tk4O1q)&ah|AJF#=PIY*PWG^}EI+y!y7YbN1Dqrp~aF2AHDO=cJn^V-H%' + '6eA@=rv`HrHw#jSaq?wb4KW3kXCqx|NmBWGQ`0^305)nek&xU~Nuh9!Lv3NHW=oc}ia^V3w6^7A8%xo)8bLT{vzfk}T8`XgiA@NR' + 'A@(Tuyfz@|b#7EAmI7fKaV!9>T*0=$lD+i!C85N3$>Ih!sc*9e?DEiF=w!B4CL44h9HKs%a&9yk`UJ;L^Q>$)Ud1ifsm4-{v$Nhl' + '4g8j#6x1qy4t$S&y28^' + 'w4K#pw}8Qsgnu-ZH=y=t0L3im`i@3o_;!4n?j2!F8$#E6tY>)(nML$8Y!=BgR!~Dhh%n|ES12%j+@rPf_}fL7}v!uC4;`#A(ywlU6-PqMXEn>Y9(H3^bQ@fD8!lA18r);YVo6A4fUux_<(2I$r_mO%0A$MDtZx{+UhQWrh3pGU=_5cias7$fbkMi-a1' + '+K7E&C!TyRN)K?%*>+-vfOv1ra9N$wGIRN%3%zPi_G*>&5P@sOm<}QvD}o7QeR#-=<6O' + '-0X&TB5Kp5YZuqXN$X6kCbz(zOqz9-`)$}}iLfDC*>to8-nL0whYbB`o9hjOM`HePT*+Y}685iS!gn+xiRWFgme7AB&Kr9IrAZ@E' + 'PtW1L?phEM`qQ2;rMyh^qsQ;wqwP!r?`_w9sZUJQ(=(V<%~YbM<=3J>bk8N$Z#R&^xJWt)1!Eay!gx*8Sqb@NqqkDrCHm2Z%+$*y' + 'EGXK3Fu#wW2^!5~5~i+|#{0EEk4mF$Iv1~chkUNt9MNuMjd`_r{eaCYEc*Dm=XVKsxBUC=uLi%&WkWlFUodZeE-vTFoER%hO-AZ`An`aFskxf1=mP9ICc?uX7' + ';vpM!<4F|sh5}t+&UXInoA>(f8qjO1TNQ*`rJ4O-88s81la+t|WN~i$QuLN^nc)*&TM=1JS-b}jlCwel`19E)+_F|m=czCkEp-+?R>K(#@HmBC*' + 'Y=(C)rPujw_P$tFp*HMtAyjhOv&{!|)TajrM~5e)^!w49;ZKLB=NHeBgNXuA5GM&Bb#{q%977$rfo7OR5>wd#!&@SSFDCEidCU|o' + 't!Xvi&Esc`(!)FHVN!Rn-5o~gG>plP3t(#;9_L<>h8RbGFVY9672Bh(M`@UIn!){k3dbXxgawk1Hn*jd)nQDZEah?A5X`aUpOE)9' + 'Nr9_`k|lqK!~>!Dj<)28p~Q$5aZIe=ws@;iLO(tpFXFzk>e%}yv0ns5PP#sh;K}6GXBIfw*MDk~E=$_vgL0%y{(9hr)C}K%hEj!l' + 'O2T}`qy!w0e`i4w9p8^y1W1_!?->P1+2qsMQEB197&6h2D>y2sQ47zJVLS5qcrQw|9-5hUu25Iso*#T9&_Z&1X6z32Ps>aZTl37QCc8a#qjklbw2I1!t*oL%S9jJEFt9NQH_)o{Lswj0d|w#fsi{TN8+K' + 'uWn``)gQK$i98)hG9ub%d9Y>s{m0=mFO6Y5o0v|aa`?i!FlIY=2MQ|3s|<5r8V%EniEK44x#>|lo5_!H*sJVAYP29+xKy<($2CH4' + 'Zya`BFJ{a6^mXi!aLuOu)PIxJRKNjAWw+FcXg=s~ba8xuojaZrfs=s(bRUb0luX@F%wy^wd>h=FY?n;9gS(&UOvCJst=I^@XIF^H' + '2z(PZ>0kBWlOHt=+RU#Y4rj@xd7cyhT_izgXFCgV63X%{T9BuWH%1Vc))YzJO^bSx6bg}G+LA3rUL5{qlE7wd_q#vZK0g6N$5`DP}OOBKXtt8>Js($&6Tb0jmjhG782;Z~tviMZ0BrB<^6#J4n8FWZOGbXDBmQq})oal)A>?{AC^)' + '8?Lz;&fp0?c>Z2nT=7t+?5`wr%-&kA(mB+I;~w@LK6!>d)V~s1Mv^LQevcs<_Cpy-V&?5zDk8os)(6CSG^97f%)p$pR@?W`(vbo|meE7-v(g&HlDnh0+-e!rY#*VcP-9UZcejN5o}epP)X9j0!~46jXoDTP1w+)Tsgb^a9uD&+lZ' + 'q{3_W#PrWDbp?JbXRW8(ZGAXBO+3+&)9-pLP_Dq#PmDJu^A6+5EpdsxLR#J)LcoeW=}H9Nrx95SI|c6nH|i0o$lK^;%0kItZs6lgNMMMiu0t{zwWS_&2hn^12;5#5P%JWK4%?+`!p' + 'EKxE4jL4YZJ38ifjga}hqhy{KDI*lzQaDf!$-bo2(?kikz9TU4+F5YQ68y)tlQ3Q;kq;g#kCsfxFRP?S;!8a5?EX#hKOw&YkP1OR' + 'v3@(A*@XR1tlzfH%FFpB0hT}Hps7y$#K%Z{f`=+QViSHu`!z($pR(QSvLy({>q2hoB@_3RRmVOw66XjG+1xn#`>%NCyR@@yf;8e?' + 'o~r9+=)GT_8)vq%x14OxzJi`bg1|RzRq6BG8ibsqdI<{k1H7>3aKkNDlz4(x{uddd9zL~^6Dn9gGnjThaRV6RpgVrf4$24d^>u}' + '!VDn2<6W`?CfIGw8w@w5W-{D{;c31sXT?ppB$GZL`|O#52bcSi0xMr%bywF@WLF;=jcsHTiO~gm-e$-!Y4e;Ve7l3M#$vWw3rfreh9{$PJksY@lXiO$h1Tw(|r-Tus+T7lan2oDX2ZBFjVQDe5#b+mf;=' + '2}w;Z>!R-o2N(RG2f8bFZqn6!UE}{_XGDI4*%)bph@jpRZI{oV9%XnbpN_=rw' + 'Hn(?G_Cn65{-x`U5n*bL(o4tPd`i~n;8JUhfU3aT#^w-+}Smwa+DG;*-d@9&2mAmu?LJ(cn0Iq-MaI%7z)^ewyzUhD=Y@zxn^*ClME8(f*' + '9$>CNkU@|Vmsme24>lWQ^FA#3a;sZuwWdtiNd0%wvV{1>#|Zdw3`WzgPPH%8SQ&+0$c2uGKMTFUmjB=qr^rFYHB`dcO=_juW^~}0' + 'Qro;8)?k@^_w$d9dxJLq7tmjNJL}5}N&A7zv%@vTb)}(-T`#vJ!Y#~!kW^#j!=pD#$=g=N#6(7lv1{~aDgL$w_G(fsQ5(fpg!45_' + 'j@4-}*5#UAR4?UAUlK@m%e6}LYi8QBbS9jI(pjxdzchsgoNE-^mh?938_C11Y16Z3!@A|w&Qey6WO`#@@?<);GM7fKOw1OCizV!%' + 'y*{{CJ7?4PdDRqHjNPS5sT56GOQ}!GJkf$~=36mwO4e_=&yp$;tt((9CI|v1GYSaJfyM`WkR~2D3Wu&~OaVt7k>4`y_u0#$9aciCY_K9|%1L6!S#6^T;gPTr+mY>_~' + 'J?^~7lNA|_>dW#GGmtQi@$&^wAF*Xy4xhj&u70>F(?};9*#^CHB}8jtETZwEMn(auu@ddjB9t+)sv8OY=@%+YsUJ_&CIkms$e2G0' + '*1dNmk7}?)tw%Bw`AZSm_Bn^?(d1DvH=xJ!g6W?Er461hvuP>=8B&Zk6z-MyG+^O}HB+)Og`*sh7%dHcRpAf`9CC~KytwXT>qdKh' + ';18ggPP}eryZ;qS3NAuiJ8#PZKku_bY|eVBrSa+Uk@2DH3xN%}{h>V&&t2x#C}vOs3xu3CI8_jH02`qL<9Aa-Ir_mmtX' + '2|K6j%$Tj_l;0T#kI_FOw9UYu6CH05r1EhbA3U1gAtWsERJ`j_?D5t+lldCDc1`@>pXb#&o#0(+(7{I*z4@xk*8vdUUWT@bKIbr?o<&ae6y|I4F1x{`' + '8N7HeVtv7i%rHkfs)u>@HU!O_YcV1pE*PkyK}69;#lA62(XgV$)qzuUB+P;~lR+=ncU)=BTygZmVyX`f9pxSbtiJiuz8i#J_%RP>' + '!_Y@&65*-Du%_c{ylbB+vRPNS7rdAX=90=U;(6YWE>z!#`zRxkb|AzsRB3DgVfpP*jLD=E|qckw-+YIAU&eS_Zs3;!=r' + 'O9KQH000080000X0E@x_u22I20PP6?022TJ0A^uxbZ~5Kb1!mrZggpFWiD`ejaE%>+c*%t>sJhPvVqixILhIy8AuNrpH()zWa*I$K-O6=m1xcf4MRC8~8B0>H;e&AW@a3z6' + '#-Iaj2Gs%ksxELIiBkQak&&4#H!9v_uiF^W#_Mgww-!lb);i6SY+VGawFFpbW~nJ`1I=*8`*8kAME&7Pj+=ep' + '++pUqqW0!`&LNG0|KOX^$aa45d5&t=U<;!Lo1}FqA=?5l?v1AiwDLO?x' + 'gpFdm;GYD@ssXEJ<;y$PiKL(I(98j<7cy{AD4w`kmmko_Ohsf;iQhyItb' + 'hP)Z;(AgOf_@)@8C*w67+7y+;*!3QF6h5YKHu@mLv%!ZtT#6l*D37~Tv)moCa~3fTj@?@xs2wwy_bFW#$U9cjOz3jaPEFNCIu~ko0*ha*s$wJlkFtQr' + 'a`XiHcYJt#`km2}#!vkJP)h*<6aW+e000O8001EX6wBD9oB{vbYEXC' + 'aCuNm0Rj{N6aWAK2mk;8Apm0DKKv;H004Xk0015U0000000000005+c&H?}cW?^%5aBOXJFJW|aWo~q7Z*DJNUukY>bYEXCaCuNm' + '0Rj{N6aWAK2mk;8Apj;jrV>XUu$J@Ze?s=Y;R+0E^v8JO928D02BZK' + '00;m803iU+&?CPA3;+NsC;$K;00000000000001_fv7D20A^uxbZ~5Kb1z|ZbY*UIX>V>XX>N37a%N#;WpgfYc~DCM0u%rg00008' + '0000X02=`jzW4zE0O17y03QGV00000000000HlHEIRF4=VRLkFY;AKdWo=?*WMpY>XLB!KUukY>bYEXCaCuNm0Rj{N6aWAK2mk;8' + 'ApmR!j0lJe005IF0015U0000000000005+c96SI3W?^%5aBOXJFJ*0FWn^S&ZfA2ZaBpvHX>MmOaCuNm0Rj{N6aWAK2mk;8Apm?p' + 'nJ1q)000^Y0RSBU0000000000005+c;70%eW?^%5aBOXJFJ*0FWn^S&ZfA2Za&>NQWpXZXc~DCM0u%rg000080000X09lsASC=;c' + '0D1!f03HAU00000000000HlGsfdBwzVRLkFY;AKdWo=?*WMpY>XLB!ebZ>HDXJsyMc~DCM0u%rg000080000X09oG$)G-GD06!T3' + '02}}S00000000000HlG9xc~rWVRLkFY;AKdWo=?*WMpY>XLB!fd2nTOE^v8JO928D02BZK00;m803iVNhLH&AQUCzdLjeF300000' + '000000001_f%w1x0A^uxbZ~5Kb1!XgWMyn~E_82gY*0%90u%rg000080000X0G4vL{O|w(08jw{02=@R00000000000HlEv6afHc' + 'VRLkFY;AKdZEs{{Y;!MPUukY>bYEXCaCuNm0Rj{N6aWAK2mk;8Apq5;RC;{`001uy0015U0000000000005+cJQV=|W?^%5aBOXJ' + 'FKusRWo&aVUuAP`Uvh7BVRCscaCuNm0Rj{N6aWAK2mk;8ApjuyFO&ok001&h001Wd0000000000005+c^%(&GW?^%5aBOXJFKusR' + 'Wo&aVV{Bn_b7^L2V_|e@Z*E_3a&KZ~E^v8JO928D02BZK00;m803iSX00002000000000e00000000000001_fj%n%0A^uxbZ~5K' + 'b1!XgWMyn~FJ*IWUvO-7b8u{Rb1z?CX>MtBUtcb8c~DCM0u%rg000080000X060MXQusXp0F*WX05t#r00000000000HlG2D**sz' + 'VRLkFY;AKdZEs{{Y;!MVb8TO6Y;|*RY;|)lZEs{{Y-w(1UuAP`UvO-7b8u{Rb1rasP)h*<6aW+e000O8001EX*jqgyAq@Zk3MT*n' + 'IRF3v0000000000q=C_C0RU!Ub98WQZF4VeZ)9a`b1!9cZC`L~b#riRb#pImZ)9a`X>MmMtBUtcb8c~DCM' + '0u%rg000080000X08g{=sZ;_007C`<03`qb00000000000HlG~c>w@sVRLkFY;AKdZEs{{Y;!MVb8TjCY-BPoVRUq5ZggpHZZ2?n' + 'P)h*<6aW+e000O8001EXS-WdA^bG(2#4!K>G5`Po0000000000q=9dI0RU!Ub98WQZF4VeZ)9a`b1!9cZDwz5WHK*fZ*FF3XLWL6' + 'bZKvHUuAP`W^Zg{GA?j=P)h*<6aW+e000O8001EX3lQLD4+a1LI1vB(^b0000000000q=B=F0RU!Ub98WQZF4VeZ)9a`b1!9c' + 'ZDwz5WHK*hZDM6)WNB_^E^v8JO928D02BZK00;m803iU9yAY675dZ+OPXGWd00000000000001_fd!EP0A^uxbZ~5Kb1!XgWMyn~' + 'FJ*IWW^Zg{GB0IwZDwz5WHMi2W@c$_Wiwr7S6f@004*z001li0000000000005+crJ?}OmVRLkFY;AKdZEs{{Y;!MV' + 'b8TjCY-BPoWpiz2Z){{TUukY|b#z~1b!lv5WpXZXc~DCM0u%rg000080000X04;act+faM0O21104x9i00000000000HlG3&jA2t' + 'VRLkFY;AKdZEs{{Y;!MVb8TjCY-BPoWpiz2Z){{TUu|V{a%p38E^v8JO928D02BZK00;m803iUXR7#%M5C8zOGyni80000000000' + '0001_fri)t0A^uxbZ~5Kb1!XgWMyn~FJ*IWW^Zg{GB0IwZDwz5WHMiEX>(&PaCuNm0Rj{N6aWAK2mk;8App~FYO+}m006!~001%o' + '0000000000005+crRf0xW?^%5aBOXJFKusRWo&aVWpiz2Z){{TFJ*IWW^Zg{GGA?NV`*kzaA9(DX>MmOaCuNm0Rj{N6aWAK2mk;8' + 'AppN4XoZ0!005nz001@s0000000000005+cPWb@003uD001Ze0000000000005+cBOU?(W?^%5aBOXJFKusRWo&aVWpiz2Z){{TFJ*IW' + 'W^Zg{GGA?TVJ>iaP)h*<6aW+e000O8001EX-=$OvzXbpQ&JF+oIRF3v0000000000q=9NM0sv-Vb98WQZF4VeZ)9a`b1!9cZDwz5' + 'WHK*hb8TjCY-BQDZF6B?W@&76WpZD0WpQ<7ZewM0E^v8JO928D02BZK00;m803iS)RN+e40{{T>2><{(00000000000001_fq*sw' + '0A^uxbZ~5Kb1!XgWMyn~FJ*IWW^Zg{GB0IwZDwz5WHMiFZ*py6Y-xIBUt@1?a%5?4VRU74E^v8JO928D02BZK00;m803iTOxO|h+' + '2LJ#o9RL6;00000000000001_fv`FP0A^uxbZ~5Kb1!XgWMyn~FJ*IWW^Zg{GB0IwZDwz5WHMiGb#!obbS`jtP)h*<6aW+e000O8' + '001EXSIg0=e+d8p(Ifx>F8}}l0000000000q=Cyq0sv-Vb98WQZF4VeZ)9a`b1!9cZDwz5WHK*hb8TjCY-BQDaA9e3Wn^D%b73xU' + 'c~DCM0u%rg000080000X0HjYs!|Mb90NV}#04x9i00000000000HlGJOacIAVRLkFY;AKdZEs{{Y;!MVb8TjCY-BPoWpiz2Z){{T' + 'UvOb^b7^j8E^v8JO928D02BZK00;m803iU_HM^@-1polW4*&o>00000000000001_fy_|?0A^uxbZ~5Kb1!XgWMyn~FJ*IWW^Zg{' + 'GB0IwZDwz5WHMiHa%E&`V{~O?Utw%%XKrO=UuAM~Z*nehc~DCM0u%rg000080000X0P&t651JzY03V+K05Sjo00000000000HlF?' + 'R{{WLVRLkFY;AKdZEs{{Y;!MVb8TjCY-BPoWpiz2Z){{TUvP3|aA9&~UukY|b#yLpc~DCM0u%rg000080000X03q06#9Ryj0CFw>' + '04@Lk00000000000HlFsdjbGvVRLkFY;AKdZEs{{Y;!MVb8TjCY-BPoWpiz2Z){{TUvP47V`X!5Z*nehc~DCM0u%rg000080000X' + '0AMka!AUd#0LkM305Sjo00000000000HlEphynm+VRLkFY;AKdZEs{{Y;!MVb8TjCY-BPoWpiz2Z){{TUvP47bY*F7Ut?%tX>KlX' + 'c~DCM0u%rg000080000X0FeA_{aY~r0As=c05kvq00000000000HlGYyaE7bVRLkFY;AKdZEs{{Y;!MVb8TjCY-BPoWpiz2Z){{T' + 'UvP47bY*F7Ut@1=aBO9GE^v8JO928D02BZK00;m803iVF11L8n2><|fApih200000000000001_fm7}R0A^uxbZ~5Kb1!XgWMyn~' + 'FJ*IWW^Zg{GB0IwZDwz5WHMiHa&L5HX>MP0baHiLbairNE^v8JO928D02BZK00;m803iSz2FU>99smHBfB*nC00000000000001_' + 'fy?&-0A^uxbZ~5Kb1!XgWMyn~FJ*IWW^Zg{GB0IwZDwz5WHMiJWpim{b!A^;Z*FsRVQzGDE^v8JO928D02BZK00;m803iUv*^dJa' + '1pok84*&o;00000000000001_fdUr;0A^uxbZ~5Kb1!XgWMyn~FJ*IWW^Zg{GB0IwZDwz5WHMiKWpQ<7Zggp3Y+qzybYWv`VRLgX' + 'aCuNm0Rj{N6aWAK2mk;8ApnO$$E1$}002M+001ih0000000000005+cW*h?mW?^%5aBOXJFKusRWo&aVWpiz2Z){{TFJ*IWW^Zg{' + 'GGB9fb97~GE^v8JO928D02BZK00;m803iUf73B0Z0RRAN0{{Rj00000000000001_fj%Dt0A^uxbZ~5Kb1!XgWMyn~FJ*IWW^Zg{' + 'GB0IwZDwz5WHMiLd2nTOE^v8JO928D02BZK00;m803iTjZ-PV`0ssIa1ONaq00000000000001_fw~|A0A^uxbZ~5Kb1!XgWMyn~' + 'FJ*IWW^Zg{GB0IwZDwz5WHMiMbZKmJUvznJWpgfYc~DCM0u%rg000080000X01RpJcs@t~053WL04)Fj00000000000HlE;BLe_t' + 'VRLkFY;AKdZEs{{Y;!MVb8TjCY-BPoZEs{{Y-w(1UuAP`W^Zg{GA?j=P)h*<6aW+e000O8001EXxRhx`;s^i$h$8?1JOBUy00000' + '00000q=Bnz0{~`Wb98WQZF4VeZ)9a`b1!9cZDwz5WHK*pZ)9a`X>MmV>WaCuNm0Rj{N6aWAK' + '2mk;8Apm^5Gl|?;004WO0RS`r0000000000005+c=X3)AW?^%5aBOXJFKusRWo&aVWpiz2Z){{TFKusRWo&6~XJ2J=ZDwz5WHMi4' + 'Z*6UFZZ2?nP)h*<6aW+e000O8001EXJ?2NjWFr6o)}#OcIsgCw0000000000q=6OB0{~`Wb98WQZF4VeZ)9a`b1!9cZDwz5WHK*p' + 'Z)9a`X>MmMO~WoBh^Wo~0-UuR`+ZDn+Fc`j;mZ*EXa0Rj{N6aWAK2mk;8Apk*H4Da#?007|~' + '001cf0000000000005+c7YPIaW?^%5aBOXJFKusRWo&aVWpiz2Z){{TFK}{ibY*F7Uv+e8Y;!Jfc~DCM0u%rg000080000X0Ch12' + 'j<*5;0M7;h04o3h00000000000HlFX5(EHdVRLkFY;AKdZEs{{Y;!MVb8TjCY-BPoa%FIGZ)A02X<}(?X>@rmaCuNm0Rj{N6aWAK' + '2mk;8Apk-}Qz-Z#004cH000&M0000000000005+cP!$9KW?^%5aBOXJFKusRWo&aVbaZqsaCuNm0Rj{N6aWAK2mk;8ApkyXf&%d~' + '008p(000yK0000000000005+cd^Q9CW?^%5aBOXJFLGsPX>)XPc`k5yP)h*<6aW+e000O8001EXi^2h}Py+w}?Fj$?6951J00000' + '00000q=BDm1OR4Xb98WQZF4Vjb#8QNZDlTSc~DCQ1^@s606PFX09_OS03dDz0000' +) diff --git a/model.safetensors b/model.safetensors new file mode 100644 index 0000000000000000000000000000000000000000..d9225c09e2f6378872db29cdcfc36ac4e67d8eda --- /dev/null +++ b/model.safetensors @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:11a53c1b4700b6c62a5a464fc3ec7076160c19e8584157f88e13275116cfd602 +size 684804708 diff --git a/modeling_fastplms.py b/modeling_fastplms.py new file mode 100644 index 0000000000000000000000000000000000000000..9acc2057146ace08bf93cbecf5cb730d04ce270e --- /dev/null +++ b/modeling_fastplms.py @@ -0,0 +1,193 @@ +"""Generated bridge to the embedded FastPLMs runtime sources.""" + +import base64 +import hashlib +import importlib +import importlib.util +import sys +import tempfile +from io import BytesIO +from pathlib import Path +from typing import ClassVar +from zipfile import ZIP_DEFLATED, ZipFile + +from .fastplms_bundle import RUNTIME_DATA, RUNTIME_HASH + +if RUNTIME_HASH != "ef1f313b70da02019bf0712085e332968c1cecf189633f721fce2554cff6b61f": + 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 _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(package, package_root): + incoming = _runtime_file_hashes(package_root) + known = getattr(package, "__fastplms_artifact_runtime_files__", None) + if not isinstance(known, dict): + raise RuntimeError( + "A non-artifact fastplms module is already loaded. Load the Hub artifact " + "in a separate Python process." + ) + 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." + ) + known = dict(known) + known.update(incoming) + package.__fastplms_artifact_runtime_files__ = known + roots = list(getattr(package, "__fastplms_artifact_runtime_roots__", ())) + if str(package_root) not in roots: + roots.append(str(package_root)) + package.__fastplms_artifact_runtime_roots__ = tuple(roots) + temporaries = list( + getattr(package, "__fastplms_artifact_runtime_temporaries__", ()) + ) + for temporary in _RUNTIME_TEMPORARIES: + if temporary not in temporaries: + temporaries.append(temporary) + package.__fastplms_artifact_runtime_temporaries__ = tuple(temporaries) + hashes = set(getattr(package, "__fastplms_artifact_runtime_hashes__", ())) + hashes.add(RUNTIME_HASH) + package.__fastplms_artifact_runtime_hashes__ = frozenset(hashes) + _extend_loaded_package_paths(package_root) + return package + +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(): + package = sys.modules.get("fastplms") + hashes = getattr(package, "__fastplms_artifact_runtime_hashes__", ()) + if RUNTIME_HASH in hashes: + return package + package_root = _ensure_runtime() + if package is not None: + return _merge_runtime(package, 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_182 = _import_without_bytecode("fastplms.models.esmfold2.configuration_esmfold2") +ESMFold2Config = _module_182.ESMFold2Config +ESMFold2Config.__module__ = __name__ +_module_185 = _import_without_bytecode("fastplms.models.esmfold2.modeling_esmfold2_classification") +ESMFold2ExperimentalForSequenceClassification = _module_185.ESMFold2ExperimentalForSequenceClassification +ESMFold2ExperimentalForSequenceClassification.__module__ = __name__ +ESMFold2ExperimentalForTokenClassification = _module_185.ESMFold2ExperimentalForTokenClassification +ESMFold2ExperimentalForTokenClassification.__module__ = __name__ +_module_190 = _import_without_bytecode("fastplms.models.esmfold2.modeling_esmfold2_experimental") +ESMFold2ExperimentalModel = _module_190.ESMFold2ExperimentalModel +ESMFold2ExperimentalModel.__module__ = __name__ diff --git a/preparation-record.json b/preparation-record.json new file mode 100644 index 0000000000000000000000000000000000000000..2d67018c172c86cda8c99a1a41ff1e80a6e99753 --- /dev/null +++ b/preparation-record.json @@ -0,0 +1,609 @@ +{ + "complete_artifact_compliance": false, + "config_contract": { + "auto_map": { + "AutoConfig": "modeling_fastplms.ESMFold2Config", + "AutoModel": "modeling_fastplms.ESMFold2ExperimentalModel", + "AutoModelForSequenceClassification": "modeling_fastplms.ESMFold2ExperimentalForSequenceClassification", + "AutoModelForTokenClassification": "modeling_fastplms.ESMFold2ExperimentalForTokenClassification" + }, + "esmc_id": "Synthyra/ESMplusplus_large", + "msa_conditioning": false + }, + "model_id": "esmfold2_600", + "official_checkpoint": { + "files": { + "config.json": { + "declared": "git-sha1:95517e555f17a1eca4b68866c89033a1f6916a5d", + "sha256": "fd8a437195f846b125f8986ded036426c1c1afaf45e92fb40b704e8d4f55a842", + "size": 2091, + "verified": "git-sha1:95517e555f17a1eca4b68866c89033a1f6916a5d" + }, + "model.safetensors": { + "declared": "sha256:11a53c1b4700b6c62a5a464fc3ec7076160c19e8584157f88e13275116cfd602", + "sha256": "11a53c1b4700b6c62a5a464fc3ec7076160c19e8584157f88e13275116cfd602", + "size": 684804708, + "verified": "sha256:11a53c1b4700b6c62a5a464fc3ec7076160c19e8584157f88e13275116cfd602" + } + }, + "repo_id": "biohub/ESMFold2-Experimental-Fast-base600M-step1500k", + "revision": "15cf2d6648692f6c17cee1297d8a285476fffa9b" + }, + "prepared_payloads": { + "LICENSES/FastPLMs-Apache-2.0.txt": { + "sha256": "2d2b50c7b1414bff1189a1db1f0cfb92e3e064b50f4c2b1019827b683e1b629a", + "size": 11601 + }, + "LICENSES/biohub-esm/LICENSE.md": { + "sha256": "b63df9ca1dd96b3b21eec226b51b236d0bd152ac20eafc43aad46bf832b48d8a", + "size": 1095 + }, + "LICENSES/biohub-esm/THIRD_PARTY_NOTICE.md": { + "sha256": "5bff8515ba4e0f53abdc43714c180b79c5b606160497d98de741a369cb9b6a23", + "size": 948 + }, + "LICENSES/biohub-transformers/LICENSE": { + "sha256": "77fd4710def9ec3c0f6225800e0235f15a425abd4a8b03559127fcd782612049", + "size": 11418 + }, + "LICENSES/protein-ttt/LICENSE": { + "sha256": "bb01e7d5554f9e2e117172e56551452f68a7818df7bc8e71cd7a776a1d4ba3df", + "size": 1071 + }, + "LICENSES/protein-ttt/SOURCE_RECORD.md": { + "sha256": "dc641c37353c2efd50ccbdb316ca4aae495ec02c1563e0e15bac92f75fc482e5", + "size": 393 + }, + "README.md": { + "sha256": "8b794bf2908f0711233c116aed4c129257a7d8ae5bdae97014bbf57d85e03380", + "size": 8784 + }, + "THIRD_PARTY_NOTICES.md": { + "sha256": "d35e506b728868b52290672d54a89c6aade09529af99dbc8a4db72d9e9ca3460", + "size": 4770 + }, + "config.json": { + "sha256": "583968e53f6512a9bc10bd7af1c902cecfc66c96b1df1ebfadee58dae404fd97", + "size": 2956 + }, + "fastplms/__init__.py": { + "sha256": "8503e02debf24abc6d33c1913ff3eb9f245e63f7e62017e220d072c7393fdf2c", + "size": 1857 + }, + "fastplms/attention/__init__.py": { + "sha256": "f81a965cda9903f1faaf938e08942a6ec63a96d702d20799083c84df6ffb244d", + "size": 1916 + }, + "fastplms/attention/_core.py": { + "sha256": "fd74f7d9d773306e966be5ee209da054b357c68db9921972a222e81030fd69a9", + "size": 32003 + }, + "fastplms/attention/_kernel_lock.py": { + "sha256": "56455eb57cee9af08438eda52b775c9cff3cf994cadb9b116ce2ee10fd5f1801", + "size": 7062 + }, + "fastplms/attention/interfaces.py": { + "sha256": "c2a7229addd655d63b0f26b82df75e12fce9457fb4e2d4d4b18cb8b5c90c7396", + "size": 10282 + }, + "fastplms/embeddings/__init__.py": { + "sha256": "cf0648de905a00ad16c0aece328eeb05f25739a4cf89b47b915e0f2b19c3ce26", + "size": 1505 + }, + "fastplms/embeddings/pooling.py": { + "sha256": "9bed8c466fce053333de9e25b04d1ab304f2954f403ffa0f5a3c749e8277b3ce", + "size": 9875 + }, + "fastplms/embeddings/runner.py": { + "sha256": "a7839763ad641a63d0e2671dcf34aa7b2a8114198950b64254188521b27453a8", + "size": 67610 + }, + "fastplms/embeddings/storage.py": { + "sha256": "98cf725d26766f32300f105345129ca6cf02ebd10628e746a7e71e107bc911a6", + "size": 66426 + }, + "fastplms/embeddings/types.py": { + "sha256": "c58f8c6a37a11a937c3f77b4e1617d5f213923d98ecaffddcdd283bd9dc42fe7", + "size": 6463 + }, + "fastplms/models.toml": { + "sha256": "9e3c99543177f548203fee4e92f0a40d5874217e7ba600d9afe4abce6bbc5b06", + "size": 82900 + }, + "fastplms/models/__init__.py": { + "sha256": "d3dd084f5e8fafcf2d0fe4a26e4ab9284d3f9414b1f9a5bd86e308a406447c43", + "size": 336 + }, + "fastplms/models/_esm_rotary.py": { + "sha256": "8e041bc4a056373e07fbf598e6acb6075f889d1bce12057715369804befb37ec", + "size": 3375 + }, + "fastplms/models/classification_probe.py": { + "sha256": "46f4d4a35c9df24925160198caaa69f97b85df7723e2a06b48464ba73b1da752", + "size": 20018 + }, + "fastplms/models/esm_plusplus/__init__.py": { + "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", + "size": 0 + }, + "fastplms/models/esm_plusplus/modeling_esm_plusplus.py": { + "sha256": "30c300f14f1a3c440bb1a98895f64fd0212124454360e61327f67c5a0737bf61", + "size": 79508 + }, + "fastplms/models/esm_plusplus/modeling_esm_plusplus_sae.py": { + "sha256": "a8bc37695c29f0a91597bc8c95ff0516615944ac7ee35d18e56c23ebcd52ab5d", + "size": 9994 + }, + "fastplms/models/esmfold2/__init__.py": { + "sha256": "e62860b2bd71fee8fb079dc5a10ee415b64526d48341a417d6191c1460ceadeb", + "size": 2029 + }, + "fastplms/models/esmfold2/attention.py": { + "sha256": "e9645c79e198c7dc2d277240cd8010010179f6aacb7e8bb511b8607bcfa793ae", + "size": 1603 + }, + "fastplms/models/esmfold2/configuration_esmfold2.py": { + "sha256": "b6c824daa8a7406b7527ba83d9229906fd8fe7c8a723209cdae4c9e735d1d641", + "size": 12740 + }, + "fastplms/models/esmfold2/embedding.py": { + "sha256": "088a6ee2cda21a34ec1b719689affcaefb9ea70b27a3b538846a21cf79d363cc", + "size": 4408 + }, + "fastplms/models/esmfold2/esmfold2_affine3d.py": { + "sha256": "f64c3326684266c033bb674bfce02cfff3fb4825c49173a11d12519e965816c7", + "size": 20401 + }, + "fastplms/models/esmfold2/esmfold2_aligner.py": { + "sha256": "a21b34fa91915d3d5bcceaacbd0676c6965a97058f46ceacb784ea75dbd94b27", + "size": 2696 + }, + "fastplms/models/esmfold2/esmfold2_atom_indexer.py": { + "sha256": "0b363b1c0fef1f4a37e82189d6d5912f669f544ad5d6b28c54c9b2fb3d3b884a", + "size": 920 + }, + "fastplms/models/esmfold2/esmfold2_conformers.py": { + "sha256": "0742d37e5776af7cf477aabc95d80ab519b8a7997b4b580e3beea1c5c654b222", + "size": 16121 + }, + "fastplms/models/esmfold2/esmfold2_constants.py": { + "sha256": "65f978329690b46c4a3cc9f4815fe439b1e355be03ef2a9134b83ab889958751", + "size": 6060 + }, + "fastplms/models/esmfold2/esmfold2_constants_esm3.py": { + "sha256": "6f9c8bbf0f80b6cd9ede136d142b51ff2b84f1b339cb6b6668c2669e8c2584b5", + "size": 2913 + }, + "fastplms/models/esmfold2/esmfold2_input_builder.py": { + "sha256": "90ca458a80095e959e5574ea93ac6598f6ab4512e54e91e67a1c9c93c8476372", + "size": 6859 + }, + "fastplms/models/esmfold2/esmfold2_metrics.py": { + "sha256": "09ddac476d7359d35cf008b749051943c63fa62c2ac1dce5e2276e32d47a0da4", + "size": 8161 + }, + "fastplms/models/esmfold2/esmfold2_misc.py": { + "sha256": "4439bb4fc427713841c556a1be5932f93bfc2b4d63c471983250386253470352", + "size": 13490 + }, + "fastplms/models/esmfold2/esmfold2_mmcif_parsing.py": { + "sha256": "29e968626ab375864037005651b41b26531924dcffb544d14a8b87c1aa8e680a", + "size": 16574 + }, + "fastplms/models/esmfold2/esmfold2_molecular_complex.py": { + "sha256": "8e2768ce766e3e655197a337d922a1d6fa7d61ea4b4c9d72032a6434c345f605", + "size": 40861 + }, + "fastplms/models/esmfold2/esmfold2_msa.py": { + "sha256": "24dd1309ab181cebb81c0f5fb21182863b75cc8d23b62aec02d3c48897a2c838", + "size": 20583 + }, + "fastplms/models/esmfold2/esmfold2_msa_filter_sequences.py": { + "sha256": "c679b63c6d0efadad32bc9cbe1157d6da87a363c1191347447ac735e6e8352a1", + "size": 3790 + }, + "fastplms/models/esmfold2/esmfold2_normalize_coordinates.py": { + "sha256": "a0d7180fee6c7bf8032c24b7e6c45a769115d2d6039ce8d8c3396283cb79e31d", + "size": 2546 + }, + "fastplms/models/esmfold2/esmfold2_output.py": { + "sha256": "1b82de5166e3229b5ff57e243721a24265bf8607293ebf0700647015a09af338", + "size": 7467 + }, + "fastplms/models/esmfold2/esmfold2_paired_msa.py": { + "sha256": "046d5fd2c98c5d88e0644d3cf7623cd31a355724c11fa4968907c49ff173a87f", + "size": 9425 + }, + "fastplms/models/esmfold2/esmfold2_parsing.py": { + "sha256": "6ca4a2d85d6158cadcde79bf8d1820c54eeb2160e00b99643c29e270d0a3f870", + "size": 3803 + }, + "fastplms/models/esmfold2/esmfold2_predicted_aligned_error.py": { + "sha256": "b76577c9e17bdef5417fea35c6ac134cff2ba949a1c9cb389323134541e52d04", + "size": 4037 + }, + "fastplms/models/esmfold2/esmfold2_prepare_input.py": { + "sha256": "f2b5f63a451555d41944e5abc8b85a978edfb456ce493821b313f8343d3eed19", + "size": 40735 + }, + "fastplms/models/esmfold2/esmfold2_processor.py": { + "sha256": "4eaed444d316bdb31a72b518c64ad11ede6554b88f22a68fd9bf109702529994", + "size": 11890 + }, + "fastplms/models/esmfold2/esmfold2_protein_chain.py": { + "sha256": "e87d7cbfab56d0cd420d8ce8b15a23872cdfe6c1036f2467555e57691af172af", + "size": 58313 + }, + "fastplms/models/esmfold2/esmfold2_protein_complex.py": { + "sha256": "ee78539192414000b78da7285175f0316d79f549c179678638d009054b4985f7", + "size": 49763 + }, + "fastplms/models/esmfold2/esmfold2_protein_structure.py": { + "sha256": "d77254171dd7dc7e269693091e51aa04a7cd8dc4a442d425fab7d922c1b1308d", + "size": 8565 + }, + "fastplms/models/esmfold2/esmfold2_residue_constants.py": { + "sha256": "f54e2856761f1d44e188ddd4682bd93b5eac6c3acac060ec8669efd5e19d49bd", + "size": 32919 + }, + "fastplms/models/esmfold2/esmfold2_sequential_dataclass.py": { + "sha256": "d1bf752de753f99673287f7b3365be84dfb68517256b1c3ecf0ff1fe2a0c4bc1", + "size": 3928 + }, + "fastplms/models/esmfold2/esmfold2_system.py": { + "sha256": "4a0b339bc7446a9ca255f6005ecb8d385c69b7a3e57170c5a48e372bb2340d13", + "size": 1600 + }, + "fastplms/models/esmfold2/esmfold2_types.py": { + "sha256": "8963440d40e3143b984679a716a5ac8327fa09692765836dea179c41e2bc9f5d", + "size": 875 + }, + "fastplms/models/esmfold2/esmfold2_utils_types.py": { + "sha256": "b846d61a383fc55091231ec8bc2b829045a4edb52dcd3907c30b81b6057fd4eb", + "size": 1058 + }, + "fastplms/models/esmfold2/modeling_esmfold2.py": { + "sha256": "7012061edaccc66d46641e46a3485d2bdb5eaa9b0716bc8b75831a267e5493c5", + "size": 80431 + }, + "fastplms/models/esmfold2/modeling_esmfold2_classification.py": { + "sha256": "36a720a20a6623dd31af5905d3db7b732dac52f4299ffed8e2895f6d0d49e5ad", + "size": 9096 + }, + "fastplms/models/esmfold2/modeling_esmfold2_common.py": { + "sha256": "b1db3d0c6f4742e861cd79ecdc481aeb2f92f1c8aa432e24109093ece155d663", + "size": 105595 + }, + "fastplms/models/esmfold2/modeling_esmfold2_experimental.py": { + "sha256": "de8d215d137744bf4d2b1c7bdf98617d7c20a425ea25d9c74f88b34d3303ca84", + "size": 42198 + }, + "fastplms/models/esmfold2/protein_reference_geometry.json": { + "sha256": "c726d81b928df0005b2d6877e2c0d5ef3ebfb11744e04ca8972388057dc565c0", + "size": 11364 + }, + "fastplms/models/esmfold2/protein_utils.py": { + "sha256": "a51cdad1769620389ca58d6782bc787ed70cadb1f310689cf3dcd1d2d2fa92be", + "size": 7649 + }, + "fastplms/models/esmfold2/reproducibility.py": { + "sha256": "2027064e3eea918f535d9b60734a3bc9288478f5b957847266d328dcfca88443", + "size": 1743 + }, + "fastplms/models/ttt.py": { + "sha256": "f1d7f9298a930f4c40f8e321ffa08b09d1c5b8492bca623bdc7bfd03ec85774d", + "size": 37757 + }, + "fastplms/registry.py": { + "sha256": "f34ef7a1cf33352925595dd3e9a82d118d3bf0d449ecc79321eb26f9bea9fbd5", + "size": 64242 + }, + "fastplms/runtime.py": { + "sha256": "9521c37fcd168bf9a4137277fb73672403ccf30d70d4845aa879fb22f42ddeb6", + "size": 2541 + }, + "fastplms_bundle.py": { + "sha256": "db039942006d836947079f08b45d7f4d791e2ebcf5a3ad9e0eb0cc725bb8e29a", + "size": 395508 + }, + "model.safetensors": { + "sha256": "11a53c1b4700b6c62a5a464fc3ec7076160c19e8584157f88e13275116cfd602", + "size": 684804708 + }, + "modeling_fastplms.py": { + "sha256": "834a9a71e64d8558f707e56dc382fe1893911636e984917538c35237167c0dd8", + "size": 7992 + }, + "requirements.txt": { + "sha256": "bc05aaba0c8bb6dfd7d0efe1eb565ee202bd25cbf6ccd90cda10711d070d0c67", + "size": 450 + } + }, + "publication_status": "pending", + "runtime_bundle_sha256": "ef1f313b70da02019bf0712085e332968c1cecf189633f721fce2554cff6b61f", + "schema_version": 1, + "scope": "local-preparation", + "source_payloads": { + "LICENSES/FastPLMs-Apache-2.0.txt": { + "sha256": "2d2b50c7b1414bff1189a1db1f0cfb92e3e064b50f4c2b1019827b683e1b629a", + "size": 11601 + }, + "LICENSES/biohub-esm/LICENSE.md": { + "sha256": "b63df9ca1dd96b3b21eec226b51b236d0bd152ac20eafc43aad46bf832b48d8a", + "size": 1095 + }, + "LICENSES/biohub-esm/THIRD_PARTY_NOTICE.md": { + "sha256": "5bff8515ba4e0f53abdc43714c180b79c5b606160497d98de741a369cb9b6a23", + "size": 948 + }, + "LICENSES/biohub-transformers/LICENSE": { + "sha256": "77fd4710def9ec3c0f6225800e0235f15a425abd4a8b03559127fcd782612049", + "size": 11418 + }, + "LICENSES/protein-ttt/LICENSE": { + "sha256": "bb01e7d5554f9e2e117172e56551452f68a7818df7bc8e71cd7a776a1d4ba3df", + "size": 1071 + }, + "LICENSES/protein-ttt/SOURCE_RECORD.md": { + "sha256": "dc641c37353c2efd50ccbdb316ca4aae495ec02c1563e0e15bac92f75fc482e5", + "size": 393 + }, + "README.md": { + "sha256": "8b794bf2908f0711233c116aed4c129257a7d8ae5bdae97014bbf57d85e03380", + "size": 8784 + }, + "THIRD_PARTY_NOTICES.md": { + "sha256": "d35e506b728868b52290672d54a89c6aade09529af99dbc8a4db72d9e9ca3460", + "size": 4770 + }, + "fastplms/__init__.py": { + "sha256": "8503e02debf24abc6d33c1913ff3eb9f245e63f7e62017e220d072c7393fdf2c", + "size": 1857 + }, + "fastplms/attention/__init__.py": { + "sha256": "f81a965cda9903f1faaf938e08942a6ec63a96d702d20799083c84df6ffb244d", + "size": 1916 + }, + "fastplms/attention/_core.py": { + "sha256": "fd74f7d9d773306e966be5ee209da054b357c68db9921972a222e81030fd69a9", + "size": 32003 + }, + "fastplms/attention/_kernel_lock.py": { + "sha256": "56455eb57cee9af08438eda52b775c9cff3cf994cadb9b116ce2ee10fd5f1801", + "size": 7062 + }, + "fastplms/attention/interfaces.py": { + "sha256": "c2a7229addd655d63b0f26b82df75e12fce9457fb4e2d4d4b18cb8b5c90c7396", + "size": 10282 + }, + "fastplms/embeddings/__init__.py": { + "sha256": "cf0648de905a00ad16c0aece328eeb05f25739a4cf89b47b915e0f2b19c3ce26", + "size": 1505 + }, + "fastplms/embeddings/pooling.py": { + "sha256": "9bed8c466fce053333de9e25b04d1ab304f2954f403ffa0f5a3c749e8277b3ce", + "size": 9875 + }, + "fastplms/embeddings/runner.py": { + "sha256": "a7839763ad641a63d0e2671dcf34aa7b2a8114198950b64254188521b27453a8", + "size": 67610 + }, + "fastplms/embeddings/storage.py": { + "sha256": "98cf725d26766f32300f105345129ca6cf02ebd10628e746a7e71e107bc911a6", + "size": 66426 + }, + "fastplms/embeddings/types.py": { + "sha256": "c58f8c6a37a11a937c3f77b4e1617d5f213923d98ecaffddcdd283bd9dc42fe7", + "size": 6463 + }, + "fastplms/models.toml": { + "sha256": "9e3c99543177f548203fee4e92f0a40d5874217e7ba600d9afe4abce6bbc5b06", + "size": 82900 + }, + "fastplms/models/__init__.py": { + "sha256": "d3dd084f5e8fafcf2d0fe4a26e4ab9284d3f9414b1f9a5bd86e308a406447c43", + "size": 336 + }, + "fastplms/models/_esm_rotary.py": { + "sha256": "8e041bc4a056373e07fbf598e6acb6075f889d1bce12057715369804befb37ec", + "size": 3375 + }, + "fastplms/models/classification_probe.py": { + "sha256": "46f4d4a35c9df24925160198caaa69f97b85df7723e2a06b48464ba73b1da752", + "size": 20018 + }, + "fastplms/models/esm_plusplus/__init__.py": { + "sha256": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", + "size": 0 + }, + "fastplms/models/esm_plusplus/modeling_esm_plusplus.py": { + "sha256": "30c300f14f1a3c440bb1a98895f64fd0212124454360e61327f67c5a0737bf61", + "size": 79508 + }, + "fastplms/models/esm_plusplus/modeling_esm_plusplus_sae.py": { + "sha256": "a8bc37695c29f0a91597bc8c95ff0516615944ac7ee35d18e56c23ebcd52ab5d", + "size": 9994 + }, + "fastplms/models/esmfold2/__init__.py": { + "sha256": "e62860b2bd71fee8fb079dc5a10ee415b64526d48341a417d6191c1460ceadeb", + "size": 2029 + }, + "fastplms/models/esmfold2/attention.py": { + "sha256": "e9645c79e198c7dc2d277240cd8010010179f6aacb7e8bb511b8607bcfa793ae", + "size": 1603 + }, + "fastplms/models/esmfold2/configuration_esmfold2.py": { + "sha256": "b6c824daa8a7406b7527ba83d9229906fd8fe7c8a723209cdae4c9e735d1d641", + "size": 12740 + }, + "fastplms/models/esmfold2/embedding.py": { + "sha256": "088a6ee2cda21a34ec1b719689affcaefb9ea70b27a3b538846a21cf79d363cc", + "size": 4408 + }, + "fastplms/models/esmfold2/esmfold2_affine3d.py": { + "sha256": "f64c3326684266c033bb674bfce02cfff3fb4825c49173a11d12519e965816c7", + "size": 20401 + }, + "fastplms/models/esmfold2/esmfold2_aligner.py": { + "sha256": "a21b34fa91915d3d5bcceaacbd0676c6965a97058f46ceacb784ea75dbd94b27", + "size": 2696 + }, + "fastplms/models/esmfold2/esmfold2_atom_indexer.py": { + "sha256": "0b363b1c0fef1f4a37e82189d6d5912f669f544ad5d6b28c54c9b2fb3d3b884a", + "size": 920 + }, + "fastplms/models/esmfold2/esmfold2_conformers.py": { + "sha256": "0742d37e5776af7cf477aabc95d80ab519b8a7997b4b580e3beea1c5c654b222", + "size": 16121 + }, + "fastplms/models/esmfold2/esmfold2_constants.py": { + "sha256": "65f978329690b46c4a3cc9f4815fe439b1e355be03ef2a9134b83ab889958751", + "size": 6060 + }, + "fastplms/models/esmfold2/esmfold2_constants_esm3.py": { + "sha256": "6f9c8bbf0f80b6cd9ede136d142b51ff2b84f1b339cb6b6668c2669e8c2584b5", + "size": 2913 + }, + "fastplms/models/esmfold2/esmfold2_input_builder.py": { + "sha256": "90ca458a80095e959e5574ea93ac6598f6ab4512e54e91e67a1c9c93c8476372", + "size": 6859 + }, + "fastplms/models/esmfold2/esmfold2_metrics.py": { + "sha256": "09ddac476d7359d35cf008b749051943c63fa62c2ac1dce5e2276e32d47a0da4", + "size": 8161 + }, + "fastplms/models/esmfold2/esmfold2_misc.py": { + "sha256": "4439bb4fc427713841c556a1be5932f93bfc2b4d63c471983250386253470352", + "size": 13490 + }, + "fastplms/models/esmfold2/esmfold2_mmcif_parsing.py": { + "sha256": "29e968626ab375864037005651b41b26531924dcffb544d14a8b87c1aa8e680a", + "size": 16574 + }, + "fastplms/models/esmfold2/esmfold2_molecular_complex.py": { + "sha256": "8e2768ce766e3e655197a337d922a1d6fa7d61ea4b4c9d72032a6434c345f605", + "size": 40861 + }, + "fastplms/models/esmfold2/esmfold2_msa.py": { + "sha256": "24dd1309ab181cebb81c0f5fb21182863b75cc8d23b62aec02d3c48897a2c838", + "size": 20583 + }, + "fastplms/models/esmfold2/esmfold2_msa_filter_sequences.py": { + "sha256": "c679b63c6d0efadad32bc9cbe1157d6da87a363c1191347447ac735e6e8352a1", + "size": 3790 + }, + "fastplms/models/esmfold2/esmfold2_normalize_coordinates.py": { + "sha256": "a0d7180fee6c7bf8032c24b7e6c45a769115d2d6039ce8d8c3396283cb79e31d", + "size": 2546 + }, + "fastplms/models/esmfold2/esmfold2_output.py": { + "sha256": "1b82de5166e3229b5ff57e243721a24265bf8607293ebf0700647015a09af338", + "size": 7467 + }, + "fastplms/models/esmfold2/esmfold2_paired_msa.py": { + "sha256": "046d5fd2c98c5d88e0644d3cf7623cd31a355724c11fa4968907c49ff173a87f", + "size": 9425 + }, + "fastplms/models/esmfold2/esmfold2_parsing.py": { + "sha256": "6ca4a2d85d6158cadcde79bf8d1820c54eeb2160e00b99643c29e270d0a3f870", + "size": 3803 + }, + "fastplms/models/esmfold2/esmfold2_predicted_aligned_error.py": { + "sha256": "b76577c9e17bdef5417fea35c6ac134cff2ba949a1c9cb389323134541e52d04", + "size": 4037 + }, + "fastplms/models/esmfold2/esmfold2_prepare_input.py": { + "sha256": "f2b5f63a451555d41944e5abc8b85a978edfb456ce493821b313f8343d3eed19", + "size": 40735 + }, + "fastplms/models/esmfold2/esmfold2_processor.py": { + "sha256": "4eaed444d316bdb31a72b518c64ad11ede6554b88f22a68fd9bf109702529994", + "size": 11890 + }, + "fastplms/models/esmfold2/esmfold2_protein_chain.py": { + "sha256": "e87d7cbfab56d0cd420d8ce8b15a23872cdfe6c1036f2467555e57691af172af", + "size": 58313 + }, + "fastplms/models/esmfold2/esmfold2_protein_complex.py": { + "sha256": "ee78539192414000b78da7285175f0316d79f549c179678638d009054b4985f7", + "size": 49763 + }, + "fastplms/models/esmfold2/esmfold2_protein_structure.py": { + "sha256": "d77254171dd7dc7e269693091e51aa04a7cd8dc4a442d425fab7d922c1b1308d", + "size": 8565 + }, + "fastplms/models/esmfold2/esmfold2_residue_constants.py": { + "sha256": "f54e2856761f1d44e188ddd4682bd93b5eac6c3acac060ec8669efd5e19d49bd", + "size": 32919 + }, + "fastplms/models/esmfold2/esmfold2_sequential_dataclass.py": { + "sha256": "d1bf752de753f99673287f7b3365be84dfb68517256b1c3ecf0ff1fe2a0c4bc1", + "size": 3928 + }, + "fastplms/models/esmfold2/esmfold2_system.py": { + "sha256": "4a0b339bc7446a9ca255f6005ecb8d385c69b7a3e57170c5a48e372bb2340d13", + "size": 1600 + }, + "fastplms/models/esmfold2/esmfold2_types.py": { + "sha256": "8963440d40e3143b984679a716a5ac8327fa09692765836dea179c41e2bc9f5d", + "size": 875 + }, + "fastplms/models/esmfold2/esmfold2_utils_types.py": { + "sha256": "b846d61a383fc55091231ec8bc2b829045a4edb52dcd3907c30b81b6057fd4eb", + "size": 1058 + }, + "fastplms/models/esmfold2/modeling_esmfold2.py": { + "sha256": "7012061edaccc66d46641e46a3485d2bdb5eaa9b0716bc8b75831a267e5493c5", + "size": 80431 + }, + "fastplms/models/esmfold2/modeling_esmfold2_classification.py": { + "sha256": "36a720a20a6623dd31af5905d3db7b732dac52f4299ffed8e2895f6d0d49e5ad", + "size": 9096 + }, + "fastplms/models/esmfold2/modeling_esmfold2_common.py": { + "sha256": "b1db3d0c6f4742e861cd79ecdc481aeb2f92f1c8aa432e24109093ece155d663", + "size": 105595 + }, + "fastplms/models/esmfold2/modeling_esmfold2_experimental.py": { + "sha256": "de8d215d137744bf4d2b1c7bdf98617d7c20a425ea25d9c74f88b34d3303ca84", + "size": 42198 + }, + "fastplms/models/esmfold2/protein_reference_geometry.json": { + "sha256": "c726d81b928df0005b2d6877e2c0d5ef3ebfb11744e04ca8972388057dc565c0", + "size": 11364 + }, + "fastplms/models/esmfold2/protein_utils.py": { + "sha256": "a51cdad1769620389ca58d6782bc787ed70cadb1f310689cf3dcd1d2d2fa92be", + "size": 7649 + }, + "fastplms/models/esmfold2/reproducibility.py": { + "sha256": "2027064e3eea918f535d9b60734a3bc9288478f5b957847266d328dcfca88443", + "size": 1743 + }, + "fastplms/models/ttt.py": { + "sha256": "f1d7f9298a930f4c40f8e321ffa08b09d1c5b8492bca623bdc7bfd03ec85774d", + "size": 37757 + }, + "fastplms/registry.py": { + "sha256": "f34ef7a1cf33352925595dd3e9a82d118d3bf0d449ecc79321eb26f9bea9fbd5", + "size": 64242 + }, + "fastplms/runtime.py": { + "sha256": "9521c37fcd168bf9a4137277fb73672403ccf30d70d4845aa879fb22f42ddeb6", + "size": 2541 + }, + "fastplms_bundle.py": { + "sha256": "db039942006d836947079f08b45d7f4d791e2ebcf5a3ad9e0eb0cc725bb8e29a", + "size": 395508 + }, + "modeling_fastplms.py": { + "sha256": "834a9a71e64d8558f707e56dc382fe1893911636e984917538c35237167c0dd8", + "size": 7992 + }, + "requirements.txt": { + "sha256": "bc05aaba0c8bb6dfd7d0efe1eb565ee202bd25cbf6ccd90cda10711d070d0c67", + "size": 450 + } + }, + "target_repository": "Synthyra/ESMFold2-600" +} diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000000000000000000000000000000000000..07d30a5623e40adeb30d564a7dda27edabe97861 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,20 @@ +# Direct runtime dependencies for Synthyra/ESMFold2-600. +# FastPLMs source is embedded in this model repository. +torch>=2.13,<2.14 +transformers>=5.13,<5.14 +huggingface-hub>=0.34,<2 +tokenizers>=0.22,<0.23 +safetensors>=0.5,<1 +numpy>=1.26,<3 +einops>=0.8,<1 +tqdm>=4.67,<5 +accelerate>=1.10,<2 +biopython>=1.85,<2 +biotite>=1.4,<2 +brotli>=1.1,<2 +msgpack>=1.1,<2 +msgpack-numpy>=0.4.8,<1 +omegaconf>=2.3,<3 +rdkit>=2025.9,<2027 +scipy>=1.15,<2 +zstandard>=0.23,<1