Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+weight/esm-main/examples/inverse_folding/illustration.png filter=lfs diff=lfs merge=lfs -text
+weight/esm-main/examples/lm-design/utils/ngram_stats/quadgram_seg.p filter=lfs diff=lfs merge=lfs -text
+weight/esm-main/examples/lm-design/utils/ngram_stats/trigram_seg.p filter=lfs diff=lfs merge=lfs -text

.mv

@@ -0,0 +1 @@
+Revision:master,CreatedAt:1784688787

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.

README.md

@@ -0,0 +1,283 @@
+---
+license: mit
+tasks:
+  - protein-structure-prediction
+frameworks:
+  - pytorch
+language:
+  - en
+  - zh
+tags:
+  - OneScience
+  - Life Sciences
+  - Protein Language Model
+  - Protein Structure Prediction
+  - Variant Effect Prediction
+  - ESM
+modelscope:
+  model: OneScience/ESM
+  data_path: data/
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">ESM</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+ESM (Evolutionary Scale Modeling) is a family of protein language models released by Meta AI / FAIR. It can be used for protein sequence representation, structure prediction, variant effect assessment, and fixed-backbone sequence design.
+
+Paper: Evolutionary-scale prediction of atomic-level protein structure with a language model  
+https://www.science.org/doi/10.1126/science.ade2574
+
+# Model Description
+
+This model package integrates PyTorch inference capabilities for ESM-1, ESM-2, MSA Transformer, ESMFold, ESM-1v, and ESM-IF1, and provides DCU runtime adaptation. The accompanying sample data is released with the ModelScope model package `OneScience/ESM`.
+
+# Applicable Scenarios
+
+| Scenario | Description |
+| :---: | :--- |
+| Protein sequence representation extraction | Input a FASTA file and output per-token, mean, BOS, or contact representations |
+| Protein structure prediction | Input one or more amino acid sequences and output the corresponding PDB structure files |
+| Variant effect scoring | Input a wild-type sequence and a DMS mutation table, and output mutation impact scores |
+| Fixed-backbone sequence design | Input a PDB / CIF structure and chain ID, and sample candidate sequences satisfying backbone constraints |
+| Structure-conditioned sequence scoring | Input a structure and candidate sequences, and compute conditional log-likelihood |
+| ModelScope / OneCode runtime | After downloading the model project, quickly verify script connectivity in a biology-domain runtime environment |
+
+
+# Usage Instructions
+
+## 1. OneCode Usage
+
+You can experience intelligent one-click AI4S programming through the OneCode online environment:
+
+[Click to experience intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware Requirements**
+
+- GPU or DCU runtime is recommended.
+- CPU can be used for import checks and small-configuration connectivity verification, but full training and inference are slow.
+- DCU users need to install DTK in advance. DTK 25.04.2 or later is recommended, or the OneScience-recommended version matching the current cluster.
+
+
+
+
+
+**Environment Check**
+
+- NVIDIA GPU:
+
+```bash
+nvidia-smi
+```
+
+- Hygon DCU:
+
+```bash
+hy-smi
+```
+
+### Download the Model Package
+
+```bash
+modelscope download --model OneScience/ESM --local_dir ./ESM
+cd ESM
+```
+
+This model package already contains a small amount of sample data that can be used directly for default workflow verification.
+
+### Install the Runtime Environment
+
+**DCU Environment**
+
+```bash
+# First activate DTK and Conda
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# Supports uv installation
+pip install onescience[bio-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+```bash
+# If libraries cannot be found, activate CUDA as shown below
+source ${ROCM_PATH}/cuda/env.sh
+export LD_LIBRARY_PATH="$CONDA_PREFIX/lib:$LD_LIBRARY_PATH"
+export LD_LIBRARY_PATH="$CONDA_PREFIX/lib/python3.11/site-packages/fastpt/torch/lib:$LD_LIBRARY_PATH"
+```
+
+After installation, return to the model package directory:
+
+```bash
+cd ./ESM
+```
+
+### Training and Inference Data Introduction
+
+The FASTA, PDB / CIF, and DMS data used by the ESM examples have been released with the [ModelScope model package OneScience/ESM](https://modelscope.cn/models/OneScience/ESM). After downloading the complete model package, they can be used directly from the `data/` directory. This model package does not include a training entry point; the data is used for example inference and workflow verification. You can also download only the data directory:
+
+```bash
+modelscope download --model OneScience/ESM ESM/data --local_dir ./data
+```
+### Training Weights
+
+The repository includes multiple ESM weights under `weight/`, which can be selected for inference as needed.
+
+### Prepare Weights
+
+Place the required ESM weights in the following directory:
+
+```text
+weight/
+  checkpoints/
+    esm2_t6_8M_UR50D.pt
+    esmfold_3B_v1.pt
+    esm1v_t33_650M_UR90S_1.pt
+    esm_if1_gvp4_t16_142M_UR50.pt
+    ...
+```
+
+If using a shared runtime directory, you can also specify the weight location through an environment variable:
+
+```bash
+export ESM_WEIGHT_DIR=/path/to/esm/weight
+```
+
+The default example reads:
+
+- `weight/checkpoints/esm2_t6_8M_UR50D.pt`
+
+The corresponding weights must be available for the ESMFold, ESM-1v, and ESM-IF1 examples.
+
+### Default Example
+
+```bash
+bash scripts/infer.sh
+```
+
+The default example reads `data/fasta/few_proteins.fasta`, uses `esm2_t6_8M_UR50D.pt` to extract protein representations, and saves the results to `outputs/embeddings/`.
+
+### Sequence Representation Extraction
+
+```bash
+python scripts/extract.py \
+  weight/checkpoints/esm2_t6_8M_UR50D.pt \
+  data/fasta/few_proteins.fasta \
+  outputs/embeddings \
+  --include mean per_tok \
+  --repr_layers 6
+```
+
+### ESMFold Structure Prediction
+
+```bash
+python scripts/fold.py \
+  -i data/fasta/few_proteins.fasta \
+  -o outputs/pdb \
+  --model-dir weight \
+  --cpu-only
+```
+
+The output directory will generate one or more `.pdb` files. For formal GPU / DCU inference, you can remove `--cpu-only` and set `--chunk-size` or `--max-tokens-per-batch` according to available memory.
+
+You can also explicitly enable ESMFold through the default script:
+
+```bash
+RUN_ESMFOLD=1 bash scripts/infer.sh
+```
+
+### Inverse Folding Sequence Sampling
+
+```bash
+python scripts/inverse_folding/sample_sequences.py \
+  data/inverse_folding/5YH2.pdb \
+  --chain A \
+  --outpath outputs/sampled_seqs.fasta \
+  --num-samples 1 \
+  --nogpu
+```
+
+### Inverse Folding Sequence Scoring
+
+```bash
+python scripts/inverse_folding/score_log_likelihoods.py \
+  data/inverse_folding/5YH2.pdb \
+  data/inverse_folding/5YH2_mutated_seqs.fasta \
+  --chain A \
+  --outpath outputs/sequence_scores.csv \
+  --nogpu
+```
+
+### Variant Effect Prediction
+
+Variant effect prediction requires providing a wild-type sequence matching the DMS mutation column:
+
+```bash
+python scripts/variant_prediction/predict.py \
+  --model-location esm1v_t33_650M_UR90S_1 \
+  --sequence "${ESM_VARIANT_SEQUENCE}" \
+  --dms-input data/variant_prediction/BLAT_ECOLX_Ranganathan2015.csv \
+  --mutation-col mutant \
+  --dms-output outputs/variant_prediction.csv \
+  --offset-idx 24 \
+  --scoring-strategy wt-marginals
+```
+
+# Data Format
+
+Sample data is stored under `data/` by default:
+
+```text
+data/
+  fasta/
+    few_proteins.fasta
+    some_proteins.fasta
+  inverse_folding/
+    5YH2.pdb
+    5YH2.cif
+    5YH2_mutated_seqs.fasta
+    example.json
+  variant_prediction/
+    BLAT_ECOLX_Ranganathan2015.csv
+    rho_pp.csv
+    aggregated_rho.csv
+    aggregated_rho_round3.csv
+```
+
+Where:
+
+- FASTA files are used for sequence representation extraction and structure prediction.
+- PDB / CIF files are used for inverse folding sampling and structure-conditioned sequence scoring.
+- The variant effect prediction CSV must contain a mutation column. The default example column name is `mutant`, and the mutation format is like `A123B`.
+- Custom DMS data must ensure that the wild-type sequence provided by `--sequence` is consistent with the wild-type amino acids in the mutation column.
+
+# Verification
+
+Static import check:
+
+```bash
+python scripts/check_import_boundaries.py
+```
+
+Syntax check:
+
+```bash
+python -B -c "import ast, pathlib; [ast.parse(p.read_text(encoding='utf-8'), filename=str(p)) for root in ['model', 'scripts', 'tests'] for p in pathlib.Path(root).rglob('*.py')]"
+```
+
+# Official OneScience Information
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citations and License
+
+- This repository is based on the ESM open-source model and provides DCU adaptation.
+- ESM-related source code uses the MIT License; see `LICENSE`. For the terms of use for model weights and data, refer to the corresponding publisher's instructions.
+- For scientific research, cite the corresponding original ESM paper according to the specific submodel. For ESM-2 / ESMFold, cite: [Evolutionary-scale prediction of atomic-level protein structure with a language model](https://www.science.org/doi/10.1126/science.ade2574).

configuration.json

@@ -0,0 +1,2 @@
+{"framework":"Pytorch","task":"other"}
+

data/fasta/few_proteins.fasta

@@ -0,0 +1,6 @@
+>UniRef50_UPI0003108055
+MPADAREYLESKHATRRFDRPAEVAGVVAFLLSDDTSFVIGAGYLVDGGYTALRAARGRLGPRRQAAPAVKLLPNTDSPR
+>UniRef50_A0A223SCH7
+MGPPRWWKGITGLAAVVHRADPEDKADLYAKMGLYLEYHPETRIVEARIKPRLHDVCESKVSEGGLEPPCP
+>UniRef50_A0A090SUK6
+MKTPEDRVFRATDEYSDFVMACRYKGNEREFVIASHDKLNEAQVETLTSYLSGEWFKKTYITGIMNDSDGVLSQHEEYGDEVFCQPLDELRVDRYIMMV

data/fasta/some_proteins.fasta

@@ -0,0 +1,30 @@
+>UniRef50_A0A1E3NP16
+AVIYYRFRSQKPDHIATIKFDGTGLTVFELKRDIILANNLLHSTDVDIVLYSTEDIQDTKSWGYQNGGSSSAGERELDDDNEVVPRSTTVLVRRTMTPKKNKGNVQRYVAGKPRLQVSGTNSVNKSISLGNNVGGTMNFGDAATNGDEDDMIKKMFSVQDEQWSQQQDVMATATRVDNFRTNVNEPVPEYYICYKCGEKGKHHIKNCPKNNDPNWEGVRVRKTTGIPKSHLKAIENPEDTIRDSNSSGNTTYMVNDEGKYVVAVADTKAWEKYQKTKKGESGGYLNGDVDVDDGELKDPETGKLWKSPVRIPCCNKIFSRKIIEDKLIDSDFTCPSCGKEQIYLDTLVADEELQAKVDEYVKNLSENKNNDGNSPKRRQVNPAGATANTSQLPQIPMMPMPPINMQMPPMNIGMPPFMPFMPMPGMNP
+>UniRef50_UPI000836A30F
+MTLRTLLALSILALAAAATVQARPGAPPCSPLGLKYQPGACEKWKREHPDVNPDGVVQTVTIVNNSSSVLGGYVTFWHANNEHTDVDLPGVKPGETWTANGSWTVGQAPYYLLYSSFQDSYGDPVTFYAVPISKYPALAKKLPPEPDNCQSNHFRMVFGDGPQYVYEQHSGAVVTGGQTDKNTLCQILGCPTGGSTAGGLNTQNRTTSRSAPQPVYFRSCDP
+>UniRef50_A0A2G8L8Y3
+MAQKRYEENLSPYAELFRSKLIEDFHLLESFDEHGKSDAPKYYSKDFEDPARQDKMMLENPHGLVKFQVYSRKEPGEHFMLVLILSNSIALGLQAEVSESDDPKFAGLKLALDIFDYCSLFLFMVEIILKWIDNFWSFWSDNWNIFDFAVTVGSFVPEIINFFAGDIGGSMVRVIVRNLRVFRILRSLKMVSRFRQVRLIALAIGKAFSAITFIMLLLFTFLYIFAITGIIFFDTYTRSERQDLKYKDSFRSLPRAMITLFQLFTLDQWYKLLNDMWKVMDSMIPLGYIILWICIGSFIFRNVFVGIMVNNFQSIRNDLFEEVKEQEAARQIIQDTEKFNEELSRQEKKLNANRRGTLYQSPTVQPPKPNQPQPSQLAGLDNSETDEQSVSQEDESNTDGQTSLSGTDSYDLLGESSDSLFRRSSDGMIDKDKLSTNWEKTVHDNLTLLTSTPSETLWPRDTLFRYFQLMESLMENLQERQDLQDLAYHSLLQIFDSFDTSA
+>UniRef50_A0A1D5ZRM3
+MPCVAHECHPRLPAANHCRSLSCLGTPAAGWSSGDDDREEDELDTKQVILNEMRNREMRKRSSRCSVDSPTLSGAFAWSFTPLHPRSSIEKVSCTEEEKEAASDSDNESEAFFSVKSFFTRSTSRAATVASSTDMDPPATWEGLRGCEGWPFGLCP
+>UniRef50_UPI0003108055
+MPADAREYLESKHATRRFDRPAEVAGVVAFLLSDDTSFVIGAGYLVDGGYTALRAARGRLGPRRQAAPAVKLLPNTDSPR
+>UniRef50_A0A223SCH7
+MGPPRWWKGITGLAAVVHRADPEDKADLYAKMGLYLEYHPETRIVEARIKPRLHDVCESKVSEGGLEPPCP
+>UniRef50_A0A090SUK6
+MKTPEDRVFRATDEYSDFVMACRYKGNEREFVIASHDKLNEAQVETLTSYLSGEWFKKTYITGIMNDSDGVLSQHEEYGDEVFCQPLDELRVDRYIMMV
+>UniRef50_V4AGU2
+MHLGSIYLMVVLLIYFAYTDDRKERENVDVINPEENLVQDDEQYVGDSTENIEKSGSEEEEEDKEAIEEEEDEEEELNYRYIPEPAQDIVDNGKKYIQVHCSFKQDESLIRHPSNCSRYFVCSYGVVEEMPVCDDGEVFSIQVSECVKKGSENDDCDKLPFDSPPEITRGTQPSLIWHPRHKSPRSQFRQPTTLQMKVPHIELESFTCSATGKVLSHHSENCAWYYNCSAHPDAVMQTFYSGFIMECPYPQLFSTETKQCEDFEDVKCGDRYEPKSPCDYRANHCHETSHCIPCWVRYASCLELPDGLNPWSELEWKPFFVECYKERTVFQGVCDKSAVFSPLTRACETPYSIPRQHGGWRPVCDGRRDGIYADEYGRCDIYYVCKGYIFTGFFRCEKGEMFNPVISICQKPEAVPYPCGDLEMPNICESSLNGYHLDMFGRCTHYFECKDQQLEGISMCPSGIFNPELQICESSRDQPKPCGNLTNLCTHKNDGFHSDENDCTKVFQCERGLTMTSYDCSGSVRTECDVCNTPTECNDKPNGLYPNLKEGVGYYYDCVRSQIQNHYKCDKEKGGPIFNPVKQRCFYPEDLCKEVFSLKIAW
+>UniRef50_R5PKX9
+MQGGEQDVFEHRQVREKVIALKNHADAAAQGAAEFERLSFEQDVAALDGFETDQAAQKRGFAAARRPENHRDFFVVKREVDAVENHSVAELLYETARFQNNVIGHFYAFHFFSRALAASDTGQHARK
+>UniRef50_UPI0003674933
+MKKKEDILNLEHTLLPWMGRTMKVLDYFIGDFLNLKGIELTKVQWILLKKLNEQNGQPQQNLAFLTNRDKASLARLITTMEKKKLVERIPSKIDGRINHIFITKHGCEILQKSAPVIEKVVGCIQEGISPEEIETVIKVMQKVNNNINRASN
+>UniRef50_E9K9Y4
+MADNVLMAYHIVHDPDERAKHVLNTKKLYKWRITEKTKGTPVVGNVALVQTQFAKRTPVMIYATKEVANDLSDLQPVKVFTNNRDQETVNQTFDDLMR
+>UniRef50_A0A2C9LWN7
+MGKITGLLFFLFTSVRVTPSMKNTVNDIYRVRKDVLTKYRNTEEYDLGQRKGPLKQMDTRETRTPYGHFNSDTAFISPKNILKATRPLAFHLGNLKNQTSPCSTWNLTVDCSYKSLDHIESSWFPSNTTVLLLNNNKLVTLHNETFAQLTNLTRLDLSSNDIRRIDAGAFQGLHNLQELNLHMHCCNFTDHYSLESVFAPLRNLRILNAMHNSDVGVLTYSYTFLTRLPLLQSLSIDFDLDTLYCGPEFNDLKNLTFLQFSGQVMYIDDRSFQNVAQLKNLSMDHLSNINNISHNAFKPLSNLKVLTMYHVLLYVQEILSLLEPFQGRNMTEITLDTTTRTLTQVNPTRNGILTNHDTKYLMNICLESFTLIDNRIFYIKPDAVQNIYTWKKCLMHLYIASNPIQGNNFALIRLFTLDNLKSFTFINMFRACHEFQPFPQSSPPATRNVASSSISSQNNHYQKDTTSNQQQMIRHSPFSPYLDMIDENPYQLNIPNYIFISPSLQYMNFQRLVMSQSFEYHFILVGAQNVTSLDISDSGFYRFNGLMEGVSAIKTLIISGNDVSVLSVSFFDTFVSLENFAISSCKLDRDFISLNSRRIFQNHTRLQELDISSNSLNYLSQNTFSYNNRLMWLNMSGNQFKDIPFDLTNTPELQFLDIRFNSLTTIDETTAQQMDHLVTKSGKLEILLEGNVLSCSCSDLSFMRWMRMTLVTFDQNGNFTCMNTDGERKYTLDYSNLDSLWRECWGSFFLYFALIMLCLYCIGVFAVFITMRNKNFIVSFFLQLFGGFKLHSRRDYPVGVYIGYSDKDYQFPCKELRSFIESSLKLKTFLIDRDLIASVDKASGIIEALNASWRILLVCSKSFLKEDDWSMFTMRSAIYTQTPANPARVVVLVHKDCLPLLPPALLSSVNDENICAVSEWAMNYEMMQMLTTRLH
+>UniRef50_Q9REE6
+MSLRGRELLTSEERLELVRIPEDISEQELGRNFTLSNFDLELIKNRRRDYNRLGFAVQLCVLRFPGWSLNDAEPIPKKVLQHLARQLHVDPDCFSLYSSREA
+>UniRef50_A0A226D4M8
+MQKVINFPIWRRYYFSECGNCNDFRPDGVKKVHPPQEKSRTMDDEILAAPEVTIPFEPSDPSEVVVNLISSEEEDDDDVIQIVEEKSVDKAERQRRRQKKKDLWAARKLQRNKGQNVPLQTAWQRGPRPQETSSFPTPPQQSGGAQQKPLSPILISTAGSPNTSGAPTPANVSQQPTPTPSFVHTTASTSTHPEISLNINSDLALLIRLGPDGRPILTRVENVEQNNTSTSTPTTTRKLPPAPPPPKISFDTQTGESLLNGELITRPIIDITTDSPPSITHAATVSPQTSRTSGPPTLSPISPPPRTTQSNPAHPPPRYEPRKSRHPPRDPLASSSSSSSSSPSPPPTSRARHTSSANIPPPLEPLFLNTTQLIHLIKTCRKCEKSFPTRCDGVIHQKKEHNRKHCPVCFLTLSRHGNTYKDHLNMYHALEGDKEMVVCPFCAVEHSFDGLYNHIGRSHLIPVESKGESEHEVIFSVAPSPQSNGHQTRSVTQGNTPPKNDTNSPPNLEKRRPGPASKTRKTTNDPVPSTSRTGLICHKYVPDVETPPSKAGRNFESRAGRNFESRNVYPPATNRLNPPRNKSPPRNKNLPRNKSPPRNKSLPRNKTPPPSSSRSSSSRSVSNNLRRKNPTPPPPTQPPPKKVAKPDEAGINEKIQAAIKAVNARVHIERSVHHPNNRSDREIPSTSRTVTSRHKVPTSKTSGNTSVRKDPSPPPPLQTPPKTTNLELSKVQKARIVEDVRSIVRDVRITRVLDDGEVPSTSRAVEEEKKKEEKKNETRARLPRSSRVGERSSGYFQMAEGIDFSPENPTPRSKDLKAIHISKLNLIYQQLKCYPDTSVIANVAKECGVEIPVVAKWFTKKHMEYCQKTQQRKRKRKPPELR
+>UniRef50_UPI000B82D8F0
+MGGLHLIELRNVNIEFDKKLIEDGTIKIYDGKITAIIGESGSGKTSLLYLLGLISSNHRYLYSFDDVTLDLSNDFEMSRIRKQKIGYIFQDNNLVENLTIFENIRLSATIAGINITDKEIKSYLEFVELGYIDSNHYPRKLSGGERQRVAIACALAKQPELILADEPTSALDTVNSEIIMGIFKKIAHKDKKKIVIATHNDRIYNEADVIYEIKNNKIQLVKGESSNESSKKEPEDYDNNVKLTPRFYFDYAIKTSRKGRFAKNLMIVLSAIAIAFSSVMYNFGDTFVKEQEKLMDAISDKEIFVVNMTAPLNTILDIDENLSIQDKDAELLRNISYVDTIYPYFEFRSIGYPLINETEASEGYIVVSKGQKEEKYTFAESKDNPYDKYVIIPYYPEQNLERRLKEKLSDESSDKVYISSQLAQLLGIENLKESVSLRVLTYVPIAQHETQMTVRPEGIVYEIDIDLSKVVELDLKIEGILDESVRNRYSNSGNNAIYVPYHKMQMILTQTQNSATIDTNLEYIEWRPSAFVVFAKSYNDVGLVIERVSSINPNFRAVSEYQDIESMNAIVKNTREIGLVIVIVILIIIFLLMSIIHMNHILDRKYEISLLKANGLTKIELTKLVSVESLRHVFLVSLISSVISLVVTKVMNLLFEEIA

data/inverse_folding/5YH2_mutated_seqs.fasta

@@ -0,0 +1,8 @@
+>mut_1
+SVLQSLFEHPLYRTVLPDLTEEDTLFNLNAEIRLYPKAASESYPNWLRFHIGINRYELYSRHNPVIAALLRDLLSQKISSVGMKSGGTQLKLIMSFQNYGQALFKPMKQTREQETPPDFFYFSDFERHNAEIAAFHLDRILDFRRVPPVAGRLVNMTREIRDVTRDKKLWRTFFVSPANNICFYGECSYYCSTEHALCGKPDQIEGSLAAFLPDLALAKRKTWRNPWRRSYHKRKKAEWEVDPDYCDEVKQTPPYDRGTRLLDIMDMTIFDFLMGNMDRHHYETFEKFGNDTFIIHLDNGRGFGKHSHDEMSILVPLTQCCRVKRSTYLRLQLLAKEEYKLSSLMEESLLQDRLVPVLIKPHLEALDRRLRLVLKVLSDCVEKDGFSAVVENDLD
+>mut_2
+DVLQSLFEHPLYRTVLPDLTEEDTLFNLNAEIRLYPKAASESYPNWLRFHIGINRYELYSRHNPVIAALLRDLLSQKISSVGMKSGGTQLKLIMSFQNYGQALFKPMKQTREQETPPDFFYFSDFERHNAEIAAFHLDRILDFRRVPPVAGRLVNMTREIRDVTRDKKLWRTFFVSPANNICFYGECSYYCSTEHALCGKPDQIEGSLAAFLPDLALAKRKTWRNPWRRSYHKRKKAEWEVDPDYCDEVKQTPPYDRGTRLLDIMDMTIFDFLMGNMDRHHYETFEKFGNDTFIIHLDNGRGFGKHSHDEMSILVPLTQCCRVKRSTYLRLQLLAKEEYKLSSLMEESLLQDRLVPVLIKPHLEALDRRLRLVLKVLSDCVEKDGFSAVVENDLD
+>mut_3
+EVLQSLFEHPLYRTVLPDLTEEDTLFNLNAEIRLYPKAASESYPNWLRFHIGINRYELYSRHNPVIAALLRDLLSQKISSVGMKSGGTQLKLIMSFQNYGQALFKPMKQTREQETPPDFFYFSDFERHNAEIAAFHLDRILDFRRVPPVAGRLVNMTREIRDVTRDKKLWRTFFVSPANNICFYGECSYYCSTEHALCGKPDQIEGSLAAFLPDLALAKRKTWRNPWRRSYHKRKKAEWEVDPDYCDEVKQTPPYDRGTRLLDIMDMTIFDFLMGNMDRHHYETFEKFGNDTFIIHLDNGRGFGKHSHDEMSILVPLTQCCRVKRSTYLRLQLLAKEEYKLSSLMEESLLQDRLVPVLIKPHLEALDRRLRLVLKVLSDCVEKDGFSAVVENDLD
+>mut_4
+VVLQSLFEHPLYRTVLPDLTEEDTLFNLNAEIRLYPKAASESYPNWLRFHIGINRYELYSRHNPVIAALLRDLLSQKISSVGMKSGGTQLKLIMSFQNYGQALFKPMKQTREQETPPDFFYFSDFERHNAEIAAFHLDRILDFRRVPPVAGRLVNMTREIRDVTRDKKLWRTFFVSPANNICFYGECSYYCSTEHALCGKPDQIEGSLAAFLPDLALAKRKTWRNPWRRSYHKRKKAEWEVDPDYCDEVKQTPPYDRGTRLLDIMDMTIFDFLMGNMDRHHYETFEKFGNDTFIIHLDNGRGFGKHSHDEMSILVPLTQCCRVKRSTYLRLQLLAKEEYKLSSLMEESLLQDRLVPVLIKPHLEALDRRLRLVLKVLSDCVEKDGFSAVVENDLD

data/variant_prediction/aggregated_rho.csv

@@ -0,0 +1,17 @@
+,valid,valid,valid,full,full,full,test,test,test
+,rho,rho_boot_mean,rho_boot_std,rho,rho_boot_mean,rho_boot_std,rho,rho_boot_mean,rho_boot_std
+prediction_name,,,,,,,,,
+DeepSequence (published),0.5602602238441781,0.5612750863332869,0.011985038405770958,0.5136039187155506,0.512778392574446,0.017143498537203965,0.4985534977063157,0.4971342978135294,0.01880751793444042
+DeepSequence (published) - single,0.5430017493589612,0.5443701641109099,0.012423838414787655,0.5020580499120493,0.5011407877592315,0.017404933590781577,0.4888504049291744,0.4871958276457871,0.01901173848626349
+DeepSequence (replicated),0.5635921424785667,0.5646995390005713,0.01208513622243054,0.5201346211408833,0.5192886065212002,0.01673076687983822,0.5061160658706628,0.5046399186246291,0.018229357414485857
+ESM-1b,0.5666981361139574,0.5676141841913277,0.012025131590764206,0.4585307340241866,0.45722552072809847,0.019420008294609183,0.42363802367264763,0.42161627444963734,0.021805452392623693
+ESM-1v (+further training),0.5979754992703865,0.5985949225433556,0.011700006848349927,0.5382829634571142,0.5368211143294155,0.01685153619316537,0.5190273067431552,0.516894079421693,0.018513319852783258
+ESM-1v (zero shot),0.593067646900253,0.5941859676949833,0.012216511625608212,0.5089550689778174,0.5077512066859666,0.01778501791633543,0.48182197932541876,0.47986902571531576,0.019581310268182918
+EVMutation (published),0.5490624796568302,0.5497402994229479,0.012972877244416833,0.508419480578392,0.5068193077204962,0.017724141740786323,0.49530883571437995,0.4929738265261569,0.01925680770735712
+EVMutation (replicated),0.5512073418060958,0.5516069925840847,0.012682710408380227,0.5112223872556211,0.509417510409684,0.01815069563293862,0.49832401481998395,0.4958080000308451,0.019914561834409074
+MSA Transformer,0.5987868160929384,0.5996093743584658,0.011619082338376548,0.5422875260142875,0.5414621273787668,0.016736474302820666,0.5240619485695612,0.5227049509337023,0.01838724590425425
+PSSM,0.4633509799617207,0.46414678680632565,0.013608999290250463,0.4604575410937561,0.45896520887271747,0.018669177520973506,0.45952417371699333,0.4572937321199404,0.020301493079271263
+ProtBERT-BFD,0.5166978551105863,0.5173010932587435,0.0135134115364477,0.42751993235504165,0.4278749478751994,0.020253227983813266,0.39875286049841424,0.39902780420308837,0.022427362321673118
+TAPE,0.16039305189488504,0.16335386054336812,0.015534139865906533,0.17144437052454356,0.1748040435344602,0.023020918331052527,0.17500931201798176,0.1784976509509415,0.025436008158518972
+UniRep,0.17009503694214162,0.1687392944334104,0.018385140300946963,0.15554362973563243,0.15406970922460442,0.022971233340804673,0.15084962741095204,0.14933758496369925,0.024450618192371674
+gt,1.0,1.0,0.0,1.0,1.0,0.0,1.0,1.0,0.0

data/variant_prediction/aggregated_rho_round3.csv

@@ -0,0 +1,17 @@
+,valid,valid,valid,full,full,full,test,test,test
+,rho,rho_boot_mean,rho_boot_std,rho,rho_boot_mean,rho_boot_std,rho,rho_boot_mean,rho_boot_std
+prediction_name,,,,,,,,,
+DeepSequence (published),0.56,0.561,0.012,0.514,0.513,0.017,0.499,0.497,0.019
+DeepSequence (published) - single,0.543,0.544,0.012,0.502,0.501,0.017,0.489,0.487,0.019
+DeepSequence (replicated),0.564,0.565,0.012,0.52,0.519,0.017,0.506,0.505,0.018
+ESM-1b,0.567,0.568,0.012,0.459,0.457,0.019,0.424,0.422,0.022
+ESM-1v (+further training),0.598,0.599,0.012,0.538,0.537,0.017,0.519,0.517,0.019
+ESM-1v (zero shot),0.593,0.594,0.012,0.509,0.508,0.018,0.482,0.48,0.02
+EVMutation (published),0.549,0.55,0.013,0.508,0.507,0.018,0.495,0.493,0.019
+EVMutation (replicated),0.551,0.552,0.013,0.511,0.509,0.018,0.498,0.496,0.02
+MSA Transformer,0.599,0.6,0.012,0.542,0.541,0.017,0.524,0.523,0.018
+PSSM,0.463,0.464,0.014,0.46,0.459,0.019,0.46,0.457,0.02
+ProtBERT-BFD,0.517,0.517,0.014,0.428,0.428,0.02,0.399,0.399,0.022
+TAPE,0.16,0.163,0.016,0.171,0.175,0.023,0.175,0.178,0.025
+UniRep,0.17,0.169,0.018,0.156,0.154,0.023,0.151,0.149,0.024
+gt,1.0,1.0,0.0,1.0,1.0,0.0,1.0,1.0,0.0

data/variant_prediction/rho_pp.csv

@@ -0,0 +1,575 @@
+,protein_name,prediction_name,rho,rho_boot_mean,rho_boot_std
+112,BLAT_ECOLX_Ranganathan2015,gt,1.0,1.0,0.0
+113,BLAT_ECOLX_Ranganathan2015,MSA Transformer,0.7734667485457659,0.7738027379896001,0.005810229642357195
+114,BLAT_ECOLX_Ranganathan2015,PSSM,0.614150958806157,0.6140369220507402,0.00914005554519492
+115,BLAT_ECOLX_Ranganathan2015,EVMutation (published),0.7176707275279989,0.7184384489293669,0.005172838583114061
+116,BLAT_ECOLX_Ranganathan2015,DeepSequence (published) - single,0.755008978954388,0.7561275888133361,0.006397359248039358
+117,BLAT_ECOLX_Ranganathan2015,DeepSequence (published),0.7755001128993768,0.7754839719498734,0.0059058513297401
+118,BLAT_ECOLX_Ranganathan2015,ESM-1v (zero shot),0.7251281855227297,0.7252663698494137,0.005817208726713974
+119,BLAT_ECOLX_Ranganathan2015,ESM-1v (+further training),0.7713005999939375,0.7711664330412893,0.005154671622503008
+120,BLAT_ECOLX_Ranganathan2015,DeepSequence (replicated),0.7907734920246668,0.7896796440415632,0.004950863463811736
+121,BLAT_ECOLX_Ranganathan2015,EVMutation (replicated),0.7409416788010579,0.7414034481390834,0.004511986638702662
+122,BLAT_ECOLX_Ranganathan2015,ProtBERT-BFD,0.5066624675266317,0.5086811085713664,0.008432655997093638
+123,BLAT_ECOLX_Ranganathan2015,TAPE,0.1437575573041793,0.14718890398398213,0.00854664176422607
+124,BLAT_ECOLX_Ranganathan2015,UniRep,0.03205236543975054,0.03109080054090812,0.01290097463153492
+125,BLAT_ECOLX_Ranganathan2015,ESM-1b,0.7244733058050135,0.7244715086369378,0.0054093172279529814
+322,PABP_YEAST_Fields2013-doubles,gt,1.0,1.0,0.0
+323,PABP_YEAST_Fields2013-doubles,MSA Transformer,0.7092749798819215,0.7083666356138891,0.0028152329308467395
+324,PABP_YEAST_Fields2013-doubles,PSSM,0.5022876340276876,0.5009430131340434,0.003914000389196363
+325,PABP_YEAST_Fields2013-doubles,EVMutation (published),0.6334326213147348,0.6319953143501386,0.0031257311261953294
+326,PABP_YEAST_Fields2013-doubles,DeepSequence (published) - single,0.697950719553739,0.6969873459722575,0.002722620907521595
+327,PABP_YEAST_Fields2013-doubles,DeepSequence (published),0.7135955622146342,0.7126803646172054,0.0026005801571960326
+328,PABP_YEAST_Fields2013-doubles,ESM-1v (zero shot),0.6808954491918657,0.6802766119879265,0.003337373999951867
+329,PABP_YEAST_Fields2013-doubles,ESM-1v (+further training),0.6724484972163407,0.6715887930967597,0.003131790856684628
+330,PABP_YEAST_Fields2013-doubles,DeepSequence (replicated),0.6690002200834452,0.6682783362133506,0.002924928420789724
+331,PABP_YEAST_Fields2013-doubles,EVMutation (replicated),0.6370081500864611,0.6365136199984236,0.0033115051939607903
+332,PABP_YEAST_Fields2013-doubles,ProtBERT-BFD,0.7018606830370693,0.7013726492421438,0.0029126542395483735
+333,PABP_YEAST_Fields2013-doubles,TAPE,0.08509310570891564,0.0848841242307348,0.0055652064474624435
+334,PABP_YEAST_Fields2013-doubles,UniRep,0.25423830212504706,0.2528601717205464,0.005044020233400653
+335,PABP_YEAST_Fields2013-doubles,ESM-1b,0.6989229474973527,0.6986292605564589,0.0033935909980897867
+0,AMIE_PSEAE_Whitehead,gt,1.0,1.0,0.0
+1,AMIE_PSEAE_Whitehead,MSA Transformer,0.6078201305487886,0.6074995010322104,0.006655174664199262
+2,AMIE_PSEAE_Whitehead,PSSM,0.3851814980540226,0.3857114471207458,0.01419582618409421
+3,AMIE_PSEAE_Whitehead,EVMutation (published),0.5482672326862741,0.5494300599441264,0.009296902066198478
+4,AMIE_PSEAE_Whitehead,DeepSequence (published) - single,0.6192004998182683,0.6235042563787871,0.009786304100603813
+5,AMIE_PSEAE_Whitehead,DeepSequence (published),0.64337902253515,0.6466208891535095,0.008118407538893896
+6,AMIE_PSEAE_Whitehead,ESM-1v (zero shot),0.6666021525370552,0.6679182566255217,0.008613681679123195
+7,AMIE_PSEAE_Whitehead,ESM-1v (+further training),0.671863659978349,0.6738279065439714,0.008391215501243852
+8,AMIE_PSEAE_Whitehead,DeepSequence (replicated),0.6459861864504847,0.6488672280630794,0.007018822987447003
+9,AMIE_PSEAE_Whitehead,EVMutation (replicated),0.5609325769621947,0.561649692444916,0.008422893320380829
+10,AMIE_PSEAE_Whitehead,ProtBERT-BFD,0.6338505619537509,0.6354477176484321,0.009554450012056636
+11,AMIE_PSEAE_Whitehead,TAPE,0.10859786063312615,0.11076584881696006,0.015015379880043802
+12,AMIE_PSEAE_Whitehead,UniRep,0.07739084999605612,0.073724376050283,0.01240093912413262
+13,AMIE_PSEAE_Whitehead,ESM-1b,0.5853972837117107,0.5836827247268739,0.009689198977483172
+476,TIM_SULSO_b0,gt,1.0,1.0,0.0
+477,TIM_SULSO_b0,MSA Transformer,0.652792451949853,0.6557516749540702,0.01600526744858049
+478,TIM_SULSO_b0,PSSM,0.5593694164359105,0.566935653048683,0.01690630663394768
+479,TIM_SULSO_b0,EVMutation (published),0.586140631338893,0.5919185710638765,0.01641131140995354
+480,TIM_SULSO_b0,DeepSequence (published) - single,0.5440966818223795,0.5487337742437185,0.015747346163591917
+481,TIM_SULSO_b0,DeepSequence (published),0.5505948722659953,0.5565402351008099,0.014555638377778407
+482,TIM_SULSO_b0,ESM-1v (zero shot),0.6424600791408842,0.647621332390191,0.016028029413886244
+483,TIM_SULSO_b0,ESM-1v (+further training),0.6528241665901272,0.6555645155980578,0.015379852940306115
+484,TIM_SULSO_b0,DeepSequence (replicated),0.58320310154665,0.5889306881826475,0.015779505439458753
+485,TIM_SULSO_b0,EVMutation (replicated),0.5887928613337577,0.5940831887831108,0.016569303240529226
+486,TIM_SULSO_b0,ProtBERT-BFD,0.6317983649729998,0.6345885687258263,0.01602936969330677
+487,TIM_SULSO_b0,TAPE,0.5532247068671533,0.5587646824504885,0.019898766536311598
+488,TIM_SULSO_b0,UniRep,0.16614520326579557,0.17309558518510257,0.03137269514815842
+489,TIM_SULSO_b0,ESM-1b,0.6237553534772002,0.6315291160092146,0.015666927222771784
+266,KKA2_KLEPN_Mikkelsen2014,gt,1.0,1.0,0.0
+267,KKA2_KLEPN_Mikkelsen2014,MSA Transformer,0.6483744764229661,0.6501362354726463,0.008392887010519141
+268,KKA2_KLEPN_Mikkelsen2014,PSSM,0.33148417825671017,0.3321644636508253,0.00954461359833698
+269,KKA2_KLEPN_Mikkelsen2014,EVMutation (published),0.6040252893227205,0.6046576561680694,0.007611417821309178
+270,KKA2_KLEPN_Mikkelsen2014,DeepSequence (published) - single,0.6428015918045534,0.646183322009227,0.006466715772852643
+271,KKA2_KLEPN_Mikkelsen2014,DeepSequence (published),0.6622233077856657,0.6647874075249056,0.007006878882782913
+272,KKA2_KLEPN_Mikkelsen2014,ESM-1v (zero shot),0.627837107524989,0.6303585274093482,0.008597860670353218
+273,KKA2_KLEPN_Mikkelsen2014,ESM-1v (+further training),0.6540013975518588,0.6568513820339432,0.008249320892729436
+274,KKA2_KLEPN_Mikkelsen2014,DeepSequence (replicated),0.6708322568508106,0.6736802365189244,0.007635987315614941
+275,KKA2_KLEPN_Mikkelsen2014,EVMutation (replicated),0.6078971023630376,0.6084473779526628,0.009849549243896777
+276,KKA2_KLEPN_Mikkelsen2014,ProtBERT-BFD,0.5441498820889966,0.5461574522717653,0.011666264975033773
+277,KKA2_KLEPN_Mikkelsen2014,TAPE,0.19644765859472593,0.19660401466689767,0.013538150129702716
+278,KKA2_KLEPN_Mikkelsen2014,UniRep,0.24229449050660715,0.2442377698228683,0.014857522651962994
+279,KKA2_KLEPN_Mikkelsen2014,ESM-1b,0.5796639591499797,0.5828984127122095,0.009462666946528803
+182,DLG4_RAT_Ranganathan2012,gt,1.0,1.0,0.0
+183,DLG4_RAT_Ranganathan2012,MSA Transformer,0.5106330149144778,0.5128948633920939,0.020360553500470727
+184,DLG4_RAT_Ranganathan2012,PSSM,0.4603978833849677,0.4625263978112141,0.023120281422982173
+185,DLG4_RAT_Ranganathan2012,EVMutation (published),0.5352653453343902,0.539345546690929,0.02299279306281621
+186,DLG4_RAT_Ranganathan2012,DeepSequence (published) - single,0.5918364209116883,0.5957477584911891,0.01995659395689766
+187,DLG4_RAT_Ranganathan2012,DeepSequence (published),0.6036625648098791,0.6070883056488765,0.020478051115311577
+188,DLG4_RAT_Ranganathan2012,ESM-1v (zero shot),0.5939968225123472,0.596847197984075,0.017630058333348964
+189,DLG4_RAT_Ranganathan2012,ESM-1v (+further training),0.5682242782581113,0.5694370162367897,0.018958485435275826
+190,DLG4_RAT_Ranganathan2012,DeepSequence (replicated),0.5765783481552431,0.5780455229510647,0.020809463738300338
+191,DLG4_RAT_Ranganathan2012,EVMutation (replicated),0.5241180720329418,0.5261248836538202,0.02186178951891566
+192,DLG4_RAT_Ranganathan2012,ProtBERT-BFD,0.5595215663542064,0.5614010576594576,0.01854665905518769
+193,DLG4_RAT_Ranganathan2012,TAPE,0.013351052100670802,0.023584538183512672,0.014457178947554198
+194,DLG4_RAT_Ranganathan2012,UniRep,0.3016230099651694,0.29971140046709394,0.027613438646412317
+195,DLG4_RAT_Ranganathan2012,ESM-1b,0.5275305911412679,0.5302881928240022,0.01742943283364256
+532,UBC9_HUMAN_Roth2017,gt,1.0,1.0,0.0
+533,UBC9_HUMAN_Roth2017,MSA Transformer,0.5257424624114496,0.5264407612968436,0.016782723059713636
+534,UBC9_HUMAN_Roth2017,PSSM,0.4252037519097337,0.42647455982971144,0.01848680680775899
+535,UBC9_HUMAN_Roth2017,EVMutation (published),0.509842177492847,0.5077129136778502,0.015575205557209867
+536,UBC9_HUMAN_Roth2017,DeepSequence (published) - single,0.5301569663083602,0.5321597182510774,0.017474789730172
+537,UBC9_HUMAN_Roth2017,DeepSequence (published),0.5451981164854786,0.5458299698225401,0.01779569747351641
+538,UBC9_HUMAN_Roth2017,ESM-1v (zero shot),0.5278002974052306,0.5289273564517684,0.0174313365381068
+539,UBC9_HUMAN_Roth2017,ESM-1v (+further training),0.530438889646823,0.5308418048769536,0.016194450129423676
+540,UBC9_HUMAN_Roth2017,DeepSequence (replicated),0.564757773228155,0.566254233293921,0.01584492457358362
+541,UBC9_HUMAN_Roth2017,EVMutation (replicated),0.5197625999541963,0.5169223464706466,0.015797268774811208
+542,UBC9_HUMAN_Roth2017,ProtBERT-BFD,0.4566745629330005,0.4597664455496374,0.01695356254093292
+543,UBC9_HUMAN_Roth2017,TAPE,0.2663201366958923,0.2655109350197859,0.020940445702932447
+544,UBC9_HUMAN_Roth2017,UniRep,0.027526620911384754,0.032226022213806015,0.0170540626567071
+545,UBC9_HUMAN_Roth2017,ESM-1b,0.42880952407240114,0.4257117461019339,0.016052256587648167
+56,BG505_env_Bloom2018,gt,1.0,1.0,0.0
+57,BG505_env_Bloom2018,MSA Transformer,0.5171538292833168,0.5171496622518762,0.006910262108734762
+58,BG505_env_Bloom2018,PSSM,0.4805043665958063,0.47934729964627765,0.006403736305762083
+59,BG505_env_Bloom2018,EVMutation (published),0.4088382668800448,0.4081773036708733,0.005327774682093342
+60,BG505_env_Bloom2018,DeepSequence (published) - single,0.17427055696697,0.17391089639228693,0.007096603606133151
+61,BG505_env_Bloom2018,DeepSequence (published),0.17901824140808725,0.17897151771830583,0.007361446332203004
+62,BG505_env_Bloom2018,ESM-1v (zero shot),0.5263381402665136,0.5258084748455352,0.007358753329943184
+63,BG505_env_Bloom2018,ESM-1v (+further training),0.5320285787368926,0.5326341655547056,0.006038819273548771
+64,BG505_env_Bloom2018,DeepSequence (replicated),0.12964060496065777,0.12952804062160664,0.009474191082563171
+65,BG505_env_Bloom2018,EVMutation (replicated),0.4183486299186111,0.41871016990435506,0.006336285035457084
+66,BG505_env_Bloom2018,ProtBERT-BFD,0.325305964247549,0.32240107597741696,0.007706083919970625
+67,BG505_env_Bloom2018,TAPE,0.028456890112386356,0.03179954870869962,0.006177994317554592
+68,BG505_env_Bloom2018,UniRep,0.20075668930719107,0.20277486226220803,0.008814544928765732
+69,BG505_env_Bloom2018,ESM-1b,0.46213634322466174,0.4615123969304536,0.007864462687152629
+462,SUMO1_HUMAN_Roth2017,gt,1.0,1.0,0.0
+463,SUMO1_HUMAN_Roth2017,MSA Transformer,0.5159144148964154,0.5204033887826773,0.016144534346014247
+464,SUMO1_HUMAN_Roth2017,PSSM,0.5068653918333971,0.5116066608485563,0.014073225487584034
+465,SUMO1_HUMAN_Roth2017,EVMutation (published),0.3943527527312187,0.39862273161443174,0.0218398217088854
+466,SUMO1_HUMAN_Roth2017,DeepSequence (published) - single,0.4546171147375774,0.4563868817701918,0.018789136591575727
+467,SUMO1_HUMAN_Roth2017,DeepSequence (published),0.4817132258726107,0.48324298006415234,0.01745562156694833
+468,SUMO1_HUMAN_Roth2017,ESM-1v (zero shot),0.4891908408067041,0.4902530587936492,0.02001929801145162
+469,SUMO1_HUMAN_Roth2017,ESM-1v (+further training),0.4870961371636435,0.48754122676040595,0.017288489612331478
+470,SUMO1_HUMAN_Roth2017,DeepSequence (replicated),0.4779470243180796,0.479960524595585,0.020147523090706877
+471,SUMO1_HUMAN_Roth2017,EVMutation (replicated),0.39100280963861933,0.3956383413944518,0.022045400156456525
+472,SUMO1_HUMAN_Roth2017,ProtBERT-BFD,0.4707613294700864,0.4702076997635277,0.020510642797657852
+473,SUMO1_HUMAN_Roth2017,TAPE,0.15068074140877044,0.14631769856923166,0.02417823419968644
+474,SUMO1_HUMAN_Roth2017,UniRep,0.08697581842056104,0.06864759400267187,0.02922305952665209
+475,SUMO1_HUMAN_Roth2017,ESM-1b,0.4208939554865267,0.4235043067669399,0.019117696248176413
+196,GAL4_YEAST_Shendure2015,gt,1.0,1.0,0.0
+197,GAL4_YEAST_Shendure2015,MSA Transformer,0.5266956520744294,0.5236482827987523,0.016313958672329276
+198,GAL4_YEAST_Shendure2015,PSSM,0.368064720312814,0.36172145092245883,0.020305140527647172
+199,GAL4_YEAST_Shendure2015,EVMutation (published),0.5527897519391798,0.5471044481198163,0.02237497642639291
+200,GAL4_YEAST_Shendure2015,DeepSequence (published) - single,0.42007796271168657,0.41396009878702655,0.019800914070488684
+201,GAL4_YEAST_Shendure2015,DeepSequence (published),0.44771721216490296,0.44150522173269036,0.018572211283338934
+202,GAL4_YEAST_Shendure2015,ESM-1v (zero shot),0.45042739409421084,0.4485824906124051,0.01733151555320306
+203,GAL4_YEAST_Shendure2015,ESM-1v (+further training),0.4395287875677813,0.43649598169067916,0.018212972219452478
+204,GAL4_YEAST_Shendure2015,DeepSequence (replicated),0.5272024171674746,0.5237709355239704,0.016265152112029255
+205,GAL4_YEAST_Shendure2015,EVMutation (replicated),0.523268936970079,0.5165768570993756,0.01812112296069153
+206,GAL4_YEAST_Shendure2015,ProtBERT-BFD,0.3363931685215737,0.33298715717786215,0.022821772133688727
+207,GAL4_YEAST_Shendure2015,TAPE,0.05800080952303017,0.06811831080338804,0.02702340073359099
+208,GAL4_YEAST_Shendure2015,UniRep,0.3119470194838534,0.3090243620686157,0.024570145461742788
+209,GAL4_YEAST_Shendure2015,ESM-1b,0.615398097573459,0.6139141766482525,0.01616576617819578
+84,BLAT_ECOLX_Ostermeier2014,gt,1.0,1.0,0.0
+85,BLAT_ECOLX_Ostermeier2014,MSA Transformer,0.7774725340080285,0.7772991120484407,0.00583115239464695
+86,BLAT_ECOLX_Ostermeier2014,PSSM,0.6266355490550037,0.6249558545704365,0.010963908722361101
+87,BLAT_ECOLX_Ostermeier2014,EVMutation (published),0.7253493163946709,0.7259111779130355,0.007055247161974134
+88,BLAT_ECOLX_Ostermeier2014,DeepSequence (published) - single,0.7573575626515066,0.7574787327674316,0.0052410410467141665
+89,BLAT_ECOLX_Ostermeier2014,DeepSequence (published),0.7780042111143465,0.7782069969302521,0.004849028422585635
+90,BLAT_ECOLX_Ostermeier2014,ESM-1v (zero shot),0.7233628349192526,0.7225479424696009,0.006683971493639809
+91,BLAT_ECOLX_Ostermeier2014,ESM-1v (+further training),0.7730745420327815,0.7726874168606355,0.0065547231293585735
+92,BLAT_ECOLX_Ostermeier2014,DeepSequence (replicated),0.7927294372907648,0.7925690818705499,0.005021676073847794
+93,BLAT_ECOLX_Ostermeier2014,EVMutation (replicated),0.747447811602467,0.7485714564837528,0.0071001482782374
+94,BLAT_ECOLX_Ostermeier2014,ProtBERT-BFD,0.4821130792800116,0.47870446935550176,0.010356473470947562
+95,BLAT_ECOLX_Ostermeier2014,TAPE,0.14731886213718737,0.1477711259950593,0.015163340123327854
+96,BLAT_ECOLX_Ostermeier2014,UniRep,0.03889602482478667,0.040737324507109665,0.01611318153937308
+97,BLAT_ECOLX_Ostermeier2014,ESM-1b,0.7251755113400309,0.7251084679688187,0.007010136942265081
+126,BLAT_ECOLX_Tenaillon2013,gt,1.0,1.0,0.0
+127,BLAT_ECOLX_Tenaillon2013,MSA Transformer,0.758747218907931,0.7560608794742473,0.014270208324360757
+128,BLAT_ECOLX_Tenaillon2013,PSSM,0.5914509988018589,0.584550802834145,0.019364487874982444
+129,BLAT_ECOLX_Tenaillon2013,EVMutation (published),0.6976808741924794,0.6939730281305263,0.017209064919944853
+130,BLAT_ECOLX_Tenaillon2013,DeepSequence (published) - single,0.7351880319385844,0.734362047270455,0.014525091995200967
+131,BLAT_ECOLX_Tenaillon2013,DeepSequence (published),0.7543453371693124,0.7522080258926729,0.012721106563380765
+132,BLAT_ECOLX_Tenaillon2013,ESM-1v (zero shot),0.705631894277465,0.6997739614439058,0.019936391958892526
+133,BLAT_ECOLX_Tenaillon2013,ESM-1v (+further training),0.7468651561611359,0.7420410512131015,0.018012297742261523
+134,BLAT_ECOLX_Tenaillon2013,DeepSequence (replicated),0.7636633966742921,0.7615069456529638,0.012651819062077643
+135,BLAT_ECOLX_Tenaillon2013,EVMutation (replicated),0.7188975120768561,0.7181426689915423,0.015603675185522175
+136,BLAT_ECOLX_Tenaillon2013,ProtBERT-BFD,0.5849974717911389,0.5759011621581035,0.02142157252628116
+137,BLAT_ECOLX_Tenaillon2013,TAPE,0.12293837281461435,0.12415027252537572,0.02966599757403504
+138,BLAT_ECOLX_Tenaillon2013,UniRep,0.018139031023161226,0.026517625360817247,0.022618042553652993
+139,BLAT_ECOLX_Tenaillon2013,ESM-1b,0.7054857406191637,0.6998800269911201,0.021248434713645674
+294,MTH3_HAEAESTABILIZED_Tawfik2015,gt,1.0,1.0,0.0
+295,MTH3_HAEAESTABILIZED_Tawfik2015,MSA Transformer,0.6835339164997478,0.6812589032616372,0.015535718954003042
+296,MTH3_HAEAESTABILIZED_Tawfik2015,PSSM,0.4416134495623428,0.44246894551424826,0.020085764953621237
+297,MTH3_HAEAESTABILIZED_Tawfik2015,EVMutation (published),0.6941903664027441,0.692016316213505,0.013094374915061331
+298,MTH3_HAEAESTABILIZED_Tawfik2015,DeepSequence (published) - single,0.7101794101084092,0.7063260455653937,0.01326393690807996
+299,MTH3_HAEAESTABILIZED_Tawfik2015,DeepSequence (published),0.7267412245284101,0.7233050341503924,0.011532195128181771
+300,MTH3_HAEAESTABILIZED_Tawfik2015,ESM-1v (zero shot),0.7028328471487627,0.6988674105054747,0.01389621848758316
+301,MTH3_HAEAESTABILIZED_Tawfik2015,ESM-1v (+further training),0.710444924607402,0.7058556345404041,0.013349806491214776
+302,MTH3_HAEAESTABILIZED_Tawfik2015,DeepSequence (replicated),0.7319377341571671,0.7284289432359438,0.011025526727213455
+303,MTH3_HAEAESTABILIZED_Tawfik2015,EVMutation (replicated),0.6881944596927662,0.6851918594602512,0.014875805426091208
+304,MTH3_HAEAESTABILIZED_Tawfik2015,ProtBERT-BFD,0.6227906481890101,0.6226488635471868,0.01368447324595496
+305,MTH3_HAEAESTABILIZED_Tawfik2015,TAPE,0.28986222371790993,0.29059002837311587,0.02293939359468894
+306,MTH3_HAEAESTABILIZED_Tawfik2015,UniRep,0.21256210891011498,0.20623088209322155,0.016982274395708104
+307,MTH3_HAEAESTABILIZED_Tawfik2015,ESM-1b,0.6013150928569109,0.6001765775002926,0.016765716391031985
+336,PABP_YEAST_Fields2013-singles,gt,1.0,1.0,0.0
+337,PABP_YEAST_Fields2013-singles,MSA Transformer,0.6686627777483488,0.6635711675242917,0.016177764258807882
+338,PABP_YEAST_Fields2013-singles,PSSM,0.42896284951548524,0.4291061605053958,0.02141741884779835
+339,PABP_YEAST_Fields2013-singles,EVMutation (published),0.6013962392775869,0.5997811886851577,0.01156645100836508
+340,PABP_YEAST_Fields2013-singles,DeepSequence (published) - single,0.6812656818721897,0.6786988061064501,0.015497646686186226
+341,PABP_YEAST_Fields2013-singles,DeepSequence (published),0.6879032274975242,0.6849609026248084,0.013722212867231664
+342,PABP_YEAST_Fields2013-singles,ESM-1v (zero shot),0.6658159648503742,0.6617537305293117,0.015156538272710452
+343,PABP_YEAST_Fields2013-singles,ESM-1v (+further training),0.6465056114571746,0.6403467410293169,0.015004313697387745
+344,PABP_YEAST_Fields2013-singles,DeepSequence (replicated),0.621237612823598,0.6166697742075684,0.015117873200550663
+345,PABP_YEAST_Fields2013-singles,EVMutation (replicated),0.6004126589511687,0.5972505301682013,0.014433400233151502
+346,PABP_YEAST_Fields2013-singles,ProtBERT-BFD,0.6590004474532792,0.6522426172055227,0.01761240430158242
+347,PABP_YEAST_Fields2013-singles,TAPE,0.08173774797871515,0.07580827883401792,0.02537818685979441
+348,PABP_YEAST_Fields2013-singles,UniRep,0.3524666785583584,0.3507280735011912,0.0214696561284629
+349,PABP_YEAST_Fields2013-singles,ESM-1b,0.6826492897030101,0.6776229323965959,0.014065470938740474
+378,POL_HV1N5-CA_Ndungu2014,gt,1.0,1.0,0.0
+379,POL_HV1N5-CA_Ndungu2014,MSA Transformer,0.8910953880622985,0.8715091885859015,0.04404390170038043
+380,POL_HV1N5-CA_Ndungu2014,PSSM,0.8027669847951657,0.7648376227138198,0.07156434462986615
+381,POL_HV1N5-CA_Ndungu2014,EVMutation (published),0.814063406437493,0.770542797914026,0.07799856588047756
+382,POL_HV1N5-CA_Ndungu2014,DeepSequence (published) - single,0.8237975570016259,0.79822473995615,0.059914597251614195
+383,POL_HV1N5-CA_Ndungu2014,DeepSequence (published),0.8152651534207193,0.7889600095175417,0.06617658502045558
+384,POL_HV1N5-CA_Ndungu2014,ESM-1v (zero shot),0.6512266902103299,0.6238138829808892,0.07766375006901825
+385,POL_HV1N5-CA_Ndungu2014,ESM-1v (+further training),0.815024804024074,0.7869525208215408,0.05466301623393627
+386,POL_HV1N5-CA_Ndungu2014,DeepSequence (replicated),0.820552840146915,0.7966023557532389,0.05723243614263114
+387,POL_HV1N5-CA_Ndungu2014,EVMutation (replicated),0.8211537136385281,0.7735182879439364,0.08095098085359952
+388,POL_HV1N5-CA_Ndungu2014,ProtBERT-BFD,0.38179501657099424,0.42193301482635154,0.14774279685397848
+389,POL_HV1N5-CA_Ndungu2014,TAPE,0.3483864504373032,0.4104118824597043,0.1470038347598428
+390,POL_HV1N5-CA_Ndungu2014,UniRep,0.4639945102236729,0.40093126150106306,0.16524010823471783
+391,POL_HV1N5-CA_Ndungu2014,ESM-1b,0.5957059795852749,0.5731935627996381,0.13224460706985164
+364,POLG_HCVJF_Sun2014,gt,1.0,1.0,0.0
+365,POLG_HCVJF_Sun2014,MSA Transformer,0.6677478668298926,0.670497369164695,0.017151824473903992
+366,POLG_HCVJF_Sun2014,PSSM,0.6494520970500091,0.6472164109404976,0.01636183923822566
+367,POLG_HCVJF_Sun2014,EVMutation (published),0.5865270173010575,0.5896466823313951,0.019303754939295415
+368,POLG_HCVJF_Sun2014,DeepSequence (published) - single,0.39251206867133814,0.39508591794495096,0.020227604662347038
+369,POLG_HCVJF_Sun2014,DeepSequence (published),0.3713820411310244,0.3795469161150128,0.023014839527838928
+370,POLG_HCVJF_Sun2014,ESM-1v (zero shot),0.6340055756536588,0.6335374345199196,0.01764641035648572
+371,POLG_HCVJF_Sun2014,ESM-1v (+further training),0.5652424039602515,0.5708409099578081,0.020101066917986344
+372,POLG_HCVJF_Sun2014,DeepSequence (replicated),0.46165721811412747,0.46807649417617975,0.022195221779086654
+373,POLG_HCVJF_Sun2014,EVMutation (replicated),0.5825890447985209,0.5852998202243922,0.022749622221884737
+374,POLG_HCVJF_Sun2014,ProtBERT-BFD,0.08677957966093554,0.07770559640019087,0.018629212118296944
+375,POLG_HCVJF_Sun2014,TAPE,0.05680250795137183,0.04875383206148289,0.02707444904306108
+376,POLG_HCVJF_Sun2014,UniRep,0.023755477003752395,0.018385016483602818,0.018204300138877978
+377,POLG_HCVJF_Sun2014,ESM-1b,0.14790101977036282,0.1451007380223277,0.018527745632610403
+308,P84126_THETH_b0,gt,1.0,1.0,0.0
+309,P84126_THETH_b0,MSA Transformer,0.5799265252107147,0.5738614544298895,0.016977861401835232
+310,P84126_THETH_b0,PSSM,0.5059938805045767,0.497875302076657,0.017450898377151254
+311,P84126_THETH_b0,EVMutation (published),0.5762676414938405,0.5713927693726344,0.015545247017306097
+312,P84126_THETH_b0,DeepSequence (published) - single,0.5267205050047176,0.5248192557448691,0.018977964279615622
+313,P84126_THETH_b0,DeepSequence (published),0.5877949812163814,0.5848736179704858,0.015886804363744445
+314,P84126_THETH_b0,ESM-1v (zero shot),0.5870635090422014,0.5811455597687403,0.01549045031439649
+315,P84126_THETH_b0,ESM-1v (+further training),0.5790524141907541,0.572625780593036,0.014736663161466617
+316,P84126_THETH_b0,DeepSequence (replicated),0.6038248556337414,0.602211740801349,0.013767217833566922
+317,P84126_THETH_b0,EVMutation (replicated),0.5700921855537731,0.5640773353834833,0.016516084944527984
+318,P84126_THETH_b0,ProtBERT-BFD,0.5721041244229742,0.5675746516094956,0.016572893593986743
+319,P84126_THETH_b0,TAPE,0.530841865819399,0.524876249458913,0.02120781628171545
+320,P84126_THETH_b0,UniRep,0.2274910400298911,0.22961913248092042,0.020560804448588488
+321,P84126_THETH_b0,ESM-1b,0.5700277792841302,0.5653475783194437,0.014015440292458613
+238,HSP82_YEAST_Bolon2016,gt,1.0,1.0,0.0
+239,HSP82_YEAST_Bolon2016,MSA Transformer,0.48917822587769644,0.4911252194110343,0.010638436729624378
+240,HSP82_YEAST_Bolon2016,PSSM,0.5080494707835187,0.5081989816141457,0.011630856503356157
+241,HSP82_YEAST_Bolon2016,EVMutation (published),0.49012823437511943,0.4936437806547088,0.009824172602805919
+242,HSP82_YEAST_Bolon2016,DeepSequence (published) - single,0.5292966677834544,0.529631038805717,0.010670673995645957
+243,HSP82_YEAST_Bolon2016,DeepSequence (published),0.5345178740223454,0.5359585453318252,0.010405888846639616
+244,HSP82_YEAST_Bolon2016,ESM-1v (zero shot),0.5770032224447383,0.5794593692641861,0.011071932164493134
+245,HSP82_YEAST_Bolon2016,ESM-1v (+further training),0.5582427433705837,0.5611201932811588,0.011556274778765427
+246,HSP82_YEAST_Bolon2016,DeepSequence (replicated),0.5447815098744462,0.5468478860295007,0.009190852805604763
+247,HSP82_YEAST_Bolon2016,EVMutation (replicated),0.4813174874985635,0.4833921506539731,0.010641726788141665
+248,HSP82_YEAST_Bolon2016,ProtBERT-BFD,0.5020206378059789,0.5038884666599158,0.010495246599296015
+249,HSP82_YEAST_Bolon2016,TAPE,0.11898211650679438,0.11418591489599068,0.0095283384478768
+250,HSP82_YEAST_Bolon2016,UniRep,0.28031483895837656,0.27995048255516847,0.016433158282085403
+251,HSP82_YEAST_Bolon2016,ESM-1b,0.5364678889354202,0.5369383571750992,0.012361055325265038
+518,TPMT_HUMAN_Fowler2018,gt,1.0,1.0,0.0
+519,TPMT_HUMAN_Fowler2018,MSA Transformer,0.5424668178781977,0.540549584398611,0.0130129824539858
+520,TPMT_HUMAN_Fowler2018,PSSM,0.48146126611262924,0.4789555004656866,0.01619133485423122
+521,TPMT_HUMAN_Fowler2018,EVMutation (published),0.5243737765811296,0.5211293318961853,0.013677282837478894
+522,TPMT_HUMAN_Fowler2018,DeepSequence (published) - single,0.5616767945109112,0.5601851812367447,0.014098853189084848
+523,TPMT_HUMAN_Fowler2018,DeepSequence (published),0.57189769074549,0.5692094068525796,0.014271609738516988
+524,TPMT_HUMAN_Fowler2018,ESM-1v (zero shot),0.5746177224716744,0.5733456843246085,0.013161147755010085
+525,TPMT_HUMAN_Fowler2018,ESM-1v (+further training),0.5667144426739071,0.5643404804061067,0.012156963082272723
+526,TPMT_HUMAN_Fowler2018,DeepSequence (replicated),0.5734186713594054,0.5725136899852928,0.01380924904516693
+527,TPMT_HUMAN_Fowler2018,EVMutation (replicated),0.5415452371370989,0.5367554768824575,0.014285595575941081
+528,TPMT_HUMAN_Fowler2018,ProtBERT-BFD,0.5755683713471709,0.5740036007868525,0.011781062164033955
+529,TPMT_HUMAN_Fowler2018,TAPE,0.21528496910537398,0.2163150200476124,0.01671186684925183
+530,TPMT_HUMAN_Fowler2018,UniRep,0.24493123064838648,0.2429991773585667,0.01737637750092445
+531,TPMT_HUMAN_Fowler2018,ESM-1b,0.585864822736713,0.5831594511762555,0.012841359293495367
+252,IF1_ECOLI_Kishony,gt,1.0,1.0,0.0
+253,IF1_ECOLI_Kishony,MSA Transformer,0.5374285372059321,0.5379972559155065,0.020194737882211888
+254,IF1_ECOLI_Kishony,PSSM,0.5052167418774901,0.5069241815649538,0.021074640419288252
+255,IF1_ECOLI_Kishony,EVMutation (published),0.5003046508987051,0.5003352234717605,0.022143569265813394
+256,IF1_ECOLI_Kishony,DeepSequence (published) - single,0.5239331603676252,0.5269853910903933,0.019737341732545446
+257,IF1_ECOLI_Kishony,DeepSequence (published),0.5232098971780704,0.5263086088632558,0.019848792789945318
+258,IF1_ECOLI_Kishony,ESM-1v (zero shot),0.5489484528277484,0.5487290835018946,0.018090689195472447
+259,IF1_ECOLI_Kishony,ESM-1v (+further training),0.5619784833217806,0.5609262853368223,0.01759996229002091
+260,IF1_ECOLI_Kishony,DeepSequence (replicated),0.5371602219910716,0.5377419211137229,0.016834767452949526
+261,IF1_ECOLI_Kishony,EVMutation (replicated),0.49106963416130595,0.487364241499556,0.02231862948055046
+262,IF1_ECOLI_Kishony,ProtBERT-BFD,0.49640551530348803,0.49572596972809546,0.021526763193477596
+263,IF1_ECOLI_Kishony,TAPE,0.3330415186026132,0.3329560370901763,0.025365922197565095
+264,IF1_ECOLI_Kishony,UniRep,0.11711947357011687,0.11208310104103972,0.02358345571859419
+265,IF1_ECOLI_Kishony,ESM-1b,0.548626388294071,0.5492417382232484,0.020649586745340585
+154,BRCA1_HUMAN_RING,gt,1.0,1.0,0.0
+155,BRCA1_HUMAN_RING,MSA Transformer,0.5193005569644538,0.5203285046503858,0.041166837523028924
+156,BRCA1_HUMAN_RING,PSSM,0.5634130110253781,0.5581058154928551,0.04362682737121433
+157,BRCA1_HUMAN_RING,EVMutation (published),0.4739893403358574,0.4727346993937392,0.036203681195866866
+158,BRCA1_HUMAN_RING,DeepSequence (published) - single,0.5543323620928434,0.5511807082846618,0.039417214029845105
+159,BRCA1_HUMAN_RING,DeepSequence (published),0.5617365720006955,0.5591185452387899,0.03842959531538541
+160,BRCA1_HUMAN_RING,ESM-1v (zero shot),0.5415424138323831,0.5341119058002446,0.03653613513934021
+161,BRCA1_HUMAN_RING,ESM-1v (+further training),0.5550269189493816,0.5537267190735496,0.03937349908877672
+162,BRCA1_HUMAN_RING,DeepSequence (replicated),0.5658118851058472,0.5627698357833636,0.037505350857757684
+163,BRCA1_HUMAN_RING,EVMutation (replicated),0.46724701283596226,0.46596752209015424,0.04012929367216169
+164,BRCA1_HUMAN_RING,ProtBERT-BFD,0.47507688987212,0.4713003309113355,0.03734723419650255
+165,BRCA1_HUMAN_RING,TAPE,0.32511391092165925,0.3259918359586091,0.05215348292668556
+166,BRCA1_HUMAN_RING,UniRep,0.4574870950916383,0.466226203934293,0.029045683651907416
+167,BRCA1_HUMAN_RING,ESM-1b,0.5729860249285059,0.5672443289637067,0.030337589962057807
+98,BLAT_ECOLX_Palzkill2012,gt,1.0,1.0,0.0
+99,BLAT_ECOLX_Palzkill2012,MSA Transformer,0.5790109592895919,0.576559464262669,0.00813860157556571
+100,BLAT_ECOLX_Palzkill2012,PSSM,0.43541874109113293,0.4346912136467994,0.011525932575641511
+101,BLAT_ECOLX_Palzkill2012,EVMutation (published),0.5035045285112302,0.5015670440233889,0.011387370179059404
+102,BLAT_ECOLX_Palzkill2012,DeepSequence (published) - single,0.5387135089739988,0.5368824370193662,0.011450482383497546
+103,BLAT_ECOLX_Palzkill2012,DeepSequence (published),0.555680249906377,0.5539233663406711,0.009685510491284273
+104,BLAT_ECOLX_Palzkill2012,ESM-1v (zero shot),0.5394874702645197,0.5391023570614875,0.01072036011365782
+105,BLAT_ECOLX_Palzkill2012,ESM-1v (+further training),0.5672174721013816,0.5661011431263104,0.010297339893877539
+106,BLAT_ECOLX_Palzkill2012,DeepSequence (replicated),0.5665462527968158,0.5646253670538975,0.00957133565445041
+107,BLAT_ECOLX_Palzkill2012,EVMutation (replicated),0.5174199245259459,0.5144044189835617,0.010884656508769686
+108,BLAT_ECOLX_Palzkill2012,ProtBERT-BFD,0.413553141253597,0.4108322854585459,0.009768034882433663
+109,BLAT_ECOLX_Palzkill2012,TAPE,0.06707069584638124,0.06961505855709685,0.01069869285554942
+110,BLAT_ECOLX_Palzkill2012,UniRep,0.0166580631994596,0.02360053325843217,0.01073550650340384
+111,BLAT_ECOLX_Palzkill2012,ESM-1b,0.54626199827611,0.5450469362296847,0.010574502773225362
+210,HG_FLU_Bloom2016,gt,1.0,1.0,0.0
+211,HG_FLU_Bloom2016,MSA Transformer,0.4807847197470188,0.48133913444729065,0.007238779061281368
+212,HG_FLU_Bloom2016,PSSM,0.467255239211276,0.4658023016610004,0.007807096034181624
+213,HG_FLU_Bloom2016,EVMutation (published),0.45827659972583357,0.45804886604872685,0.008079452114456188
+214,HG_FLU_Bloom2016,DeepSequence (published) - single,0.38185154496925244,0.3815243480167388,0.009700262344961326
+215,HG_FLU_Bloom2016,DeepSequence (published),0.408559044093847,0.4088655809693629,0.007996666160114359
+216,HG_FLU_Bloom2016,ESM-1v (zero shot),0.5377581449296356,0.5381755966315286,0.00849950538034111
+217,HG_FLU_Bloom2016,ESM-1v (+further training),0.508004961821385,0.50881255466271,0.008423294599454892
+218,HG_FLU_Bloom2016,DeepSequence (replicated),0.425338497638415,0.425397348539834,0.007829596385814108
+219,HG_FLU_Bloom2016,EVMutation (replicated),0.4745801474914726,0.4740523393070549,0.00772632329264063
+220,HG_FLU_Bloom2016,ProtBERT-BFD,0.01216830054805121,0.014371682213071179,0.008595363873197908
+221,HG_FLU_Bloom2016,TAPE,0.011506602346493722,0.011595399204682536,0.006272208060584064
+222,HG_FLU_Bloom2016,UniRep,0.11238855588872076,0.11416182055768724,0.009338414613763222
+223,HG_FLU_Bloom2016,ESM-1b,0.13466161180017006,0.13615942676936274,0.010191230998347687
+224,HIS7_YEAST_Kondrashov2017,gt,1.0,1.0,0.0
+225,HIS7_YEAST_Kondrashov2017,MSA Transformer,0.30171610633916407,0.301852298319144,0.0012676098127205092
+226,HIS7_YEAST_Kondrashov2017,PSSM,0.46025981696434587,0.4602850849857719,0.0008742604182146493
+227,HIS7_YEAST_Kondrashov2017,EVMutation (published),0.5865226545066141,0.5868249726401576,0.0008233711922239475
+228,HIS7_YEAST_Kondrashov2017,DeepSequence (published) - single,0.4924067186635951,0.49266995692585214,0.0009319735566991191
+229,HIS7_YEAST_Kondrashov2017,DeepSequence (published),0.539985999849893,0.540227447923642,0.0008648831057151851
+230,HIS7_YEAST_Kondrashov2017,ESM-1v (zero shot),0.529358182381123,0.5295580913555231,0.0007957702421518582
+231,HIS7_YEAST_Kondrashov2017,ESM-1v (+further training),0.5600853174526501,0.5602849321317114,0.0008034030458959766
+232,HIS7_YEAST_Kondrashov2017,DeepSequence (replicated),0.5771729957674867,0.5774905109783149,0.0007170541729409019
+233,HIS7_YEAST_Kondrashov2017,EVMutation (replicated),0.5839282060088532,0.5842957417091784,0.0008208185316044222
+234,HIS7_YEAST_Kondrashov2017,ProtBERT-BFD,0.5529258544618092,0.5529676964174914,0.0007360080980740549
+235,HIS7_YEAST_Kondrashov2017,TAPE,0.1014260339205887,0.101752697842131,0.0014218010200125605
+236,HIS7_YEAST_Kondrashov2017,UniRep,0.1815849109979771,0.18158251774490938,0.0014496180282351032
+237,HIS7_YEAST_Kondrashov2017,ESM-1b,0.5000466702129686,0.5002672548811231,0.0008300778727114249
+70,BG_STRSQ_hmmerbit,gt,1.0,1.0,0.0
+71,BG_STRSQ_hmmerbit,MSA Transformer,0.6555958579973886,0.659326801766521,0.00913574784080978
+72,BG_STRSQ_hmmerbit,PSSM,0.5431285035494208,0.5491891696554254,0.01382628634556949
+73,BG_STRSQ_hmmerbit,EVMutation (published),0.6297133475442988,0.6326270443670063,0.006525349201880503
+74,BG_STRSQ_hmmerbit,DeepSequence (published) - single,0.6663981684713202,0.667051777591926,0.008814659279799686
+75,BG_STRSQ_hmmerbit,DeepSequence (published),0.671353366524517,0.6721885348190513,0.009076474761061238
+76,BG_STRSQ_hmmerbit,ESM-1v (zero shot),0.528482255049547,0.5356316829302887,0.013445455516506805
+77,BG_STRSQ_hmmerbit,ESM-1v (+further training),0.6516055107633002,0.6555062061803748,0.009214695345811378
+78,BG_STRSQ_hmmerbit,DeepSequence (replicated),0.6751486915624705,0.6760984429469532,0.009352478062896172
+79,BG_STRSQ_hmmerbit,EVMutation (replicated),0.6403197196758124,0.6434637152906157,0.007241339119188467
+80,BG_STRSQ_hmmerbit,ProtBERT-BFD,0.6365069784319516,0.6396483776091306,0.007811080769658407
+81,BG_STRSQ_hmmerbit,TAPE,0.43294588799591494,0.44008364687109375,0.01682011773736568
+82,BG_STRSQ_hmmerbit,UniRep,0.19989240659304872,0.20517404401460504,0.014963964273707563
+83,BG_STRSQ_hmmerbit,ESM-1b,0.6009580023829318,0.6080513205167206,0.013233900564372743
+546,UBE4B_MOUSE_Klevit2013-singles,gt,1.0,1.0,0.0
+547,UBE4B_MOUSE_Klevit2013-singles,MSA Transformer,0.49832943520981005,0.49905990983648263,0.03105986611446269
+548,UBE4B_MOUSE_Klevit2013-singles,PSSM,0.48860473816872396,0.49337529571134364,0.031233090528237636
+549,UBE4B_MOUSE_Klevit2013-singles,EVMutation (published),0.4906841388190489,0.49199509841786193,0.03340124290266937
+550,UBE4B_MOUSE_Klevit2013-singles,DeepSequence (published) - single,0.5312205433544338,0.5329703058977546,0.029417950250717785
+551,UBE4B_MOUSE_Klevit2013-singles,DeepSequence (published),0.5395815016518124,0.5430445683834082,0.02805620658685915
+552,UBE4B_MOUSE_Klevit2013-singles,ESM-1v (zero shot),0.5220724759142595,0.5231561587068752,0.024550717856015463
+553,UBE4B_MOUSE_Klevit2013-singles,ESM-1v (+further training),0.5218576853965179,0.5198018293944963,0.029197148702440747
+554,UBE4B_MOUSE_Klevit2013-singles,DeepSequence (replicated),0.5223215278319518,0.5271163507848048,0.02943995962728437
+555,UBE4B_MOUSE_Klevit2013-singles,EVMutation (replicated),0.4846034274372467,0.48559313231244233,0.035130463278636635
+556,UBE4B_MOUSE_Klevit2013-singles,ProtBERT-BFD,0.5185431412393127,0.5236656965174686,0.028751086492193833
+557,UBE4B_MOUSE_Klevit2013-singles,TAPE,0.01730835573288986,0.03524679123603523,0.02987730394986271
+558,UBE4B_MOUSE_Klevit2013-singles,UniRep,0.10151869130428068,0.09383999274532392,0.03556604553816262
+559,UBE4B_MOUSE_Klevit2013-singles,ESM-1b,0.4877718340726251,0.48675329873590983,0.03538514294193213
+14,B3VI55_LIPSTSTABLE,gt,1.0,1.0,0.0
+15,B3VI55_LIPSTSTABLE,MSA Transformer,0.5251005167937696,0.5275545723143799,0.009041770337910218
+16,B3VI55_LIPSTSTABLE,PSSM,0.33176316087515695,0.3329724203391258,0.011406640278605069
+17,B3VI55_LIPSTSTABLE,EVMutation (published),0.43366434910484786,0.43528511611104986,0.010773070057405095
+18,B3VI55_LIPSTSTABLE,DeepSequence (published) - single,0.4747977554498681,0.4764432728845286,0.00955452160837742
+19,B3VI55_LIPSTSTABLE,DeepSequence (published),0.48780043789141697,0.489281157459712,0.008995219832666802
+20,B3VI55_LIPSTSTABLE,ESM-1v (zero shot),0.5144594346026741,0.516476518473857,0.009568231035205155
+21,B3VI55_LIPSTSTABLE,ESM-1v (+further training),0.5408988777571562,0.5429418957469877,0.008374931117004644
+22,B3VI55_LIPSTSTABLE,DeepSequence (replicated),0.5108241872870117,0.511210509948324,0.008885420291063774
+23,B3VI55_LIPSTSTABLE,EVMutation (replicated),0.45220927299548147,0.45292929103074864,0.01073069192782447
+24,B3VI55_LIPSTSTABLE,ProtBERT-BFD,0.3909053736573994,0.39106874014118553,0.010538767559809602
+25,B3VI55_LIPSTSTABLE,TAPE,0.1114444287805763,0.1120507952365121,0.012548143306758398
+26,B3VI55_LIPSTSTABLE,UniRep,0.1050036846541006,0.10502881984575867,0.013417162356959514
+27,B3VI55_LIPSTSTABLE,ESM-1b,0.4153440056253525,0.4178700751301828,0.011532303943529923
+42,BF520_env_Bloom2018,gt,1.0,1.0,0.0
+43,BF520_env_Bloom2018,MSA Transformer,0.53045503362033,0.528360683768996,0.007312684938262093
+44,BF520_env_Bloom2018,PSSM,0.48128944926535605,0.47846350663914833,0.007277527859925501
+45,BF520_env_Bloom2018,EVMutation (published),0.3961853129701546,0.3926921454653422,0.00906487418575762
+46,BF520_env_Bloom2018,DeepSequence (published) - single,0.0033350478364570987,0.009194378034284379,0.008689823275268422
+47,BF520_env_Bloom2018,DeepSequence (published),0.0029335219050744344,0.008967370879028378,0.005401668464782235
+48,BF520_env_Bloom2018,ESM-1v (zero shot),0.5137302974636209,0.5110743977804137,0.006968499056122589
+49,BF520_env_Bloom2018,ESM-1v (+further training),0.5260276733897689,0.524307610096477,0.006318717908534139
+50,BF520_env_Bloom2018,DeepSequence (replicated),0.03940344835961486,0.03571128763490593,0.010040975563315926
+51,BF520_env_Bloom2018,EVMutation (replicated),0.41701008243014626,0.4131323889976084,0.009699036294956902
+52,BF520_env_Bloom2018,ProtBERT-BFD,0.35920295750142894,0.356294565641996,0.008888654031775356
+53,BF520_env_Bloom2018,TAPE,0.03296454518403542,0.03347884384026924,0.010604055757164086
+54,BF520_env_Bloom2018,UniRep,0.19509435241603182,0.19534025730349494,0.008044823411980718
+55,BF520_env_Bloom2018,ESM-1b,0.4644044299599039,0.4621942590259476,0.0072479072681551
+490,TIM_THEMA_b0,gt,1.0,1.0,0.0
+491,TIM_THEMA_b0,MSA Transformer,0.476641570693168,0.4772084109762894,0.021262538481700146
+492,TIM_THEMA_b0,PSSM,0.3987702061334132,0.39949709692713725,0.022351322419246027
+493,TIM_THEMA_b0,EVMutation (published),0.43189426495468003,0.4320501737162544,0.023786620230565422
+494,TIM_THEMA_b0,DeepSequence (published) - single,0.4429769169797119,0.4459763974635589,0.023367347171715933
+495,TIM_THEMA_b0,DeepSequence (published),0.43793312514525096,0.43932643176681746,0.02400129226573652
+496,TIM_THEMA_b0,ESM-1v (zero shot),0.5026024874119385,0.5028606317050863,0.0196194686907316
+497,TIM_THEMA_b0,ESM-1v (+further training),0.49669513157740486,0.49715001129465347,0.021252079717001313
+498,TIM_THEMA_b0,DeepSequence (replicated),0.4065555906798227,0.4088509392059708,0.02383945792859995
+499,TIM_THEMA_b0,EVMutation (replicated),0.4372508700326144,0.43818153131258974,0.022249944638541977
+500,TIM_THEMA_b0,ProtBERT-BFD,0.5030640016728883,0.505061273209954,0.02004593130312488
+501,TIM_THEMA_b0,TAPE,0.426213240694561,0.42842760017164616,0.025447870266352153
+502,TIM_THEMA_b0,UniRep,0.15476155586059467,0.14652965312897737,0.02996974334997726
+503,TIM_THEMA_b0,ESM-1b,0.45916689313021203,0.45909489666370673,0.021672817044030984
+560,YAP1_HUMAN_Fields2012-singles,gt,1.0,1.0,0.0
+561,YAP1_HUMAN_Fields2012-singles,MSA Transformer,0.5887070642875368,0.5706519861099527,0.040386258549904866
+562,YAP1_HUMAN_Fields2012-singles,PSSM,0.5690684565019047,0.5451477752262303,0.04238593338136975
+563,YAP1_HUMAN_Fields2012-singles,EVMutation (published),0.5844091425595583,0.5694702589374069,0.03313034558652812
+564,YAP1_HUMAN_Fields2012-singles,DeepSequence (published) - single,0.6529398136915281,0.6360722049001328,0.032593234109222484
+565,YAP1_HUMAN_Fields2012-singles,DeepSequence (published),0.6456339137145388,0.6294381645066613,0.03481028251531764
+566,YAP1_HUMAN_Fields2012-singles,ESM-1v (zero shot),0.493433729941941,0.48010349687654763,0.03994836553901359
+567,YAP1_HUMAN_Fields2012-singles,ESM-1v (+further training),0.6376365184641369,0.623715503200643,0.03529100768352496
+568,YAP1_HUMAN_Fields2012-singles,DeepSequence (replicated),0.6447076152922595,0.6272717205426198,0.033734031749303436
+569,YAP1_HUMAN_Fields2012-singles,EVMutation (replicated),0.557456385066247,0.543521991089175,0.035472281995781466
+570,YAP1_HUMAN_Fields2012-singles,ProtBERT-BFD,0.5482997903990557,0.5417356076069597,0.03807717806310085
+571,YAP1_HUMAN_Fields2012-singles,TAPE,0.14957625858401177,0.1801038227952158,0.05696901517750361
+572,YAP1_HUMAN_Fields2012-singles,UniRep,0.029127721955503238,0.05432100409858017,0.03754188708662973
+573,YAP1_HUMAN_Fields2012-singles,ESM-1b,0.49074132050160213,0.47904274853673057,0.033063480562130175
+140,BRCA1_HUMAN_BRCT,gt,1.0,1.0,0.0
+141,BRCA1_HUMAN_BRCT,MSA Transformer,0.5928951868322753,0.5960239232158157,0.024771243739877567
+142,BRCA1_HUMAN_BRCT,PSSM,0.45174214496735526,0.4661431870435194,0.02242218107567778
+143,BRCA1_HUMAN_BRCT,EVMutation (published),0.575430631741363,0.5736725403805653,0.02241134787829384
+144,BRCA1_HUMAN_BRCT,DeepSequence (published) - single,0.56817101100728,0.5677360233280984,0.023388317280203356
+145,BRCA1_HUMAN_BRCT,DeepSequence (published),0.5887397586618573,0.5891289275127845,0.025473578386602996
+146,BRCA1_HUMAN_BRCT,ESM-1v (zero shot),0.43099744949022667,0.43820237951190066,0.026042969300316396
+147,BRCA1_HUMAN_BRCT,ESM-1v (+further training),0.5902839735003296,0.5924129041526023,0.022139044593609002
+148,BRCA1_HUMAN_BRCT,DeepSequence (replicated),0.6048855640279864,0.6073185827311438,0.022386535818213062
+149,BRCA1_HUMAN_BRCT,EVMutation (replicated),0.5943999321364778,0.5914745443226089,0.021113952368322746
+150,BRCA1_HUMAN_BRCT,ProtBERT-BFD,0.3807653003587669,0.38891615029062426,0.025074152619186402
+151,BRCA1_HUMAN_BRCT,TAPE,0.05185305485548559,0.06300777252497933,0.028006890598497248
+152,BRCA1_HUMAN_BRCT,UniRep,0.001009857244270908,0.024731061663019048,0.015154925933309927
+153,BRCA1_HUMAN_BRCT,ESM-1b,0.5015315617812091,0.5064634555923619,0.023953542534915467
+392,PTEN_HUMAN_Fowler2018,gt,1.0,1.0,0.0
+393,PTEN_HUMAN_Fowler2018,MSA Transformer,0.47723620402030725,0.4753707747929236,0.02041283060478458
+394,PTEN_HUMAN_Fowler2018,PSSM,0.4856705789707024,0.48540163299274086,0.01456452062108191
+395,PTEN_HUMAN_Fowler2018,EVMutation (published),0.4719118285668927,0.471271072166362,0.0171946566182038
+396,PTEN_HUMAN_Fowler2018,DeepSequence (published) - single,0.3920133344336003,0.38858890670495655,0.021275758912234644
+397,PTEN_HUMAN_Fowler2018,DeepSequence (published),0.41987506342546416,0.4166267407337688,0.02048001735899819
+398,PTEN_HUMAN_Fowler2018,ESM-1v (zero shot),0.4229117460619157,0.4216791615081128,0.022889108923404226
+399,PTEN_HUMAN_Fowler2018,ESM-1v (+further training),0.39745119841802773,0.3945916752986646,0.021297362107723237
+400,PTEN_HUMAN_Fowler2018,DeepSequence (replicated),0.4305140142957644,0.4275242282868633,0.019725711174546957
+401,PTEN_HUMAN_Fowler2018,EVMutation (replicated),0.4562780374489617,0.4570343049276496,0.017322527192949605
+402,PTEN_HUMAN_Fowler2018,ProtBERT-BFD,0.33431387346212144,0.33179215244793264,0.020748503501032432
+403,PTEN_HUMAN_Fowler2018,TAPE,0.0815737781117076,0.07995492505287519,0.010503276183057134
+404,PTEN_HUMAN_Fowler2018,UniRep,0.12983006902884325,0.12588348859114926,0.014127656261654718
+405,PTEN_HUMAN_Fowler2018,ESM-1b,0.4531834407900847,0.4504022413646627,0.019849910120253392
+406,RASH_HUMAN_Kuriyan,gt,1.0,1.0,0.0
+407,RASH_HUMAN_Kuriyan,MSA Transformer,0.4803453272496495,0.47796060028078263,0.017296002593559925
+408,RASH_HUMAN_Kuriyan,PSSM,0.45403466334542103,0.45183098094027907,0.01787397783493158
+409,RASH_HUMAN_Kuriyan,EVMutation (published),0.426341747437809,0.4229653964620471,0.01796419071897713
+410,RASH_HUMAN_Kuriyan,DeepSequence (published) - single,0.4728843219749578,0.468646862737559,0.01795985229559378
+411,RASH_HUMAN_Kuriyan,DeepSequence (published),0.46630499110930534,0.4614459892414701,0.01844189066007253
+412,RASH_HUMAN_Kuriyan,ESM-1v (zero shot),0.40727875438677086,0.40514547710795606,0.019133507919327063
+413,RASH_HUMAN_Kuriyan,ESM-1v (+further training),0.4147935141888191,0.4135324566538796,0.019219555375441195
+414,RASH_HUMAN_Kuriyan,DeepSequence (replicated),0.4718463706052156,0.4693887363231767,0.01665373741632862
+415,RASH_HUMAN_Kuriyan,EVMutation (replicated),0.41905311471869217,0.4152533339275954,0.0186456211620091
+416,RASH_HUMAN_Kuriyan,ProtBERT-BFD,0.3754077679946978,0.37224319469530265,0.020828472200020876
+417,RASH_HUMAN_Kuriyan,TAPE,0.3410037149796282,0.33821305583243977,0.02006172500654845
+418,RASH_HUMAN_Kuriyan,UniRep,0.28316890117187576,0.2804521504745546,0.016931648162857683
+419,RASH_HUMAN_Kuriyan,ESM-1b,0.34359820576597994,0.3393198546702645,0.02233342505800977
+420,RL401_YEAST_Bolon2013,gt,1.0,1.0,0.0
+421,RL401_YEAST_Bolon2013,MSA Transformer,0.5393318004129806,0.539097206955352,0.021107284547232343
+422,RL401_YEAST_Bolon2013,PSSM,0.4115446821389685,0.4029967757216066,0.02694607790699914
+423,RL401_YEAST_Bolon2013,EVMutation (published),0.45052851516991577,0.44375165687449697,0.020730159359455567
+424,RL401_YEAST_Bolon2013,DeepSequence (published) - single,0.4749118165434772,0.4687082939023236,0.02394781384346245
+425,RL401_YEAST_Bolon2013,DeepSequence (published),0.4891069305943915,0.4832628867355531,0.020837235335480923
+426,RL401_YEAST_Bolon2013,ESM-1v (zero shot),0.3780305382545389,0.3758572950839938,0.024438101597142162
+427,RL401_YEAST_Bolon2013,ESM-1v (+further training),0.36982341765304827,0.36644542392899027,0.02407282307497699
+428,RL401_YEAST_Bolon2013,DeepSequence (replicated),0.4875479135136087,0.48185494209379953,0.0224889402600134
+429,RL401_YEAST_Bolon2013,EVMutation (replicated),0.4649054471564614,0.45886647149033805,0.023302945044134975
+430,RL401_YEAST_Bolon2013,ProtBERT-BFD,0.39888508524198757,0.39548316549478596,0.024888924809883948
+431,RL401_YEAST_Bolon2013,TAPE,0.14136362276356793,0.1330418322501083,0.03239530659402373
+432,RL401_YEAST_Bolon2013,UniRep,0.11132107785812106,0.10371964011546289,0.027489697230632073
+433,RL401_YEAST_Bolon2013,ESM-1b,0.24192958338903614,0.24085320243486236,0.026555927080409752
+504,TPK1_HUMAN_Roth2017,gt,1.0,1.0,0.0
+505,TPK1_HUMAN_Roth2017,MSA Transformer,0.3415408511396077,0.33848473630002246,0.021212292675018753
+506,TPK1_HUMAN_Roth2017,PSSM,0.22072657020579714,0.2217644870830699,0.02369863818079656
+507,TPK1_HUMAN_Roth2017,EVMutation (published),0.2469522178641623,0.24730841435573284,0.01871959166700927
+508,TPK1_HUMAN_Roth2017,DeepSequence (published) - single,0.2575167904390149,0.2558063144550207,0.02055476962559519
+509,TPK1_HUMAN_Roth2017,DeepSequence (published),0.25614519361128574,0.2543757481978049,0.021967216506252808
+510,TPK1_HUMAN_Roth2017,ESM-1v (zero shot),0.33312805710397364,0.327155968484034,0.019777338079982523
+511,TPK1_HUMAN_Roth2017,ESM-1v (+further training),0.3317814621900089,0.3260704249321161,0.020979034647608977
+512,TPK1_HUMAN_Roth2017,DeepSequence (replicated),0.254473186445843,0.2514295290110877,0.02177016005083741
+513,TPK1_HUMAN_Roth2017,EVMutation (replicated),0.23867041174072273,0.23850268855356252,0.019193703578397097
+514,TPK1_HUMAN_Roth2017,ProtBERT-BFD,0.24102052288064413,0.2413615342747228,0.021738163307345192
+515,TPK1_HUMAN_Roth2017,TAPE,0.08255844412157543,0.08001237804591124,0.014507842152561057
+516,TPK1_HUMAN_Roth2017,UniRep,0.03311972123416161,0.03305706490328286,0.021578991252520243
+517,TPK1_HUMAN_Roth2017,ESM-1b,0.31345297835126973,0.3093139122995233,0.01991362706272715
+448,RL401_YEAST_Fraser2016,gt,1.0,1.0,0.0
+449,RL401_YEAST_Fraser2016,MSA Transformer,0.47482528061086265,0.4755124342582121,0.022569966034399314
+450,RL401_YEAST_Fraser2016,PSSM,0.3774016009912691,0.3773571687306486,0.015993333803246766
+451,RL401_YEAST_Fraser2016,EVMutation (published),0.40924774479083326,0.414388594923377,0.023024753763437884
+452,RL401_YEAST_Fraser2016,DeepSequence (published) - single,0.4318166476420829,0.4331294439053056,0.01910722980495821
+453,RL401_YEAST_Fraser2016,DeepSequence (published),0.4461365210878178,0.44899827665090486,0.019914780607228476
+454,RL401_YEAST_Fraser2016,ESM-1v (zero shot),0.32760495384990906,0.3291822704728004,0.02640133685107885
+455,RL401_YEAST_Fraser2016,ESM-1v (+further training),0.32990765588855403,0.3301471135676151,0.027322775648250797
+456,RL401_YEAST_Fraser2016,DeepSequence (replicated),0.4347770704966268,0.44167537164330406,0.017646832557385868
+457,RL401_YEAST_Fraser2016,EVMutation (replicated),0.4137359503211986,0.42342848446020753,0.023604416626258205
+458,RL401_YEAST_Fraser2016,ProtBERT-BFD,0.352256182908315,0.35488121411289225,0.026464030488060585
+459,RL401_YEAST_Fraser2016,TAPE,0.11246047580054895,0.11267381281460156,0.029451868723884767
+460,RL401_YEAST_Fraser2016,UniRep,0.09228239837675453,0.0855139650342207,0.02787982622396566
+461,RL401_YEAST_Fraser2016,ESM-1b,0.20980461592629226,0.21406696973008782,0.028067024964270798
+168,CALM1_HUMAN_Roth2017,gt,1.0,1.0,0.0
+169,CALM1_HUMAN_Roth2017,MSA Transformer,0.23711821254687918,0.23943485564742276,0.03113662255833698
+170,CALM1_HUMAN_Roth2017,PSSM,0.1940025702015548,0.19913438415655305,0.02630867989950736
+171,CALM1_HUMAN_Roth2017,EVMutation (published),0.23244288732891077,0.23757839644300488,0.02612586314740137
+172,CALM1_HUMAN_Roth2017,DeepSequence (published) - single,0.22492618953594645,0.22906576202846649,0.030735041631348013
+173,CALM1_HUMAN_Roth2017,DeepSequence (published),0.22789905079105036,0.22969671581656934,0.029946806999461846
+174,CALM1_HUMAN_Roth2017,ESM-1v (zero shot),0.2617141623664426,0.2652270689648074,0.025228781176764935
+175,CALM1_HUMAN_Roth2017,ESM-1v (+further training),0.2721918164887156,0.2732936184582442,0.027420080243011954
+176,CALM1_HUMAN_Roth2017,DeepSequence (replicated),0.23307246773532486,0.23633457440270692,0.029771707595051294
+177,CALM1_HUMAN_Roth2017,EVMutation (replicated),0.23487698274177318,0.24137564501443426,0.02702788568104191
+178,CALM1_HUMAN_Roth2017,ProtBERT-BFD,0.29084218774026066,0.2918048272665779,0.027753008367247246
+179,CALM1_HUMAN_Roth2017,TAPE,0.12219793057831531,0.12415304964391952,0.019957547541863924
+180,CALM1_HUMAN_Roth2017,UniRep,0.07689880828775296,0.08034276883486069,0.028587576009811506
+181,CALM1_HUMAN_Roth2017,ESM-1b,0.24496865240057522,0.2451996389337371,0.02464309631183189
+434,RL401_YEAST_Bolon2014,gt,1.0,1.0,0.0
+435,RL401_YEAST_Bolon2014,MSA Transformer,0.473016586760016,0.4787538735227928,0.01635240282545988
+436,RL401_YEAST_Bolon2014,PSSM,0.4689225734723619,0.4737125791144366,0.018513489042818982
+437,RL401_YEAST_Bolon2014,EVMutation (published),0.39726026462632547,0.40239056557397346,0.0228080866519743
+438,RL401_YEAST_Bolon2014,DeepSequence (published) - single,0.4295635873940029,0.4359091171923386,0.020125832508895207
+439,RL401_YEAST_Bolon2014,DeepSequence (published),0.4226902253289641,0.4296887661396075,0.02042130768745852
+440,RL401_YEAST_Bolon2014,ESM-1v (zero shot),0.2555902509515176,0.26273723347289263,0.023433435512064758
+441,RL401_YEAST_Bolon2014,ESM-1v (+further training),0.30819485583382966,0.3149369605358795,0.022303272256952115
+442,RL401_YEAST_Bolon2014,DeepSequence (replicated),0.4208601802656798,0.427057708121087,0.022928982939476122
+443,RL401_YEAST_Bolon2014,EVMutation (replicated),0.3795164863252923,0.38353041172591634,0.022737389271498636
+444,RL401_YEAST_Bolon2014,ProtBERT-BFD,0.3305475177656725,0.33333306438483723,0.024066951141161263
+445,RL401_YEAST_Bolon2014,TAPE,0.27031483697506303,0.2748860791185031,0.023958009305355937
+446,RL401_YEAST_Bolon2014,UniRep,0.16611039235432376,0.16134631003484823,0.029548951247265007
+447,RL401_YEAST_Bolon2014,ESM-1b,0.18258453677698241,0.18424186541936532,0.02190939840845001
+28,B3VI55_LIPST_Whitehead2015,gt,1.0,1.0,0.0
+29,B3VI55_LIPST_Whitehead2015,MSA Transformer,0.25670663558953755,0.2568287164790948,0.014303617378529112
+30,B3VI55_LIPST_Whitehead2015,PSSM,0.17356487805783094,0.17229178651491212,0.014476649812023814
+31,B3VI55_LIPST_Whitehead2015,EVMutation (published),0.22415909454397737,0.22201195664151258,0.015151263735652506
+32,B3VI55_LIPST_Whitehead2015,DeepSequence (published) - single,0.2314762900256378,0.23010900979180499,0.015350637015825532
+33,B3VI55_LIPST_Whitehead2015,DeepSequence (published),0.2314408826509719,0.23018066970458023,0.015030815308284325
+34,B3VI55_LIPST_Whitehead2015,ESM-1v (zero shot),0.2361417829414803,0.235599751959246,0.01380734986895981
+35,B3VI55_LIPST_Whitehead2015,ESM-1v (+further training),0.2424504346933426,0.24107272979662708,0.013833009304666605
+36,B3VI55_LIPST_Whitehead2015,DeepSequence (replicated),0.24479443697061962,0.2433108240939475,0.015917153531149808
+37,B3VI55_LIPST_Whitehead2015,EVMutation (replicated),0.2251732484740546,0.2236926065862732,0.015150897486335867
+38,B3VI55_LIPST_Whitehead2015,ProtBERT-BFD,0.22029752268681757,0.22189350964142437,0.012296008345577645
+39,B3VI55_LIPST_Whitehead2015,TAPE,0.03217727698102332,0.033398754623545045,0.014930614521878988
+40,B3VI55_LIPST_Whitehead2015,UniRep,0.052410727620420754,0.051318677282283995,0.012063912744717856
+41,B3VI55_LIPST_Whitehead2015,ESM-1b,0.21414102345131894,0.21411325444647447,0.01376144270479094
+280,MK01_HUMAN_Johannessen,gt,1.0,1.0,0.0
+281,MK01_HUMAN_Johannessen,MSA Transformer,0.21600719562316348,0.21324362844208555,0.010869477909858859
+282,MK01_HUMAN_Johannessen,PSSM,0.24003601283569995,0.23781768411144993,0.012873462567914572
+283,MK01_HUMAN_Johannessen,EVMutation (published),0.2411806504637287,0.23555029839403527,0.014919109500043859
+284,MK01_HUMAN_Johannessen,DeepSequence (published) - single,0.21622583201496393,0.2118857875656635,0.01330985778159382
+285,MK01_HUMAN_Johannessen,DeepSequence (published),0.2262833888426422,0.2222242751773241,0.013845563303298817
+286,MK01_HUMAN_Johannessen,ESM-1v (zero shot),0.18610638097061274,0.1807887324699657,0.011421581908870595
+287,MK01_HUMAN_Johannessen,ESM-1v (+further training),0.1711612951379414,0.1656998628430867,0.011097399419126037
+288,MK01_HUMAN_Johannessen,DeepSequence (replicated),0.2262833888426422,0.2222242751773241,0.013845563303298817
+289,MK01_HUMAN_Johannessen,EVMutation (replicated),0.2411806504637287,0.23555029839403527,0.014919109500043859
+290,MK01_HUMAN_Johannessen,ProtBERT-BFD,0.03761357865712906,0.032009850323277314,0.013532258402167076
+291,MK01_HUMAN_Johannessen,TAPE,0.24566726233768266,0.2425196910807727,0.014722071012130074
+292,MK01_HUMAN_Johannessen,UniRep,0.16466166691457634,0.15865978272069822,0.012603347408739064
+293,MK01_HUMAN_Johannessen,ESM-1b,0.020968477238634434,0.014843912213619878,0.011595884022987897
+350,PA_FLU_Sun2015,gt,1.0,1.0,0.0
+351,PA_FLU_Sun2015,MSA Transformer,0.40499549570010246,0.4071708283839041,0.020127599355418007
+352,PA_FLU_Sun2015,PSSM,0.4870284991956425,0.48503558622411713,0.017254863079323307
+353,PA_FLU_Sun2015,EVMutation (published),0.4799931222249091,0.47963201439188896,0.017318908492685997
+354,PA_FLU_Sun2015,DeepSequence (published) - single,0.47395691140007384,0.47172219190050396,0.017516562617318692
+355,PA_FLU_Sun2015,DeepSequence (published),0.4782770520849889,0.4776150037730753,0.01692698104707
+356,PA_FLU_Sun2015,ESM-1v (zero shot),0.10354167707274801,0.10113956150869738,0.014997098538970815
+357,PA_FLU_Sun2015,ESM-1v (+further training),0.5736052915722709,0.5754278729569327,0.011947354137920797
+358,PA_FLU_Sun2015,DeepSequence (replicated),0.49574925840401,0.49600755923376527,0.014212454786638254
+359,PA_FLU_Sun2015,EVMutation (replicated),0.5055094042813111,0.5062333117389017,0.014772450707935335
+360,PA_FLU_Sun2015,ProtBERT-BFD,0.025567814891833348,0.028868599359005404,0.017475521452476987
+361,PA_FLU_Sun2015,TAPE,0.023351679974442957,0.02740069503679209,0.01716926448528921
+362,PA_FLU_Sun2015,UniRep,0.03233737793643953,0.030453300705534824,0.023348423732335796
+363,PA_FLU_Sun2015,ESM-1b,0.0350533539652243,0.03379222480788328,0.01958723862748923

model/README.md

@@ -0,0 +1,10 @@
+# Local Model Source
+
+This directory contains OneScience `models` sources that ESM imports directly.
+
+- `esm/`: local copy of `onescience.models.esm`
+- `openfold/`: local copy of `onescience.models.openfold`, required by ESMFold
+- `protenix/layer_norm/`: optional fused layer norm dependency used by OpenFold when `LAYERNORM_TYPE=fast_layernorm`
+
+Other OneScience namespaces such as `onescience.datapipes`, `onescience.modules`, and `onescience.utils` are intentionally imported from the installed OneScience environment.
+

model/__init__.py

@@ -0,0 +1,2 @@
+"""Local model packages copied from OneScience for the standalone ESM project."""
+

model/esm/__init__.py

@@ -0,0 +1,13 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from onescience.datapipes import esm as data  # noqa
+from onescience.datapipes.esm import Alphabet, BatchConverter, FastaBatchedDataset  # noqa
+
+from .version import version as __version__  # noqa
+from .esm1 import ProteinBertModel  # noqa
+from .esm2 import ESM2  # noqa
+from .msa_transformer import MSATransformer  #noqa
+from . import pretrained  # noqa

model/esm/esm1.py

@@ -0,0 +1,200 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from onescience.modules.esm import (
+    TransformerLayer,
+    LearnedPositionalEmbedding,
+    SinusoidalPositionalEmbedding,
+    RobertaLMHead,
+    ESM1bLayerNorm,
+    ContactPredictionHead,
+)
+
+
+class ProteinBertModel(nn.Module):
+    @classmethod
+    def add_args(cls, parser):
+        parser.add_argument(
+            "--num_layers", default=36, type=int, metavar="N", help="number of layers"
+        )
+        parser.add_argument(
+            "--embed_dim", default=1280, type=int, metavar="N", help="embedding dimension"
+        )
+        parser.add_argument(
+            "--logit_bias", action="store_true", help="whether to apply bias to logits"
+        )
+        parser.add_argument(
+            "--ffn_embed_dim",
+            default=5120,
+            type=int,
+            metavar="N",
+            help="embedding dimension for FFN",
+        )
+        parser.add_argument(
+            "--attention_heads",
+            default=20,
+            type=int,
+            metavar="N",
+            help="number of attention heads",
+        )
+
+    def __init__(self, args, alphabet):
+        super().__init__()
+        self.args = args
+        self.alphabet_size = len(alphabet)
+        self.padding_idx = alphabet.padding_idx
+        self.mask_idx = alphabet.mask_idx
+        self.cls_idx = alphabet.cls_idx
+        self.eos_idx = alphabet.eos_idx
+        self.prepend_bos = alphabet.prepend_bos
+        self.append_eos = alphabet.append_eos
+        self.emb_layer_norm_before = getattr(self.args, "emb_layer_norm_before", False)
+        if self.args.arch == "roberta_large":
+            self.model_version = "ESM-1b"
+            self._init_submodules_esm1b()
+        else:
+            self.model_version = "ESM-1"
+            self._init_submodules_esm1()
+
+    def _init_submodules_common(self):
+        self.embed_tokens = nn.Embedding(
+            self.alphabet_size, self.args.embed_dim, padding_idx=self.padding_idx
+        )
+        self.layers = nn.ModuleList(
+            [
+                TransformerLayer(
+                    self.args.embed_dim,
+                    self.args.ffn_embed_dim,
+                    self.args.attention_heads,
+                    add_bias_kv=(self.model_version != "ESM-1b"),
+                    use_esm1b_layer_norm=(self.model_version == "ESM-1b"),
+                )
+                for _ in range(self.args.layers)
+            ]
+        )
+
+        self.contact_head = ContactPredictionHead(
+            self.args.layers * self.args.attention_heads,
+            self.prepend_bos,
+            self.append_eos,
+            eos_idx=self.eos_idx,
+        )
+
+    def _init_submodules_esm1b(self):
+        self._init_submodules_common()
+        self.embed_scale = 1
+        self.embed_positions = LearnedPositionalEmbedding(
+            self.args.max_positions, self.args.embed_dim, self.padding_idx
+        )
+        self.emb_layer_norm_before = (
+            ESM1bLayerNorm(self.args.embed_dim) if self.emb_layer_norm_before else None
+        )
+        self.emb_layer_norm_after = ESM1bLayerNorm(self.args.embed_dim)
+        self.lm_head = RobertaLMHead(
+            embed_dim=self.args.embed_dim,
+            output_dim=self.alphabet_size,
+            weight=self.embed_tokens.weight,
+        )
+
+    def _init_submodules_esm1(self):
+        self._init_submodules_common()
+        self.embed_scale = math.sqrt(self.args.embed_dim)
+        self.embed_positions = SinusoidalPositionalEmbedding(self.args.embed_dim, self.padding_idx)
+        self.embed_out = nn.Parameter(torch.zeros((self.alphabet_size, self.args.embed_dim)))
+        self.embed_out_bias = None
+        if self.args.final_bias:
+            self.embed_out_bias = nn.Parameter(torch.zeros(self.alphabet_size))
+
+    def forward(self, tokens, repr_layers=[], need_head_weights=False, return_contacts=False):
+        if return_contacts:
+            need_head_weights = True
+
+        assert tokens.ndim == 2
+        padding_mask = tokens.eq(self.padding_idx)  # B, T
+
+        x = self.embed_scale * self.embed_tokens(tokens)
+
+        if getattr(self.args, "token_dropout", False):
+            x.masked_fill_((tokens == self.mask_idx).unsqueeze(-1), 0.0)
+            # x: B x T x C
+            mask_ratio_train = 0.15 * 0.8
+            src_lengths = (~padding_mask).sum(-1)
+            mask_ratio_observed = (tokens == self.mask_idx).sum(-1).float() / src_lengths
+            x = x * (1 - mask_ratio_train) / (1 - mask_ratio_observed)[:, None, None]
+
+        x = x + self.embed_positions(tokens)
+
+        if self.model_version == "ESM-1b":
+            if self.emb_layer_norm_before:
+                x = self.emb_layer_norm_before(x)
+            if padding_mask is not None:
+                x = x * (1 - padding_mask.unsqueeze(-1).type_as(x))
+
+        repr_layers = set(repr_layers)
+        hidden_representations = {}
+        if 0 in repr_layers:
+            hidden_representations[0] = x
+
+        if need_head_weights:
+            attn_weights = []
+
+        # (B, T, E) => (T, B, E)
+        x = x.transpose(0, 1)
+
+        if not padding_mask.any():
+            padding_mask = None
+
+        for layer_idx, layer in enumerate(self.layers):
+            x, attn = layer(
+                x, self_attn_padding_mask=padding_mask, need_head_weights=need_head_weights
+            )
+            if (layer_idx + 1) in repr_layers:
+                hidden_representations[layer_idx + 1] = x.transpose(0, 1)
+            if need_head_weights:
+                # (H, B, T, T) => (B, H, T, T)
+                attn_weights.append(attn.transpose(1, 0))
+
+        if self.model_version == "ESM-1b":
+            x = self.emb_layer_norm_after(x)
+            x = x.transpose(0, 1)  # (T, B, E) => (B, T, E)
+
+            # last hidden representation should have layer norm applied
+            if (layer_idx + 1) in repr_layers:
+                hidden_representations[layer_idx + 1] = x
+            x = self.lm_head(x)
+        else:
+            x = F.linear(x, self.embed_out, bias=self.embed_out_bias)
+            x = x.transpose(0, 1)  # (T, B, E) => (B, T, E)
+
+        result = {"logits": x, "representations": hidden_representations}
+        if need_head_weights:
+            # attentions: B x L x H x T x T
+            attentions = torch.stack(attn_weights, 1)
+            if self.model_version == "ESM-1":
+                # ESM-1 models have an additional null-token for attention, which we remove
+                attentions = attentions[..., :-1]
+            if padding_mask is not None:
+                attention_mask = 1 - padding_mask.type_as(attentions)
+                attention_mask = attention_mask.unsqueeze(1) * attention_mask.unsqueeze(2)
+                attentions = attentions * attention_mask[:, None, None, :, :]
+            result["attentions"] = attentions
+            if return_contacts:
+                contacts = self.contact_head(tokens, attentions)
+                result["contacts"] = contacts
+
+        return result
+
+    def predict_contacts(self, tokens):
+        return self(tokens, return_contacts=True)["contacts"]
+
+    @property
+    def num_layers(self):
+        return self.args.layers

model/esm/esm2.py

@@ -0,0 +1,152 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Union
+import torch
+import torch.nn as nn
+
+from onescience.datapipes.esm import Alphabet
+from onescience.modules.esm import (
+    ContactPredictionHead,
+    ESM1bLayerNorm,
+    RobertaLMHead,
+    TransformerLayer,
+)
+
+
+class ESM2(nn.Module):
+    def __init__(
+        self,
+        num_layers: int = 33,
+        embed_dim: int = 1280,
+        attention_heads: int = 20,
+        alphabet: Union[Alphabet, str] = "ESM-1b",
+        token_dropout: bool = True,
+    ):
+        super().__init__()
+        self.num_layers = num_layers
+        self.embed_dim = embed_dim
+        self.attention_heads = attention_heads
+        if not isinstance(alphabet, Alphabet):
+            alphabet = Alphabet.from_architecture(alphabet)
+        self.alphabet = alphabet
+        self.alphabet_size = len(alphabet)
+        self.padding_idx = alphabet.padding_idx
+        self.mask_idx = alphabet.mask_idx
+        self.cls_idx = alphabet.cls_idx
+        self.eos_idx = alphabet.eos_idx
+        self.prepend_bos = alphabet.prepend_bos
+        self.append_eos = alphabet.append_eos
+        self.token_dropout = token_dropout
+
+        self._init_submodules()
+
+    def _init_submodules(self):
+        self.embed_scale = 1
+        self.embed_tokens = nn.Embedding(
+            self.alphabet_size,
+            self.embed_dim,
+            padding_idx=self.padding_idx,
+        )
+
+        self.layers = nn.ModuleList(
+            [
+                TransformerLayer(
+                    self.embed_dim,
+                    4 * self.embed_dim,
+                    self.attention_heads,
+                    add_bias_kv=False,
+                    use_esm1b_layer_norm=True,
+                    use_rotary_embeddings=True,
+                )
+                for _ in range(self.num_layers)
+            ]
+        )
+
+        self.contact_head = ContactPredictionHead(
+            self.num_layers * self.attention_heads,
+            self.prepend_bos,
+            self.append_eos,
+            eos_idx=self.eos_idx,
+        )
+        self.emb_layer_norm_after = ESM1bLayerNorm(self.embed_dim)
+
+        self.lm_head = RobertaLMHead(
+            embed_dim=self.embed_dim,
+            output_dim=self.alphabet_size,
+            weight=self.embed_tokens.weight,
+        )
+
+    def forward(self, tokens, repr_layers=[], need_head_weights=False, return_contacts=False):
+        if return_contacts:
+            need_head_weights = True
+
+        assert tokens.ndim == 2
+        padding_mask = tokens.eq(self.padding_idx)  # B, T
+
+        x = self.embed_scale * self.embed_tokens(tokens)
+
+        if self.token_dropout:
+            x.masked_fill_((tokens == self.mask_idx).unsqueeze(-1), 0.0)
+            # x: B x T x C
+            mask_ratio_train = 0.15 * 0.8
+            src_lengths = (~padding_mask).sum(-1)
+            mask_ratio_observed = (tokens == self.mask_idx).sum(-1).to(x.dtype) / src_lengths
+            x = x * (1 - mask_ratio_train) / (1 - mask_ratio_observed)[:, None, None]
+
+        if padding_mask is not None:
+            x = x * (1 - padding_mask.unsqueeze(-1).type_as(x))
+
+        repr_layers = set(repr_layers)
+        hidden_representations = {}
+        if 0 in repr_layers:
+            hidden_representations[0] = x
+
+        if need_head_weights:
+            attn_weights = []
+
+        # (B, T, E) => (T, B, E)
+        x = x.transpose(0, 1)
+
+        if not padding_mask.any():
+            padding_mask = None
+
+        for layer_idx, layer in enumerate(self.layers):
+            x, attn = layer(
+                x,
+                self_attn_padding_mask=padding_mask,
+                need_head_weights=need_head_weights,
+            )
+            if (layer_idx + 1) in repr_layers:
+                hidden_representations[layer_idx + 1] = x.transpose(0, 1)
+            if need_head_weights:
+                # (H, B, T, T) => (B, H, T, T)
+                attn_weights.append(attn.transpose(1, 0))
+
+        x = self.emb_layer_norm_after(x)
+        x = x.transpose(0, 1)  # (T, B, E) => (B, T, E)
+
+        # last hidden representation should have layer norm applied
+        if (layer_idx + 1) in repr_layers:
+            hidden_representations[layer_idx + 1] = x
+        x = self.lm_head(x)
+
+        result = {"logits": x, "representations": hidden_representations}
+        if need_head_weights:
+            # attentions: B x L x H x T x T
+            attentions = torch.stack(attn_weights, 1)
+            if padding_mask is not None:
+                attention_mask = 1 - padding_mask.type_as(attentions)
+                attention_mask = attention_mask.unsqueeze(1) * attention_mask.unsqueeze(2)
+                attentions = attentions * attention_mask[:, None, None, :, :]
+            result["attentions"] = attentions
+            if return_contacts:
+                contacts = self.contact_head(tokens, attentions)
+                result["contacts"] = contacts
+
+        return result
+
+    def predict_contacts(self, tokens):
+        return self(tokens, return_contacts=True)["contacts"]

model/esm/esmfold/v1/categorical_mixture.py

@@ -0,0 +1,43 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+import torch
+
+
+class CategoricalMixture:
+    def __init__(self, param, bins=50, start=0, end=1):
+        # All tensors are of shape ..., bins.
+        self.logits = param
+        bins = torch.linspace(
+            start, end, bins + 1, device=self.logits.device, dtype=self.logits.dtype
+        )
+        self.v_bins = (bins[:-1] + bins[1:]) / 2
+
+    def log_prob(self, true):
+        # Shapes are:
+        #     self.probs: ... x bins
+        #     true      : ...
+        true_index = (
+            (
+                true.unsqueeze(-1)
+                - self.v_bins[
+                    [
+                        None,
+                    ]
+                    * true.ndim
+                ]
+            )
+            .abs()
+            .argmin(-1)
+        )
+        nll = self.logits.log_softmax(-1)
+        return torch.take_along_dim(nll, true_index.unsqueeze(-1), dim=-1).squeeze(-1)
+
+    def mean(self):
+        return (self.logits.softmax(-1) @ self.v_bins.unsqueeze(1)).squeeze(-1)
+
+
+def categorical_lddt(logits, bins=50):
+    # Logits are ..., 37, bins.
+    return CategoricalMixture(logits, bins=bins).mean()

model/esm/esmfold/v1/esmfold.py

@@ -0,0 +1,364 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+import typing as T
+from dataclasses import dataclass, field
+from functools import partial
+
+import torch
+import torch.nn as nn
+from torch import nn
+from torch.nn import LayerNorm
+
+import model.esm as esm
+from model.esm import Alphabet
+from model.esm.esmfold.v1.categorical_mixture import categorical_lddt
+from model.esm.esmfold.v1.misc import (
+    batch_encode_sequences,
+    collate_dense_tensors,
+    output_to_pdb,
+)
+from model.esm.esmfold.v1.trunk import FoldingTrunk, FoldingTrunkConfig
+from onescience.datapipes.openfold.data_transforms import make_atom14_masks
+from onescience.utils.openfold.np import residue_constants
+from onescience.utils.openfold.loss import compute_predicted_aligned_error, compute_tm
+
+
+@dataclass
+class ESMFoldConfig:
+    trunk: T.Any = field(default_factory=FoldingTrunkConfig)
+    lddt_head_hid_dim: int = 128
+
+
+load_fn = esm.pretrained.load_model_and_alphabet
+esm_registry = {
+    "esm2_8M": partial(load_fn, "esm2_t6_8M_UR50D_500K"),
+    "esm2_8M_270K": esm.pretrained.esm2_t6_8M_UR50D,
+    "esm2_35M": partial(load_fn, "esm2_t12_35M_UR50D_500K"),
+    "esm2_35M_270K": esm.pretrained.esm2_t12_35M_UR50D,
+    "esm2_150M": partial(load_fn, "esm2_t30_150M_UR50D_500K"),
+    "esm2_150M_270K": partial(load_fn, "esm2_t30_150M_UR50D_270K"),
+    "esm2_650M": esm.pretrained.esm2_t33_650M_UR50D,
+    "esm2_650M_270K": partial(load_fn, "esm2_t33_650M_270K_UR50D"),
+    "esm2_3B": esm.pretrained.esm2_t36_3B_UR50D,
+    "esm2_3B_270K": partial(load_fn, "esm2_t36_3B_UR50D_500K"),
+    "esm2_15B": esm.pretrained.esm2_t48_15B_UR50D,
+}
+
+
+class ESMFold(nn.Module):
+    def __init__(self, esmfold_config=None, **kwargs):
+        super().__init__()
+
+        self.cfg = esmfold_config if esmfold_config else ESMFoldConfig(**kwargs)
+        cfg = self.cfg
+
+        self.distogram_bins = 64
+
+        self.esm, self.esm_dict = esm_registry.get(cfg.esm_type)()
+
+        self.esm.requires_grad_(False)
+        self.esm.half()
+
+        self.esm_feats = self.esm.embed_dim
+        self.esm_attns = self.esm.num_layers * self.esm.attention_heads
+        self.register_buffer("af2_to_esm", ESMFold._af2_to_esm(self.esm_dict))
+        self.esm_s_combine = nn.Parameter(torch.zeros(self.esm.num_layers + 1))
+
+        c_s = cfg.trunk.sequence_state_dim
+        c_z = cfg.trunk.pairwise_state_dim
+
+        self.esm_s_mlp = nn.Sequential(
+            LayerNorm(self.esm_feats),
+            nn.Linear(self.esm_feats, c_s),
+            nn.ReLU(),
+            nn.Linear(c_s, c_s),
+        )
+        if cfg.use_esm_attn_map:
+            self.esm_z_mlp = nn.Sequential(
+                LayerNorm(self.esm_attns),
+                nn.Linear(self.esm_attns, c_z),
+                nn.ReLU(),
+                nn.Linear(c_z, c_z),
+            )
+
+        # 0 is padding, N is unknown residues, N + 1 is mask.
+        self.n_tokens_embed = residue_constants.restype_num + 3
+        self.pad_idx = 0
+        self.unk_idx = self.n_tokens_embed - 2
+        self.mask_idx = self.n_tokens_embed - 1
+        self.embedding = nn.Embedding(self.n_tokens_embed, c_s, padding_idx=0)
+
+        self.trunk = FoldingTrunk(**cfg.trunk)
+
+        self.distogram_head = nn.Linear(c_z, self.distogram_bins)
+        self.ptm_head = nn.Linear(c_z, self.distogram_bins)
+        self.lm_head = nn.Linear(c_s, self.n_tokens_embed)
+        self.lddt_bins = 50
+        self.lddt_head = nn.Sequential(
+            nn.LayerNorm(cfg.trunk.structure_module.c_s),
+            nn.Linear(cfg.trunk.structure_module.c_s, cfg.lddt_head_hid_dim),
+            nn.Linear(cfg.lddt_head_hid_dim, cfg.lddt_head_hid_dim),
+            nn.Linear(cfg.lddt_head_hid_dim, 37 * self.lddt_bins),
+        )
+
+    @staticmethod
+    def _af2_to_esm(d: Alphabet):
+        # Remember that t is shifted from residue_constants by 1 (0 is padding).
+        esm_reorder = [d.padding_idx] + [
+            d.get_idx(v) for v in residue_constants.restypes_with_x
+        ]
+        return torch.tensor(esm_reorder)
+
+    def _af2_idx_to_esm_idx(self, aa, mask):
+        aa = (aa + 1).masked_fill(mask != 1, 0)
+        return self.af2_to_esm[aa]
+
+    def _compute_language_model_representations(
+        self, esmaa: torch.Tensor
+    ) -> torch.Tensor:
+        """Adds bos/eos tokens for the language model, since the structure module doesn't use these."""
+        batch_size = esmaa.size(0)
+
+        bosi, eosi = self.esm_dict.cls_idx, self.esm_dict.eos_idx
+        bos = esmaa.new_full((batch_size, 1), bosi)
+        eos = esmaa.new_full((batch_size, 1), self.esm_dict.padding_idx)
+        esmaa = torch.cat([bos, esmaa, eos], dim=1)
+        # Use the first padding index as eos during inference.
+        esmaa[range(batch_size), (esmaa != 1).sum(1)] = eosi
+
+        res = self.esm(
+            esmaa,
+            repr_layers=range(self.esm.num_layers + 1),
+            need_head_weights=self.cfg.use_esm_attn_map,
+        )
+        esm_s = torch.stack(
+            [v for _, v in sorted(res["representations"].items())], dim=2
+        )
+        esm_s = esm_s[:, 1:-1]  # B, L, nLayers, C
+        esm_z = (
+            res["attentions"].permute(0, 4, 3, 1, 2).flatten(3, 4)[:, 1:-1, 1:-1, :]
+            if self.cfg.use_esm_attn_map
+            else None
+        )
+        return esm_s, esm_z
+
+    def _mask_inputs_to_esm(self, esmaa, pattern):
+        new_esmaa = esmaa.clone()
+        new_esmaa[pattern == 1] = self.esm_dict.mask_idx
+        return new_esmaa
+
+    def forward(
+        self,
+        aa: torch.Tensor,
+        mask: T.Optional[torch.Tensor] = None,
+        residx: T.Optional[torch.Tensor] = None,
+        masking_pattern: T.Optional[torch.Tensor] = None,
+        num_recycles: T.Optional[int] = None,
+    ):
+        """Runs a forward pass given input tokens. Use `model.infer` to
+        run inference from a sequence.
+
+        Args:
+            aa (torch.Tensor): Tensor containing indices corresponding to amino acids. Indices match
+                onescience.utils.openfold.np.residue_constants.restype_order_with_x.
+            mask (torch.Tensor): Binary tensor with 1 meaning position is unmasked and 0 meaning position is masked.
+            residx (torch.Tensor): Residue indices of amino acids. Will assume contiguous if not provided.
+            masking_pattern (torch.Tensor): Optional masking to pass to the input. Binary tensor of the same size
+                as `aa`. Positions with 1 will be masked. ESMFold sometimes produces different samples when
+                different masks are provided.
+            num_recycles (int): How many recycle iterations to perform. If None, defaults to training max
+                recycles, which is 3.
+        """
+
+        if mask is None:
+            mask = torch.ones_like(aa)
+
+        B = aa.shape[0]
+        L = aa.shape[1]
+        device = aa.device
+
+        if residx is None:
+            residx = torch.arange(L, device=device).expand_as(aa)
+
+        # === ESM ===
+        esmaa = self._af2_idx_to_esm_idx(aa, mask)
+
+        if masking_pattern is not None:
+            esmaa = self._mask_inputs_to_esm(esmaa, masking_pattern)
+
+        esm_s, esm_z = self._compute_language_model_representations(esmaa)
+
+        # Convert esm_s to the precision used by the trunk and
+        # the structure module. These tensors may be a lower precision if, for example,
+        # we're running the language model in fp16 precision.
+        esm_s = esm_s.to(self.esm_s_combine.dtype)
+        esm_s = esm_s.detach()
+
+        # === preprocessing ===
+        esm_s = (self.esm_s_combine.softmax(0).unsqueeze(0) @ esm_s).squeeze(2)
+
+        s_s_0 = self.esm_s_mlp(esm_s)
+        if self.cfg.use_esm_attn_map:
+            esm_z = esm_z.to(self.esm_s_combine.dtype)
+            esm_z = esm_z.detach()
+            s_z_0 = self.esm_z_mlp(esm_z)
+        else:
+            s_z_0 = s_s_0.new_zeros(B, L, L, self.cfg.trunk.pairwise_state_dim)
+
+        s_s_0 += self.embedding(aa)
+
+        structure: dict = self.trunk(
+            s_s_0, s_z_0, aa, residx, mask, no_recycles=num_recycles
+        )
+        # Documenting what we expect:
+        structure = {
+            k: v
+            for k, v in structure.items()
+            if k
+            in [
+                "s_z",
+                "s_s",
+                "frames",
+                "sidechain_frames",
+                "unnormalized_angles",
+                "angles",
+                "positions",
+                "states",
+            ]
+        }
+
+        disto_logits = self.distogram_head(structure["s_z"])
+        disto_logits = (disto_logits + disto_logits.transpose(1, 2)) / 2
+        structure["distogram_logits"] = disto_logits
+
+        lm_logits = self.lm_head(structure["s_s"])
+        structure["lm_logits"] = lm_logits
+
+        structure["aatype"] = aa
+        make_atom14_masks(structure)
+
+        for k in [
+            "atom14_atom_exists",
+            "atom37_atom_exists",
+        ]:
+            structure[k] *= mask.unsqueeze(-1)
+        structure["residue_index"] = residx
+
+        lddt_head = self.lddt_head(structure["states"]).reshape(
+            structure["states"].shape[0], B, L, -1, self.lddt_bins
+        )
+        structure["lddt_head"] = lddt_head
+        plddt = categorical_lddt(lddt_head[-1], bins=self.lddt_bins)
+        structure["plddt"] = (
+            100 * plddt
+        )  # we predict plDDT between 0 and 1, scale to be between 0 and 100.
+
+        ptm_logits = self.ptm_head(structure["s_z"])
+
+        seqlen = mask.type(torch.int64).sum(1)
+        structure["ptm_logits"] = ptm_logits
+        structure["ptm"] = torch.stack(
+            [
+                compute_tm(
+                    batch_ptm_logits[None, :sl, :sl],
+                    max_bins=31,
+                    no_bins=self.distogram_bins,
+                )
+                for batch_ptm_logits, sl in zip(ptm_logits, seqlen)
+            ]
+        )
+        structure.update(
+            compute_predicted_aligned_error(
+                ptm_logits, max_bin=31, no_bins=self.distogram_bins
+            )
+        )
+
+        return structure
+
+    @torch.no_grad()
+    def infer(
+        self,
+        sequences: T.Union[str, T.List[str]],
+        residx=None,
+        masking_pattern: T.Optional[torch.Tensor] = None,
+        num_recycles: T.Optional[int] = None,
+        residue_index_offset: T.Optional[int] = 512,
+        chain_linker: T.Optional[str] = "G" * 25,
+    ):
+        """Runs a forward pass given input sequences.
+
+        Args:
+            sequences (Union[str, List[str]]): A list of sequences to make predictions for. Multimers can also be passed in,
+                each chain should be separated by a ':' token (e.g. "<chain1>:<chain2>:<chain3>").
+            residx (torch.Tensor): Residue indices of amino acids. Will assume contiguous if not provided.
+            masking_pattern (torch.Tensor): Optional masking to pass to the input. Binary tensor of the same size
+                as `aa`. Positions with 1 will be masked. ESMFold sometimes produces different samples when
+                different masks are provided.
+            num_recycles (int): How many recycle iterations to perform. If None, defaults to training max
+                recycles (cfg.trunk.max_recycles), which is 4.
+            residue_index_offset (int): Residue index separation between chains if predicting a multimer. Has no effect on
+                single chain predictions. Default: 512.
+            chain_linker (str): Linker to use between chains if predicting a multimer. Has no effect on single chain
+                predictions. Default: length-25 poly-G ("G" * 25).
+        """
+        if isinstance(sequences, str):
+            sequences = [sequences]
+
+        aatype, mask, _residx, linker_mask, chain_index = batch_encode_sequences(
+            sequences, residue_index_offset, chain_linker
+        )
+
+        if residx is None:
+            residx = _residx
+        elif not isinstance(residx, torch.Tensor):
+            residx = collate_dense_tensors(residx)
+
+        aatype, mask, residx, linker_mask = map(
+            lambda x: x.to(self.device), (aatype, mask, residx, linker_mask)
+        )
+
+        output = self.forward(
+            aatype,
+            mask=mask,
+            residx=residx,
+            masking_pattern=masking_pattern,
+            num_recycles=num_recycles,
+        )
+
+        output["atom37_atom_exists"] = output[
+            "atom37_atom_exists"
+        ] * linker_mask.unsqueeze(2)
+
+        output["mean_plddt"] = (output["plddt"] * output["atom37_atom_exists"]).sum(
+            dim=(1, 2)
+        ) / output["atom37_atom_exists"].sum(dim=(1, 2))
+        output["chain_index"] = chain_index
+
+        return output
+
+    def output_to_pdb(self, output: T.Dict) -> T.List[str]:
+        """Returns the pbd (file) string from the model given the model output."""
+        return output_to_pdb(output)
+
+    def infer_pdbs(self, seqs: T.List[str], *args, **kwargs) -> T.List[str]:
+        """Returns list of pdb (files) strings from the model given a list of input sequences."""
+        output = self.infer(seqs, *args, **kwargs)
+        return self.output_to_pdb(output)
+
+    def infer_pdb(self, sequence: str, *args, **kwargs) -> str:
+        """Returns the pdb (file) string from the model given an input sequence."""
+        return self.infer_pdbs([sequence], *args, **kwargs)[0]
+
+    def set_chunk_size(self, chunk_size: T.Optional[int]):
+        # This parameter means the axial attention will be computed
+        # in a chunked manner. This should make the memory used more or less O(L) instead of O(L^2).
+        # It's equivalent to running a for loop over chunks of the dimension we're iterative over,
+        # where the chunk_size is the size of the chunks, so 128 would mean to parse 128-lengthed chunks.
+        # Setting the value to None will return to default behavior, disable chunking.
+        self.trunk.set_chunk_size(chunk_size)
+
+    @property
+    def device(self):
+        return self.esm_s_combine.device

model/esm/esmfold/v1/misc.py

@@ -0,0 +1,309 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+import typing as T
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from einops import rearrange, repeat
+from torch import nn
+from onescience.utils.openfold.np import residue_constants
+from onescience.utils.openfold.np.protein import Protein as OFProtein
+from onescience.utils.openfold.np.protein import to_pdb
+from onescience.utils.openfold.feats import atom14_to_atom37
+
+
+def encode_sequence(
+    seq: str,
+    residue_index_offset: T.Optional[int] = 512,
+    chain_linker: T.Optional[str] = "G" * 25,
+) -> T.Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    if chain_linker is None:
+        chain_linker = ""
+    if residue_index_offset is None:
+        residue_index_offset = 0
+
+    chains = seq.split(":")
+    seq = chain_linker.join(chains)
+
+    unk_idx = residue_constants.restype_order_with_x["X"]
+    encoded = torch.tensor(
+        [residue_constants.restype_order_with_x.get(aa, unk_idx) for aa in seq]
+    )
+    residx = torch.arange(len(encoded))
+
+    if residue_index_offset > 0:
+        start = 0
+        for i, chain in enumerate(chains):
+            residx[start : start + len(chain) + len(chain_linker)] += (
+                i * residue_index_offset
+            )
+            start += len(chain) + len(chain_linker)
+
+    linker_mask = torch.ones_like(encoded, dtype=torch.float32)
+    chain_index = []
+    offset = 0
+    for i, chain in enumerate(chains):
+        if i > 0:
+            chain_index.extend([i - 1] * len(chain_linker))
+        chain_index.extend([i] * len(chain))
+        offset += len(chain)
+        linker_mask[offset : offset + len(chain_linker)] = 0
+        offset += len(chain_linker)
+
+    chain_index = torch.tensor(chain_index, dtype=torch.int64)
+
+    return encoded, residx, linker_mask, chain_index
+
+
+def batch_encode_sequences(
+    sequences: T.Sequence[str],
+    residue_index_offset: T.Optional[int] = 512,
+    chain_linker: T.Optional[str] = "G" * 25,
+) -> T.Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+
+    aatype_list = []
+    residx_list = []
+    linker_mask_list = []
+    chain_index_list = []
+    for seq in sequences:
+        aatype_seq, residx_seq, linker_mask_seq, chain_index_seq = encode_sequence(
+            seq,
+            residue_index_offset=residue_index_offset,
+            chain_linker=chain_linker,
+        )
+        aatype_list.append(aatype_seq)
+        residx_list.append(residx_seq)
+        linker_mask_list.append(linker_mask_seq)
+        chain_index_list.append(chain_index_seq)
+
+    aatype = collate_dense_tensors(aatype_list)
+    mask = collate_dense_tensors(
+        [aatype.new_ones(len(aatype_seq)) for aatype_seq in aatype_list]
+    )
+    residx = collate_dense_tensors(residx_list)
+    linker_mask = collate_dense_tensors(linker_mask_list)
+    chain_index_list = collate_dense_tensors(chain_index_list, -1)
+
+    return aatype, mask, residx, linker_mask, chain_index_list
+
+
+def output_to_pdb(output: T.Dict) -> T.List[str]:
+    """Returns the pbd (file) string from the model given the model output."""
+    # atom14_to_atom37 must be called first, as it fails on latest numpy if the
+    # input is a numpy array. It will work if the input is a torch tensor.
+    final_atom_positions = atom14_to_atom37(output["positions"][-1], output)
+    output = {k: v.to("cpu").numpy() for k, v in output.items()}
+    final_atom_positions = final_atom_positions.cpu().numpy()
+    final_atom_mask = output["atom37_atom_exists"]
+    pdbs = []
+    for i in range(output["aatype"].shape[0]):
+        aa = output["aatype"][i]
+        pred_pos = final_atom_positions[i]
+        mask = final_atom_mask[i]
+        resid = output["residue_index"][i] + 1
+        pred = OFProtein(
+            aatype=aa,
+            atom_positions=pred_pos,
+            atom_mask=mask,
+            residue_index=resid,
+            b_factors=output["plddt"][i],
+            chain_index=output["chain_index"][i] if "chain_index" in output else None,
+        )
+        pdbs.append(to_pdb(pred))
+    return pdbs
+
+
+def collate_dense_tensors(
+    samples: T.List[torch.Tensor], pad_v: float = 0
+) -> torch.Tensor:
+    """
+    Takes a list of tensors with the following dimensions:
+        [(d_11,       ...,           d_1K),
+         (d_21,       ...,           d_2K),
+         ...,
+         (d_N1,       ...,           d_NK)]
+    and stack + pads them into a single tensor of:
+    (N, max_i=1,N { d_i1 }, ..., max_i=1,N {diK})
+    """
+    if len(samples) == 0:
+        return torch.Tensor()
+    if len(set(x.dim() for x in samples)) != 1:
+        raise RuntimeError(
+            f"Samples has varying dimensions: {[x.dim() for x in samples]}"
+        )
+    (device,) = tuple(set(x.device for x in samples))  # assumes all on same device
+    max_shape = [max(lst) for lst in zip(*[x.shape for x in samples])]
+    result = torch.empty(
+        len(samples), *max_shape, dtype=samples[0].dtype, device=device
+    )
+    result.fill_(pad_v)
+    for i in range(len(samples)):
+        result_i = result[i]
+        t = samples[i]
+        result_i[tuple(slice(0, k) for k in t.shape)] = t
+    return result
+
+
+class Attention(nn.Module):
+    def __init__(self, embed_dim, num_heads, head_width, gated=False):
+        super().__init__()
+        assert embed_dim == num_heads * head_width
+
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_width = head_width
+
+        self.proj = nn.Linear(embed_dim, embed_dim * 3, bias=False)
+        self.o_proj = nn.Linear(embed_dim, embed_dim, bias=True)
+        self.gated = gated
+        if gated:
+            self.g_proj = nn.Linear(embed_dim, embed_dim)
+            torch.nn.init.zeros_(self.g_proj.weight)
+            torch.nn.init.ones_(self.g_proj.bias)
+
+        self.rescale_factor = self.head_width**-0.5
+
+        torch.nn.init.zeros_(self.o_proj.bias)
+
+    def forward(self, x, mask=None, bias=None, indices=None):
+        """
+        Basic self attention with optional mask and external pairwise bias.
+        To handle sequences of different lengths, use mask.
+
+        Inputs:
+          x: batch of input sequneces (.. x L x C)
+          mask: batch of boolean masks where 1=valid, 0=padding position (.. x L_k). optional.
+          bias: batch of scalar pairwise attention biases (.. x Lq x Lk x num_heads). optional.
+
+        Outputs:
+          sequence projection (B x L x embed_dim), attention maps (B x L x L x num_heads)
+        """
+
+        t = rearrange(self.proj(x), "... l (h c) -> ... h l c", h=self.num_heads)
+        q, k, v = t.chunk(3, dim=-1)
+
+        q = self.rescale_factor * q
+        a = torch.einsum("...qc,...kc->...qk", q, k)
+
+        # Add external attention bias.
+        if bias is not None:
+            a = a + rearrange(bias, "... lq lk h -> ... h lq lk")
+
+        # Do not attend to padding tokens.
+        if mask is not None:
+            mask = repeat(
+                mask, "... lk -> ... h lq lk", h=self.num_heads, lq=q.shape[-2]
+            )
+            a = a.masked_fill(mask == False, -np.inf)
+
+        a = F.softmax(a, dim=-1)
+
+        y = torch.einsum("...hqk,...hkc->...qhc", a, v)
+        y = rearrange(y, "... h c -> ... (h c)", h=self.num_heads)
+
+        if self.gated:
+            y = self.g_proj(x).sigmoid() * y
+        y = self.o_proj(y)
+
+        return y, rearrange(a, "... lq lk h -> ... h lq lk")
+
+
+class Dropout(nn.Module):
+    """
+    Implementation of dropout with the ability to share the dropout mask
+    along a particular dimension.
+    """
+
+    def __init__(self, r: float, batch_dim: T.Union[int, T.List[int]]):
+        super(Dropout, self).__init__()
+
+        self.r = r
+        if type(batch_dim) == int:
+            batch_dim = [batch_dim]
+        self.batch_dim = batch_dim
+        self.dropout = nn.Dropout(self.r)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        shape = list(x.shape)
+        if self.batch_dim is not None:
+            for bd in self.batch_dim:
+                shape[bd] = 1
+        return x * self.dropout(x.new_ones(shape))
+
+
+class SequenceToPair(nn.Module):
+    def __init__(self, sequence_state_dim, inner_dim, pairwise_state_dim):
+        super().__init__()
+
+        self.layernorm = nn.LayerNorm(sequence_state_dim)
+        self.proj = nn.Linear(sequence_state_dim, inner_dim * 2, bias=True)
+        self.o_proj = nn.Linear(2 * inner_dim, pairwise_state_dim, bias=True)
+
+        torch.nn.init.zeros_(self.proj.bias)
+        torch.nn.init.zeros_(self.o_proj.bias)
+
+    def forward(self, sequence_state):
+        """
+        Inputs:
+          sequence_state: B x L x sequence_state_dim
+
+        Output:
+          pairwise_state: B x L x L x pairwise_state_dim
+
+        Intermediate state:
+          B x L x L x 2*inner_dim
+        """
+
+        assert len(sequence_state.shape) == 3
+
+        s = self.layernorm(sequence_state)
+        s = self.proj(s)
+        q, k = s.chunk(2, dim=-1)
+
+        prod = q[:, None, :, :] * k[:, :, None, :]
+        diff = q[:, None, :, :] - k[:, :, None, :]
+
+        x = torch.cat([prod, diff], dim=-1)
+        x = self.o_proj(x)
+
+        return x
+
+
+class PairToSequence(nn.Module):
+    def __init__(self, pairwise_state_dim, num_heads):
+        super().__init__()
+
+        self.layernorm = nn.LayerNorm(pairwise_state_dim)
+        self.linear = nn.Linear(pairwise_state_dim, num_heads, bias=False)
+
+    def forward(self, pairwise_state):
+        """
+        Inputs:
+          pairwise_state: B x L x L x pairwise_state_dim
+
+        Output:
+          pairwise_bias: B x L x L x num_heads
+        """
+        assert len(pairwise_state.shape) == 4
+        z = self.layernorm(pairwise_state)
+        pairwise_bias = self.linear(z)
+        return pairwise_bias
+
+
+class ResidueMLP(nn.Module):
+    def __init__(self, embed_dim, inner_dim, norm=nn.LayerNorm, dropout=0):
+        super().__init__()
+
+        self.mlp = nn.Sequential(
+            norm(embed_dim),
+            nn.Linear(embed_dim, inner_dim),
+            nn.ReLU(),
+            nn.Linear(inner_dim, embed_dim),
+            nn.Dropout(dropout),
+        )
+
+    def forward(self, x):
+        return x + self.mlp(x)

model/esm/esmfold/v1/pretrained.py

@@ -0,0 +1,197 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from pathlib import Path
+
+import torch
+
+from model.esm.esmfold.v1.esmfold import ESMFold
+
+def _upgrade_state_dict(model_state):
+    """Map legacy ESMFold checkpoint keys to the current OpenFold modules."""
+    point_projection_prefixes = (
+        "trunk.structure_module.ipa.linear_q_points",
+        "trunk.structure_module.ipa.linear_kv_points",
+        "trunk.structure_module.ipa.linear_k_points",
+        "trunk.structure_module.ipa.linear_v_points",
+    )
+    for prefix in point_projection_prefixes:
+        for suffix in ("weight", "bias"):
+            old_key = f"{prefix}.{suffix}"
+            new_key = f"{prefix}.linear.{suffix}"
+            if old_key in model_state and new_key not in model_state:
+                model_state[new_key] = model_state.pop(old_key)
+    return model_state
+
+def _load_model(model_name):
+    if model_name.endswith(".pt"):  # local, treat as filepath
+        model_path = Path(model_name)
+        model_data = torch.load(str(model_path), map_location="cpu")
+    else:  # load from hub
+        url = f"https://dl.fbaipublicfiles.com/fair-esm/models/{model_name}.pt"
+        model_data = torch.hub.load_state_dict_from_url(url, progress=False, map_location="cpu")
+
+    cfg = model_data["cfg"]["model"]
+    model_state = _upgrade_state_dict(model_data["model"])
+    # model_state = model_data["model"]
+    model = ESMFold(esmfold_config=cfg)
+
+    expected_keys = set(model.state_dict().keys())
+    found_keys = set(model_state.keys())
+
+    missing_essential_keys = []
+    for missing_key in expected_keys - found_keys:
+        if not missing_key.startswith("esm."):
+            missing_essential_keys.append(missing_key)
+
+    if missing_essential_keys:
+        raise RuntimeError(f"Keys '{', '.join(missing_essential_keys)}' are missing.")
+
+    model.load_state_dict(model_state, strict=False)
+
+    return model
+
+
+def esmfold_v0():
+    """
+    ESMFold v0 model with 3B ESM-2, 48 folding blocks.
+    This version was used for the paper (Lin et al, 2022). It was trained
+    on all PDB chains until 2020-05, to ensure temporal holdout with CASP14
+    and the CAMEO validation and test set reported there.
+    """
+    return _load_model("esmfold_3B_v0")
+
+
+def esmfold_v1():
+    """
+    ESMFold v1 model using 3B ESM-2, 48 folding blocks.
+    ESMFold provides fast high accuracy atomic level structure prediction
+    directly from the individual sequence of a protein. ESMFold uses the ESM2
+    protein language model to extract meaningful representations from the
+    protein sequence.
+    """
+    return _load_model("esmfold_3B_v1")
+
+
+def esmfold_structure_module_only_8M():
+    """
+    ESMFold baseline model using 8M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_8M")
+
+
+def esmfold_structure_module_only_8M_270K():
+    """
+    ESMFold baseline model using 8M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_8M_270K")
+
+
+def esmfold_structure_module_only_35M():
+    """
+    ESMFold baseline model using 35M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_35M")
+
+
+def esmfold_structure_module_only_35M_270K():
+    """
+    ESMFold baseline model using 35M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_35M_270K")
+
+
+def esmfold_structure_module_only_150M():
+    """
+    ESMFold baseline model using 150M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_150M")
+
+
+def esmfold_structure_module_only_150M_270K():
+    """
+    ESMFold baseline model using 150M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_150M_270K")
+
+
+def esmfold_structure_module_only_650M():
+    """
+    ESMFold baseline model using 650M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_650M")
+
+
+def esmfold_structure_module_only_650M_270K():
+    """
+    ESMFold baseline model using 650M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_650M_270K")
+
+
+def esmfold_structure_module_only_3B():
+    """
+    ESMFold baseline model using 3B ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_3B")
+
+
+def esmfold_structure_module_only_3B_270K():
+    """
+    ESMFold baseline model using 3B ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_3B_270K")
+
+
+def esmfold_structure_module_only_15B():
+    """
+    ESMFold baseline model using 15B ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    The 15B parameter ESM-2 was not trained out to 500K updates
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    return _load_model("esmfold_structure_module_only_15B")

model/esm/esmfold/v1/tri_self_attn_block.py

@@ -0,0 +1,160 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+import torch
+from model.openfold.triangular_attention import (
+    TriangleAttentionEndingNode,
+    TriangleAttentionStartingNode,
+)
+from model.openfold.triangular_multiplicative_update import (
+    TriangleMultiplicationIncoming,
+    TriangleMultiplicationOutgoing,
+)
+from torch import nn
+
+from model.esm.esmfold.v1.misc import (
+    Attention,
+    Dropout,
+    PairToSequence,
+    ResidueMLP,
+    SequenceToPair,
+)
+
+
+class TriangularSelfAttentionBlock(nn.Module):
+    def __init__(
+        self,
+        sequence_state_dim,
+        pairwise_state_dim,
+        sequence_head_width,
+        pairwise_head_width,
+        dropout=0,
+        **__kwargs,
+    ):
+        super().__init__()
+
+        assert sequence_state_dim % sequence_head_width == 0
+        assert pairwise_state_dim % pairwise_head_width == 0
+        sequence_num_heads = sequence_state_dim // sequence_head_width
+        pairwise_num_heads = pairwise_state_dim // pairwise_head_width
+        assert sequence_state_dim == sequence_num_heads * sequence_head_width
+        assert pairwise_state_dim == pairwise_num_heads * pairwise_head_width
+        assert pairwise_state_dim % 2 == 0
+
+        self.sequence_state_dim = sequence_state_dim
+        self.pairwise_state_dim = pairwise_state_dim
+
+        self.layernorm_1 = nn.LayerNorm(sequence_state_dim)
+
+        self.sequence_to_pair = SequenceToPair(
+            sequence_state_dim, pairwise_state_dim // 2, pairwise_state_dim
+        )
+        self.pair_to_sequence = PairToSequence(pairwise_state_dim, sequence_num_heads)
+
+        self.seq_attention = Attention(
+            sequence_state_dim, sequence_num_heads, sequence_head_width, gated=True
+        )
+        self.tri_mul_out = TriangleMultiplicationOutgoing(
+            pairwise_state_dim,
+            pairwise_state_dim,
+        )
+        self.tri_mul_in = TriangleMultiplicationIncoming(
+            pairwise_state_dim,
+            pairwise_state_dim,
+        )
+        self.tri_att_start = TriangleAttentionStartingNode(
+            pairwise_state_dim,
+            pairwise_head_width,
+            pairwise_num_heads,
+            inf=1e9,
+        )  # type: ignore
+        self.tri_att_end = TriangleAttentionEndingNode(
+            pairwise_state_dim,
+            pairwise_head_width,
+            pairwise_num_heads,
+            inf=1e9,
+        )  # type: ignore
+
+        self.mlp_seq = ResidueMLP(sequence_state_dim, 4 * sequence_state_dim, dropout=dropout)
+        self.mlp_pair = ResidueMLP(pairwise_state_dim, 4 * pairwise_state_dim, dropout=dropout)
+
+        assert dropout < 0.4
+        self.drop = nn.Dropout(dropout)
+        self.row_drop = Dropout(dropout * 2, 2)
+        self.col_drop = Dropout(dropout * 2, 1)
+
+        torch.nn.init.zeros_(self.tri_mul_in.linear_z.weight)
+        torch.nn.init.zeros_(self.tri_mul_in.linear_z.bias)
+        torch.nn.init.zeros_(self.tri_mul_out.linear_z.weight)
+        torch.nn.init.zeros_(self.tri_mul_out.linear_z.bias)
+        torch.nn.init.zeros_(self.tri_att_start.mha.linear_o.weight)
+        torch.nn.init.zeros_(self.tri_att_start.mha.linear_o.bias)
+        torch.nn.init.zeros_(self.tri_att_end.mha.linear_o.weight)
+        torch.nn.init.zeros_(self.tri_att_end.mha.linear_o.bias)
+
+        torch.nn.init.zeros_(self.sequence_to_pair.o_proj.weight)
+        torch.nn.init.zeros_(self.sequence_to_pair.o_proj.bias)
+        torch.nn.init.zeros_(self.pair_to_sequence.linear.weight)
+        torch.nn.init.zeros_(self.seq_attention.o_proj.weight)
+        torch.nn.init.zeros_(self.seq_attention.o_proj.bias)
+        torch.nn.init.zeros_(self.mlp_seq.mlp[-2].weight)
+        torch.nn.init.zeros_(self.mlp_seq.mlp[-2].bias)
+        torch.nn.init.zeros_(self.mlp_pair.mlp[-2].weight)
+        torch.nn.init.zeros_(self.mlp_pair.mlp[-2].bias)
+
+    def forward(self, sequence_state, pairwise_state, mask=None, chunk_size=None, **__kwargs):
+        """
+        Inputs:
+          sequence_state: B x L x sequence_state_dim
+          pairwise_state: B x L x L x pairwise_state_dim
+          mask: B x L boolean tensor of valid positions
+
+        Output:
+          sequence_state: B x L x sequence_state_dim
+          pairwise_state: B x L x L x pairwise_state_dim
+        """
+        assert len(sequence_state.shape) == 3
+        assert len(pairwise_state.shape) == 4
+        if mask is not None:
+            assert len(mask.shape) == 2
+
+        batch_dim, seq_dim, sequence_state_dim = sequence_state.shape
+        pairwise_state_dim = pairwise_state.shape[3]
+        assert sequence_state_dim == self.sequence_state_dim
+        assert pairwise_state_dim == self.pairwise_state_dim
+        assert batch_dim == pairwise_state.shape[0]
+        assert seq_dim == pairwise_state.shape[1]
+        assert seq_dim == pairwise_state.shape[2]
+
+        # Update sequence state
+        bias = self.pair_to_sequence(pairwise_state)
+
+        # Self attention with bias + mlp.
+        y = self.layernorm_1(sequence_state)
+        y, _ = self.seq_attention(y, mask=mask, bias=bias)
+        sequence_state = sequence_state + self.drop(y)
+        sequence_state = self.mlp_seq(sequence_state)
+
+        # Update pairwise state
+        pairwise_state = pairwise_state + self.sequence_to_pair(sequence_state)
+
+        # Axial attention with triangular bias.
+        tri_mask = mask.unsqueeze(2) * mask.unsqueeze(1) if mask is not None else None
+        pairwise_state = pairwise_state + self.row_drop(
+            self.tri_mul_out(pairwise_state, mask=tri_mask)
+        )
+        pairwise_state = pairwise_state + self.col_drop(
+            self.tri_mul_in(pairwise_state, mask=tri_mask)
+        )
+        pairwise_state = pairwise_state + self.row_drop(
+            self.tri_att_start(pairwise_state, mask=tri_mask, chunk_size=chunk_size)
+        )
+        pairwise_state = pairwise_state + self.col_drop(
+            self.tri_att_end(pairwise_state, mask=tri_mask, chunk_size=chunk_size)
+        )
+
+        # MLP over pairs.
+        pairwise_state = self.mlp_pair(pairwise_state)
+
+        return sequence_state, pairwise_state

model/esm/esmfold/v1/trunk.py

@@ -0,0 +1,245 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+import typing as T
+from contextlib import ExitStack
+# from dataclasses import dataclass
+from dataclasses import dataclass, field
+
+import torch
+import torch.nn as nn
+from model.openfold.structure_module import StructureModule
+
+from model.esm.esmfold.v1.tri_self_attn_block import TriangularSelfAttentionBlock
+
+
+@dataclass
+class StructureModuleConfig:
+    c_s: int = 384
+    c_z: int = 128
+    c_ipa: int = 16
+    c_resnet: int = 128
+    no_heads_ipa: int = 12
+    no_qk_points: int = 4
+    no_v_points: int = 8
+    dropout_rate: float = 0.1
+    no_blocks: int = 8
+    no_transition_layers: int = 1
+    no_resnet_blocks: int = 2
+    no_angles: int = 7
+    trans_scale_factor: int = 10
+    epsilon: float = 1e-8
+    inf: float = 1e5
+
+
+@dataclass
+class FoldingTrunkConfig:
+    _name: str = "FoldingTrunkConfig"
+    num_blocks: int = 48
+    sequence_state_dim: int = 1024
+    pairwise_state_dim: int = 128
+    sequence_head_width: int = 32
+    pairwise_head_width: int = 32
+    position_bins: int = 32
+    dropout: float = 0
+    layer_drop: float = 0
+    cpu_grad_checkpoint: bool = False
+
+    max_recycles: int = 4
+    chunk_size: T.Optional[int] = None
+
+    # structure_module: StructureModuleConfig = StructureModuleConfig()
+    structure_module: StructureModuleConfig = field(default_factory=StructureModuleConfig)
+
+
+def get_axial_mask(mask):
+    """
+    Helper to convert B x L mask of valid positions to axial mask used
+    in row column attentions.
+
+    Input:
+      mask: B x L tensor of booleans
+
+    Output:
+      mask: B x L x L tensor of booleans
+    """
+
+    if mask is None:
+        return None
+    assert len(mask.shape) == 2
+    batch_dim, seq_dim = mask.shape
+    m = mask.unsqueeze(1).expand(batch_dim, seq_dim, seq_dim)
+    m = m.reshape(batch_dim * seq_dim, seq_dim)
+    return m
+
+
+class RelativePosition(nn.Module):
+    def __init__(self, bins, pairwise_state_dim):
+        super().__init__()
+        self.bins = bins
+
+        # Note an additional offset is used so that the 0th position
+        # is reserved for masked pairs.
+        self.embedding = torch.nn.Embedding(2 * bins + 2, pairwise_state_dim)
+
+    def forward(self, residue_index, mask=None):
+        """
+        Input:
+          residue_index: B x L tensor of indices (dytpe=torch.long)
+          mask: B x L tensor of booleans
+
+        Output:
+          pairwise_state: B x L x L x pairwise_state_dim tensor of embeddings
+        """
+
+        assert residue_index.dtype == torch.long
+        if mask is not None:
+            assert residue_index.shape == mask.shape
+
+        diff = residue_index[:, None, :] - residue_index[:, :, None]
+        diff = diff.clamp(-self.bins, self.bins)
+        diff = diff + self.bins + 1  # Add 1 to adjust for padding index.
+
+        if mask is not None:
+            mask = mask[:, None, :] * mask[:, :, None]
+            diff[mask == False] = 0
+
+        output = self.embedding(diff)
+        return output
+
+
+class FoldingTrunk(nn.Module):
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.cfg = FoldingTrunkConfig(**kwargs)
+        assert self.cfg.max_recycles > 0
+
+        c_s = self.cfg.sequence_state_dim
+        c_z = self.cfg.pairwise_state_dim
+
+        assert c_s % self.cfg.sequence_head_width == 0
+        assert c_z % self.cfg.pairwise_head_width == 0
+        block = TriangularSelfAttentionBlock
+
+        self.pairwise_positional_embedding = RelativePosition(self.cfg.position_bins, c_z)
+
+        self.blocks = nn.ModuleList(
+            [
+                block(
+                    sequence_state_dim=c_s,
+                    pairwise_state_dim=c_z,
+                    sequence_head_width=self.cfg.sequence_head_width,
+                    pairwise_head_width=self.cfg.pairwise_head_width,
+                    dropout=self.cfg.dropout,
+                )
+                for i in range(self.cfg.num_blocks)
+            ]
+        )
+
+        self.recycle_bins = 15
+        self.recycle_s_norm = nn.LayerNorm(c_s)
+        self.recycle_z_norm = nn.LayerNorm(c_z)
+        self.recycle_disto = nn.Embedding(self.recycle_bins, c_z)
+        self.recycle_disto.weight[0].detach().zero_()
+
+        self.structure_module = StructureModule(**self.cfg.structure_module)  # type: ignore
+        self.trunk2sm_s = nn.Linear(c_s, self.structure_module.c_s)
+        self.trunk2sm_z = nn.Linear(c_z, self.structure_module.c_z)
+
+        self.chunk_size = self.cfg.chunk_size
+
+    def set_chunk_size(self, chunk_size):
+        # This parameter means the axial attention will be computed
+        # in a chunked manner. This should make the memory used more or less O(L) instead of O(L^2).
+        # It's equivalent to running a for loop over chunks of the dimension we're iterative over,
+        # where the chunk_size is the size of the chunks, so 128 would mean to parse 128-lengthed chunks.
+        self.chunk_size = chunk_size
+
+    def forward(self, seq_feats, pair_feats, true_aa, residx, mask, no_recycles: T.Optional[int] = None):
+        """
+        Inputs:
+          seq_feats:     B x L x C            tensor of sequence features
+          pair_feats:    B x L x L x C        tensor of pair features
+          residx:        B x L                long tensor giving the position in the sequence
+          mask:          B x L                boolean tensor indicating valid residues
+
+        Output:
+          predicted_structure: B x L x (num_atoms_per_residue * 3) tensor wrapped in a Coordinates object
+        """
+
+        device = seq_feats.device
+        s_s_0 = seq_feats
+        s_z_0 = pair_feats
+
+        if no_recycles is None:
+            no_recycles = self.cfg.max_recycles
+        else:
+            assert no_recycles >= 0, "Number of recycles must not be negative."
+            no_recycles += 1  # First 'recycle' is just the standard forward pass through the model.
+
+        def trunk_iter(s, z, residx, mask):
+            z = z + self.pairwise_positional_embedding(residx, mask=mask)
+
+            for block in self.blocks:
+                s, z = block(s, z, mask=mask, residue_index=residx, chunk_size=self.chunk_size)
+            return s, z
+
+        s_s = s_s_0
+        s_z = s_z_0
+        recycle_s = torch.zeros_like(s_s)
+        recycle_z = torch.zeros_like(s_z)
+        recycle_bins = torch.zeros(*s_z.shape[:-1], device=device, dtype=torch.int64)
+
+        assert no_recycles > 0
+        for recycle_idx in range(no_recycles):
+            with ExitStack() if recycle_idx == no_recycles - 1 else torch.no_grad():
+                # === Recycling ===
+                recycle_s = self.recycle_s_norm(recycle_s.detach())
+                recycle_z = self.recycle_z_norm(recycle_z.detach())
+                recycle_z += self.recycle_disto(recycle_bins.detach())
+
+                s_s, s_z = trunk_iter(s_s_0 + recycle_s, s_z_0 + recycle_z, residx, mask)
+
+                # === Structure module ===
+                structure = self.structure_module(
+                    {"single": self.trunk2sm_s(s_s), "pair": self.trunk2sm_z(s_z)},
+                    true_aa,
+                    mask.float(),
+                )
+
+                recycle_s = s_s
+                recycle_z = s_z
+                # Distogram needs the N, CA, C coordinates, and bin constants same as alphafold.
+                recycle_bins = FoldingTrunk.distogram(
+                    structure["positions"][-1][:, :, :3],
+                    3.375,
+                    21.375,
+                    self.recycle_bins,
+                )
+
+        assert isinstance(structure, dict)  # type: ignore
+        structure["s_s"] = s_s
+        structure["s_z"] = s_z
+
+        return structure
+
+    @staticmethod
+    def distogram(coords, min_bin, max_bin, num_bins):
+        # Coords are [... L x 3 x 3], where it's [N, CA, C] x 3 coordinates.
+        boundaries = torch.linspace(
+            min_bin,
+            max_bin,
+            num_bins - 1,
+            device=coords.device,
+        )
+        boundaries = boundaries**2
+        N, CA, C = [x.squeeze(-2) for x in coords.chunk(3, dim=-2)]
+        # Infer CB coordinates.
+        b = CA - N
+        c = C - CA
+        a = b.cross(c, dim=-1)
+        CB = -0.58273431 * a + 0.56802827 * b - 0.54067466 * c + CA
+        dists = (CB[..., None, :, :] - CB[..., :, None, :]).pow(2).sum(dim=-1, keepdims=True)
+        bins = torch.sum(dists > boundaries, dim=-1)  # [..., L, L]
+        return bins

model/esm/inverse_folding/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from . import gvp_transformer
+from . import util
+from . import multichain_util

model/esm/inverse_folding/features.py

@@ -0,0 +1,352 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+# 
+# Portions of this file were adapted from the open source code for the following
+# two papers:
+#
+#   Ingraham, J., Garg, V., Barzilay, R., & Jaakkola, T. (2019). Generative
+#   models for graph-based protein design. Advances in Neural Information
+#   Processing Systems, 32.
+#
+#   Jing, B., Eismann, S., Suriana, P., Townshend, R. J. L., & Dror, R. (2020).
+#   Learning from Protein Structure with Geometric Vector Perceptrons. In
+#   International Conference on Learning Representations.
+#
+# MIT License
+# 
+# Copyright (c) 2020 Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael Townshend, Ron Dror
+# 
+# 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.
+# 
+# ================================================================
+# The below license applies to the portions of the code (parts of 
+# src/datasets.py and src/models.py) adapted from Ingraham, et al.
+# ================================================================
+# 
+# MIT License
+# 
+# Copyright (c) 2019 John Ingraham, Vikas Garg, Regina Barzilay, Tommi Jaakkola
+# 
+# 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.
+
+import math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .gvp_utils import flatten_graph
+from .gvp_modules import GVP, LayerNorm
+from .util import normalize, norm, nan_to_num, rbf
+
+
+class GVPInputFeaturizer(nn.Module):
+
+    @staticmethod
+    def get_node_features(coords, coord_mask, with_coord_mask=True):
+        # scalar features
+        node_scalar_features = GVPInputFeaturizer._dihedrals(coords)
+        if with_coord_mask:
+            node_scalar_features = torch.cat([
+                node_scalar_features,
+                coord_mask.float().unsqueeze(-1)
+            ], dim=-1) 
+        # vector features
+        X_ca = coords[:, :, 1]
+        orientations = GVPInputFeaturizer._orientations(X_ca)
+        sidechains = GVPInputFeaturizer._sidechains(coords)
+        node_vector_features = torch.cat([orientations, sidechains.unsqueeze(-2)], dim=-2)
+        return node_scalar_features, node_vector_features
+
+    @staticmethod
+    def _orientations(X):
+        forward = normalize(X[:, 1:] - X[:, :-1])
+        backward = normalize(X[:, :-1] - X[:, 1:])
+        forward = F.pad(forward, [0, 0, 0, 1])
+        backward = F.pad(backward, [0, 0, 1, 0])
+        return torch.cat([forward.unsqueeze(-2), backward.unsqueeze(-2)], -2)
+    
+    @staticmethod
+    def _sidechains(X):
+        n, origin, c = X[:, :, 0], X[:, :, 1], X[:, :, 2]
+        c, n = normalize(c - origin), normalize(n - origin)
+        bisector = normalize(c + n)
+        perp = normalize(torch.cross(c, n, dim=-1))
+        vec = -bisector * math.sqrt(1 / 3) - perp * math.sqrt(2 / 3)
+        return vec 
+
+    @staticmethod
+    def _dihedrals(X, eps=1e-7):
+        X = torch.flatten(X[:, :, :3], 1, 2)
+        bsz = X.shape[0]
+        dX = X[:, 1:] - X[:, :-1]
+        U = normalize(dX, dim=-1)
+        u_2 = U[:, :-2]
+        u_1 = U[:, 1:-1]
+        u_0 = U[:, 2:]
+    
+        # Backbone normals
+        n_2 = normalize(torch.cross(u_2, u_1, dim=-1), dim=-1)
+        n_1 = normalize(torch.cross(u_1, u_0, dim=-1), dim=-1)
+    
+        # Angle between normals
+        cosD = torch.sum(n_2 * n_1, -1)
+        cosD = torch.clamp(cosD, -1 + eps, 1 - eps)
+        D = torch.sign(torch.sum(u_2 * n_1, -1)) * torch.acos(cosD)
+    
+        # This scheme will remove phi[0], psi[-1], omega[-1]
+        D = F.pad(D, [1, 2]) 
+        D = torch.reshape(D, [bsz, -1, 3])
+        # Lift angle representations to the circle
+        D_features = torch.cat([torch.cos(D), torch.sin(D)], -1)
+        return D_features
+
+    @staticmethod
+    def _positional_embeddings(edge_index, 
+                               num_embeddings=None,
+                               num_positional_embeddings=16,
+                               period_range=[2, 1000]):
+        # From https://github.com/jingraham/neurips19-graph-protein-design
+        num_embeddings = num_embeddings or num_positional_embeddings
+        d = edge_index[0] - edge_index[1]
+     
+        frequency = torch.exp(
+            torch.arange(0, num_embeddings, 2, dtype=torch.float32,
+                device=edge_index.device)
+            * -(np.log(10000.0) / num_embeddings)
+        )
+        angles = d.unsqueeze(-1) * frequency
+        E = torch.cat((torch.cos(angles), torch.sin(angles)), -1)
+        return E
+
+    @staticmethod
+    def _dist(X, coord_mask, padding_mask, top_k_neighbors, eps=1e-8):
+        """ Pairwise euclidean distances """
+        bsz, maxlen = X.size(0), X.size(1)
+        coord_mask_2D = torch.unsqueeze(coord_mask,1) * torch.unsqueeze(coord_mask,2)
+        residue_mask = ~padding_mask
+        residue_mask_2D = torch.unsqueeze(residue_mask,1) * torch.unsqueeze(residue_mask,2)
+        dX = torch.unsqueeze(X,1) - torch.unsqueeze(X,2)
+        D = coord_mask_2D * norm(dX, dim=-1)
+    
+        # sorting preference: first those with coords, then among the residues that
+        # exist but are masked use distance in sequence as tie breaker, and then the
+        # residues that came from padding are last
+        seqpos = torch.arange(maxlen, device=X.device)
+        Dseq = torch.abs(seqpos.unsqueeze(1) - seqpos.unsqueeze(0)).repeat(bsz, 1, 1)
+        D_adjust = nan_to_num(D) + (~coord_mask_2D) * (1e8 + Dseq*1e6) + (
+            ~residue_mask_2D) * (1e10)
+    
+        if top_k_neighbors == -1:
+            D_neighbors = D_adjust
+            E_idx = seqpos.repeat(
+                    *D_neighbors.shape[:-1], 1)
+        else:
+            # Identify k nearest neighbors (including self)
+            k = min(top_k_neighbors, X.size(1))
+            D_neighbors, E_idx = torch.topk(D_adjust, k, dim=-1, largest=False)
+    
+        coord_mask_neighbors = (D_neighbors < 5e7)
+        residue_mask_neighbors = (D_neighbors < 5e9)
+        return D_neighbors, E_idx, coord_mask_neighbors, residue_mask_neighbors
+
+
+class Normalize(nn.Module):
+    def __init__(self, features, epsilon=1e-6):
+        super(Normalize, self).__init__()
+        self.gain = nn.Parameter(torch.ones(features))
+        self.bias = nn.Parameter(torch.zeros(features))
+        self.epsilon = epsilon
+
+    def forward(self, x, dim=-1):
+        mu = x.mean(dim, keepdim=True)
+        sigma = torch.sqrt(x.var(dim, keepdim=True) + self.epsilon)
+        gain = self.gain
+        bias = self.bias
+        # Reshape
+        if dim != -1:
+            shape = [1] * len(mu.size())
+            shape[dim] = self.gain.size()[0]
+            gain = gain.view(shape)
+            bias = bias.view(shape)
+        return gain * (x - mu) / (sigma + self.epsilon) + bias
+
+
+class DihedralFeatures(nn.Module):
+    def __init__(self, node_embed_dim):
+        """ Embed dihedral angle features. """
+        super(DihedralFeatures, self).__init__()
+        # 3 dihedral angles; sin and cos of each angle
+        node_in = 6
+        # Normalization and embedding
+        self.node_embedding = nn.Linear(node_in,  node_embed_dim, bias=True)
+        self.norm_nodes = Normalize(node_embed_dim)
+
+    def forward(self, X):
+        """ Featurize coordinates as an attributed graph """
+        V = self._dihedrals(X)
+        V = self.node_embedding(V)
+        V = self.norm_nodes(V)
+        return V
+
+    @staticmethod
+    def _dihedrals(X, eps=1e-7, return_angles=False):
+        # First 3 coordinates are N, CA, C
+        X = X[:,:,:3,:].reshape(X.shape[0], 3*X.shape[1], 3)
+
+        # Shifted slices of unit vectors
+        dX = X[:,1:,:] - X[:,:-1,:]
+        U = F.normalize(dX, dim=-1)
+        u_2 = U[:,:-2,:]
+        u_1 = U[:,1:-1,:]
+        u_0 = U[:,2:,:]
+        # Backbone normals
+        n_2 = F.normalize(torch.cross(u_2, u_1, dim=-1), dim=-1)
+        n_1 = F.normalize(torch.cross(u_1, u_0, dim=-1), dim=-1)
+
+        # Angle between normals
+        cosD = (n_2 * n_1).sum(-1)
+        cosD = torch.clamp(cosD, -1+eps, 1-eps)
+        D = torch.sign((u_2 * n_1).sum(-1)) * torch.acos(cosD)
+
+        # This scheme will remove phi[0], psi[-1], omega[-1]
+        D = F.pad(D, (1,2), 'constant', 0)
+        D = D.view((D.size(0), int(D.size(1)/3), 3))
+        phi, psi, omega = torch.unbind(D,-1)
+
+        if return_angles:
+            return phi, psi, omega
+
+        # Lift angle representations to the circle
+        D_features = torch.cat((torch.cos(D), torch.sin(D)), 2)
+        return D_features
+
+
+class GVPGraphEmbedding(GVPInputFeaturizer):
+
+    def __init__(self, args):
+        super().__init__()
+        self.top_k_neighbors = args.top_k_neighbors
+        self.num_positional_embeddings = 16
+        self.remove_edges_without_coords = True
+        node_input_dim = (7, 3)
+        edge_input_dim = (34, 1)
+        node_hidden_dim = (args.node_hidden_dim_scalar,
+                args.node_hidden_dim_vector)
+        edge_hidden_dim = (args.edge_hidden_dim_scalar,
+                args.edge_hidden_dim_vector)
+        self.embed_node = nn.Sequential(
+            GVP(node_input_dim, node_hidden_dim, activations=(None, None)),
+            LayerNorm(node_hidden_dim, eps=1e-4)
+        )
+        self.embed_edge = nn.Sequential(
+            GVP(edge_input_dim, edge_hidden_dim, activations=(None, None)),
+            LayerNorm(edge_hidden_dim, eps=1e-4)
+        )
+        self.embed_confidence = nn.Linear(16, args.node_hidden_dim_scalar)
+
+    def forward(self, coords, coord_mask, padding_mask, confidence):
+        with torch.no_grad():
+            node_features = self.get_node_features(coords, coord_mask)
+            edge_features, edge_index = self.get_edge_features(
+                coords, coord_mask, padding_mask)
+        node_embeddings_scalar, node_embeddings_vector = self.embed_node(node_features)
+        edge_embeddings = self.embed_edge(edge_features)
+
+        rbf_rep = rbf(confidence, 0., 1.)
+        node_embeddings = (
+            node_embeddings_scalar + self.embed_confidence(rbf_rep),
+            node_embeddings_vector
+        )
+
+        node_embeddings, edge_embeddings, edge_index = flatten_graph(
+            node_embeddings, edge_embeddings, edge_index)
+        return node_embeddings, edge_embeddings, edge_index
+
+    def get_edge_features(self, coords, coord_mask, padding_mask):
+        X_ca = coords[:, :, 1]
+        # Get distances to the top k neighbors
+        E_dist, E_idx, E_coord_mask, E_residue_mask = GVPInputFeaturizer._dist(
+                X_ca, coord_mask, padding_mask, self.top_k_neighbors)
+        # Flatten the graph to be batch size 1 for torch_geometric package 
+        dest = E_idx
+        B, L, k = E_idx.shape[:3]
+        src = torch.arange(L, device=E_idx.device).view([1, L, 1]).expand(B, L, k)
+        # After flattening, [2, B, E]
+        edge_index = torch.stack([src, dest], dim=0).flatten(2, 3)
+        # After flattening, [B, E]
+        E_dist = E_dist.flatten(1, 2)
+        E_coord_mask = E_coord_mask.flatten(1, 2).unsqueeze(-1)
+        E_residue_mask = E_residue_mask.flatten(1, 2)
+        # Calculate relative positional embeddings and distance RBF 
+        pos_embeddings = GVPInputFeaturizer._positional_embeddings(
+            edge_index,
+            num_positional_embeddings=self.num_positional_embeddings,
+        )
+        D_rbf = rbf(E_dist, 0., 20.)
+        # Calculate relative orientation 
+        X_src = X_ca.unsqueeze(2).expand(-1, -1, k, -1).flatten(1, 2)
+        X_dest = torch.gather(
+            X_ca,
+            1,
+            edge_index[1, :, :].unsqueeze(-1).expand([B, L*k, 3])
+        )
+        coord_mask_src = coord_mask.unsqueeze(2).expand(-1, -1, k).flatten(1, 2)
+        coord_mask_dest = torch.gather(
+            coord_mask,
+            1,
+            edge_index[1, :, :].expand([B, L*k])
+        )
+        E_vectors = X_src - X_dest
+        # For the ones without coordinates, substitute in the average vector
+        E_vector_mean = torch.sum(E_vectors * E_coord_mask, dim=1,
+                keepdims=True) / torch.sum(E_coord_mask, dim=1, keepdims=True)
+        E_vectors = E_vectors * E_coord_mask + E_vector_mean * ~(E_coord_mask)
+        # Normalize and remove nans 
+        edge_s = torch.cat([D_rbf, pos_embeddings], dim=-1)
+        edge_v = normalize(E_vectors).unsqueeze(-2)
+        edge_s, edge_v = map(nan_to_num, (edge_s, edge_v))
+        # Also add indications of whether the coordinates are present 
+        edge_s = torch.cat([
+            edge_s,
+            (~coord_mask_src).float().unsqueeze(-1),
+            (~coord_mask_dest).float().unsqueeze(-1),
+        ], dim=-1)
+        edge_index[:, ~E_residue_mask] = -1
+        if self.remove_edges_without_coords:
+            edge_index[:, ~E_coord_mask.squeeze(-1)] = -1
+        return (edge_s, edge_v), edge_index.transpose(0, 1) 

model/esm/inverse_folding/gvp_encoder.py

@@ -0,0 +1,56 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from argparse import Namespace
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .features import GVPGraphEmbedding
+from .gvp_modules import GVPConvLayer, LayerNorm 
+from .gvp_utils import unflatten_graph
+
+
+
+class GVPEncoder(nn.Module):
+
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+        self.embed_graph = GVPGraphEmbedding(args)
+
+        node_hidden_dim = (args.node_hidden_dim_scalar,
+                args.node_hidden_dim_vector)
+        edge_hidden_dim = (args.edge_hidden_dim_scalar,
+                args.edge_hidden_dim_vector)
+        
+        conv_activations = (F.relu, torch.sigmoid)
+        self.encoder_layers = nn.ModuleList(
+                GVPConvLayer(
+                    node_hidden_dim,
+                    edge_hidden_dim,
+                    drop_rate=args.dropout,
+                    vector_gate=True,
+                    attention_heads=0,
+                    n_message=3,
+                    conv_activations=conv_activations,
+                    n_edge_gvps=0,
+                    eps=1e-4,
+                    layernorm=True,
+                ) 
+            for i in range(args.num_encoder_layers)
+        )
+
+    def forward(self, coords, coord_mask, padding_mask, confidence):
+        node_embeddings, edge_embeddings, edge_index = self.embed_graph(
+                coords, coord_mask, padding_mask, confidence)
+        
+        for i, layer in enumerate(self.encoder_layers):
+            node_embeddings, edge_embeddings = layer(node_embeddings,
+                    edge_index, edge_embeddings)
+
+        node_embeddings = unflatten_graph(node_embeddings, coords.shape[0])
+        return node_embeddings

model/esm/inverse_folding/gvp_modules.py

@@ -0,0 +1,475 @@
+# Contents of this file are from the open source code for
+#
+#   Jing, B., Eismann, S., Suriana, P., Townshend, R. J. L., & Dror, R. (2020).
+#   Learning from Protein Structure with Geometric Vector Perceptrons. In
+#   International Conference on Learning Representations.
+#
+# MIT License
+# 
+# Copyright (c) 2020 Bowen Jing, Stephan Eismann, Patricia Suriana, Raphael Townshend, Ron Dror
+# 
+# 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.
+
+import typing as T
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch_geometric.nn import MessagePassing
+
+def tuple_size(tp):
+    return tuple([0 if a is None else a.size() for a in tp])
+
+def tuple_sum(tp1, tp2):
+    s1, v1 = tp1
+    s2, v2 = tp2
+    if v2 is None and v2 is None:
+        return (s1 + s2, None)
+    return (s1 + s2, v1 + v2)
+
+def tuple_cat(*args, dim=-1):
+    '''
+    Concatenates any number of tuples (s, V) elementwise.
+    
+    :param dim: dimension along which to concatenate when viewed
+                as the `dim` index for the scalar-channel tensors.
+                This means that `dim=-1` will be applied as
+                `dim=-2` for the vector-channel tensors.
+    '''
+    dim %= len(args[0][0].shape)
+    s_args, v_args = list(zip(*args))
+    return torch.cat(s_args, dim=dim), torch.cat(v_args, dim=dim)
+
+def tuple_index(x, idx):
+    '''
+    Indexes into a tuple (s, V) along the first dimension.
+    
+    :param idx: any object which can be used to index into a `torch.Tensor`
+    '''
+    return x[0][idx], x[1][idx]
+
+def randn(n, dims, device="cpu"):
+    '''
+    Returns random tuples (s, V) drawn elementwise from a normal distribution.
+    
+    :param n: number of data points
+    :param dims: tuple of dimensions (n_scalar, n_vector)
+    
+    :return: (s, V) with s.shape = (n, n_scalar) and
+             V.shape = (n, n_vector, 3)
+    '''
+    return torch.randn(n, dims[0], device=device), \
+            torch.randn(n, dims[1], 3, device=device)
+
+def _norm_no_nan(x, axis=-1, keepdims=False, eps=1e-8, sqrt=True):
+    '''
+    L2 norm of tensor clamped above a minimum value `eps`.
+    
+    :param sqrt: if `False`, returns the square of the L2 norm
+    '''
+    # clamp is slow
+    # out = torch.clamp(torch.sum(torch.square(x), axis, keepdims), min=eps)
+    out = torch.sum(torch.square(x), axis, keepdims) + eps
+    return torch.sqrt(out) if sqrt else out
+
+def _split(x, nv):
+    '''
+    Splits a merged representation of (s, V) back into a tuple. 
+    Should be used only with `_merge(s, V)` and only if the tuple 
+    representation cannot be used.
+    
+    :param x: the `torch.Tensor` returned from `_merge`
+    :param nv: the number of vector channels in the input to `_merge`
+    '''
+    v = torch.reshape(x[..., -3*nv:], x.shape[:-1] + (nv, 3))
+    s = x[..., :-3*nv]
+    return s, v
+
+def _merge(s, v):
+    '''
+    Merges a tuple (s, V) into a single `torch.Tensor`, where the
+    vector channels are flattened and appended to the scalar channels.
+    Should be used only if the tuple representation cannot be used.
+    Use `_split(x, nv)` to reverse.
+    '''
+    v = torch.reshape(v, v.shape[:-2] + (3*v.shape[-2],))
+    return torch.cat([s, v], -1)
+
+class GVP(nn.Module):
+    '''
+    Geometric Vector Perceptron. See manuscript and README.md
+    for more details.
+    
+    :param in_dims: tuple (n_scalar, n_vector)
+    :param out_dims: tuple (n_scalar, n_vector)
+    :param h_dim: intermediate number of vector channels, optional
+    :param activations: tuple of functions (scalar_act, vector_act)
+    :param tuple_io: whether to keep accepting tuple inputs and outputs when vi
+    or vo = 0
+    '''
+    def __init__(self, in_dims, out_dims, h_dim=None, vector_gate=False,
+                 activations=(F.relu, torch.sigmoid), tuple_io=True,
+                 eps=1e-8):
+        super(GVP, self).__init__()
+        self.si, self.vi = in_dims
+        self.so, self.vo = out_dims
+        self.tuple_io = tuple_io
+        if self.vi: 
+            self.h_dim = h_dim or max(self.vi, self.vo) 
+            self.wh = nn.Linear(self.vi, self.h_dim, bias=False)
+            self.ws = nn.Linear(self.h_dim + self.si, self.so)
+            if self.vo:
+                self.wv = nn.Linear(self.h_dim, self.vo, bias=False)
+                if vector_gate:
+                    self.wg = nn.Linear(self.so, self.vo)
+        else:
+            self.ws = nn.Linear(self.si, self.so)
+        
+        self.vector_gate = vector_gate
+        self.scalar_act, self.vector_act = activations
+        self.eps = eps
+        
+    def forward(self, x):
+        '''
+        :param x: tuple (s, V) of `torch.Tensor`, 
+                  or (if vectors_in is 0), a single `torch.Tensor`
+        :return: tuple (s, V) of `torch.Tensor`,
+                 or (if vectors_out is 0), a single `torch.Tensor`
+        '''
+        if self.vi:
+            s, v = x
+            v = torch.transpose(v, -1, -2)
+            vh = self.wh(v)    
+            vn = _norm_no_nan(vh, axis=-2, eps=self.eps)
+            s = self.ws(torch.cat([s, vn], -1))
+            if self.scalar_act:
+                s = self.scalar_act(s)
+            if self.vo: 
+                v = self.wv(vh) 
+                v = torch.transpose(v, -1, -2)
+                if self.vector_gate:
+                    g = self.wg(s).unsqueeze(-1)
+                else:
+                    g = _norm_no_nan(v, axis=-1, keepdims=True, eps=self.eps)
+                if self.vector_act:
+                    g = self.vector_act(g)
+                    v = v * g
+        else:
+            if self.tuple_io:
+                assert x[1] is None
+                x = x[0]
+            s = self.ws(x)
+            if self.scalar_act:
+                s = self.scalar_act(s)
+            if self.vo:
+                v = torch.zeros(list(s.shape)[:-1] + [self.vo, 3],
+                        device=s.device)
+        
+        if self.vo:
+            return (s, v)
+        elif self.tuple_io:
+            return (s, None)
+        else:
+            return s
+
+
+class _VDropout(nn.Module):
+    '''
+    Vector channel dropout where the elements of each
+    vector channel are dropped together.
+    '''
+    def __init__(self, drop_rate):
+        super(_VDropout, self).__init__()
+        self.drop_rate = drop_rate
+
+    def forward(self, x):
+        '''
+        :param x: `torch.Tensor` corresponding to vector channels
+        '''
+        if x is None:
+            return None
+        device = x.device
+        if not self.training:
+            return x
+        mask = torch.bernoulli(
+            (1 - self.drop_rate) * torch.ones(x.shape[:-1], device=device)
+        ).unsqueeze(-1)
+        x = mask * x / (1 - self.drop_rate)
+        return x
+
+class Dropout(nn.Module):
+    '''
+    Combined dropout for tuples (s, V).
+    Takes tuples (s, V) as input and as output.
+    '''
+    def __init__(self, drop_rate):
+        super(Dropout, self).__init__()
+        self.sdropout = nn.Dropout(drop_rate)
+        self.vdropout = _VDropout(drop_rate)
+
+    def forward(self, x):
+        '''
+        :param x: tuple (s, V) of `torch.Tensor`,
+                  or single `torch.Tensor` 
+                  (will be assumed to be scalar channels)
+        '''
+        if type(x) is torch.Tensor:
+            return self.sdropout(x)
+        s, v = x
+        return self.sdropout(s), self.vdropout(v)
+
+class LayerNorm(nn.Module):
+    '''
+    Combined LayerNorm for tuples (s, V).
+    Takes tuples (s, V) as input and as output.
+    '''
+    def __init__(self, dims, tuple_io=True, eps=1e-8):
+        super(LayerNorm, self).__init__()
+        self.tuple_io = tuple_io
+        self.s, self.v = dims
+        self.scalar_norm = nn.LayerNorm(self.s)
+        self.eps = eps
+        
+    def forward(self, x):
+        '''
+        :param x: tuple (s, V) of `torch.Tensor`,
+                  or single `torch.Tensor` 
+                  (will be assumed to be scalar channels)
+        '''
+        if not self.v:
+            if self.tuple_io:
+                return self.scalar_norm(x[0]), None
+            return self.scalar_norm(x)
+        s, v = x
+        vn = _norm_no_nan(v, axis=-1, keepdims=True, sqrt=False, eps=self.eps)
+        nonzero_mask = (vn > 2 * self.eps)
+        vn = torch.sum(vn * nonzero_mask, dim=-2, keepdim=True
+            ) / (self.eps + torch.sum(nonzero_mask, dim=-2, keepdim=True))
+        vn = torch.sqrt(vn + self.eps)
+        v = nonzero_mask * (v / vn)
+        return self.scalar_norm(s), v
+
+class GVPConv(MessagePassing):
+    '''
+    Graph convolution / message passing with Geometric Vector Perceptrons.
+    Takes in a graph with node and edge embeddings,
+    and returns new node embeddings.
+    
+    This does NOT do residual updates and pointwise feedforward layers
+    ---see `GVPConvLayer`.
+    
+    :param in_dims: input node embedding dimensions (n_scalar, n_vector)
+    :param out_dims: output node embedding dimensions (n_scalar, n_vector)
+    :param edge_dims: input edge embedding dimensions (n_scalar, n_vector)
+    :param n_layers: number of GVPs in the message function
+    :param module_list: preconstructed message function, overrides n_layers
+    :param aggr: should be "add" if some incoming edges are masked, as in
+                 a masked autoregressive decoder architecture
+    '''
+    def __init__(self, in_dims, out_dims, edge_dims, n_layers=3,
+            vector_gate=False, module_list=None, aggr="mean", eps=1e-8,
+            activations=(F.relu, torch.sigmoid)):
+        super(GVPConv, self).__init__(aggr=aggr)
+        self.eps = eps
+        self.si, self.vi = in_dims
+        self.so, self.vo = out_dims
+        self.se, self.ve = edge_dims
+        
+        module_list = module_list or []
+        if not module_list:
+            if n_layers == 1:
+                module_list.append(
+                    GVP((2*self.si + self.se, 2*self.vi + self.ve), 
+                        (self.so, self.vo), activations=(None, None)))
+            else:
+                module_list.append(
+                    GVP((2*self.si + self.se, 2*self.vi + self.ve), out_dims,
+                        vector_gate=vector_gate, activations=activations)
+                )
+                for i in range(n_layers - 2):
+                    module_list.append(GVP(out_dims, out_dims,
+                        vector_gate=vector_gate))
+                module_list.append(GVP(out_dims, out_dims,
+                                       activations=(None, None)))
+        self.message_func = nn.Sequential(*module_list)
+
+    def forward(self, x, edge_index, edge_attr):
+        '''
+        :param x: tuple (s, V) of `torch.Tensor`
+        :param edge_index: array of shape [2, n_edges]
+        :param edge_attr: tuple (s, V) of `torch.Tensor`
+        '''
+        x_s, x_v = x
+        message = self.propagate(edge_index, 
+                    s=x_s, v=x_v.reshape(x_v.shape[0], 3*x_v.shape[1]),
+                    edge_attr=edge_attr)
+        return _split(message, self.vo) 
+
+    def message(self, s_i, v_i, s_j, v_j, edge_attr):
+        v_j = v_j.view(v_j.shape[0], v_j.shape[1]//3, 3)
+        v_i = v_i.view(v_i.shape[0], v_i.shape[1]//3, 3)
+        message = tuple_cat((s_j, v_j), edge_attr, (s_i, v_i))
+        message = self.message_func(message)
+        return _merge(*message)
+
+
+class GVPConvLayer(nn.Module):
+    '''
+    Full graph convolution / message passing layer with 
+    Geometric Vector Perceptrons. Residually updates node embeddings with
+    aggregated incoming messages, applies a pointwise feedforward 
+    network to node embeddings, and returns updated node embeddings.
+    
+    To only compute the aggregated messages, see `GVPConv`.
+    
+    :param node_dims: node embedding dimensions (n_scalar, n_vector)
+    :param edge_dims: input edge embedding dimensions (n_scalar, n_vector)
+    :param n_message: number of GVPs to use in message function
+    :param n_feedforward: number of GVPs to use in feedforward function
+    :param drop_rate: drop probability in all dropout layers
+    :param autoregressive: if `True`, this `GVPConvLayer` will be used
+           with a different set of input node embeddings for messages
+           where src >= dst
+    '''
+    def __init__(self, node_dims, edge_dims, vector_gate=False,
+                 n_message=3, n_feedforward=2, drop_rate=.1,
+                 autoregressive=False, attention_heads=0,
+                 conv_activations=(F.relu, torch.sigmoid),
+                 n_edge_gvps=0, layernorm=True, eps=1e-8):
+        
+        super(GVPConvLayer, self).__init__()
+        if attention_heads == 0:
+            self.conv = GVPConv(
+                    node_dims, node_dims, edge_dims, n_layers=n_message,
+                    vector_gate=vector_gate,
+                    aggr="add" if autoregressive else "mean",
+                    activations=conv_activations, 
+                    eps=eps,
+            )
+        else:
+            raise NotImplementedError
+        if layernorm:
+            self.norm = nn.ModuleList([LayerNorm(node_dims, eps=eps) for _ in range(2)])
+        else:
+            self.norm = nn.ModuleList([nn.Identity() for _ in range(2)])
+        self.dropout = nn.ModuleList([Dropout(drop_rate) for _ in range(2)])
+
+        ff_func = []
+        if n_feedforward == 1:
+            ff_func.append(GVP(node_dims, node_dims, activations=(None, None)))
+        else:
+            hid_dims = 4*node_dims[0], 2*node_dims[1]
+            ff_func.append(GVP(node_dims, hid_dims, vector_gate=vector_gate))
+            for i in range(n_feedforward-2):
+                ff_func.append(GVP(hid_dims, hid_dims, vector_gate=vector_gate))
+            ff_func.append(GVP(hid_dims, node_dims, activations=(None, None)))
+        self.ff_func = nn.Sequential(*ff_func)
+
+        self.edge_message_func = None
+        if n_edge_gvps > 0:
+            si, vi = node_dims
+            se, ve = edge_dims
+            module_list = [
+                GVP((2*si + se, 2*vi + ve), edge_dims, vector_gate=vector_gate)
+            ]
+            for i in range(n_edge_gvps - 2):
+                module_list.append(GVP(edge_dims, edge_dims,
+                    vector_gate=vector_gate))
+            if n_edge_gvps > 1:
+                module_list.append(GVP(edge_dims, edge_dims,
+                    activations=(None, None)))
+            self.edge_message_func = nn.Sequential(*module_list)
+            if layernorm:
+                self.edge_norm = LayerNorm(edge_dims, eps=eps)
+            else:
+                self.edge_norm = nn.Identity()
+            self.edge_dropout = Dropout(drop_rate)
+
+    def forward(self, x, edge_index, edge_attr,
+                autoregressive_x=None, node_mask=None):
+        '''
+        :param x: tuple (s, V) of `torch.Tensor`
+        :param edge_index: array of shape [2, n_edges]
+        :param edge_attr: tuple (s, V) of `torch.Tensor`
+        :param autoregressive_x: tuple (s, V) of `torch.Tensor`. 
+                If not `None`, will be used as srcqq node embeddings
+                for forming messages where src >= dst. The corrent node 
+                embeddings `x` will still be the base of the update and the 
+                pointwise feedforward.
+        :param node_mask: array of type `bool` to index into the first
+                dim of node embeddings (s, V). If not `None`, only
+                these nodes will be updated.
+        '''
+        if self.edge_message_func:
+            src, dst = edge_index
+            if autoregressive_x is None:
+                x_src = x[0][src], x[1][src]
+            else: 
+                mask = (src < dst).unsqueeze(-1)
+                x_src = (
+                    torch.where(mask, x[0][src], autoregressive_x[0][src]),
+                    torch.where(mask.unsqueeze(-1), x[1][src],
+                        autoregressive_x[1][src])
+                )
+            x_dst = x[0][dst], x[1][dst]
+            x_edge = (
+                torch.cat([x_src[0], edge_attr[0], x_dst[0]], dim=-1),
+                torch.cat([x_src[1], edge_attr[1], x_dst[1]], dim=-2)
+            )
+            edge_attr_dh = self.edge_message_func(x_edge)
+            edge_attr = self.edge_norm(tuple_sum(edge_attr,
+                self.edge_dropout(edge_attr_dh)))
+        
+        if autoregressive_x is not None:
+            # Guarding this import here to remove the dependency on torch_scatter, since this isn't used
+            # in ESM-IF1
+            from torch_scatter import scatter_add
+            src, dst = edge_index
+            mask = src < dst
+            edge_index_forward = edge_index[:, mask]
+            edge_index_backward = edge_index[:, ~mask]
+            edge_attr_forward = tuple_index(edge_attr, mask)
+            edge_attr_backward = tuple_index(edge_attr, ~mask)
+            
+            dh = tuple_sum(
+                self.conv(x, edge_index_forward, edge_attr_forward),
+                self.conv(autoregressive_x, edge_index_backward, edge_attr_backward)
+            )
+                        
+            count = scatter_add(torch.ones_like(dst), dst,
+                        dim_size=dh[0].size(0)).clamp(min=1).unsqueeze(-1)
+            
+            dh = dh[0] / count, dh[1] / count.unsqueeze(-1)
+
+        else:
+            dh = self.conv(x, edge_index, edge_attr)
+        
+        if node_mask is not None:
+            x_ = x
+            x, dh = tuple_index(x, node_mask), tuple_index(dh, node_mask)
+            
+        x = self.norm[0](tuple_sum(x, self.dropout[0](dh)))
+        
+        dh = self.ff_func(x)
+        x = self.norm[1](tuple_sum(x, self.dropout[1](dh)))
+        
+        if node_mask is not None:
+            x_[0][node_mask], x_[1][node_mask] = x[0], x[1]
+            x = x_
+
+        return x, edge_attr

model/esm/inverse_folding/gvp_transformer.py

@@ -0,0 +1,140 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+from typing import Any, Dict, List, Optional, Tuple, NamedTuple
+import torch
+from torch import nn
+from torch import Tensor
+import torch.nn.functional as F
+from scipy.spatial import transform
+
+from onescience.datapipes.esm import Alphabet
+
+from .features import DihedralFeatures
+from .gvp_encoder import GVPEncoder
+from .gvp_utils import unflatten_graph
+from .gvp_transformer_encoder import GVPTransformerEncoder
+from .transformer_decoder import TransformerDecoder
+from .util import rotate, CoordBatchConverter 
+
+
+class GVPTransformerModel(nn.Module):
+    """
+    GVP-Transformer inverse folding model.
+
+    Architecture: Geometric GVP-GNN as initial layers, followed by
+    sequence-to-sequence Transformer encoder and decoder.
+    """
+
+    def __init__(self, args, alphabet):
+        super().__init__()
+        encoder_embed_tokens = self.build_embedding(
+            args, alphabet, args.encoder_embed_dim,
+        )
+        decoder_embed_tokens = self.build_embedding(
+            args, alphabet, args.decoder_embed_dim, 
+        )
+        encoder = self.build_encoder(args, alphabet, encoder_embed_tokens)
+        decoder = self.build_decoder(args, alphabet, decoder_embed_tokens)
+        self.args = args
+        self.encoder = encoder
+        self.decoder = decoder
+
+    @classmethod
+    def build_encoder(cls, args, src_dict, embed_tokens):
+        encoder = GVPTransformerEncoder(args, src_dict, embed_tokens)
+        return encoder
+
+    @classmethod
+    def build_decoder(cls, args, tgt_dict, embed_tokens):
+        decoder = TransformerDecoder(
+            args,
+            tgt_dict,
+            embed_tokens,
+        )
+        return decoder
+
+    @classmethod
+    def build_embedding(cls, args, dictionary, embed_dim):
+        num_embeddings = len(dictionary)
+        padding_idx = dictionary.padding_idx
+        emb = nn.Embedding(num_embeddings, embed_dim, padding_idx)
+        nn.init.normal_(emb.weight, mean=0, std=embed_dim ** -0.5)
+        nn.init.constant_(emb.weight[padding_idx], 0)
+        return emb
+
+    def forward(
+        self,
+        coords,
+        padding_mask,
+        confidence,
+        prev_output_tokens,
+        return_all_hiddens: bool = False,
+        features_only: bool = False,
+    ):
+        encoder_out = self.encoder(coords, padding_mask, confidence,
+            return_all_hiddens=return_all_hiddens)
+        logits, extra = self.decoder(
+            prev_output_tokens,
+            encoder_out=encoder_out,
+            features_only=features_only,
+            return_all_hiddens=return_all_hiddens,
+        )
+        return logits, extra
+    
+    def sample(self, coords, partial_seq=None, temperature=1.0, confidence=None, device=None):
+        """
+        Samples sequences based on multinomial sampling (no beam search).
+
+        Args:
+            coords: L x 3 x 3 list representing one backbone
+            partial_seq: Optional, partial sequence with mask tokens if part of
+                the sequence is known
+            temperature: sampling temperature, use low temperature for higher
+                sequence recovery and high temperature for higher diversity
+            confidence: optional length L list of confidence scores for coordinates
+        """
+        L = len(coords)
+        # Convert to batch format
+        batch_converter = CoordBatchConverter(self.decoder.dictionary)
+        batch_coords, confidence, _, _, padding_mask = (
+            batch_converter([(coords, confidence, None)], device=device)
+        )
+        
+        # Start with prepend token
+        mask_idx = self.decoder.dictionary.get_idx('<mask>')
+        sampled_tokens = torch.full((1, 1+L), mask_idx, dtype=int)
+        sampled_tokens[0, 0] = self.decoder.dictionary.get_idx('<cath>')
+        if partial_seq is not None:
+            for i, c in enumerate(partial_seq):
+                sampled_tokens[0, i+1] = self.decoder.dictionary.get_idx(c)
+            
+        # Save incremental states for faster sampling
+        incremental_state = dict()
+        
+        # Run encoder only once
+        encoder_out = self.encoder(batch_coords, padding_mask, confidence)
+        
+        # Make sure all tensors are on the same device if a GPU is present
+        if device:
+            sampled_tokens = sampled_tokens.to(device)
+        
+        # Decode one token at a time
+        for i in range(1, L+1):
+            logits, _ = self.decoder(
+                sampled_tokens[:, :i], 
+                encoder_out,
+                incremental_state=incremental_state,
+            )
+            logits = logits[0].transpose(0, 1)
+            logits /= temperature
+            probs = F.softmax(logits, dim=-1)
+            if sampled_tokens[0, i] == mask_idx:
+                sampled_tokens[:, i] = torch.multinomial(probs, 1).squeeze(-1)
+        sampled_seq = sampled_tokens[0, 1:]
+        
+        # Convert back to string via lookup
+        return ''.join([self.decoder.dictionary.get_tok(a) for a in sampled_seq])

model/esm/inverse_folding/gvp_transformer_encoder.py

@@ -0,0 +1,184 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# Contents of this file were adapted from the open source fairseq repository.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+import math
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from onescience.modules.esm import SinusoidalPositionalEmbedding
+from .features import GVPInputFeaturizer, DihedralFeatures
+from .gvp_encoder import GVPEncoder
+from .transformer_layer import TransformerEncoderLayer
+from .util import nan_to_num, get_rotation_frames, rotate, rbf
+
+
+class GVPTransformerEncoder(nn.Module):
+    """
+    Transformer encoder consisting of *args.encoder.layers* layers. Each layer
+    is a :class:`TransformerEncoderLayer`.
+
+    Args:
+        args (argparse.Namespace): parsed command-line arguments
+        dictionary (~fairseq.data.Dictionary): encoding dictionary
+        embed_tokens (torch.nn.Embedding): input embedding
+    """
+
+    def __init__(self, args, dictionary, embed_tokens):
+        super().__init__()
+        self.args = args
+        self.dictionary = dictionary
+
+        self.dropout_module = nn.Dropout(args.dropout)
+
+        embed_dim = embed_tokens.embedding_dim
+        self.padding_idx = embed_tokens.padding_idx
+
+        self.embed_tokens = embed_tokens
+        self.embed_scale = math.sqrt(embed_dim)
+        self.embed_positions = SinusoidalPositionalEmbedding(
+            embed_dim,
+            self.padding_idx,
+        )
+        self.embed_gvp_input_features = nn.Linear(15, embed_dim)
+        self.embed_confidence = nn.Linear(16, embed_dim)
+        self.embed_dihedrals = DihedralFeatures(embed_dim)
+
+        gvp_args = argparse.Namespace()
+        for k, v in vars(args).items():
+            if k.startswith("gvp_"):
+                setattr(gvp_args, k[4:], v)
+        self.gvp_encoder = GVPEncoder(gvp_args)
+        gvp_out_dim = gvp_args.node_hidden_dim_scalar + (3 *
+                gvp_args.node_hidden_dim_vector)
+        self.embed_gvp_output = nn.Linear(gvp_out_dim, embed_dim)
+
+        self.layers = nn.ModuleList([])
+        self.layers.extend(
+            [self.build_encoder_layer(args) for i in range(args.encoder_layers)]
+        )
+        self.num_layers = len(self.layers)
+        self.layer_norm = nn.LayerNorm(embed_dim)
+
+    def build_encoder_layer(self, args):
+        return TransformerEncoderLayer(args)
+
+    def forward_embedding(self, coords, padding_mask, confidence):
+        """
+        Args:
+            coords: N, CA, C backbone coordinates in shape length x 3 (atoms) x 3 
+            padding_mask: boolean Tensor (true for padding) of shape length
+            confidence: confidence scores between 0 and 1 of shape length
+        """
+        components = dict()
+        coord_mask = torch.all(torch.all(torch.isfinite(coords), dim=-1), dim=-1)
+        coords = nan_to_num(coords)
+        mask_tokens = (
+            padding_mask * self.dictionary.padding_idx + 
+            ~padding_mask * self.dictionary.get_idx("<mask>")
+        )
+        components["tokens"] = self.embed_tokens(mask_tokens) * self.embed_scale
+        components["diherals"] = self.embed_dihedrals(coords)
+
+        # GVP encoder
+        gvp_out_scalars, gvp_out_vectors = self.gvp_encoder(coords,
+                coord_mask, padding_mask, confidence)
+        R = get_rotation_frames(coords)
+        # Rotate to local rotation frame for rotation-invariance
+        gvp_out_features = torch.cat([
+            gvp_out_scalars,
+            rotate(gvp_out_vectors, R.transpose(-2, -1)).flatten(-2, -1),
+        ], dim=-1)
+        components["gvp_out"] = self.embed_gvp_output(gvp_out_features)
+
+        components["confidence"] = self.embed_confidence(
+             rbf(confidence, 0., 1.))
+
+        # In addition to GVP encoder outputs, also directly embed GVP input node
+        # features to the Transformer
+        scalar_features, vector_features = GVPInputFeaturizer.get_node_features(
+            coords, coord_mask, with_coord_mask=False)
+        features = torch.cat([
+            scalar_features,
+            rotate(vector_features, R.transpose(-2, -1)).flatten(-2, -1),
+        ], dim=-1)
+        components["gvp_input_features"] = self.embed_gvp_input_features(features)
+
+        embed = sum(components.values())
+        # for k, v in components.items():
+        #     print(k, torch.mean(v, dim=(0,1)), torch.std(v, dim=(0,1)))
+
+        x = embed
+        x = x + self.embed_positions(mask_tokens)
+        x = self.dropout_module(x)
+        return x, components 
+
+    def forward(
+        self,
+        coords,
+        encoder_padding_mask,
+        confidence,
+        return_all_hiddens: bool = False,
+    ):
+        """
+        Args:
+            coords (Tensor): backbone coordinates
+                shape batch_size x num_residues x num_atoms (3 for N, CA, C) x 3
+            encoder_padding_mask (ByteTensor): the positions of
+                  padding elements of shape `(batch_size x num_residues)`
+            confidence (Tensor): the confidence score of shape (batch_size x
+                num_residues). The value is between 0. and 1. for each residue
+                coordinate, or -1. if no coordinate is given
+            return_all_hiddens (bool, optional): also return all of the
+                intermediate hidden states (default: False).
+
+        Returns:
+            dict:
+                - **encoder_out** (Tensor): the last encoder layer's output of
+                  shape `(num_residues, batch_size, embed_dim)`
+                - **encoder_padding_mask** (ByteTensor): the positions of
+                  padding elements of shape `(batch_size, num_residues)`
+                - **encoder_embedding** (Tensor): the (scaled) embedding lookup
+                  of shape `(batch_size, num_residues, embed_dim)`
+                - **encoder_states** (List[Tensor]): all intermediate
+                  hidden states of shape `(num_residues, batch_size, embed_dim)`.
+                  Only populated if *return_all_hiddens* is True.
+        """
+        x, encoder_embedding = self.forward_embedding(coords,
+                encoder_padding_mask, confidence)
+        # account for padding while computing the representation
+        x = x * (1 - encoder_padding_mask.unsqueeze(-1).type_as(x))
+
+        # B x T x C -> T x B x C
+        x = x.transpose(0, 1)
+
+        encoder_states = []
+
+        if return_all_hiddens:
+            encoder_states.append(x)
+
+        # encoder layers
+        for layer in self.layers:
+            x = layer(
+                x, encoder_padding_mask=encoder_padding_mask
+            )
+            if return_all_hiddens:
+                assert encoder_states is not None
+                encoder_states.append(x)
+
+        if self.layer_norm is not None:
+            x = self.layer_norm(x)
+
+        return {
+            "encoder_out": [x],  # T x B x C
+            "encoder_padding_mask": [encoder_padding_mask],  # B x T
+            "encoder_embedding": [encoder_embedding],  # dictionary
+            "encoder_states": encoder_states,  # List[T x B x C]
+        }

model/esm/inverse_folding/gvp_utils.py

@@ -0,0 +1,68 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+
+
+def flatten_graph(node_embeddings, edge_embeddings, edge_index):
+    """
+    Flattens the graph into a batch size one (with disconnected subgraphs for
+    each example) to be compatible with pytorch-geometric package.
+    Args:
+        node_embeddings: node embeddings in tuple form (scalar, vector)
+                - scalar: shape batch size x nodes x node_embed_dim
+                - vector: shape batch size x nodes x node_embed_dim x 3
+        edge_embeddings: edge embeddings of in tuple form (scalar, vector)
+                - scalar: shape batch size x edges x edge_embed_dim
+                - vector: shape batch size x edges x edge_embed_dim x 3
+        edge_index: shape batch_size x 2 (source node and target node) x edges
+    Returns:
+        node_embeddings: node embeddings in tuple form (scalar, vector)
+                - scalar: shape batch total_nodes x node_embed_dim
+                - vector: shape batch total_nodes x node_embed_dim x 3
+        edge_embeddings: edge embeddings of in tuple form (scalar, vector)
+                - scalar: shape batch total_edges x edge_embed_dim
+                - vector: shape batch total_edges x edge_embed_dim x 3
+        edge_index: shape 2 x total_edges
+    """
+    x_s, x_v = node_embeddings
+    e_s, e_v = edge_embeddings
+    batch_size, N = x_s.shape[0], x_s.shape[1]
+    node_embeddings = (torch.flatten(x_s, 0, 1), torch.flatten(x_v, 0, 1))
+    edge_embeddings = (torch.flatten(e_s, 0, 1), torch.flatten(e_v, 0, 1))
+
+    edge_mask = torch.any(edge_index != -1, dim=1)
+    # Re-number the nodes by adding batch_idx * N to each batch
+    edge_index = edge_index + (torch.arange(batch_size, device=edge_index.device) *
+            N).unsqueeze(-1).unsqueeze(-1)
+    edge_index = edge_index.permute(1, 0, 2).flatten(1, 2)
+    edge_mask = edge_mask.flatten()
+    edge_index = edge_index[:, edge_mask] 
+    edge_embeddings = (
+        edge_embeddings[0][edge_mask, :],
+        edge_embeddings[1][edge_mask, :]
+    )
+    return node_embeddings, edge_embeddings, edge_index 
+
+
+def unflatten_graph(node_embeddings, batch_size):
+    """
+    Unflattens node embeddings.
+    Args:
+        node_embeddings: node embeddings in tuple form (scalar, vector)
+                - scalar: shape batch total_nodes x node_embed_dim
+                - vector: shape batch total_nodes x node_embed_dim x 3
+        batch_size: int
+    Returns:
+        node_embeddings: node embeddings in tuple form (scalar, vector)
+                - scalar: shape batch size x nodes x node_embed_dim
+                - vector: shape batch size x nodes x node_embed_dim x 3
+    """
+    x_s, x_v = node_embeddings
+    x_s = x_s.reshape(batch_size, -1, x_s.shape[1])
+    x_v = x_v.reshape(batch_size, -1, x_v.shape[1], x_v.shape[2])
+    return (x_s, x_v)
+
+

model/esm/inverse_folding/multichain_util.py

@@ -0,0 +1,152 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import biotite.structure
+import numpy as np
+import torch
+from typing import Sequence, Tuple, List
+
+from model.esm.inverse_folding.util import (
+    load_structure,
+    extract_coords_from_structure,
+    load_coords,
+    get_sequence_loss,
+    get_encoder_output,
+)
+
+
+def extract_coords_from_complex(structure: biotite.structure.AtomArray):
+    """
+    Args:
+        structure: biotite AtomArray
+    Returns:
+        Tuple (coords_list, seq_list)
+        - coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+          coordinates representing the backbone of each chain
+        - seqs: Dictionary mapping chain ids to native sequences of each chain
+    """
+    coords = {}
+    seqs = {}
+    all_chains = biotite.structure.get_chains(structure)
+    for chain_id in all_chains:
+        chain = structure[structure.chain_id == chain_id]
+        coords[chain_id], seqs[chain_id] = extract_coords_from_structure(chain)
+    return coords, seqs
+
+
+def load_complex_coords(fpath, chains):
+    """
+    Args:
+        fpath: filepath to either pdb or cif file
+        chains: the chain ids (the order matters for autoregressive model)
+    Returns:
+        Tuple (coords_list, seq_list)
+        - coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+          coordinates representing the backbone of each chain
+        - seqs: Dictionary mapping chain ids to native sequences of each chain
+    """
+    structure = load_structure(fpath, chains)
+    return extract_coords_from_complex(structure)
+
+
+def _concatenate_coords(coords, target_chain_id, padding_length=10):
+    """
+    Args:
+        coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+            coordinates representing the backbone of each chain
+        target_chain_id: The chain id to sample sequences for
+        padding_length: Length of padding between concatenated chains
+    Returns:
+        Tuple (coords, seq)
+            - coords is an L x 3 x 3 array for N, CA, C coordinates, a
+              concatenation of the chains with padding in between
+            - seq is the extracted sequence, with padding tokens inserted
+              between the concatenated chains
+    """
+    pad_coords = np.full((padding_length, 3, 3), np.nan, dtype=np.float32)
+    # For best performance, put the target chain first in concatenation.
+    coords_list = [coords[target_chain_id]]
+    for chain_id in coords:
+        if chain_id == target_chain_id:
+            continue
+        coords_list.append(pad_coords)
+        coords_list.append(coords[chain_id])
+    coords_concatenated = np.concatenate(coords_list, axis=0)
+    return coords_concatenated
+
+
+def sample_sequence_in_complex(model, coords, target_chain_id, temperature=1.,
+        padding_length=10):
+    """
+    Samples sequence for one chain in a complex.
+    Args:
+        model: An instance of the GVPTransformer model
+        coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+            coordinates representing the backbone of each chain
+        target_chain_id: The chain id to sample sequences for
+        padding_length: padding length in between chains
+    Returns:
+        Sampled sequence for the target chain
+    """
+    target_chain_len = coords[target_chain_id].shape[0]
+    all_coords = _concatenate_coords(coords, target_chain_id)
+    device = next(model.parameters()).device
+
+    # Supply padding tokens for other chains to avoid unused sampling for speed
+    padding_pattern = ['<pad>'] * all_coords.shape[0]
+    for i in range(target_chain_len):
+        padding_pattern[i] = '<mask>'
+    sampled = model.sample(all_coords, partial_seq=padding_pattern,
+            temperature=temperature, device=device)
+    sampled = sampled[:target_chain_len]
+    return sampled
+
+
+def score_sequence_in_complex(model, alphabet, coords, target_chain_id,
+        target_seq, padding_length=10):
+    """
+    Scores sequence for one chain in a complex.
+    Args:
+        model: An instance of the GVPTransformer model
+        alphabet: Alphabet for the model
+        coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+            coordinates representing the backbone of each chain
+        target_chain_id: The chain id to sample sequences for
+        target_seq: Target sequence for the target chain for scoring.
+        padding_length: padding length in between chains
+    Returns:
+        Tuple (ll_fullseq, ll_withcoord)
+        - ll_fullseq: Average log-likelihood over the full target chain
+        - ll_withcoord: Average log-likelihood in target chain excluding those
+            residues without coordinates
+    """
+    all_coords = _concatenate_coords(coords, target_chain_id)
+
+    loss, target_padding_mask = get_sequence_loss(model, alphabet, all_coords,
+            target_seq)
+    ll_fullseq = -np.sum(loss * ~target_padding_mask) / np.sum(
+            ~target_padding_mask)
+
+    # Also calculate average when excluding masked portions
+    coord_mask = np.all(np.isfinite(coords[target_chain_id]), axis=(-1, -2))
+    ll_withcoord = -np.sum(loss * coord_mask) / np.sum(coord_mask)
+    return ll_fullseq, ll_withcoord
+
+
+def get_encoder_output_for_complex(model, alphabet, coords, target_chain_id):
+    """
+    Args:
+        model: An instance of the GVPTransformer model
+        alphabet: Alphabet for the model
+        coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+            coordinates representing the backbone of each chain
+        target_chain_id: The chain id to sample sequences for
+    Returns:
+        Dictionary mapping chain id to encoder output for each chain
+    """
+    all_coords = _concatenate_coords(coords, target_chain_id)
+    all_rep = get_encoder_output(model, alphabet, all_coords)
+    target_chain_len = coords[target_chain_id].shape[0]
+    return all_rep[:target_chain_len]

model/esm/inverse_folding/transformer_decoder.py

@@ -0,0 +1,228 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# Contents of this file were adapted from the open source fairseq repository.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from typing import Any, Dict, List, Optional
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from onescience.modules.esm import SinusoidalPositionalEmbedding
+from .transformer_layer import TransformerDecoderLayer
+
+
+def fill_with_neg_inf(t):
+    """FP16-compatible function that fills a tensor with -inf."""
+    return t.float().fill_(float("-inf")).type_as(t)
+
+
+class TransformerDecoder(nn.Module):
+    """
+    Transformer decoder consisting of *args.decoder.layers* layers. Each layer
+    is a :class:`TransformerDecoderLayer`.
+
+    Args:
+        args (argparse.Namespace): parsed command-line arguments
+        dictionary (~fairseq.data.Dictionary): decoding dictionary
+        embed_tokens (torch.nn.Embedding): output embedding
+        no_encoder_attn (bool, optional): whether to attend to encoder outputs
+            (default: False).
+    """
+
+    def __init__(
+        self,
+        args,
+        dictionary,
+        embed_tokens,
+    ):
+        super().__init__()
+        self.args = args
+        self.dictionary = dictionary
+        self._future_mask = torch.empty(0)
+
+        self.dropout_module = nn.Dropout(args.dropout)
+
+        input_embed_dim = embed_tokens.embedding_dim
+        embed_dim = args.decoder_embed_dim
+        self.embed_dim = embed_dim
+
+        self.padding_idx = embed_tokens.padding_idx
+
+        self.embed_tokens = embed_tokens
+        self.embed_scale = math.sqrt(embed_dim)
+
+        self.project_in_dim = (
+            nn.Linear(input_embed_dim, embed_dim, bias=False)
+            if embed_dim != input_embed_dim
+            else None
+        )
+        self.embed_positions = SinusoidalPositionalEmbedding(
+            embed_dim,
+            self.padding_idx,
+        )
+
+        self.layers = nn.ModuleList([])
+        self.layers.extend(
+            [
+                self.build_decoder_layer(args)
+                for _ in range(args.decoder_layers)
+            ]
+        )
+        self.num_layers = len(self.layers)
+        self.layer_norm = nn.LayerNorm(embed_dim)
+
+        self.build_output_projection(args, dictionary)
+
+    def build_output_projection(self, args, dictionary):
+        self.output_projection = nn.Linear(
+            args.decoder_embed_dim, len(dictionary), bias=False
+        )
+        nn.init.normal_(
+            self.output_projection.weight, mean=0, std=args.decoder_embed_dim ** -0.5
+        )
+
+    def build_decoder_layer(self, args):
+        return TransformerDecoderLayer(args)
+
+    def forward(
+        self,
+        prev_output_tokens,
+        encoder_out: Optional[Dict[str, List[Tensor]]] = None,
+        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
+        features_only: bool = False,
+        return_all_hiddens: bool = False,
+    ):
+        """
+        Args:
+            prev_output_tokens (LongTensor): previous decoder outputs of shape
+                `(batch, tgt_len)`, for teacher forcing
+            encoder_out (optional): output from the encoder, used for
+                encoder-side attention, should be of size T x B x C
+            incremental_state (dict): dictionary used for storing state during
+                :ref:`Incremental decoding`
+            features_only (bool, optional): only return features without
+                applying output layer (default: False).
+
+        Returns:
+            tuple:
+                - the decoder's output of shape `(batch, tgt_len, vocab)`
+                - a dictionary with any model-specific outputs
+        """
+
+        x, extra = self.extract_features(
+            prev_output_tokens,
+            encoder_out=encoder_out,
+            incremental_state=incremental_state,
+        )
+
+        if not features_only:
+            x = self.output_layer(x)
+        x = x.transpose(1, 2) # B x T x C -> B x C x T
+        return x, extra
+
+    def extract_features(
+        self,
+        prev_output_tokens,
+        encoder_out: Optional[Dict[str, List[Tensor]]],
+        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
+    ):
+        """
+        Similar to *forward* but only return features.
+
+        Includes several features from "Jointly Learning to Align and
+        Translate with Transformer Models" (Garg et al., EMNLP 2019).
+
+        Returns:
+            tuple:
+                - the decoder's features of shape `(batch, tgt_len, embed_dim)`
+                - a dictionary with any model-specific outputs
+        """
+        bs, slen = prev_output_tokens.size()
+
+        enc: Optional[Tensor] = None
+        padding_mask: Optional[Tensor] = None
+        if encoder_out is not None and len(encoder_out["encoder_out"]) > 0:
+            enc = encoder_out["encoder_out"][0]
+            assert (
+                enc.size()[1] == bs
+            ), f"Expected enc.shape == (t, {bs}, c) got {enc.shape}"
+        if encoder_out is not None and len(encoder_out["encoder_padding_mask"]) > 0:
+            padding_mask = encoder_out["encoder_padding_mask"][0]
+
+        # embed positions
+        positions = self.embed_positions(
+            prev_output_tokens
+        )
+
+        if incremental_state is not None:
+            prev_output_tokens = prev_output_tokens[:, -1:]
+            positions = positions[:, -1:]
+
+        # embed tokens and positions
+        x = self.embed_scale * self.embed_tokens(prev_output_tokens)
+
+        if self.project_in_dim is not None:
+            x = self.project_in_dim(x)
+
+        x += positions
+
+        x = self.dropout_module(x)
+
+        # B x T x C -> T x B x C
+        x = x.transpose(0, 1)
+
+        self_attn_padding_mask: Optional[Tensor] = None
+        if prev_output_tokens.eq(self.padding_idx).any():
+            self_attn_padding_mask = prev_output_tokens.eq(self.padding_idx)
+
+        # decoder layers
+        attn: Optional[Tensor] = None
+        inner_states: List[Optional[Tensor]] = [x]
+        for idx, layer in enumerate(self.layers):
+            if incremental_state is None:
+                self_attn_mask = self.buffered_future_mask(x)
+            else:
+                self_attn_mask = None
+
+            x, layer_attn, _ = layer(
+                x,
+                enc,
+                padding_mask,
+                incremental_state,
+                self_attn_mask=self_attn_mask,
+                self_attn_padding_mask=self_attn_padding_mask,
+                need_attn=False,
+                need_head_weights=False,
+            )
+            inner_states.append(x)
+
+        if self.layer_norm is not None:
+            x = self.layer_norm(x)
+
+        # T x B x C -> B x C x T
+        x = x.transpose(0, 1)
+
+        return x, {"inner_states": inner_states}
+
+    def output_layer(self, features):
+        """Project features to the vocabulary size."""
+        return self.output_projection(features)
+
+    def buffered_future_mask(self, tensor):
+        dim = tensor.size(0)
+        # self._future_mask.device != tensor.device is not working in TorchScript. This is a workaround.
+        if (
+            self._future_mask.size(0) == 0
+            or (not self._future_mask.device == tensor.device)
+            or self._future_mask.size(0) < dim
+        ):
+            self._future_mask = torch.triu(
+                fill_with_neg_inf(torch.zeros([dim, dim])), 1
+            )
+        self._future_mask = self._future_mask.to(tensor)
+        return self._future_mask[:dim, :dim]

model/esm/inverse_folding/transformer_layer.py

@@ -0,0 +1,304 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# Contents of this file were adapted from the open source fairseq repository.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Dict, List, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from onescience.modules.attention import MultiheadAttention
+from torch import Tensor
+
+
+class TransformerEncoderLayer(nn.Module):
+    """Encoder layer block.
+    `layernorm -> dropout -> add residual`
+
+    Args:
+        args (argparse.Namespace): parsed command-line arguments
+    """
+
+    def __init__(self, args):
+        super().__init__()
+        self.args = args
+        self.embed_dim = args.encoder_embed_dim
+        self.self_attn = self.build_self_attention(self.embed_dim, args)
+        self.self_attn_layer_norm = torch.nn.LayerNorm(self.embed_dim)
+        self.dropout_module = nn.Dropout(args.dropout)
+        self.activation_fn = F.relu
+        self.fc1 = self.build_fc1(
+            self.embed_dim,
+            args.encoder_ffn_embed_dim,
+        )
+        self.fc2 = self.build_fc2(
+            args.encoder_ffn_embed_dim,
+            self.embed_dim,
+        )
+
+        self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+
+    def build_fc1(self, input_dim, output_dim):
+        return nn.Linear(input_dim, output_dim)
+
+    def build_fc2(self, input_dim, output_dim):
+        return nn.Linear(input_dim, output_dim)
+
+    def build_self_attention(self, embed_dim, args):
+        return MultiheadAttention(
+            embed_dim,
+            args.encoder_attention_heads,
+            dropout=args.attention_dropout,
+            self_attention=True,
+        )
+
+    def residual_connection(self, x, residual):
+        return residual + x
+
+    def forward(
+        self,
+        x,
+        encoder_padding_mask: Optional[Tensor],
+        attn_mask: Optional[Tensor] = None,
+    ):
+        """
+        Args:
+            x (Tensor): input to the layer of shape `(seq_len, batch, embed_dim)`
+            encoder_padding_mask (ByteTensor): binary ByteTensor of shape
+                `(batch, seq_len)` where padding elements are indicated by ``1``.
+            attn_mask (ByteTensor): binary tensor of shape `(tgt_len, src_len)`,
+                where `tgt_len` is the length of output and `src_len` is the
+                length of input, though here both are equal to `seq_len`.
+                `attn_mask[tgt_i, src_j] = 1` means that when calculating the
+                embedding for `tgt_i`, we exclude (mask out) `src_j`. This is
+                useful for strided self-attention.
+
+        Returns:
+            encoded output of shape `(seq_len, batch, embed_dim)`
+        """
+        # anything in original attn_mask = 1, becomes -1e8
+        # anything in original attn_mask = 0, becomes 0
+        # Note that we cannot use -inf here, because at some edge cases,
+        # the attention weight (before softmax) for some padded element in query
+        # will become -inf, which results in NaN in model parameters
+        if attn_mask is not None:
+            attn_mask = attn_mask.masked_fill(
+                attn_mask.to(torch.bool), -1e8 if x.dtype == torch.float32 else -1e4
+            )
+
+        residual = x
+        x = self.self_attn_layer_norm(x)
+        x, _ = self.self_attn(
+            query=x,
+            key=x,
+            value=x,
+            key_padding_mask=encoder_padding_mask,
+            need_weights=False,
+            attn_mask=attn_mask,
+        )
+        x = self.dropout_module(x)
+        x = self.residual_connection(x, residual)
+
+        residual = x
+        x = self.final_layer_norm(x)
+        x = self.activation_fn(self.fc1(x))
+        x = self.fc2(x)
+        x = self.dropout_module(x)
+        x = self.residual_connection(x, residual)
+        return x
+
+
+class TransformerDecoderLayer(nn.Module):
+    """Decoder layer block.
+    `layernorm -> dropout -> add residual`
+
+    Args:
+        args (argparse.Namespace): parsed command-line arguments
+        no_encoder_attn (bool, optional): whether to attend to encoder outputs
+            (default: False).
+    """
+
+    def __init__(
+        self, args, no_encoder_attn=False, add_bias_kv=False, add_zero_attn=False
+    ):
+        super().__init__()
+        self.embed_dim = args.decoder_embed_dim
+        self.dropout_module = nn.Dropout(args.dropout)
+
+        self.self_attn = self.build_self_attention(
+            self.embed_dim,
+            args,
+            add_bias_kv=add_bias_kv,
+            add_zero_attn=add_zero_attn,
+        )
+        self.nh = self.self_attn.num_heads
+        self.head_dim = self.self_attn.head_dim
+
+        self.activation_fn = F.relu
+
+        self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+
+        if no_encoder_attn:
+            self.encoder_attn = None
+            self.encoder_attn_layer_norm = None
+        else:
+            self.encoder_attn = self.build_encoder_attention(self.embed_dim, args)
+            self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)
+
+        self.ffn_layernorm = (
+            LayerNorm(args.decoder_ffn_embed_dim)
+            if getattr(args, "scale_fc", False)
+            else None
+        )
+        self.w_resid = (
+            nn.Parameter(
+                torch.ones(
+                    self.embed_dim,
+                ),
+                requires_grad=True,
+            )
+            if getattr(args, "scale_resids", False)
+            else None
+        )
+
+        self.fc1 = self.build_fc1(
+            self.embed_dim,
+            args.decoder_ffn_embed_dim,
+        )
+        self.fc2 = self.build_fc2(
+            args.decoder_ffn_embed_dim,
+            self.embed_dim,
+        )
+
+        self.final_layer_norm = nn.LayerNorm(self.embed_dim)
+        self.need_attn = True
+
+    def build_fc1(self, input_dim, output_dim):
+        return nn.Linear(input_dim, output_dim)
+
+    def build_fc2(self, input_dim, output_dim):
+        return nn.Linear(input_dim, output_dim)
+
+    def build_self_attention(
+        self, embed_dim, args, add_bias_kv=False, add_zero_attn=False
+    ):
+        return MultiheadAttention(
+            embed_dim,
+            args.decoder_attention_heads,
+            dropout=args.attention_dropout,
+            add_bias_kv=add_bias_kv,
+            add_zero_attn=add_zero_attn,
+            self_attention=True,
+        )
+
+    def build_encoder_attention(self, embed_dim, args):
+        return MultiheadAttention(
+            embed_dim,
+            args.decoder_attention_heads,
+            kdim=args.encoder_embed_dim,
+            vdim=args.encoder_embed_dim,
+            dropout=args.attention_dropout,
+            encoder_decoder_attention=True,
+        )
+
+    def residual_connection(self, x, residual):
+        return residual + x
+
+    def forward(
+        self,
+        x,
+        encoder_out: Optional[torch.Tensor] = None,
+        encoder_padding_mask: Optional[torch.Tensor] = None,
+        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
+        prev_self_attn_state: Optional[List[torch.Tensor]] = None,
+        prev_attn_state: Optional[List[torch.Tensor]] = None,
+        self_attn_mask: Optional[torch.Tensor] = None,
+        self_attn_padding_mask: Optional[torch.Tensor] = None,
+        need_attn: bool = False,
+        need_head_weights: bool = False,
+    ):
+        """
+        Args:
+            x (Tensor): input to the layer of shape `(seq_len, batch, embed_dim)`
+            encoder_padding_mask (ByteTensor, optional): binary
+                ByteTensor of shape `(batch, src_len)` where padding
+                elements are indicated by ``1``.
+            need_attn (bool, optional): return attention weights
+            need_head_weights (bool, optional): return attention weights
+                for each head (default: return average over heads).
+
+        Returns:
+            encoded output of shape `(seq_len, batch, embed_dim)`
+        """
+        if need_head_weights:
+            need_attn = True
+
+        residual = x
+        x = self.self_attn_layer_norm(x)
+        if prev_self_attn_state is not None:
+            prev_key, prev_value = prev_self_attn_state[:2]
+            saved_state: Dict[str, Optional[Tensor]] = {
+                "prev_key": prev_key,
+                "prev_value": prev_value,
+            }
+            if len(prev_self_attn_state) >= 3:
+                saved_state["prev_key_padding_mask"] = prev_self_attn_state[2]
+            assert incremental_state is not None
+            self.self_attn._set_input_buffer(incremental_state, saved_state)
+        _self_attn_input_buffer = self.self_attn._get_input_buffer(incremental_state)
+        y = x
+
+        x, attn = self.self_attn(
+            query=x,
+            key=y,
+            value=y,
+            key_padding_mask=self_attn_padding_mask,
+            incremental_state=incremental_state,
+            need_weights=False,
+            attn_mask=self_attn_mask,
+        )
+        x = self.dropout_module(x)
+        x = self.residual_connection(x, residual)
+
+        if self.encoder_attn is not None and encoder_out is not None:
+            residual = x
+            x = self.encoder_attn_layer_norm(x)
+            if prev_attn_state is not None:
+                prev_key, prev_value = prev_attn_state[:2]
+                saved_state: Dict[str, Optional[Tensor]] = {
+                    "prev_key": prev_key,
+                    "prev_value": prev_value,
+                }
+                if len(prev_attn_state) >= 3:
+                    saved_state["prev_key_padding_mask"] = prev_attn_state[2]
+                assert incremental_state is not None
+                self.encoder_attn._set_input_buffer(incremental_state, saved_state)
+
+            x, attn = self.encoder_attn(
+                query=x,
+                key=encoder_out,
+                value=encoder_out,
+                key_padding_mask=encoder_padding_mask,
+                incremental_state=incremental_state,
+                static_kv=True,
+                need_weights=need_attn or (not self.training and self.need_attn),
+                need_head_weights=need_head_weights,
+            )
+            x = self.dropout_module(x)
+            x = self.residual_connection(x, residual)
+
+        residual = x
+        x = self.final_layer_norm(x)
+
+        x = self.activation_fn(self.fc1(x))
+        if self.ffn_layernorm is not None:
+            x = self.ffn_layernorm(x)
+        x = self.fc2(x)
+        x = self.dropout_module(x)
+        if self.w_resid is not None:
+            residual = torch.mul(self.w_resid, residual)
+        x = self.residual_connection(x, residual)
+        return x, attn, None

model/esm/inverse_folding/util.py

@@ -0,0 +1,323 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import json
+import math
+
+import biotite.structure
+from biotite.structure.io import pdbx, pdb
+from biotite.structure.residues import get_residues
+from biotite.structure import filter_backbone
+from biotite.structure import get_chains
+from biotite.sequence import ProteinSequence
+import numpy as np
+from scipy.spatial import transform
+from scipy.stats import special_ortho_group
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.data as data
+from typing import Sequence, Tuple, List
+
+from onescience.datapipes.esm import BatchConverter
+
+
+def load_structure(fpath, chain=None):
+    """
+    Args:
+        fpath: filepath to either pdb or cif file
+        chain: the chain id or list of chain ids to load
+    Returns:
+        biotite.structure.AtomArray
+    """
+    if fpath.endswith('cif'):
+        with open(fpath) as fin:
+            pdbxf = pdbx.PDBxFile.read(fin)
+        structure = pdbx.get_structure(pdbxf, model=1)
+    elif fpath.endswith('pdb'):
+        with open(fpath) as fin:
+            pdbf = pdb.PDBFile.read(fin)
+        structure = pdb.get_structure(pdbf, model=1)
+    bbmask = filter_backbone(structure)
+    structure = structure[bbmask]
+    all_chains = get_chains(structure)
+    if len(all_chains) == 0:
+        raise ValueError('No chains found in the input file.')
+    if chain is None:
+        chain_ids = all_chains
+    elif isinstance(chain, list):
+        chain_ids = chain
+    else:
+        chain_ids = [chain] 
+    for chain in chain_ids:
+        if chain not in all_chains:
+            raise ValueError(f'Chain {chain} not found in input file')
+    chain_filter = [a.chain_id in chain_ids for a in structure]
+    structure = structure[chain_filter]
+    return structure
+
+
+def extract_coords_from_structure(structure: biotite.structure.AtomArray):
+    """
+    Args:
+        structure: An instance of biotite AtomArray
+    Returns:
+        Tuple (coords, seq)
+            - coords is an L x 3 x 3 array for N, CA, C coordinates
+            - seq is the extracted sequence
+    """
+    coords = get_atom_coords_residuewise(["N", "CA", "C"], structure)
+    residue_identities = get_residues(structure)[1]
+    seq = ''.join([ProteinSequence.convert_letter_3to1(r) for r in residue_identities])
+    return coords, seq
+
+
+def load_coords(fpath, chain):
+    """
+    Args:
+        fpath: filepath to either pdb or cif file
+        chain: the chain id
+    Returns:
+        Tuple (coords, seq)
+            - coords is an L x 3 x 3 array for N, CA, C coordinates
+            - seq is the extracted sequence
+    """
+    structure = load_structure(fpath, chain)
+    return extract_coords_from_structure(structure)
+
+
+def get_atom_coords_residuewise(atoms: List[str], struct: biotite.structure.AtomArray):
+    """
+    Example for atoms argument: ["N", "CA", "C"]
+    """
+    def filterfn(s, axis=None):
+        filters = np.stack([s.atom_name == name for name in atoms], axis=1)
+        sum = filters.sum(0)
+        if not np.all(sum <= np.ones(filters.shape[1])):
+            raise RuntimeError("structure has multiple atoms with same name")
+        index = filters.argmax(0)
+        coords = s[index].coord
+        coords[sum == 0] = float("nan")
+        return coords
+
+    return biotite.structure.apply_residue_wise(struct, struct, filterfn)
+
+
+def get_sequence_loss(model, alphabet, coords, seq):
+    device = next(model.parameters()).device
+    batch_converter = CoordBatchConverter(alphabet)
+    batch = [(coords, None, seq)]
+    coords, confidence, strs, tokens, padding_mask = batch_converter(
+        batch, device=device)
+
+    prev_output_tokens = tokens[:, :-1].to(device)
+    target = tokens[:, 1:]
+    target_padding_mask = (target == alphabet.padding_idx)
+    logits, _ = model.forward(coords, padding_mask, confidence, prev_output_tokens)
+    loss = F.cross_entropy(logits, target, reduction='none')
+    loss = loss[0].cpu().detach().numpy()
+    target_padding_mask = target_padding_mask[0].cpu().numpy()
+    return loss, target_padding_mask
+
+
+def score_sequence(model, alphabet, coords, seq):
+    loss, target_padding_mask = get_sequence_loss(model, alphabet, coords, seq)
+    ll_fullseq = -np.sum(loss * ~target_padding_mask) / np.sum(~target_padding_mask)
+    # Also calculate average when excluding masked portions
+    coord_mask = np.all(np.isfinite(coords), axis=(-1, -2))
+    ll_withcoord = -np.sum(loss * coord_mask) / np.sum(coord_mask)
+    return ll_fullseq, ll_withcoord
+
+
+def get_encoder_output(model, alphabet, coords):
+    device = next(model.parameters()).device
+    batch_converter = CoordBatchConverter(alphabet)
+    batch = [(coords, None, None)]
+    coords, confidence, strs, tokens, padding_mask = batch_converter(
+        batch, device=device)
+    encoder_out = model.encoder.forward(coords, padding_mask, confidence,
+            return_all_hiddens=False)
+    # remove beginning and end (bos and eos tokens)
+    return encoder_out['encoder_out'][0][1:-1, 0]
+
+
+def rotate(v, R):
+    """
+    Rotates a vector by a rotation matrix.
+    
+    Args:
+        v: 3D vector, tensor of shape (length x batch_size x channels x 3)
+        R: rotation matrix, tensor of shape (length x batch_size x 3 x 3)
+
+    Returns:
+        Rotated version of v by rotation matrix R.
+    """
+    R = R.unsqueeze(-3)
+    v = v.unsqueeze(-1)
+    return torch.sum(v * R, dim=-2)
+
+
+def get_rotation_frames(coords):
+    """
+    Returns a local rotation frame defined by N, CA, C positions.
+
+    Args:
+        coords: coordinates, tensor of shape (batch_size x length x 3 x 3)
+        where the third dimension is in order of N, CA, C
+
+    Returns:
+        Local relative rotation frames in shape (batch_size x length x 3 x 3)
+    """
+    v1 = coords[:, :, 2] - coords[:, :, 1]
+    v2 = coords[:, :, 0] - coords[:, :, 1]
+    e1 = normalize(v1, dim=-1)
+    u2 = v2 - e1 * torch.sum(e1 * v2, dim=-1, keepdim=True)
+    e2 = normalize(u2, dim=-1)
+    e3 = torch.cross(e1, e2, dim=-1)
+    R = torch.stack([e1, e2, e3], dim=-2)
+    return R
+
+
+def nan_to_num(ts, val=0.0):
+    """
+    Replaces nans in tensor with a fixed value.    
+    """
+    val = torch.tensor(val, dtype=ts.dtype, device=ts.device)
+    return torch.where(~torch.isfinite(ts), val, ts)
+
+
+def rbf(values, v_min, v_max, n_bins=16):
+    """
+    Returns RBF encodings in a new dimension at the end.
+    """
+    rbf_centers = torch.linspace(v_min, v_max, n_bins, device=values.device)
+    rbf_centers = rbf_centers.view([1] * len(values.shape) + [-1])
+    rbf_std = (v_max - v_min) / n_bins
+    v_expand = torch.unsqueeze(values, -1)
+    z = (values.unsqueeze(-1) - rbf_centers) / rbf_std
+    return torch.exp(-z ** 2)
+
+
+def norm(tensor, dim, eps=1e-8, keepdim=False):
+    """
+    Returns L2 norm along a dimension.
+    """
+    return torch.sqrt(
+            torch.sum(torch.square(tensor), dim=dim, keepdim=keepdim) + eps)
+
+
+def normalize(tensor, dim=-1):
+    """
+    Normalizes a tensor along a dimension after removing nans.
+    """
+    return nan_to_num(
+        torch.div(tensor, norm(tensor, dim=dim, keepdim=True))
+    )
+
+
+class CoordBatchConverter(BatchConverter):
+    def __call__(self, raw_batch: Sequence[Tuple[Sequence, str]], device=None):
+        """
+        Args:
+            raw_batch: List of tuples (coords, confidence, seq)
+            In each tuple,
+                coords: list of floats, shape L x 3 x 3
+                confidence: list of floats, shape L; or scalar float; or None
+                seq: string of length L
+        Returns:
+            coords: Tensor of shape batch_size x L x 3 x 3
+            confidence: Tensor of shape batch_size x L
+            strs: list of strings
+            tokens: LongTensor of shape batch_size x L
+            padding_mask: ByteTensor of shape batch_size x L
+        """
+        self.alphabet.cls_idx = self.alphabet.get_idx("<cath>") 
+        batch = []
+        for coords, confidence, seq in raw_batch:
+            if confidence is None:
+                confidence = 1.
+            if isinstance(confidence, float) or isinstance(confidence, int):
+                confidence = [float(confidence)] * len(coords)
+            if seq is None:
+                seq = 'X' * len(coords)
+            batch.append(((coords, confidence), seq))
+
+        coords_and_confidence, strs, tokens = super().__call__(batch)
+
+        # pad beginning and end of each protein due to legacy reasons
+        coords = [
+            F.pad(torch.tensor(cd), (0, 0, 0, 0, 1, 1), value=np.inf)
+            for cd, _ in coords_and_confidence
+        ]
+        confidence = [
+            F.pad(torch.tensor(cf), (1, 1), value=-1.)
+            for _, cf in coords_and_confidence
+        ]
+        coords = self.collate_dense_tensors(coords, pad_v=np.nan)
+        confidence = self.collate_dense_tensors(confidence, pad_v=-1.)
+        if device is not None:
+            coords = coords.to(device)
+            confidence = confidence.to(device)
+            tokens = tokens.to(device)
+        padding_mask = torch.isnan(coords[:,:,0,0])
+        coord_mask = torch.isfinite(coords.sum(-2).sum(-1))
+        confidence = confidence * coord_mask + (-1.) * padding_mask
+        return coords, confidence, strs, tokens, padding_mask
+
+    def from_lists(self, coords_list, confidence_list=None, seq_list=None, device=None):
+        """
+        Args:
+            coords_list: list of length batch_size, each item is a list of
+            floats in shape L x 3 x 3 to describe a backbone
+            confidence_list: one of
+                - None, default to highest confidence
+                - list of length batch_size, each item is a scalar
+                - list of length batch_size, each item is a list of floats of
+                    length L to describe the confidence scores for the backbone
+                    with values between 0. and 1.
+            seq_list: either None or a list of strings
+        Returns:
+            coords: Tensor of shape batch_size x L x 3 x 3
+            confidence: Tensor of shape batch_size x L
+            strs: list of strings
+            tokens: LongTensor of shape batch_size x L
+            padding_mask: ByteTensor of shape batch_size x L
+        """
+        batch_size = len(coords_list)
+        if confidence_list is None:
+            confidence_list = [None] * batch_size
+        if seq_list is None:
+            seq_list = [None] * batch_size
+        raw_batch = zip(coords_list, confidence_list, seq_list)
+        return self.__call__(raw_batch, device)
+
+    @staticmethod
+    def collate_dense_tensors(samples, pad_v):
+        """
+        Takes a list of tensors with the following dimensions:
+            [(d_11,       ...,           d_1K),
+             (d_21,       ...,           d_2K),
+             ...,
+             (d_N1,       ...,           d_NK)]
+        and stack + pads them into a single tensor of:
+        (N, max_i=1,N { d_i1 }, ..., max_i=1,N {diK})
+        """
+        if len(samples) == 0:
+            return torch.Tensor()
+        if len(set(x.dim() for x in samples)) != 1:
+            raise RuntimeError(
+                f"Samples has varying dimensions: {[x.dim() for x in samples]}"
+            )
+        (device,) = tuple(set(x.device for x in samples))  # assumes all on same device
+        max_shape = [max(lst) for lst in zip(*[x.shape for x in samples])]
+        result = torch.empty(
+            len(samples), *max_shape, dtype=samples[0].dtype, device=device
+        )
+        result.fill_(pad_v)
+        for i in range(len(samples)):
+            result_i = result[i]
+            t = samples[i]
+            result_i[tuple(slice(0, k) for k in t.shape)] = t
+        return result

model/esm/msa_transformer.py

@@ -0,0 +1,238 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn as nn
+
+from onescience.modules.esm import (
+    AxialTransformerLayer,
+    LearnedPositionalEmbedding,
+    RobertaLMHead,
+    ESM1bLayerNorm,
+    ContactPredictionHead,
+)
+
+from onescience.modules.attention import RowSelfAttention, ColumnSelfAttention
+
+
+
+class MSATransformer(nn.Module):
+    @classmethod
+    def add_args(cls, parser):
+        # fmt: off
+        parser.add_argument(
+            "--num_layers",
+            default=12,
+            type=int,
+            metavar="N",
+            help="number of layers"
+        )
+        parser.add_argument(
+            "--embed_dim",
+            default=768,
+            type=int,
+            metavar="N",
+            help="embedding dimension"
+        )
+        parser.add_argument(
+            "--logit_bias",
+            action="store_true",
+            help="whether to apply bias to logits"
+        )
+        parser.add_argument(
+            "--ffn_embed_dim",
+            default=3072,
+            type=int,
+            metavar="N",
+            help="embedding dimension for FFN",
+        )
+        parser.add_argument(
+            "--attention_heads",
+            default=12,
+            type=int,
+            metavar="N",
+            help="number of attention heads",
+        )
+        parser.add_argument(
+            "--dropout",
+            default=0.1,
+            type=float,
+            help="Dropout to apply."
+        )
+        parser.add_argument(
+            "--attention_dropout",
+            default=0.1,
+            type=float,
+            help="Dropout to apply."
+        )
+        parser.add_argument(
+            "--activation_dropout",
+            default=0.1,
+            type=float,
+            help="Dropout to apply."
+        )
+        parser.add_argument(
+            "--max_tokens_per_msa",
+            default=2 ** 14,
+            type=int,
+            help=(
+                "Used during inference to batch attention computations in a single "
+                "forward pass. This allows increased input sizes with less memory."
+            ),
+        )
+        # fmt: on
+
+    def __init__(self, args, alphabet):
+        super().__init__()
+        self.args = args
+        self.alphabet_size = len(alphabet)
+        self.padding_idx = alphabet.padding_idx
+        self.mask_idx = alphabet.mask_idx
+        self.cls_idx = alphabet.cls_idx
+        self.eos_idx = alphabet.eos_idx
+        self.prepend_bos = alphabet.prepend_bos
+        self.append_eos = alphabet.append_eos
+
+        self.embed_tokens = nn.Embedding(
+            self.alphabet_size, self.args.embed_dim, padding_idx=self.padding_idx
+        )
+
+        if getattr(self.args, "embed_positions_msa", False):
+            emb_dim = getattr(self.args, "embed_positions_msa_dim", self.args.embed_dim)
+            self.msa_position_embedding = nn.Parameter(
+                0.01 * torch.randn(1, 1024, 1, emb_dim),
+                requires_grad=True,
+            )
+        else:
+            self.register_parameter("msa_position_embedding", None)
+
+        self.dropout_module = nn.Dropout(self.args.dropout)
+        self.layers = nn.ModuleList(
+            [
+                AxialTransformerLayer(
+                    self.args.embed_dim,
+                    self.args.ffn_embed_dim,
+                    self.args.attention_heads,
+                    self.args.dropout,
+                    self.args.attention_dropout,
+                    self.args.activation_dropout,
+                    getattr(self.args, "max_tokens_per_msa", self.args.max_tokens),
+                )
+                for _ in range(self.args.layers)
+            ]
+        )
+
+        self.contact_head = ContactPredictionHead(
+            self.args.layers * self.args.attention_heads,
+            self.prepend_bos,
+            self.append_eos,
+            eos_idx=self.eos_idx,
+        )
+        self.embed_positions = LearnedPositionalEmbedding(
+            self.args.max_positions,
+            self.args.embed_dim,
+            self.padding_idx,
+        )
+        self.emb_layer_norm_before = ESM1bLayerNorm(self.args.embed_dim)
+        self.emb_layer_norm_after = ESM1bLayerNorm(self.args.embed_dim)
+        self.lm_head = RobertaLMHead(
+            embed_dim=self.args.embed_dim,
+            output_dim=self.alphabet_size,
+            weight=self.embed_tokens.weight,
+        )
+
+    def forward(self, tokens, repr_layers=[], need_head_weights=False, return_contacts=False):
+        if return_contacts:
+            need_head_weights = True
+
+        assert tokens.ndim == 3
+        batch_size, num_alignments, seqlen = tokens.size()
+        padding_mask = tokens.eq(self.padding_idx)  # B, R, C
+        if not padding_mask.any():
+            padding_mask = None
+
+        x = self.embed_tokens(tokens)
+        x += self.embed_positions(tokens.view(batch_size * num_alignments, seqlen)).view(x.size())
+        if self.msa_position_embedding is not None:
+            if x.size(1) > 1024:
+                raise RuntimeError(
+                    "Using model with MSA position embedding trained on maximum MSA "
+                    f"depth of 1024, but received {x.size(1)} alignments."
+                )
+            x += self.msa_position_embedding[:, :num_alignments]
+
+        x = self.emb_layer_norm_before(x)
+
+        x = self.dropout_module(x)
+
+        if padding_mask is not None:
+            x = x * (1 - padding_mask.unsqueeze(-1).type_as(x))
+
+        repr_layers = set(repr_layers)
+        hidden_representations = {}
+        if 0 in repr_layers:
+            hidden_representations[0] = x
+
+        if need_head_weights:
+            row_attn_weights = []
+            col_attn_weights = []
+
+        # B x R x C x D -> R x C x B x D
+        x = x.permute(1, 2, 0, 3)
+
+        for layer_idx, layer in enumerate(self.layers):
+            x = layer(
+                x,
+                self_attn_padding_mask=padding_mask,
+                need_head_weights=need_head_weights,
+            )
+            if need_head_weights:
+                x, col_attn, row_attn = x
+                # H x C x B x R x R -> B x H x C x R x R
+                col_attn_weights.append(col_attn.permute(2, 0, 1, 3, 4))
+                # H x B x C x C -> B x H x C x C
+                row_attn_weights.append(row_attn.permute(1, 0, 2, 3))
+            if (layer_idx + 1) in repr_layers:
+                hidden_representations[layer_idx + 1] = x.permute(2, 0, 1, 3)
+
+        x = self.emb_layer_norm_after(x)
+        x = x.permute(2, 0, 1, 3)  # R x C x B x D -> B x R x C x D
+
+        # last hidden representation should have layer norm applied
+        if (layer_idx + 1) in repr_layers:
+            hidden_representations[layer_idx + 1] = x
+        x = self.lm_head(x)
+
+        result = {"logits": x, "representations": hidden_representations}
+        if need_head_weights:
+            # col_attentions: B x L x H x C x R x R
+            col_attentions = torch.stack(col_attn_weights, 1)
+            # row_attentions: B x L x H x C x C
+            row_attentions = torch.stack(row_attn_weights, 1)
+            result["col_attentions"] = col_attentions
+            result["row_attentions"] = row_attentions
+            if return_contacts:
+                contacts = self.contact_head(tokens, row_attentions)
+                result["contacts"] = contacts
+
+        return result
+
+    def predict_contacts(self, tokens):
+        return self(tokens, return_contacts=True)["contacts"]
+
+    @property
+    def num_layers(self):
+        return self.args.layers
+
+    def max_tokens_per_msa_(self, value: int) -> None:
+        """The MSA Transformer automatically batches attention computations when
+        gradients are disabled to allow you to pass in larger MSAs at test time than
+        you can fit in GPU memory. By default this occurs when more than 2^14 tokens
+        are passed in the input MSA. You can set this value to infinity to disable
+        this behavior.
+        """
+        for module in self.modules():
+            if isinstance(module, (RowSelfAttention, ColumnSelfAttention)):
+                module.max_tokens_per_msa = value

model/esm/pretrained.py

@@ -0,0 +1,553 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+import re
+import urllib
+import warnings
+from argparse import Namespace
+from pathlib import Path
+
+import torch
+
+import model.esm as esm
+from onescience.datapipes.esm import Alphabet
+from model.esm.esm2 import ESM2
+
+
+def _has_regression_weights(model_name):
+    """Return whether we expect / require regression weights;
+    Right now that is all models except ESM-1v, ESM-IF, and partially trained ESM2 models"""
+    return not ("esm1v" in model_name or "esm_if" in model_name or "270K" in model_name or "500K" in model_name)
+
+
+def load_model_and_alphabet(model_name):
+    if model_name.endswith(".pt"):  # treat as filepath
+        return load_model_and_alphabet_local(model_name)
+    else:
+        return load_model_and_alphabet_hub(model_name)
+
+
+def load_hub_workaround(url):
+    try:
+        data = torch.hub.load_state_dict_from_url(url, progress=False, map_location="cpu")
+    except RuntimeError:
+        # Pytorch version issue - see https://github.com/pytorch/pytorch/issues/43106
+        fn = Path(url).name
+        data = torch.load(
+            f"{torch.hub.get_dir()}/checkpoints/{fn}",
+            map_location="cpu",
+        )
+    except urllib.error.HTTPError as e:
+        raise Exception(f"Could not load {url}, check if you specified a correct model name?")
+    return data
+
+
+def load_regression_hub(model_name):
+    url = f"https://dl.fbaipublicfiles.com/fair-esm/regression/{model_name}-contact-regression.pt"
+    regression_data = load_hub_workaround(url)
+    return regression_data
+
+
+def _download_model_and_regression_data(model_name):
+    url = f"https://dl.fbaipublicfiles.com/fair-esm/models/{model_name}.pt"
+    model_data = load_hub_workaround(url)
+    if _has_regression_weights(model_name):
+        regression_data = load_regression_hub(model_name)
+    else:
+        regression_data = None
+    return model_data, regression_data
+
+
+def load_model_and_alphabet_hub(model_name):
+    model_data, regression_data = _download_model_and_regression_data(model_name)
+    return load_model_and_alphabet_core(model_name, model_data, regression_data)
+
+
+def load_model_and_alphabet_local(model_location):
+    """Load from local path. The regression weights need to be co-located"""
+    model_location = Path(model_location)
+    model_data = torch.load(str(model_location), map_location="cpu")
+    model_name = model_location.stem
+    if _has_regression_weights(model_name):
+        regression_location = str(model_location.with_suffix("")) + "-contact-regression.pt"
+        regression_data = torch.load(regression_location, map_location="cpu")
+    else:
+        regression_data = None
+    return load_model_and_alphabet_core(model_name, model_data, regression_data)
+
+
+def has_emb_layer_norm_before(model_state):
+    """Determine whether layer norm needs to be applied before the encoder"""
+    return any(k.startswith("emb_layer_norm_before") for k, param in model_state.items())
+
+
+def _load_model_and_alphabet_core_v1(model_data):
+    import model.esm as esm  # since esm.inverse_folding is imported below, you actually have to re-import esm here
+
+    alphabet = esm.Alphabet.from_architecture(model_data["args"].arch)
+
+    if model_data["args"].arch == "roberta_large":
+        # upgrade state dict
+        pra = lambda s: "".join(s.split("encoder_")[1:] if "encoder" in s else s)
+        prs1 = lambda s: "".join(s.split("encoder.")[1:] if "encoder" in s else s)
+        prs2 = lambda s: "".join(
+            s.split("sentence_encoder.")[1:] if "sentence_encoder" in s else s
+        )
+        model_args = {pra(arg[0]): arg[1] for arg in vars(model_data["args"]).items()}
+        model_state = {prs1(prs2(arg[0])): arg[1] for arg in model_data["model"].items()}
+        model_state["embed_tokens.weight"][alphabet.mask_idx].zero_()  # For token drop
+        model_args["emb_layer_norm_before"] = has_emb_layer_norm_before(model_state)
+        model_type = esm.ProteinBertModel
+
+    elif model_data["args"].arch == "protein_bert_base":
+
+        # upgrade state dict
+        pra = lambda s: "".join(s.split("decoder_")[1:] if "decoder" in s else s)
+        prs = lambda s: "".join(s.split("decoder.")[1:] if "decoder" in s else s)
+        model_args = {pra(arg[0]): arg[1] for arg in vars(model_data["args"]).items()}
+        model_state = {prs(arg[0]): arg[1] for arg in model_data["model"].items()}
+        model_type = esm.ProteinBertModel
+    elif model_data["args"].arch == "msa_transformer":
+
+        # upgrade state dict
+        pra = lambda s: "".join(s.split("encoder_")[1:] if "encoder" in s else s)
+        prs1 = lambda s: "".join(s.split("encoder.")[1:] if "encoder" in s else s)
+        prs2 = lambda s: "".join(
+            s.split("sentence_encoder.")[1:] if "sentence_encoder" in s else s
+        )
+        prs3 = lambda s: s.replace("row", "column") if "row" in s else s.replace("column", "row")
+        model_args = {pra(arg[0]): arg[1] for arg in vars(model_data["args"]).items()}
+        model_state = {prs1(prs2(prs3(arg[0]))): arg[1] for arg in model_data["model"].items()}
+        if model_args.get("embed_positions_msa", False):
+            emb_dim = model_state["msa_position_embedding"].size(-1)
+            model_args["embed_positions_msa_dim"] = emb_dim  # initial release, bug: emb_dim==1
+
+        model_type = esm.MSATransformer
+
+    elif "invariant_gvp" in model_data["args"].arch:
+        import model.esm.inverse_folding
+
+        model_type = esm.inverse_folding.gvp_transformer.GVPTransformerModel
+        model_args = vars(model_data["args"])  # convert Namespace -> dict
+
+        def update_name(s):
+            # Map the module names in checkpoints trained with internal code to
+            # the updated module names in open source code
+            s = s.replace("W_v", "embed_graph.embed_node")
+            s = s.replace("W_e", "embed_graph.embed_edge")
+            s = s.replace("embed_scores.0", "embed_confidence")
+            s = s.replace("embed_score.", "embed_graph.embed_confidence.")
+            s = s.replace("seq_logits_projection.", "")
+            s = s.replace("embed_ingraham_features", "embed_dihedrals")
+            s = s.replace("embed_gvp_in_local_frame.0", "embed_gvp_output")
+            s = s.replace("embed_features_in_local_frame.0", "embed_gvp_input_features")
+            return s
+
+        model_state = {
+            update_name(sname): svalue
+            for sname, svalue in model_data["model"].items()
+            if "version" not in sname
+        }
+
+    else:
+        raise ValueError("Unknown architecture selected")
+
+    model = model_type(
+        Namespace(**model_args),
+        alphabet,
+    )
+
+    return model, alphabet, model_state
+
+
+def _load_model_and_alphabet_core_v2(model_data):
+    def upgrade_state_dict(state_dict):
+        """Removes prefixes 'model.encoder.sentence_encoder.' and 'model.encoder.'."""
+        prefixes = ["encoder.sentence_encoder.", "encoder."]
+        pattern = re.compile("^" + "|".join(prefixes))
+        state_dict = {pattern.sub("", name): param for name, param in state_dict.items()}
+        return state_dict
+
+    cfg = model_data["cfg"]["model"]
+    state_dict = model_data["model"]
+    state_dict = upgrade_state_dict(state_dict)
+    alphabet = Alphabet.from_architecture("ESM-1b")
+    model = ESM2(
+        num_layers=cfg.encoder_layers,
+        embed_dim=cfg.encoder_embed_dim,
+        attention_heads=cfg.encoder_attention_heads,
+        alphabet=alphabet,
+        token_dropout=cfg.token_dropout,
+    )
+    return model, alphabet, state_dict
+
+
+def load_model_and_alphabet_core(model_name, model_data, regression_data=None):
+    if regression_data is not None:
+        model_data["model"].update(regression_data["model"])
+
+    if model_name.startswith("esm2"):
+        model, alphabet, model_state = _load_model_and_alphabet_core_v2(model_data)
+    else:
+        model, alphabet, model_state = _load_model_and_alphabet_core_v1(model_data)
+
+    expected_keys = set(model.state_dict().keys())
+    found_keys = set(model_state.keys())
+
+    if regression_data is None:
+        expected_missing = {"contact_head.regression.weight", "contact_head.regression.bias"}
+        error_msgs = []
+        missing = (expected_keys - found_keys) - expected_missing
+        if missing:
+            error_msgs.append(f"Missing key(s) in state_dict: {missing}.")
+        unexpected = found_keys - expected_keys
+        if unexpected:
+            error_msgs.append(f"Unexpected key(s) in state_dict: {unexpected}.")
+
+        if error_msgs:
+            raise RuntimeError(
+                "Error(s) in loading state_dict for {}:\n\t{}".format(
+                    model.__class__.__name__, "\n\t".join(error_msgs)
+                )
+            )
+        if expected_missing - found_keys:
+            warnings.warn(
+                "Regression weights not found, predicting contacts will not produce correct results."
+            )
+
+    model.load_state_dict(model_state, strict=regression_data is not None)
+
+    return model, alphabet
+
+
+def esm1_t34_670M_UR50S():
+    """34 layer transformer model with 670M params, trained on Uniref50 Sparse.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1_t34_670M_UR50S")
+
+
+def esm1_t34_670M_UR50D():
+    """34 layer transformer model with 670M params, trained on Uniref50 Dense.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1_t34_670M_UR50D")
+
+
+def esm1_t34_670M_UR100():
+    """34 layer transformer model with 670M params, trained on Uniref100.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1_t34_670M_UR100")
+
+
+def esm1_t12_85M_UR50S():
+    """12 layer transformer model with 85M params, trained on Uniref50 Sparse.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1_t12_85M_UR50S")
+
+
+def esm1_t6_43M_UR50S():
+    """6 layer transformer model with 43M params, trained on Uniref50 Sparse.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1_t6_43M_UR50S")
+
+
+def esm1b_t33_650M_UR50S():
+    """33 layer transformer model with 650M params, trained on Uniref50 Sparse.
+    This is our best performing model, which will be described in a future publication.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1b_t33_650M_UR50S")
+
+
+def esm_msa1_t12_100M_UR50S():
+    warnings.warn(
+        "This model had a minor bug in the positional embeddings, "
+        "please use ESM-MSA-1b: esm.pretrained.esm_msa1b_t12_100M_UR50S()",
+    )
+    return load_model_and_alphabet_hub("esm_msa1_t12_100M_UR50S")
+
+
+def esm_msa1b_t12_100M_UR50S():
+    return load_model_and_alphabet_hub("esm_msa1b_t12_100M_UR50S")
+
+
+def esm1v_t33_650M_UR90S():
+    """33 layer transformer model with 650M params, trained on Uniref90.
+    This is model 1 of a 5 model ensemble.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1v_t33_650M_UR90S_1")
+
+
+def esm1v_t33_650M_UR90S_1():
+    """33 layer transformer model with 650M params, trained on Uniref90.
+    This is model 1 of a 5 model ensemble.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1v_t33_650M_UR90S_1")
+
+
+def esm1v_t33_650M_UR90S_2():
+    """33 layer transformer model with 650M params, trained on Uniref90.
+    This is model 2 of a 5 model ensemble.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1v_t33_650M_UR90S_2")
+
+
+def esm1v_t33_650M_UR90S_3():
+    """33 layer transformer model with 650M params, trained on Uniref90.
+    This is model 3 of a 5 model ensemble.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1v_t33_650M_UR90S_3")
+
+
+def esm1v_t33_650M_UR90S_4():
+    """33 layer transformer model with 650M params, trained on Uniref90.
+    This is model 4 of a 5 model ensemble.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1v_t33_650M_UR90S_4")
+
+
+def esm1v_t33_650M_UR90S_5():
+    """33 layer transformer model with 650M params, trained on Uniref90.
+    This is model 5 of a 5 model ensemble.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm1v_t33_650M_UR90S_5")
+
+
+def esm_if1_gvp4_t16_142M_UR50():
+    """Inverse folding model with 142M params, with 4 GVP-GNN layers, 8
+    Transformer encoder layers, and 8 Transformer decoder layers, trained on
+    CATH structures and 12 million alphafold2 predicted structures from UniRef50
+    sequences.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm_if1_gvp4_t16_142M_UR50")
+
+
+def esm2_t6_8M_UR50D():
+    """6 layer ESM-2 model with 8M params, trained on UniRef50.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm2_t6_8M_UR50D")
+
+
+def esm2_t12_35M_UR50D():
+    """12 layer ESM-2 model with 35M params, trained on UniRef50.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm2_t12_35M_UR50D")
+
+
+def esm2_t30_150M_UR50D():
+    """30 layer ESM-2 model with 150M params, trained on UniRef50.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm2_t30_150M_UR50D")
+
+
+def esm2_t33_650M_UR50D():
+    """33 layer ESM-2 model with 650M params, trained on UniRef50.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm2_t33_650M_UR50D")
+
+
+def esm2_t36_3B_UR50D():
+    """36 layer ESM-2 model with 3B params, trained on UniRef50.
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm2_t36_3B_UR50D")
+
+
+def esm2_t48_15B_UR50D():
+    """48 layer ESM-2 model with 15B params, trained on UniRef50.
+    If you have OOM while loading this model, please refer to README
+    on how to employ FSDP and ZeRO CPU offloading
+
+    Returns a tuple of (Model, Alphabet).
+    """
+    return load_model_and_alphabet_hub("esm2_t48_15B_UR50D")
+
+
+def esmfold_v0():
+    """
+    ESMFold v0 model with 3B ESM-2, 48 folding blocks.
+    This version was used for the paper (Lin et al, 2022). It was trained 
+    on all PDB chains until 2020-05, to ensure temporal holdout with CASP14
+    and the CAMEO validation and test set reported there.
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_v0()
+
+
+def esmfold_v1():
+    """
+    ESMFold v1 model using 3B ESM-2, 48 folding blocks.
+    ESMFold provides fast high accuracy atomic level structure prediction
+    directly from the individual sequence of a protein. ESMFold uses the ESM2
+    protein language model to extract meaningful representations from the
+    protein sequence.
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_v1()
+
+def esmfold_structure_module_only_8M():
+    """
+    ESMFold baseline model using 8M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_8M()
+
+
+def esmfold_structure_module_only_8M_270K():
+    """
+    ESMFold baseline model using 8M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_8M_270K()
+
+
+def esmfold_structure_module_only_35M():
+    """
+    ESMFold baseline model using 35M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_35M()
+
+
+def esmfold_structure_module_only_35M_270K():
+    """
+    ESMFold baseline model using 35M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_35M_270K()
+
+
+def esmfold_structure_module_only_150M():
+    """
+    ESMFold baseline model using 150M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_150M()
+
+
+def esmfold_structure_module_only_150M_270K():
+    """
+    ESMFold baseline model using 150M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_150M_270K()
+
+
+def esmfold_structure_module_only_650M():
+    """
+    ESMFold baseline model using 650M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_650M()
+
+
+def esmfold_structure_module_only_650M_270K():
+    """
+    ESMFold baseline model using 650M ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_650M_270K()
+
+
+def esmfold_structure_module_only_3B():
+    """
+    ESMFold baseline model using 3B ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 500K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_3B()
+
+
+def esmfold_structure_module_only_3B_270K():
+    """
+    ESMFold baseline model using 3B ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_3B_270K()
+
+
+def esmfold_structure_module_only_15B():
+    """
+    ESMFold baseline model using 15B ESM-2, 0 folding blocks.
+    ESM-2 here is trained out to 270K updates.
+    The 15B parameter ESM-2 was not trained out to 500K updates
+    This is a model designed to test the capabilities of the language model
+    when ablated for number of parameters in the language model.
+    See table S1 in (Lin et al, 2022).
+    """
+    import model.esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_15B()

model/esm/version.py

@@ -0,0 +1,6 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+version = "2.0.1"