Upload 480 files

.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
+esm/source/examples/inverse_folding/illustration.png filter=lfs diff=lfs merge=lfs -text
+esm/source/examples/lm-design/utils/ngram_stats/quadgram_seg.p filter=lfs diff=lfs merge=lfs -text
+esm/source/examples/lm-design/utils/ngram_stats/trigram_seg.p filter=lfs diff=lfs merge=lfs -text

esm/source/.flake8

@@ -0,0 +1,10 @@
+[flake8]
+max-line-length = 99
+ignore = E203,W503
+exclude =
+    .git,
+    __pycache__,
+    build,
+    dist,
+    experimental
+    third_party

esm/source/.git-blame-ignore-revs

@@ -0,0 +1,2 @@
+# Migrate code style to Black
+8bc7e948cd9bf0b6d1f2113e221ef548ef663377

esm/source/.github/ISSUE_TEMPLATE/bug.md

@@ -0,0 +1,27 @@
+---
+name: "[Bug Report]"
+about: "Create a bug report. For other questions: see Discussions tab."
+
+---
+
+NOTE: if this is not a bug report, please use the [GitHub Discussions](https://github.com/facebookresearch/esm/discussions) for support questions (How do I do X?), feature requests, ideas, showcasing new applications, etc.
+
+
+**Bug description**
+Please enter a clear and concise description of what the bug is.
+
+**Reproduction steps**
+Enter steps to reproduce the behavior.
+
+**Expected behavior**
+Give a clear and concise description of what you expected to happen.
+
+**Logs**
+Please paste the command line output:
+
+```
+Output goes here
+```
+
+**Additional context**
+Add any other context about the problem here. (like proxy settings, network setup, overall goals, etc.)

esm/source/.gitignore

@@ -0,0 +1,31 @@
+# tensor dumps
+*.pt
+# Compiler Output #
+###################
+*.py[cod]
+*.so
+*.o
+*.exe
+*.class
+
+# Folders #
+###########
+bin/
+build/
+dist/
+local/
+tmp/
+__pycache__/
+*.egg-info/
+.idea/
+.ipynb_checkpoints/
+.vscode/
+esm/dev
+
+# Junk #
+########
+.DS_Store*
+.*.swp
+*.swp
+*.log
+*~

esm/source/CODE_OF_CONDUCT.rst

@@ -0,0 +1,6 @@
+Code of Conduct
+===============
+
+Facebook has adopted a Code of Conduct that we expect project participants to adhere to. Please `read the full text`__ so that you can understand what actions will and will not be tolerated.
+
+__ https://code.facebook.com/codeofconduct

esm/source/CONTRIBUTING.md

@@ -0,0 +1,31 @@
+# Contributing to esm
+We want to make contributing to this project as easy and transparent as
+possible.
+
+## Pull Requests
+We actively welcome your pull requests.
+
+1. Fork the repo and create your branch from `master`.
+2. If you've added code that should be tested, add tests.
+3. If you've changed APIs, update the documentation.
+4. Ensure the test suite passes.
+5. Make sure your code lints.
+6. If you haven't already, complete the Contributor License Agreement ("CLA").
+
+## Contributor License Agreement ("CLA")
+In order to accept your pull request, we need you to submit a CLA. You only need
+to do this once to work on any of Facebook's open source projects.
+
+Complete your CLA here: <https://code.facebook.com/cla>
+
+## Issues
+We use GitHub issues to track public bugs. Please ensure your description is
+clear and has sufficient instructions to be able to reproduce the issue.
+
+Facebook has a [bounty program](https://www.facebook.com/whitehat/) for the safe
+disclosure of security bugs. In those cases, please go through the process
+outlined on that page and do not file a public issue.
+
+## License
+By contributing to icp-block-mdp, you agree that your contributions will be licensed
+under the LICENSE file in the root directory of this source tree.

esm/source/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.

esm/source/README.md

@@ -0,0 +1,795 @@
+# Evolutionary Scale Modeling
+
+[![atlas](https://user-images.githubusercontent.com/3605224/199301187-a9e38b3f-71a7-44be-94f4-db0d66143c53.png)](https://esmatlas.com)
+
+***Update April 2023:*** Code for the two simultaneous preprints on protein design is now released! Code for "Language models generalize beyond natural proteins" is under [examples/lm-design/](examples/lm-design/). Code for "A high-level programming language for generative protein design" is under [examples/protein-programming-language/](examples/protein-programming-language/).
+
+This repository contains code and pre-trained weights for **Transformer protein language models** from the Meta Fundamental AI Research Protein Team (FAIR), including our state-of-the-art [**ESM-2** and **ESMFold**](#esmfold), as well as [**MSA Transformer**](https://www.biorxiv.org/content/10.1101/2021.02.12.430858v1), [**ESM-1v**](#zs_variant) for predicting variant effects and [**ESM-IF1**](#invf) for inverse folding.
+Transformer protein language models were introduced in the [2019 preprint](https://doi.org/10.1101/622803) of the paper ["Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences"](https://doi.org/10.1073/pnas.2016239118).
+ESM-2 outperforms all tested single-sequence protein language models across a range of structure prediction tasks.
+ESMFold harnesses the ESM-2 language model to generate accurate structure predictions end to end directly from the sequence of a protein.
+
+In November 2022, we released `v0` of the [ESM Metagenomic Atlas](https://esmatlas.com), an open atlas of 617 million predicted metagenomic protein structures.
+The Atlas was updated in March 2023 in collaboration with EBI. The new `v2023_02` adds another 150 million predicted structures to the Atlas, as well as pre-computed ESM2 embeddings.
+Bulk download, blog post and the resources provided on the Atlas website are documented [on this README](#atlas).
+
+In December 2022, we released two simultaneous preprints on protein design.
+* "Language models generalize beyond natural proteins" ([PAPER](https://doi.org/10.1101/2022.12.21.521521), [CODE](examples/lm-design/)) uses ESM2 to design de novo proteins. The code and data associated with the preprint can be found [here](examples/lm-design/).
+* "A high-level programming language for generative protein design" ([PAPER](https://doi.org/10.1101/2022.12.21.521526), [CODE](examples/protein-programming-language/)) uses ESMFold to design proteins according to a high-level programming language.
+
+
+
+<details><summary><b>Citation</b></summary>
+For ESM2, ESMFold and ESM Atlas:
+```bibtex
+@article{lin2023evolutionary,
+title = {Evolutionary-scale prediction of atomic-level protein structure with a language model},
+author = {Zeming Lin  and Halil Akin  and Roshan Rao  and Brian Hie  and Zhongkai Zhu  and Wenting Lu  and Nikita Smetanin  and Robert Verkuil  and Ori Kabeli  and Yaniv Shmueli  and Allan dos Santos Costa  and Maryam Fazel-Zarandi  and Tom Sercu  and Salvatore Candido  and Alexander Rives },
+journal = {Science},
+volume = {379},
+number = {6637},
+pages = {1123-1130},
+year = {2023},
+doi = {10.1126/science.ade2574},
+URL = {https://www.science.org/doi/abs/10.1126/science.ade2574},
+note={Earlier versions as preprint: bioRxiv 2022.07.20.500902},
+}
+```
+
+For transformer protein language models:
+```bibtex
+@article{rives2021biological,
+  title={Biological structure and function emerge from scaling unsupervised learning to 250 million protein sequences},
+  author={Rives, Alexander and Meier, Joshua and Sercu, Tom and Goyal, Siddharth and Lin, Zeming and Liu, Jason and Guo, Demi and Ott, Myle and Zitnick, C Lawrence and Ma, Jerry and others},
+  journal={Proceedings of the National Academy of Sciences},
+  volume={118},
+  number={15},
+  pages={e2016239118},
+  year={2021},
+  publisher={National Acad Sciences},
+  note={bioRxiv 10.1101/622803},
+  doi={10.1073/pnas.2016239118},
+  url={https://www.pnas.org/doi/full/10.1073/pnas.2016239118},
+}
+```
+</details>
+
+<details open><summary><b>Table of contents</b></summary>
+
+- [Main models you should use](#main-models)
+- [Usage](#usage)
+  - [Quick Start](#quickstart)
+  - [Getting Started with this repository](#repostart)
+  - [ESMFold Structure Prediction](#esmfold)
+  - [Compute embeddings in bulk from FASTA](#bulk_fasta)
+  - [CPU offloading for inference with large models](#fsdp)
+  - [Zero-shot variant prediction](#zs_variant)
+  - [Inverse folding](#invf)
+- [ESM Metagenomic Atlas](#atlas)
+- [Notebooks](#notebooks)
+- [Available Models and Datasets](#available)
+  - [Pre-trained Models](#available-models)
+  - [ESM Structural Split Dataset](#available-esmssd)
+  - [Pre-training Dataset Split](#available-pretraining-split)
+  - [Comparison to related works](#perf_related)
+- [Citations](#citations)
+- [License](#license)
+</details>
+
+<details><summary><b>What's New</b></summary>
+
+- April 2023: Code for the protein design preprints released under [examples/lm-design/](examples/lm-design/).
+- March 2023: We release an update to the ESM Metagenomic Atlas, `v2023_02`. See [website](https://esmatlas.com/) and [bulk download details](#atlas).
+- December 2022: The Meta Fundamental AI Research Protein Team (FAIR) released two simultaneous preprints on protein design:
+["Language models generalize beyond natural proteins" (Verkuil, Kabeli, et al., 2022)](https://doi.org/10.1101/2022.12.21.521521), and ["A high-level programming language for generative protein design" (Hie, Candido, et al., 2022)](https://doi.org/10.1101/2022.12.21.521521).
+- November 2022: ESM Metagenomic Atlas, a repository of 600M+ metagenomics structures released, see [website](https://esmatlas.com/) and [bulk download details](#atlas)
+- November 2022: ESMFold - new end-to-end structure prediction model released (see [Lin et al. 2022](https://www.science.org/doi/abs/10.1126/science.ade2574))
+- August 2022: ESM-2 - new SOTA Language Models released (see [Lin et al. 2022](https://www.science.org/doi/abs/10.1126/science.ade2574))
+- April 2022: New inverse folding model ESM-IF1 released, trained on CATH and UniRef50 predicted structures.
+- August 2021: Added flexibility to tokenizer to allow for spaces and special tokens (like `<mask>`) in sequence.
+- July 2021: New pre-trained model ESM-1v released, trained on UniRef90 (see [Meier et al. 2021](https://doi.org/10.1101/2021.07.09.450648)).
+- July 2021: New MSA Transformer released, with a minor fix in the row positional embeddings (`ESM-MSA-1b`).
+- Feb 2021: MSA Transformer added (see [Rao et al. 2021](https://www.biorxiv.org/content/10.1101/2021.02.12.430858v1)). Example usage in [notebook](#notebooks).
+- Dec 2020: [Self-Attention Contacts](#notebooks) for all pre-trained models (see [Rao et al. 2020](https://doi.org/10.1101/2020.12.15.422761))
+- Dec 2020: Added new pre-trained model [ESM-1b](#perf_related) (see [Rives et al. 2019](https://doi.org/10.1101/622803) Appendix B)
+- Dec 2020: [ESM Structural Split Dataset](#available-esmssd) (see [Rives et al. 2019](https://doi.org/10.1101/622803) Appendix A.10)
+
+</details>
+
+## Main models you should use <a name="main-models"></a>
+
+| Shorthand | `esm.pretrained.`           | Dataset | Description  |
+|-----------|-----------------------------|---------|--------------|
+| ESM-2    | `esm2_t36_3B_UR50D()` `esm2_t48_15B_UR50D()`       | UR50 (sample UR90)  | SOTA general-purpose protein language model. Can be used to predict structure, function and other protein properties directly from individual sequences. Released with [Lin et al. 2022](https://www.science.org/doi/abs/10.1126/science.ade2574) (Aug 2022 update). |
+| ESMFold   | `esmfold_v1()`         | PDB + UR50 | End-to-end single sequence 3D structure predictor (Nov 2022 update). |
+| ESM-MSA-1b| `esm_msa1b_t12_100M_UR50S()` |  UR50 + MSA  | MSA Transformer language model. Can be used to extract embeddings from an MSA. Enables SOTA inference of structure. Released with [Rao et al. 2021](https://www.biorxiv.org/content/10.1101/2021.02.12.430858v2) (ICML'21 version, June 2021).  |
+| ESM-1v    | `esm1v_t33_650M_UR90S_1()` ... `esm1v_t33_650M_UR90S_5()`| UR90  | Language model specialized for prediction of variant effects. Enables SOTA zero-shot prediction of the functional effects of sequence variations. Same architecture as ESM-1b, but trained on UniRef90. Released with [Meier et al. 2021](https://doi.org/10.1101/2021.07.09.450648). |
+| ESM-IF1  | `esm_if1_gvp4_t16_142M_UR50()` | CATH + UR50 | Inverse folding model. Can be used to design sequences for given structures, or to predict functional effects of sequence variation for given structures. Enables SOTA fixed backbone sequence design. Released with [Hsu et al. 2022](https://doi.org/10.1101/2022.04.10.487779). |
+
+For a complete list of available models, with details and release notes, see [Pre-trained Models](#available-models).
+
+
+## Usage <a name="usage"></a>
+
+### Quick start <a name="quickstart"></a>
+
+An easy way to get started is to load ESM or ESMFold through the [HuggingFace transformers library](https://huggingface.co/docs/transformers/model_doc/esm),
+which has simplified the ESMFold dependencies and provides a standardized API and tools to work with state-of-the-art pretrained models.
+
+Alternatively, [ColabFold](https://colab.research.google.com/github/sokrypton/ColabFold/blob/main/ESMFold.ipynb) has integrated ESMFold so that you can 
+easily run it directly in the browser on a Google Colab instance.
+
+We also provide an API which you can access through curl or on [the ESM Metagenomic Atlas web page](https://esmatlas.com/resources?action=fold).
+```
+curl -X POST --data "KVFGRCELAAAMKRHGLDNYRGYSLGNWVCAAKFESNFNTQATNRNTDGSTDYGILQINSRWWCNDGRTPGSRNLCNIPCSALLSSDITASVNCAKKIVSDGNGMNAWVAWRNRCKGTDVQAWIRGCRL" https://api.esmatlas.com/foldSequence/v1/pdb/
+```
+
+For ESM-MSA-1b, ESM-IF1, or any of the other models you can use the original implementation from our repo directly via the instructions below.
+
+### Getting started with this repo <a name="repostart"></a>
+
+As a prerequisite, you must have PyTorch installed to use this repository.
+
+You can use this one-liner for installation, using the latest release of esm:
+
+```bash
+pip install fair-esm  # latest release, OR:
+pip install git+https://github.com/facebookresearch/esm.git  # bleeding edge, current repo main branch
+```
+
+To use the ESMFold model, make sure you start from an environment with python <= 3.9 and pytorch installed.
+Then add the `[esmfold]` option to your pip install, which will install the dependencies for OpenFold
+automatically. Openfold installation requires `nvcc`.
+
+```bash
+pip install "fair-esm[esmfold]"
+# OpenFold and its remaining dependency
+pip install 'dllogger @ git+https://github.com/NVIDIA/dllogger.git'
+pip install 'openfold @ git+https://github.com/aqlaboratory/openfold.git@4b41059694619831a7db195b7e0988fc4ff3a307'
+```
+
+**NOTE**: If openfold installation fails, please double check that `nvcc` is available and that a cuda-compatable version of PyTorch has been installed.
+
+Alternatively, we provide the `esmfold` conda environment, which can be built via `conda env create -f environment.yml`.
+
+We also support PyTorch Hub, which removes the need to clone and/or install this repository yourself:
+
+```python
+import torch
+model, alphabet = torch.hub.load("facebookresearch/esm:main", "esm2_t33_650M_UR50D")
+```
+
+After pip install, you can load and use a pretrained model as follows:
+
+```python
+import torch
+import esm
+
+# Load ESM-2 model
+model, alphabet = esm.pretrained.esm2_t33_650M_UR50D()
+batch_converter = alphabet.get_batch_converter()
+model.eval()  # disables dropout for deterministic results
+
+# Prepare data (first 2 sequences from ESMStructuralSplitDataset superfamily / 4)
+data = [
+    ("protein1", "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"),
+    ("protein2", "KALTARQQEVFDLIRDHISQTGMPPTRAEIAQRLGFRSPNAAEEHLKALARKGVIEIVSGASRGIRLLQEE"),
+    ("protein2 with mask","KALTARQQEVFDLIRD<mask>ISQTGMPPTRAEIAQRLGFRSPNAAEEHLKALARKGVIEIVSGASRGIRLLQEE"),
+    ("protein3",  "K A <mask> I S Q"),
+]
+batch_labels, batch_strs, batch_tokens = batch_converter(data)
+batch_lens = (batch_tokens != alphabet.padding_idx).sum(1)
+
+# Extract per-residue representations (on CPU)
+with torch.no_grad():
+    results = model(batch_tokens, repr_layers=[33], return_contacts=True)
+token_representations = results["representations"][33]
+
+# Generate per-sequence representations via averaging
+# NOTE: token 0 is always a beginning-of-sequence token, so the first residue is token 1.
+sequence_representations = []
+for i, tokens_len in enumerate(batch_lens):
+    sequence_representations.append(token_representations[i, 1 : tokens_len - 1].mean(0))
+
+# Look at the unsupervised self-attention map contact predictions
+import matplotlib.pyplot as plt
+for (_, seq), tokens_len, attention_contacts in zip(data, batch_lens, results["contacts"]):
+    plt.matshow(attention_contacts[: tokens_len, : tokens_len])
+    plt.title(seq)
+    plt.show()
+```
+
+
+### ESMFold Structure Prediction <a name="esmfold"></a>
+
+After installing with the `[esmfold]` option, you can use the ESMFold structure prediction model as follows:
+
+```python
+import torch
+import esm
+
+model = esm.pretrained.esmfold_v1()
+model = model.eval().cuda()
+
+# Optionally, uncomment to set a chunk size for axial attention. This can help reduce memory.
+# Lower sizes will have lower memory requirements at the cost of increased speed.
+# model.set_chunk_size(128)
+
+sequence = "MKTVRQERLKSIVRILERSKEPVSGAQLAEELSVSRQVIVQDIAYLRSLGYNIVATPRGYVLAGG"
+# Multimer prediction can be done with chains separated by ':'
+
+with torch.no_grad():
+    output = model.infer_pdb(sequence)
+
+with open("result.pdb", "w") as f:
+    f.write(output)
+
+import biotite.structure.io as bsio
+struct = bsio.load_structure("result.pdb", extra_fields=["b_factor"])
+print(struct.b_factor.mean())  # this will be the pLDDT
+# 88.3
+```
+
+
+Besides `esm.pretrained.esmfold_v1()` which is the best performing model we recommend using, we
+also provide `esm.pretrained.esmfold_v0()` which was used for the experiments in
+[Lin et al. 2022](https://www.science.org/doi/abs/10.1126/science.ade2574).
+
+We also provide a command line interface (`esm-fold`) that efficiently predicts structures in bulk from a FASTA file using ESMFold:
+```
+usage: esm-fold [-h] -i FASTA -o PDB [--num-recycles NUM_RECYCLES]
+                [--max-tokens-per-batch MAX_TOKENS_PER_BATCH]
+                [--chunk-size CHUNK_SIZE] [--cpu-only] [--cpu-offload]
+
+optional arguments:
+  -h, --help            show this help message and exit
+  -i FASTA, --fasta FASTA
+                        Path to input FASTA file
+  -o PDB, --pdb PDB     Path to output PDB directory
+  --num-recycles NUM_RECYCLES
+                        Number of recycles to run. Defaults to number used in
+                        training (4).
+  --max-tokens-per-batch MAX_TOKENS_PER_BATCH
+                        Maximum number of tokens per gpu forward-pass. This
+                        will group shorter sequences together for batched
+                        prediction. Lowering this can help with out of memory
+                        issues, if these occur on short sequences.
+  --chunk-size CHUNK_SIZE
+                        Chunks axial attention computation to reduce memory
+                        usage from O(L^2) to O(L). Equivalent to running a for
+                        loop over chunks of of each dimension. Lower values
+                        will result in lower memory usage at the cost of
+                        speed. Recommended values: 128, 64, 32. Default: None.
+  --cpu-only            CPU only
+  --cpu-offload         Enable CPU offloading
+```
+
+The command will make one prediction for every sequence in the fasta file. Multimers can be predicted and should be entered in the fasta file as a single sequence, with chains seprated by a ":" character.
+
+By default, predictions will be batched together so that shorter sequences are predicted simultaneously. This can be disabled by setting `--max-tokens-per-batch=0`. Batching can significantly improve prediction speed on shorter sequences.
+
+The `--cpu-offload` flag can be useful for making predictions on longer sequences. It will attempt to offload some parameters to the CPU RAM, rather than storing on GPU.
+
+Finally, the ablation experiments for LMs of varying sizes [Lin et al. 2022 table S1](https://www.science.org/doi/abs/10.1126/science.ade2574) are released as `esm.pretrained.esmfold_structure_module_only_*()`. We don't recommend using these models for structure prediction.
+
+
+### Compute embeddings in bulk from FASTA <a name="bulk_fasta"></a>
+
+We provide a command line interface (`esm-extract`) that efficiently extracts embeddings in bulk for a FASTA file from the ESM:
+```
+usage: esm-extract [-h] [--toks_per_batch TOKS_PER_BATCH]
+                   [--repr_layers REPR_LAYERS [REPR_LAYERS ...]] --include
+                   {mean,per_tok,bos,contacts}
+                   [{mean,per_tok,bos,contacts} ...]
+                   [--truncation_seq_length TRUNCATION_SEQ_LENGTH]
+                   model_location fasta_file output_dir
+
+Extract per-token representations and model outputs for sequences in a FASTA
+file
+
+positional arguments:
+  model_location        PyTorch model file OR name of pretrained model to
+                        download (see README for models)
+  fasta_file            FASTA file on which to extract representations
+  output_dir            output directory for extracted representations
+
+optional arguments:
+  -h, --help            show this help message and exit
+  --toks_per_batch TOKS_PER_BATCH
+                        maximum batch size
+  --repr_layers REPR_LAYERS [REPR_LAYERS ...]
+                        layers indices from which to extract representations
+                        (0 to num_layers, inclusive)
+  --include {mean,per_tok,bos,contacts} [{mean,per_tok,bos,contacts} ...]
+                        specify which representations to return
+  --truncation_seq_length TRUNCATION_SEQ_LENGTH
+                        truncate sequences longer than the given value
+```
+
+The following commands allow the extraction of the final-layer embedding for a FASTA file from the ESM-2 model:
+
+```bash
+esm-extract esm2_t33_650M_UR50D examples/data/some_proteins.fasta \
+  examples/data/some_proteins_emb_esm2 --repr_layers 0 32 33 --include
+```
+```bash
+python scripts/extract.py esm2_t33_650M_UR50D examples/data/some_proteins.fasta \
+  examples/data/some_proteins_emb_esm2 --repr_layers 0 32 33 --include mean per_tok
+```
+
+A cuda device is optional and will be auto-detected.
+
+Directory `some_proteins_emb_esm2/` now contains one `.pt` file per FASTA sequence; use `torch.load()` to load them.
+`scripts/extract.py` has flags that determine what's included in the `.pt` file:
+* `--repr-layers` (default: final only) selects which layers to include embeddings from.
+* `--include` specifies what embeddings to save. You can use the following:
+  * `per_tok` includes the full sequence, with an embedding per amino acid (seq_len x hidden_dim).
+  * `mean` includes the embeddings averaged over the full sequence, per layer.
+  * `bos` includes the embeddings from the beginning-of-sequence token.
+  (NOTE: Don't use with the pre-trained models - we trained without bos-token supervision)
+
+
+### CPU offloading for inference with large models <a name="fsdp"></a>
+If you want to load very large models like 15B and/or do inference on long sequences on your machine, regular GPU inference may lead to OOM errors.
+We show how to load the model with Fairscale's [Fully Sharded Data Parallel (FSDP)](https://fairscale.readthedocs.io/en/stable/api/nn/fsdp.html) and
+use its CPU offloading feature.
+This allows to do inference of large models on a single GPU.
+Please check out `examples/esm2_infer_fairscale_fsdp_cpu_offloading.py` for more details.
+
+### Zero-shot variant prediction <a name="zs_variant"></a>
+See "[examples/variant-prediction/](examples/variant-prediction/)" for code and pre-trained weights for the ESM-1v models described in
+[Language models enable zero-shot prediction of the effects of mutations on protein function. (Meier et al. 2021)](https://doi.org/10.1101/2021.07.09.450648).
+
+Note that ESM-2 could be used for variant prediction as well, and is expected to have similar performance to ESM-1v.
+
+### Inverse folding <a name="invf"></a>
+See "[examples/inverse_folding/](examples/inverse_folding/)" for detailed user guide. The ESM-IF1 model is described as `GVPTransformer` in [Learning inverse folding from millions of predicted structures. (Hsu et al. 2022)](https://doi.org/10.1101/2022.04.10.487779).
+
+We also provide a colab notebook for the sequence design and sequence scoring functionalities.
+
+[<img src="https://colab.research.google.com/assets/colab-badge.svg">](https://colab.research.google.com/github/facebookresearch/esm/blob/main/examples/inverse_folding/notebook_multichain.ipynb)
+
+The ESM-IF1 inverse folding model is built for predicting protein sequences
+from their backbone atom coordinates. We provide scripts here 1) to sample sequence
+designs for a given structure and 2) to score sequences for a given structure.
+
+Trained with 12M protein structures predicted by AlphaFold2, the ESM-IF1
+model consists of invariant geometric input processing layers followed by a
+sequence-to-sequence transformer, and achieves 51% native sequence recovery on
+structurally held-out backbones with 72% recovery for buried residues.
+The model is also trained with span masking to tolerate missing backbone
+coordinates and therefore can predict sequences for partially masked structures.
+
+#### Sample sequence designs for a given structure
+The environment setup is described in [this subsection of examples/inverse_folding](examples/inverse_folding#recommended-environment).
+
+To sample sequences for a given structure in PDB or mmCIF format, use the
+`sample_sequences.py` script. The input file can have either `.pdb` or
+`.cif` as suffix.
+
+For example, to sample 3 sequence designs for the golgi casein kinase structure
+(PDB [5YH2](https://www.rcsb.org/structure/5yh2); [PDB Molecule of the Month
+from January 2022](https://pdb101.rcsb.org/motm/265)), we can run the following
+command from the esm root directory:
+```bash
+python examples/inverse_folding/sample_sequences.py examples/inverse_folding/data/5YH2.pdb \
+  --chain C --temperature 1 --num-samples 3 --outpath examples/inverse_folding/output/sampled_sequences.fasta
+```
+
+The sampled sequences will be saved in a fasta format to the specified output file.
+
+The temperature parameter controls the sharpness of the probability
+distribution for sequence sampling. Higher sampling temperatures yield more
+diverse sequences but likely with lower native sequence recovery.
+The default sampling temperature is 1. To optimize for native sequence
+recovery, we recommend sampling with low temperature such as 1e-6.
+
+#### Scoring sequences
+To score the conditional log-likelihoods for sequences conditioned on a given
+structure, use the `score_log_likelihoods.py` script.
+
+For example, to score the sequences in `examples/inverse_folding/data/5YH2_mutated_seqs.fasta`
+according to the structure in `examples/inverse_folding/data/5YH2.pdb`, we can run
+the following command from the esm root directory:
+```
+python examples/inverse_folding/score_log_likelihoods.py examples/inverse_folding/data/5YH2.pdb \
+  examples/inverse_folding/data/5YH2_mutated_seqs.fasta --chain C \
+  --outpath examples/inverse_folding/output/5YH2_mutated_seqs_scores.csv
+```
+
+The conditional log-likelihoods are saved in a csv format in the specified output path.
+The output values are the average log-likelihoods averaged over all amino acids in a sequence.
+
+For more information, see "[./examples/inverse_folding/](examples/inverse_folding/)" for detailed user guide.
+
+## ESM Metagenomic Atlas <a name="atlas"></a>
+
+Please visit the [ESM Metagenomic Atlas](https://esmatlas.com/) website, and
+see our [blog post](https://ai.facebook.com/blog/protein-folding-esmfold-metagenomics/) to learn more.
+
+Bulk download instructions available at a seperate README [here](scripts/atlas/README.md).
+
+The Atlas resources include a page to [fold a sequence using ESMFold](https://esmatlas.com/resources?action=fold),
+searching a subset of the ESM Atlas by [structure](https://esmatlas.com/resources?action=search_structure) or 
+[sequence](https://esmatlas.com/resources?action=search_sequence),
+as well as an [API](https://esmatlas.com/about#api) to access those resources programmatically.
+
+Foldseek provides search against the Atlas without the length limitation [here](https://search.foldseek.com/search).
+
+
+## Notebooks <a name="notebooks"></a>
+
+### Inverse folding - predicting or scoring sequences based on backbone structures
+
+[<img src="https://colab.research.google.com/assets/colab-badge.svg">](https://colab.research.google.com/github/facebookresearch/esm/blob/main/examples/inverse_folding/notebook.ipynb)
+
+The ESM-IF1 inverse folding model predicts protein sequences from their backbone atom coordinates, trained with 12M protein structures predicted by AlphaFold2.
+This notetook guide you through examples of sampling sequences, calculating conditional log-likelihoods, and extracting encoder output as structure representation.
+
+### Supervised variant prediction - training a classifier on the embeddings
+
+[<img src="https://colab.research.google.com/assets/colab-badge.svg">](https://colab.research.google.com/github/facebookresearch/esm/blob/main/examples/sup_variant_prediction.ipynb)
+
+
+To help you get started with using the embeddings, this [jupyter notebook tutorial](examples/sup_variant_prediction.ipynb) shows how to train a supervised variant predictor using embeddings from ESM-1.
+You can adopt a similar protocol to train a model for any downstream task, even with limited data.
+First you can obtain the embeddings for ``examples/data/P62593.fasta`` either by [downloading the precomputed](https://dl.fbaipublicfiles.com/fair-esm/examples/P62593_reprs.tar.gz) embeddings
+as instructed in the notebook or by running the following:
+
+```bash
+# Obtain the embeddings
+python scripts/extract.py esm1v_t33_650M_UR90S_1 examples/data/P62593.fasta \
+  examples/data/P62593_emb_esm1v --repr_layers 33 --include mean
+```
+
+Then, follow the remaining instructions in the tutorial. You can also run the tutorial in a [colab notebook](https://colab.research.google.com/github/facebookresearch/esm/blob/main/examples/sup_variant_prediction.ipynb).
+
+**Note, alternatively use [the newer instructions for zero-shot variant prediction](examples/variant-prediction/),
+which predicts mutational effects without any supervised training.**
+
+
+### Unsupervised contact prediction
+[<img src="https://colab.research.google.com/assets/colab-badge.svg">](https://colab.research.google.com/github/facebookresearch/esm/blob/main/examples/contact_prediction.ipynb)
+
+This [jupyter notebook tutorial](examples/contact_prediction.ipynb) demonstrates contact prediction with both the ESM-2 and MSA Transformer (ESM-MSA-1) models.
+Contact prediction is based on a logistic regression over the model's attention maps.
+This methodology is based on our ICLR 2021 paper,
+[Transformer protein language models are unsupervised structure learners. (Rao et al. 2020)](https://doi.org/10.1101/2020.12.15.422761)
+The MSA Transformer (ESM-MSA-1) takes a multiple sequence alignment (MSA) as input, and uses the tied row self-attention maps in the same way.
+See [MSA Transformer. (Rao et al. 2021)](https://www.biorxiv.org/content/10.1101/2021.02.12.430858v1).
+
+To get unsupervised attention-based contacts, call `model.predict_contacts(tokens)` or `model(tokens, return_contacts=True)`.
+
+
+### ESMStructuralSplitDataset and self-attention contact prediction
+[<img src="https://colab.research.google.com/assets/colab-badge.svg">](https://colab.research.google.com/github/facebookresearch/esm/blob/main/examples/esm_structural_dataset.ipynb)
+
+And this [jupyter notebook tutorial](examples/esm_structural_dataset.ipynb) shows how to load and index the `ESMStructuralSplitDataset`,
+and computes the self-attention map unsupervised contact predictions using ESM-2.
+
+
+## Available Models and Datasets <a name="available"></a>
+
+### Pre-trained Models <a name="available-models"></a>
+
+| Shorthand | `esm.pretrained.`           | #layers | #params | Dataset | Embedding Dim |  Model URL (automatically downloaded to `~/.cache/torch/hub/checkpoints`) |
+|-----------|---------------------|---------|-------------|---------|---------------|-----------------------------------------------------------------------|
+| ESM-2     | `esm2_t48_15B_UR50D`         | 48           | 15B         | UR50/D 2021_04                           | 5120 |  https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t48_15B_UR50D.pt          |
+|           | `esm2_t36_3B_UR50D`          | 36           | 3B          | UR50/D 2021_04                           | 2560 |  https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t36_3B_UR50D.pt           |
+|           | `esm2_t33_650M_UR50D`        | 33           | 650M        | UR50/D 2021_04                           | 1280 |  https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t33_650M_UR50D.pt         |
+|           | `esm2_t30_150M_UR50D`        | 30           | 150M        | UR50/D 2021_04                           | 640  |  https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t30_150M_UR50D.pt         |
+|           | `esm2_t12_35M_UR50D`         | 12           | 35M         | UR50/D 2021_04                           | 480  |  https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t12_35M_UR50D.pt          |
+|           | `esm2_t6_8M_UR50D`           | 6            | 8M          | UR50/D 2021_04                           | 320  |  https://dl.fbaipublicfiles.com/fair-esm/models/esm2_t6_8M_UR50D.pt            |
+| ESMFold   | `esmfold_v1`                 | 48 (+36)     | 690M (+3B)  | UR50/D 2021_04                           | -    |  https://dl.fbaipublicfiles.com/fair-esm/models/esmfold_3B_v1.pt               |
+|           | `esmfold_v0`                 | 48 (+36)     | 690M (+3B)  | UR50/D 2021_04                           | -    |  https://dl.fbaipublicfiles.com/fair-esm/models/esmfold_3B_v0.pt               |
+|           | `esmfold_structure_module_only_*`              | 0 (+various) | various     | UR50/D 2021_04                           | -    |  https://dl.fbaipublicfiles.com/fair-esm/models/esmfold_structure_module_only_*                  |
+| ESM-IF1   | `esm_if1_gvp4_t16_142M_UR50` | 20           | 124M        | CATH 4.3 + predicted structures for UR50 | 512  | https://dl.fbaipublicfiles.com/fair-esm/models/esm_if1_gvp4_t16_142M_UR50.pt   |
+| ESM-1v    | `esm1v_t33_650M_UR90S_[1-5]` | 33           | 650M        | UR90/S 2020_03                           | 1280 | https://dl.fbaipublicfiles.com/fair-esm/models/esm1v_t33_650M_UR90S_1.pt       |
+| ESM-MSA-1b| `esm_msa1b_t12_100M_UR50S`   | 12           | 100M        | UR50/S + MSA 2018_03                     | 768  | https://dl.fbaipublicfiles.com/fair-esm/models/esm_msa1b_t12_100M_UR50S.pt     |
+| ESM-MSA-1 | `esm_msa1_t12_100M_UR50S`    | 12           | 100M        | UR50/S + MSA 2018_03                     | 768  | https://dl.fbaipublicfiles.com/fair-esm/models/esm_msa1_t12_100M_UR50S.pt      |
+| ESM-1b    | `esm1b_t33_650M_UR50S`       | 33           | 650M        | UR50/S 2018_03                           | 1280 | https://dl.fbaipublicfiles.com/fair-esm/models/esm1b_t33_650M_UR50S.pt         |
+| ESM-1     | `esm1_t34_670M_UR50S`        | 34           | 670M        | UR50/S 2018_03                           | 1280 |  https://dl.fbaipublicfiles.com/fair-esm/models/esm1_t34_670M_UR50S.pt         |
+|           | `esm1_t34_670M_UR50D`        | 34           | 670M        | UR50/D 2018_03                           | 1280 |  https://dl.fbaipublicfiles.com/fair-esm/models/esm1_t34_670M_UR50D.pt         |
+|           | `esm1_t34_670M_UR100`        | 34           | 670M        | UR100 2018_03                            | 1280 |  https://dl.fbaipublicfiles.com/fair-esm/models/esm1_t34_670M_UR100.pt         |
+|           | `esm1_t12_85M_UR50S`         | 12           | 85M         | UR50/S 2018_03                           | 768  |  https://dl.fbaipublicfiles.com/fair-esm/models/esm1_t12_85M_UR50S.pt          |
+|           | `esm1_t6_43M_UR50S`          | 6            | 43M         | UR50/S 2018_03                           | 768  |  https://dl.fbaipublicfiles.com/fair-esm/models/esm1_t6_43M_UR50S.pt           |
+
+
+Here is a chronological list of the released models and the paper they were introduced in:
+
+| Shorthand  | Release Notes |
+|------------|---------------|
+| ESM-1      | Released with Rives et al. 2019 (Aug 2020 update). |
+| ESM-1b     | Released with Rives et al. 2019 (Dec 2020 update). See Appendix B. |
+| ESM-MSA-1  | Released with Rao et al. 2021 (Preprint v1). |
+| ESM-MSA-1b | Released with Rao et al. 2021 (ICML'21 version, June 2021). |
+| ESM-1v     | Released with Meier et al. 2021. |
+| ESM-IF1    | Released with Hsu et al. 2022. |
+| ESM-2      | Released with Lin et al. 2022. |
+
+### ESM Structural Split Dataset <a name="available-esmssd"></a>
+This is a five-fold cross validation dataset of protein domain structures that can be used to measure generalization of representations
+across different levels of structural dissimilarity.
+The dataset implements structural holdouts at the family, superfamily, and fold
+level. The SCOPe database is used to classify domains. Independently for each level of structural hold-out,
+the domains are split into 5 equal sets, i.e. five sets of folds, superfamilies, or families. This ensures
+that for each of the five partitions, structures having the same classification do not appear in both the
+train and test sets. For a given classification level each structure appears in a test set once, so that
+in the cross validation experiment each of the structures will be evaluated exactly once.
+
+The dataset provides 3d coordinates, distance maps, and secondary structure labels.
+For further details on the construction of the dataset
+see [Rives et al. 2019](https://doi.org/10.1101/622803) Appendix A.10.
+
+This [jupyter notebook tutorial](examples/esm_structural_dataset.ipynb) shows how to load and index the `ESMStructuralSplitDataset`.
+
+`ESMStructuralSplitDataset`, upon initializing, will download `splits` and `pkl`.
+We also provide `msas` for each of the domains. The data can be directly downloaded below.
+
+| Name   | Description                                                                   | URL                                                                   |
+|--------|-------------------------------------------------------------------------------|-----------------------------------------------------------------------|
+| splits | train/valid splits                                                            | https://dl.fbaipublicfiles.com/fair-esm/structural-data/splits.tar.gz |
+| pkl    | pkl objects containing sequence, SSP labels, distance map, and 3d coordinates | https://dl.fbaipublicfiles.com/fair-esm/structural-data/pkl.tar.gz    |
+| msas   | a3m files containing MSA for each domain                                      | https://dl.fbaipublicfiles.com/fair-esm/structural-data/msas.tar.gz   |
+
+### Pre-training Dataset Split  <a name="available-pretraining-split"></a>
+The split files establishing which UniRef50 clusters were used as held-out evaluation set for pre-training
+in [Rives et al. 2019](https://doi.org/10.1101/622803) and [Rao et al. 2021](https://doi.org/10.1101/2021.02.12.430858) can be found here:
+* [UniRef50 IDs of evaluation set](https://dl.fbaipublicfiles.com/fair-esm/pretraining-data/uniref201803_ur50_valid_headers.txt.gz): 3.016 M clusters
+* [UniRef100 IDs of evaluation set](https://dl.fbaipublicfiles.com/fair-esm/pretraining-data/uniref201803_ur100_valid_headers.txt.gz): 13.745 M proteins, expanding the same UniRef50 clusters.
+
+These files only contain only the UniRef50 IDs and UniRef100 IDs corresponding to the [UniRef database, 2018-03 release](https://ftp.uniprot.org/pub/databases/uniprot/previous_releases/release-2018_03/uniref/)
+which is released by the UniProt Consortium under a [Creative Commons Attribution (CC BY 4.0) License](https://www.uniprot.org/help/license).
+
+
+### Comparison to related works <a name="perf_related"></a>
+<!--
+DO NOT EDIT THIS TABLE! This is the source of truth:
+https://docs.google.com/spreadsheets/d/1RPvWF47rIMEr-Jg-SRCoGElHcwCl5d7RyEeSyPgp59A/edit#gid=0
+exported via https://www.tablesgenerator.com/html_tables
+-->
+
+<table class="tg">
+<thead>
+  <tr>
+    <th class="tg-0thz"><span style="font-weight:bold">Task</span></th>
+    <th class="tg-j6zm" colspan="3"><span style="font-weight:bold">Unsupervised contact prediction</span></th>
+    <th class="tg-j6zm" colspan="2"><span style="font-weight:bold">Structure Prediction</span></th>
+  </tr>
+</thead>
+<tbody>
+  <tr>
+    <td class="tg-j6zm"><span style="font-weight:bold">Test set</span></td>
+    <td class="tg-j6zm"><span style="font-weight:bold">Large valid</span></td>
+    <td class="tg-j6zm"><span style="font-weight:bold">CASP14</span></td>
+    <td class="tg-j6zm"><span style="font-weight:bold">CAMEO (Apr-Jun 2022)</span></td>
+    <td class="tg-j6zm"><span style="font-weight:bold">CASP14</span></td>
+    <td class="tg-j6zm"><span style="font-weight:bold">CAMEO (Apr-Jun 2022)</span></td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">Gremlin (Potts)</td>
+    <td class="tg-7zrl">39.3</td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">TAPE</td>
+    <td class="tg-7zrl">11.2</td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">ProtBert-BFD</td>
+    <td class="tg-7zrl">34.1</td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">Prot-T5-XL-BFD</td>
+    <td class="tg-7zrl">35.6</td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-2b7s">46.1</td>
+    <td class="tg-2b7s">62.6</td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">Prot-T5-XL-Ur50 (3B)</td>
+    <td class="tg-7zrl">47.9</td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-2b7s">49.8</td>
+    <td class="tg-2b7s">69.4</td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">ESM-1</td>
+    <td class="tg-7zrl">33.7</td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">ESM-1b</td>
+    <td class="tg-7zrl">41.1</td>
+    <td class="tg-7zrl">24.4</td>
+    <td class="tg-7zrl">39</td>
+    <td class="tg-2b7s">41.6</td>
+    <td class="tg-2b7s">64.5</td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">ESM-1v</td>
+    <td class="tg-7zrl">35.3</td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">ESM-MSA-1b</td>
+    <td class="tg-7zrl">57.4</td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+    <td class="tg-7zrl"></td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">ESM-2 (8M)</td>
+    <td class="tg-7zrl">15.9</td>
+    <td class="tg-7zrl">9.8</td>
+    <td class="tg-7zrl">15.7</td>
+    <td class="tg-2b7s">36.7</td>
+    <td class="tg-2b7s">48.1</td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">ESM-2 (35M)</td>
+    <td class="tg-7zrl">28.8</td>
+    <td class="tg-7zrl">16.4</td>
+    <td class="tg-7zrl">28.4</td>
+    <td class="tg-2b7s">41.4</td>
+    <td class="tg-2b7s">56.4</td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">ESM-2 (150M)</td>
+    <td class="tg-7zrl">42.2</td>
+    <td class="tg-7zrl">26.8</td>
+    <td class="tg-7zrl">40.1</td>
+    <td class="tg-2b7s">49.0</td>
+    <td class="tg-2b7s">64.9</td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">ESM-2 (700M)</td>
+    <td class="tg-7zrl">50.1</td>
+    <td class="tg-7zrl">32.5</td>
+    <td class="tg-7zrl">47.6</td>
+    <td class="tg-2b7s">51.3</td>
+    <td class="tg-2b7s">70.1</td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">ESM-2 (3B)</td>
+    <td class="tg-7zrl">52.7</td>
+    <td class="tg-7zrl">34.0</td>
+    <td class="tg-7zrl">49.9</td>
+    <td class="tg-2b7s">52.5</td>
+    <td class="tg-2b7s">71.8</td>
+  </tr>
+  <tr>
+    <td class="tg-7zrl">ESM-2 (15B)</td>
+    <td class="tg-7zrl">54.5</td>
+    <td class="tg-7zrl">37.0</td>
+    <td class="tg-7zrl">51.7</td>
+    <td class="tg-2b7s">55.4</td>
+    <td class="tg-2b7s">72.1</td>
+  </tr>
+</tbody>
+</table>
+
+Comparison to related protein language models on structure prediction tasks.
+
+* All contact numbers are the top-L,LR precision metric, where long range means sequence separation of at least 24 residues
+* For unsupervised contact prediction, a sparse linear combination of the attention heads is used to directly predict protein contacts,
+fitted with logistic regression on 20 structures.
+For more details on the method, see [Rao et al. 2020](https://doi.org/10.1101/2020.12.15.422761).
+* For structure prediction, an AlphaFold2 structure module is trained directly from the frozen language model embeddings.
+For more details on the method, see [Lin et al. 2022](https://www.science.org/doi/abs/10.1126/science.ade2574).
+* Direct coupling analysis methods (Gremlin, mfDCA, Psicov) and ESM-MSA-1 use the [trRosetta MSAs](https://yanglab.nankai.edu.cn/trRosetta/benchmark/), while other methods predict from single sequence.
+
+
+## Citations <a name="citations"></a>
+
+If you find the models useful in your research, we ask that you cite the relevant paper:
+
+```bibtex
+@article{rives2019biological,
+  author={Rives, Alexander and Meier, Joshua and Sercu, Tom and Goyal, Siddharth and Lin, Zeming and Liu, Jason and Guo, Demi and Ott, Myle and Zitnick, C. Lawrence and Ma, Jerry and Fergus, Rob},
+  title={Biological Structure and Function Emerge from Scaling Unsupervised Learning to 250 Million Protein Sequences},
+  year={2019},
+  doi={10.1101/622803},
+  url={https://www.biorxiv.org/content/10.1101/622803v4},
+  journal={PNAS}
+}
+```
+
+For the self-attention contact prediction:
+
+```bibtex
+@article{rao2020transformer,
+  author = {Rao, Roshan M and Meier, Joshua and Sercu, Tom and Ovchinnikov, Sergey and Rives, Alexander},
+  title={Transformer protein language models are unsupervised structure learners},
+  year={2020},
+  doi={10.1101/2020.12.15.422761},
+  url={https://www.biorxiv.org/content/10.1101/2020.12.15.422761v1},
+  journal={bioRxiv}
+}
+```
+
+For the MSA Transformer:
+
+```bibtex
+@article{rao2021msa,
+  author = {Rao, Roshan and Liu, Jason and Verkuil, Robert and Meier, Joshua and Canny, John F. and Abbeel, Pieter and Sercu, Tom and Rives, Alexander},
+  title={MSA Transformer},
+  year={2021},
+  doi={10.1101/2021.02.12.430858},
+  url={https://www.biorxiv.org/content/10.1101/2021.02.12.430858v1},
+  journal={bioRxiv}
+}
+```
+
+For variant prediction using ESM-1v:
+
+```bibtex
+@article{meier2021language,
+  author = {Meier, Joshua and Rao, Roshan and Verkuil, Robert and Liu, Jason and Sercu, Tom and Rives, Alexander},
+  title = {Language models enable zero-shot prediction of the effects of mutations on protein function},
+  year={2021},
+  doi={10.1101/2021.07.09.450648},
+  url={https://www.biorxiv.org/content/10.1101/2021.07.09.450648v1},
+  journal={bioRxiv}
+}
+```
+
+For inverse folding using ESM-IF1:
+
+```bibtex
+@article{hsu2022learning,
+	author = {Hsu, Chloe and Verkuil, Robert and Liu, Jason and Lin, Zeming and Hie, Brian and Sercu, Tom and Lerer, Adam and Rives, Alexander},
+	title = {Learning inverse folding from millions of predicted structures},
+	year = {2022},
+	doi = {10.1101/2022.04.10.487779},
+	url = {https://www.biorxiv.org/content/early/2022/04/10/2022.04.10.487779},
+	journal = {ICML}
+}
+```
+
+For the ESM-2 language model and ESMFold:
+
+```bibtex
+@article{lin2022language,
+  title={Language models of protein sequences at the scale of evolution enable accurate structure prediction},
+  author={Lin, Zeming and Akin, Halil and Rao, Roshan and Hie, Brian and Zhu, Zhongkai and Lu, Wenting and Smetanin, Nikita and dos Santos Costa, Allan and Fazel-Zarandi, Maryam and Sercu, Tom and Candido, Sal and others},
+  journal={bioRxiv},
+  year={2022},
+  publisher={Cold Spring Harbor Laboratory}
+}
+```
+
+Much of this code builds on the [fairseq](https://github.com/pytorch/fairseq) sequence modeling framework. We use fairseq internally for our protein language modeling research. We highly recommend trying it out if you'd like to pre-train protein language models from scratch.
+
+Additionally, if you would like to use the variant prediction benchmark from Meier et al. (2021), we provide a bibtex file with citations for all data in [./examples/variant-prediction/mutation_data.bib](./examples/variant-prediction/mutation_data.bib). You can cite each paper individually, or add all citations in bulk using the LaTeX command:
+
+```tex
+\nocite{wrenbeck2017deep,klesmith2015comprehensive,haddox2018mapping,romero2015dissecting,firnberg2014comprehensive,deng2012deep,stiffler2015evolvability,jacquier2013capturing,findlay2018comprehensive,mclaughlin2012spatial,kitzman2015massively,doud2016accurate,pokusaeva2019experimental,mishra2016systematic,kelsic2016rna,melnikov2014comprehensive,brenan2016phenotypic,rockah2015systematic,wu2015functional,aakre2015evolving,qi2014quantitative,matreyek2018multiplex,bandaru2017deconstruction,roscoe2013analyses,roscoe2014systematic,mavor2016determination,chan2017correlation,melamed2013deep,starita2013activity,araya2012fundamental}
+```
+
+## License <a name="license"></a>
+
+This source code is licensed under the MIT license found in the `LICENSE` file
+in the root directory of this source tree.
+
+ESM Metagenomic Atlas (also referred to as “ESM Metagenomic Structure Atlas” or “ESM Atlas”) data is available under a CC BY 4.0 license for academic and commercial use. Copyright (c) Meta Platforms, Inc. All Rights Reserved. Use of the ESM Metagenomic Atlas data is subject to the Meta Open Source [Terms of Use](https://opensource.fb.com/legal/terms/) and [Privacy Policy](https://opensource.fb.com/legal/privacy/).

esm/source/__init__.py

@@ -0,0 +1,4 @@
+# -*- coding: utf-8 -*-
+"""
+esm Project Package Initialization File
+"""

esm/source/environment.yml

@@ -0,0 +1,36 @@
+name: esmfold
+channels:
+  - conda-forge
+  - bioconda
+  - pytorch
+dependencies:
+  - conda-forge::python=3.7
+  - conda-forge::setuptools=59.5.0
+  - conda-forge::pip
+  - conda-forge::openmm=7.5.1
+  - conda-forge::pdbfixer
+  - conda-forge::cudatoolkit==11.3.*
+  - conda-forge::einops
+  - conda-forge::fairscale
+  - conda-forge::omegaconf
+  - conda-forge::hydra-core
+  - conda-forge::pandas
+  - conda-forge::pytest
+  - bioconda::hmmer==3.3.2
+  - bioconda::hhsuite==3.3.0
+  - bioconda::kalign2==2.04
+  - pytorch::pytorch=1.12.*
+  - pip:
+      - biopython==1.79
+      - deepspeed==0.5.9
+      - dm-tree==0.1.6
+      - ml-collections==0.1.0
+      - numpy==1.21.2
+      - PyYAML==5.4.1
+      - requests==2.26.0
+      - scipy==1.7.1
+      - tqdm==4.62.2
+      - typing-extensions==3.10.0.2
+      - pytorch_lightning==1.5.10
+      - wandb==0.12.21
+      - git+https://github.com/NVIDIA/dllogger.git

esm/source/esm/__init__.py

@@ -0,0 +1,12 @@
+# 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 .version import version as __version__  # noqa
+
+from .data import Alphabet, BatchConverter, FastaBatchedDataset  # noqa
+from .model.esm1 import ProteinBertModel  # noqa
+from .model.esm2 import ESM2  # noqa
+from .model.msa_transformer import MSATransformer  #noqa
+from . import pretrained  # noqa

esm/source/esm/axial_attention.py

@@ -0,0 +1,239 @@
+# 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
+
+
+class RowSelfAttention(nn.Module):
+    """Compute self-attention over rows of a 2D input."""
+
+    def __init__(
+        self,
+        embed_dim,
+        num_heads,
+        dropout=0.0,
+        max_tokens_per_msa: int = 2 ** 16,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        self.scaling = self.head_dim ** -0.5
+        self.max_tokens_per_msa = max_tokens_per_msa
+        self.attn_shape = "hnij"
+
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+        self.dropout_module = nn.Dropout(dropout)
+
+    def align_scaling(self, q):
+        num_rows = q.size(0)
+        return self.scaling / math.sqrt(num_rows)
+
+    def _batched_forward(
+        self,
+        x,
+        self_attn_mask=None,
+        self_attn_padding_mask=None,
+    ):
+        num_rows, num_cols, batch_size, embed_dim = x.size()
+        max_rows = max(1, self.max_tokens_per_msa // num_cols)
+        attns = 0
+        scaling = self.align_scaling(x)
+        for start in range(0, num_rows, max_rows):
+            attn_weights = self.compute_attention_weights(
+                x[start : start + max_rows],
+                scaling,
+                self_attn_mask=self_attn_mask,
+                self_attn_padding_mask=self_attn_padding_mask[:, start : start + max_rows]
+                if self_attn_padding_mask is not None
+                else None,
+            )
+            attns += attn_weights
+        attn_probs = attns.softmax(-1)
+        attn_probs = self.dropout_module(attn_probs)
+
+        outputs = []
+        for start in range(0, num_rows, max_rows):
+            output = self.compute_attention_update(x[start : start + max_rows], attn_probs)
+            outputs.append(output)
+
+        output = torch.cat(outputs, 0)
+        return output, attn_probs
+
+    def compute_attention_weights(
+        self,
+        x,
+        scaling: float,
+        self_attn_mask=None,
+        self_attn_padding_mask=None,
+    ):
+        num_rows, num_cols, batch_size, embed_dim = x.size()
+        q = self.q_proj(x).view(num_rows, num_cols, batch_size, self.num_heads, self.head_dim)
+        k = self.k_proj(x).view(num_rows, num_cols, batch_size, self.num_heads, self.head_dim)
+        q *= scaling
+        if self_attn_padding_mask is not None:
+            # Zero out any padded aligned positions - this is important since
+            # we take a sum across the alignment axis.
+            q *= 1 - self_attn_padding_mask.permute(1, 2, 0).unsqueeze(3).unsqueeze(4).to(q)
+
+        attn_weights = torch.einsum(f"rinhd,rjnhd->{self.attn_shape}", q, k)
+
+        if self_attn_mask is not None:
+            raise NotImplementedError
+            # Mask Size: [B x R x C], Weights Size: [H x B x C x C]
+
+        if self_attn_padding_mask is not None:
+            attn_weights = attn_weights.masked_fill(
+                self_attn_padding_mask[:, 0].unsqueeze(0).unsqueeze(2),
+                -10000,
+            )
+
+        return attn_weights
+
+    def compute_attention_update(
+        self,
+        x,
+        attn_probs,
+    ):
+        num_rows, num_cols, batch_size, embed_dim = x.size()
+        v = self.v_proj(x).view(num_rows, num_cols, batch_size, self.num_heads, self.head_dim)
+        context = torch.einsum(f"{self.attn_shape},rjnhd->rinhd", attn_probs, v)
+        context = context.contiguous().view(num_rows, num_cols, batch_size, embed_dim)
+        output = self.out_proj(context)
+        return output
+
+    def forward(
+        self,
+        x,
+        self_attn_mask=None,
+        self_attn_padding_mask=None,
+    ):
+        num_rows, num_cols, batch_size, embed_dim = x.size()
+        if (num_rows * num_cols > self.max_tokens_per_msa) and not torch.is_grad_enabled():
+            return self._batched_forward(x, self_attn_mask, self_attn_padding_mask)
+        else:
+            scaling = self.align_scaling(x)
+            attn_weights = self.compute_attention_weights(
+                x, scaling, self_attn_mask, self_attn_padding_mask
+            )
+            attn_probs = attn_weights.softmax(-1)
+            attn_probs = self.dropout_module(attn_probs)
+            output = self.compute_attention_update(x, attn_probs)
+            return output, attn_probs
+
+
+class ColumnSelfAttention(nn.Module):
+    """Compute self-attention over columns of a 2D input."""
+
+    def __init__(
+        self,
+        embed_dim,
+        num_heads,
+        dropout=0.0,
+        max_tokens_per_msa: int = 2 ** 16,
+    ):
+        super().__init__()
+
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        self.scaling = self.head_dim ** -0.5
+        self.max_tokens_per_msa = max_tokens_per_msa
+
+        self.k_proj = nn.Linear(embed_dim, embed_dim)
+        self.v_proj = nn.Linear(embed_dim, embed_dim)
+        self.q_proj = nn.Linear(embed_dim, embed_dim)
+
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+        self.dropout_module = nn.Dropout(dropout)
+
+    def _batched_forward(
+        self,
+        x,
+        self_attn_mask=None,
+        self_attn_padding_mask=None,
+    ):
+        num_rows, num_cols, batch_size, embed_dim = x.size()
+        max_cols = max(1, self.max_tokens_per_msa // num_rows)
+        outputs = []
+        attns = []
+        for start in range(0, num_cols, max_cols):
+            output, attn = self(
+                x[:, start : start + max_cols],
+                self_attn_mask=self_attn_mask,
+                self_attn_padding_mask=self_attn_padding_mask[:, :, start : start + max_cols]
+                if self_attn_padding_mask is not None
+                else None,
+            )
+            outputs.append(output)
+            attns.append(attn)
+        output = torch.cat(outputs, 1)
+        attns = torch.cat(attns, 1)
+        return output, attns
+
+    def compute_attention_update(
+        self,
+        x,
+        self_attn_mask=None,
+        self_attn_padding_mask=None,
+    ):
+        num_rows, num_cols, batch_size, embed_dim = x.size()
+        if num_rows == 1:
+            # if there is only 1 position, this is equivalent and doesn't break with padding
+            attn_probs = torch.ones(
+                self.num_heads,
+                num_cols,
+                batch_size,
+                num_rows,
+                num_rows,
+                device=x.device,
+                dtype=x.dtype,
+            )
+            output = self.out_proj(self.v_proj(x))
+        else:
+            q = self.q_proj(x).view(num_rows, num_cols, batch_size, self.num_heads, self.head_dim)
+            k = self.k_proj(x).view(num_rows, num_cols, batch_size, self.num_heads, self.head_dim)
+            v = self.v_proj(x).view(num_rows, num_cols, batch_size, self.num_heads, self.head_dim)
+            q *= self.scaling
+
+            attn_weights = torch.einsum("icnhd,jcnhd->hcnij", q, k)
+
+            if self_attn_mask is not None:
+                raise NotImplementedError
+            if self_attn_padding_mask is not None:
+                attn_weights = attn_weights.masked_fill(
+                    self_attn_padding_mask.permute(2, 0, 1).unsqueeze(0).unsqueeze(3),
+                    -10000,
+                )
+
+            attn_probs = attn_weights.softmax(-1)
+            attn_probs = self.dropout_module(attn_probs)
+            context = torch.einsum("hcnij,jcnhd->icnhd", attn_probs, v)
+            context = context.contiguous().view(num_rows, num_cols, batch_size, embed_dim)
+            output = self.out_proj(context)
+        return output, attn_probs
+
+    def forward(
+        self,
+        x,
+        self_attn_mask=None,
+        self_attn_padding_mask=None,
+    ):
+        num_rows, num_cols, batch_size, embed_dim = x.size()
+        # if False and num_rows * num_cols > 2 ** 14 and not torch.is_grad_enabled():
+        if (num_rows * num_cols) > self.max_tokens_per_msa and not torch.is_grad_enabled():
+            return self._batched_forward(
+                x,
+                self_attn_mask,
+                self_attn_padding_mask,
+            )
+        else:
+            return self.compute_attention_update(x, self_attn_mask, self_attn_padding_mask)

esm/source/esm/constants.py

@@ -0,0 +1,10 @@
+# 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.
+
+# fmt: off
+proteinseq_toks = {
+    'toks': ['L', 'A', 'G', 'V', 'S', 'E', 'R', 'T', 'I', 'D', 'P', 'K', 'Q', 'N', 'F', 'Y', 'M', 'H', 'W', 'C', 'X', 'B', 'U', 'Z', 'O', '.', '-']
+}
+# fmt: on

esm/source/esm/data.py

@@ -0,0 +1,493 @@
+# 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 itertools
+import os
+from typing import Sequence, Tuple, List, Union
+import pickle
+import re
+import shutil
+import torch
+from pathlib import Path
+from esm.constants import proteinseq_toks
+
+RawMSA = Sequence[Tuple[str, str]]
+
+
+class FastaBatchedDataset(object):
+    def __init__(self, sequence_labels, sequence_strs):
+        self.sequence_labels = list(sequence_labels)
+        self.sequence_strs = list(sequence_strs)
+
+    @classmethod
+    def from_file(cls, fasta_file):
+        sequence_labels, sequence_strs = [], []
+        cur_seq_label = None
+        buf = []
+
+        def _flush_current_seq():
+            nonlocal cur_seq_label, buf
+            if cur_seq_label is None:
+                return
+            sequence_labels.append(cur_seq_label)
+            sequence_strs.append("".join(buf))
+            cur_seq_label = None
+            buf = []
+
+        with open(fasta_file, "r") as infile:
+            for line_idx, line in enumerate(infile):
+                if line.startswith(">"):  # label line
+                    _flush_current_seq()
+                    line = line[1:].strip()
+                    if len(line) > 0:
+                        cur_seq_label = line
+                    else:
+                        cur_seq_label = f"seqnum{line_idx:09d}"
+                else:  # sequence line
+                    buf.append(line.strip())
+
+        _flush_current_seq()
+
+        assert len(set(sequence_labels)) == len(
+            sequence_labels
+        ), "Found duplicate sequence labels"
+
+        return cls(sequence_labels, sequence_strs)
+
+    def __len__(self):
+        return len(self.sequence_labels)
+
+    def __getitem__(self, idx):
+        return self.sequence_labels[idx], self.sequence_strs[idx]
+
+    def get_batch_indices(self, toks_per_batch, extra_toks_per_seq=0):
+        sizes = [(len(s), i) for i, s in enumerate(self.sequence_strs)]
+        sizes.sort()
+        batches = []
+        buf = []
+        max_len = 0
+
+        def _flush_current_buf():
+            nonlocal max_len, buf
+            if len(buf) == 0:
+                return
+            batches.append(buf)
+            buf = []
+            max_len = 0
+
+        for sz, i in sizes:
+            sz += extra_toks_per_seq
+            if max(sz, max_len) * (len(buf) + 1) > toks_per_batch:
+                _flush_current_buf()
+            max_len = max(max_len, sz)
+            buf.append(i)
+
+        _flush_current_buf()
+        return batches
+
+
+class Alphabet(object):
+    def __init__(
+        self,
+        standard_toks: Sequence[str],
+        prepend_toks: Sequence[str] = ("<null_0>", "<pad>", "<eos>", "<unk>"),
+        append_toks: Sequence[str] = ("<cls>", "<mask>", "<sep>"),
+        prepend_bos: bool = True,
+        append_eos: bool = False,
+        use_msa: bool = False,
+    ):
+        self.standard_toks = list(standard_toks)
+        self.prepend_toks = list(prepend_toks)
+        self.append_toks = list(append_toks)
+        self.prepend_bos = prepend_bos
+        self.append_eos = append_eos
+        self.use_msa = use_msa
+
+        self.all_toks = list(self.prepend_toks)
+        self.all_toks.extend(self.standard_toks)
+        for i in range((8 - (len(self.all_toks) % 8)) % 8):
+            self.all_toks.append(f"<null_{i  + 1}>")
+        self.all_toks.extend(self.append_toks)
+
+        self.tok_to_idx = {tok: i for i, tok in enumerate(self.all_toks)}
+
+        self.unk_idx = self.tok_to_idx["<unk>"]
+        self.padding_idx = self.get_idx("<pad>")
+        self.cls_idx = self.get_idx("<cls>")
+        self.mask_idx = self.get_idx("<mask>")
+        self.eos_idx = self.get_idx("<eos>")
+        self.all_special_tokens = ['<eos>', '<unk>', '<pad>', '<cls>', '<mask>']
+        self.unique_no_split_tokens = self.all_toks
+
+    def __len__(self):
+        return len(self.all_toks)
+
+    def get_idx(self, tok):
+        return self.tok_to_idx.get(tok, self.unk_idx)
+
+    def get_tok(self, ind):
+        return self.all_toks[ind]
+
+    def to_dict(self):
+        return self.tok_to_idx.copy()
+
+    def get_batch_converter(self, truncation_seq_length: int = None):
+        if self.use_msa:
+            return MSABatchConverter(self, truncation_seq_length)
+        else:
+            return BatchConverter(self, truncation_seq_length)
+
+    @classmethod
+    def from_architecture(cls, name: str) -> "Alphabet":
+        if name in ("ESM-1", "protein_bert_base"):
+            standard_toks = proteinseq_toks["toks"]
+            prepend_toks: Tuple[str, ...] = ("<null_0>", "<pad>", "<eos>", "<unk>")
+            append_toks: Tuple[str, ...] = ("<cls>", "<mask>", "<sep>")
+            prepend_bos = True
+            append_eos = False
+            use_msa = False
+        elif name in ("ESM-1b", "roberta_large"):
+            standard_toks = proteinseq_toks["toks"]
+            prepend_toks = ("<cls>", "<pad>", "<eos>", "<unk>")
+            append_toks = ("<mask>",)
+            prepend_bos = True
+            append_eos = True
+            use_msa = False
+        elif name in ("MSA Transformer", "msa_transformer"):
+            standard_toks = proteinseq_toks["toks"]
+            prepend_toks = ("<cls>", "<pad>", "<eos>", "<unk>")
+            append_toks = ("<mask>",)
+            prepend_bos = True
+            append_eos = False
+            use_msa = True
+        elif "invariant_gvp" in name.lower():
+            standard_toks = proteinseq_toks["toks"]
+            prepend_toks = ("<null_0>", "<pad>", "<eos>", "<unk>")
+            append_toks = ("<mask>", "<cath>", "<af2>")
+            prepend_bos = True
+            append_eos = False
+            use_msa = False
+        else:
+            raise ValueError("Unknown architecture selected")
+        return cls(standard_toks, prepend_toks, append_toks, prepend_bos, append_eos, use_msa)
+
+    def _tokenize(self, text) -> str:
+        return text.split()
+
+    def tokenize(self, text, **kwargs) -> List[str]:
+        """
+        Inspired by https://github.com/huggingface/transformers/blob/master/src/transformers/tokenization_utils.py
+        Converts a string in a sequence of tokens, using the tokenizer.
+
+        Args:
+            text (:obj:`str`):
+                The sequence to be encoded.
+
+        Returns:
+            :obj:`List[str]`: The list of tokens.
+        """
+
+        def split_on_token(tok, text):
+            result = []
+            split_text = text.split(tok)
+            for i, sub_text in enumerate(split_text):
+                # AddedToken can control whitespace stripping around them.
+                # We use them for GPT2 and Roberta to have different behavior depending on the special token
+                # Cf. https://github.com/huggingface/transformers/pull/2778
+                # and https://github.com/huggingface/transformers/issues/3788
+                # We strip left and right by default
+                if i < len(split_text) - 1:
+                    sub_text = sub_text.rstrip()
+                if i > 0:
+                    sub_text = sub_text.lstrip()
+
+                if i == 0 and not sub_text:
+                    result.append(tok)
+                elif i == len(split_text) - 1:
+                    if sub_text:
+                        result.append(sub_text)
+                    else:
+                        pass
+                else:
+                    if sub_text:
+                        result.append(sub_text)
+                    result.append(tok)
+            return result
+
+        def split_on_tokens(tok_list, text):
+            if not text.strip():
+                return []
+
+            tokenized_text = []
+            text_list = [text]
+            for tok in tok_list:
+                tokenized_text = []
+                for sub_text in text_list:
+                    if sub_text not in self.unique_no_split_tokens:
+                        tokenized_text.extend(split_on_token(tok, sub_text))
+                    else:
+                        tokenized_text.append(sub_text)
+                text_list = tokenized_text
+
+            return list(
+                itertools.chain.from_iterable(
+                    (
+                        self._tokenize(token)
+                        if token not in self.unique_no_split_tokens
+                        else [token]
+                        for token in tokenized_text
+                    )
+                )
+            )
+
+        no_split_token = self.unique_no_split_tokens
+        tokenized_text = split_on_tokens(no_split_token, text)
+        return tokenized_text
+
+    def encode(self, text):
+        return [self.tok_to_idx[tok] for tok in self.tokenize(text)]
+
+
+class BatchConverter(object):
+    """Callable to convert an unprocessed (labels + strings) batch to a
+    processed (labels + tensor) batch.
+    """
+
+    def __init__(self, alphabet, truncation_seq_length: int = None):
+        self.alphabet = alphabet
+        self.truncation_seq_length = truncation_seq_length
+
+    def __call__(self, raw_batch: Sequence[Tuple[str, str]]):
+        # RoBERTa uses an eos token, while ESM-1 does not.
+        batch_size = len(raw_batch)
+        batch_labels, seq_str_list = zip(*raw_batch)
+        seq_encoded_list = [self.alphabet.encode(seq_str) for seq_str in seq_str_list]
+        if self.truncation_seq_length:
+            seq_encoded_list = [seq_str[:self.truncation_seq_length] for seq_str in seq_encoded_list]
+        max_len = max(len(seq_encoded) for seq_encoded in seq_encoded_list)
+        tokens = torch.empty(
+            (
+                batch_size,
+                max_len + int(self.alphabet.prepend_bos) + int(self.alphabet.append_eos),
+            ),
+            dtype=torch.int64,
+        )
+        tokens.fill_(self.alphabet.padding_idx)
+        labels = []
+        strs = []
+
+        for i, (label, seq_str, seq_encoded) in enumerate(
+            zip(batch_labels, seq_str_list, seq_encoded_list)
+        ):
+            labels.append(label)
+            strs.append(seq_str)
+            if self.alphabet.prepend_bos:
+                tokens[i, 0] = self.alphabet.cls_idx
+            seq = torch.tensor(seq_encoded, dtype=torch.int64)
+            tokens[
+                i,
+                int(self.alphabet.prepend_bos) : len(seq_encoded)
+                + int(self.alphabet.prepend_bos),
+            ] = seq
+            if self.alphabet.append_eos:
+                tokens[i, len(seq_encoded) + int(self.alphabet.prepend_bos)] = self.alphabet.eos_idx
+
+        return labels, strs, tokens
+
+
+class MSABatchConverter(BatchConverter):
+    def __call__(self, inputs: Union[Sequence[RawMSA], RawMSA]):
+        if isinstance(inputs[0][0], str):
+            # Input is a single MSA
+            raw_batch: Sequence[RawMSA] = [inputs]  # type: ignore
+        else:
+            raw_batch = inputs  # type: ignore
+
+        batch_size = len(raw_batch)
+        max_alignments = max(len(msa) for msa in raw_batch)
+        max_seqlen = max(len(msa[0][1]) for msa in raw_batch)
+
+        tokens = torch.empty(
+            (
+                batch_size,
+                max_alignments,
+                max_seqlen + int(self.alphabet.prepend_bos) + int(self.alphabet.append_eos),
+            ),
+            dtype=torch.int64,
+        )
+        tokens.fill_(self.alphabet.padding_idx)
+        labels = []
+        strs = []
+
+        for i, msa in enumerate(raw_batch):
+            msa_seqlens = set(len(seq) for _, seq in msa)
+            if not len(msa_seqlens) == 1:
+                raise RuntimeError(
+                    "Received unaligned sequences for input to MSA, all sequence "
+                    "lengths must be equal."
+                )
+            msa_labels, msa_strs, msa_tokens = super().__call__(msa)
+            labels.append(msa_labels)
+            strs.append(msa_strs)
+            tokens[i, : msa_tokens.size(0), : msa_tokens.size(1)] = msa_tokens
+
+        return labels, strs, tokens
+
+
+def read_fasta(
+    path,
+    keep_gaps=True,
+    keep_insertions=True,
+    to_upper=False,
+):
+    with open(path, "r") as f:
+        for result in read_alignment_lines(
+            f, keep_gaps=keep_gaps, keep_insertions=keep_insertions, to_upper=to_upper
+        ):
+            yield result
+
+
+def read_alignment_lines(
+    lines,
+    keep_gaps=True,
+    keep_insertions=True,
+    to_upper=False,
+):
+    seq = desc = None
+
+    def parse(s):
+        if not keep_gaps:
+            s = re.sub("-", "", s)
+        if not keep_insertions:
+            s = re.sub("[a-z]", "", s)
+        return s.upper() if to_upper else s
+
+    for line in lines:
+        # Line may be empty if seq % file_line_width == 0
+        if len(line) > 0 and line[0] == ">":
+            if seq is not None:
+                yield desc, parse(seq)
+            desc = line.strip().lstrip(">")
+            seq = ""
+        else:
+            assert isinstance(seq, str)
+            seq += line.strip()
+    assert isinstance(seq, str) and isinstance(desc, str)
+    yield desc, parse(seq)
+
+
+class ESMStructuralSplitDataset(torch.utils.data.Dataset):
+    """
+    Structural Split Dataset as described in section A.10 of the supplement of our paper.
+    https://doi.org/10.1101/622803
+
+    We use the full version of SCOPe 2.07, clustered at 90% sequence identity,
+    generated on January 23, 2020.
+
+    For each SCOPe domain:
+        - We extract the sequence from the corresponding PDB file
+        - We extract the 3D coordinates of the Carbon beta atoms, aligning them
+          to the sequence. We put NaN where Cb atoms are missing.
+        - From the 3D coordinates, we calculate a pairwise distance map, based
+          on L2 distance
+        - We use DSSP to generate secondary structure labels for the corresponding
+          PDB file. This is also aligned to the sequence. We put - where SSP
+          labels are missing.
+
+    For each SCOPe classification level of family/superfamily/fold (in order of difficulty),
+    we have split the data into 5 partitions for cross validation. These are provided
+    in a downloaded splits folder, in the format:
+            splits/{split_level}/{cv_partition}/{train|valid}.txt
+    where train is the partition and valid is the concatentation of the remaining 4.
+
+    For each SCOPe domain, we provide a pkl dump that contains:
+        - seq    : The domain sequence, stored as an L-length string
+        - ssp    : The secondary structure labels, stored as an L-length string
+        - dist   : The distance map, stored as an LxL numpy array
+        - coords : The 3D coordinates, stored as an Lx3 numpy array
+
+    """
+
+    base_folder = "structural-data"
+    file_list = [
+        #  url  tar filename   filename      MD5 Hash
+        (
+            "https://dl.fbaipublicfiles.com/fair-esm/structural-data/splits.tar.gz",
+            "splits.tar.gz",
+            "splits",
+            "456fe1c7f22c9d3d8dfe9735da52411d",
+        ),
+        (
+            "https://dl.fbaipublicfiles.com/fair-esm/structural-data/pkl.tar.gz",
+            "pkl.tar.gz",
+            "pkl",
+            "644ea91e56066c750cd50101d390f5db",
+        ),
+    ]
+
+    def __init__(
+        self,
+        split_level,
+        cv_partition,
+        split,
+        root_path=os.path.expanduser("~/.cache/torch/data/esm"),
+        download=False,
+    ):
+        super().__init__()
+        assert split in [
+            "train",
+            "valid",
+        ], "train_valid must be 'train' or 'valid'"
+        self.root_path = root_path
+        self.base_path = os.path.join(self.root_path, self.base_folder)
+
+        # check if root path has what you need or else download it
+        if download:
+            self.download()
+
+        self.split_file = os.path.join(
+            self.base_path, "splits", split_level, cv_partition, f"{split}.txt"
+        )
+        self.pkl_dir = os.path.join(self.base_path, "pkl")
+        self.names = []
+        with open(self.split_file) as f:
+            self.names = f.read().splitlines()
+
+    def __len__(self):
+        return len(self.names)
+
+    def _check_exists(self) -> bool:
+        for (_, _, filename, _) in self.file_list:
+            fpath = os.path.join(self.base_path, filename)
+            if not os.path.exists(fpath) or not os.path.isdir(fpath):
+                return False
+        return True
+
+    def download(self):
+
+        if self._check_exists():
+            print("Files already downloaded and verified")
+            return
+
+        from torchvision.datasets.utils import download_url
+
+        for url, tar_filename, filename, md5_hash in self.file_list:
+            download_path = os.path.join(self.base_path, tar_filename)
+            download_url(url=url, root=self.base_path, filename=tar_filename, md5=md5_hash)
+            shutil.unpack_archive(download_path, self.base_path)
+
+    def __getitem__(self, idx):
+        """
+        Returns a dict with the following entires
+         - seq : Str (domain sequence)
+         - ssp : Str (SSP labels)
+         - dist : np.array (distance map)
+         - coords : np.array (3D coordinates)
+        """
+        name = self.names[idx]
+        pkl_fname = os.path.join(self.pkl_dir, name[1:3], f"{name}.pkl")
+        with open(pkl_fname, "rb") as f:
+            obj = pickle.load(f)
+        return obj

esm/source/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()

esm/source/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
+from functools import partial
+
+import torch
+import torch.nn as nn
+from torch import nn
+from torch.nn import LayerNorm
+
+import esm
+from esm import Alphabet
+from esm.esmfold.v1.categorical_mixture import categorical_lddt
+from esm.esmfold.v1.misc import (
+    batch_encode_sequences,
+    collate_dense_tensors,
+    output_to_pdb,
+)
+from esm.esmfold.v1.trunk import FoldingTrunk, FoldingTrunkConfig
+from openfold.data.data_transforms import make_atom14_masks
+from openfold.np import residue_constants
+from openfold.utils.loss import compute_predicted_aligned_error, compute_tm
+
+
+@dataclass
+class ESMFoldConfig:
+    trunk: T.Any = 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
+                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

esm/source/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 openfold.np import residue_constants
+from openfold.np.protein import Protein as OFProtein
+from openfold.np.protein import to_pdb
+from openfold.utils.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)

esm/source/esm/esmfold/v1/pretrained.py

@@ -0,0 +1,181 @@
+# 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 esm.esmfold.v1.esmfold import ESMFold
+
+
+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 = 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")

esm/source/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 openfold.model.triangular_attention import (
+    TriangleAttentionEndingNode,
+    TriangleAttentionStartingNode,
+)
+from openfold.model.triangular_multiplicative_update import (
+    TriangleMultiplicationIncoming,
+    TriangleMultiplicationOutgoing,
+)
+from torch import nn
+
+from 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

esm/source/esm/esmfold/v1/trunk.py

@@ -0,0 +1,243 @@
+# 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
+
+import torch
+import torch.nn as nn
+from openfold.model.structure_module import StructureModule
+
+from 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()
+
+
+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

esm/source/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

esm/source/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) 

esm/source/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

esm/source/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

esm/source/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 esm.data 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])

esm/source/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 esm.modules 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]
+        }

esm/source/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)
+
+

esm/source/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 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]

esm/source/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 esm.modules 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]

esm/source/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 esm.multihead_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

esm/source/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 esm.data 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

esm/source/esm/model/__init__.py

@@ -0,0 +1 @@
+

esm/source/esm/model/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 ..modules 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

esm/source/esm/model/esm2.py

@@ -0,0 +1,147 @@
+# 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
+
+import esm
+from esm.modules 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[esm.data.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, esm.data.Alphabet):
+            alphabet = esm.data.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"]

esm/source/esm/model/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 ..modules import (
+    AxialTransformerLayer,
+    LearnedPositionalEmbedding,
+    RobertaLMHead,
+    ESM1bLayerNorm,
+    ContactPredictionHead,
+)
+
+from ..axial_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

esm/source/esm/modules.py

@@ -0,0 +1,418 @@
+# 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
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .multihead_attention import MultiheadAttention  # noqa
+from .axial_attention import ColumnSelfAttention, RowSelfAttention
+
+
+def gelu(x):
+    """Implementation of the gelu activation function.
+
+    For information: OpenAI GPT's gelu is slightly different
+    (and gives slightly different results):
+    0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
+    """
+    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))
+
+
+def symmetrize(x):
+    "Make layer symmetric in final two dimensions, used for contact prediction."
+    return x + x.transpose(-1, -2)
+
+
+def apc(x):
+    "Perform average product correct, used for contact prediction."
+    a1 = x.sum(-1, keepdims=True)
+    a2 = x.sum(-2, keepdims=True)
+    a12 = x.sum((-1, -2), keepdims=True)
+
+    avg = a1 * a2
+    avg.div_(a12)  # in-place to reduce memory
+    normalized = x - avg
+    return normalized
+
+
+class ESM1LayerNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-12, affine=True):
+        """Construct a layernorm layer in the TF style (eps inside the sqrt)."""
+        super().__init__()
+        self.hidden_size = (hidden_size,) if isinstance(hidden_size, int) else tuple(hidden_size)
+        self.eps = eps
+        self.affine = bool(affine)
+        if self.affine:
+            self.weight = nn.Parameter(torch.ones(hidden_size))
+            self.bias = nn.Parameter(torch.zeros(hidden_size))
+        else:
+            self.weight, self.bias = None, None
+
+    def forward(self, x):
+        dims = tuple(-(i + 1) for i in range(len(self.hidden_size)))
+        means = x.mean(dims, keepdim=True)
+        x_zeromean = x - means
+        variances = x_zeromean.pow(2).mean(dims, keepdim=True)
+        x = x_zeromean / torch.sqrt(variances + self.eps)
+        if self.affine:
+            x = (self.weight * x) + self.bias
+        return x
+
+
+try:
+    from apex.normalization import FusedLayerNorm as _FusedLayerNorm
+
+    class ESM1bLayerNorm(_FusedLayerNorm):
+        @torch.jit.unused
+        def forward(self, x):
+            if not x.is_cuda:
+                return super().forward(x)
+            else:
+                with torch.cuda.device(x.device):
+                    return super().forward(x)
+
+except ImportError:
+    from torch.nn import LayerNorm as ESM1bLayerNorm
+
+
+class TransformerLayer(nn.Module):
+    """Transformer layer block."""
+
+    def __init__(
+        self,
+        embed_dim,
+        ffn_embed_dim,
+        attention_heads,
+        add_bias_kv=True,
+        use_esm1b_layer_norm=False,
+        use_rotary_embeddings: bool = False,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.ffn_embed_dim = ffn_embed_dim
+        self.attention_heads = attention_heads
+        self.use_rotary_embeddings = use_rotary_embeddings
+        self._init_submodules(add_bias_kv, use_esm1b_layer_norm)
+
+    def _init_submodules(self, add_bias_kv, use_esm1b_layer_norm):
+        BertLayerNorm = ESM1bLayerNorm if use_esm1b_layer_norm else ESM1LayerNorm
+
+        self.self_attn = MultiheadAttention(
+            self.embed_dim,
+            self.attention_heads,
+            add_bias_kv=add_bias_kv,
+            add_zero_attn=False,
+            use_rotary_embeddings=self.use_rotary_embeddings,
+        )
+        self.self_attn_layer_norm = BertLayerNorm(self.embed_dim)
+
+        self.fc1 = nn.Linear(self.embed_dim, self.ffn_embed_dim)
+        self.fc2 = nn.Linear(self.ffn_embed_dim, self.embed_dim)
+
+        self.final_layer_norm = BertLayerNorm(self.embed_dim)
+
+    def forward(
+        self, x, self_attn_mask=None, self_attn_padding_mask=None, need_head_weights=False
+    ):
+        residual = x
+        x = self.self_attn_layer_norm(x)
+        x, attn = self.self_attn(
+            query=x,
+            key=x,
+            value=x,
+            key_padding_mask=self_attn_padding_mask,
+            need_weights=True,
+            need_head_weights=need_head_weights,
+            attn_mask=self_attn_mask,
+        )
+        x = residual + x
+
+        residual = x
+        x = self.final_layer_norm(x)
+        x = gelu(self.fc1(x))
+        x = self.fc2(x)
+        x = residual + x
+
+        return x, attn
+
+
+class AxialTransformerLayer(nn.Module):
+    """Implements an Axial MSA Transformer block."""
+
+    def __init__(
+        self,
+        embedding_dim: int = 768,
+        ffn_embedding_dim: int = 3072,
+        num_attention_heads: int = 8,
+        dropout: float = 0.1,
+        attention_dropout: float = 0.1,
+        activation_dropout: float = 0.1,
+        max_tokens_per_msa: int = 2**14,
+    ) -> None:
+        super().__init__()
+
+        # Initialize parameters
+        self.embedding_dim = embedding_dim
+        self.dropout_prob = dropout
+
+        row_self_attention = RowSelfAttention(
+            embedding_dim,
+            num_attention_heads,
+            dropout=dropout,
+            max_tokens_per_msa=max_tokens_per_msa,
+        )
+
+        column_self_attention = ColumnSelfAttention(
+            embedding_dim,
+            num_attention_heads,
+            dropout=dropout,
+            max_tokens_per_msa=max_tokens_per_msa,
+        )
+
+        feed_forward_layer = FeedForwardNetwork(
+            embedding_dim,
+            ffn_embedding_dim,
+            activation_dropout=activation_dropout,
+            max_tokens_per_msa=max_tokens_per_msa,
+        )
+
+        self.row_self_attention = self.build_residual(row_self_attention)
+        self.column_self_attention = self.build_residual(column_self_attention)
+        self.feed_forward_layer = self.build_residual(feed_forward_layer)
+
+    def build_residual(self, layer: nn.Module):
+        return NormalizedResidualBlock(
+            layer,
+            self.embedding_dim,
+            self.dropout_prob,
+        )
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        self_attn_mask: Optional[torch.Tensor] = None,
+        self_attn_padding_mask: Optional[torch.Tensor] = None,
+        need_head_weights: bool = False,
+    ):
+        """
+        LayerNorm is applied either before or after the self-attention/ffn
+        modules similar to the original Transformer implementation.
+        """
+        x, row_attn = self.row_self_attention(
+            x,
+            self_attn_mask=self_attn_mask,
+            self_attn_padding_mask=self_attn_padding_mask,
+        )
+        x, column_attn = self.column_self_attention(
+            x,
+            self_attn_mask=self_attn_mask,
+            self_attn_padding_mask=self_attn_padding_mask,
+        )
+        x = self.feed_forward_layer(x)
+        if need_head_weights:
+            return x, column_attn, row_attn
+        else:
+            return x
+
+
+class LearnedPositionalEmbedding(nn.Embedding):
+    """
+    This module learns positional embeddings up to a fixed maximum size.
+    Padding ids are ignored by either offsetting based on padding_idx
+    or by setting padding_idx to None and ensuring that the appropriate
+    position ids are passed to the forward function.
+    """
+
+    def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int):
+        if padding_idx is not None:
+            num_embeddings_ = num_embeddings + padding_idx + 1
+        else:
+            num_embeddings_ = num_embeddings
+        super().__init__(num_embeddings_, embedding_dim, padding_idx)
+        self.max_positions = num_embeddings
+
+    def forward(self, input: torch.Tensor):
+        """Input is expected to be of size [bsz x seqlen]."""
+        if input.size(1) > self.max_positions:
+            raise ValueError(
+                f"Sequence length {input.size(1)} above maximum "
+                f" sequence length of {self.max_positions}"
+            )
+        mask = input.ne(self.padding_idx).int()
+        positions = (torch.cumsum(mask, dim=1).type_as(mask) * mask).long() + self.padding_idx
+        return F.embedding(
+            positions,
+            self.weight,
+            self.padding_idx,
+            self.max_norm,
+            self.norm_type,
+            self.scale_grad_by_freq,
+            self.sparse,
+        )
+
+
+class SinusoidalPositionalEmbedding(nn.Module):
+    def __init__(self, embed_dim, padding_idx, learned=False):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.padding_idx = padding_idx
+        self.register_buffer("_float_tensor", torch.FloatTensor(1))
+        self.weights = None
+
+    def forward(self, x):
+        bsz, seq_len = x.shape
+        max_pos = self.padding_idx + 1 + seq_len
+        if self.weights is None or max_pos > self.weights.size(0):
+            self.weights = self.get_embedding(max_pos)
+        self.weights = self.weights.type_as(self._float_tensor)
+
+        positions = self.make_positions(x)
+        return self.weights.index_select(0, positions.view(-1)).view(bsz, seq_len, -1).detach()
+
+    def make_positions(self, x):
+        mask = x.ne(self.padding_idx)
+        range_buf = torch.arange(x.size(1), device=x.device).expand_as(x) + self.padding_idx + 1
+        positions = range_buf.expand_as(x)
+        return positions * mask.long() + self.padding_idx * (1 - mask.long())
+
+    def get_embedding(self, num_embeddings):
+        half_dim = self.embed_dim // 2
+        emb = math.log(10000) / (half_dim - 1)
+        emb = torch.exp(torch.arange(half_dim, dtype=torch.float) * -emb)
+        emb = torch.arange(num_embeddings, dtype=torch.float).unsqueeze(1) * emb.unsqueeze(0)
+        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1)
+        if self.embed_dim % 2 == 1:
+            # zero pad
+            emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)
+        if self.padding_idx is not None:
+            emb[self.padding_idx, :] = 0
+        return emb
+
+
+class RobertaLMHead(nn.Module):
+    """Head for masked language modeling."""
+
+    def __init__(self, embed_dim, output_dim, weight):
+        super().__init__()
+        self.dense = nn.Linear(embed_dim, embed_dim)
+        self.layer_norm = ESM1bLayerNorm(embed_dim)
+        self.weight = weight
+        self.bias = nn.Parameter(torch.zeros(output_dim))
+
+    def forward(self, features):
+        x = self.dense(features)
+        x = gelu(x)
+        x = self.layer_norm(x)
+        # project back to size of vocabulary with bias
+        x = F.linear(x, self.weight) + self.bias
+        return x
+
+
+class ContactPredictionHead(nn.Module):
+    """Performs symmetrization, apc, and computes a logistic regression on the output features"""
+
+    def __init__(
+        self,
+        in_features: int,
+        prepend_bos: bool,
+        append_eos: bool,
+        bias=True,
+        eos_idx: Optional[int] = None,
+    ):
+        super().__init__()
+        self.in_features = in_features
+        self.prepend_bos = prepend_bos
+        self.append_eos = append_eos
+        if append_eos and eos_idx is None:
+            raise ValueError("Using an alphabet with eos token, but no eos token was passed in.")
+        self.eos_idx = eos_idx
+        self.regression = nn.Linear(in_features, 1, bias)
+        self.activation = nn.Sigmoid()
+
+    def forward(self, tokens, attentions):
+        # remove eos token attentions
+        if self.append_eos:
+            eos_mask = tokens.ne(self.eos_idx).to(attentions)
+            eos_mask = eos_mask.unsqueeze(1) * eos_mask.unsqueeze(2)
+            attentions = attentions * eos_mask[:, None, None, :, :]
+            attentions = attentions[..., :-1, :-1]
+        # remove cls token attentions
+        if self.prepend_bos:
+            attentions = attentions[..., 1:, 1:]
+        batch_size, layers, heads, seqlen, _ = attentions.size()
+        attentions = attentions.view(batch_size, layers * heads, seqlen, seqlen)
+
+        # features: B x C x T x T
+        attentions = attentions.to(
+            self.regression.weight.device
+        )  # attentions always float32, may need to convert to float16
+        attentions = apc(symmetrize(attentions))
+        attentions = attentions.permute(0, 2, 3, 1)
+        return self.activation(self.regression(attentions).squeeze(3))
+
+
+class NormalizedResidualBlock(nn.Module):
+    def __init__(
+        self,
+        layer: nn.Module,
+        embedding_dim: int,
+        dropout: float = 0.1,
+    ):
+        super().__init__()
+        self.embedding_dim = embedding_dim
+
+        self.layer = layer
+        self.dropout_module = nn.Dropout(
+            dropout,
+        )
+        self.layer_norm = ESM1bLayerNorm(self.embedding_dim)
+
+    def forward(self, x, *args, **kwargs):
+        residual = x
+        x = self.layer_norm(x)
+        outputs = self.layer(x, *args, **kwargs)
+        if isinstance(outputs, tuple):
+            x, *out = outputs
+        else:
+            x = outputs
+            out = None
+
+        x = self.dropout_module(x)
+        x = residual + x
+
+        if out is not None:
+            return (x,) + tuple(out)
+        else:
+            return x
+
+
+class FeedForwardNetwork(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: int,
+        ffn_embedding_dim: int,
+        activation_dropout: float = 0.1,
+        max_tokens_per_msa: int = 2**14,
+    ):
+        super().__init__()
+        self.embedding_dim = embedding_dim
+        self.ffn_embedding_dim = ffn_embedding_dim
+        self.max_tokens_per_msa = max_tokens_per_msa
+        self.activation_fn = nn.GELU()
+        self.activation_dropout_module = nn.Dropout(
+            activation_dropout,
+        )
+        self.fc1 = nn.Linear(embedding_dim, ffn_embedding_dim)
+        self.fc2 = nn.Linear(ffn_embedding_dim, embedding_dim)
+
+    def forward(self, x):
+        x = self.activation_fn(self.fc1(x))
+        x = self.activation_dropout_module(x)
+        x = self.fc2(x)
+        return x

esm/source/esm/multihead_attention.py

@@ -0,0 +1,508 @@
+# 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
+from typing import Dict, Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor, nn
+from torch.nn import Parameter
+from esm.rotary_embedding import RotaryEmbedding
+
+import uuid
+
+
+def utils_softmax(x, dim: int, onnx_trace: bool = False):
+    if onnx_trace:
+        return F.softmax(x.float(), dim=dim)
+    else:
+        return F.softmax(x, dim=dim, dtype=torch.float32)
+
+
+class FairseqIncrementalState(object):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.init_incremental_state()
+
+    def init_incremental_state(self):
+        self._incremental_state_id = str(uuid.uuid4())
+
+    def _get_full_incremental_state_key(self, key: str) -> str:
+        return "{}.{}".format(self._incremental_state_id, key)
+
+    def get_incremental_state(
+        self,
+        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]],
+        key: str,
+    ) -> Optional[Dict[str, Optional[Tensor]]]:
+        """Helper for getting incremental state for an nn.Module."""
+        full_key = self._get_full_incremental_state_key(key)
+        if incremental_state is None or full_key not in incremental_state:
+            return None
+        return incremental_state[full_key]
+
+    def set_incremental_state(
+        self,
+        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]],
+        key: str,
+        value: Dict[str, Optional[Tensor]],
+    ) -> Optional[Dict[str, Dict[str, Optional[Tensor]]]]:
+        """Helper for setting incremental state for an nn.Module."""
+        if incremental_state is not None:
+            full_key = self._get_full_incremental_state_key(key)
+            incremental_state[full_key] = value
+        return incremental_state
+
+
+def with_incremental_state(cls):
+    cls.__bases__ = (FairseqIncrementalState,) + tuple(
+        b for b in cls.__bases__ if b != FairseqIncrementalState
+    )
+    return cls
+
+
+@with_incremental_state
+class MultiheadAttention(nn.Module):
+    """Multi-headed attention.
+
+    See "Attention Is All You Need" for more details.
+    """
+
+    def __init__(
+        self,
+        embed_dim,
+        num_heads,
+        kdim=None,
+        vdim=None,
+        dropout=0.0,
+        bias=True,
+        add_bias_kv: bool = False,
+        add_zero_attn: bool = False,
+        self_attention: bool = False,
+        encoder_decoder_attention: bool = False,
+        use_rotary_embeddings: bool = False,
+    ):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.kdim = kdim if kdim is not None else embed_dim
+        self.vdim = vdim if vdim is not None else embed_dim
+        self.qkv_same_dim = self.kdim == embed_dim and self.vdim == embed_dim
+
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        assert (
+            self.head_dim * num_heads == self.embed_dim
+        ), "embed_dim must be divisible by num_heads"
+        self.scaling = self.head_dim**-0.5
+
+        self.self_attention = self_attention
+        self.encoder_decoder_attention = encoder_decoder_attention
+
+        assert not self.self_attention or self.qkv_same_dim, (
+            "Self-attention requires query, key and " "value to be of the same size"
+        )
+
+        self.k_proj = nn.Linear(self.kdim, embed_dim, bias=bias)
+        self.v_proj = nn.Linear(self.vdim, embed_dim, bias=bias)
+        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
+
+        if add_bias_kv:
+            self.bias_k = Parameter(torch.Tensor(1, 1, embed_dim))
+            self.bias_v = Parameter(torch.Tensor(1, 1, embed_dim))
+        else:
+            self.bias_k = self.bias_v = None
+
+        self.add_zero_attn = add_zero_attn
+
+        self.reset_parameters()
+
+        self.onnx_trace = False
+        self.rot_emb = None
+        if use_rotary_embeddings:
+            self.rot_emb = RotaryEmbedding(dim=self.head_dim)
+
+        self.enable_torch_version = False
+        if hasattr(F, "multi_head_attention_forward"):
+            self.enable_torch_version = True
+        else:
+            self.enable_torch_version = False
+
+    def prepare_for_onnx_export_(self):
+        self.onnx_trace = True
+
+    def reset_parameters(self):
+        if self.qkv_same_dim:
+            # Empirically observed the convergence to be much better with
+            # the scaled initialization
+            nn.init.xavier_uniform_(self.k_proj.weight, gain=1 / math.sqrt(2))
+            nn.init.xavier_uniform_(self.v_proj.weight, gain=1 / math.sqrt(2))
+            nn.init.xavier_uniform_(self.q_proj.weight, gain=1 / math.sqrt(2))
+        else:
+            nn.init.xavier_uniform_(self.k_proj.weight)
+            nn.init.xavier_uniform_(self.v_proj.weight)
+            nn.init.xavier_uniform_(self.q_proj.weight)
+
+        nn.init.xavier_uniform_(self.out_proj.weight)
+        if self.out_proj.bias is not None:
+            nn.init.constant_(self.out_proj.bias, 0.0)
+        if self.bias_k is not None:
+            nn.init.xavier_normal_(self.bias_k)
+        if self.bias_v is not None:
+            nn.init.xavier_normal_(self.bias_v)
+
+    def forward(
+        self,
+        query,
+        key: Optional[Tensor],
+        value: Optional[Tensor],
+        key_padding_mask: Optional[Tensor] = None,
+        incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None,
+        need_weights: bool = True,
+        static_kv: bool = False,
+        attn_mask: Optional[Tensor] = None,
+        before_softmax: bool = False,
+        need_head_weights: bool = False,
+    ) -> Tuple[Tensor, Optional[Tensor]]:
+        """Input shape: Time x Batch x Channel
+
+        Args:
+            key_padding_mask (ByteTensor, optional): mask to exclude
+                keys that are pads, of shape `(batch, src_len)`, where
+                padding elements are indicated by 1s.
+            need_weights (bool, optional): return the attention weights,
+                averaged over heads (default: False).
+            attn_mask (ByteTensor, optional): typically used to
+                implement causal attention, where the mask prevents the
+                attention from looking forward in time (default: None).
+            before_softmax (bool, optional): return the raw attention
+                weights and values before the attention softmax.
+            need_head_weights (bool, optional): return the attention
+                weights for each head. Implies *need_weights*. Default:
+                return the average attention weights over all heads.
+        """
+        if need_head_weights:
+            need_weights = True
+
+        tgt_len, bsz, embed_dim = query.size()
+        assert embed_dim == self.embed_dim
+        assert list(query.size()) == [tgt_len, bsz, embed_dim]
+
+        if (
+            not self.rot_emb
+            and self.enable_torch_version
+            and not self.onnx_trace
+            and incremental_state is None
+            and not static_kv
+            # A workaround for quantization to work. Otherwise JIT compilation
+            # treats bias in linear module as method.
+            and not torch.jit.is_scripting()
+            and not need_head_weights
+        ):
+            assert key is not None and value is not None
+            return F.multi_head_attention_forward(
+                query,
+                key,
+                value,
+                self.embed_dim,
+                self.num_heads,
+                torch.empty([0]),
+                torch.cat((self.q_proj.bias, self.k_proj.bias, self.v_proj.bias)),
+                self.bias_k,
+                self.bias_v,
+                self.add_zero_attn,
+                self.dropout,
+                self.out_proj.weight,
+                self.out_proj.bias,
+                self.training,
+                key_padding_mask,
+                need_weights,
+                attn_mask,
+                use_separate_proj_weight=True,
+                q_proj_weight=self.q_proj.weight,
+                k_proj_weight=self.k_proj.weight,
+                v_proj_weight=self.v_proj.weight,
+            )
+        if incremental_state is not None:
+            saved_state = self._get_input_buffer(incremental_state)
+            if saved_state is not None and "prev_key" in saved_state:
+                # previous time steps are cached - no need to recompute
+                # key and value if they are static
+                if static_kv:
+                    assert self.encoder_decoder_attention and not self.self_attention
+                    key = value = None
+        else:
+            saved_state = None
+
+        if self.self_attention:
+            q = self.q_proj(query)
+            k = self.k_proj(query)
+            v = self.v_proj(query)
+        elif self.encoder_decoder_attention:
+            # encoder-decoder attention
+            q = self.q_proj(query)
+            if key is None:
+                assert value is None
+                k = v = None
+            else:
+                k = self.k_proj(key)
+                v = self.v_proj(key)
+
+        else:
+            assert key is not None and value is not None
+            q = self.q_proj(query)
+            k = self.k_proj(key)
+            v = self.v_proj(value)
+        q *= self.scaling
+
+        if self.bias_k is not None:
+            assert self.bias_v is not None
+            k = torch.cat([k, self.bias_k.repeat(1, bsz, 1)])
+            v = torch.cat([v, self.bias_v.repeat(1, bsz, 1)])
+            if attn_mask is not None:
+                attn_mask = torch.cat(
+                    [attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1
+                )
+            if key_padding_mask is not None:
+                key_padding_mask = torch.cat(
+                    [
+                        key_padding_mask,
+                        key_padding_mask.new_zeros(key_padding_mask.size(0), 1),
+                    ],
+                    dim=1,
+                )
+
+        q = q.contiguous().view(tgt_len, bsz * self.num_heads, self.head_dim).transpose(0, 1)
+        if k is not None:
+            k = k.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)
+        if v is not None:
+            v = v.contiguous().view(-1, bsz * self.num_heads, self.head_dim).transpose(0, 1)
+
+        if saved_state is not None:
+            # saved states are stored with shape (bsz, num_heads, seq_len, head_dim)
+            if "prev_key" in saved_state:
+                _prev_key = saved_state["prev_key"]
+                assert _prev_key is not None
+                prev_key = _prev_key.view(bsz * self.num_heads, -1, self.head_dim)
+                if static_kv:
+                    k = prev_key
+                else:
+                    assert k is not None
+                    k = torch.cat([prev_key, k], dim=1)
+            if "prev_value" in saved_state:
+                _prev_value = saved_state["prev_value"]
+                assert _prev_value is not None
+                prev_value = _prev_value.view(bsz * self.num_heads, -1, self.head_dim)
+                if static_kv:
+                    v = prev_value
+                else:
+                    assert v is not None
+                    v = torch.cat([prev_value, v], dim=1)
+            prev_key_padding_mask: Optional[Tensor] = None
+            if "prev_key_padding_mask" in saved_state:
+                prev_key_padding_mask = saved_state["prev_key_padding_mask"]
+            assert k is not None and v is not None
+            key_padding_mask = MultiheadAttention._append_prev_key_padding_mask(
+                key_padding_mask=key_padding_mask,
+                prev_key_padding_mask=prev_key_padding_mask,
+                batch_size=bsz,
+                src_len=k.size(1),
+                static_kv=static_kv,
+            )
+
+            saved_state["prev_key"] = k.view(bsz, self.num_heads, -1, self.head_dim)
+            saved_state["prev_value"] = v.view(bsz, self.num_heads, -1, self.head_dim)
+            saved_state["prev_key_padding_mask"] = key_padding_mask
+            # In this branch incremental_state is never None
+            assert incremental_state is not None
+            incremental_state = self._set_input_buffer(incremental_state, saved_state)
+        assert k is not None
+        src_len = k.size(1)
+
+        # This is part of a workaround to get around fork/join parallelism
+        # not supporting Optional types.
+        if key_padding_mask is not None and key_padding_mask.dim() == 0:
+            key_padding_mask = None
+
+        if key_padding_mask is not None:
+            assert key_padding_mask.size(0) == bsz
+            assert key_padding_mask.size(1) == src_len
+
+        if self.add_zero_attn:
+            assert v is not None
+            src_len += 1
+            k = torch.cat([k, k.new_zeros((k.size(0), 1) + k.size()[2:])], dim=1)
+            v = torch.cat([v, v.new_zeros((v.size(0), 1) + v.size()[2:])], dim=1)
+            if attn_mask is not None:
+                attn_mask = torch.cat(
+                    [attn_mask, attn_mask.new_zeros(attn_mask.size(0), 1)], dim=1
+                )
+            if key_padding_mask is not None:
+                key_padding_mask = torch.cat(
+                    [
+                        key_padding_mask,
+                        torch.zeros(key_padding_mask.size(0), 1).type_as(key_padding_mask),
+                    ],
+                    dim=1,
+                )
+
+        if self.rot_emb:
+            q, k = self.rot_emb(q, k)
+
+        attn_weights = torch.bmm(q, k.transpose(1, 2))
+        attn_weights = MultiheadAttention.apply_sparse_mask(attn_weights, tgt_len, src_len, bsz)
+
+        assert list(attn_weights.size()) == [bsz * self.num_heads, tgt_len, src_len]
+
+        if attn_mask is not None:
+            attn_mask = attn_mask.unsqueeze(0)
+            if self.onnx_trace:
+                attn_mask = attn_mask.repeat(attn_weights.size(0), 1, 1)
+            attn_weights += attn_mask
+
+        if key_padding_mask is not None:
+            # don't attend to padding symbols
+            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
+            attn_weights = attn_weights.masked_fill(
+                key_padding_mask.unsqueeze(1).unsqueeze(2).to(torch.bool), float("-inf")
+            )
+            attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
+
+        if before_softmax:
+            return attn_weights, v
+
+        attn_weights_float = utils_softmax(attn_weights, dim=-1, onnx_trace=self.onnx_trace)
+        attn_weights = attn_weights_float.type_as(attn_weights)
+        attn_probs = F.dropout(
+            attn_weights_float.type_as(attn_weights),
+            p=self.dropout,
+            training=self.training,
+        )
+        assert v is not None
+        attn = torch.bmm(attn_probs, v)
+        assert list(attn.size()) == [bsz * self.num_heads, tgt_len, self.head_dim]
+        if self.onnx_trace and attn.size(1) == 1:
+            # when ONNX tracing a single decoder step (sequence length == 1)
+            # the transpose is a no-op copy before view, thus unnecessary
+            attn = attn.contiguous().view(tgt_len, bsz, embed_dim)
+        else:
+            attn = attn.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
+        attn = self.out_proj(attn)
+        attn_weights: Optional[Tensor] = None
+        if need_weights:
+            attn_weights = attn_weights_float.view(
+                bsz, self.num_heads, tgt_len, src_len
+            ).type_as(attn).transpose(1, 0)
+            if not need_head_weights:
+                # average attention weights over heads
+                attn_weights = attn_weights.mean(dim=0)
+
+        return attn, attn_weights
+
+    @staticmethod
+    def _append_prev_key_padding_mask(
+        key_padding_mask: Optional[Tensor],
+        prev_key_padding_mask: Optional[Tensor],
+        batch_size: int,
+        src_len: int,
+        static_kv: bool,
+    ) -> Optional[Tensor]:
+        # saved key padding masks have shape (bsz, seq_len)
+        if prev_key_padding_mask is not None and static_kv:
+            new_key_padding_mask = prev_key_padding_mask
+        elif prev_key_padding_mask is not None and key_padding_mask is not None:
+            new_key_padding_mask = torch.cat(
+                [prev_key_padding_mask.float(), key_padding_mask.float()], dim=1
+            )
+        # During incremental decoding, as the padding token enters and
+        # leaves the frame, there will be a time when prev or current
+        # is None
+        elif prev_key_padding_mask is not None:
+            filler = torch.zeros(
+                (batch_size, src_len - prev_key_padding_mask.size(1)),
+                device=prev_key_padding_mask.device,
+            )
+            new_key_padding_mask = torch.cat(
+                [prev_key_padding_mask.float(), filler.float()], dim=1
+            )
+        elif key_padding_mask is not None:
+            filler = torch.zeros(
+                (batch_size, src_len - key_padding_mask.size(1)),
+                device=key_padding_mask.device,
+            )
+            new_key_padding_mask = torch.cat([filler.float(), key_padding_mask.float()], dim=1)
+        else:
+            new_key_padding_mask = prev_key_padding_mask
+        return new_key_padding_mask
+
+    @torch.jit.export
+    def reorder_incremental_state(
+        self, incremental_state: Dict[str, Dict[str, Optional[Tensor]]], new_order: Tensor
+    ):
+        """Reorder buffered internal state (for incremental generation)."""
+        input_buffer = self._get_input_buffer(incremental_state)
+        if input_buffer is not None:
+            for k in input_buffer.keys():
+                input_buffer_k = input_buffer[k]
+                if input_buffer_k is not None:
+                    if self.encoder_decoder_attention and input_buffer_k.size(0) == new_order.size(
+                        0
+                    ):
+                        break
+                    input_buffer[k] = input_buffer_k.index_select(0, new_order)
+            incremental_state = self._set_input_buffer(incremental_state, input_buffer)
+        return incremental_state
+
+    def _get_input_buffer(
+        self, incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]]
+    ) -> Dict[str, Optional[Tensor]]:
+        result = self.get_incremental_state(incremental_state, "attn_state")
+        if result is not None:
+            return result
+        else:
+            empty_result: Dict[str, Optional[Tensor]] = {}
+            return empty_result
+
+    def _set_input_buffer(
+        self,
+        incremental_state: Dict[str, Dict[str, Optional[Tensor]]],
+        buffer: Dict[str, Optional[Tensor]],
+    ):
+        return self.set_incremental_state(incremental_state, "attn_state", buffer)
+
+    def apply_sparse_mask(attn_weights, tgt_len: int, src_len: int, bsz: int):
+        return attn_weights
+
+    def upgrade_state_dict_named(self, state_dict, name):
+        prefix = name + "." if name != "" else ""
+        items_to_add = {}
+        keys_to_remove = []
+        for k in state_dict.keys():
+            if k.endswith(prefix + "in_proj_weight"):
+                # in_proj_weight used to be q + k + v with same dimensions
+                dim = int(state_dict[k].shape[0] / 3)
+                items_to_add[prefix + "q_proj.weight"] = state_dict[k][:dim]
+                items_to_add[prefix + "k_proj.weight"] = state_dict[k][dim : 2 * dim]
+                items_to_add[prefix + "v_proj.weight"] = state_dict[k][2 * dim :]
+
+                keys_to_remove.append(k)
+
+                k_bias = prefix + "in_proj_bias"
+                if k_bias in state_dict.keys():
+                    dim = int(state_dict[k].shape[0] / 3)
+                    items_to_add[prefix + "q_proj.bias"] = state_dict[k_bias][:dim]
+                    items_to_add[prefix + "k_proj.bias"] = state_dict[k_bias][dim : 2 * dim]
+                    items_to_add[prefix + "v_proj.bias"] = state_dict[k_bias][2 * dim :]
+
+                    keys_to_remove.append(prefix + "in_proj_bias")
+
+        for k in keys_to_remove:
+            del state_dict[k]
+
+        for key, value in items_to_add.items():
+            state_dict[key] = value

esm/source/esm/pretrained.py

@@ -0,0 +1,552 @@
+# 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 esm
+from esm.model.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 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 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 = esm.data.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 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 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 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 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 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 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 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 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 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 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 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 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 esm.esmfold.v1.pretrained
+    return esm.esmfold.v1.pretrained.esmfold_structure_module_only_15B()

esm/source/esm/rotary_embedding.py

@@ -0,0 +1,69 @@
+# 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 Tuple
+
+import torch
+
+
+def rotate_half(x):
+    x1, x2 = x.chunk(2, dim=-1)
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(x, cos, sin):
+    cos = cos[:, : x.shape[-2], :]
+    sin = sin[:, : x.shape[-2], :]
+
+    return (x * cos) + (rotate_half(x) * sin)
+
+
+class RotaryEmbedding(torch.nn.Module):
+    """
+    The rotary position embeddings from RoFormer_ (Su et. al).
+    A crucial insight from the method is that the query and keys are
+    transformed by rotation matrices which depend on the relative positions.
+    Other implementations are available in the Rotary Transformer repo_ and in
+    GPT-NeoX_, GPT-NeoX was an inspiration
+    .. _RoFormer: https://arxiv.org/abs/2104.09864
+    .. _repo: https://github.com/ZhuiyiTechnology/roformer
+    .. _GPT-NeoX: https://github.com/EleutherAI/gpt-neox
+    .. warning: Please note that this embedding is not registered on purpose, as it is transformative
+        (it does not create the embedding dimension) and will likely be picked up (imported) on a ad-hoc basis
+    """
+
+    def __init__(self, dim: int, *_, **__):
+        super().__init__()
+        # Generate and save the inverse frequency buffer (non trainable)
+        inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2).float() / dim))
+        self.register_buffer("inv_freq", inv_freq)
+
+        self._seq_len_cached = None
+        self._cos_cached = None
+        self._sin_cached = None
+
+    def _update_cos_sin_tables(self, x, seq_dimension=1):
+        seq_len = x.shape[seq_dimension]
+
+        # Reset the tables if the sequence length has changed,
+        # or if we're on a new device (possibly due to tracing for instance)
+        if seq_len != self._seq_len_cached or self._cos_cached.device != x.device:
+            self._seq_len_cached = seq_len
+            t = torch.arange(x.shape[seq_dimension], device=x.device).type_as(self.inv_freq)
+            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
+
+            self._cos_cached = emb.cos()[None, :, :]
+            self._sin_cached = emb.sin()[None, :, :]
+
+        return self._cos_cached, self._sin_cached
+
+    def forward(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        self._cos_cached, self._sin_cached = self._update_cos_sin_tables(k, seq_dimension=-2)
+
+        return (
+            apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached),
+            apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached),
+        )

esm/source/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"

esm/source/examples/README.md

@@ -0,0 +1,11 @@
+# What's in this directory
+
+* The notebooks are introduced and summarized in `../README.md`
+* `data/some_proteins.fasta` and its smaller version, `data/few_proteins.fasta` are a random selection of UniRef50 sequences used in the second example of `../README.md`
+* `data/1a3a_1_A.a3m`, `data/1xcr_1_A.a3m`, `data/5ahw_1_A.a3m` are MSAs distributed with trRosetta, used in `contact_prediction.ipynb`
+* `data/P62593.fasta` is introduced and used in `sup_variant_prediction.ipynb`
+* Example MSAs genereated in the same way as the MSAs used for MSA Transformer pre-training:
+  - `data/UniRef50_E9K9Y4.a3m`, `data/UniRef50_UPI0003108055.a3m`, `data/UniRef50_UPI0003674933.a3m`, from the same sequences as trRosetta:
+  `data/hhblits_uniclust_2017_10_1a3a_1_A.a3m`, `data/hhblits_uniclust_2017_10_1xcr_1_A.a3m`, `data/hhblits_uniclust_2017_10_5ahw_1_A.a3m`.
+  - Generated with: `hhblits -i UniRef50_$id.fas -oa3m UniRef50_$id.a3m -n 3 -d /uniclust30_2017_10/uniclust30_2017_10`.
+* `esm2_infer_fairscale_fsdp_cpu_offloading.py` shows how to load the ESM-2 15B model with Fairscale's FSDP's CPU offloading capability on a single GPU