Add computational release from Zenodo 22960654

.gitattributes

@@ -37,3 +37,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.png filter=lfs diff=lfs merge=lfs -text
 figures/*.png filter=lfs diff=lfs merge=lfs -text
 figures/*.jpg filter=lfs diff=lfs merge=lfs -text
+Computational/package/data/MbPA/mbpa_source_export.csv filter=lfs diff=lfs merge=lfs -text
+Computational/package/data/MbPA/mbpa_training_sequences.csv filter=lfs diff=lfs merge=lfs -text
+Computational/validation/figures/physicochemical_properties_full_reference.csv filter=lfs diff=lfs merge=lfs -text

Computational/README.md

@@ -0,0 +1,111 @@
+---
+title: Metalorian / MetaLATTE
+
+---
+
+# Repo structure
+
+```text
+train.py                 Train either stage with --stage
+sample_direct.py         Direct sampling
+sample_prototype.py      Prototype-seeded sampling
+metalorian/              Models and support functions
+data/metalatte/           Training and validation split
+data/metalorian/          Diffusion sequence-label training and test split
+data/MbPA/                Original MbPA CSV exports
+checkpoints/             Predictor and diffusion checkpoints
+assets/                  Local ESM-2 backbone and tokenizer
+requirements.txt
+```
+
+## Install
+
+Use Python 3.10. Run from this directory:
+
+```bash
+python3.10 -m venv .venv
+source .venv/bin/activate
+python -m pip install -r requirements.txt
+```
+
+Weights and tokenizer files are included for offline use after installation.
+Use a CUDA-capable PyTorch installation. For CPU checks, pass `--device cpu`.
+
+## Data
+
+| Split | Records |
+| --- | ---: |
+| `data/metalatte/train` | 50,694 groups |
+| `data/metalatte/validation` | 12,729 groups |
+| `data/metalorian/train` | 38,123 sequence-label records |
+| `data/metalorian/test` | 9,531 sequence-label records |
+
+MetaLATTE groups are stored in 1,852 training pre-batches and 460 validation
+pre-batches. 
+Metalorian uses the separate sequence-label splits listed above.
+
+`data/MbPA/` contains eight original CSV exports, aggregated from 14 different metal labels. File names, row counts and
+hashes are listed in `provenance/mbpa_inventory.csv`.
+
+## Train either stage
+
+MetaLATTE, initialized from pretrained ESM-2:
+
+```bash
+python train.py --stage metalatte --epochs 30 --device cuda:0 \
+  --output runs/metalatte
+```
+
+To start from the supplied predictor, add
+`--init-checkpoint checkpoints/metalatte_predictor.ckpt`.
+To resume training and optimizer state, use `--resume PATH`.
+
+Metalorian, initialized from the supplied MetaLATTE predictor:
+
+```bash
+python train.py --stage metalorian --epochs 1 --batch-size 2 --device cuda:0 \
+  --output runs/metalorian
+```
+
+Use `--init-checkpoint PATH` for another MetaLATTE checkpoint, or
+`--resume checkpoints/metalorian_diffusion.pt` to resume diffusion training.
+Here, `--epochs` counts additional passes through the training set.
+Outputs are `metalorian_latest.pt` in the output directory.
+
+Short forward/backward checks for both stages:
+
+```bash
+python train.py --stage metalatte --smoke-test --device cpu \
+  --init-checkpoint checkpoints/metalatte_predictor.ckpt --output runs/check_metalatte
+python train.py --stage metalorian --smoke-test --batch-size 2 --device cpu \
+  --output runs/check_metalorian
+```
+
+`--smoke-test` uses at most two records and 64 tokens for one update and writes
+no model checkpoint. Training uses one device; see `VALIDATION.md` for checks.
+
+## Direct sampling
+
+```bash
+python sample_direct.py --target_idx 3 --seq_len 30 --batch 1 \
+  --desired 1 --max_rounds 20 --output outputs/Cu_L30.csv
+```
+
+## Prototype-seeded sampling
+
+```bash
+python sample_prototype.py --target_idx 1 --seq_len 30 --batch 1 \
+  --proto MDPETCPCPSGGSCTCADSCKCEGCKCTSCKKSCCSCCPAECEKCAKDCVCKGGEAAEAEAEKCSCCQ \
+  --t_start 15 --desired 1 --max_rounds 20 --output outputs/Cd_L30.csv
+```
+
+```text
+Metal indices: 0 Ag, 1 Cd, 2 Co, 3 Cu, 4 Fe, 5 Hg, 6 Mn, 7 Mo,
+               8 Ni, 9 Pb, 10 Pt, 11 V, 12 W, 13 Zn, 14 NA
+```
+
+Both scripts load `checkpoints/metalatte_predictor.ckpt` and
+`checkpoints/metalorian_diffusion.pt`. Use `--help` for sampling options.
+
+The CSV contains unique accepted sequences; the adjacent JSON records settings
+and counts. A run may return fewer than `--desired` at `--max_rounds`.

Computational/package/APACHE-2.0.txt

@@ -0,0 +1,201 @@
+                                 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.

Computational/package/THIRD_PARTY_NOTICES.md

@@ -0,0 +1,10 @@
+# Third-party materials
+
+The ESM-2 backbone and tokenizer are from
+https://huggingface.co/facebook/esm2_t33_650M_UR50D.
+The ESM project license is included at `assets/esm2_t33_650M_UR50D/LICENSE`.
+
+The layer and tabular network modules retain the original Google Research
+copyright and Apache License 2.0 headers; the license text is in `APACHE-2.0.txt`.
+
+A project-wide license for the research code and data has not been assigned.

Computational/package/assets/esm2_t33_650M_UR50D/LICENSE

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

Computational/package/assets/esm2_t33_650M_UR50D/config.json

@@ -0,0 +1,29 @@
+{
+  "architectures": [
+    "EsmForMaskedLM"
+  ],
+  "attention_probs_dropout_prob": 0.0,
+  "classifier_dropout": null,
+  "emb_layer_norm_before": false,
+  "esmfold_config": null,
+  "hidden_act": "gelu",
+  "hidden_dropout_prob": 0.0,
+  "hidden_size": 1280,
+  "initializer_range": 0.02,
+  "intermediate_size": 5120,
+  "is_folding_model": false,
+  "layer_norm_eps": 1e-05,
+  "mask_token_id": 32,
+  "max_position_embeddings": 1026,
+  "model_type": "esm",
+  "num_attention_heads": 20,
+  "num_hidden_layers": 33,
+  "pad_token_id": 1,
+  "position_embedding_type": "rotary",
+  "token_dropout": true,
+  "torch_dtype": "float32",
+  "transformers_version": "4.25.0.dev0",
+  "use_cache": true,
+  "vocab_list": null,
+  "vocab_size": 33
+}

Computational/package/assets/esm2_t33_650M_UR50D/model.safetensors

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

Computational/package/assets/esm2_t33_650M_UR50D/special_tokens_map.json

@@ -0,0 +1,7 @@
+{
+  "cls_token": "<cls>",
+  "eos_token": "<eos>",
+  "mask_token": "<mask>",
+  "pad_token": "<pad>",
+  "unk_token": "<unk>"
+}

Computational/package/assets/esm2_t33_650M_UR50D/tokenizer_config.json

@@ -0,0 +1,4 @@
+{
+  "model_max_length": 1000000000000000019884624838656,
+  "tokenizer_class": "EsmTokenizer"
+}

Computational/package/assets/esm2_t33_650M_UR50D/vocab.txt

@@ -0,0 +1,33 @@
+<cls>
+<pad>
+<eos>
+<unk>
+L
+A
+G
+V
+S
+E
+R
+T
+I
+D
+P
+K
+Q
+N
+F
+Y
+M
+H
+W
+C
+X
+B
+U
+Z
+O
+.
+-
+<null_1>
+<mask>

Computational/package/checkpoints/metalatte_predictor.ckpt

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

Computational/package/checkpoints/metalorian_diffusion.pt

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

Computational/package/data/MbPA/mbpa_source_export.csv

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

Computational/package/data/MbPA/mbpa_training_sequences.csv

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

Computational/package/data/metalatte/train/data-00000-of-00004.arrow

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

Computational/package/data/metalatte/train/data-00001-of-00004.arrow

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

Computational/package/data/metalatte/train/data-00002-of-00004.arrow

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

Computational/package/data/metalatte/train/data-00003-of-00004.arrow

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

Computational/package/data/metalatte/train/dataset_info.json

@@ -0,0 +1,128 @@
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "anchor_input_ids": {
+      "feature": {
+        "feature": {
+          "dtype": "int64",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "anchor_attention_mask": {
+      "feature": {
+        "feature": {
+          "dtype": "int64",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "positive_input_ids": {
+      "feature": {
+        "feature": {
+          "dtype": "int64",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "positive_attention_mask": {
+      "feature": {
+        "feature": {
+          "dtype": "int64",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "negative_input_ids": {
+      "feature": {
+        "feature": {
+          "dtype": "int64",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "negative_attention_mask": {
+      "feature": {
+        "feature": {
+          "dtype": "int64",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "hard_negative_input_ids": {
+      "feature": {
+        "feature": {
+          "dtype": "int64",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "hard_negative_attention_mask": {
+      "feature": {
+        "feature": {
+          "dtype": "int64",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "anchor_labels": {
+      "feature": {
+        "feature": {
+          "dtype": "float32",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "positive_labels": {
+      "feature": {
+        "feature": {
+          "dtype": "float32",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "negative_labels": {
+      "feature": {
+        "feature": {
+          "dtype": "float32",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "hard_negative_labels": {
+      "feature": {
+        "feature": {
+          "dtype": "float32",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

Computational/package/data/metalatte/train/state.json

@@ -0,0 +1,22 @@
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00004.arrow"
+    },
+    {
+      "filename": "data-00001-of-00004.arrow"
+    },
+    {
+      "filename": "data-00002-of-00004.arrow"
+    },
+    {
+      "filename": "data-00003-of-00004.arrow"
+    }
+  ],
+  "_fingerprint": "3566914580fe3dd1",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

Computational/package/data/metalatte/validation/data-00000-of-00001.arrow

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

Computational/package/data/metalatte/validation/dataset_info.json

@@ -0,0 +1,128 @@
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "anchor_input_ids": {
+      "feature": {
+        "feature": {
+          "dtype": "int64",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "anchor_attention_mask": {
+      "feature": {
+        "feature": {
+          "dtype": "int64",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "positive_input_ids": {
+      "feature": {
+        "feature": {
+          "dtype": "int64",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "positive_attention_mask": {
+      "feature": {
+        "feature": {
+          "dtype": "int64",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "negative_input_ids": {
+      "feature": {
+        "feature": {
+          "dtype": "int64",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "negative_attention_mask": {
+      "feature": {
+        "feature": {
+          "dtype": "int64",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "hard_negative_input_ids": {
+      "feature": {
+        "feature": {
+          "dtype": "int64",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "hard_negative_attention_mask": {
+      "feature": {
+        "feature": {
+          "dtype": "int64",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "anchor_labels": {
+      "feature": {
+        "feature": {
+          "dtype": "float32",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "positive_labels": {
+      "feature": {
+        "feature": {
+          "dtype": "float32",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "negative_labels": {
+      "feature": {
+        "feature": {
+          "dtype": "float32",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    },
+    "hard_negative_labels": {
+      "feature": {
+        "feature": {
+          "dtype": "float32",
+          "_type": "Value"
+        },
+        "_type": "List"
+      },
+      "_type": "List"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

Computational/package/data/metalatte/validation/state.json

@@ -0,0 +1,13 @@
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "f4740a503b2e7630",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

Computational/package/data/metalorian/test/data-00000-of-00001.arrow

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

Computational/package/data/metalorian/test/dataset_info.json

@@ -0,0 +1,29 @@
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "input_ids": {
+      "feature": {
+        "dtype": "int32",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "attention_mask": {
+      "feature": {
+        "dtype": "int8",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "labels": {
+      "feature": {
+        "dtype": "int64",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

Computational/package/data/metalorian/test/state.json

@@ -0,0 +1,13 @@
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "4513d78e9e12f6ec",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

Computational/package/data/metalorian/train/data-00000-of-00001.arrow

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

Computational/package/data/metalorian/train/dataset_info.json

@@ -0,0 +1,29 @@
+{
+  "citation": "",
+  "description": "",
+  "features": {
+    "input_ids": {
+      "feature": {
+        "dtype": "int32",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "attention_mask": {
+      "feature": {
+        "dtype": "int8",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    },
+    "labels": {
+      "feature": {
+        "dtype": "int64",
+        "_type": "Value"
+      },
+      "_type": "Sequence"
+    }
+  },
+  "homepage": "",
+  "license": ""
+}

Computational/package/data/metalorian/train/state.json

@@ -0,0 +1,13 @@
+{
+  "_data_files": [
+    {
+      "filename": "data-00000-of-00001.arrow"
+    }
+  ],
+  "_fingerprint": "23bf04a78c6463c4",
+  "_format_columns": null,
+  "_format_kwargs": {},
+  "_format_type": null,
+  "_output_all_columns": false,
+  "_split": null
+}

Computational/package/metalorian/__init__.py

@@ -0,0 +1 @@
+

Computational/package/metalorian/cli.py

@@ -0,0 +1,83 @@
+import argparse
+import csv
+import json
+from pathlib import Path
+
+
+def parser(mode):
+    p = argparse.ArgumentParser(description=f'Metalorian {mode} sampling')
+    p.add_argument('--target_idx', type=int, required=True, help='Metal index; Cu=3, Cd=1, Zn=13, Ni=8')
+    p.add_argument('--seq_len', type=int, default=30)
+    p.add_argument('--batch', type=int, default=1)
+    p.add_argument('--desired', type=int, default=1)
+    p.add_argument('--max_rounds', type=int, default=150)
+    p.add_argument('--grad_steps', type=int, default=4)
+    p.add_argument('--guidance_scale', type=float, default=6.0)
+    p.add_argument('--penalty', type=float, default=0.1)
+    p.add_argument('--tau', type=float, default=1.0)
+    p.add_argument('--label_clamp', type=float, default=0.8)
+    p.add_argument('--best_decode', type=int, default=16, help='Candidate decodes; 0 uses argmax')
+    p.add_argument('--decode_temp', type=float, default=1.0)
+    p.add_argument('--device', default='cuda:0')
+    p.add_argument('--seed', type=int, default=42)
+    p.add_argument('--output', type=Path, required=True)
+    if mode == 'prototype':
+        p.add_argument('--proto', required=True)
+        p.add_argument('--t_start', type=int, default=15)
+    p.set_defaults(guide_on='x0', no_accum=False, track_best=False, accept_all=False,
+                   amp=False, out_prefix=None, cys_logit_bias=0.0, cys_target=0.0,
+                   cys_weight=0.0, min_cys=0, max_aa_frac=1.0, relax_after=10**9,
+                   init_noise=0.3)
+    return p
+
+
+def collect(args, sampler_class, mode):
+    from .sampling import FLAGS, METAL_LABELS
+    if not 0 <= args.target_idx < 15 or not 1 <= args.seq_len <= 1022:
+        raise ValueError('target_idx must be 0–14; seq_len must be 1–1022')
+    if min(args.batch, args.desired, args.max_rounds) < 1:
+        raise ValueError('batch, desired and max_rounds must be positive')
+    if mode == 'prototype':
+        if not 0 < args.t_start < 50:
+            raise ValueError('t_start must be 1–49')
+        if not args.proto or set(args.proto)-set('ACDEFGHIKLMNPQRSTVWY'):
+            raise ValueError('proto must contain standard uppercase amino acids')
+    import random
+    import numpy as np
+    import torch
+    random.seed(args.seed)
+    np.random.seed(args.seed)
+    torch.manual_seed(args.seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(args.seed)
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.backends.cudnn.allow_tf32 = True
+    sampler = sampler_class(args, device=args.device)
+    rows, seen, candidates = [], set(), 0
+    for round_index in range(1, args.max_rounds+1):
+        sequences, _ = sampler.run_round()
+        sequences = [s for s in sequences if s]
+        candidates += len(sequences)
+        hits, _ = sampler.accept(sequences, None) if sequences else ([], None)
+        for hit in hits:
+            if hit['sequence'] in seen:
+                continue
+            seen.add(hit['sequence'])
+            rows.append(dict(metal=METAL_LABELS[args.target_idx], **hit))
+            if len(rows) == args.desired:
+                break
+        print(f'round={round_index} candidates={candidates} unique_accepted={len(rows)}/{args.desired}', flush=True)
+        if len(rows) >= args.desired:
+            break
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    with args.output.open('w', newline='') as f:
+        writer = csv.DictWriter(f, fieldnames=['metal','sequence','score','n_cys','n_his'])
+        writer.writeheader()
+        writer.writerows(rows)
+    metadata = dict(mode=mode, arguments={k:str(v) if isinstance(v,Path) else v for k,v in vars(args).items()},
+                    candidates=candidates, unique_accepted=len(rows), rounds=round_index,
+                    target_threshold=sampler.threshold, diffusion_steps=int(FLAGS.T),
+                    beta_start=float(FLAGS.beta_1), beta_end=float(FLAGS.beta_T),
+                    predictor='checkpoints/metalatte_predictor.ckpt', diffusion='checkpoints/metalorian_diffusion.pt')
+    args.output.with_suffix('.metadata.json').write_text(json.dumps(metadata, indent=2)+'\n')
+    print(f'Wrote {len(rows)} sequences to {args.output}; requested {args.desired}.')

Computational/package/metalorian/continuous.py

@@ -0,0 +1,202 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import math
+from .layers import get_timestep_embedding
+
+def extract(v, t, x_shape):
+    out = torch.gather(v, index=t, dim=0).float()
+    return out.view([t.shape[0]] + [1] * (len(x_shape) - 1))
+
+
+class GaussianDiffusionTrainer(nn.Module):
+    def __init__(self, model, beta_1, beta_T, T):
+        super().__init__()
+
+        self.model = model
+        self.T = T
+        betas = torch.linspace(beta_1, beta_T, T, dtype=torch.float64).double()
+        alphas = 1. - betas
+        self.register_buffer('betas', betas)
+        alphas_bar = torch.cumprod(alphas, dim=0)
+
+        self.register_buffer(
+            'sqrt_alphas_bar', torch.sqrt(alphas_bar))
+        self.register_buffer(
+            'sqrt_one_minus_alphas_bar', torch.sqrt(1. - alphas_bar))
+        self.register_buffer(
+            'sqrt_recip_alphas_bar', torch.sqrt(1. / alphas_bar))
+        self.register_buffer(
+            'sqrt_recipm1_alphas_bar', torch.sqrt(1. / alphas_bar - 1))
+
+    def make_x_t(self, x_0_con, t, noise):
+        x_t_con = (
+            extract(self.sqrt_alphas_bar, t, x_0_con.shape) * x_0_con +
+            extract(self.sqrt_one_minus_alphas_bar, t, x_0_con.shape) * noise)
+        return x_t_con
+
+    def predict_xstart_from_eps(self, x_t, t, eps):
+        assert x_t.shape == eps.shape
+        return (
+            extract(self.sqrt_recip_alphas_bar, t, x_t.shape) * x_t -
+            extract(self.sqrt_recipm1_alphas_bar, t, x_t.shape) * eps
+        )
+
+
+class GaussianDiffusionSampler(nn.Module):
+    def __init__(self, model, beta_1, beta_T, T,
+                 mean_type='eps', var_type='fixedlarge'):
+        assert mean_type in ['xprev' 'xstart', 'epsilon']
+        assert var_type in ['fixedlarge', 'fixedsmall']
+        super().__init__()
+
+        self.model = model
+        self.T = T
+        self.mean_type = mean_type
+        self.var_type = var_type
+
+        betas = torch.linspace(beta_1, beta_T, T, dtype=torch.float64).double()
+
+        alphas = 1. - betas
+        self.register_buffer(
+            'betas', betas)
+        alphas_bar = torch.cumprod(alphas, dim=0)
+        alphas_bar_prev = F.pad(alphas_bar, [1, 0], value=1)[:T]
+
+        self.register_buffer(
+            'sqrt_alphas_bar', torch.sqrt(alphas_bar))
+        self.register_buffer(
+            'sqrt_one_minus_alphas_bar', torch.sqrt(1. - alphas_bar))
+
+        self.register_buffer(
+            'sqrt_recip_alphas_bar', torch.sqrt(1. / alphas_bar))
+        self.register_buffer(
+            'sqrt_recipm1_alphas_bar', torch.sqrt(1. / alphas_bar - 1))
+
+        self.register_buffer(
+            'posterior_var',
+            self.betas * (1. - alphas_bar_prev) / (1. - alphas_bar))
+        # Avoid log(0) at the first diffusion step.
+        self.register_buffer(
+            'posterior_log_var_clipped',
+            torch.log(
+                torch.cat([self.posterior_var[1:2], self.posterior_var[1:]])))
+        self.register_buffer(
+            'posterior_mean_coef1',
+            torch.sqrt(alphas_bar_prev) * self.betas / (1. - alphas_bar))
+        self.register_buffer(
+            'posterior_mean_coef2',
+            torch.sqrt(alphas) * (1. - alphas_bar_prev) / (1. - alphas_bar))
+
+    def q_mean_variance(self, x_0, x_t, t):
+        """Mean and log variance of q(x_{t-1} | x_t, x_0)."""
+        assert x_0.shape == x_t.shape
+        posterior_mean = (
+            extract(self.posterior_mean_coef1, t, x_t.shape) * x_0 +
+            extract(self.posterior_mean_coef2, t, x_t.shape) * x_t
+        )
+        posterior_log_var_clipped = extract(
+            self.posterior_log_var_clipped, t, x_t.shape)
+        return posterior_mean, posterior_log_var_clipped
+
+    def predict_xstart_from_eps(self, x_t, t, eps):
+        assert x_t.shape == eps.shape
+        return (
+            extract(self.sqrt_recip_alphas_bar, t, x_t.shape) * x_t -
+            extract(self.sqrt_recipm1_alphas_bar, t, x_t.shape) * eps
+        )
+
+
+    def p_mean_variance(self, x_t, t, cond, trans):
+        model_log_var = {
+            'fixedlarge': torch.log(torch.cat([self.posterior_var[1:2],
+                                               self.betas[1:]])),
+            'fixedsmall': self.posterior_log_var_clipped,
+        }[self.var_type]
+        model_log_var = extract(model_log_var, t, x_t.shape)
+
+        if self.mean_type == 'epsilon':
+            eps = self.model(x_t, t, cond)
+            x_0 = self.predict_xstart_from_eps(x_t, t, eps=eps)
+            model_mean, _ = self.q_mean_variance(x_0, x_t, t)
+        else:
+            raise NotImplementedError(self.mean_type)
+
+        return model_mean, model_log_var
+
+
+class DiffusionProteinModel(nn.Module):
+    def __init__(self, base_model, FLAGS):
+        super().__init__()
+        self.base_model = base_model
+        self._freeze_base_model_layers()
+        self.embed_dim = FLAGS.nf
+
+        self.time_embed = nn.Sequential(
+            nn.Linear(self.embed_dim, self.embed_dim * 4),
+            nn.SiLU(),
+            nn.Linear(self.embed_dim * 4, 1280)
+        )
+
+        self.cond_proj = nn.Linear(FLAGS.cond_size, 1280)
+
+        self.output_proj = nn.Sequential(
+            nn.Linear(1280, FLAGS.output_size * 2),
+            nn.SiLU(),
+            nn.Linear(FLAGS.output_size * 2, FLAGS.output_size)
+        )
+    def _freeze_base_model_layers(self):
+        for param in self.base_model.parameters():
+            param.requires_grad = False
+
+        for i in range(23, 33):
+            for param in self.base_model.esm_model.encoder.layer[i].parameters():
+                param.requires_grad = True
+
+        for component in [self.base_model.lm_head, self.base_model.attn_head,
+                         self.base_model.attn_ln, self.base_model.attn_skip,
+                         self.base_model.linear_layers, self.base_model.reduction_layers,
+                         self.base_model.clf_ln]:
+            for param in component.parameters():
+                param.requires_grad = True
+
+        for param in self.base_model.esm_model.embeddings.parameters():
+            param.requires_grad = False
+    def decode_embeddings(self, embeddings):
+        mlm_logits = self.base_model.lm_head(embeddings)
+        return mlm_logits
+
+    def get_embeddings(self, x, detach=False):
+        _, _, x_combined, _ = self.base_model(
+            input_ids=x['input_ids'],
+            attention_masks=x['attention_mask']
+        )
+        if detach:
+            return x_combined.detach()
+        return x_combined
+
+    def forward(self, x, time_cond, cond):
+        if isinstance(x, dict):
+            _, _, x_combined, _ = self.base_model(
+                input_ids=x['input_ids'],
+                attention_masks=x['attention_mask']
+            )
+            x = x_combined
+        seq_length = x.size(1)
+
+        temb = get_timestep_embedding(time_cond, self.embed_dim)
+        temb = self.time_embed(temb)
+        temb = temb.unsqueeze(1).expand(-1, seq_length, -1)
+
+        if len(cond.shape) == 3:
+            cond = self.cond_proj(cond)
+        else:
+            cond = self.cond_proj(cond)
+            cond = cond.unsqueeze(1).expand(-1, seq_length, -1)
+
+        x = x + temb + cond
+        return x
+
+    def get_token_embeddings(self):
+        return self.base_model.esm_model.embeddings.word_embeddings.weight

Computational/package/metalorian/data.py

@@ -0,0 +1,35 @@
+import torch
+
+def collate_fn_hard(batch):
+    # Each stored row is a pre-batch.
+    item = batch[0]
+    return {
+        'anchor_input_ids': torch.tensor(item['anchor_input_ids']),
+        'anchor_attention_mask': torch.tensor(item['anchor_attention_mask']),
+        'positive_input_ids': torch.tensor(item['positive_input_ids']),
+        'positive_attention_mask': torch.tensor(item['positive_attention_mask']),
+        'negative_input_ids': torch.tensor(item['negative_input_ids']),
+        'negative_attention_mask': torch.tensor(item['negative_attention_mask']),
+        'anchor_labels': torch.tensor(item['anchor_labels']),
+        'positive_labels': torch.tensor(item['positive_labels']),
+        'negative_labels': torch.tensor(item['negative_labels']),
+        'hard_negative_input_ids': torch.tensor(item['hard_negative_input_ids']),
+        'hard_negative_attention_mask': torch.tensor(item['hard_negative_attention_mask']),
+        'hard_negative_labels': torch.tensor(item['hard_negative_labels'])
+    }
+
+
+def collate_fn(batch):
+    # Each stored row is a pre-batch.
+    item = batch[0]
+    return {
+        'anchor_input_ids': torch.tensor(item['anchor_input_ids']),
+        'anchor_attention_mask': torch.tensor(item['anchor_attention_mask']),
+        'positive_input_ids': torch.tensor(item['positive_input_ids']),
+        'positive_attention_mask': torch.tensor(item['positive_attention_mask']),
+        'negative_input_ids': torch.tensor(item['negative_input_ids']),
+        'negative_attention_mask': torch.tensor(item['negative_attention_mask']),
+        'anchor_labels': torch.tensor(item['anchor_labels']),
+        'positive_labels': torch.tensor(item['positive_labels']),
+        'negative_labels': torch.tensor(item['negative_labels']),
+    }

Computational/package/metalorian/discrete.py

@@ -0,0 +1,276 @@
+import torch
+import torch.nn.functional as F
+import numpy as np
+from inspect import isfunction
+
+
+"""
+Based in part on: https://github.com/lucidrains/denoising-diffusion-pytorch/blob/5989f4c77eafcdc6be0fb4739f0f277a6dd7f7d8/denoising_diffusion_pytorch/denoising_diffusion_pytorch.py#L281
+"""
+eps = 1e-8
+
+def log_1_min_a(a):
+    return torch.log(1 - a.exp() + 1e-40)
+
+def log_add_exp(a, b):
+    maximum = torch.max(a, b)
+    return maximum + torch.log(torch.exp(a - maximum) + torch.exp(b - maximum))
+
+
+def extract(a, t, x_shape):
+    b, *_ = t.shape
+    out = a.gather(-1, t)
+    return out.reshape(b, *((1,) * (len(x_shape) - 1)))
+
+def index_to_log_onehot(x, num_classes):
+    log_x = torch.log(x.float().clamp(min=1e-30))
+
+    return log_x
+
+
+class MultinomialDiffusion(torch.nn.Module):
+    def __init__(self, num_classes, shape, denoise_fn, FLAGS, timesteps=1000,
+                 loss_type='vb_stochastic', parametrization='x0'):
+        super(MultinomialDiffusion, self).__init__()
+        assert loss_type in ('vb_stochastic', 'vb_all')
+        assert parametrization in ('x0', 'direct')
+
+        if loss_type == 'vb_all':
+            print('Computing the loss using the bound on _all_ timesteps.'
+                  ' This is expensive both in terms of memory and computation.')
+        self.num_classes = num_classes
+        self._denoise_fn = denoise_fn
+        self.loss_type = loss_type
+        self.shape = shape
+        self.num_timesteps = timesteps
+        self.parametrization = parametrization
+
+        betas = torch.linspace(FLAGS.beta_1, FLAGS.beta_T, FLAGS.T, dtype=torch.float64).double()
+        alphas = 1. - betas
+
+        alphas = np.sqrt(alphas)
+        betas = 1. - alphas
+
+        log_alpha = np.log(alphas)
+        log_cumprod_alpha = np.cumsum(log_alpha)
+        log_1_min_alpha = log_1_min_a(log_alpha)
+        log_1_min_cumprod_alpha = log_1_min_a(log_cumprod_alpha)
+        self.num_classes_column = np.concatenate([self.num_classes[i].repeat(self.num_classes[i]) for i in range(len(self.num_classes))])
+        assert log_add_exp(log_alpha, log_1_min_alpha).abs().sum().item() < 1.e-5
+        assert log_add_exp(log_cumprod_alpha, log_1_min_cumprod_alpha).abs().sum().item() < 1e-5
+        assert (np.cumsum(log_alpha) - log_cumprod_alpha).abs().sum().item() < 1.e-5
+
+        self.register_buffer('log_alpha', log_alpha.float())
+        self.register_buffer('log_1_min_alpha', log_1_min_alpha.float())
+        self.register_buffer('log_cumprod_alpha', log_cumprod_alpha.float())
+        self.register_buffer('log_1_min_cumprod_alpha', log_1_min_cumprod_alpha.float())
+
+        self.register_buffer('Lt_history', torch.zeros(timesteps))
+        self.register_buffer('Lt_count', torch.zeros(timesteps))
+
+    def multinomial_kl(self, log_prob1, log_prob2):
+        kl = (log_prob1.exp() * (log_prob1 - log_prob2))
+        k=0
+        kl_list = []
+        for i in self.num_classes:
+            sub = kl[:, k:i+k].mean(dim=1)
+            kl_list.append(sub)
+            k+=i
+        kl = torch.stack(kl_list, 1)
+        return kl
+
+    def log_categorical(self, log_x_start, log_prob):
+        kl = (log_x_start.exp() * log_prob)
+        k=0
+        kl_list = []
+        for i in self.num_classes:
+            sub =  kl[:, k:i+k].mean(dim=1)
+            kl_list.append(sub)
+            k+=i
+        kl = torch.stack(kl_list, 1)
+
+        return kl
+
+    def q_pred_one_timestep(self, log_x_t, t):
+        log_alpha_t = extract(self.log_alpha, t, log_x_t.shape)
+        log_1_min_alpha_t = extract(self.log_1_min_alpha, t, log_x_t.shape)
+
+        log_probs = log_add_exp(
+            log_x_t + log_alpha_t,
+            log_1_min_alpha_t -torch.tensor(np.log(self.num_classes_column)).to(log_1_min_alpha_t.device)
+        )
+
+        return log_probs
+
+    def q_pred(self, log_x_start, t):
+
+        log_cumprod_alpha_t = extract(self.log_cumprod_alpha, t, log_x_start.shape)
+        log_1_min_cumprod_alpha = extract(self.log_1_min_cumprod_alpha, t, log_x_start.shape)
+        log_probs = log_add_exp(
+            log_x_start + log_cumprod_alpha_t,
+            log_1_min_cumprod_alpha - torch.tensor(np.log(self.num_classes_column)).to(log_1_min_cumprod_alpha.device)
+        )
+
+        return log_probs
+
+    def predict_start(self, log_x_t, t, cond_con):
+        x_t = log_x_t
+        out = self._denoise_fn(x_t, t, cond_con)
+        assert out.size(0) == x_t.size(0)
+
+        k=0
+        log_pred = torch.empty_like(out)
+        full_sample=[]
+        for i in range(len(self.num_classes)):
+            out_column = out[:, k:self.num_classes[i]+k]
+            log_pred[:, k:self.num_classes[i]+k] = F.log_softmax(out_column, dim=1)
+            k+=self.num_classes[i]
+
+        return log_pred
+
+
+    def q_posterior(self, log_x_start, log_x_t, t):
+
+        t_minus_1 = t - 1
+        t_minus_1 = torch.where(t_minus_1 < 0, torch.zeros_like(t_minus_1), t_minus_1)
+
+        log_EV_qxtmin_x0 = self.q_pred(log_x_start, t_minus_1)
+
+        num_axes = (1,) * (len(log_x_start.size()) - 1)
+        t_broadcast = t.view(-1, *num_axes) * torch.ones_like(log_x_start)
+        log_EV_qxtmin_x0 = torch.where(t_broadcast == 0,
+                                       log_x_start.to(torch.float64),
+                                       log_EV_qxtmin_x0)
+
+
+        unnormed_logprobs = log_EV_qxtmin_x0 + self.q_pred_one_timestep(log_x_t, t)
+        k=0
+        unnormed_logprobs_column_list=[]
+        for i in range(len(self.num_classes)):
+            unnormed_logprobs_column = unnormed_logprobs[:,k:self.num_classes[i]+k]
+            k+=self.num_classes[i]
+            for j in range(self.num_classes[i]):
+                unnormed_logprobs_column_list.append(torch.logsumexp(unnormed_logprobs_column, dim=1, keepdim=True))
+        unnormed_logprobs_column_ = torch.stack(unnormed_logprobs_column_list,1).squeeze()
+
+        log_EV_xtmin_given_xt_given_xstart = \
+            unnormed_logprobs - unnormed_logprobs_column_
+
+        return log_EV_xtmin_given_xt_given_xstart
+
+    def p_pred(self, log_x, t, cond_con):
+        if self.parametrization == 'x0':
+            log_x_recon = self.predict_start(log_x, t=t, cond_con = cond_con)
+            log_model_pred = self.q_posterior(
+                log_x_start=log_x_recon, log_x_t=log_x, t=t)
+        elif self.parametrization == 'direct':
+            log_model_pred = self.predict_start(log_x, t=t, cond_con = cond_con)
+        else:
+            raise ValueError
+        return log_model_pred, log_x_recon
+
+    @torch.no_grad()
+    def p_sample(self, log_x, t, cond_con):
+        model_log_prob, log_x_recon = self.p_pred(log_x=log_x, t=t, cond_con=cond_con)
+        out = self.log_sample_categorical(model_log_prob).to(log_x.device)
+        return out
+    def log_sample_categorical_with_temp(self, logits, temperature=1.0):
+        full_sample = []
+        k = 0
+        for i in range(len(self.num_classes)):
+            logits_column = logits[:, k:self.num_classes[i]+k] / temperature
+            k += self.num_classes[i]
+
+            uniform = torch.rand_like(logits_column)
+            gumbel_noise = -torch.log(-torch.log(uniform + 1e-30) + 1e-30)
+
+            sample = (gumbel_noise + logits_column).argmax(dim=1)
+
+            col_t = torch.zeros_like(logits_column)
+            col_t.scatter_(1, sample.unsqueeze(1), 1)
+            full_sample.append(col_t)
+
+        full_sample = torch.cat(full_sample, dim=1)
+        log_sample = index_to_log_onehot(full_sample, self.num_classes)
+        return log_sample
+
+    def p_sample_with_temp(self, log_x, t, cond_con, temperature=1.0):
+        model_log_prob, log_x_recon = self.p_pred(log_x=log_x, t=t, cond_con=cond_con)
+        out = self.log_sample_categorical_with_temp(model_log_prob, temperature).to(log_x.device)
+        return out
+
+    def log_sample_categorical(self, logits):
+        full_sample = []
+        k=0
+        for i in range(len(self.num_classes)):
+            logits_column = logits[:,k:self.num_classes[i]+k]
+            k+=self.num_classes[i]
+            uniform = torch.rand_like(logits_column)
+            gumbel_noise = -torch.log(-torch.log(uniform+1e-30)+1e-30)
+            sample = (gumbel_noise + logits_column).argmax(dim=1)
+            col_t =np.zeros(logits_column.shape)
+            col_t[np.arange(logits_column.shape[0]), sample.detach().cpu()] = 1
+            full_sample.append(col_t)
+        full_sample = torch.tensor(np.concatenate(full_sample, axis=1))
+        log_sample = index_to_log_onehot(full_sample, self.num_classes)
+        return log_sample
+
+
+    def q_sample(self, log_x_start, t):
+        log_EV_qxt_x0 = self.q_pred(log_x_start, t)
+        log_sample = self.log_sample_categorical(log_EV_qxt_x0).to(log_EV_qxt_x0.device)
+        return log_sample
+
+
+    def kl_prior(self, log_x_start):
+        b = log_x_start.size(0)
+        device = log_x_start.device
+        ones = torch.ones(b, device=device).long()
+
+        log_qxT_prob = self.q_pred(log_x_start, t=(self.num_timesteps - 1) * ones)
+        log_half_prob = -torch.log(torch.tensor(self.num_classes_column, device=device) * torch.ones_like(log_qxT_prob))
+
+        kl_prior = self.multinomial_kl(log_qxT_prob, log_half_prob).mean(dim=1)
+        return kl_prior
+
+    def compute_Lt(self, log_x_start, log_x_t, t, cond_con, detach_mean=False):
+        log_true_prob = self.q_posterior(
+            log_x_start=log_x_start, log_x_t=log_x_t, t=t)
+        log_model_prob, log_x_recon = self.p_pred(log_x=log_x_t, t=t, cond_con=cond_con)
+
+        if detach_mean:
+            log_model_prob = log_model_prob.detach()
+        kl = self.multinomial_kl(log_true_prob, log_model_prob).mean(dim=1)
+
+        decoder_nll = -self.log_categorical(log_x_start, log_model_prob).mean(dim=1)
+
+        mask = (t == torch.zeros_like(t)).float()
+        loss = mask * decoder_nll + (1. - mask) * kl
+
+        return loss, log_x_recon
+
+    def compute_Lt_esm(self, log_x_start, log_x_t, t, cond_con, detach_mean=False):
+        if len(cond_con.shape) == 3:
+            x_t_con_pooled = cond_con.mean(dim=1)
+        else:
+            x_t_con_pooled = cond_con
+        log_true_prob = self.q_posterior(
+            log_x_start=log_x_start, log_x_t=log_x_t, t=t)
+        log_model_prob, log_x_recon = self.p_pred(log_x=log_x_t, t=t, cond_con=x_t_con_pooled)
+
+        if detach_mean:
+            log_model_prob = log_model_prob.detach()
+        kl = self.multinomial_kl(log_true_prob, log_model_prob).mean(dim=1)
+
+        decoder_nll = -self.log_categorical(log_x_start, log_model_prob).mean(dim=1)
+
+        direct_cls_loss = F.binary_cross_entropy_with_logits(
+            log_x_recon,
+            log_x_start.exp(),
+            reduction='none'
+        ).mean(dim=1)
+
+        mask = (t == torch.zeros_like(t)).float()
+        loss = mask * decoder_nll + (1. - mask) * kl + 0.1 * direct_cls_loss
+
+        return loss, log_x_recon

Computational/package/metalorian/layers.py

@@ -0,0 +1,164 @@
+# coding=utf-8
+# Copyright 2020 The Google Research Authors.
+#
+# 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.
+
+# pylint: skip-file
+import math
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+import numpy as np
+
+def get_act(FLAGS):
+  if FLAGS.activation.lower() == 'elu':
+    return nn.ELU()
+  elif FLAGS.activation.lower() == 'relu':
+    return nn.ReLU()
+  elif FLAGS.activation.lower() == 'lrelu':
+    return nn.LeakyReLU(negative_slope=0.2)
+  elif FLAGS.activation.lower() == 'swish':
+    return nn.SiLU()
+  elif FLAGS.activation.lower() == 'tanh':
+    return nn.Tanh()
+  elif FLAGS.activation.lower() == 'softplus':
+    return nn.Softplus()
+  else:
+    raise NotImplementedError('activation function does not exist!')
+
+def variance_scaling(scale, mode, distribution,
+                     in_axis=1, out_axis=0,
+                     dtype=torch.float32,
+                     device='cpu'):
+  def _compute_fans(shape, in_axis=1, out_axis=0):
+    receptive_field_size = np.prod(shape) / shape[in_axis] / shape[out_axis]
+    fan_in = shape[in_axis] * receptive_field_size
+    fan_out = shape[out_axis] * receptive_field_size
+    return fan_in, fan_out
+
+  def init(shape, dtype=dtype, device=device):
+    fan_in, fan_out = _compute_fans(shape, in_axis, out_axis)
+    if mode == "fan_in":
+      denominator = fan_in
+    elif mode == "fan_out":
+      denominator = fan_out
+    elif mode == "fan_avg":
+      denominator = (fan_in + fan_out) / 2
+    else:
+      raise ValueError(
+        "invalid mode for variance scaling initializer: {}".format(mode))
+    variance = scale / denominator
+    if distribution == "normal":
+      return torch.randn(*shape, dtype=dtype, device=device) * np.sqrt(variance)
+    elif distribution == "uniform":
+      return (torch.rand(*shape, dtype=dtype, device=device) * 2. - 1.) * np.sqrt(3 * variance)
+    else:
+      raise ValueError("invalid distribution for variance scaling initializer")
+
+  return init
+
+def default_init(scale=1.):
+  """DDPM initialization."""
+  scale = 1e-10 if scale == 0 else scale
+  return variance_scaling(scale, 'fan_avg', 'uniform')
+
+def get_timestep_embedding(timesteps, embedding_dim, max_positions=10000):
+  assert len(timesteps.shape) == 1
+  half_dim = embedding_dim // 2
+  emb = math.log(max_positions) / (half_dim - 1)
+  emb = torch.exp(torch.arange(half_dim, dtype=torch.float32, device=timesteps.device) * -emb)
+  emb = timesteps.float()[:, None] * emb[None, :]
+  emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1)
+  if embedding_dim % 2 == 1:
+    emb = F.pad(emb, (0, 1), mode='constant')
+  assert emb.shape == (timesteps.shape[0], embedding_dim)
+  return emb
+
+class Encoder(nn.Module):
+  def __init__(self, encoder_dim, tdim, FLAGS):
+    super(Encoder, self).__init__()
+    self.encoding_blocks = nn.ModuleList()
+    for i in range(len(encoder_dim)):
+      if (i+1)==len(encoder_dim): break
+      encoding_block = EncodingBlock(encoder_dim[i], encoder_dim[i+1], tdim, FLAGS)
+      self.encoding_blocks.append(encoding_block)
+
+  def forward(self, x, t):
+    skip_connections = []
+    for encoding_block in self.encoding_blocks:
+      x, skip_connection = encoding_block(x, t)
+      skip_connections.append(skip_connection)
+    return skip_connections, x
+
+class EncodingBlock(nn.Module):
+  def __init__(self, dim_in, dim_out, tdim, FLAGS):
+    super(EncodingBlock, self).__init__()
+    self.layer1 = nn.Sequential(
+        nn.Linear(dim_in, dim_out),
+        get_act(FLAGS)
+    )
+    self.temb_proj = nn.Sequential(
+        nn.Linear(tdim, dim_out),
+        get_act(FLAGS)
+    )
+    self.layer2 = nn.Sequential(
+        nn.Linear(dim_out, dim_out),
+        get_act(FLAGS)
+    )
+
+  def forward(self, x, t):
+    x = self.layer1(x).clone()
+    x += self.temb_proj(t)
+    x = self.layer2(x)
+    skip_connection = x
+    return x, skip_connection
+
+class Decoder(nn.Module):
+  def __init__(self, decoder_dim, tdim, FLAGS):
+    super(Decoder, self).__init__()
+    self.decoding_blocks = nn.ModuleList()
+    for i in range(len(decoder_dim)):
+      if (i+1)==len(decoder_dim): break
+      decoding_block = DecodingBlock(decoder_dim[i], decoder_dim[i+1], tdim, FLAGS)
+      self.decoding_blocks.append(decoding_block)
+
+  def forward(self, skip_connections, x, t):
+    zipped = zip(reversed(skip_connections), self.decoding_blocks)
+    for skip_connection, decoding_block in zipped:
+      x = decoding_block(skip_connection, x, t)
+    return x
+
+class DecodingBlock(nn.Module):
+  def __init__(self, dim_in, dim_out, tdim, FLAGS):
+    super(DecodingBlock, self).__init__()
+    self.layer1 = nn.Sequential(
+        nn.Linear(dim_in*2, dim_in),
+        get_act(FLAGS)
+    )
+    self.temb_proj = nn.Sequential(
+        nn.Linear(tdim, dim_in),
+        get_act(FLAGS)
+    )
+    self.layer2 = nn.Sequential(
+        nn.Linear(dim_in, dim_out),
+        get_act(FLAGS)
+    )
+
+  def forward(self, skip_connection, x, t):
+
+    x = torch.cat((skip_connection, x), dim=1)
+    x = self.layer1(x).clone()
+    x += self.temb_proj(t)
+    x = self.layer2(x)
+
+    return x

Computational/package/metalorian/objectives.py

@@ -0,0 +1,76 @@
+import torch
+import torch.nn.functional as F
+
+def warmup_lr(step):
+    return min(step, 5000) / 5000
+
+
+def training_with_esm(x_0_con, x_0_dis, trainer, trainer_dis, ns_con, ns_dis, trans, FLAGS, cls_head, class_weights):
+
+    t = torch.randint(FLAGS.T, size=(x_0_con['input_ids'].shape[0], ), device=x_0_con['input_ids'].device)
+    pt = torch.ones_like(t).float() / FLAGS.T
+
+    embeddings = trainer.model.get_embeddings(x_0_con)
+
+    noise = torch.randn_like(embeddings)
+    x_t_con = trainer.make_x_t(embeddings, t, noise)
+    torch.cuda.empty_cache()
+    log_x_start = torch.log(x_0_dis.float().clamp(min=1e-30))
+    x_t_dis = trainer_dis.q_sample(log_x_start=log_x_start, t=t)
+
+    eps = trainer.model(x_t_con, t, x_t_dis.to(x_t_con.device))
+    ps_0_con = trainer.predict_xstart_from_eps(x_t_con, t, eps=eps)
+    con_loss = F.mse_loss(eps, noise, reduction='none')
+    con_loss = con_loss.mean()
+    kl, ps_0_dis = trainer_dis.compute_Lt_esm(log_x_start, x_t_dis, t, x_t_con)
+    ps_0_dis = torch.exp(ps_0_dis)
+    kl_prior = trainer_dis.kl_prior(log_x_start)
+    dis_loss = (kl / pt + kl_prior).mean()
+    torch.cuda.empty_cache()
+    noise_ns = torch.randn_like(ns_con)
+    ns_t_con = trainer.make_x_t(ns_con, t, noise_ns)
+    log_ns_start = torch.log(ns_dis.float().clamp(min=1e-30))
+    ns_t_dis = trainer_dis.q_sample(log_x_start=log_ns_start, t=t)
+    torch.cuda.empty_cache()
+    eps_ns = trainer.model(x_t_con, t, ns_t_dis.to(ns_t_dis.device))
+    ns_0_con = trainer.predict_xstart_from_eps(x_t_con, t, eps=eps_ns)
+    _, ns_0_dis = trainer_dis.compute_Lt_esm(log_x_start, x_t_dis, t, ns_t_con)
+    ns_0_dis = torch.exp(ns_0_dis)
+
+
+    torch.cuda.empty_cache()
+    triplet_loss = torch.nn.TripletMarginLoss(margin=0.5, p=2)
+    triplet_con = triplet_loss(embeddings, ps_0_con, ns_0_con)
+    total_loss = con_loss
+    st=0
+    triplet_dis = []
+    for item in trans.output_info:
+        ed = st + item[0]
+        ps_dis = F.cross_entropy(ps_0_dis[:, st:ed], torch.argmax(x_0_dis[:, st:ed], dim=-1).long(), reduction='none')
+        ns_dis = F.cross_entropy(ns_0_dis[:, st:ed], torch.argmax(x_0_dis[:, st:ed], dim=-1).long(), reduction='none')
+
+        triplet_dis.append(max((ps_dis-ns_dis).mean()+1,0))
+        st = ed
+    triplet_dis = sum(triplet_dis)/len(triplet_dis)
+
+
+    return total_loss, triplet_con, dis_loss, triplet_dis
+
+
+def make_negative_condition_esm(x_0_con, x_0_dis, base_model):
+    device = x_0_dis.device
+
+    with torch.no_grad():
+        _, _, embeddings, _ = base_model(
+            input_ids=x_0_con['input_ids'],
+            attention_masks=x_0_con['attention_mask']
+        )
+
+    B = embeddings.shape[0]
+
+    indices = torch.randperm(B)
+
+    ns_con = embeddings[indices].clone()
+    ns_dis = x_0_dis[indices].clone()
+
+    return ns_con, ns_dis

Computational/package/metalorian/paths.py

@@ -0,0 +1,35 @@
+from pathlib import Path
+import os
+import sys
+
+ROOT = Path(__file__).resolve().parents[1]
+BACKBONE = ROOT/'assets/esm2_t33_650M_UR50D'
+PREDICTOR = ROOT/'checkpoints/metalatte_predictor.ckpt'
+DIFFUSION = ROOT/'checkpoints/metalorian_diffusion.pt'
+TRAIN = ROOT/'data/metalatte/train'
+VALIDATION = ROOT/'data/metalatte/validation'
+os.environ['USE_TF'] = '0'
+os.environ.setdefault('HF_HUB_OFFLINE', '1')
+os.environ.setdefault('WANDB_MODE', 'disabled')
+os.environ.setdefault('MPLCONFIGDIR', str(ROOT/'.cache/matplotlib'))
+sys.dont_write_bytecode = True
+_configured = False
+
+
+def configure():
+    global _configured
+    if _configured:
+        return
+    import torch
+    original = torch.load
+
+    def load(file, *args, **kwargs):
+        candidate = file if isinstance(file, (str, Path)) else getattr(file, 'name', getattr(file, 'path', None))
+        if isinstance(candidate, (str, Path)) and Path(candidate).resolve() in {PREDICTOR.resolve(), DIFFUSION.resolve()}:
+            file = str(Path(candidate).resolve())
+            kwargs.setdefault('weights_only', False)
+            kwargs.setdefault('mmap', True)
+        return original(file, *args, **kwargs)
+
+    torch.load = load
+    _configured = True

Computational/package/metalorian/predictor.py

@@ -0,0 +1,585 @@
+from .paths import BACKBONE
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from transformers import AutoTokenizer
+from torch.optim import AdamW
+import numpy as np
+from sklearn.metrics import f1_score
+from typing import Tuple
+import lightning.pytorch as pl
+import gc
+from torch.optim.lr_scheduler import _LRScheduler
+from sklearn.metrics import f1_score, roc_auc_score, accuracy_score, matthews_corrcoef, precision_score, recall_score
+
+
+class GELU(nn.Module):
+    def forward(self, x):
+        return 0.5 * x * (1.0 + torch.erf(x / math.sqrt(2.0)))
+
+
+class FocalLoss(nn.Module):
+    """Focal loss.
+
+    Sources:
+    https://github.com/fcakyon/balanced-loss/blob/main/balanced_loss/losses.py
+    https://github.com/vandit15/Class-balanced-loss-pytorch/blob/master/class_balanced_loss.py
+    """
+    def __init__(self, gamma=2):
+        super(FocalLoss, self).__init__()
+        self.gamma = gamma
+
+    def forward(self, inputs, targets, sample_weights):
+        BCLoss = F.binary_cross_entropy_with_logits(input = inputs, target = targets,reduction = "none")
+
+        if self.gamma == 0.0:
+            modulator = 1.0
+        else:
+            modulator = torch.exp(-self.gamma * targets * inputs - self.gamma * torch.log(1 +
+                torch.exp(-1.0 * inputs)))
+
+        loss = modulator * BCLoss
+
+        weighted_loss = sample_weights * loss
+        focal_loss = torch.sum(weighted_loss)
+
+        focal_loss /= torch.sum(targets)
+        return focal_loss
+
+
+class F1Loss(nn.Module):
+    def __init__(self, epsilon=1e-7):
+        super().__init__()
+        self.epsilon = epsilon
+
+    def forward(self, y_pred, y_true):
+        assert y_pred.ndim == 2
+        assert y_true.ndim == 2
+        y_true = y_true.float()
+        y_pred = torch.sigmoid(y_pred)
+
+        tp = (y_true * y_pred).sum(dim=0)
+        fp = ((1 - y_true) * y_pred).sum(dim=0)
+        fn = (y_true * (1 - y_pred)).sum(dim=0)
+
+        precision = tp / (tp + fp + self.epsilon)
+        recall = tp / (tp + fn + self.epsilon)
+
+        f1 = 2 * (precision * recall) / (precision + recall + self.epsilon)
+        f1 = f1.clamp(min=self.epsilon, max=1-self.epsilon)
+        return 1 - f1.mean()
+
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(x, cos, sin):
+    # x: (B, L, H, D); cos and sin: (1, L, D).
+    cos = cos.unsqueeze(2)
+    sin = sin.unsqueeze(2)
+    return (x * cos) + (rotate_half(x) * sin)
+
+
+class RotaryEmbedding(torch.nn.Module):
+    """Rotary position embeddings (Su et al.), inspired by GPT-NeoX.
+
+    https://arxiv.org/abs/2104.09864
+    https://github.com/ZhuiyiTechnology/roformer
+    https://github.com/EleutherAI/gpt-neox
+    """
+
+    def __init__(self, dim: int, *_, **__):
+        super().__init__()
+        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+
+        self._seq_len_cached = None
+        self._cos_cached = None
+        self._sin_cached = None
+
+    def _update_cos_sin_tables(self, x, seq_dimension=1):
+        seq_len = x.shape[seq_dimension]
+
+        if seq_len != self._seq_len_cached or self._cos_cached.device != x.device:
+            self._seq_len_cached = seq_len
+            t = torch.arange(x.shape[seq_dimension], device=x.device).type_as(self.inv_freq)
+            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
+
+            self._cos_cached = emb.cos()[None, :, :]
+            self._sin_cached = emb.sin()[None, :, :]
+
+        return self._cos_cached, self._sin_cached
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        self._cos_cached, self._sin_cached = self._update_cos_sin_tables(k)
+
+        return (
+            apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached),
+            apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached),
+        )
+
+
+def macro_f1(y_true, y_pred, thresholds):
+    y_pred_binary = (y_pred >= thresholds).float()
+    tp = (y_true * y_pred_binary).sum(dim=0)
+    fp = ((1 - y_true) * y_pred_binary).sum(dim=0)
+    fn = (y_true * (1 - y_pred_binary)).sum(dim=0)
+    precision = tp / (tp + fp + 1e-7)
+    recall = tp / (tp + fn + 1e-7)
+    f1 = 2 * precision * recall / (precision + recall + 1e-7)
+    macro_f1 = f1.mean()
+    return macro_f1
+
+
+def safeguard_softmax(logits, dim=-1):
+    max_logits, _ = logits.max(dim=dim, keepdim=True)
+    exp_logits = torch.exp(logits - max_logits)
+    exp_sum = exp_logits.sum(dim=dim, keepdim=True)
+    probs = exp_logits / (exp_sum + 1e-7)
+    return probs
+
+
+class PositionalAttentionHead(nn.Module):
+    def __init__(self, hidden_dim, n_heads):
+        super(PositionalAttentionHead, self).__init__()
+        self.n_heads = n_heads
+        self.hidden_dim = hidden_dim
+        self.head_dim = hidden_dim // n_heads
+        self.preattn_ln = nn.LayerNorm(self.head_dim)
+        self.Q = nn.Linear(self.head_dim, self.head_dim, bias=False)
+        self.K = nn.Linear(self.head_dim, self.head_dim, bias=False)
+        self.V = nn.Linear(self.head_dim, self.head_dim, bias=False)
+        self.rot_emb = RotaryEmbedding(self.head_dim)
+
+    def forward(self, x, attention_mask):
+        batch_size, seq_len, _ = x.size()
+        x = x.view(batch_size, seq_len, self.n_heads, self.head_dim)
+        x = self.preattn_ln(x)
+
+        q = self.Q(x)
+        k = self.K(x)
+        v = self.V(x)
+
+        q, k = self.rot_emb(q, k)
+        gc.collect()
+        torch.cuda.empty_cache()
+
+        attn_scores = torch.einsum('bqhd,bkhd->bhqk', q, k) / math.sqrt(self.head_dim)
+        attn_scores = attn_scores.masked_fill(torch.logical_not(attention_mask.unsqueeze(1).unsqueeze(1)), float("-inf"))
+
+        attn_probs = safeguard_softmax(attn_scores, dim=-1)
+
+        x = torch.einsum('bhqk,bkhd->bqhd', attn_probs, v)
+        x = x.reshape(batch_size, seq_len, self.hidden_dim)
+        gc.collect()
+        torch.cuda.empty_cache()
+        return x, attn_probs
+
+
+class CosineAnnealingWithWarmup(_LRScheduler):
+    def __init__(self, optimizer, warmup_steps, total_steps, eta_ratio=0.1, last_epoch=-1):
+        self.warmup_steps = warmup_steps
+        self.total_steps = total_steps
+        self.eta_ratio = eta_ratio
+        super(CosineAnnealingWithWarmup, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        if self.last_epoch < self.warmup_steps:
+            return [base_lr * self.last_epoch / self.warmup_steps for base_lr in self.base_lrs]
+
+        progress = (self.last_epoch - self.warmup_steps) / (self.total_steps - self.warmup_steps)
+        cosine_decay = 0.5 * (1 + np.cos(np.pi * progress))
+        decayed_lr = (1 - self.eta_ratio) * cosine_decay + self.eta_ratio
+
+        return [decayed_lr * base_lr for base_lr in self.base_lrs]
+
+
+class AdaptiveMarginTripletLoss(nn.Module):
+    def __init__(self, initial_margin=0.05, scale_factor=0.1, min_margin=0.05, max_margin=2.0):
+        super(AdaptiveMarginTripletLoss, self).__init__()
+        self.initial_margin = initial_margin
+        self.scale_factor = scale_factor
+        self.min_margin = min_margin
+        self.max_margin = max_margin
+
+    def forward(self, anchor, positive, negative):
+        pos_dist = F.pairwise_distance(anchor, positive, p=2)
+        neg_dist = F.pairwise_distance(anchor, negative, p=2)
+
+        adaptive_margins = self.initial_margin + self.scale_factor * (neg_dist - pos_dist)
+
+        adaptive_margins = torch.clamp(adaptive_margins, min=self.min_margin, max=self.max_margin)
+
+        losses = F.relu(pos_dist - neg_dist + adaptive_margins)
+
+        return losses.mean()
+
+
+class RobertaLMHead(nn.Module):
+    def __init__(self, embed_dim, output_dim, weight):
+        super().__init__()
+        self.dense = nn.Linear(embed_dim, embed_dim)
+        self.layer_norm = nn.LayerNorm(embed_dim)
+        self.weight = weight
+        self.gelu = GELU()
+        self.bias = nn.Parameter(torch.zeros(output_dim))
+    def forward(self, features):
+        x = self.dense(features)
+        x = self.gelu(x)
+        x = self.layer_norm(x)
+        x = F.linear(x, self.weight) + self.bias
+        return x
+
+
+def pad_to_max_length(tensor, max_length, pad_value=1):
+    return F.pad(tensor, (0, max_length - tensor.shape[1]), value=pad_value)
+
+
+def compute_ece(probs, labels, n_bins=10):
+    """Per-class calibration error for probabilities and labels of shape (B, C)."""
+    bin_boundaries = torch.linspace(0, 1, n_bins + 1)
+    bin_lowers = bin_boundaries[:-1]
+    bin_uppers = bin_boundaries[1:]
+
+    n_classes = probs.size(1)
+    ece_per_class = torch.zeros(n_classes, device=probs.device)
+
+    for class_idx in range(n_classes):
+        class_probs = probs[:, class_idx]
+        class_labels = labels[:, class_idx]
+
+        ece = torch.zeros(1, device=probs.device)
+        for bin_lower, bin_upper in zip(bin_lowers, bin_uppers):
+            in_bin = (class_probs.gt(bin_lower.item()) * class_probs.le(bin_upper.item())).float()
+            prop_in_bin = in_bin.mean()
+
+            if prop_in_bin > 0:
+                accuracy_in_bin = (class_labels * in_bin).sum() / in_bin.sum()
+                avg_confidence_in_bin = (class_probs * in_bin).sum() / in_bin.sum()
+                ece += torch.abs(avg_confidence_in_bin - accuracy_in_bin) * prop_in_bin
+
+        ece_per_class[class_idx] = ece.item()
+
+    return ece_per_class
+
+
+class MultitaskProteinModel(pl.LightningModule):
+    def __init__(self, hparams, esm_model, num_layers_to_finetune, input_dim=1280, output_dim=12, class_weights=None, hidden_dim=512, n_heads=10):
+        super(MultitaskProteinModel, self).__init__()
+        self.save_hyperparameters(hparams)
+        self.esm_model = esm_model
+        for param in self.esm_model.parameters():
+            param.requires_grad = False
+        for i in range(num_layers_to_finetune):
+            for param in self.esm_model.encoder.layer[-i-1].parameters():
+                param.requires_grad = True
+
+        self.lm_head = RobertaLMHead(embed_dim = 1280, output_dim=33, weight=self.esm_model.embeddings.word_embeddings.weight)
+        self.attn_head = PositionalAttentionHead(input_dim, n_heads)
+        self.attn_ln = nn.LayerNorm(input_dim)
+        self.attn_skip = nn.Linear(input_dim, input_dim)
+        self.linear_layers = nn.ModuleList()
+        for _ in range(self.hparams.num_linear_layers):
+            self.linear_layers.append(nn.Linear(input_dim, input_dim))
+        self.reduction_layers = nn.Sequential(
+            nn.Linear(input_dim, hidden_dim),
+            GELU(),
+            nn.Linear(hidden_dim, output_dim)
+        )
+        self.clf_ln = nn.LayerNorm(input_dim)
+        self.class_weights = class_weights
+        self.classification_thresholds = nn.Parameter(torch.tensor([0.5]*output_dim))
+        self.historical_memory = 0.9
+        self.training_step_outputs = []
+        self.training_step_true_labels = []
+        self.testing_step_predictions = []
+        self.testing_step_labels = []
+        self.test_dataloader = None
+        self.tokenizer = AutoTokenizer.from_pretrained(str(BACKBONE), local_files_only=True)
+        self.corruption_epoch_threshold = 5
+
+    def setup(self, stage=None):
+        self.__build_loss()
+
+    def __build_loss(self):
+        self.contrastive_loss = AdaptiveMarginTripletLoss()
+        focal_gamma = torch.tensor(self.hparams.focal_gamma).to(self.device)
+        self.class_weights = self.class_weights.to(self.device)
+        self.focal_loss = FocalLoss(gamma=focal_gamma)
+        self.f1_loss = F1Loss()
+        self.reconstruction_loss = nn.CrossEntropyLoss(ignore_index=self.tokenizer.pad_token_id)
+
+    def forward(self, input_ids, attention_masks=None):
+        outputs = self.esm_model(input_ids=input_ids, attention_mask=attention_masks, output_hidden_states=True)
+        embeddings = outputs.last_hidden_state
+        attention_masks = attention_masks
+
+        x_pool, x_attns = self.attn_head(embeddings, attention_masks)
+        x_pool = self.attn_ln(x_pool + self.attn_skip(x_pool))
+
+        for linear_layer in self.linear_layers:
+            residue = x_pool
+            x_pool = linear_layer(x_pool)
+            x_pool = F.silu(x_pool)
+            x_pool = x_pool + residue
+
+        x_weighted = torch.einsum('bhlk,bld->bhld', x_attns, x_pool)
+        x_combined = x_weighted.mean(dim=1)
+        x_combined = self.clf_ln(x_combined)
+        mlm_logits = self.lm_head(x_combined)
+        attention_masks = attention_masks.unsqueeze(-1).float()
+        attention_sum = attention_masks.sum(dim=1, keepdim=True)
+        x_combined_masked = (x_combined * attention_masks).sum(dim=1) / attention_sum.squeeze(1)
+
+        x_pred = self.reduction_layers(x_combined_masked)
+        gc.collect()
+        torch.cuda.empty_cache()
+        return x_pred, x_attns, x_combined_masked, mlm_logits
+    def forward_from_embeds(self, inputs_embeds, attention_masks=None):
+        esm = self.esm_model
+        # Disable token dropout for embedding input.
+        if getattr(esm.embeddings, "token_dropout", False):
+            esm.embeddings.token_dropout = False
+
+        outs = esm(
+            inputs_embeds=inputs_embeds,
+            attention_mask=attention_masks.to(torch.long),
+            output_hidden_states=True,
+        )
+        embeddings = outs.last_hidden_state
+
+        x_pool, x_attns = self.attn_head(embeddings, attention_masks)
+        x_pool = self.attn_ln(x_pool + self.attn_skip(x_pool))
+        for linear_layer in self.linear_layers:
+            residue = x_pool
+            x_pool = F.silu(linear_layer(x_pool))
+            x_pool = x_pool + residue
+
+        x_weighted = torch.einsum('bhlk,bld->bhld', x_attns, x_pool).mean(dim=1)
+        x_weighted = self.clf_ln(x_weighted)
+
+        mask_f = attention_masks.unsqueeze(-1).float()
+        x_combined_masked = (x_weighted * mask_f).sum(dim=1) / mask_f.sum(dim=1).clamp_min(1e-6)
+
+        x_pred = self.reduction_layers(x_combined_masked)
+        return x_pred, x_attns, x_combined_masked, None
+    def predict(self, input_ids, attention_masks=None):
+        x_pred, _, _, _ = self.forward(input_ids, attention_masks)
+        classification_output = torch.sigmoid(x_pred)
+        predictions = (classification_output >= self.classification_thresholds).float()
+
+        for i, pred in enumerate(predictions):
+            if pred.sum() == 0:
+                weighted_probs = classification_output[i]
+                max_class = torch.argmax(weighted_probs)
+                predictions[i, max_class] = 1.0
+
+        confidence_scores = torch.abs(classification_output - self.classification_thresholds) / self.classification_thresholds
+        max_confidence_scores = confidence_scores.max(dim=1).values
+
+        return classification_output, predictions, max_confidence_scores
+
+    def training_step(self, batch, batch_idx):
+        anchor_input_ids, anchor_attention_masks = batch['anchor_input_ids'], batch['anchor_attention_mask']
+        positive_input_ids, positive_attention_masks = batch['positive_input_ids'], batch['positive_attention_mask']
+        negative_input_ids, negative_attention_masks = batch['negative_input_ids'], batch['negative_attention_mask']
+        anchor_labels = batch['anchor_labels']
+        positive_labels = batch['positive_labels']
+        negative_labels = batch['negative_labels']
+
+        anchor_x_pred, _, anchor_x_combined, mlm_logits = self.forward(anchor_input_ids, anchor_attention_masks)
+        positive_x_pred, _, positive_x_combined, _ = self.forward(positive_input_ids, positive_attention_masks)
+        negative_x_pred, _, negative_x_combined, _ = self.forward(negative_input_ids, negative_attention_masks)
+
+        reconstruction_loss = self.reconstruction_loss(mlm_logits.view(-1, mlm_logits.size(-1)), anchor_input_ids.view(-1))
+
+        anchor_focal_loss = self.focal_loss(anchor_x_pred, anchor_labels, self.class_weights)
+        positive_focal_loss = self.focal_loss(positive_x_pred, positive_labels, self.class_weights)
+
+        semi_hard_negative_ce_loss = F.binary_cross_entropy_with_logits(input=negative_x_pred, target=negative_labels)
+
+        focal_loss = (anchor_focal_loss + positive_focal_loss) / 2
+
+        anchor_f1_loss = self.f1_loss(anchor_x_pred, anchor_labels)
+        positive_f1_loss = self.f1_loss(positive_x_pred, positive_labels)
+        semi_hard_f1_loss = self.f1_loss(negative_x_pred, negative_labels)
+
+        semi_hard_triplet_loss  = self.contrastive_loss(anchor_x_combined, positive_x_combined, negative_x_combined)
+
+        if self.current_epoch > self.corruption_epoch_threshold:
+            hard_negative_input_ids, hard_negative_attention_masks = batch['hard_negative_input_ids'], batch['hard_negative_attention_mask']
+            hard_negative_labels = batch['hard_negative_labels']
+            if self.current_epoch < 15:
+                epoch_ratio = min(1.0, (self.current_epoch - self.corruption_epoch_threshold) / (self.hparams.epochs - self.corruption_epoch_threshold))
+                hard_weight = epoch_ratio
+                semi_hard_weight = 1 - hard_weight
+            else:
+                semi_hard_weight = 0.0
+                hard_weight = 1.0
+            hard_negative_x_pred, _, hard_negative_x_combined, _ = self.forward(hard_negative_input_ids, hard_negative_attention_masks)
+            hard_negative_ce_loss = F.binary_cross_entropy_with_logits(input=hard_negative_x_pred, target=hard_negative_labels)
+            hard_triplet_loss = self.contrastive_loss(anchor_x_combined, positive_x_combined, hard_negative_x_combined)
+            hard_f1_loss = self.f1_loss(hard_negative_x_pred, hard_negative_labels)
+            negative_ce_loss = semi_hard_weight * semi_hard_negative_ce_loss + hard_weight * hard_negative_ce_loss
+            triplet_loss = semi_hard_weight * semi_hard_triplet_loss + hard_weight * hard_triplet_loss
+            f1_loss = anchor_f1_loss + positive_f1_loss + semi_hard_weight * semi_hard_triplet_loss + hard_weight * hard_f1_loss
+        else:
+            negative_ce_loss = semi_hard_negative_ce_loss
+            triplet_loss = semi_hard_triplet_loss
+            f1_loss = anchor_f1_loss + positive_f1_loss + semi_hard_f1_loss
+
+        loss = focal_loss * 2 + negative_ce_loss + f1_loss + triplet_loss * 2 + reconstruction_loss
+
+        self.training_step_outputs.append(anchor_x_pred)
+        self.training_step_outputs.append(positive_x_pred)
+        self.training_step_outputs.append(negative_x_pred)
+        self.training_step_true_labels.append(anchor_labels)
+        self.training_step_true_labels.append(positive_labels)
+        self.training_step_true_labels.append(negative_labels)
+        macro_f1_val = macro_f1(anchor_labels, torch.sigmoid(anchor_x_pred), self.classification_thresholds)
+        self.log('train_macro_f1', macro_f1_val, on_step=True, on_epoch=True, prog_bar=True, logger=True, sync_dist=True)
+        gc.collect()
+        torch.cuda.empty_cache()
+        gc.collect()
+        torch.cuda.empty_cache()
+        self.log('train_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True, sync_dist=True)
+        return loss
+
+    def on_train_epoch_end(self):
+        epoch_outputs = torch.cat(self.training_step_outputs, dim=0)
+        epoch_true_labels = torch.cat(self.training_step_true_labels, dim=0)
+        epoch_outputs = torch.sigmoid(epoch_outputs)
+        self.update_class_thresholds(epoch_outputs, epoch_true_labels)
+        threshold_dict = {f"threshold_{i}": self.classification_thresholds[i].item() for i in range(15)}
+        self.log_dict(threshold_dict, on_epoch=True)
+        self.training_step_outputs.clear()
+        self.training_step_true_labels.clear()
+        gc.collect()
+        torch.cuda.empty_cache()
+
+    def validation_step(self, batch, batch_idx):
+        anchor_input_ids, anchor_attention_masks = batch['anchor_input_ids'], batch['anchor_attention_mask']
+        positive_input_ids, positive_attention_masks = batch['positive_input_ids'], batch['positive_attention_mask']
+        negative_input_ids, negative_attention_masks = batch['negative_input_ids'], batch['negative_attention_mask']
+        anchor_labels = batch['anchor_labels']
+        positive_labels = batch['positive_labels']
+        negative_labels = batch['negative_labels']
+
+        anchor_x_pred, _, anchor_x_combined, mlm_logits = self.forward(anchor_input_ids, anchor_attention_masks)
+        positive_x_pred, _, positive_x_combined, _ = self.forward(positive_input_ids, positive_attention_masks)
+        negative_x_pred, _, negative_x_combined, _ = self.forward(negative_input_ids, negative_attention_masks)
+
+        reconstruction_loss = self.reconstruction_loss(mlm_logits.view(-1, mlm_logits.size(-1)), anchor_input_ids.view(-1))
+        anchor_focal_loss = self.focal_loss(anchor_x_pred, anchor_labels , self.class_weights)
+        positive_focal_loss = self.focal_loss(positive_x_pred, positive_labels, self.class_weights)
+
+        negative_ce_loss = F.binary_cross_entropy_with_logits(input=negative_x_pred, target=negative_labels)
+
+        focal_loss = (anchor_focal_loss + positive_focal_loss) / 2
+
+        f1_loss = self.f1_loss(anchor_x_pred, anchor_labels)
+
+        triplet_loss = self.contrastive_loss(anchor_x_combined, positive_x_combined, negative_x_combined)
+
+        loss = focal_loss * 2 + negative_ce_loss + f1_loss + triplet_loss * 2 + reconstruction_loss
+
+        macro_f1_val = macro_f1(anchor_labels, torch.sigmoid(anchor_x_pred), self.classification_thresholds)
+        self.log('val_macro_f1', macro_f1_val, on_step=True, on_epoch=True, prog_bar=True, logger=True, sync_dist=True)
+        gc.collect()
+        torch.cuda.empty_cache()
+        self.log('val_loss', loss, on_step=True, on_epoch=True, prog_bar=True, logger=True, sync_dist=True)
+
+
+    def test_step(self, batch, batch_idx):
+        input_ids, labels, attention_masks = batch['input_ids'], batch['labels'], batch['attention_masks']
+        x_pred, _, _ = self.forward_test(input_ids, attention_masks)
+        predictions = torch.sigmoid(x_pred)
+        self.testing_step_predictions.append(predictions)
+        self.testing_step_labels.append(labels)
+
+    def on_test_epoch_end(self,):
+        predictions = torch.cat(self.testing_step_predictions, dim=0)
+        labels = torch.cat(self.testing_step_labels, dim=0)
+
+        predictions_binary = (predictions >= self.classification_thresholds).float().cpu().numpy()
+        labels_binary = labels.cpu().numpy()
+
+        f1 = f1_score(labels_binary, predictions_binary, average='macro')
+        roc_auc = roc_auc_score(labels_binary, predictions.cpu().to(torch.float32).numpy(), average='micro')
+        accuracy = accuracy_score(labels_binary.flatten(), predictions_binary.flatten())
+        mcc = matthews_corrcoef(labels_binary.flatten(), predictions_binary.flatten())
+        precision = precision_score(labels_binary, predictions_binary, average='macro')
+        recall = recall_score(labels_binary, predictions_binary, average='macro')
+
+        results = {
+            'test_F1_macro': f1,
+            'test_AUCROC': roc_auc,
+            'test_Accuracy': accuracy,
+            'test_MCC': mcc,
+            'test_Precision': precision,
+            'test_Recall': recall,
+        }
+        self.log_dict(results)
+        self.testing_step_predictions.clear()
+        self.testing_step_labels.clear()
+        gc.collect()
+        torch.cuda.empty_cache()
+
+
+    def configure_optimizers(self):
+        optimizer = AdamW(self.parameters(), lr=self.hparams.lr, betas=(0.9, 0.95), weight_decay=1e-5)
+        schedulers = CosineAnnealingWithWarmup(optimizer, warmup_steps=100, eta_ratio=0.1, total_steps=9000)
+        lr_schedulers = {
+                "scheduler": schedulers,
+                "name": 'learning_rate_logs',
+                "interval": 'step',
+                'frequency': 1
+        }
+        return [optimizer], [lr_schedulers]
+
+
+    def update_class_thresholds(self, inputs, targets):
+        with torch.no_grad():
+            min_threshold_value = 0.001
+            thresholds = torch.linspace(0.0, 1.0, steps=1000, device=inputs.device)
+            num_classes = targets.size(1)
+
+            best_f1_scores = torch.zeros(num_classes, device=inputs.device)
+            best_thresholds = torch.full((num_classes,), min_threshold_value, device=inputs.device)
+
+            for threshold in thresholds:
+                binary_predictions = (inputs >= threshold).float()
+
+                tp = (binary_predictions * targets).sum(dim=0)
+                fp = (binary_predictions * (1 - targets)).sum(dim=0)
+                fn = ((1 - binary_predictions) * targets).sum(dim=0)
+
+                precision = tp / (tp + fp + 1e-7)
+                recall = tp / (tp + fn + 1e-7)
+                f1_scores = 2 * precision * recall / (precision + recall + 1e-7)
+
+                is_better = f1_scores > best_f1_scores
+                best_f1_scores = torch.where(is_better, f1_scores, best_f1_scores)
+                best_thresholds = torch.where(is_better, threshold.expand_as(best_f1_scores), best_thresholds)
+
+            updated_thresholds  = self.historical_memory * self.classification_thresholds.data + (1 - self.historical_memory) * best_thresholds
+
+            self.classification_thresholds = nn.Parameter(torch.clamp(updated_thresholds, min=min_threshold_value))
+            gc.collect()
+            torch.cuda.empty_cache()
+
+
+class ResidueEncoder(MultitaskProteinModel):
+    """Return per-residue clf_ln outputs for diffusion training."""
+    def forward(self, input_ids, attention_masks=None):
+        captured = []
+        hook = self.clf_ln.register_forward_hook(lambda module, inputs, output: captured.append(output))
+        try:
+            logits, attention, pooled, mlm_logits = super().forward(input_ids, attention_masks)
+        finally:
+            hook.remove()
+        return logits, attention, captured[0], mlm_logits

Computational/package/metalorian/sampling.py

@@ -0,0 +1,173 @@
+from .paths import ROOT, BACKBONE, PREDICTOR, DIFFUSION, TRAIN, configure
+configure()
+import argparse
+import numpy as np
+import torch
+from absl import flags
+from transformers import AutoTokenizer, EsmModel
+from datasets import load_from_disk
+from .continuous import GaussianDiffusionSampler, DiffusionProteinModel
+from .discrete import MultinomialDiffusion
+from .unet import tabularUnet
+from .predictor import MultitaskProteinModel
+FLAGS = flags.FLAGS
+METAL_LABELS = ['Ag','Cd','Co','Cu','Fe','Hg','Mn','Mo','Ni','Pb','Pt','V','W','Zn','NA']
+TOKEN_DROPOUT_SCALE = 0.88
+CODI_CKPT, PRED_CKPT, TRAIN_DS, ESM_NAME = map(str, (DIFFUSION,PREDICTOR,TRAIN,BACKBONE))
+
+def build_hparams():
+    ns = argparse.Namespace()
+    ns.date = 'test'
+    ns.attempts = '1'
+    ns.output_model_path = str(ROOT/'runs')
+    ns.num_workers = 8
+    ns.num_linear_layers = 3
+    ns.batch_size = 5
+    ns.devices = 1
+    ns.special_prefix = None
+    ns.focal_gamma = 1
+    ns.lr = 1e-3
+    ns.eps = 0.4
+    ns.n_splits = 2
+    ns.epochs = 100
+    ns.save_top_k = 3
+    return ns
+
+
+def define_flags():
+    if not FLAGS.is_parsed():
+        FLAGS([""])
+    d = {
+        'encoder_dim_con': ('string', "64,128,256"), 'encoder_dim_dis': ('string', "64,128,256"),
+        'nf_con': ('int', 16), 'nf_dis': ('int', 64),
+        'input_size': ('int', None), 'cond_size': ('int', None), 'output_size': ('int', None),
+        'activation': ('string', 'relu'), 'nf': ('int', None),
+        'hidden_size': ('int', 768), 'num_layers': ('int', 12), 'num_heads': ('int', 12),
+        'intermediate_size': ('int', 4 * 768), 'max_seq_length': ('int', 1280),
+        'T': ('int', 50), 'beta_1': ('float', 0.00001), 'beta_T': ('float', 0.02),
+    }
+    for name, (kind, default) in d.items():
+        if not hasattr(FLAGS, name):
+            {'string': flags.DEFINE_string, 'int': flags.DEFINE_integer,
+             'float': flags.DEFINE_float}[kind](name, default, help=name)
+    if not hasattr(FLAGS, 'encoder_dim'):
+        flags.DEFINE_multi_integer('encoder_dim', None, help='encoder_dim')
+    if not hasattr(FLAGS, 'mean_type'):
+        flags.DEFINE_enum('mean_type', 'epsilon', ['xprev', 'xstart', 'epsilon'], help='mean type')
+    if not hasattr(FLAGS, 'var_type'):
+        flags.DEFINE_enum('var_type', 'fixedsmall', ['fixedlarge', 'fixedsmall'], help='var type')
+
+
+def load_diffusion_models(checkpoint_path, base_model, device):
+    ckpt = torch.load(checkpoint_path, map_location=device)
+
+    model_con_state = ckpt['model_con']
+    base_keys = [k for k in model_con_state if 'base_model' in k]
+    if base_keys:
+        base_state = {k.replace('base_model.', ''): model_con_state[k] for k in base_keys}
+        base_model.load_state_dict(base_state, strict=False)
+        print("Loaded base model weights from CoDi checkpoint")
+
+    def strip(sd):
+        return {k.replace('module.', ''): v for k, v in sd.items()}
+    if 'module.' in next(iter(ckpt['model_con'])):
+        ckpt['model_con'] = strip(ckpt['model_con'])
+    if 'module.' in next(iter(ckpt['model_dis'])):
+        ckpt['model_dis'] = strip(ckpt['model_dis'])
+
+    FLAGS.input_size, FLAGS.cond_size, FLAGS.output_size = 1280, 15, 1280
+    FLAGS.encoder_dim = [int(x) for x in FLAGS.encoder_dim_con.split(',')]
+    FLAGS.nf = FLAGS.nf_con
+    FLAGS.hidden_size, FLAGS.num_layers, FLAGS.num_heads = 768, 12, 12
+    FLAGS.intermediate_size, FLAGS.max_seq_length = 4 * 768, 1280
+
+    model_con = DiffusionProteinModel(base_model, FLAGS).to(device)
+    filtered = {}
+    for k, v in ckpt['model_con'].items():
+        if not k.startswith('base_model.'):
+            filtered[k] = v
+        else:
+            bk = k.replace('base_model.', '')
+            if bk in base_model.state_dict() and v.shape == base_model.state_dict()[bk].shape:
+                filtered[k] = v
+    model_con.load_state_dict(filtered, strict=False)
+
+    FLAGS.input_size, FLAGS.cond_size, FLAGS.output_size = 15, 1280, 15
+    FLAGS.encoder_dim = [int(x) for x in FLAGS.encoder_dim_dis.split(',')]
+    FLAGS.nf = FLAGS.nf_dis
+    model_dis = tabularUnet(FLAGS).to(device)
+    model_dis.load_state_dict(ckpt['model_dis'])
+
+    trainer_dis = MultinomialDiffusion(
+        num_classes=np.array([15]), shape=torch.Size([10, 15]), denoise_fn=model_dis,
+        FLAGS=FLAGS, timesteps=FLAGS.T, loss_type='vb_stochastic').to(device)
+    net_sampler = GaussianDiffusionSampler(
+        model_con, FLAGS.beta_1, FLAGS.beta_T, FLAGS.T,
+        mean_type=FLAGS.mean_type, var_type=FLAGS.var_type).to(device)
+    return model_con, model_dis, trainer_dis, net_sampler
+
+
+class Sampler:
+    def __init__(self, cfg, device='cuda:0'):
+        self.cfg = cfg
+        self.device = torch.device(device)
+        define_flags()
+
+        self.tokenizer = AutoTokenizer.from_pretrained(ESM_NAME)
+        self.cys_id = self.tokenizer.convert_tokens_to_ids('C')
+        self.cls_id = self.tokenizer.cls_token_id
+        self.eos_id = self.tokenizer.eos_token_id
+
+        ds = load_from_disk(TRAIN_DS)
+        labels = np.stack([l for b in ds['anchor_labels'] for l in b])
+        cw = 1.0 / (labels.sum(axis=0) + 1e-6)
+        cw = torch.tensor(cw / cw.sum() * labels.shape[1], dtype=torch.float32).to(self.device)
+
+        hp = build_hparams()
+        esm = EsmModel.from_pretrained(ESM_NAME).cpu()
+        base_model = MultitaskProteinModel.load_from_checkpoint(
+            PRED_CKPT, hparams=hp, esm_model=esm, num_layers_to_finetune=2,
+            input_dim=1280, output_dim=15, class_weights=cw).cpu()
+
+        esm2 = EsmModel.from_pretrained(ESM_NAME).cpu()
+        self.predictor = MultitaskProteinModel.load_from_checkpoint(
+            PRED_CKPT, hparams=hp, esm_model=esm2, num_layers_to_finetune=2,
+            input_dim=1280, output_dim=15, class_weights=cw).to(self.device).eval()
+
+        self.model_con, self.model_dis, self.trainer_dis, self.net_sampler = \
+            load_diffusion_models(CODI_CKPT, base_model, self.device)
+        self.model_con.eval()
+        self.model_dis.eval()
+
+        # Gradients are needed only with respect to the latent.
+        for m in (self.predictor, self.model_con, self.model_dis):
+            for p in m.parameters():
+                p.requires_grad_(False)
+
+        self.W = self.predictor.esm_model.embeddings.word_embeddings.weight
+        self.threshold = self.predictor.classification_thresholds[cfg.target_idx].item()
+
+        self.vocab_mask = torch.full((33,), float('-inf'), device=self.device)
+        self.vocab_mask[4:24] = 0.0
+        self.aa_bias = torch.zeros(33, device=self.device)
+        self.aa_bias[self.cys_id] = cfg.cys_logit_bias
+
+        self.others = [i for i in range(15) if i != cfg.target_idx]
+        print(f"target={METAL_LABELS[cfg.target_idx]} idx={cfg.target_idx} "
+              f"threshold={self.threshold:.4f} cys_token={self.cys_id}")
+
+    @torch.no_grad()
+    def hard_decode(self, x, attn_len):
+        logits = self.model_con.decode_embeddings(x[:, :attn_len])
+        biased = logits + self.vocab_mask[None, None, :] + self.aa_bias[None, None, :]
+        toks = torch.argmax(biased, dim=-1)
+        return [self.tokenizer.decode(t, skip_special_tokens=True).replace(' ', '')
+                for t in toks]
+
+    @torch.no_grad()
+    def score(self, seqs):
+        enc = self.tokenizer(seqs, return_tensors='pt', padding=True)
+        enc = {k: v.to(self.device) for k, v in enc.items()}
+        probs, fires, _ = self.predictor.predict(input_ids=enc['input_ids'],
+                                                 attention_masks=enc['attention_mask'])
+        return probs, fires

Computational/package/metalorian/unet.py

@@ -0,0 +1,93 @@
+# coding=utf-8
+# Copyright 2020 The Google Research Authors.
+#
+# 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.
+
+# pylint: skip-file
+
+
+from . import layers
+import torch
+import torch.nn as nn
+
+get_act = layers.get_act
+default_initializer = layers.default_init
+
+class tabularUnet(nn.Module):
+  def __init__(self, FLAGS):
+    super().__init__()
+
+    self.embed_dim = FLAGS.nf
+    tdim = self.embed_dim*4
+    self.act = get_act(FLAGS)
+
+    modules = []
+    modules.append(nn.Linear(self.embed_dim, tdim))
+    modules[-1].weight.data = default_initializer()(modules[-1].weight.shape)
+    nn.init.zeros_(modules[-1].bias)
+    modules.append(nn.Linear(tdim, tdim))
+    modules[-1].weight.data = default_initializer()(modules[-1].weight.shape)
+    nn.init.zeros_(modules[-1].bias)
+
+    cond = FLAGS.cond_size
+    cond_out = FLAGS.input_size
+    if cond_out < 2:
+      cond_out = FLAGS.input_size
+    modules.append(nn.Linear(cond, cond_out))
+    modules[-1].weight.data = default_initializer()(modules[-1].weight.shape)
+    nn.init.zeros_(modules[-1].bias)
+
+    self.all_modules = nn.ModuleList(modules)
+
+    dim_in = FLAGS.input_size + cond_out
+    dim_out = list(FLAGS.encoder_dim)[0]
+    self.inputs = nn.Linear(dim_in, dim_out)
+
+    self.encoder = layers.Encoder(list(FLAGS.encoder_dim), tdim, FLAGS)
+
+    dim_in = list(FLAGS.encoder_dim)[-1]
+    dim_out = list(FLAGS.encoder_dim)[-1]
+    self.bottom_block = nn.Linear(dim_in, dim_out)
+
+    self.decoder = layers.Decoder(list(reversed(FLAGS.encoder_dim)), tdim, FLAGS)
+
+    dim_in = list(FLAGS.encoder_dim)[0]
+    dim_out = FLAGS.output_size
+    self.outputs = nn.Linear(dim_in, dim_out)
+
+
+  def forward(self, x, time_cond, cond):
+    if len(x.shape) == 3:
+        x = x.mean(dim=1)
+    if len(cond.shape) == 3:
+        cond = cond.mean(dim=1)
+    modules = self.all_modules
+    m_idx = 0
+
+    temb = layers.get_timestep_embedding(time_cond, self.embed_dim)
+    temb = modules[m_idx](temb)
+    m_idx += 1
+    temb= self.act(temb)
+    temb = modules[m_idx](temb)
+    m_idx += 1
+    cond = modules[m_idx](cond)
+    m_idx += 1
+    x = torch.cat([x, cond], dim=1).float()
+    inputs = self.inputs(x)
+    skip_connections, encoding = self.encoder(inputs, temb)
+    encoding = self.bottom_block(encoding)
+    encoding = self.act(encoding)
+    x = self.decoder(skip_connections, encoding, temb)
+    outputs = self.outputs(x)
+
+    return outputs

Computational/package/requirements.txt

@@ -0,0 +1,14 @@
+# Python 3.10
+torch==2.8.0
+transformers==4.46.0
+lightning==2.5.5
+pytorch-lightning==2.5.5
+datasets==4.0.0
+pyarrow==21.0.0
+numpy==1.24.4
+pandas==2.0.3
+scipy==1.13.0
+scikit-learn==1.7.1
+absl-py==2.3.1
+safetensors==0.6.2
+tqdm==4.67.1

Computational/package/sample_direct.py

@@ -0,0 +1,174 @@
+"""Direct Metalorian sampling."""
+from metalorian.sampling import Sampler, FLAGS, METAL_LABELS, TOKEN_DROPOUT_SCALE
+from metalorian.cli import parser, collect
+import torch
+import torch.nn.functional as F
+CANVAS = 1024
+
+class FixedGradSampler(Sampler):
+
+    def __init__(self, cfg, device='cuda:0'):
+        super().__init__(cfg, device)
+        assert float(self.aa_bias.abs().max()) == 0.0, "aa_bias must be zero (no chemistry)"
+        print(f"guide_on={cfg.guide_on}  accumulate={not cfg.no_accum}  "
+              f"grad_steps={cfg.grad_steps}  scale={cfg.guidance_scale}")
+
+    def x0_hat(self, x_t, t, cond):
+        eps = self.net_sampler.model(x_t, t, cond)
+        return self.net_sampler.predict_xstart_from_eps(x_t, t, eps=eps)
+
+    def target_logit(self, lat, attn_len):
+        cfg = self.cfg
+        logits = self.model_con.decode_embeddings(lat[:, :attn_len])
+        seg = logits + self.vocab_mask[None, None, :]
+        soft = F.softmax(seg / cfg.tau, dim=-1)
+
+        emb = soft @ self.W
+        B = emb.shape[0]
+        cls_e = self.W[self.cls_id].view(1, 1, -1).expand(B, 1, -1)
+        eos_e = self.W[self.eos_id].view(1, 1, -1).expand(B, 1, -1)
+        inp = torch.cat([cls_e, emb, eos_e], dim=1) * TOKEN_DROPOUT_SCALE
+        am = torch.ones(B, attn_len + 2, device=self.device)
+
+        x_pred, _, _, _ = self.predictor.forward_from_embeds(inp, am)
+        tgt = x_pred[:, cfg.target_idx]
+        oth = x_pred[:, self.others].max(dim=1).values
+        return tgt - cfg.penalty * oth, soft
+
+    def init_latent(self, B, lmask):
+        return torch.randn(B, CANVAS, 1280, device=self.device) * lmask
+
+    @torch.no_grad()
+    def _encode_perpos(self, seqs):
+        """Encode sequences as residue latents of shape (N, S, 1280)."""
+        S = self.cfg.seq_len
+        m = self.model_con.base_model
+        enc = self.tokenizer(seqs, return_tensors='pt', padding='max_length',
+                             max_length=S + 2, truncation=True)
+        ids = enc['input_ids'].to(self.device); am = enc['attention_mask'].to(self.device)
+        emb = m.esm_model(input_ids=ids, attention_mask=am).last_hidden_state
+        x_pool, x_attns = m.attn_head(emb, am)
+        x_pool = m.attn_ln(x_pool + m.attn_skip(x_pool))
+        for lin in m.linear_layers:
+            r = x_pool; x_pool = F.silu(lin(x_pool)); x_pool = x_pool + r
+        x_w = torch.einsum('bhlk,bld->bhld', x_attns, x_pool).mean(dim=1)
+        return m.clf_ln(x_w)[:, 1:1 + S, :]
+
+    @torch.no_grad()
+    def _best_decode_one(self, lat_S, K, temp):
+        """Select among K decodes by mean positional cosine similarity."""
+        S = self.cfg.seq_len
+        lg = self.model_con.decode_embeddings(lat_S.unsqueeze(0))[0] + self.vocab_mask[None, :]
+        probs = F.softmax(lg / temp, dim=-1)
+        toks = torch.multinomial(probs, K, replacement=True).T
+        cand = [self.tokenizer.decode(t, skip_special_tokens=True).replace(' ', '') for t in toks]
+        cand = [c for c in cand if len(c) == S]
+        if not cand:
+            return self.tokenizer.decode(lg.argmax(-1), skip_special_tokens=True).replace(' ', '')
+        xc = F.normalize(self._encode_perpos(cand), dim=-1)
+        Ln = F.normalize(lat_S, dim=-1)
+        cos = (xc * Ln[None]).sum(-1).mean(-1)
+        return cand[int(cos.argmax())]
+
+    def decode(self, x, attn_len):
+        if getattr(self.cfg, 'best_decode', 0) and self.cfg.best_decode > 0:
+            return [self._best_decode_one(x[i, :attn_len], self.cfg.best_decode,
+                                          self.cfg.decode_temp) for i in range(x.shape[0])]
+        return self.hard_decode(x, attn_len)
+
+    def run_round(self):
+        cfg, dev = self.cfg, self.device
+        B, S = cfg.batch, cfg.seq_len
+
+        lens = torch.full((B,), S, device=dev, dtype=torch.long)
+        attn = torch.arange(CANVAS, device=dev)[None, :] < lens[:, None]
+        lmask = attn.unsqueeze(-1).float()
+
+        x = self.init_latent(B, lmask)
+
+        tgt1h = torch.zeros(B, 15, device=dev)
+        tgt1h[:, cfg.target_idx] = 1.0
+        log_dis = torch.log(tgt1h + 1e-10)
+
+        best_score = [-1.0] * B
+        best_seq = [None] * B
+
+        cys_trace = []
+        for step in reversed(range(FLAGS.T)):
+            t = torch.full((B,), step, device=dev, dtype=torch.long)
+
+            for k in range(cfg.grad_steps):
+                x = x.detach().requires_grad_(True)
+                if cfg.guide_on == 'x0':
+                    lat = self.x0_hat(x, t, log_dis)
+                else:
+                    lat = x
+                obj, soft = self.target_logit(lat, S)
+                g, = torch.autograd.grad(obj.sum(), x)
+                with torch.no_grad():
+                    gn = g.flatten(1).norm(dim=1).clamp_min(1e-8).view(-1, 1, 1)
+                    stepv = cfg.guidance_scale * (g / gn) * lmask
+                    if cfg.no_accum:
+                        mean, _ = self.net_sampler.p_mean_variance(
+                            x_t=x.detach(), t=t, cond=log_dis, trans=None)
+                        x = (mean * lmask + stepv).clamp(-1, 1)
+                    else:
+                        x = (x.detach() + stepv).clamp(-1, 1)
+                    x = x * lmask
+
+            with torch.no_grad():
+                mean, log_var = self.net_sampler.p_mean_variance(
+                    x_t=x, t=t, cond=log_dis, trans=None)
+                mean = (mean * lmask).float()
+                noise = torch.randn_like(x) if step > 0 else torch.zeros_like(x)
+                x = (mean + torch.exp(0.5 * log_var.float()) * noise * lmask).clamp(-1, 1)
+                x = x * lmask
+
+                log_dis = self.trainer_dis.p_sample(log_dis, t, x)
+                if step % 2 == 0:
+                    c = cfg.label_clamp
+                    log_dis = (1 - c) * log_dis + c * torch.log(tgt1h + 1e-10)
+
+            if cfg.track_best:
+                with torch.no_grad():
+                    cur = self.hard_decode(x, S)
+                    ok = [i for i, q in enumerate(cur) if len(q) == S]
+                    if ok:
+                        pr, _ = self.score([cur[i] for i in ok])
+                        sc = pr[:, cfg.target_idx]
+                        for j, i in enumerate(ok):
+                            if sc[j].item() > best_score[i]:
+                                best_score[i] = sc[j].item()
+                                best_seq[i] = cur[i]
+
+            if step % 10 == 0:
+                with torch.no_grad():
+                    seqs = self.hard_decode(x, S)
+                    cys_trace.append(sum(q.count('C') for q in seqs) / max(len(seqs), 1))
+
+        if cfg.track_best:
+            final = self.decode(x, S)
+            out = [best_seq[i] if best_seq[i] is not None else final[i] for i in range(B)]
+            return out, cys_trace
+        return self.decode(x, S), cys_trace
+
+    def accept(self, seqs, _unused):
+        cfg = self.cfg
+        probs, fires = self.score(seqs)
+        out = []
+        for i, s in enumerate(seqs):
+            if len(s) != cfg.seq_len:
+                continue
+            sc = probs[i, cfg.target_idx].item()
+            if cfg.accept_all:
+                out.append({'sequence': s, 'score': sc, 'n_cys': s.count('C'),
+                            'n_his': s.count('H')})
+                continue
+            if sc >= self.threshold and fires[i, cfg.target_idx].item() >= 0.5:
+                out.append({'sequence': s, 'score': sc, 'n_cys': s.count('C'),
+                            'n_his': s.count('H')})
+        return out, probs
+
+
+if __name__ == '__main__':
+    collect(parser('direct').parse_args(), FixedGradSampler, 'direct')

Computational/package/sample_prototype.py

@@ -0,0 +1,129 @@
+"""Prototype-seeded Metalorian sampling."""
+from sample_direct import FixedGradSampler, CANVAS
+from metalorian.sampling import FLAGS
+from metalorian.cli import parser, collect
+import torch
+
+class ProtoSampler(FixedGradSampler):
+
+    def __init__(self, cfg, device='cuda:0'):
+        super().__init__(cfg, device)
+        self.proto_lat = self._encode_prototype(cfg.proto)
+
+    @torch.no_grad()
+    def _encode_prototype(self, seq):
+        """Encode the prototype using the diffusion model's residue representation."""
+        import torch.nn.functional as F
+        S = self.cfg.seq_len
+        seq = (seq * ((S // max(len(seq), 1)) + 1))[:S] if len(seq) < S else seq[:S]
+        enc = self.tokenizer(seq, return_tensors='pt')
+        ids = enc['input_ids'].to(self.device)
+        am = enc['attention_mask'].to(self.device)
+
+        m = self.model_con.base_model
+        emb = m.esm_model(input_ids=ids, attention_mask=am).last_hidden_state
+        x_pool, x_attns = m.attn_head(emb, am)
+        x_pool = m.attn_ln(x_pool + m.attn_skip(x_pool))
+        for lin in m.linear_layers:
+            res = x_pool
+            x_pool = F.silu(lin(x_pool))
+            x_pool = x_pool + res
+        x_comb = torch.einsum('bhlk,bld->bhld', x_attns, x_pool).mean(dim=1)
+        x_comb = m.clf_ln(x_comb)
+        h = x_comb[:, 1:-1, :]                          # Remove CLS/EOS.
+
+        lat = torch.zeros(1, CANVAS, 1280, device=self.device)
+        lat[:, :h.shape[1], :] = h
+        lat = lat.clamp(-1, 1)
+
+        rt = self.hard_decode(lat, S)[0]
+        ident = sum(a == b for a, b in zip(rt, seq))
+        print(f"prototype  '{seq}'")
+        print(f"round-trip '{rt}'   identity={ident}/{S}  Cys={rt.count('C')}")
+        if rt.count('C') == 0:
+            print("Warning: the decoded seed contains no cysteine.")
+        return lat
+
+    def init_latent(self, B, lmask):
+        cfg = self.cfg
+        seed = self.proto_lat.expand(B, -1, -1).clone()
+        noise = torch.randn(B, CANVAS, 1280, device=self.device)
+        x = seed + cfg.init_noise * noise
+        return (x * lmask).clamp(-1, 1)
+
+    def _alphas_bar(self):
+        betas = torch.linspace(FLAGS.beta_1, FLAGS.beta_T, FLAGS.T, device=self.device)
+        return torch.cumprod(1.0 - betas, dim=0)
+
+    def run_round(self):
+        cfg, dev = self.cfg, self.device
+        if cfg.t_start <= 0 or cfg.t_start >= FLAGS.T:
+            return super().run_round()
+
+        import torch.nn.functional as F
+        B, S = cfg.batch, cfg.seq_len
+        lens = torch.full((B,), S, device=dev, dtype=torch.long)
+        attn = torch.arange(CANVAS, device=dev)[None, :] < lens[:, None]
+        lmask = attn.unsqueeze(-1).float()
+
+        ab = self._alphas_bar()
+        t0 = int(cfg.t_start)
+        seed = self.proto_lat.expand(B, -1, -1).clone() * lmask
+        noise = torch.randn(B, CANVAS, 1280, device=dev)
+        x = (ab[t0].sqrt() * seed + (1 - ab[t0]).sqrt() * noise) * lmask
+        x = x.clamp(-1, 1)
+
+        tgt1h = torch.zeros(B, 15, device=dev)
+        tgt1h[:, cfg.target_idx] = 1.0
+        log_dis = torch.log(tgt1h + 1e-10)
+
+        best_score = [-1.0] * B
+        best_seq = [None] * B
+
+        cys_trace = []
+        for step in reversed(range(t0 + 1)):
+            t = torch.full((B,), step, device=dev, dtype=torch.long)
+            for _ in range(cfg.grad_steps):
+                x = x.detach().requires_grad_(True)
+                lat = self.x0_hat(x, t, log_dis) if cfg.guide_on == 'x0' else x
+                obj, _ = self.target_logit(lat, S)
+                g, = torch.autograd.grad(obj.sum(), x)
+                with torch.no_grad():
+                    gn = g.flatten(1).norm(dim=1).clamp_min(1e-8).view(-1, 1, 1)
+                    x = (x.detach() + cfg.guidance_scale * (g / gn) * lmask).clamp(-1, 1) * lmask
+            with torch.no_grad():
+                mean, log_var = self.net_sampler.p_mean_variance(
+                    x_t=x, t=t, cond=log_dis, trans=None)
+                mean = (mean * lmask).float()
+                nz = torch.randn_like(x) if step > 0 else torch.zeros_like(x)
+                x = (mean + torch.exp(0.5 * log_var.float()) * nz * lmask).clamp(-1, 1) * lmask
+                log_dis = self.trainer_dis.p_sample(log_dis, t, x)
+                if step % 2 == 0:
+                    c = cfg.label_clamp
+                    log_dis = (1 - c) * log_dis + c * torch.log(tgt1h + 1e-10)
+            if cfg.track_best:
+                with torch.no_grad():
+                    cur = self.hard_decode(x, S)
+                    ok = [i for i, q in enumerate(cur) if len(q) == S]
+                    if ok:
+                        pr, _ = self.score([cur[i] for i in ok])
+                        sc = pr[:, cfg.target_idx]
+                        for j, i in enumerate(ok):
+                            if sc[j].item() > best_score[i]:
+                                best_score[i] = sc[j].item()
+                                best_seq[i] = cur[i]
+
+            if step % 5 == 0:
+                with torch.no_grad():
+                    sq = self.hard_decode(x, S)
+                    cys_trace.append(sum(q.count('C') for q in sq) / max(len(sq), 1))
+
+        if cfg.track_best:
+            final = self.decode(x, S)
+            return [best_seq[i] if best_seq[i] is not None else final[i]
+                    for i in range(B)], cys_trace
+        return self.decode(x, S), cys_trace
+
+
+if __name__ == '__main__':
+    collect(parser('prototype').parse_args(), ProtoSampler, 'prototype')

Computational/package/train.py

@@ -0,0 +1,219 @@
+"""Train MetaLATTE or Metalorian."""
+import argparse
+import csv
+import json
+from pathlib import Path
+
+
+def arguments():
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument('--stage', choices=['metalatte', 'metalorian'], required=True)
+    p.add_argument('--output', type=Path, required=True)
+    p.add_argument('--device', default='cuda:0')
+    p.add_argument('--epochs', type=int, default=30)
+    p.add_argument('--batch-size', type=int, default=2,
+                   help='Sequence batch size for Metalorian; MetaLATTE uses its saved pre-batches')
+    p.add_argument('--workers', type=int, default=0)
+    p.add_argument('--threads', type=int, default=4)
+    p.add_argument('--seed', type=int, default=42)
+    p.add_argument('--lr', type=float, default=2e-4)
+    p.add_argument('--init-checkpoint', type=Path,
+                   help='MetaLATTE weights to initialize from; defaults to the supplied predictor for Metalorian')
+    p.add_argument('--resume', type=Path, help='Resume a checkpoint for the selected stage')
+    p.add_argument('--smoke-test', action='store_true', help='One update using at most two records and 64 tokens; saves no model')
+    a = p.parse_args()
+    if min(a.epochs, a.batch_size, a.threads) < 1 or a.workers < 0:
+        p.error('epochs, batch-size and threads must be positive; workers cannot be negative')
+    if a.smoke_test and a.resume:
+        p.error('smoke-test starts a short test run; omit resume')
+    return a
+
+
+def class_weights():
+    import numpy as np
+    import torch
+    from datasets import load_from_disk
+    from metalorian.paths import TRAIN
+    ds = load_from_disk(str(TRAIN))
+    counts = np.zeros(15, dtype=np.float64)
+    for row in ds.select_columns(['anchor_labels']):
+        counts += np.asarray(row['anchor_labels']).sum(axis=0)
+    weights = 1.0/(counts+1e-6)
+    return torch.tensor(weights/weights.sum()*15, dtype=torch.float32)
+
+
+def predictor(args, checkpoint=None, residue_encoder=False):
+    import torch
+    from transformers import EsmModel, EsmConfig
+    from metalorian.paths import BACKBONE, PREDICTOR
+    from metalorian.predictor import MultitaskProteinModel, ResidueEncoder
+    metadata = torch.load(PREDICTOR, map_location='cpu', weights_only=False, mmap=True)
+    hp = argparse.Namespace(**metadata['hyper_parameters'])
+    hp.lr, hp.epochs = args.lr, args.epochs
+    hp.output_model_path = str(args.output.resolve())
+    del metadata
+    esm = (EsmModel(EsmConfig.from_pretrained(str(BACKBONE), local_files_only=True))
+           if checkpoint else EsmModel.from_pretrained(str(BACKBONE), local_files_only=True))
+    cls = ResidueEncoder if residue_encoder else MultitaskProteinModel
+    model = cls(hparams=hp, esm_model=esm, num_layers_to_finetune=2,
+                input_dim=1280, output_dim=15, class_weights=class_weights())
+    if checkpoint:
+        state = torch.load(checkpoint, map_location='cpu', weights_only=False, mmap=True)
+        model.load_state_dict(state['state_dict'], strict=True)
+    return model
+
+
+def smoke_group(row):
+    result = {}
+    for key, values in row.items():
+        result[key] = values[:2] if key.endswith('labels') else [x[:64] for x in values[:2]]
+    return result
+
+
+def train_metalatte(args):
+    import torch
+    from datasets import load_from_disk
+    from torch.utils.data import DataLoader
+    from lightning.pytorch import Trainer
+    from lightning.pytorch.callbacks import ModelCheckpoint, StochasticWeightAveraging
+    from lightning.pytorch.loggers import CSVLogger
+    from metalorian.paths import TRAIN, VALIDATION
+    from metalorian.data import collate_fn_hard, collate_fn
+    train, valid = load_from_disk(str(TRAIN)), load_from_disk(str(VALIDATION))
+    if args.smoke_test:
+        train, valid = [smoke_group(train[0])], [smoke_group(valid[0])]
+    model = predictor(args, args.init_checkpoint or args.resume)
+    train_loader = DataLoader(train, batch_size=1, shuffle=False, collate_fn=collate_fn_hard, num_workers=args.workers)
+    valid_loader = DataLoader(valid, batch_size=1, shuffle=False, collate_fn=collate_fn, num_workers=args.workers)
+    gpu = args.device.startswith('cuda')
+    device_index = int(args.device.split(':')[1]) if ':' in args.device else 0
+    callbacks = [] if args.smoke_test else [
+        ModelCheckpoint(dirpath=args.output/'checkpoints', save_last=True, save_top_k=1,
+                        monitor='val_macro_f1', mode='max', filename='metalatte-{epoch:02d}'),
+        StochasticWeightAveraging(swa_lrs=1e-2)]
+    trainer = Trainer(max_epochs=args.epochs, accelerator='gpu' if gpu else 'cpu',
+                      devices=[device_index] if gpu else 1, precision='bf16-mixed' if gpu else '32-true',
+                      logger=CSVLogger(str(args.output), name='metrics'), callbacks=callbacks,
+                      gradient_clip_val=0.5, enable_checkpointing=not args.smoke_test,
+                      enable_model_summary=False, fast_dev_run=args.smoke_test,
+                      default_root_dir=str(args.output))
+    trainer.fit(model, train_dataloaders=train_loader, val_dataloaders=valid_loader,
+                ckpt_path=str(args.resume) if args.resume else None)
+    return {'stage': 'metalatte', 'status': 'completed', 'steps': int(trainer.global_step),
+            'smoke_test': args.smoke_test}
+
+
+def train_metalorian(args):
+    import torch
+    import numpy as np
+    from datasets import load_from_disk
+    from torch.utils.data import DataLoader, DistributedSampler
+    from types import SimpleNamespace
+    from metalorian.paths import ROOT, PREDICTOR
+    from metalorian.continuous import DiffusionProteinModel, GaussianDiffusionTrainer
+    from metalorian.discrete import MultinomialDiffusion
+    from metalorian.unet import tabularUnet
+    from metalorian.objectives import training_with_esm, make_negative_condition_esm, warmup_lr
+
+    flags = SimpleNamespace(T=50, beta_1=1e-5, beta_T=0.02, input_size=1280,
+        cond_size=15, output_size=1280, encoder_dim=[64,128,256], nf=16, activation='relu',
+        hidden_size=768, num_layers=12, num_heads=12, intermediate_size=3072,
+        max_seq_length=1280, lambda_con=0.2, lambda_dis=0.2)
+    base = predictor(args, args.init_checkpoint or PREDICTOR, residue_encoder=True)
+    con = DiffusionProteinModel(base, flags).to(args.device)
+    flags.input_size, flags.cond_size, flags.output_size, flags.nf = 15, 1280, 15, 64
+    dis = tabularUnet(flags).to(args.device)
+    con_trainer = GaussianDiffusionTrainer(con, flags.beta_1, flags.beta_T, flags.T).to(args.device)
+    data = load_from_disk(str(ROOT/'data/metalorian/train'))
+    if args.smoke_test:
+        data = [{k:(v[:64] if k in {'input_ids','attention_mask'} else v) for k,v in data[i].items()}
+                for i in range(min(2, len(data)))]
+    data = TokenRows(data)
+    # Preserve the training sampler order: seed 0, epoch 0.
+    order = DistributedSampler(data, num_replicas=1, rank=0, shuffle=True, seed=0)
+    loader = DataLoader(data, batch_size=args.batch_size, sampler=order, num_workers=args.workers)
+    dis_trainer = MultinomialDiffusion(np.array([15]), (len(data),15), dis, flags,
+        timesteps=flags.T, loss_type='vb_stochastic').to(args.device)
+    optim_con = torch.optim.Adam(con.parameters(), lr=args.lr)
+    optim_dis = torch.optim.Adam(dis.parameters(), lr=args.lr)
+    sched_con = torch.optim.lr_scheduler.LambdaLR(optim_con, warmup_lr)
+    sched_dis = torch.optim.lr_scheduler.LambdaLR(optim_dis, warmup_lr)
+    start_epoch, step = 0, 0
+    if args.resume:
+        ckpt = torch.load(args.resume, map_location='cpu', weights_only=False, mmap=True)
+        strip = lambda state: {k.removeprefix('module.'):v for k,v in state.items()}
+        con.load_state_dict(strip(ckpt['model_con']), strict=True)
+        dis.load_state_dict(strip(ckpt['model_dis']), strict=True)
+        optim_con.load_state_dict(ckpt['optim_con']); optim_dis.load_state_dict(ckpt['optim_dis'])
+        sched_con.load_state_dict(ckpt['sched_con']); sched_dis.load_state_dict(ckpt['sched_dis'])
+        start_epoch, step = int(ckpt.get('epoch',0)), int(ckpt.get('step',0))+1
+        del ckpt
+    con_params = [p for p in con.parameters() if p.requires_grad]
+    dis_params = [p for p in dis.parameters() if p.requires_grad]
+    transform = SimpleNamespace(output_info=[(15,'softmax')])
+    weights = class_weights().to(args.device)
+    with (args.output/'losses.csv').open('w', newline='') as handle:
+        log = csv.writer(handle)
+        log.writerow(['epoch','step','continuous_loss','continuous_contrastive','discrete_loss','discrete_contrastive'])
+        for epoch in range(start_epoch, start_epoch+args.epochs):
+            con.train(); dis.train()
+            for batch in loader:
+                x = {k:batch[k].to(args.device) for k in ['input_ids','attention_mask']}
+                labels = batch['labels'].float().to(args.device)
+                ns_con, ns_dis = make_negative_condition_esm(x, labels, base)
+                values = training_with_esm(x, labels, con_trainer, dis_trainer, ns_con, ns_dis,
+                                          transform, flags, None, weights)
+                loss_con = values[0]+flags.lambda_con*values[1]
+                loss_dis = values[2]+flags.lambda_dis*values[3]
+                if not torch.isfinite(loss_con+loss_dis):
+                    raise RuntimeError('Nonfinite training loss')
+                # Compute both gradients before updating the shared forward graph.
+                optim_con.zero_grad(set_to_none=True); optim_dis.zero_grad(set_to_none=True)
+                gc = torch.autograd.grad(loss_con, con_params, retain_graph=True, allow_unused=True)
+                gd = torch.autograd.grad(loss_dis, dis_params, allow_unused=True)
+                for p,g in zip(con_params,gc): p.grad=g
+                for p,g in zip(dis_params,gd): p.grad=g
+                torch.nn.utils.clip_grad_norm_(con_params, 1.0)
+                torch.nn.utils.clip_grad_norm_(dis_params, 1.0)
+                optim_con.step(); sched_con.step()
+                optim_dis.step(); sched_dis.step()
+                numbers = [float(v.detach()) if torch.is_tensor(v) else float(v) for v in values]
+                log.writerow([epoch,step,*numbers]); handle.flush()
+                print(f'epoch={epoch} step={step} continuous={float(loss_con.detach()):.6f} discrete={float(loss_dis.detach()):.6f}', flush=True)
+                step += 1
+                if args.smoke_test:
+                    return {'stage':'metalorian','status':'completed','steps':1,'smoke_test':True,'losses':numbers}
+            torch.save(dict(model_con=con.state_dict(), model_dis=dis.state_dict(),
+                optim_con=optim_con.state_dict(), optim_dis=optim_dis.state_dict(),
+                sched_con=sched_con.state_dict(), sched_dis=sched_dis.state_dict(),
+                step=step-1, epoch=epoch+1), args.output/'metalorian_latest.pt')
+    return {'stage':'metalorian','status':'completed','steps':step,'smoke_test':False}
+
+
+class TokenRows:
+    def __init__(self, data): self.data=data
+    def __len__(self): return len(self.data)
+    def __getitem__(self, index):
+        import torch
+        row=self.data[index]
+        return {k:torch.tensor(row[k]) for k in ['input_ids','attention_mask','labels']}
+
+
+def main():
+    args = arguments()
+    from metalorian.paths import configure
+    configure()
+    import torch
+    from lightning.pytorch import seed_everything
+    seed_everything(args.seed, workers=True)
+    torch.set_num_threads(args.threads)
+    args.output.mkdir(parents=True, exist_ok=True)
+    (args.output/'settings.json').write_text(json.dumps({k:str(v) if isinstance(v,Path) else v
+                                                       for k,v in vars(args).items()}, indent=2)+'\n')
+    result = train_metalatte(args) if args.stage=='metalatte' else train_metalorian(args)
+    (args.output/'result.json').write_text(json.dumps(result, indent=2)+'\n')
+    print(json.dumps(result, indent=2))
+
+
+if __name__ == '__main__':
+    main()

Computational/validation/figures/amino_acid_composition_difference_full_reference_predictor_gated_n500.csv

@@ -0,0 +1,81 @@
+metal,AA,generated_percent,full_reference_percent,difference_percentage_points,generated_n,full_reference_n,reference_mode
+Cu,A,15.179714285714288,8.253464474116626,6.926249811597662,500,722,all deduplicated canonical labeled sequences; no length filter
+Cu,C,4.784714285714286,1.7030293368818004,3.0816849488324856,500,722,all deduplicated canonical labeled sequences; no length filter
+Cu,D,3.7777857142857116,5.662140562080337,-1.8843548477946253,500,722,all deduplicated canonical labeled sequences; no length filter
+Cu,E,3.2550714285714277,5.208001313139308,-1.95292988456788,500,722,all deduplicated canonical labeled sequences; no length filter
+Cu,F,5.276190476190477,4.33157251405727,0.9446179621332069,500,722,all deduplicated canonical labeled sequences; no length filter
+Cu,G,2.389214285714285,9.233285808964153,-6.844071523249868,500,722,all deduplicated canonical labeled sequences; no length filter
+Cu,H,17.78226190476191,3.4977835883197557,14.284478316442154,500,722,all deduplicated canonical labeled sequences; no length filter
+Cu,I,1.0133333333333336,5.140004122271833,-4.126670788938499,500,722,all deduplicated canonical labeled sequences; no length filter
+Cu,K,10.510690476190476,4.957583187361847,5.553107288828628,500,722,all deduplicated canonical labeled sequences; no length filter
+Cu,L,5.250500000000002,7.812511951883257,-2.562011951883255,500,722,all deduplicated canonical labeled sequences; no length filter
+Cu,M,8.446428571428573,2.4922300040473155,5.954198567381257,500,722,all deduplicated canonical labeled sequences; no length filter
+Cu,N,2.312333333333333,4.622110010058687,-2.3097766767253542,500,722,all deduplicated canonical labeled sequences; no length filter
+Cu,P,2.3733571428571416,5.646073830664104,-3.2727166878069625,500,722,all deduplicated canonical labeled sequences; no length filter
+Cu,Q,3.8210476190476177,3.019580967529003,0.8014666515186146,500,722,all deduplicated canonical labeled sequences; no length filter
+Cu,R,1.3417857142857146,3.7191257767039034,-2.377340062418189,500,722,all deduplicated canonical labeled sequences; no length filter
+Cu,S,5.754333333333336,6.325280933433413,-0.5709476001000766,500,722,all deduplicated canonical labeled sequences; no length filter
+Cu,T,4.632785714285715,6.443409990551853,-1.8106242762661378,500,722,all deduplicated canonical labeled sequences; no length filter
+Cu,V,0.5072857142857139,7.794299592156853,-7.287013877871139,500,722,all deduplicated canonical labeled sequences; no length filter
+Cu,W,0.4050238095238093,1.3277750447432664,-0.9227512352194571,500,722,all deduplicated canonical labeled sequences; no length filter
+Cu,Y,1.186142857142857,2.8107369910354243,-1.6245941338925673,500,722,all deduplicated canonical labeled sequences; no length filter
+Cd,A,13.550928571428567,8.496912513658536,5.054016057770031,500,158,all deduplicated canonical labeled sequences; no length filter
+Cd,C,28.13461904761906,2.90845505077731,25.226163996841752,500,158,all deduplicated canonical labeled sequences; no length filter
+Cd,D,1.3285000000000007,5.777820408205994,-4.449320408205993,500,158,all deduplicated canonical labeled sequences; no length filter
+Cd,E,1.0580952380952384,6.735951990218785,-5.677856752123547,500,158,all deduplicated canonical labeled sequences; no length filter
+Cd,F,0.8873809523809523,3.624294735065239,-2.7369137826842866,500,158,all deduplicated canonical labeled sequences; no length filter
+Cd,G,0.21890476190476182,7.537380486687105,-7.318475724782343,500,158,all deduplicated canonical labeled sequences; no length filter
+Cd,H,1.8600000000000003,2.169815582343022,-0.30981558234302176,500,158,all deduplicated canonical labeled sequences; no length filter
+Cd,I,0.5347619047619048,5.5484634938451824,-5.0137015890832775,500,158,all deduplicated canonical labeled sequences; no length filter
+Cd,K,12.108476190476187,6.889414083395107,5.21906210708108,500,158,all deduplicated canonical labeled sequences; no length filter
+Cd,L,0.5774285714285713,8.319675467711688,-7.742246896283117,500,158,all deduplicated canonical labeled sequences; no length filter
+Cd,M,1.9189761904761924,2.3005032704500055,-0.3815270799738131,500,158,all deduplicated canonical labeled sequences; no length filter
+Cd,N,0.6128809523809523,4.2105386760265775,-3.597657723645625,500,158,all deduplicated canonical labeled sequences; no length filter
+Cd,P,12.505142857142854,4.4208136568506236,8.084329200292231,500,158,all deduplicated canonical labeled sequences; no length filter
+Cd,Q,7.871976190476187,3.609113466561349,4.262862723914838,500,158,all deduplicated canonical labeled sequences; no length filter
+Cd,R,0.4346904761904761,4.382222241972551,-3.9475317657820748,500,158,all deduplicated canonical labeled sequences; no length filter
+Cd,S,1.1151904761904763,6.254286863063874,-5.139096386873398,500,158,all deduplicated canonical labeled sequences; no length filter
+Cd,T,12.808642857142862,5.87294310794069,6.935699749202172,500,158,all deduplicated canonical labeled sequences; no length filter
+Cd,V,1.0143095238095234,6.694234937216741,-5.679925413407218,500,158,all deduplicated canonical labeled sequences; no length filter
+Cd,W,0.23940476190476198,1.1352977659664403,-0.8958930040616783,500,158,all deduplicated canonical labeled sequences; no length filter
+Cd,Y,1.2196904761904768,3.1118622020431728,-1.892171725852696,500,158,all deduplicated canonical labeled sequences; no length filter
+Zn,A,9.557857142857143,8.083222676194968,1.4746344666621756,500,17094,all deduplicated canonical labeled sequences; no length filter
+Zn,C,12.880357142857148,1.9417560484397594,10.93860109441739,500,17094,all deduplicated canonical labeled sequences; no length filter
+Zn,D,2.8502619047619047,5.750286885020571,-2.900024980258666,500,17094,all deduplicated canonical labeled sequences; no length filter
+Zn,E,2.1406428571428577,7.187470003911243,-5.046827146768385,500,17094,all deduplicated canonical labeled sequences; no length filter
+Zn,F,4.354476190476189,3.9175415402120577,0.4369346502641309,500,17094,all deduplicated canonical labeled sequences; no length filter
+Zn,G,0.2592380952380954,7.367717120991302,-7.108479025753207,500,17094,all deduplicated canonical labeled sequences; no length filter
+Zn,H,8.881214285714282,2.8099397109243713,6.071274574789911,500,17094,all deduplicated canonical labeled sequences; no length filter
+Zn,I,1.1118095238095242,5.863585955880062,-4.751776432070537,500,17094,all deduplicated canonical labeled sequences; no length filter
+Zn,K,16.81728571428571,6.022020919922009,10.7952647943637,500,17094,all deduplicated canonical labeled sequences; no length filter
+Zn,L,8.739642857142861,9.015234213023932,-0.27559135588107075,500,17094,all deduplicated canonical labeled sequences; no length filter
+Zn,M,4.8561428571428555,2.408384549980984,2.4477583071618714,500,17094,all deduplicated canonical labeled sequences; no length filter
+Zn,N,1.1791428571428582,3.876948690019201,-2.697805832876343,500,17094,all deduplicated canonical labeled sequences; no length filter
+Zn,P,3.4906190476190475,4.590996319680774,-1.1003772720617264,500,17094,all deduplicated canonical labeled sequences; no length filter
+Zn,Q,6.338071428571425,3.664161825260114,2.673909603311311,500,17094,all deduplicated canonical labeled sequences; no length filter
+Zn,R,0.7605952380952381,5.649111106864038,-4.8885158687688,500,17094,all deduplicated canonical labeled sequences; no length filter
+Zn,S,1.8969047619047636,5.949762082324762,-4.0528573204199985,500,17094,all deduplicated canonical labeled sequences; no length filter
+Zn,T,9.747642857142855,5.181499373617499,4.566143483525356,500,17094,all deduplicated canonical labeled sequences; no length filter
+Zn,V,1.3977380952380967,6.639014939543678,-5.241276844305581,500,17094,all deduplicated canonical labeled sequences; no length filter
+Zn,W,0.2297142857142858,1.0493511419736743,-0.8196368562593885,500,17094,all deduplicated canonical labeled sequences; no length filter
+Zn,Y,2.510642857142857,3.031994896215333,-0.5213520390724762,500,17094,all deduplicated canonical labeled sequences; no length filter
+Ni,A,16.60880952380952,8.690803978647393,7.918005545162126,500,657,all deduplicated canonical labeled sequences; no length filter
+Ni,C,0.7498095238095238,1.7965676833282913,-1.0467581595187676,500,657,all deduplicated canonical labeled sequences; no length filter
+Ni,D,5.812261904761905,6.235482314489284,-0.42322040972737884,500,657,all deduplicated canonical labeled sequences; no length filter
+Ni,E,1.7889047619047629,7.681116691784482,-5.8922119298797195,500,657,all deduplicated canonical labeled sequences; no length filter
+Ni,F,17.295357142857135,3.4049905754311363,13.890366567426,500,657,all deduplicated canonical labeled sequences; no length filter
+Ni,G,2.0798333333333336,7.859951665750784,-5.78011833241745,500,657,all deduplicated canonical labeled sequences; no length filter
+Ni,H,9.26166666666667,3.493432428620067,5.768234238046603,500,657,all deduplicated canonical labeled sequences; no length filter
+Ni,I,2.907500000000002,6.706331318902961,-3.7988313189029594,500,657,all deduplicated canonical labeled sequences; no length filter
+Ni,K,4.857428571428573,4.754340691032745,0.10308788039582861,500,657,all deduplicated canonical labeled sequences; no length filter
+Ni,L,10.335404761904762,8.629904050710053,1.7055007111947091,500,657,all deduplicated canonical labeled sequences; no length filter
+Ni,M,1.7150952380952393,2.6432846868016315,-0.9281894487063922,500,657,all deduplicated canonical labeled sequences; no length filter
+Ni,N,4.7038095238095226,3.2841711826656286,1.419638341143894,500,657,all deduplicated canonical labeled sequences; no length filter
+Ni,P,0.6057142857142859,4.251372796204826,-3.6456585104905406,500,657,all deduplicated canonical labeled sequences; no length filter
+Ni,Q,7.510761904761904,3.54265907397933,3.968102830782574,500,657,all deduplicated canonical labeled sequences; no length filter
+Ni,R,2.0985238095238095,5.169522496366449,-3.0709986868426395,500,657,all deduplicated canonical labeled sequences; no length filter
+Ni,S,0.4001904761904762,5.279528374057049,-4.879337897866574,500,657,all deduplicated canonical labeled sequences; no length filter
+Ni,T,5.401714285714286,5.684247317678673,-0.2825330319643866,500,657,all deduplicated canonical labeled sequences; no length filter
+Ni,V,4.9288333333333325,7.393953634343296,-2.4651203010099634,500,657,all deduplicated canonical labeled sequences; no length filter
+Ni,W,0.05726190476190476,1.059226726729418,-1.0019648219675132,500,657,all deduplicated canonical labeled sequences; no length filter
+Ni,Y,0.8811190476190476,2.439112312476509,-1.5579932648574613,500,657,all deduplicated canonical labeled sequences; no length filter

Computational/validation/figures/amino_acid_composition_difference_overall_full_reference_predictor_gated_n500.csv

@@ -0,0 +1,21 @@
+AA,generated_percent,full_reference_percent_equal_metal_weight,difference_percentage_points,metals,generated_n_total,generated_n_per_metal,reference_weighting
+A,13.72432738095238,8.38110091065438,5.343226470297999,"Cu,Cd,Zn,Ni",2000,500,equal weight per metal; all deduplicated canonical labeled sequences within metal
+C,11.637375000000006,2.0874520298567902,9.549922970143216,"Cu,Cd,Zn,Ni",2000,500,equal weight per metal; all deduplicated canonical labeled sequences within metal
+D,3.4422023809523807,5.8564325424490455,-2.414230161496665,"Cu,Cd,Zn,Ni",2000,500,equal weight per metal; all deduplicated canonical labeled sequences within metal
+E,2.0606785714285714,6.703134999763455,-4.642456428334883,"Cu,Cd,Zn,Ni",2000,500,equal weight per metal; all deduplicated canonical labeled sequences within metal
+F,6.953351190476188,3.819599841191426,3.133751349284762,"Cu,Cd,Zn,Ni",2000,500,equal weight per metal; all deduplicated canonical labeled sequences within metal
+G,1.236797619047619,7.999583770598336,-6.762786151550717,"Cu,Cd,Zn,Ni",2000,500,equal weight per metal; all deduplicated canonical labeled sequences within metal
+H,9.446285714285715,2.9927428275518038,6.453542886733911,"Cu,Cd,Zn,Ni",2000,500,equal weight per metal; all deduplicated canonical labeled sequences within metal
+I,1.3918511904761912,5.81459622272501,-4.422745032248819,"Cu,Cd,Zn,Ni",2000,500,equal weight per metal; all deduplicated canonical labeled sequences within metal
+K,11.073470238095236,5.655839720427927,5.417630517667309,"Cu,Cd,Zn,Ni",2000,500,equal weight per metal; all deduplicated canonical labeled sequences within metal
+L,6.225744047619049,8.444331420832233,-2.218587373213184,"Cu,Cd,Zn,Ni",2000,500,equal weight per metal; all deduplicated canonical labeled sequences within metal
+M,4.234160714285714,2.461100627819984,1.7730600864657302,"Cu,Cd,Zn,Ni",2000,500,equal weight per metal; all deduplicated canonical labeled sequences within metal
+N,2.2020416666666662,3.9984421396925236,-1.7964004730258574,"Cu,Cd,Zn,Ni",2000,500,equal weight per metal; all deduplicated canonical labeled sequences within metal
+P,4.743708333333332,4.7273141508500816,0.016394182483250752,"Cu,Cd,Zn,Ni",2000,500,equal weight per metal; all deduplicated canonical labeled sequences within metal
+Q,6.385464285714283,3.458878833332449,2.926585452381834,"Cu,Cd,Zn,Ni",2000,500,equal weight per metal; all deduplicated canonical labeled sequences within metal
+R,1.1588988095238095,4.7299954054767355,-3.571096595952926,"Cu,Cd,Zn,Ni",2000,500,equal weight per metal; all deduplicated canonical labeled sequences within metal
+S,2.291654761904763,5.9522145632197745,-3.6605598013150114,"Cu,Cd,Zn,Ni",2000,500,equal weight per metal; all deduplicated canonical labeled sequences within metal
+T,8.147696428571429,5.795524947447179,2.35217148112425,"Cu,Cd,Zn,Ni",2000,500,equal weight per metal; all deduplicated canonical labeled sequences within metal
+V,1.9620416666666667,7.130375775815143,-5.168334109148476,"Cu,Cd,Zn,Ni",2000,500,equal weight per metal; all deduplicated canonical labeled sequences within metal
+W,0.23285119047619043,1.1429126698531997,-0.9100614793770093,"Cu,Cd,Zn,Ni",2000,500,equal weight per metal; all deduplicated canonical labeled sequences within metal
+Y,1.4493988095238095,2.8484266004426098,-1.3990277909188003,"Cu,Cd,Zn,Ni",2000,500,equal weight per metal; all deduplicated canonical labeled sequences within metal

Computational/validation/figures/physicochemical_properties_full_reference.csv

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

Computational/validation/figures/physicochemical_summary_predictor_gated_n500.csv

@@ -0,0 +1,49 @@
+metal,group,property,n,mean,sd,median,q1,q3,min,max
+Cu,WT/reference (full),molecular_weight_Da,722,43090.45546994471,28203.941041231414,36932.138300000064,15896.76067499999,66060.54120000014,2365.6033,137715.5366000003
+Cu,WT/reference (full),net_charge_pH7_4,722,-8.25361403234684,12.4030167462157,-5.50461928539662,-11.331205552654342,-1.3497848576496594,-78.13935885230536,22.76992820373951
+Cu,WT/reference (full),gravy_KyteDoolittle,722,-0.17670491017620815,0.3255128630139271,-0.20901702786377702,-0.39409619285379927,-0.018678380443086378,-1.487247474747477,0.8271573604060918
+Cu,WT/reference (full),instability_index,722,34.62494980142827,11.405359537675217,34.1490756302521,27.83105587487564,40.529098466082374,5.6118421052631575,89.19934306569343
+Cu,WT/reference (full),isoelectric_point,722,6.293408877301414,1.3432803957025612,5.978262710571289,5.447429180145263,6.817334079742432,4.0500284194946286,10.410951042175292
+Cu,WT/reference (full),aromaticity_fraction,722,0.08470084549835961,0.036741741922848396,0.09090909090909091,0.0548243036784581,0.10772732578498892,0.0,0.18840579710144928
+Cd,WT/reference (full),molecular_weight_Da,158,38162.80735189883,29985.143310554795,28397.248999999996,16767.83745,49090.314475000065,2330.6007999999993,138218.4704000007
+Cd,WT/reference (full),net_charge_pH7_4,158,-5.18745702946965,10.657151393684089,-3.612803138259399,-9.650048583039897,0.5578937284123242,-76.51336223204387,18.45704204492644
+Cd,WT/reference (full),gravy_KyteDoolittle,158,-0.27185842475346467,0.32504225635572787,-0.27876919468720684,-0.4955614183199661,-0.08339659884603691,-1.0413333333333337,0.6153488372093022
+Cd,WT/reference (full),instability_index,158,36.816472953023585,13.556098919629193,33.659181148212234,28.149665327978575,44.14624934282756,10.112987012987015,106.2375
+Cd,WT/reference (full),isoelectric_point,158,6.527721004244647,1.5629290349761467,6.04476432800293,5.286006450653076,7.892574214935303,4.0500284194946286,11.395060920715334
+Cd,WT/reference (full),aromaticity_fraction,158,0.07871454703074855,0.033234909113898765,0.08107772192454289,0.058969132207338376,0.09881661969723421,0.0,0.15348837209302324
+Zn,WT/reference (full),molecular_weight_Da,17094,50365.21003442743,30972.964836527746,41483.126850000015,27411.95415000001,71259.21570000025,4309.194600000001,142330.48110000076
+Zn,WT/reference (full),net_charge_pH7_4,17094,-8.195455008890075,13.890057948376617,-7.3908581197670475,-15.524508282469954,-0.6151202899890471,-127.18133881427826,117.35403972706803
+Zn,WT/reference (full),gravy_KyteDoolittle,17094,-0.2966314668557234,0.2554964850257074,-0.2878875291375293,-0.4388888888888886,-0.1464622641509436,-1.9647746243739599,1.205241935483873
+Zn,WT/reference (full),instability_index,17094,39.0336336336893,9.790887210488817,38.172547174447175,33.003702578755274,44.10148512736877,3.819354838709676,132.48106060606062
+Zn,WT/reference (full),isoelectric_point,17094,6.434931484372822,1.5255823948469378,5.892577934265136,5.373964786529541,7.05678253173828,4.0500284194946286,11.999967765808105
+Zn,WT/reference (full),aromaticity_fraction,17094,0.07998887578401025,0.022271011939281448,0.07878300691805515,0.06451612903225806,0.0941797200104344,0.0,0.21666666666666667
+Ni,WT/reference (full),molecular_weight_Da,657,38937.599573059444,27114.09229114452,24870.277900000005,15803.043599999995,60890.174900000195,9049.837699999996,138461.7130000007
+Ni,WT/reference (full),net_charge_pH7_4,657,-11.406012049401001,8.476514075696667,-10.675186860662173,-16.34818370744476,-6.642544224721824,-50.63535184763144,29.893302208509866
+Ni,WT/reference (full),gravy_KyteDoolittle,657,-0.1914087666458152,0.22978860657998554,-0.1661654135338346,-0.3254054054054056,-0.056179775280898854,-1.0196531791907517,0.7126086956521743
+Ni,WT/reference (full),instability_index,657,36.69795705742055,10.86026519405585,34.26344827586205,30.09354275741709,41.79953051643192,2.4646076794657734,84.0031719532561
+Ni,WT/reference (full),isoelectric_point,657,5.622122509113003,0.9833549580541419,5.437638664245605,5.045449638366699,5.794644355773925,4.0500284194946286,10.522094535827637
+Ni,WT/reference (full),aromaticity_fraction,657,0.06903329614637067,0.023885101007221323,0.06451612903225806,0.05272407732864675,0.08482142857142858,0.017094017094017092,0.19148936170212766
+Cu,Generated (gated n=500),molecular_weight_Da,500,5735.773334799998,1839.599909094262,5652.8439499999995,4361.48805,7027.385850000001,3207.6695,9530.235899999992
+Cu,Generated (gated n=500),net_charge_pH7_4,500,2.177348899931101,2.9945467305726434,1.8331160455255957,-0.03857126382054177,3.9026786485555056,-6.729755249327814,12.011673374838622
+Cu,Generated (gated n=500),gravy_KyteDoolittle,500,-0.6777146190476191,0.47979901270675285,-0.7075000000000002,-0.9887500000000007,-0.4314999999999999,-2.133333333333334,0.9979999999999991
+Cu,Generated (gated n=500),instability_index,500,45.18942392857143,21.356961236948408,44.536833333333334,30.851000000000003,58.26525000000001,-13.610000000000003,121.28
+Cu,Generated (gated n=500),isoelectric_point,500,8.550102472686767,1.3040879626857107,8.79901580810547,7.377226543426513,9.51532402038574,4.397996711730957,11.999967765808105
+Cu,Generated (gated n=500),aromaticity_fraction,500,0.06867357142857142,0.042392192246260334,0.06666666666666667,0.0375,0.1,0.0,0.2666666666666667
+Cd,Generated (gated n=500),molecular_weight_Da,500,5493.7174532,1732.143948919548,5334.160750000001,4171.450199999999,7053.215525000005,3003.5239999999994,8905.773500000003
+Cd,Generated (gated n=500),net_charge_pH7_4,500,4.51393227305999,3.60080299668978,4.246191182297465,2.1320600271788384,6.3576947067504355,-3.611645777855447,18.01270896203495
+Cd,Generated (gated n=500),gravy_KyteDoolittle,500,-0.15216604761904762,0.4158202977951427,-0.17124999999999996,-0.4129761904761904,0.08200000000000005,-1.4033333333333333,1.2699999999999996
+Cd,Generated (gated n=500),instability_index,500,65.46233939999999,20.6266938625509,65.34999999999998,52.273333333333326,77.97723214285715,-4.176666666666666,142.56
+Cd,Generated (gated n=500),isoelectric_point,500,8.404394290161134,0.887774666714031,8.611477279663086,8.214319801330568,8.865547180175781,4.0500284194946286,10.113816261291504
+Cd,Generated (gated n=500),aromaticity_fraction,500,0.02346476190476191,0.02711355897301278,0.016666666666666666,0.0,0.03333333333333333,0.0,0.15
+Zn,Generated (gated n=500),molecular_weight_Da,500,6605.982959399998,1756.9879616805197,6826.571400000001,4912.413500000001,8093.869899999996,3269.275499999999,9678.349299999998
+Zn,Generated (gated n=500),net_charge_pH7_4,500,6.555086110494916,4.7360009669730445,6.299141294665338,3.2239116039552886,9.547982906097477,-7.714779902996944,24.33435124373645
+Zn,Generated (gated n=500),gravy_KyteDoolittle,500,-0.43800007142857145,0.5332998416900696,-0.43499999999999994,-0.8028749999999999,-0.0946874999999999,-2.1100000000000003,1.2449999999999994
+Zn,Generated (gated n=500),instability_index,500,46.64874702142856,17.700449944911732,46.44541666666666,34.27791666666667,58.32916666666666,6.147499999999999,110.19666666666666
+Zn,Generated (gated n=500),isoelectric_point,500,9.02513862915039,0.9833703279437971,9.169870185852053,8.663841724395754,9.63151216506958,4.0500284194946286,11.389581108093264
+Zn,Generated (gated n=500),aromaticity_fraction,500,0.07094833333333334,0.0375784126587158,0.06666666666666667,0.04285714285714286,0.1,0.0,0.2
+Ni,Generated (gated n=500),molecular_weight_Da,500,5514.623654799999,1466.8061038915648,5662.897049999998,4539.1118,6815.410549999999,3126.432099999999,9563.039099999993
+Ni,Generated (gated n=500),net_charge_pH7_4,500,-0.46191225026998384,2.490171046797122,-0.46541276040318746,-2.2824402144574085,0.7739741389243606,-8.348135432634312,10.401790276072498
+Ni,Generated (gated n=500),gravy_KyteDoolittle,500,0.2203628571428572,0.3181122698122251,0.22549999999999995,0.014000000000000234,0.4293750000000003,-1.063333333333334,1.1649999999999998
+Ni,Generated (gated n=500),instability_index,500,16.09240997142857,12.640068138370005,15.558000000000002,7.160089285714286,24.033437499999998,-14.759999999999998,63.05
+Ni,Generated (gated n=500),isoelectric_point,500,7.335778021240234,1.746538088466558,6.961150932312011,6.013289737701415,8.631720352172852,4.0500284194946286,11.999967765808105
+Ni,Generated (gated n=500),aromaticity_fraction,500,0.18233738095238095,0.06417890948974825,0.18333333333333335,0.14,0.225,0.02,0.3666666666666667

Computational/validation/sampling_smoke.slurm

@@ -0,0 +1,18 @@
+#!/bin/bash
+#SBATCH --job-name=zenodo_smoke
+#SBATCH --partition=b200-mig45
+#SBATCH --qos=mig-max
+#SBATCH --nodes=1
+#SBATCH --ntasks=1
+#SBATCH --gres=gpu:1
+#SBATCH --cpus-per-task=6
+#SBATCH --mem=48G
+#SBATCH --time=00:30:00
+#SBATCH --output=/vast/projects/pranam/lab/yz927/projects/metalmd2/06_analysis/manuscript_binding_20260921/zenodo_draft_20260925/validation/sampling_%j.log
+set -euo pipefail
+cd /vast/projects/pranam/lab/yz927/projects/metalmd2/06_analysis/manuscript_binding_20260921/zenodo_draft_20260925/package
+export PYTHONDONTWRITEBYTECODE=1
+export OMP_NUM_THREADS=6
+export USE_TF=0
+/vast/projects/pranam/lab/yz927/envs/metal/bin/python -u tools/sample.py --mode direct --target_idx 3 --seq_len 30 --batch 1 --desired 1 --max_rounds 1 --best_decode 16 --output ../validation/sample_direct.csv
+/vast/projects/pranam/lab/yz927/envs/metal/bin/python -u tools/sample.py --mode prototype --target_idx 1 --seq_len 30 --batch 1 --desired 1 --max_rounds 1 --best_decode 16 --t_start 15 --proto MDPETCPCPSGGSCTCADSCKCEGCKCTSCKKSCCSCCPAECEKCAKDCVCKGGEAAEAEAEKCSCCQ --output ../validation/sample_prototype.csv