Initial commit

.gitattributes

@@ -0,0 +1,35 @@
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz 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

Dyna-1/LICENSE.txt

@@ -0,0 +1,73 @@
+Non-Commercial License Agreement.
+
+Copyright (c) 2025 Brandeis University, Gina El Nesr, Hannah Wayment-Steele.
+
+IMPORTANT: PLEASE READ THIS NON-COMMERCIAL LICENSE AGREEMENT ("AGREEMENT") CAREFULLY BEFORE USING THE SOFTWARE. BY DOWNLOADING, INSTALLING, OR USING THE SOFTWARE, YOU (THE "LICENSEE") AGREE TO BE BOUND BY THE TERMS AND CONDITIONS OF THIS AGREEMENT. IF YOU DO NOT AGREE TO THE TERMS AND CONDITIONS OF THIS AGREEMENT, DO NOT DOWNLOAD, INSTALL, OR USE THE SOFTWARE.
+
+WHEREAS, Brandeis University (the "Licensor") are the licensor of the Dyna-1 model (the "Model");
+
+DEFINITIONS, in addition to other terms defined elsewhere have the following meanings:
+
+THE MODEL (the "Model" or the "Software") means the Dyna-1 released model as a combination of the Dyna-1 Model Code and the Dyna-1 Model Weights and all software, algorithms, machine-learning model code, inference-enabling code, training-enabling code, fine-tuning enabling code, and other elements of the foregoing distributed and made available by the Licensor on the GitHub Page (https://github.com/WaymentSteeleLab/Dyna-1) or other code expressly listed within the GitHub Page including but not limited to Google Colab Notebooks, each as may be updated and amended from time to time, whether in source or object form. Any instances of any source code or weights that depend explicitly on ESM-3 Models are licensed under the EvolutionaryScale Cambrian Non-Commercial License Agreement and subject to the terms that are more restrictive.
+
+COMMERCIAL USE means any use or activity intended for or directed towards commercial advantage or monetary compensation, including without limitation, the development of any product, method, or service intended to be sold or made available for a fee. 
+
+COMMERCIAL ENTITY means any individual, entity, or organization involved in any capacity in Commercial Use. For the purposes of this Agreement, references to a Commercial Entity expressly exclude any universities, non-profit organizations, not-for-profit entities, research institutions, educational bodies, and government bodies. 
+
+NON-COMMERCIAL USE means use not leading to or directed towards commercial advantage or monetary compensations, or the facilitation of development of any product, method, or service to be sold or made available for a fee. Non-commercial use excludes any activities conducted by commercial entities, even if those activities themselves do not directly generate revenue. Examples of non-commercial use include academic research, personal projects, and use by non-profit organizations.
+
+DERIVATIVE WORKS means any work, in source or object form, that is based on or derived from the Model.
+
+OUTPUT DERIVATIVES means any outputs resulting from use of the Model. This includes, but is not limited to, any predictions that result from the Model.
+
+NOW, THEREFORE, the Licensor and the Licensee hereto agree as follows:
+
+GRANT OF LICENSE.
+
+    Non-Commercial Use. The Model is provided under the MIT License for non-commercial use. The MIT License is included below. 
+
+    Commercial Use. If the Licensee is or represents a Commercial Entity and wishes to use the Model for Commercial Use, the Licensee must first contact the Licensor for a commercial license. Examples of Commercial Use include but are not limited to the:
+
+        Selling the Model, either as is or as part of a larger software package;
+
+        Using the Model as an interface or component of a service where customers pay to use it;
+
+        Providing the Model, either as is or as part of a larger software package;
+
+        Incorporating the Model, in part or in whole, into a any commercial operation or workflow that may lead to commercial use of a product, method, or service;
+
+        Use of any derivative works or output derivatives towards any commercial operation or workflow that may lead to commercial use.
+
+    RESTRICTIONS. The Licensee shall not:
+
+        Use, modify, or distribute the Model for commercial purposes without receiving a commercial license.
+
+        Sublicense, assign, or transfer the Model without Licensor's prior written consent.
+
+        Use the Model in any manner that infringes upon any third-party rights.
+
+        Use the Model in any way that is not expressly permitted by this Agreement.
+
+OWNERSHIP. Licensor retains all right, title, and interest in and to the Model, including all intellectual property rights therein. No rights are granted to Licensee other than those expressly set forth in this Agreement.
+
+TERM AND TERMINATION. This Agreement shall commence upon Licensee's initial use of the Model and continue unless terminated by either party upon providing written notice to the other party. Upon termination, Licensee shall immediately cease using the Model and destroy all copies thereof.
+
+LIABILITY. Licensor shall not be liable for any direct, indirect, incidental, special, consequential, or exemplary damages, including but not limited to damages for loss of profits, goodwill, use, data, or other intangible losses, resulting from the use of the Model.
+
+GOVERNING LAW. This Agreement shall be governed by and construed in accordance with the laws of the United States of America.
+
+MIT LICENSE
+
+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.
+
+By using the Software, Licensee acknowledges that they have read this Agreement, understood it, and agreed to be bound by its terms and conditions.
+
+CONTACT INFORMATION. 
+For commercial license inquiries, please contact: Brandeis University Office of Technology Licensing at [email protected]
+For technical inquiries, please contact: [email protected], [email protected]
+
+[END OF AGREEMENT]

Dyna-1/README.md

@@ -0,0 +1,115 @@
+# Dyna-1
+[![Requires Python 3.10+](https://img.shields.io/badge/Python-3.10+-blue.svg?logo=python&logoColor=white)](https://python.org/downloads)
+[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/WaymentSteeleLab/Dyna-1/blob/main/colab/Dyna_1.ipynb)
+
+![image](assets/dyna1.png)
+
+Dyna-1 is a model introduced in our paper, ["Learning millisecond protein dynamics from what is missing in NMR spectra"](https://www.biorxiv.org/content/10.1101/2025.03.19.642801v1).
+
+Given a sequence and/or structure, Dyna-1 will predict the probability that each residue experiences micro-millisecond motions.
+
+Dyna-1 was achieved using the `esm3-sm-open-v1` weights from ESM-3. Inference with this model is subject to the EvolutionaryScale Cambrian Non-Commercial License Agreement of the ESM-3 Model and requires read permission of the weights found [here](https://huggingface.co/EvolutionaryScale/esm3-sm-open-v1). We also make available an alternate version of Dyna-1 that uses ESM-2 embeddings; use of this model is subject to a Non-Commercial License Agreement. 
+
+To make Dyna-1 readily accessible for research purposes, we also provide a [Google Colab](https://colab.research.google.com/github/WaymentSteeleLab/Dyna-1/blob/main/colab/Dyna_1.ipynb).
+
+We provide the curated datasets used to evaluate Dyna-1: 133 curated R1/R2/NOE datasets "RelaxDB" and 10 relaxation-dispersion Carr-Purcell-Meiboom-Gill datasets "RelaxDB-CPMG". 
+
+In this repository we provide:
+* [Installation](#installation)
+* [Inference code for Dyna-1](#inference)
+* [PyMol visualization](#visualization)
+* [RelaxDB and RelaxDB-CPMG datasets](#datasets)
+* [Training code for Dyna-1](#training)
+* [Citation](#citation)
+* [Acknowledgements](#acknowledgement)
+
+If you have any questions not covered here, please create an issue and we will get back to you ASAP.
+
+# Installation 
+To run the scripts in this repository, we recommend using a conda environment. First, clone this repository. Then navigate to the root directory and run:
+```
+conda create -n dyna1 python=3.11
+conda activate dyna1
+```
+This package requires PyTorch, ideally with GPU support. For more information, follow instructions from https://pytorch.org/get-started/locally/ for your system and CUDA version. We used PyTorch 2.5.0 with CUDA 12.4. To install all of the requirements:
+```
+pip install -r requirements.txt
+```
+Then, download the model weights and upload them to the `model/weights` folder. The weights can be found on 🤗HuggingFace at <a href='https://huggingface.co/gelnesr/Dyna-1'>gelnesr/Dyna-1</a>. More information on how to download them can be found <a href='https://github.com/gelnesr/Dyna-1-public/blob/main/model/weights/README.md'>here</a>. 
+
+# Inference
+
+The best-performing Dyna-1 is based on ESM-3. To run this version, you will have to request access to the ESM-3 `esm3-sm-open-v1` weights at HuggingFace [here](https://huggingface.co/EvolutionaryScale/esm3-sm-open-v1). Follow the steps to agree to their License terms and receive your access token to the model weights. 
+
+> [!NOTE]
+> If this is your first time requesting access to the ESM-3 weights, you may need to set up your access token. For more information on how to set up an SSH token, please consult <a href='https://huggingface.co/docs/hub/en/security-git-ssh'>this</a> tutorial. Alternatively, you can use the huggingface login prompt, which will prompt you for the access token each time you re-instantiate Dyna-1. This can be configured by adding the following code to the inference script: `from huggingface_hub import login; login()` 
+
+To run inference using our best-performing model, run:
+
+```
+python dyna1.py --pdb <PDB CODE or PATH> --chain <CHAIN> --name <NAME> --use_pdb_seq --write_to_pdb
+```
+
+We provide three options for running inference: sequence and structure input (best performance!), sequence only, or structure only. Additionally, we make it possible to test different sequences for the same backbone. Examples on how to run each of these modes can be found in the `scripts/` folder. 
+
+To output the probabilities to the input structure, make sure to pass the `--write_to_pdb` flag.
+
+Alternatively, we also provide a version of Dyna-1 based on ESM-2. To run inference using this version of the model, run:
+
+```
+python dyna1-esm2.py --sequence <SEQUENCE> --name <NAME>
+```
+
+# Visualization
+
+We visualize probabilities of exchange on protein structures using [PyMol](https://www.pymol.org). To re-create the putty visualization on your protein, import the pdb file into PyMol and copy-paste the following commands into the PyMol command line:
+
+```
+cartoon putty; set cartoon_putty_transform, 6; set cartoon_putty_radius, 0.25; set cartoon_putty_range, 0.1; set cartoon_putty_scale_max, 10
+```
+
+Annotated:
+```
+cartoon putty
+set cartoon_putty_transform, 6 #scaled linear transformation
+set cartoon_putty_radius, 0.25 # min radius for p=0
+set cartoon_putty_range, 0.1 # min_radius / max_radius, sets max_radius=2.5
+set cartoon_putty_scale_max, 10 #max_radius / min_radius
+```
+
+# Datasets
+
+*RelaxDB* contains 133 R1/R2/NOE datasets curated from the [BMRB](https://bmrb.io/) and from literature.
+
+*RelaxDB-CPMG* contains motion labels derived from 10 CPMG relaxation-dispersion datasets curated from literature.
+
+These datasets are made available on 🤗HuggingFace at <a href='https://huggingface.co/datasets/gelnesr/RelaxDB'>datasets/gelnesr/RelaxDB</a>.
+
+In this repo, you can find: 
+- data formatted for input into Dyna-1 is in `data/RelaxDB_pkls_22jan2025.zip`
+- datasets in json format is in `data/RelaxDB_datasets/`
+- demo notebooks for demo notebooks for visualizing and using datasets to evaluate model outputs in `analysis/`
+
+# Training
+
+Training code will be made available upon journal publication.
+
+# Citation
+
+If you are using our code, datasets, or model, please use the following citation:
+```bibtex
+@article {Dyna-1,
+    author = {Wayment-Steele, Hannah K. and El Nesr, Gina and Hettiarachchi, Ramith and Kariyawasam, Hasindu and Ovchinnikov, Sergey and Kern, Dorothee},
+    title = {Learning millisecond protein dynamics from what is missing in NMR spectra},
+    year = {2025},
+    doi = {10.1101/2025.03.19.642801},
+    journal = {bioRxiv}
+}
+```
+# Acknowledgements
+
+We would like to acknowledge the Evolutionary Scale Team for their contributions to the field with ESM-3. The code in `esm` is imported from `evolutionaryscale/esm` with all modifications identified and includes the associated LICENSE terms for the ESM-3 model.
+
+We would also like to acknowledge the FAIR Team for their contributions to the field with ESM-2. The ESM-2 model is called using the HuggingFace API call.
+
+We thank Katie Henzler-Wildman, Magnus Wolf-Watz, Elan Eisenmesser, J. Patrick Loria, Marcellus Ubbelink, George Lisi, Sam Butcher, and Nicolas Doucet for sharing data. We thank Martin Stone for sharing the Indiana Dynamics Database data his group curated in 2000.

Dyna-1/configs/af2.yml

@@ -0,0 +1,31 @@
+model_name: 'af2'
+train:
+  epochs: 500
+  batchsize: 4
+  accum_steps: 4
+  lr: 0.00001
+  dropout: 0.1
+  num_workers: 1
+model:
+  hidden_size: 128
+  res_count: 32
+  length: 400
+  nheads: 8
+  nlayers: 12
+dir:
+  save_dir: '/path/to/save/dir'
+data:
+  pair_rep: None
+  sample_clusters: True
+  relaxdb: 
+    split: 'relaxdb'
+    crop_len: 400
+    type: 'rex'
+  cpmg: 
+    split: 'relaxdb-cpmg'
+    crop_len: 400
+    type: 'cpmg'
+wandb:
+  project: "New-Project"
+  team: 'my-project-team'
+  dir: '/wandb/log/dir'

Dyna-1/configs/baseline.yml

@@ -0,0 +1,28 @@
+model_name: 'baseline'
+train:
+  epochs: 500
+  batchsize: 16
+  accum_steps: 1
+  lr: 0.000001
+  dropout: 0.1
+  num_workers: 2
+model:
+  nheads: 6
+  nlayers: 10
+dir:
+  save_dir: '/path/to/save/dir'
+data:
+  pair_rep: None
+  sample_clusters: True
+  relaxdb: 
+    split: 'relaxdb'
+    crop_len: 367
+    type: 'rex'
+  cpmg: 
+    split: 'relaxdb-cpmg'
+    crop_len: 367
+    type: 'cpmg'
+wandb:
+  project: "New-Project"
+  team: 'my-project-team'
+  dir: '/wandb/log/dir'

Dyna-1/configs/esm2.yml

@@ -0,0 +1,28 @@
+model_name: 'esm2_t30_150M_UR50D'
+train:
+  epochs: 500
+  batchsize: 16
+  accum_steps: 1
+  lr: 0.000001
+  dropout: 0.1
+  num_workers: 2
+model:
+  nheads: 8
+  nlayers: 12
+dir:
+  save_dir: '/path/to/save/dir'
+data:
+  pair_rep: None
+  sample_clusters: True
+  relaxdb: 
+    split: 'relaxdb'
+    crop_len: 367
+    type: 'rex'
+  cpmg: 
+    split: 'relaxdb-cpmg'
+    crop_len: 367
+    type: 'cpmg'
+wandb:
+  project: "New-Project"
+  team: 'my-project-team'
+  dir: '/wandb/log/dir'

Dyna-1/configs/esm3.yml

@@ -0,0 +1,28 @@
+model_name: 'esm3'
+train:
+  epochs: 500
+  batchsize: 16
+  accum_steps: 1
+  lr: 0.000001
+  dropout: 0.1
+  num_workers: 2
+model:
+  nheads: 6
+  nlayers: 12
+dir:
+  save_dir: '/path/to/save/dir'
+data:
+  pair_rep: None
+  sample_clusters: True
+  relaxdb: 
+    split: 'relaxdb'
+    crop_len: 367
+    type: 'rex'
+  cpmg: 
+    split: 'relaxdb-cpmg'
+    crop_len: 367
+    type: 'cpmg'
+wandb:
+  project: "New-Project"
+  team: 'my-project-team'
+  dir: '/wandb/log/dir'

Dyna-1/data/dataloader.py

@@ -0,0 +1,341 @@
+import numpy as np
+import warnings
+import random
+import torch
+import os
+import pickle
+import data.vocab as vocab
+import pandas as pd
+from typing import Tuple, List, Any
+from esm.sdk.api import ESMProtein
+from transformers import AutoTokenizer
+import utils
+
+class DynaData(torch.utils.data.Dataset):
+    """
+    For each protein, we use a pkl file that contains:
+        seq: The domain sequence, stored as an L-length string
+        assns: string containing labels of dynamics type
+    """
+
+    def __init__(
+        self,
+        split,
+        type = 'missing',
+        sample_clusters = False,
+        cluster_file = None,
+        crop_len = 300,
+        missing_only = False,
+        rex_only = False,
+        unsuppressed = False,
+        method = None,
+        pair_rep = None,
+        return_dssp = False
+    ):
+        super().__init__()
+
+        self.return_dssp = return_dssp
+        self.crop_len = crop_len
+        self.sample_clusters = sample_clusters
+        self.label_tokenizer = vocab.label_tokenizer(type = type, 
+                                                     missing_only = missing_only,
+                                                     rex_only = rex_only,
+                                                     unsuppressed = unsuppressed)
+        
+        # tokenization is the same for all ESM models, use the lightest one 
+        self.tokenizer = AutoTokenizer.from_pretrained(f"facebook/esm2_t6_8M_UR50D")
+        self.proline = self.tokenizer.get_vocab()['P']
+  
+        self.method = method[0]
+        self.model = method[1]
+
+        if isinstance(split, str):
+            split = [split]
+                
+        self.all_names, self.names = [], []
+
+        # read in all pdb names
+        for fil in split:
+            filename = f'data/split_files/{fil}.txt'
+            with open(filename,'r') as f:
+                self.all_names.extend(f.read().splitlines())
+
+        # set up cluster sampling
+        if self.sample_clusters:
+            self.cluster_info = pd.read_csv(f'data/{cluster_file}.tsv', sep='\t')
+            self.cluster_info['cluster'] = self.cluster_info.apply(lambda row: row['cluster'], axis=1)
+            for nm in self.all_names:
+                subset = self.cluster_info.loc[self.cluster_info.entry_ID==nm]
+                if len(subset) == 0:
+                    print('NO!', nm)
+                cluster_ind = subset['cluster'].iloc[0]
+                if cluster_ind not in self.names:
+                    self.names.append(cluster_ind)
+        else:
+            self.names = self.all_names
+
+        self.pair_rep_dir = pair_rep
+
+    def __len__(self):
+        return len(self.names)
+
+    def __baseline_get_item__(self, name, obj, crop_start):
+        if crop_start > -1:
+            sequence = obj['sequence'][crop_start:crop_start+self.crop_len]
+        else:
+            sequence = obj['sequence'][:self.crop_len]
+        
+        sequence_tokens = self.tokenizer.encode(sequence, 
+                                           add_special_tokens=False,
+                                           padding='max_length',
+                                           max_length=self.crop_len,
+                                           return_tensors='np').squeeze()
+        
+        # set mask to 1 for length of seq, padded tokens then are 0
+        eval_mask = np.zeros_like(sequence_tokens)
+        eval_mask[:len(sequence)] = 1
+
+        sequence_id = sequence_tokens != 0
+
+        # mask prolines in eval
+        eval_mask[sequence_tokens==self.proline] = 0
+        
+        return sequence_tokens, sequence_id, eval_mask
+
+    def __af2_get_item__(self, name, obj, crop_start):
+        """
+        Prepares input for the AF2-pair model
+        """
+        
+        pair_rep = np.load(f"{self.pair_rep_dir}/{name}.npy")
+        labels, seq = obj['label'], obj['sequence']
+
+        if crop_start > -1:
+            pair_rep = pair_rep[crop_start:crop_start+self.crop_len,
+                crop_start:crop_start+self.crop_len, :]
+            labels = labels[crop_start:crop_start+self.crop_len]
+            seq = seq[crop_start:crop_start+self.crop_len]
+
+        eval_mask = np.zeros((pair_rep.shape[0],))
+        
+        prolines = [i for i, aa in enumerate(seq) if aa == 'P']
+        eval_mask[:len(labels)] = 1
+
+        sequence_id = eval_mask != 0
+        eval_mask[prolines] = 0
+        x = pair_rep.shape[0]
+
+        eval_mask = np.pad(eval_mask, (0, self.crop_len - len(eval_mask)), mode='constant')
+        sequence_id = np.pad(sequence_id, (0, self.crop_len - len(sequence_id)), mode='constant')
+        if x < self.crop_len:
+            pair_rep = np.pad(pair_rep, ((0, self.crop_len - x), (0, self.crop_len - x), (0, 0)), mode='constant', constant_values=0)
+            
+        return pair_rep, sequence_id, eval_mask
+
+    def __esm3_get_item__(self, name, crop_start, data_path = 'esm3_data/'):
+        """
+        Prepares input for the ESM3 model
+        """
+        pkl_fname = os.path.join(data_path, f"{name}.pkl")
+        
+        try:
+            with open(pkl_fname, "rb") as f:
+                esm_data = pickle.load(f)
+        except:
+            print(f'writing pkl for {name} {crop_start}')
+            pdb_path = f'pdbs/{name}.pdb'
+            protein = ESMProtein.from_pdb(pdb_path)
+            
+            self.model.eval()
+            encoder = self.model.model.encode(protein)
+            self.model.train()
+
+            seq = encoder.sequence.cpu().detach()[1:-1][:700]
+            struct = encoder.structure.cpu().detach()[1:-1][:700]
+
+            sequence_tokens = np.full(700, 1, dtype=np.int32) ## sequence pad token is 1
+            structure_tokens = np.full(700, 4099, dtype=np.int32) ## structure pad token is 4099
+            
+            sequence_tokens[:len(seq)] = seq
+            structure_tokens[:len(struct)] = struct
+
+            sequence_id = sequence_tokens != 1
+            
+            obj ={'name': name, 'len': len(seq), 'seq_tokens': sequence_tokens, 
+                  'struct_tokens': structure_tokens, 'sequence_id': sequence_id}
+
+            with open(pkl_fname, 'wb') as f:
+                pickle.dump(obj, f)
+        
+        with open(pkl_fname, "rb") as f:
+            esm_data = pickle.load(f)
+
+        if crop_start > -1:
+            sequence_tokens = esm_data['seq_tokens'][crop_start:crop_start+self.crop_len]
+            structure_tokens = esm_data['struct_tokens'][crop_start:crop_start+self.crop_len]
+            sequence_id = esm_data['sequence_id'][crop_start:crop_start+self.crop_len]
+        else:
+            sequence_tokens = esm_data['seq_tokens'][:self.crop_len]
+            structure_tokens = esm_data['struct_tokens'][:self.crop_len]
+            sequence_id = esm_data['sequence_id'][:self.crop_len]
+
+        eval_mask = np.zeros_like(sequence_tokens)
+        eval_mask[:esm_data['len']] = 1
+        eval_mask[sequence_tokens==self.proline] = 0
+
+        return sequence_tokens, structure_tokens, sequence_id, eval_mask
+
+    def __esm2_get_item__(self, obj, crop_start):
+        """
+        Prepares input for the ESM2 model
+        """
+        sequence = obj['sequence'].replace(' ','')
+        if crop_start > -1:
+            sequence = sequence[crop_start:crop_start+self.crop_len]
+
+        sequence_tokens = self.tokenizer.encode(sequence, 
+                                           add_special_tokens=False,
+                                           padding='max_length',
+                                           max_length=self.crop_len,
+                                           return_tensors='np').squeeze()
+        
+        # Set mask to 1 for length of sequence, padded tokens then are 0
+        eval_mask = np.zeros_like(sequence_tokens)
+        eval_mask[:len(sequence)] = 1
+        sequence_id = eval_mask != 0
+        eval_mask[sequence_tokens==self.proline] = 0
+        
+        return sequence_tokens, sequence_id, eval_mask
+    
+    def __get_dssp__(self, name, crop_start):
+        """
+        Prepares DSSP information for a given sequence
+        """
+        try:
+            dssp_csv = pd.read_csv('data/dssp.csv')
+            entry = dssp_csv.loc[dssp_csv.PDB == str(name)].iloc[0]
+        except:
+            entry = {}
+            entry['DSSP'] = utils.calc_dssp(f'pdbs/{name}.pdb')
+
+        assert len(entry) > 0
+        if crop_start ==-1:
+            dssp_data = entry['DSSP'].replace(' ','')[:self.crop_len]
+        else:
+            dssp_data = entry['DSSP'].replace(' ','')[crop_start:crop_start + self.crop_len]
+
+        dssp = np.zeros(self.crop_len)
+        inds = [i for i, char in enumerate(dssp_data) if char=='C']
+        dssp[inds] = 1.0
+        
+        inds = [i for i, char in enumerate(dssp_data) if char=='H']
+        dssp[inds] = 2.0
+
+        return dssp
+    
+    def __getitem__(self, idx):
+        """
+        Returns a dict with the appropriate entries for each model
+        """
+        exists = -1
+        while exists == -1:
+            name = self.names[idx]
+            if self.sample_clusters:
+                roptions = list(self.cluster_info.loc[self.cluster_info.cluster==name]['entry_ID'].values)
+                options = [opt for opt in roptions if opt in self.all_names]
+                name = random.choice(options)
+            pkl_fname = f"data/mBMRB_data/{name}.pkl"
+            
+            try:
+                with open(pkl_fname, "rb") as f:
+                    obj = pickle.load(f)
+                exists = 1
+            except:
+                print(f'{pkl_fname} not found')
+
+        assns = obj['label']
+        assns = vocab.mask_termini(assns)
+
+        crop_start = -1
+        if len(assns) > self.crop_len:
+            crop_start = np.random.choice(range(0, len(assns)-self.crop_len))   
+            assns = assns[crop_start:crop_start + self.crop_len]
+
+        labels = self.label_tokenizer.convert_tokens_to_ids(assns, pad_to_length=self.crop_len)
+        labels = np.asarray(labels, np.int64)
+
+        dssp = None
+        if self.return_dssp:
+            dssp = self.__get_dssp__(name, crop_start)
+        if 'esm3' in self.method:
+            sequence, structure, sequence_id, eval_mask = self.__esm3_get_item__(name, crop_start)
+        elif 'esm2' in self.method:
+            sequence, sequence_id, eval_mask = self.__esm2_get_item__(obj, crop_start)
+        elif 'af2' in self.method:
+            pair_rep, sequence_id, eval_mask = self.__af2_get_item__(name, obj, crop_start)
+        elif 'baseline' in self.method:
+            sequence, sequence_id, eval_mask = self.__baseline_get_item__(name, obj, crop_start)
+        
+        # Mask termini for eval. A -1 label corresponds to indices that are getting masked in vocabs
+        eval_mask[labels==-1] = 0
+
+        if 'esm2' in self.method:
+            return sequence, sequence_id, eval_mask, labels, name, dssp
+        elif 'esm3' in self.method:
+            return sequence, structure, sequence_id, eval_mask, labels, name, dssp
+        elif 'af2' in self.method:
+            return pair_rep, labels, sequence_id, eval_mask, name, dssp
+        elif 'baseline' in self.method:
+            return sequence, labels, sequence_id, eval_mask, name, dssp
+
+    def __collate_fn__(self, batch: List[Tuple[Any, ...]]):
+
+        if 'baseline' in self.method:
+            seqs, labels, sequence_id, eval_mask, names, dssp = tuple(zip(*batch))
+            seqs = torch.tensor(np.array(seqs))
+
+            labels = torch.from_numpy(np.array(labels)).float()
+            eval_mask = torch.from_numpy(np.array(eval_mask))    
+            if self.return_dssp:
+                dssp = torch.from_numpy(np.array(dssp))
+            sequence_id = torch.from_numpy(np.array(sequence_id))
+
+            output = {'names': names, 'seqs': seqs, 'seq_id': sequence_id, 
+                        'targets': labels, 'eval_mask': eval_mask, 'dssp': dssp}
+            return output
+        
+        elif 'af2' in self.method:
+            pair_reps, labels, sequence_id, eval_mask, names, dssp = tuple(zip(*batch))
+
+            pair_reps = torch.from_numpy(np.array(pair_reps, dtype=np.float64))
+            labels = torch.from_numpy(np.array(labels)).float()
+            eval_mask = torch.from_numpy(np.array(eval_mask))
+            sequence_id = torch.from_numpy(np.array(sequence_id, dtype=bool))
+
+            if self.return_dssp:
+                dssp = torch.from_numpy(np.array(dssp))
+            output = {'names': names, 'pair_reps': pair_reps, 'targets': labels, "seq_id": sequence_id, 
+                      'eval_mask': eval_mask, 'dssp': dssp}
+            return output
+        
+        elif 'esm2' in self.method:
+            seqs, sequence_id, eval_mask, label, names, dssp = tuple(zip(*batch))
+            seqs = torch.from_numpy(np.array(seqs))
+            structs = None
+            sequence_id = torch.from_numpy(np.array(sequence_id))
+
+        elif 'esm3' in self.method:
+            seqs, structs, sequence_id, eval_mask, label, names, dssp = tuple(zip(*batch))
+            seqs = torch.from_numpy(np.array(seqs))
+            structs = torch.from_numpy(np.array(structs))
+            sequence_id = torch.from_numpy(np.array(sequence_id))
+
+        eval_mask = torch.from_numpy(np.array(eval_mask))
+        label = torch.from_numpy(np.array(label)).float()
+        if self.return_dssp:
+            dssp = torch.from_numpy(np.array(dssp))
+
+        output = {'seqs': seqs, "structs": structs, "seq_id": sequence_id, "eval_mask": eval_mask, 
+                     'targets': label, 'names': names, 'dssp': dssp}
+
+        return output

Dyna-1/data/vocab.py

@@ -0,0 +1,104 @@
+from collections import OrderedDict
+from typing import List
+
+VOCAB_MISSING = OrderedDict([
+    ('t',-1),   # no data bc disorderd terminus
+    ('x',-1),   # no data, R1/R2/NOE not reported
+    ('p',0),    # proline, not evaluated
+    ('A', 0),   # nothing
+    ('v',0),    # fast motion
+    ('.',1),    # missing
+    ('b',0),    # both fast and slow
+    ('^',0)     # rex
+])
+
+VOCAB_REX = OrderedDict([
+    ('t',-1),   # no data bc disorderd terminus
+    ('x',-1),   # no data, R1/R2/NOE not reported
+    ('p',0),    # proline, not evaluated
+    ('A', 0),   # nothing
+    ('v',0),    # fast motion
+    ('.',1),    # missing
+    ('b',1),    # both fast and slow
+    ('^',1)     # rex
+])
+
+VOCAB_CPMG = OrderedDict([
+    ('t',-1),   # no data bc disordered terminus
+    ('P',0),    # proline, not evaluated
+    ('N',-1),   # no data, assned but CPMG not reported
+    ('A', 0),   # nothing
+    ('.',1),    # missing
+    ('X',1),    # exchange from Rex definition
+    ('Y',0)     # exchange from unsuppressed R2
+])
+
+def mask_termini(seq):
+    """
+    Mask the termini of a sequence
+    """
+    seq = seq.lstrip('.p').rjust(len(seq), 't')
+    seq = seq.rstrip('.p').ljust(len(seq), 't')
+    return seq
+
+class label_tokenizer():
+    def __init__(self, 
+                 type = 'missing', 
+                 missing_only = False, 
+                 rex_only = False, 
+                 unsuppressed = False):
+        """
+        Tokenize the data labeling for BMRB, REX, and CPMG
+        
+        Args:
+            type: (str) which type of experiment
+            missing_only: (bool) only return residues with missing peaks
+            rex_only: (bool) only return residues with Rex
+            unsuppressed: (bool) return residues with unsuppressed Rex
+        """
+        if type == 'missing':
+            self.vocab = VOCAB_MISSING.copy()
+        elif type == 'rex':
+            self.vocab = VOCAB_REX.copy()
+            if missing_only:
+                self.vocab['b'] = 0
+                self.vocab['^'] = 0
+        elif type == 'cpmg':
+            self.vocab = VOCAB_CPMG.copy()
+            if missing_only:
+                self.vocab['X'] = 0
+            if unsuppressed:
+                self.vocab['Y'] = 1
+        if rex_only:
+            self.vocab['.'] = -1
+        self.tokens = list(self.vocab.keys())
+
+    @property
+    def vocab_size(self) -> int:
+        return len(self.vocab)
+
+    def convert_token_to_id(self, token: str) -> int:
+        """ Converts a token (str/unicode) in an id using the vocab. """
+        try:
+            return self.vocab[token]
+        except KeyError:
+            raise KeyError(f"Unrecognized token: '{token}'")
+
+    def convert_tokens_to_ids(self, tokens: List[str], pad_to_length = None) -> List[int]:
+        """Converts a list of tokens (str/unicode) into a list of ids (int) using the vocab. """
+
+        if pad_to_length is None:
+          return [self.convert_token_to_id(token) for token in tokens]
+        else:
+          return [self.convert_token_to_id(token) for token in tokens] + [0] * (pad_to_length - len(tokens))
+
+    def convert_id_to_token(self, index: int) -> str:
+        """Converts an index (integer) in a token (string/unicode) using the vocab."""
+        try:
+            return self.tokens[index]
+        except IndexError:
+            raise IndexError(f"Unrecognized index: '{index}'")
+
+    def convert_ids_to_tokens(self, indices: List[int]) -> List[str]:
+        """Converts a list of indices (integer) into a list of tokens (string/unicode) using the vocab."""
+        return [self.convert_id_to_token(id_) for id_ in indices]

Dyna-1/esm/__init__.py

@@ -0,0 +1,2 @@
+__version__ = "3.1.1"
+

Dyna-1/esm/layers/attention.py

@@ -0,0 +1,76 @@
+import functools
+
+import einops
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from esm.layers.rotary import RotaryEmbedding
+
+
+class MultiHeadAttention(nn.Module):
+    def __init__(
+        self, d_model: int, n_heads: int, bias: bool = False, qk_layernorm: bool = True
+    ):
+        super().__init__()
+
+        self.d_model = d_model
+        self.n_heads = n_heads
+
+        self.d_head = self.d_model // self.n_heads
+        self.layernorm_qkv = nn.Sequential(
+            nn.LayerNorm(d_model), nn.Linear(d_model, d_model * 3, bias=bias)
+        )
+        self.out_proj = nn.Linear(d_model, d_model, bias=bias)
+
+        if qk_layernorm:
+            self.q_ln = nn.LayerNorm(d_model, bias=bias)
+            self.k_ln = nn.LayerNorm(d_model, bias=bias)
+        else:
+            self.q_ln = nn.Identity()
+            self.k_ln = nn.Identity()
+
+        self.rotary = RotaryEmbedding(d_model // n_heads)
+
+    def _apply_rotary(self, q: torch.Tensor, k: torch.Tensor):
+        q = q.unflatten(-1, (self.n_heads, self.d_head))
+        k = k.unflatten(-1, (self.n_heads, self.d_head))
+        q, k = self.rotary(q, k)
+        q = q.flatten(-2, -1)
+        k = k.flatten(-2, -1)
+        return q, k
+
+    def forward(self, x, seq_id):
+        qkv_BLD3 = self.layernorm_qkv(x)
+        query_BLD, key_BLD, value_BLD = torch.chunk(qkv_BLD3, 3, dim=-1)
+        query_BLD, key_BLD = (
+            self.q_ln(query_BLD).to(query_BLD.dtype),
+            self.k_ln(key_BLD).to(query_BLD.dtype),
+        )
+        query_BLD, key_BLD = self._apply_rotary(query_BLD, key_BLD)
+
+        n_heads = self.n_heads
+        reshaper = functools.partial(
+            einops.rearrange, pattern="b s (h d) -> b h s d", h=n_heads
+        )
+
+        query_BHLD, key_BHLD, value_BHLD = map(
+            reshaper, (query_BLD, key_BLD, value_BLD)
+        )
+
+        if seq_id is not None:
+            # Where True, enable participation in attention.
+            mask_BLL = seq_id.unsqueeze(-1) == seq_id.unsqueeze(-2)
+            mask_BHLL = mask_BLL.unsqueeze(1)
+
+            context_BHLD = F.scaled_dot_product_attention(
+                query_BHLD, key_BHLD, value_BHLD, mask_BHLL
+            )
+        else:
+            # Shortcut, if we don't use attention biases then torch
+            # will autoselect flashattention as the implementation
+            context_BHLD = F.scaled_dot_product_attention(
+                query_BHLD, key_BHLD, value_BHLD
+            )
+        context_BLD = einops.rearrange(context_BHLD, "b h s d -> b s (h d)")
+        return self.out_proj(context_BLD)

Dyna-1/esm/layers/blocks.py

@@ -0,0 +1,153 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from esm.layers.attention import MultiHeadAttention
+from esm.layers.geom_attention import (
+    GeometricReasoningOriginalImpl,
+)
+from esm.utils.structure.affine3d import Affine3D
+
+
+def swiglu_correction_fn(expansion_ratio: float, d_model: int) -> int:
+    # set hidden dimesion to nearest multiple of 256 after expansion ratio
+    return int(((expansion_ratio * d_model) + 255) // 256 * 256)
+
+
+class SwiGLU(nn.Module):
+    """
+    SwiGLU activation function as an nn.Module, allowing it to be used within nn.Sequential.
+    This module splits the input tensor along the last dimension and applies the SiLU (Swish)
+    activation function to the first half, then multiplies it by the second half.
+    """
+
+    def __init__(self):
+        super(SwiGLU, self).__init__()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x1, x2 = x.chunk(2, dim=-1)
+        return F.silu(x1) * x2
+
+
+def swiglu_ln_ffn(d_model: int, expansion_ratio: float, bias: bool):
+    return nn.Sequential(
+        nn.LayerNorm(d_model),
+        nn.Linear(
+            d_model, swiglu_correction_fn(expansion_ratio, d_model) * 2, bias=bias
+        ),
+        SwiGLU(),
+        nn.Linear(swiglu_correction_fn(expansion_ratio, d_model), d_model, bias=bias),
+    )
+
+
+def gelu_ln_ffn(d_model: int, expansion_ratio: float, bias: bool):
+    hidden_dim = int(expansion_ratio * d_model)
+    return nn.Sequential(
+        nn.LayerNorm(d_model),
+        nn.Linear(d_model, hidden_dim, bias=bias),
+        nn.GELU(),
+        nn.Linear(hidden_dim, d_model, bias=bias),
+    )
+
+
+class UnifiedTransformerBlock(nn.Module):
+    """
+    A unified transformer block that can optionally incorporate geometric attention.
+
+    This class defines a transformer block that can be configured to use geometric attention
+    alongside the standard multi-head attention mechanism. It is designed to be a flexible
+    component of transformer-based models, allowing for the integration of geometric reasoning.
+
+    Parameters
+    ----------
+    d_model : int
+        The dimensionality of the input and output features of the transformer block.
+    n_heads : int
+        The number of attention heads in the multi-head attention mechanism.
+    n_layers : int
+        The number of layers in the transformer block.
+    use_geom_attn : bool, optional
+        Whether to use geometric attention in addition to the standard multi-head attention. Defaults to False.
+    v_heads : int, optional
+        The number of heads to use for the geometric attention mechanism, if enabled. Must be specified if `use_geom_attn` is True.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        n_heads: int,
+        use_geom_attn: bool = False,
+        use_plain_attn: bool = True,
+        v_heads: int | None = None,
+        bias: bool = False,
+        expansion_ratio: float = 4.0,
+        residue_scaling_factor: float = 1,
+        mask_and_zero_frameless: bool = False,
+        qk_layernorm: bool = True,
+        ffn_type: str = "swiglu",  # swiglu | gelu
+    ):
+        super().__init__()
+        self.use_plain_attn = use_plain_attn
+        if self.use_plain_attn:
+            self.attn = MultiHeadAttention(
+                d_model, n_heads, bias, qk_layernorm=qk_layernorm
+            )
+        self.use_geom_attn = use_geom_attn
+        if self.use_geom_attn:
+            if v_heads is None:
+                raise ValueError("v_heads must be specified when use_geom_attn is True")
+            self.geom_attn = GeometricReasoningOriginalImpl(
+                c_s=d_model,
+                v_heads=v_heads,
+                bias=bias,
+                mask_and_zero_frameless=mask_and_zero_frameless,
+            )
+        if ffn_type == "swiglu":
+            self.ffn = swiglu_ln_ffn(d_model, expansion_ratio, bias)
+        elif ffn_type == "gelu":
+            self.ffn = gelu_ln_ffn(d_model, expansion_ratio, bias)
+        else:
+            raise ValueError(f"Unknown ffn_type: {ffn_type}")
+        self.scaling_factor = residue_scaling_factor
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        sequence_id: torch.Tensor,
+        frames: Affine3D,
+        frames_mask: torch.Tensor,
+        chain_id: torch.Tensor,
+    ) -> torch.Tensor:
+        """
+        Forward pass for the UnifiedTransformerBlock.
+
+        Parameters
+        ----------
+        x : torch.Tensor[float]
+            Input tensor to the transformer block, typically the output from the previous layer.
+        sequence_id : torch.Tensor[int]
+            Tensor containing sequence IDs for each element in the batch, used for attention masking.
+        frames : Affine3D
+            Affine3D containing geometric frame information for geometric attention.
+        frames_mask : torch.Tensor[bool]
+            Boolean mask tensor indicating valid frames for geometric attention.
+        chain_id : torch.Tensor[int]
+            Tensor containing chain IDs for each element, used for attention masking in geometric attention.
+
+        Returns
+        -------
+        torch.Tensor[float]
+            The output tensor after applying the transformer block operations.
+        """
+        if self.use_plain_attn:
+            r1 = self.attn(x, sequence_id)
+            x = x + r1 / self.scaling_factor
+
+        if self.use_geom_attn:
+            r2 = self.geom_attn(x, frames, frames_mask, sequence_id, chain_id)
+            x = x + r2 / self.scaling_factor
+
+        r3 = self.ffn(x) / self.scaling_factor
+        x = x + r3
+
+        return x

Dyna-1/esm/layers/codebook.py

@@ -0,0 +1,88 @@
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class EMACodebook(nn.Module):
+    def __init__(
+        self,
+        n_codes,
+        embedding_dim,
+        no_random_restart=True,
+        restart_thres=1.0,
+        ema_decay=0.99,
+    ):
+        super().__init__()
+        self.register_buffer("embeddings", torch.randn(n_codes, embedding_dim))
+        self.register_buffer("N", torch.zeros(n_codes))
+        self.register_buffer("z_avg", self.embeddings.data.clone())
+
+        self.n_codes = n_codes
+        self.embedding_dim = embedding_dim
+        self._need_init = True
+        self.no_random_restart = no_random_restart
+        self.restart_thres = restart_thres
+        self.freeze_codebook = False
+        self.ema_decay = ema_decay
+
+    def reset_parameters(self):
+        # For meta init
+        pass
+
+    def _tile(self, x):
+        d, ew = x.shape
+        if d < self.n_codes:
+            n_repeats = (self.n_codes + d - 1) // d
+            std = 0.01 / np.sqrt(ew)
+            x = x.repeat(n_repeats, 1)
+            x = x + torch.randn_like(x) * std
+        return x
+
+    def _init_embeddings(self, z):
+        # z: [b, t, c]
+        self._need_init = False
+        flat_inputs = z.view(-1, self.embedding_dim)
+        y = self._tile(flat_inputs)
+
+        y.shape[0]
+        _k_rand = y[torch.randperm(y.shape[0])][: self.n_codes]
+        if dist.is_initialized():
+            dist.broadcast(_k_rand, 0)
+        self.embeddings.data.copy_(_k_rand)
+        self.z_avg.data.copy_(_k_rand)
+        self.N.data.copy_(torch.ones(self.n_codes))
+
+    def forward(self, z):
+        # z: [b, t, c]
+        if self._need_init and self.training and not self.freeze_codebook:
+            self._init_embeddings(z)
+        # z is of shape [batch_size, sequence length, channels]
+        flat_inputs = z.view(-1, self.embedding_dim)
+        distances = (
+            (flat_inputs**2).sum(dim=1, keepdim=True)
+            - 2 * flat_inputs @ self.embeddings.t()
+            + (self.embeddings.t() ** 2).sum(dim=0, keepdim=True)
+        )  # [bt, c]
+
+        encoding_indices = torch.argmin(distances, dim=1)
+        encoding_indices = encoding_indices.view(*z.shape[:2])  # [b, t, ncode]
+
+        embeddings = F.embedding(encoding_indices, self.embeddings)  # [b, t, c]
+
+        commitment_loss = 0.25 * F.mse_loss(z, embeddings.detach())
+
+        # EMA codebook update
+        if self.training and not self.freeze_codebook:
+            assert False, "Not implemented"
+        embeddings_st = (embeddings - z).detach() + z
+
+        return embeddings_st, encoding_indices, commitment_loss
+
+    def dictionary_lookup(self, encodings):
+        embeddings = F.embedding(encodings, self.embeddings)
+        return embeddings
+
+    def soft_codebook_lookup(self, weights: torch.Tensor) -> torch.Tensor:
+        return weights @ self.embeddings

Dyna-1/esm/layers/ffn.py

@@ -0,0 +1,29 @@
+import torch.nn as nn
+import torch.nn.functional as F
+from torch import Tensor
+
+# NOT CURRENTLY USED
+
+
+class SwiGLU(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def forward(self, x: Tensor) -> Tensor:
+        x1, x2 = x.chunk(2, dim=-1)
+        hidden = F.silu(x1) * x2
+        return hidden
+
+
+class FFN(nn.Module):
+    def __init__(self, in_proj, activation, out_proj) -> None:
+        super().__init__()
+        self.in_proj = in_proj
+        self.activation = activation
+        self.out_proj = out_proj
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = self.in_proj(x)
+        x = self.activation(x)
+        x = self.out_proj(x)
+        return x

Dyna-1/esm/layers/geom_attention.py

@@ -0,0 +1,149 @@
+from math import sqrt
+
+import torch
+from einops import rearrange
+from torch import nn
+from torch.nn import functional as F
+
+
+class GeometricReasoningOriginalImpl(nn.Module):
+    def __init__(
+        self,
+        c_s: int,
+        v_heads: int,
+        num_vector_messages: int = 1,
+        mask_and_zero_frameless: bool = True,
+        divide_residual_by_depth: bool = False,
+        bias: bool = False,
+    ):
+        """Approximate implementation:
+
+        ATTN(A, v) := (softmax_j A_ij) v_j
+        make_rot_vectors(x) := R(i->g) Linear(x).reshape(..., 3)
+        make_vectors(x) := T(i->g) Linear(x).reshape(..., 3)
+
+        v <- make_rot_vectors(x)
+        q_dir, k_dir <- make_rot_vectors(x)
+        q_dist, k_dist <- make_vectors(x)
+
+        A_ij       <- dot(q_dir_i, k_dir_j) -||q_dist_i - k_dist_j||^2
+        x          <- x + Linear(T(g->i) ATTN(A, v))
+        """
+        super().__init__()
+        self.c_s = c_s
+        self.v_heads = v_heads
+        self.num_vector_messages = num_vector_messages
+        self.mask_and_zero_frameless = mask_and_zero_frameless
+
+        self.s_norm = nn.LayerNorm(c_s, bias=bias)
+        dim_proj = (
+            4 * self.v_heads * 3 + self.v_heads * 3 * self.num_vector_messages
+        )  # 2 x (q, k) * number of heads * (x, y, z) + number of heads * number of vector messages * (x, y, z)
+        self.proj = nn.Linear(c_s, dim_proj, bias=bias)
+        channels_out = self.v_heads * 3 * self.num_vector_messages
+        self.out_proj = nn.Linear(channels_out, c_s, bias=bias)
+
+        # The basic idea is for some attention heads to pay more or less attention to rotation versus distance,
+        # as well as to control the sharpness of the softmax (i.e., should this head only attend to those residues
+        # very nearby or should there be shallower dropoff in attention weight?)
+        self.distance_scale_per_head = nn.Parameter(torch.zeros((self.v_heads)))
+        self.rotation_scale_per_head = nn.Parameter(torch.zeros((self.v_heads)))
+
+    def forward(self, s, affine, affine_mask, sequence_id, chain_id):
+        if sequence_id is None:
+            sequence_id = torch.zeros_like(s[..., 0], dtype=torch.int64)
+        attn_bias = sequence_id.unsqueeze(-1) == sequence_id.unsqueeze(-2)
+        attn_bias = attn_bias.unsqueeze(1).float()
+        attn_bias = attn_bias.masked_fill(
+            ~affine_mask[:, None, None, :], torch.finfo(attn_bias.dtype).min
+        )
+        chain_id_mask = chain_id.unsqueeze(1) != chain_id.unsqueeze(2)
+        attn_bias = attn_bias.masked_fill(
+            chain_id_mask.unsqueeze(1), torch.finfo(s.dtype).min
+        )
+
+        ns = self.s_norm(s)
+        vec_rot, vec_dist = self.proj(ns).split(
+            [
+                self.v_heads * 2 * 3 + self.v_heads * 3 * self.num_vector_messages,
+                self.v_heads * 2 * 3,
+            ],
+            dim=-1,
+        )
+
+        # Rotate the queries and keys for the rotation term. We also rotate the values.
+        # NOTE(zeming, thayes): Values are only rotated, not translated. We may wish to change
+        # this in the future.
+        query_rot, key_rot, value = (
+            affine.rot[..., None]
+            .apply(rearrange(vec_rot, "... (h c) -> ... h c", c=3))
+            .split(
+                [self.v_heads, self.v_heads, self.v_heads * self.num_vector_messages],
+                dim=-2,
+            )
+        )
+
+        # Rotate and translate the queries and keys for the distance term
+        # NOTE(thayes): a simple speedup would be to apply all rotations together, then
+        # separately apply the translations.
+        query_dist, key_dist = (
+            affine[..., None]
+            .apply(rearrange(vec_dist, "... (h c) -> ... h c", c=3))
+            .chunk(2, dim=-2)
+        )
+
+        query_dist = rearrange(query_dist, "b s h d -> b h s 1 d")
+        key_dist = rearrange(key_dist, "b s h d -> b h 1 s d")
+        query_rot = rearrange(query_rot, "b s h d -> b h s d")
+        key_rot = rearrange(key_rot, "b s h d -> b h d s")
+        value = rearrange(
+            value, "b s (h m) d -> b h s (m d)", m=self.num_vector_messages
+        )
+
+        distance_term = (query_dist - key_dist).norm(dim=-1) / sqrt(3)
+        rotation_term = query_rot.matmul(key_rot) / sqrt(3)
+        distance_term_weight = rearrange(
+            F.softplus(self.distance_scale_per_head), "h -> h 1 1"
+        )
+        rotation_term_weight = rearrange(
+            F.softplus(self.rotation_scale_per_head), "h -> h 1 1"
+        )
+
+        attn_weight = (
+            rotation_term * rotation_term_weight - distance_term * distance_term_weight
+        )
+
+        if attn_bias is not None:
+            # we can re-use the attention bias from the transformer layers
+            # NOTE(thayes): This attention bias is expected to handle two things:
+            # 1. Masking attention on padding tokens
+            # 2. Masking cross sequence attention in the case of bin packing
+            s_q = attn_weight.size(2)
+            s_k = attn_weight.size(3)
+            _s_q = max(0, attn_bias.size(2) - s_q)
+            _s_k = max(0, attn_bias.size(3) - s_k)
+            attn_bias = attn_bias[:, :, _s_q:, _s_k:]
+            attn_weight = attn_weight + attn_bias
+
+        attn_weight = torch.softmax(attn_weight, dim=-1)
+
+        attn_out = attn_weight.matmul(value)
+
+        attn_out = (
+            affine.rot[..., None]
+            .invert()
+            .apply(
+                rearrange(
+                    attn_out, "b h s (m d) -> b s (h m) d", m=self.num_vector_messages
+                )
+            )
+        )
+
+        attn_out = rearrange(
+            attn_out, "b s (h m) d -> b s (h m d)", m=self.num_vector_messages
+        )
+        if self.mask_and_zero_frameless:
+            attn_out = attn_out.masked_fill(~affine_mask[..., None], 0.0)
+        s = self.out_proj(attn_out)
+
+        return s

Dyna-1/esm/layers/regression_head.py

@@ -0,0 +1,22 @@
+import torch.nn as nn
+
+
+def RegressionHead(
+    d_model: int, output_dim: int, hidden_dim: int | None = None
+) -> nn.Module:
+    """Single-hidden layer MLP for supervised output.
+
+    Args:
+        d_model: input dimension
+        output_dim: dimensionality of the output.
+        hidden_dim: optional dimension of hidden layer, defaults to d_model.
+    Returns:
+        output MLP module.
+    """
+    hidden_dim = hidden_dim if hidden_dim is not None else d_model
+    return nn.Sequential(
+        nn.Linear(d_model, hidden_dim),
+        nn.GELU(),
+        nn.LayerNorm(hidden_dim),
+        nn.Linear(hidden_dim, output_dim),
+    )

Dyna-1/esm/layers/rotary.py

@@ -0,0 +1,221 @@
+# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.
+#
+# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX
+# and OPT implementations in this library. It has been modified from its
+# original forms to accommodate minor architectural differences compared
+# to GPT-NeoX and OPT used by the Meta AI team that trained the model.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# NOTE: this implementation is from LLaMA 2:
+# https://huggingface.co/togethercomputer/LLaMA-2-7B-32K/blob/08639a72e17836184096ae6a7e2766f2a34c3e36/modeling_flash_llama.py#L114
+# Flash attention rotary implementation can be installed like so: `pip install git+https://github.com/HazyResearch/flash-attention.git#subdirectory=csrc/rotary`
+
+from typing import Tuple
+
+import torch
+from einops import rearrange, repeat
+
+
+def rotate_half(x, interleaved=False):
+    if not interleaved:
+        x1, x2 = x.chunk(2, dim=-1)
+        return torch.cat((-x2, x1), dim=-1)
+    else:
+        x1, x2 = x[..., ::2], x[..., 1::2]
+        return rearrange(
+            torch.stack((-x2, x1), dim=-1), "... d two -> ... (d two)", two=2
+        )
+
+
+def apply_rotary_emb_torch(x, cos, sin, interleaved=False, _inplace=False):
+    """
+    x: (batch_size, seqlen, nheads, headdim)
+    cos, sin: (seqlen, rotary_dim / 2)
+    """
+    ro_dim = cos.shape[-1] * 2
+    assert ro_dim <= x.shape[-1]
+    seqlen = x.size(1)
+    cos = cos[:seqlen]
+    sin = sin[:seqlen]
+    cos = repeat(cos, "s d -> s 1 (2 d)")
+    sin = repeat(sin, "s d -> s 1 (2 d)")
+    return torch.cat(
+        [
+            x[..., :ro_dim] * cos + rotate_half(x[..., :ro_dim], interleaved) * sin,
+            x[..., ro_dim:],
+        ],
+        dim=-1,
+    )
+
+
+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
+    If scale_base is not None, this implements XPos (Sun et al., https://arxiv.org/abs/2212.10554).
+    A recommended value for scale_base is 512: https://github.com/HazyResearch/flash-attention/issues/96
+    Reference: https://github.com/sunyt32/torchscale/blob/main/torchscale/component/xpos_relative_position.py
+    """
+
+    def __init__(
+        self,
+        dim: int,
+        base=10000.0,
+        interleaved=False,
+        scale_base=None,
+        scaling_factor=1.0,
+        pos_idx_in_fp32=True,
+        device=None,
+    ):
+        """
+        interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead
+            of 1st half and 2nd half (GPT-NeoX style).
+        pos_idx_in_fp32: if True, the position indices [0.0, ..., seqlen - 1] are in fp32,
+            otherwise they might be in lower precision.
+            This option was added because previously (before 2023-07-02), when we construct
+            the position indices, we use the dtype of self.inv_freq. In most cases this would
+            be fp32, but if the model is trained in pure bf16 (not mixed precision), then
+            self.inv_freq would be bf16, and the position indices are also in bf16.
+            Because of the limited precision of bf16 (e.g. 1995.0 is rounded to 2000.0), the
+            embeddings for some positions will coincide.
+            To maintain compatibility with models previously trained in pure bf16,
+            we add this option.
+        scaling_factor: RotaryEmbedding extended with linear scaling.
+        """
+        super().__init__()
+        self.dim = dim
+        self.base = float(base)
+        self.pos_idx_in_fp32 = pos_idx_in_fp32
+        # Generate and save the inverse frequency buffer (non trainable)
+        self.interleaved = interleaved
+        self.scale_base = scale_base
+        self.scaling_factor = scaling_factor
+        self.device = device
+
+        self._seq_len_cached = 0
+        self._cos_cached = None
+        self._sin_cached = None
+        self._cos_k_cached = None
+        self._sin_k_cached = None
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        inv_freq = self._compute_inv_freq(self.device)
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+        arange = torch.arange(0, self.dim, 2, device=self.device, dtype=torch.float32)
+        scale = (
+            (arange + 0.4 * self.dim) / (1.4 * self.dim)
+            if self.scale_base is not None
+            else None
+        )
+        self.register_buffer("scale", scale)
+
+    def _compute_inv_freq(self, device=None):
+        return 1 / (
+            self.base
+            ** (
+                torch.arange(0, self.dim, 2, device=device, dtype=torch.float32)
+                / self.dim
+            )
+        )
+
+    def _update_cos_sin_cache(self, seqlen, device=None, dtype=None):
+        # Reset the tables if the sequence length has changed,
+        # if we're on a new device (possibly due to tracing for instance),
+        # or if we're switching from inference mode to training
+        if (
+            seqlen > self._seq_len_cached
+            or self._cos_cached is None
+            or self._cos_cached.device != device
+            or self._cos_cached.dtype != dtype
+            or (self.training and self._cos_cached.is_inference())
+        ):
+            self._seq_len_cached = seqlen
+            # We want fp32 here, not self.inv_freq.dtype, since the model could be loaded in bf16
+            # And the output of arange can be quite large, so bf16 would lose a lot of precision.
+            # However, for compatibility reason, we add an option to use the dtype of self.inv_freq.
+            if self.pos_idx_in_fp32:
+                t = torch.arange(seqlen, device=device, dtype=torch.float32)
+                t /= self.scaling_factor
+                # We want fp32 here as well since inv_freq will be multiplied with t, and the output
+                # will be large. Having it in bf16 will lose a lot of precision and cause the
+                # cos & sin output to change significantly.
+                # We want to recompute self.inv_freq if it was not loaded in fp32
+                if self.inv_freq.dtype != torch.float32:
+                    inv_freq = self.inv_freq.to(torch.float32)
+                else:
+                    inv_freq = self.inv_freq
+            else:
+                t = torch.arange(seqlen, device=device, dtype=self.inv_freq.dtype)
+                t /= self.scaling_factor
+                inv_freq = self.inv_freq
+            # Don't do einsum, it converts fp32 to fp16 under AMP
+            # freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+            freqs = torch.outer(t, inv_freq)
+
+            if self.scale is None:
+                self._cos_cached = torch.cos(freqs).to(dtype)
+                self._sin_cached = torch.sin(freqs).to(dtype)
+            else:
+                power = (
+                    torch.arange(
+                        seqlen, dtype=self.scale.dtype, device=self.scale.device
+                    )
+                    - seqlen // 2
+                ) / self.scale_base
+                scale = self.scale.to(device=power.device) ** power.unsqueeze(-1)
+                # We want the multiplication by scale to happen in fp32
+                self._cos_cached = (torch.cos(freqs) * scale).to(dtype)
+                self._sin_cached = (torch.sin(freqs) * scale).to(dtype)
+                self._cos_k_cached = (torch.cos(freqs) / scale).to(dtype)
+                self._sin_k_cached = (torch.sin(freqs) / scale).to(dtype)
+
+    def forward(
+        self, q: torch.Tensor, k: torch.Tensor, seqlen_offset: int = 0
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        q: (batch, seqlen, nheads, headdim)
+        k: (batch, seqlen, nheads, headdim)
+        seqlen_offset: can be used in generation where the qkv being passed in is only the last
+        token in the batch.
+        """
+        self._update_cos_sin_cache(
+            q.shape[1] + seqlen_offset, device=q.device, dtype=q.dtype
+        )
+        assert self._cos_cached is not None
+        assert self._sin_cached is not None
+        if self.scale is None:
+            return (
+                apply_rotary_emb_torch(
+                    q,
+                    self._cos_cached[seqlen_offset:],
+                    self._sin_cached[seqlen_offset:],
+                    self.interleaved,
+                    True,  # inplace=True
+                ),
+                apply_rotary_emb_torch(
+                    k,
+                    self._cos_cached[seqlen_offset:],
+                    self._sin_cached[seqlen_offset:],
+                    self.interleaved,
+                    True,  # inplace=True
+                ),
+            )  # type: ignore
+        else:
+            assert False

Dyna-1/esm/layers/structure_proj.py

@@ -0,0 +1,66 @@
+import torch
+import torch.nn as nn
+
+from esm.utils.constants.physics import BB_COORDINATES
+from esm.utils.structure.affine3d import (
+    Affine3D,
+    RotationMatrix,
+)
+
+
+class Dim6RotStructureHead(nn.Module):
+    # Normally, AF2 uses quaternions to specify rotations. There's some evidence that
+    # other representations are more well behaved - the best one according to
+    # https://openaccess.thecvf.com/content_CVPR_2019/papers/Zhou_On_the_Continuity_of_Rotation_Representations_in_Neural_Networks_CVPR_2019_paper.pdf
+    # is using graham schmidt on 2 vectors, which is implemented here.
+    def __init__(
+        self,
+        input_dim: int,
+        trans_scale_factor: float = 10,
+        norm_type: str = "layernorm",
+        activation_fn: str = "esm_gelu",
+        predict_torsion_angles: bool = True,
+    ):
+        super().__init__()
+        self.ffn1 = nn.Linear(input_dim, input_dim)
+        self.activation_fn = nn.GELU()
+        self.norm = nn.LayerNorm(input_dim)
+        self.proj = nn.Linear(input_dim, 9 + 7 * 2)
+        self.trans_scale_factor = trans_scale_factor
+        self.predict_torsion_angles = predict_torsion_angles
+        self.bb_local_coords = torch.tensor(BB_COORDINATES).float()
+
+    def forward(self, x, affine, affine_mask, **kwargs):
+        if affine is None:
+            rigids = Affine3D.identity(
+                x.shape[:-1],
+                dtype=x.dtype,
+                device=x.device,
+                requires_grad=self.training,
+                rotation_type=RotationMatrix,
+            )
+        else:
+            rigids = affine
+
+        # [*, N]
+        x = self.ffn1(x)
+        x = self.activation_fn(x)
+        x = self.norm(x)
+        trans, x, y, angles = self.proj(x).split([3, 3, 3, 7 * 2], dim=-1)
+        trans = trans * self.trans_scale_factor
+        x = x / (x.norm(dim=-1, keepdim=True) + 1e-5)
+        y = y / (y.norm(dim=-1, keepdim=True) + 1e-5)
+        update = Affine3D.from_graham_schmidt(x + trans, trans, y + trans)
+        rigids = rigids.compose(update.mask(affine_mask))
+        affine = rigids.tensor
+
+        # We approximate the positions of the backbone atoms in the global frame by applying the rigid
+        # transformation to the mean of the backbone atoms in the local frame.
+        all_bb_coords_local = (
+            self.bb_local_coords[None, None, :, :]
+            .expand(*x.shape[:-1], 3, 3)
+            .to(x.device)
+        )
+        pred_xyz = rigids[..., None].apply(all_bb_coords_local)
+
+        return affine, pred_xyz

Dyna-1/esm/layers/transformer_stack.py

@@ -0,0 +1,93 @@
+import math
+
+import torch
+import torch.nn as nn
+
+from esm.layers.blocks import UnifiedTransformerBlock
+from esm.utils.structure.affine3d import Affine3D
+
+
+class TransformerStack(nn.Module):
+    """
+    A stack of transformer blocks used in the ESM-3 model. Each block is a UnifiedTransformerBlock,
+    which can either be geometric attention or standard multi-head attention.
+
+    Args:
+        d_model (int): The dimensionality of the input and output feature vectors.
+        n_heads (int): The number of attention heads.
+        v_heads (int): The number of voting heads.
+        n_layers (int): The number of transformer blocks in the stack.
+        n_layers_geom (int, optional): The number of transformer blocks that use geometric attention.
+        scale_residue (bool, optional): Whether to scale the residue connections in each transformer block.
+        mask_and_zero_frameless (bool, optional): Whether to mask and zero frameless positions in the input.
+            Only applies in the geometric attention blocks, which is conditioned on the structure
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        n_heads: int,
+        v_heads: int | None,
+        n_layers: int,
+        n_layers_geom: int = 1,
+        scale_residue: bool = True,
+        mask_and_zero_frameless: bool = False,
+        bias: bool = False,
+        qk_layernorm: bool = True,
+        ffn_type: str = "swiglu",  # swiglu | gelu
+        expansion_ratio: float = 8 / 3,
+    ):
+        super().__init__()
+        self.blocks = nn.ModuleList(
+            [
+                UnifiedTransformerBlock(
+                    d_model,
+                    n_heads,
+                    v_heads=v_heads,
+                    use_geom_attn=i < n_layers_geom,
+                    residue_scaling_factor=(
+                        math.sqrt(n_layers / 36) if scale_residue else 1.0
+                    ),
+                    expansion_ratio=expansion_ratio,
+                    mask_and_zero_frameless=mask_and_zero_frameless,
+                    bias=bias,
+                    qk_layernorm=qk_layernorm,
+                    ffn_type=ffn_type,
+                )
+                for i in range(n_layers)
+            ]
+        )
+        self.norm = nn.LayerNorm(d_model, bias=False)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        sequence_id: torch.Tensor | None = None,
+        affine: Affine3D | None = None,
+        affine_mask: torch.Tensor | None = None,
+        chain_id: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Forward pass of the TransformerStack.
+
+        Args:
+            x (torch.Tensor): The input tensor of shape (batch_size, sequence_length, d_model).
+            sequence_id (torch.Tensor): The sequence ID tensor of shape (batch_size, sequence_length).
+            affine (Affine3D | None): The affine transformation tensor or None.
+            affine_mask (torch.Tensor | None): The affine mask tensor or None.
+            chain_id (torch.Tensor): The protein chain tensor of shape (batch_size, sequence_length).
+                Only used in geometric attention.
+
+        Returns:
+            post_norm: The output tensor of shape (batch_size, sequence_length, d_model).
+            pre_norm: The embedding of shape (batch_size, sequence_length, d_model).
+        """
+        *batch_dims, _ = x.shape
+        if chain_id is None:
+            chain_id = torch.ones(size=batch_dims, dtype=torch.int64, device=x.device)
+        hiddens = []
+        for block in self.blocks:
+            x = block(x, sequence_id, affine, affine_mask, chain_id)
+            hiddens.append(x)
+        hiddens = torch.stack(hiddens, dim=0)
+        return self.norm(x), x, hiddens

Dyna-1/esm/models/esm3.py

@@ -0,0 +1,606 @@
+from __future__ import annotations
+
+import contextlib
+from functools import partial
+from typing import Callable
+
+import attr
+import einops
+import torch
+import torch.nn as nn
+from attr import dataclass
+
+from esm.layers.regression_head import RegressionHead
+from esm.layers.transformer_stack import TransformerStack
+from esm.models.function_decoder import FunctionTokenDecoder
+from esm.models.vqvae import (
+    StructureTokenDecoder,
+    StructureTokenEncoder,
+)
+from esm.sdk.api import (
+    ESM3InferenceClient,
+    ESMProtein,
+    ESMProteinTensor,
+    ForwardAndSampleOutput,
+    ForwardTrackData,
+    GenerationConfig,
+    LogitsConfig,
+    LogitsOutput,
+    ProteinType,
+    SamplingConfig,
+)
+from esm.tokenization import TokenizerCollectionProtocol
+from esm.utils import encoding
+from esm.utils.constants import esm3 as C
+from esm.utils.constants.models import (
+    ESM3_OPEN_SMALL,
+    normalize_model_name,
+)
+from esm.utils.decoding import decode_protein_tensor
+from esm.utils.generation import (
+    _batch_forward,
+    _sample_per_prompt,
+    _slice_tensor_dataclass,
+    iterative_sampling_raw,
+    iterative_sampling_tokens,
+)
+from esm.utils.misc import rbf
+from esm.utils.sampling import (
+    _BatchedESMProteinTensor,
+    get_default_sampling_config,
+    validate_sampling_config,
+)
+from esm.utils.structure.affine3d import (
+    build_affine3d_from_coordinates,
+)
+
+
+@dataclass
+class ESMOutput:
+    sequence_logits: torch.Tensor
+    structure_logits: torch.Tensor
+    secondary_structure_logits: torch.Tensor
+    sasa_logits: torch.Tensor
+    function_logits: torch.Tensor
+    residue_logits: torch.Tensor
+    embeddings: torch.Tensor
+
+
+class EncodeInputs(nn.Module):
+    """
+    Module for encoding input features in the ESM-3 model.
+
+    Args:
+        d_model (int): The dimensionality of the model's hidden states.
+    """
+
+    def __init__(self, d_model: int):
+        super().__init__()
+
+        # Sequence
+        self.sequence_embed = nn.Embedding(64, d_model)
+        # Mandatory information
+        self.plddt_projection = nn.Linear(16, d_model)
+        self.structure_per_res_plddt_projection = nn.Linear(16, d_model)
+
+        # Structure
+        self.structure_tokens_embed = nn.Embedding(4096 + 5, d_model)
+
+        # "Structural" features
+        self.ss8_embed = nn.Embedding(8 + 3, d_model)
+        self.sasa_embed = nn.Embedding(16 + 3, d_model)
+
+        # "Functional" features
+        self.function_embed = nn.ModuleList(
+            [nn.Embedding(260, d_model // 8, padding_idx=0) for _ in range(8)]
+        )
+
+        self.residue_embed = nn.EmbeddingBag(1478, d_model, mode="sum", padding_idx=0)
+
+    def forward(
+        self,
+        sequence_tokens: torch.Tensor,
+        structure_tokens: torch.Tensor,
+        average_plddt: torch.Tensor,
+        per_res_plddt: torch.Tensor,
+        ss8_tokens: torch.Tensor,
+        sasa_tokens: torch.Tensor,
+        function_tokens: torch.Tensor,
+        residue_annotation_tokens: torch.Tensor,
+    ) -> torch.Tensor:
+        sequence_embed = self.sequence_embed(sequence_tokens)
+
+        rbf_16_fn = partial(rbf, v_min=0.0, v_max=1.0, n_bins=16)
+        # the `masked_fill(padding_mask.unsqueeze(2), 0)` for the two below is unnecessary
+        # as pad tokens never even interact with the "real" tokens (due to sequence_id)
+        plddt_embed = self.plddt_projection(rbf_16_fn(average_plddt))
+        structure_per_res_plddt = self.structure_per_res_plddt_projection(
+            rbf_16_fn(per_res_plddt)
+        )
+
+        # Structure + "structural features" embeds
+        structure_embed = self.structure_tokens_embed(structure_tokens)
+        ss8_embed = self.ss8_embed(ss8_tokens)
+        sasa_embed = self.sasa_embed(sasa_tokens)
+
+        # "Functional" features embeds
+        function_embed = torch.cat(
+            [
+                embed_fn(funcs)
+                for embed_fn, funcs in zip(
+                    self.function_embed, function_tokens.unbind(-1)
+                )
+            ],
+            -1,
+        )
+
+        # Residue embeds
+        B, L, N = residue_annotation_tokens.shape
+        residue_embed = self.residue_embed(
+            einops.rearrange(
+                residue_annotation_tokens, "B L N -> (B L) N", B=B, L=L, N=N
+            )
+        )
+        residue_embed = einops.rearrange(residue_embed, "(B L) D -> B L D", B=B, L=L)
+
+        return (
+            sequence_embed
+            + plddt_embed
+            + structure_per_res_plddt
+            + structure_embed
+            + ss8_embed
+            + sasa_embed
+            + function_embed
+            + residue_embed
+        )
+
+
+class OutputHeads(nn.Module):
+    def __init__(self, d_model: int):
+        super().__init__()
+        self.sequence_head = RegressionHead(d_model, 64)
+        self.structure_head = RegressionHead(d_model, 4096)
+        self.ss8_head = RegressionHead(d_model, 8 + 3)
+        self.sasa_head = RegressionHead(d_model, 16 + 3)
+        self.function_head = RegressionHead(d_model, 260 * 8)
+        self.residue_head = RegressionHead(d_model, 1478)
+
+    def forward(self, x: torch.Tensor, embed: torch.Tensor) -> ESMOutput:
+        sequence_logits = self.sequence_head(x)
+        structure_logits = self.structure_head(x)
+        secondary_structure_logits = self.ss8_head(x)
+        sasa_logits = self.sasa_head(x)
+        function_logits = self.function_head(x)
+        function_logits = einops.rearrange(function_logits, "... (k v) -> ... k v", k=8)
+
+        residue_logits = self.residue_head(x)
+
+        return ESMOutput(
+            sequence_logits=sequence_logits,
+            structure_logits=structure_logits,
+            secondary_structure_logits=secondary_structure_logits,
+            sasa_logits=sasa_logits,
+            function_logits=function_logits,
+            residue_logits=residue_logits,
+            embeddings=embed,
+        )
+
+
+class ESM3(nn.Module, ESM3InferenceClient):
+    """
+    ESM3 model implementation.
+
+    Args:
+        d_model (int): The dimensionality of the input and output feature vectors.
+        n_heads (int): The number of attention heads in the transformer layers.
+        v_heads (int): The number of attention heads in the variational transformer layers.
+        n_layers (int): The number of transformer layers.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        n_heads: int,
+        v_heads: int,
+        n_layers: int,
+        structure_encoder_fn: Callable[[torch.device | str], StructureTokenEncoder],
+        structure_decoder_fn: Callable[[torch.device | str], StructureTokenDecoder],
+        function_decoder_fn: Callable[[torch.device | str], FunctionTokenDecoder],
+        tokenizers: TokenizerCollectionProtocol,
+    ):
+        super().__init__()
+        self.encoder = EncodeInputs(d_model)
+        self.transformer = TransformerStack(
+            d_model, n_heads, v_heads, n_layers, mask_and_zero_frameless=True
+        )
+        self.output_heads = OutputHeads(d_model)
+
+        self.structure_encoder_fn = structure_encoder_fn
+        self.structure_decoder_fn = structure_decoder_fn
+        self.function_decoder_fn = function_decoder_fn
+
+        self._structure_encoder = None
+        self._structure_decoder = None
+        self._function_decoder = None
+
+        self.tokenizers = tokenizers
+
+    @classmethod
+    def from_pretrained(
+        cls, model_name: str = ESM3_OPEN_SMALL, device: torch.device | None = None
+    ) -> ESM3:
+        from esm.pretrained import load_local_model
+
+        model_name = normalize_model_name(model_name)
+        if not model_name:
+            raise ValueError(f"Model name {model_name} is not a valid ESM3 model name.")
+        if device is None:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model = load_local_model(model_name, device=device)
+        if device.type != "cpu":
+            model = model.to(torch.bfloat16)
+        assert isinstance(model, ESM3)
+        return model
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property
+    def raw_model(self):
+        return self
+
+    def get_structure_encoder(self) -> StructureTokenEncoder:
+        if self._structure_encoder is None:
+            self._structure_encoder = self.structure_encoder_fn(self.device)
+        return self._structure_encoder
+
+    def get_structure_decoder(self) -> StructureTokenDecoder:
+        if self._structure_decoder is None:
+            self._structure_decoder = self.structure_decoder_fn(self.device)
+        return self._structure_decoder
+
+    def get_function_decoder(self) -> FunctionTokenDecoder:
+        if self._function_decoder is None:
+            self._function_decoder = self.function_decoder_fn(self.device)
+        return self._function_decoder
+
+    def forward(
+        self,
+        *,
+        sequence_tokens: torch.Tensor | None = None,
+        structure_tokens: torch.Tensor | None = None,
+        ss8_tokens: torch.Tensor | None = None,
+        sasa_tokens: torch.Tensor | None = None,
+        function_tokens: torch.Tensor | None = None,
+        residue_annotation_tokens: torch.Tensor | None = None,
+        average_plddt: torch.Tensor | None = None,
+        per_res_plddt: torch.Tensor | None = None,
+        structure_coords: torch.Tensor | None = None,
+        chain_id: torch.Tensor | None = None,
+        sequence_id: torch.Tensor | None = None,
+    ) -> ESMOutput:
+        """
+        Performs forward pass through the ESM3 model. Check utils to see how to tokenize inputs from raw data.
+
+        Args:
+            sequence_tokens (torch.Tensor, optional): The amino acid tokens.
+            structure_tokens (torch.Tensor, optional): The structure tokens.
+            ss8_tokens (torch.Tensor, optional): The secondary structure tokens.
+            sasa_tokens (torch.Tensor, optional): The solvent accessible surface area tokens.
+            function_tokens (torch.Tensor, optional): The function tokens.
+            residue_annotation_tokens (torch.Tensor, optional): The residue annotation tokens.
+            average_plddt (torch.Tensor, optional): The average plddt across the entire sequence.
+            per_res_plddt (torch.Tensor, optional): The per residue plddt, if you want to specify exact plddts, use this,
+                otherwise, use average_plddt.
+            structure_coords (torch.Tensor, optional): The structure coordinates, in the form of (B, L, 3, 3).
+            chain_id (torch.Tensor, optional): The chain ID
+            sequence_id (torch.Tensor, optional): The sequence ID.
+
+        Returns:
+            ESMOutput: The output of the ESM3 model.
+
+        Raises:
+            ValueError: If at least one of the inputs is None.
+
+        """
+        # Reasonable defaults:
+        try:
+            L, device = next(
+                (x.shape[1], x.device)
+                for x in [
+                    sequence_tokens,
+                    structure_tokens,
+                    ss8_tokens,
+                    sasa_tokens,
+                    structure_coords,
+                    function_tokens,
+                    residue_annotation_tokens,
+                ]
+                if x is not None
+            )
+        except StopIteration:
+            raise ValueError("At least one of the inputs must be non-None")
+
+        t = self.tokenizers
+        defaults = lambda x, tok: (
+            torch.full((1, L), tok, dtype=torch.long, device=device) if x is None else x
+        )
+        sequence_tokens = defaults(sequence_tokens, t.sequence.mask_token_id)
+        ss8_tokens = defaults(ss8_tokens, C.SS8_PAD_TOKEN)
+        sasa_tokens = defaults(sasa_tokens, C.SASA_PAD_TOKEN)
+        average_plddt = defaults(average_plddt, 1).float()
+        per_res_plddt = defaults(per_res_plddt, 0).float()
+        chain_id = defaults(chain_id, 0)
+
+        if residue_annotation_tokens is None:
+            residue_annotation_tokens = torch.full(
+                (1, L, 16), C.RESIDUE_PAD_TOKEN, dtype=torch.long, device=device
+            )
+
+        if function_tokens is None:
+            function_tokens = torch.full(
+                (1, L, 8), C.INTERPRO_PAD_TOKEN, dtype=torch.long, device=device
+            )
+
+        if structure_coords is None:
+            structure_coords = torch.full(
+                (1, L, 3, 3), float("nan"), dtype=torch.float, device=device
+            )
+
+        structure_coords = structure_coords[
+            ..., :3, :
+        ]  # In case we pass in an atom14 or atom37 repr
+        affine, affine_mask = build_affine3d_from_coordinates(structure_coords)
+
+        structure_tokens = defaults(structure_tokens, C.STRUCTURE_MASK_TOKEN)
+        assert structure_tokens is not None
+        structure_tokens = (
+            structure_tokens.masked_fill(structure_tokens == -1, C.STRUCTURE_MASK_TOKEN)
+            .masked_fill(sequence_tokens == C.SEQUENCE_BOS_TOKEN, C.STRUCTURE_BOS_TOKEN)
+            .masked_fill(sequence_tokens == C.SEQUENCE_PAD_TOKEN, C.STRUCTURE_PAD_TOKEN)
+            .masked_fill(sequence_tokens == C.SEQUENCE_EOS_TOKEN, C.STRUCTURE_EOS_TOKEN)
+            .masked_fill(
+                sequence_tokens == C.SEQUENCE_CHAINBREAK_TOKEN,
+                C.STRUCTURE_CHAINBREAK_TOKEN,
+            )
+        )
+
+        x = self.encoder(
+            sequence_tokens,
+            structure_tokens,
+            average_plddt,
+            per_res_plddt,
+            ss8_tokens,
+            sasa_tokens,
+            function_tokens,
+            residue_annotation_tokens,
+        )
+        x, embedding, hidden = self.transformer(
+            x, sequence_id, affine, affine_mask, chain_id
+        )
+        return self.output_heads(x, embedding), hidden # MODIFIED FOR DYNA-1
+
+    # The following methods are for the ESM3InferenceClient interface
+    def generate(self, input: ProteinType, config: GenerationConfig) -> ProteinType:
+        """Wrap around batched generation."""
+        proteins = self.batch_generate([input], [config])
+        assert len(proteins) == 1
+        return proteins[0]
+
+    def batch_generate(
+        self, inputs: list[ProteinType], configs: list[GenerationConfig]
+    ) -> list[ProteinType]:
+        assert len(inputs) == len(
+            configs
+        ), "Must have the same number of prompts and configs."
+
+        if inputs == []:
+            # Nothing to do.
+            return []
+
+        # Make sure prompts are of the same type.
+        t = type(inputs[0])
+        for i in range(1, len(inputs)):
+            assert isinstance(inputs[i], t), (
+                "Prompts must have the same type. Got "
+                f"{t.__name__ and type(inputs[i]).__name__} instead."
+            )
+
+        if isinstance(inputs[0], ESMProtein):
+            return iterative_sampling_raw(self, inputs, configs)  # type: ignore
+        elif isinstance(inputs[0], ESMProteinTensor):
+            return iterative_sampling_tokens(
+                self,
+                inputs,  # type: ignore
+                configs,
+                self.tokenizers,  # type: ignore
+            )
+        else:
+            raise ValueError("Input must be an ESMProtein or ESMProteinTensor")
+
+    def encode(self, input: ESMProtein) -> ESMProteinTensor:
+        input = attr.evolve(input)  # Make a copy
+
+        sequence_tokens = None
+        structure_tokens = None
+        secondary_structure_tokens = None
+        sasa_tokens = None
+        function_tokens = None
+        residue_annotation_tokens = None
+
+        coordinates = None
+
+        if input.sequence is not None:
+            sequence_tokens = encoding.tokenize_sequence(
+                input.sequence, self.tokenizers.sequence, add_special_tokens=True
+            )
+        if input.secondary_structure is not None:
+            secondary_structure_tokens = encoding.tokenize_secondary_structure(
+                input.secondary_structure,
+                self.tokenizers.secondary_structure,
+                add_special_tokens=True,
+            )
+        if input.sasa is not None:
+            sasa_tokens = encoding.tokenize_sasa(
+                input.sasa, self.tokenizers.sasa, add_special_tokens=True
+            )
+
+        # Infer input length
+        sequence_length = -1
+        if sequence_tokens is not None:
+            sequence_length = len(sequence_tokens)
+        elif secondary_structure_tokens is not None:
+            sequence_length = len(secondary_structure_tokens)
+        elif sasa_tokens is not None:
+            sequence_length = len(sasa_tokens)
+
+        # Try to infer input length from structure data
+        if input.coordinates is not None:
+            coordinates, _, structure_tokens = encoding.tokenize_structure(
+                input.coordinates,
+                self.get_structure_encoder(),
+                structure_tokenizer=self.tokenizers.structure,
+                reference_sequence=input.sequence or "",
+                add_special_tokens=True,
+            )
+            if sequence_length == -1:
+                sequence_length = len(structure_tokens)
+
+        if sequence_length == -1:
+            raise ValueError(
+                "Cannot infer input length from input data. Please provide one of: sequence, structure, secondary_structure, sasa.\n"
+                "To condition on sequence length only, use ESM3LocalInferenceClient.get_default_sequence(sequence_length) to generate a default sequence input."
+            )
+
+        # Function and Residue annotations
+        if input.function_annotations is not None:
+            if input.sequence is None:
+                reference_sequence = encoding.get_default_sequence(sequence_length - 2)
+            else:
+                reference_sequence = input.sequence
+            (function_tokens, residue_annotation_tokens) = (
+                encoding.tokenize_function_annotations(
+                    input.function_annotations,
+                    reference_sequence=reference_sequence,
+                    function_tokenizer=self.tokenizers.function,
+                    residue_annotation_tokenizer=self.tokenizers.residue_annotations,
+                    add_special_tokens=True,
+                )
+            )
+
+        return ESMProteinTensor(
+            sequence=sequence_tokens,
+            structure=structure_tokens,
+            secondary_structure=secondary_structure_tokens,
+            sasa=sasa_tokens,
+            function=function_tokens,
+            residue_annotations=residue_annotation_tokens,
+            coordinates=coordinates,
+        ).to(next(self.parameters()).device)
+
+    def decode(self, input: ESMProteinTensor) -> ESMProtein:
+        return decode_protein_tensor(
+            input=input,
+            tokenizers=self.tokenizers,
+            structure_token_decoder=self.get_structure_decoder(),
+            function_token_decoder=self.get_function_decoder(),
+        )
+
+    def logits(
+        self,
+        input: ESMProteinTensor | _BatchedESMProteinTensor,
+        config: LogitsConfig = LogitsConfig(),
+    ) -> LogitsOutput:
+        if not isinstance(input, _BatchedESMProteinTensor):
+            # Create batch dimension if necessary.
+            input = _BatchedESMProteinTensor.from_protein_tensor(input)
+
+        device = torch.device(input.device)
+
+        # Default plddt conditioning for inference. 1s where coordinates are provided.
+        if input.coordinates is None:
+            per_res_plddt = None
+        else:
+            # 1.0 if all coordinates at specific indices have valid non-nan values.
+            per_res_plddt = input.coordinates.isfinite().all(dim=-1).any(dim=-1).float()
+
+        with (
+            torch.no_grad(),  # Assume no gradients for now...
+            torch.autocast(enabled=True, device_type=device.type, dtype=torch.bfloat16)  # type: ignore
+            if device.type == "cuda"
+            else contextlib.nullcontext(),
+        ):
+            output = self.forward(
+                sequence_tokens=input.sequence,
+                structure_tokens=input.structure,
+                ss8_tokens=input.secondary_structure,
+                sasa_tokens=input.sasa,
+                function_tokens=input.function,
+                residue_annotation_tokens=input.residue_annotations,
+                average_plddt=torch.tensor(1.0, device=input.device),
+                per_res_plddt=per_res_plddt,
+                structure_coords=input.coordinates,
+                chain_id=None,
+                sequence_id=None,
+            )
+
+        output = ESMOutput(
+            **{k: v.to(device).to(torch.float32) for k, v in vars(output).items()}
+        )
+
+        return LogitsOutput(
+            logits=ForwardTrackData(
+                sequence=output.sequence_logits if config.sequence else None,
+                structure=output.structure_logits if config.structure else None,
+                secondary_structure=output.secondary_structure_logits
+                if config.secondary_structure
+                else None,
+                sasa=output.sasa_logits if config.sasa else None,
+                function=output.function_logits if config.function else None,
+            ),
+            residue_annotation_logits=output.residue_logits
+            if config.residue_annotations
+            else None,
+            embeddings=output.embeddings if config.return_embeddings else None,
+        )
+
+    def forward_and_sample(
+        self, input: ESMProteinTensor, sampling_configuration: SamplingConfig
+    ) -> ForwardAndSampleOutput:
+        validate_sampling_config(sampling_configuration, on_invalid="warn")
+
+        protein_tensor = attr.evolve(input)  # Make a copy
+
+        device = next(self.parameters()).device
+
+        sampling_config = sampling_configuration
+        if sampling_config is None:
+            sampling_config = get_default_sampling_config(self.tokenizers)
+
+        # Initialize default values for missing tracks
+        default_protein_tensor = ESMProteinTensor.empty(
+            len(input) - 2, tokenizers=self.tokenizers, device=input.device
+        )
+        for track in attr.fields(ESMProteinTensor):
+            if getattr(protein_tensor, track.name, None) is None:
+                setattr(
+                    protein_tensor,
+                    track.name,
+                    getattr(default_protein_tensor, track.name, None),
+                )
+
+        if len(protein_tensor) <= 0:
+            raise ValueError("No input data provided")
+
+        # Move input protein to proper device.
+        batched_protein = _BatchedESMProteinTensor.from_protein_tensor(protein_tensor)
+        batched_protein.to(device)
+
+        logits_output: LogitsOutput = _batch_forward(self, batched_protein)
+        forward_and_sample_out: ForwardAndSampleOutput = _sample_per_prompt(
+            batched_protein, logits_output, sampling_config, self.tokenizers
+        )
+
+        # There is only 1 prompt to sample for.
+        return _slice_tensor_dataclass(forward_and_sample_out, 0)

Dyna-1/esm/models/esmc.py

@@ -0,0 +1,164 @@
+from __future__ import annotations
+
+import contextlib
+
+import attr
+import torch
+import torch.nn as nn
+from attr import dataclass
+
+from esm.layers.regression_head import RegressionHead
+from esm.layers.transformer_stack import TransformerStack
+from esm.sdk.api import (
+    ESMCInferenceClient,
+    ESMProtein,
+    ESMProteinTensor,
+    ForwardTrackData,
+    LogitsConfig,
+    LogitsOutput,
+)
+from esm.tokenization import EsmSequenceTokenizer
+from esm.utils import encoding
+from esm.utils.constants.models import ESMC_600M
+from esm.utils.decoding import decode_sequence
+from esm.utils.misc import stack_variable_length_tensors
+from esm.utils.sampling import _BatchedESMProteinTensor
+
+
+@dataclass
+class ESMCOutput:
+    sequence_logits: torch.Tensor
+    embeddings: torch.Tensor | None
+    hidden_states: torch.Tensor | None
+
+
+class ESMC(nn.Module, ESMCInferenceClient):
+    """
+    ESMC model implementation.
+
+    Args:
+        d_model (int): The dimensionality of the input and output feature vectors.
+        n_heads (int): The number of attention heads in the transformer layers.
+        n_layers (int): The number of transformer layers.
+    """
+
+    def __init__(
+        self, d_model: int, n_heads: int, n_layers: int, tokenizer: EsmSequenceTokenizer
+    ):
+        super().__init__()
+        self.embed = nn.Embedding(64, d_model)
+        self.transformer = TransformerStack(
+            d_model, n_heads, None, n_layers, n_layers_geom=0
+        )
+        self.sequence_head = RegressionHead(d_model, 64)
+        self.tokenizer = tokenizer
+
+    @classmethod
+    def from_pretrained(
+        cls, model_name: str = ESMC_600M, device: torch.device | None = None
+    ) -> ESMC:
+        from esm.pretrained import load_local_model
+
+        if device is None:
+            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        model = load_local_model(model_name, device=device)
+        if device.type != "cpu":
+            model = model.to(torch.bfloat16)
+        assert isinstance(model, ESMC)
+        return model
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    @property
+    def raw_model(self):
+        return self
+
+    def _tokenize(self, sequence: list[str]) -> torch.Tensor:
+        pad = self.tokenizer.pad_token_id
+        assert pad is not None
+        return stack_variable_length_tensors(
+            [
+                encoding.tokenize_sequence(x, self.tokenizer, add_special_tokens=True)
+                for x in sequence
+            ],
+            constant_value=pad,
+        ).to(next(self.parameters()).device)
+
+    def _detokenize(self, sequence: torch.Tensor) -> list[str]:
+        pad = self.tokenizer.pad_token_id
+        assert pad is not None
+        assert sequence.ndim == 2
+        return [decode_sequence(x[x != pad][1:-1], self.tokenizer) for x in sequence]
+
+    def forward(
+        self,
+        sequence_tokens: torch.Tensor | None = None,
+        sequence_id: torch.Tensor | None = None,
+    ) -> ESMCOutput:
+        """
+        Performs forward pass through the ESMC model. Check utils to see how to tokenize inputs from raw data.
+
+        Args:
+            sequence_tokens (torch.Tensor, optional): The amino acid tokens.
+            sequence_id (torch.Tensor, optional): The sequence ID.
+
+        Returns:
+            ESMCOutput: The output of the ESMC model.
+
+        """
+        if sequence_id is None:
+            sequence_id = sequence_tokens == self.tokenizer.pad_token_id
+
+        x = self.embed(sequence_tokens)
+        x, _, hidden = self.transformer(x, sequence_id=sequence_id)
+        sequence_logits = self.sequence_head(x)
+        output = ESMCOutput(
+            sequence_logits=sequence_logits, embeddings=x, hidden_states=hidden
+        )
+        return output, hidden #MODIFIED FOR DYNA-1
+
+    def encode(self, input: ESMProtein) -> ESMProteinTensor:
+        input = attr.evolve(input)  # Make a copy
+        sequence_tokens = None
+
+        if input.sequence is not None:
+            sequence_tokens = self._tokenize([input.sequence])[0]
+        return ESMProteinTensor(sequence=sequence_tokens).to(
+            next(self.parameters()).device
+        )
+
+    def decode(self, input: ESMProteinTensor) -> ESMProtein:
+        input = attr.evolve(input)  # Make a copy
+
+        assert input.sequence is not None
+        sequence = self._detokenize(input.sequence)[0]
+
+        return ESMProtein(sequence=sequence)
+
+    def logits(
+        self,
+        input: ESMProteinTensor | _BatchedESMProteinTensor,
+        config: LogitsConfig = LogitsConfig(),
+    ) -> LogitsOutput:
+        if not isinstance(input, _BatchedESMProteinTensor):
+            # Create batch dimension if necessary.
+            input = _BatchedESMProteinTensor.from_protein_tensor(input)
+
+        device = torch.device(input.device)
+
+        with (
+            torch.no_grad(),
+            torch.autocast(enabled=True, device_type=device.type, dtype=torch.bfloat16)  # type: ignore
+            if device.type == "cuda"
+            else contextlib.nullcontext(),
+        ):
+            output = self.forward(sequence_tokens=input.sequence)
+
+        return LogitsOutput(
+            logits=ForwardTrackData(
+                sequence=output.sequence_logits if config.sequence else None
+            ),
+            embeddings=output.embeddings if config.return_embeddings else None,
+        )

Dyna-1/esm/models/function_decoder.py

@@ -0,0 +1,306 @@
+"""Function Token Decoder."""
+
+from collections import defaultdict
+from dataclasses import dataclass, field
+
+import numpy as np
+import pandas as pd
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from cloudpathlib import AnyPath
+
+from esm.layers.regression_head import RegressionHead
+from esm.layers.transformer_stack import TransformerStack
+from esm.tokenization.function_tokenizer import (
+    InterProQuantizedTokenizer,
+)
+from esm.utils.constants import esm3 as C
+from esm.utils.misc import merge_annotations, merge_ranges
+from esm.utils.types import FunctionAnnotation
+
+
+@dataclass(frozen=True)
+class FunctionTokenDecoderConfig:
+    """Configures function token decoder."""
+
+    # Embedding dimension of decoder.
+    d_model: int = 1024
+    # Number of attention heads of decoder.
+    n_heads: int = 8
+    # Number of layers of decoder.
+    n_layers: int = 3
+    # Number of integer values that function tokens may assume.
+    function_token_vocab_size: int = 260
+    # Number of function tokens at each position.
+    function_token_depth: int = 8
+    # Number of InterPro labels that can be decoded.
+    num_interpro_classes: int = 29026
+    # Number of function keywords that can be decoded.
+    keyword_vocabulary_size: int = 58641
+    # List of supported InterPro ids.
+    interpro_entry_list: str = field(default_factory=lambda: str(C.INTERPRO_ENTRY))
+    # Path to keywords vocabulary.
+    keyword_vocabulary_path: str = field(
+        default_factory=lambda: str(C.data_root("esm3") / C.KEYWORDS_VOCABULARY)
+    )
+    # Whether to unpack LSH bits into single-bit tokens.
+    unpack_lsh_bits: bool = True
+    # The number of special tokens in the function tokenizer vocabulary which come
+    # before the LSH tokens.
+    num_special_tokens: int = 4
+    # The number of bits per LSH token in the function tokenizer.
+    bits_per_token: int = 8
+
+
+class FunctionTokenDecoder(nn.Module):
+    def __init__(self, config: FunctionTokenDecoderConfig | None = None):
+        """Constructs function token decoder."""
+        super().__init__()
+        if config is None:
+            config = FunctionTokenDecoderConfig()
+        self.config = config
+
+        # Get the supported set of InterPro ids.
+        with AnyPath(config.interpro_entry_list).open("r") as f:
+            df = pd.read_csv(f, sep="\t")
+        self.interpro_ids = sorted(df.ENTRY_AC)
+        self.interpro2index = {
+            interpro_id: i for i, interpro_id in enumerate(self.interpro_ids)
+        }
+        assert len(self.interpro_ids) == config.num_interpro_classes
+
+        with AnyPath(config.keyword_vocabulary_path).open("r") as f:
+            self.keywords_vocabulary: list[str] = list(f.read().strip().split("\n"))
+            assert len(self.keywords_vocabulary) == config.keyword_vocabulary_size
+
+        if config.unpack_lsh_bits:
+            vocab_size = 2 * config.function_token_depth * config.bits_per_token
+        else:
+            # Function-token id's re-use the same token ids at each position along the depth
+            # dimension, despite distinct meanings. The decoder should take this into
+            # account so create distinct embeddings for tokens at each position.
+            vocab_size = (
+                self.config.function_token_depth * self.config.function_token_vocab_size
+            )
+
+        self.embedding = nn.Embedding(
+            # Function-token id's re-use the same token ids at each position along the
+            # depth dimension, despite distinct meanings. The decoder should take this
+            # into account so create distinct embeddings for tokens at each position.
+            num_embeddings=(vocab_size),
+            embedding_dim=config.d_model,
+        )
+        self.decoder = TransformerStack(
+            d_model=config.d_model,
+            n_heads=config.n_heads,
+            v_heads=None,
+            n_layers=config.n_layers,
+            n_layers_geom=0,
+            scale_residue=False,
+            bias=True,
+            qk_layernorm=False,
+            ffn_type="gelu",
+            expansion_ratio=4,
+        )
+        self.heads = nn.ModuleDict(
+            {
+                # Binary classification head predicting which keywords are present.
+                "keyword_logits": RegressionHead(
+                    d_model=config.d_model,
+                    output_dim=config.keyword_vocabulary_size,
+                    hidden_dim=4 * config.d_model,
+                ),
+                # Regresses the TF-IDF value of each present keyword.
+                "keyword_tfidf": RegressionHead(
+                    d_model=config.d_model,
+                    output_dim=config.keyword_vocabulary_size,
+                    hidden_dim=4 * config.d_model,
+                ),
+                # Predicts which InterPro annotations are present.
+                "interpro_logits": RegressionHead(
+                    d_model=config.d_model,
+                    output_dim=config.num_interpro_classes,
+                    hidden_dim=4 * config.d_model,
+                ),
+            }
+        )
+
+    def forward(self, token_ids: torch.Tensor) -> dict[str, torch.Tensor]:
+        """Forward pass through function token decoder.
+
+        Args:
+            token_ids: <int>[batch_size, function_token_depth] batch of function tokens
+                ids to decode.
+        Returns:
+            interpro_logits: binary classification logits tensor of shape
+                <float>[batch_size, num_interpro_classes]
+        """
+        assert token_ids.ndim == 2
+        assert token_ids.shape[1] == self.config.function_token_depth
+        batch_size, depth = token_ids.shape
+
+        if self.config.unpack_lsh_bits:
+            # Shift values into [0, 2^bits/token)
+            lsh_bits = token_ids - self.config.num_special_tokens
+            # extract each bit. (hob stands for highest-order bit)
+            bits = torch.concat(
+                [
+                    torch.bitwise_and(lsh_bits, 1 << hob).gt(0).to(torch.int32)
+                    for hob in range(self.config.bits_per_token)
+                ],
+                dim=1,
+            )
+            assert bits.shape == (batch_size, depth * self.config.bits_per_token)
+
+            # Shift each bit into individual vocabulary ranges, so they get distinct
+            # embeddings.
+            vocab_offsets = 2 * torch.arange(
+                depth * self.config.bits_per_token, device=token_ids.device
+            )
+            inputs = vocab_offsets[None, :] + bits
+
+            # zero-out special tokens, i.e. non LSH tokens.
+            where_special = token_ids < self.config.num_special_tokens
+            inputs = torch.where(where_special.any(dim=1, keepdim=True), 0, inputs)
+        else:
+            # Apply depth-position offset to use distinct vocabs. See __init__ for
+            # explaination.
+            vocab_offsets = self.config.function_token_vocab_size * torch.arange(
+                self.config.function_token_depth, device=token_ids.device
+            )
+            inputs = token_ids + vocab_offsets[None, :]
+
+        embed = self.embedding(inputs)
+        encoding, _, _ = self.decoder(embed)
+        pooled = torch.mean(encoding, dim=1)
+
+        return {name: head(pooled) for name, head in self.heads.items()}
+
+    @property
+    def device(self) -> torch.device:
+        return next(self.parameters()).device
+
+    def decode(
+        self,
+        function_token_ids: torch.Tensor,
+        tokenizer: InterProQuantizedTokenizer,
+        decode_annotations: bool = True,
+        annotation_threshold: float = 0.1,
+        decode_keywords=True,
+        keywords_threshold: float = 0.5,
+        annotation_min_length: int | None = 5,
+        annotation_gap_merge_max: int | None = 3,
+    ):
+        """Decodes function tokens into predicted annotations and keywords.
+
+        Args:
+            function_token_ids: <int>[length, depth] function token ids. NOTE:
+                without <bos>/<eos> prefix
+            tokenizer: function tokenizer.
+            decode_annotations: whether to decode InterPro annotations.
+            annotation_threshold: threshold for emitting a function annotation.
+            decode_keywords: whether to decode function keywords.
+            keywords_threshold: threshold for emitting a keyword.
+            annotation_min_length: optional minimum length of predicted annotations for
+                size filtering.
+            annotation_gap_merge_max: optional merge adjacent annotation of the same type
+        Returns:
+            Decoder outputs:
+            - "interpro_logits": <float>[length, num_interpro] predicted interpro logits.
+            - "interpro_preds": <bool>[length, num_interpro] predicted intepro labels.
+            - "interpro_annotations": list[FunctionAnnotation] predicted InterPro
+                annotations
+            - "keyword_logits": <float>[length, keyword_vocabulary] binary prediciton
+              logits for keywrods.
+            - "function_keywords": list[FunctionAnnotation] predicted function keyword
+                ranges.
+        """
+        assert function_token_ids.ndim == 2
+        assert function_token_ids.shape[1] == tokenizer.depth
+        assert self.config.function_token_depth == tokenizer.depth
+
+        outputs = {}
+
+        outputs = self(function_token_ids.to(self.device))
+
+        # Only decode in positions that have function tokens.
+        where_decode = torch.all(
+            (function_token_ids != tokenizer.vocab_to_index["<pad>"])
+            & (function_token_ids != tokenizer.vocab_to_index["<none>"])
+            & (function_token_ids != tokenizer.vocab_to_index["<unk>"]),
+            dim=1,
+        )
+
+        # Decode InterPro annotations ranges.
+        interpro_preds = F.sigmoid(outputs["interpro_logits"])
+        interpro_preds = interpro_preds >= annotation_threshold
+        interpro_preds[~where_decode, :] = False
+        outputs["interpro_preds"] = interpro_preds
+        if decode_annotations:
+            annotations: list[FunctionAnnotation] = []
+            preds: np.ndarray = interpro_preds.detach().cpu().numpy()
+            for position_index, class_index in zip(*preds.nonzero()):
+                interpro_id = self.interpro_ids[class_index]
+                annotation = FunctionAnnotation(
+                    label=interpro_id,
+                    start=position_index,  # one-index inclusive (BOS shifts indexes +1)
+                    end=position_index,  # one-index inclusive
+                )
+                annotations.append(annotation)
+
+            annotations = merge_annotations(
+                annotations, merge_gap_max=annotation_gap_merge_max
+            )
+
+            # Drop very small annotations.
+            if annotation_min_length is not None:
+                annotations = [
+                    annotation
+                    for annotation in annotations
+                    if annotation.end - annotation.start + 1 >= annotation_min_length
+                ]
+
+            outputs["interpro_annotations"] = annotations
+
+        # Decode function keyword ranges.
+        keyword_logits = outputs["keyword_logits"]
+        keyword_logits[~where_decode, :] = -torch.inf
+        if decode_keywords:
+            keyword_preds = F.sigmoid(keyword_logits) >= keywords_threshold
+            outputs["function_keywords"] = self._preds_to_keywords(
+                keyword_preds.detach().cpu().numpy()
+            )
+
+        return outputs
+
+    def _preds_to_keywords(self, keyword_preds: np.ndarray) -> list[FunctionAnnotation]:
+        """Converts output log-TFDF to predicted keywords over the sequence.
+
+        Args:
+            keyword_precs: <bool>[length, keyword_vocab] positional predictions of
+              function keywords from the keyword prediction head.
+        Returns:
+            Non-overlapping keyword annotated ranges along the sequence. Note that indices
+            will index into the *sequence*, not the function token array which has a
+            <pad> prefix.
+        """
+        assert keyword_preds.ndim == 2
+        assert keyword_preds.shape[1] == self.config.keyword_vocabulary_size
+
+        keyword_positions: dict[str, list[range]] = defaultdict(list)
+        for position, keyword_id in zip(*np.nonzero(keyword_preds)):
+            keyword = self.keywords_vocabulary[keyword_id]
+            keyword_positions[keyword].append(range(position, position + 1))
+
+        annotations: list[FunctionAnnotation] = []
+        for keyword, ranges in keyword_positions.items():
+            for range_ in merge_ranges(ranges):
+                annotation = FunctionAnnotation(
+                    label=keyword,
+                    start=range_.start,  # one-index inclusive  (BOS shifts indexes +1)
+                    end=range_.stop - 1,  # one-index exclusive -> one-index inclusive
+                )
+                annotations.append(annotation)
+
+        return annotations

Dyna-1/esm/models/vqvae.py

@@ -0,0 +1,440 @@
+import torch
+import torch.nn as nn
+
+from esm.layers.blocks import UnifiedTransformerBlock
+from esm.layers.codebook import EMACodebook
+from esm.layers.structure_proj import Dim6RotStructureHead
+from esm.layers.transformer_stack import TransformerStack
+from esm.utils.constants import esm3 as C
+from esm.utils.misc import knn_graph
+from esm.utils.structure.affine3d import (
+    Affine3D,
+    build_affine3d_from_coordinates,
+)
+from esm.utils.structure.predicted_aligned_error import (
+    compute_predicted_aligned_error,
+    compute_tm,
+)
+
+
+class RelativePositionEmbedding(nn.Module):
+    """
+    Embedding layer for relative position embeddings. `bins` is the number of positions relative
+    to the query position that are considered before clipping. For instance, if `bins=10`, then
+    the relative position embedding will have 21 positions, [-10, 10].
+    """
+
+    def __init__(self, bins, embedding_dim, init_std=0.02):
+        super().__init__()
+        self.bins = bins
+
+        self.embedding = torch.nn.Embedding(2 * bins + 2, embedding_dim)
+        self.embedding.weight.data.normal_(0, init_std)
+
+    def forward(self, query_residue_index, key_residue_index):
+        """
+        Input:
+          query_residue_index: (B, ) tensor of source indices (dytpe=torch.long)
+          key_residue_index: (B, L) tensor of target indices (dytpe=torch.long)
+        Output:
+          embeddings: B x L x embedding_dim tensor of embeddings
+        """
+
+        assert query_residue_index.dtype == torch.long
+        assert key_residue_index.dtype == torch.long
+        assert query_residue_index.ndim == 1
+        assert key_residue_index.ndim == 2
+
+        diff = key_residue_index - query_residue_index.unsqueeze(1)
+        diff = diff.clamp(-self.bins, self.bins)
+        diff = diff + self.bins + 1  # add 1 to adjust for padding index
+        output = self.embedding(diff)
+        return output
+
+
+class PairwisePredictionHead(nn.Module):
+    def __init__(
+        self,
+        input_dim: int,
+        downproject_dim: int,
+        hidden_dim: int,
+        n_bins: int,
+        bias: bool = True,
+        pairwise_state_dim: int = 0,
+    ):
+        super().__init__()
+        self.downproject = nn.Linear(input_dim, downproject_dim, bias=bias)
+        self.linear1 = nn.Linear(
+            downproject_dim + pairwise_state_dim, hidden_dim, bias=bias
+        )
+        self.activation_fn = nn.GELU()
+        self.norm = nn.LayerNorm(hidden_dim)
+        self.linear2 = nn.Linear(hidden_dim, n_bins, bias=bias)
+
+    def forward(self, x, pairwise: torch.Tensor | None = None):
+        """
+        Args:
+            x: [B x L x D]
+
+        Output:
+            [B x L x L x K]
+        """
+        x = self.downproject(x)
+        # Let x_i be a vector of size (B, D).
+        # Input is {x_1, ..., x_L} of size (B, L, D)
+        # Output is 2D where x_ij = cat([x_i * x_j, x_i - x_j])
+        q, k = x.chunk(2, dim=-1)
+
+        prod = q[:, None, :, :] * k[:, :, None, :]
+        diff = q[:, None, :, :] - k[:, :, None, :]
+        x_2d = [prod, diff]
+        if pairwise is not None:
+            x_2d.append(pairwise)
+        x = torch.cat(x_2d, dim=-1)
+        x = self.linear1(x)
+        x = self.activation_fn(x)
+        x = self.norm(x)
+        x = self.linear2(x)
+        return x
+
+
+class RegressionHead(nn.Module):
+    def __init__(self, embed_dim: int, output_dim: int):
+        super().__init__()
+        self.dense = nn.Linear(embed_dim, embed_dim)
+        self.activation_fn = nn.GELU()
+        self.norm = nn.LayerNorm(embed_dim)
+        self.output = nn.Linear(embed_dim, output_dim)
+
+    def forward(self, features):
+        x = self.dense(features)
+        x = self.activation_fn(x)
+        x = self.norm(x)
+        x = self.output(x)
+        return x
+
+
+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=torch.float32
+        )
+        self.v_bins = (bins[:-1] + bins[1:]) / 2
+
+    def log_prob(self, true):
+        # Shapes are:
+        #     self.probs: ... x bins
+        #     true      : ... (floating point # for target)
+        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.to(self.v_bins.dtype).softmax(-1) @ self.v_bins.unsqueeze(1)
+        ).squeeze(-1)
+
+    def median(self):
+        return self.v_bins[self.logits.max(-1).indices]
+
+
+class GeometricEncoderStack(TransformerStack):
+    def __init__(self, d_model, n_heads, v_heads, n_layers):
+        super().__init__(d_model, n_heads, v_heads, 0)
+        self.blocks = nn.ModuleList(
+            [
+                UnifiedTransformerBlock(
+                    d_model,
+                    n_heads,
+                    v_heads=v_heads,
+                    use_geom_attn=True,
+                    use_plain_attn=False,
+                    expansion_ratio=4,
+                    bias=True,
+                )
+                for i in range(n_layers)
+            ]
+        )
+        self.norm = nn.Identity()
+
+
+def batched_gather(data, inds, dim=0, no_batch_dims=0):
+    ranges = []
+    for i, s in enumerate(data.shape[:no_batch_dims]):
+        r = torch.arange(s)
+        r = r.view(*(*((1,) * i), -1, *((1,) * (len(inds.shape) - i - 1))))
+        ranges.append(r)
+
+    remaining_dims = [slice(None) for _ in range(len(data.shape) - no_batch_dims)]
+    remaining_dims[dim - no_batch_dims if dim >= 0 else dim] = inds
+    ranges.extend(remaining_dims)
+    return data[ranges]
+
+
+def node_gather(s: torch.Tensor, edges: torch.Tensor) -> torch.Tensor:
+    return batched_gather(s.unsqueeze(-3), edges, -2, no_batch_dims=len(s.shape) - 1)
+
+
+class StructureTokenEncoder(nn.Module):
+    def __init__(self, d_model, n_heads, v_heads, n_layers, d_out, n_codes):
+        super().__init__()
+        # We only support fully-geometric structure token encoders for now...
+        # setting n_layers_geom to something that's not n_layers won't work because
+        # sequence ID isn't supported fully in this repo for plain-old transformers
+        self.transformer = GeometricEncoderStack(d_model, n_heads, v_heads, n_layers)
+        self.pre_vq_proj = nn.Linear(d_model, d_out)
+        self.codebook = EMACodebook(n_codes, d_out)
+        self.relative_positional_embedding = RelativePositionEmbedding(
+            32, d_model, init_std=0.02
+        )
+        self.knn = 16
+
+    def encode_local_structure(
+        self,
+        coords: torch.Tensor,
+        affine: Affine3D,
+        attention_mask: torch.Tensor,
+        sequence_id: torch.Tensor | None,
+        affine_mask: torch.Tensor,
+        residue_index: torch.Tensor | None = None,
+    ):
+        """This function allows for a multi-layered encoder to encode tokens with a local receptive fields. The implementation is as follows:
+
+        1. Starting with (B, L) frames, we find the KNN in structure space. This now gives us (B, L, K) where the last dimension is the local
+        neighborhood of all (B, L) residues.
+        2. We reshape these frames to (B*L, K) so now we have a large batch of a bunch of local neighborhoods.
+        3. Pass the (B*L, K) local neighborhoods through a stack of geometric reasoning blocks, effectively getting all to all communication between
+        all frames in the local neighborhood.
+        4. This gives (B*L, K, d_model) embeddings, from which we need to get a single embedding per local neighborhood. We do this by simply
+        taking the embedding corresponding to the query node. This gives us (B*L, d_model) embeddings.
+        5. Reshape back to (B, L, d_model) embeddings
+        """
+        assert coords.size(-1) == 3 and coords.size(-2) == 3, "need N, CA, C"
+        with torch.no_grad():
+            knn_edges, _ = self.find_knn_edges(
+                coords,
+                ~attention_mask,
+                coord_mask=affine_mask,
+                sequence_id=sequence_id,
+                knn=self.knn,
+            )
+            B, L, E = knn_edges.shape
+
+            affine_tensor = affine.tensor  # for easier manipulation
+            T_D = affine_tensor.size(-1)
+            knn_affine_tensor = node_gather(affine_tensor, knn_edges)
+            knn_affine_tensor = knn_affine_tensor.view(-1, E, T_D).contiguous()
+            affine = Affine3D.from_tensor(knn_affine_tensor)
+            knn_sequence_id = (
+                node_gather(sequence_id.unsqueeze(-1), knn_edges).view(-1, E)
+                if sequence_id is not None
+                else torch.zeros(B * L, E, dtype=torch.int64, device=coords.device)
+            )
+            knn_affine_mask = node_gather(affine_mask.unsqueeze(-1), knn_edges).view(
+                -1, E
+            )
+            knn_chain_id = torch.zeros(
+                B * L, E, dtype=torch.int64, device=coords.device
+            )
+
+            if residue_index is None:
+                res_idxs = knn_edges.view(-1, E)
+            else:
+                res_idxs = node_gather(residue_index.unsqueeze(-1), knn_edges).view(
+                    -1, E
+                )
+
+        z = self.relative_positional_embedding(res_idxs[:, 0], res_idxs)
+
+        z, _, _ = self.transformer.forward(
+            x=z,
+            sequence_id=knn_sequence_id,
+            affine=affine,
+            affine_mask=knn_affine_mask,
+            chain_id=knn_chain_id,
+        )
+
+        # Unflatten the output and take the query node embedding, which will always be the first one because
+        # a node has distance 0 with itself and the KNN are sorted.
+        z = z.view(B, L, E, -1)
+        z = z[:, :, 0, :]
+
+        return z
+
+    @staticmethod
+    def find_knn_edges(
+        coords,
+        padding_mask,
+        coord_mask,
+        sequence_id: torch.Tensor | None = None,
+        knn: int | None = None,
+    ) -> tuple:
+        assert knn is not None, "Must specify a non-null knn to find_knn_edges"
+        # Coords are N, CA, C
+        coords = coords.clone()
+        coords[~coord_mask] = 0
+
+        if sequence_id is None:
+            sequence_id = torch.zeros(
+                (coords.shape[0], coords.shape[1]), device=coords.device
+            ).long()
+
+        with torch.no_grad(), torch.cuda.amp.autocast(enabled=False):  # type: ignore
+            ca = coords[..., 1, :]
+            edges, edge_mask = knn_graph(
+                ca, coord_mask, padding_mask, sequence_id, no_knn=knn
+            )
+
+        return edges, edge_mask
+
+    def encode(
+        self,
+        coords: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        sequence_id: torch.Tensor | None = None,
+        residue_index: torch.Tensor | None = None,
+    ):
+        coords = coords[..., :3, :]
+        affine, affine_mask = build_affine3d_from_coordinates(coords=coords)
+
+        if attention_mask is None:
+            attention_mask = torch.ones_like(affine_mask, dtype=torch.bool)
+        attention_mask = attention_mask.bool()
+
+        if sequence_id is None:
+            sequence_id = torch.zeros_like(affine_mask, dtype=torch.int64)
+
+        z = self.encode_local_structure(
+            coords=coords,
+            affine=affine,
+            attention_mask=attention_mask,
+            sequence_id=sequence_id,
+            affine_mask=affine_mask,
+            residue_index=residue_index,
+        )
+
+        z = z.masked_fill(~affine_mask.unsqueeze(2), 0)
+        z = self.pre_vq_proj(z)
+
+        z_q, min_encoding_indices, _ = self.codebook(z)
+
+        return z_q, min_encoding_indices
+
+
+class StructureTokenDecoder(nn.Module):
+    def __init__(self, d_model, n_heads, n_layers):
+        super().__init__()
+        self.decoder_channels = d_model
+
+        self.vqvae_codebook_size = C.VQVAE_CODEBOOK_SIZE
+        self.special_tokens = C.VQVAE_SPECIAL_TOKENS
+        self.max_pae_bin = C.VQVAE_MAX_PAE_BIN
+
+        self.embed = nn.Embedding(
+            self.vqvae_codebook_size + len(self.special_tokens), d_model
+        )
+        self.decoder_stack = TransformerStack(
+            d_model, n_heads, 1, n_layers, scale_residue=False, n_layers_geom=0
+        )
+
+        self.affine_output_projection = Dim6RotStructureHead(
+            self.decoder_channels, 10, predict_torsion_angles=False
+        )
+
+        direction_loss_bins = C.VQVAE_DIRECTION_LOSS_BINS
+        pae_bins = C.VQVAE_PAE_BINS
+        self.pairwise_bins = [
+            64,  # distogram
+            direction_loss_bins * 6,  # direction bins
+            pae_bins,  # predicted aligned error
+        ]
+        self.pairwise_classification_head = PairwisePredictionHead(
+            self.decoder_channels,
+            downproject_dim=128,
+            hidden_dim=128,
+            n_bins=sum(self.pairwise_bins),
+            bias=False,
+        )
+
+        plddt_bins = C.VQVAE_PLDDT_BINS
+        self.plddt_head = RegressionHead(
+            embed_dim=self.decoder_channels, output_dim=plddt_bins
+        )
+
+    def decode(
+        self,
+        structure_tokens: torch.Tensor,
+        attention_mask: torch.Tensor | None = None,
+        sequence_id: torch.Tensor | None = None,
+    ):
+        if attention_mask is None:
+            attention_mask = torch.ones_like(structure_tokens, dtype=torch.bool)
+
+        attention_mask = attention_mask.bool()
+        if sequence_id is None:
+            sequence_id = torch.zeros_like(structure_tokens, dtype=torch.int64)
+        # not supported for now
+        chain_id = torch.zeros_like(structure_tokens, dtype=torch.int64)
+
+        # check that BOS and EOS are set correctly
+        assert (
+            structure_tokens[:, 0].eq(self.special_tokens["BOS"]).all()
+        ), "First token in structure_tokens must be BOS token"
+        assert (
+            structure_tokens[
+                torch.arange(structure_tokens.shape[0]), attention_mask.sum(1) - 1
+            ]
+            .eq(self.special_tokens["EOS"])
+            .all()
+        ), "Last token in structure_tokens must be EOS token"
+        assert (
+            (structure_tokens < 0).sum() == 0
+        ), "All structure tokens set to -1 should be replaced with BOS, EOS, PAD, or MASK tokens by now, but that isn't the case!"
+
+        x = self.embed(structure_tokens)
+        # !!! NOTE: Attention mask is actually unused here so watch out
+        x, _, _ = self.decoder_stack.forward(
+            x, affine=None, affine_mask=None, sequence_id=sequence_id, chain_id=chain_id
+        )
+
+        tensor7_affine, bb_pred = self.affine_output_projection(
+            x, affine=None, affine_mask=torch.zeros_like(attention_mask)
+        )
+
+        pae, ptm = None, None
+        pairwise_logits = self.pairwise_classification_head(x)
+        _, _, pae_logits = [
+            (o if o.numel() > 0 else None)
+            for o in pairwise_logits.split(self.pairwise_bins, dim=-1)
+        ]
+
+        special_tokens_mask = structure_tokens >= min(self.special_tokens.values())
+        pae = compute_predicted_aligned_error(
+            pae_logits,  # type: ignore
+            aa_mask=~special_tokens_mask,
+            sequence_id=sequence_id,
+            max_bin=self.max_pae_bin,
+        )
+        # This might be broken for chainbreak tokens? We might align to the chainbreak
+        ptm = compute_tm(
+            pae_logits,  # type: ignore
+            aa_mask=~special_tokens_mask,
+            max_bin=self.max_pae_bin,
+        )
+
+        plddt_logits = self.plddt_head(x)
+        plddt_value = CategoricalMixture(
+            plddt_logits, bins=plddt_logits.shape[-1]
+        ).mean()
+
+        return dict(
+            tensor7_affine=tensor7_affine,
+            bb_pred=bb_pred,
+            plddt=plddt_value,
+            ptm=ptm,
+            predicted_aligned_error=pae,
+        )

Dyna-1/esm/pretrained.py

@@ -0,0 +1,132 @@
+from typing import Callable
+
+import torch
+import torch.nn as nn
+
+from esm.models.esm3 import ESM3
+from esm.models.esmc import ESMC
+from esm.models.function_decoder import FunctionTokenDecoder
+from esm.models.vqvae import (
+    StructureTokenDecoder,
+    StructureTokenEncoder,
+)
+from esm.tokenization import (
+    get_esm3_model_tokenizers,
+    get_esmc_model_tokenizers,
+)
+from esm.utils.constants.esm3 import data_root
+from esm.utils.constants.models import (
+    ESM3_FUNCTION_DECODER_V0,
+    ESM3_OPEN_SMALL,
+    ESM3_STRUCTURE_DECODER_V0,
+    ESM3_STRUCTURE_ENCODER_V0,
+    ESMC_300M,
+    ESMC_600M,
+)
+
+ModelBuilder = Callable[[torch.device | str], nn.Module]
+
+
+def ESM3_structure_encoder_v0(device: torch.device | str = "cpu"):
+    with torch.device(device):
+        model = StructureTokenEncoder(
+            d_model=1024, n_heads=1, v_heads=128, n_layers=2, d_out=128, n_codes=4096
+        ).eval()
+    state_dict = torch.load(
+        data_root("esm3") / "data/weights/esm3_structure_encoder_v0.pth",
+        map_location=device,
+    )
+    model.load_state_dict(state_dict)
+    return model
+
+
+def ESM3_structure_decoder_v0(device: torch.device | str = "cpu"):
+    with torch.device(device):
+        model = StructureTokenDecoder(d_model=1280, n_heads=20, n_layers=30).eval()
+    state_dict = torch.load(
+        data_root("esm3") / "data/weights/esm3_structure_decoder_v0.pth",
+        map_location=device,
+    )
+    model.load_state_dict(state_dict)
+    return model
+
+
+def ESM3_function_decoder_v0(device: torch.device | str = "cpu"):
+    with torch.device(device):
+        model = FunctionTokenDecoder().eval()
+    state_dict = torch.load(
+        data_root("esm3") / "data/weights/esm3_function_decoder_v0.pth",
+        map_location=device,
+    )
+    model.load_state_dict(state_dict)
+    return model
+
+
+def ESMC_300M_202412(device: torch.device | str = "cpu"):
+    with torch.device(device):
+        model = ESMC(
+            d_model=960, n_heads=15, n_layers=30, tokenizer=get_esmc_model_tokenizers()
+        ).eval()
+    state_dict = torch.load(
+        data_root("esmc-300") / "data/weights/esmc_300m_2024_12_v0.pth",
+        map_location=device,
+    )
+    model.load_state_dict(state_dict)
+
+    return model
+
+
+def ESMC_600M_202412(device: torch.device | str = "cpu"):
+    with torch.device(device):
+        model = ESMC(
+            d_model=1152, n_heads=18, n_layers=36, tokenizer=get_esmc_model_tokenizers()
+        ).eval()
+    state_dict = torch.load(
+        data_root("esmc-600") / "data/weights/esmc_600m_2024_12_v0.pth",
+        map_location=device,
+    )
+    model.load_state_dict(state_dict)
+
+    return model
+
+
+def ESM3_sm_open_v0(device: torch.device | str = "cpu"):
+    with torch.device(device):
+        model = ESM3(
+            d_model=1536,
+            n_heads=24,
+            v_heads=256,
+            n_layers=48,
+            structure_encoder_fn=ESM3_structure_encoder_v0,
+            structure_decoder_fn=ESM3_structure_decoder_v0,
+            function_decoder_fn=ESM3_function_decoder_v0,
+            tokenizers=get_esm3_model_tokenizers(ESM3_OPEN_SMALL),
+        ).eval()
+    state_dict = torch.load(
+        data_root("esm3") / "data/weights/esm3_sm_open_v1.pth", map_location=device
+    )
+    model.load_state_dict(state_dict)
+    return model
+
+
+LOCAL_MODEL_REGISTRY: dict[str, ModelBuilder] = {
+    ESM3_OPEN_SMALL: ESM3_sm_open_v0,
+    ESM3_STRUCTURE_ENCODER_V0: ESM3_structure_encoder_v0,
+    ESM3_STRUCTURE_DECODER_V0: ESM3_structure_decoder_v0,
+    ESM3_FUNCTION_DECODER_V0: ESM3_function_decoder_v0,
+    ESMC_600M: ESMC_600M_202412,
+    ESMC_300M: ESMC_300M_202412,
+}
+
+
+def load_local_model(
+    model_name: str, device: torch.device = torch.device("cpu")
+) -> nn.Module:
+    if model_name not in LOCAL_MODEL_REGISTRY:
+        raise ValueError(f"Model {model_name} not found in local model registry.")
+    return LOCAL_MODEL_REGISTRY[model_name](device)
+
+
+# Register custom versions of ESM3 for use with the local inference API
+def register_local_model(model_name: str, model_builder: ModelBuilder) -> None:
+    LOCAL_MODEL_REGISTRY[model_name] = model_builder

Dyna-1/esm/sdk/__init__.py

@@ -0,0 +1,22 @@
+import os
+
+from esm.sdk.forge import ESM3ForgeInferenceClient
+
+# Note: please do not import ESM3SageMakerClient here since that requires AWS SDK.
+
+
+def client(
+    model="esm3-sm-open-v1",
+    url="https://forge.evolutionaryscale.ai",
+    token=os.environ.get("ESM_API_KEY", ""),
+    request_timeout=None,
+):
+    """
+    Args:
+        model: Name of the model to use.
+        url: URL of a forge server.
+        token: User's API token.
+        request_timeout: Amount of time to wait for a request to finish.
+            Default is wait indefinitely.
+    """
+    return ESM3ForgeInferenceClient(model, url, token, request_timeout)

Dyna-1/esm/sdk/api.py

@@ -0,0 +1,445 @@
+from __future__ import annotations
+
+from abc import ABC
+from typing import List, Sequence
+
+import attr
+import torch
+from attr import asdict, define
+
+import esm.utils.constants.api as C
+from esm.tokenization import (
+    TokenizerCollectionProtocol,
+    get_esm3_model_tokenizers,
+)
+from esm.utils import encoding
+from esm.utils.constants.models import ESM3_OPEN_SMALL
+from esm.utils.misc import (
+    get_chainbreak_boundaries_from_sequence,
+)
+from esm.utils.structure.protein_chain import ProteinChain
+from esm.utils.structure.protein_complex import ProteinComplex
+from esm.utils.types import FunctionAnnotation, PathOrBuffer
+
+
+class ProteinType(ABC): ...
+
+
+## Basic Types
+@define
+class ESMProtein(ProteinType):
+    # Tracks
+    sequence: str | None = None
+    secondary_structure: str | None = None
+    sasa: list[float | None] | None = None
+    function_annotations: list[FunctionAnnotation] | None = None
+    coordinates: torch.Tensor | None = None
+
+    # Metrics
+    plddt: torch.Tensor | None = None
+    ptm: torch.Tensor | None = None
+
+
+    # When calling EvolutionaryScale API, use this flag to disclose any
+    # sequences that may potentially have concerns.
+    # Such sequences may not go through standard safety filter for approved users.
+    # Reach out if interested in using this.
+    potential_sequence_of_concern: bool = False
+
+    def __len__(self):
+        if self.sequence is not None:
+            return len(self.sequence)
+        elif self.secondary_structure is not None:
+            return len(self.secondary_structure)
+        elif self.sasa is not None:
+            return len(self.sasa)
+        elif self.coordinates is not None:
+            return self.coordinates.size(0)
+        else:
+            raise ValueError("No track to determine length from.")
+
+    @classmethod
+    def from_pdb(
+        cls,
+        path: PathOrBuffer,
+        chain_id: str = "detect",
+        id: str | None = None,
+        is_predicted: bool = False,
+    ) -> ESMProtein:
+        protein_chain = ProteinChain.from_pdb(
+            path=path, chain_id=chain_id, id=id, is_predicted=is_predicted
+        )
+        return cls.from_protein_chain(protein_chain)
+
+    @classmethod
+    def from_protein_chain(
+        cls, protein_chain: ProteinChain, with_annotations: bool = False
+    ) -> ESMProtein:
+        # By default, we don't annotate with DSSP / SASA, which are expensive.
+        # If mkdssp is installed, we can annotate with a flag.
+        if with_annotations:
+            return ESMProtein(
+                sequence=protein_chain.sequence,
+                secondary_structure=protein_chain.dssp().tolist(),
+                sasa=protein_chain.sasa().tolist(),
+                function_annotations=None,
+                coordinates=torch.tensor(protein_chain.atom37_positions),
+            )
+        else:
+            return ESMProtein(
+                sequence=protein_chain.sequence,
+                secondary_structure=None,
+                sasa=None,
+                function_annotations=None,
+                coordinates=torch.tensor(protein_chain.atom37_positions),
+            )
+
+    @classmethod
+    def from_protein_complex(
+        cls, protein_complex: ProteinComplex, with_annotations: bool = False
+    ) -> ESMProtein:
+        if with_annotations:
+            raise NotImplementedError(
+                "Annotations are not supported for ProteinComplex yet."
+            )
+
+        return ESMProtein(
+            sequence=protein_complex.sequence,
+            secondary_structure=None,
+            sasa=None,
+            function_annotations=None,
+            coordinates=torch.tensor(protein_complex.atom37_positions),
+        )
+
+    def to_pdb(self, pdb_path: PathOrBuffer) -> None:
+        # Note: Will work for single chains as well and produce same pdb file
+        protein_complex = self.to_protein_complex().infer_oxygen()
+        protein_complex.to_pdb(pdb_path)
+
+    def to_pdb_string(self) -> str:
+        protein_chain = self.to_protein_chain()
+        return protein_chain.to_pdb_string()
+
+    def to_protein_chain(self) -> ProteinChain:
+        if self.coordinates is None:
+            raise ValueError("Coordinates are required to convert to a ProteinChain.")
+        protein_chain = ProteinChain.from_atom37(
+            atom37_positions=self.coordinates.to("cpu").numpy(),
+            id=None,
+            sequence=None if self.sequence is None else self.sequence.replace("_", "X"),
+            chain_id=None,
+            entity_id=None,
+            residue_index=None,
+            insertion_code=None,
+            confidence=None
+            if self.plddt is None
+            else self.plddt.detach().cpu().numpy(),
+        )
+        return protein_chain
+
+    def to_protein_complex(
+        self, copy_annotations_from_ground_truth: ProteinComplex | None = None
+    ) -> ProteinComplex:
+        assert (
+            self.sequence is not None
+        ), "ESMProtein must have a sequence to convert to ProteinComplex"
+        assert (
+            self.coordinates is not None
+        ), "ESMProtein must have coordinates to convert to ProteinComplex"
+        coords = self.coordinates.to("cpu").numpy()
+
+        chain_boundaries = get_chainbreak_boundaries_from_sequence(self.sequence)
+        if copy_annotations_from_ground_truth is not None:
+            gt_chains = list(copy_annotations_from_ground_truth.chain_iter())
+        else:
+            gt_chains = None
+        pred_chains = []
+        for i, (start, end) in enumerate(chain_boundaries):
+            pred_chain = ProteinChain.from_atom37(
+                atom37_positions=coords[start:end],
+                sequence=self.sequence[start:end],
+                chain_id=gt_chains[i].chain_id if gt_chains is not None else None,
+                entity_id=gt_chains[i].entity_id if gt_chains is not None else None,
+            )
+            pred_chains.append(pred_chain)
+        return ProteinComplex.from_chains(pred_chains)
+
+
+@define
+class ESMProteinTensor(ProteinType):
+    sequence: torch.Tensor | None = None
+    structure: torch.Tensor | None = None
+    secondary_structure: torch.Tensor | None = None
+    sasa: torch.Tensor | None = None
+    function: torch.Tensor | None = None
+    residue_annotations: torch.Tensor | None = None
+    coordinates: torch.Tensor | None = None
+
+    # When calling EvolutionaryScale API, use this flag to disclose any
+    # sequences that may potentially have concerns.
+    # Such sequences may not go through standard safety filter for approved users.
+    # Reach out if interested in using this.
+    potential_sequence_of_concern: bool = False
+    # Control vectors are vectors added to each layer of the model to nudge hidden states to the desired direction.
+    # len(control_vectors)  == number of blocks in the model. Each vector in the list have the shape of (batch size, sequence length, hidden dim)
+    # so it can be added to the corresponding layer in the model
+
+    def _detect_attribute(self, func, msg):
+        mapped = {
+            k: func(k, v)
+            for k, v in asdict(self).items()
+            if isinstance(v, torch.Tensor)
+        }
+        s = set(mapped.values())
+        if len(s) <= 0:
+            return None
+        if len(s) != 1:
+            raise ValueError(f"Either no tracks or inconsistent {msg}: {mapped}")
+        return next(iter(s))
+
+    def __len__(self) -> int:
+        l = self._detect_attribute(lambda _, x: x.size(0), "length")
+        return l if l is not None else 0
+
+    @property
+    def device(self) -> str | torch.device:
+        d = self._detect_attribute(lambda _, x: x.device, "device")
+        assert d is not None
+        return d
+
+    def to(self, device_or_dtype: str | torch.device | torch.dtype) -> ESMProteinTensor:
+        def _to(name):
+            v = getattr(self, name)
+            if v is not None and isinstance(v, torch.Tensor):
+                setattr(self, name, v.to(device_or_dtype))
+
+        for n in attr.fields(ESMProteinTensor):
+            _to(n.name)
+
+        return self
+
+    @classmethod
+    def empty(
+        cls,
+        length: int,
+        tokenizers: TokenizerCollectionProtocol | None = None,
+        device: torch.device | str = "cpu",
+    ) -> ESMProteinTensor:
+        if tokenizers is None:
+            tokenizers = get_esm3_model_tokenizers(ESM3_OPEN_SMALL)
+
+        return ESMProteinTensor(
+            sequence=encoding.get_default_sequence_tokens(
+                length, tokenizers.sequence
+            ).to(device),
+            structure=encoding.get_default_structure_tokens(
+                length, tokenizers.structure
+            ).to(device),
+            secondary_structure=encoding.get_default_secondary_structure_tokens(
+                length, tokenizers.secondary_structure
+            ).to(device),
+            sasa=encoding.get_default_sasa_tokens(length, tokenizers.sasa).to(device),
+            function=encoding.get_default_function_tokens(
+                length, tokenizers.function
+            ).to(device),
+            residue_annotations=encoding.get_default_residue_annotation_tokens(
+                length, tokenizers.residue_annotations
+            ).to(device),
+        )
+
+
+@define
+class ESMProteinError(Exception, ProteinType):
+    error_code: int  # Error code follows HTTP convention, i.e., 404 NotFoundError, 500 InternalError.
+    error_msg: str
+
+
+## High Level Endpoint Types
+@define
+class GenerationConfig:
+    track: str = ""
+    invalid_ids: Sequence[int] = []
+    # Controls the number of tokens to unmask during each round of iterative generation.
+    schedule: str = attr.field(
+        validator=attr.validators.in_(["cosine", "linear"]), default="cosine"
+    )
+    # Controls which tokens to unmask during each round of iterative generation.
+    # "random" will unmask a correct number of tokens randomly.
+    # "entropy" will unmask the tokens with the lowest logit entropy first.
+    strategy: str = attr.field(
+        validator=attr.validators.in_(["random", "entropy"]), default="entropy"
+    )
+    # Set this to a higher value for better generation results.
+    # Note that this needs to be less than or equal to the sequence length.
+    num_steps: int = 1
+    temperature: float = 1.0
+    temperature_annealing: bool = False
+    top_p: float = 1.0
+    condition_on_coordinates_only: bool = True
+
+    def use_entropy_based_unmasking_strategy(self):
+        """Use entropy based unmasking strategy during generation."""
+        self.schedule = "cosine"
+        self.strategy = "entropy"
+        self.temperature_annealing = False
+
+    def use_generative_unmasking_strategy(self):
+        """Use an unmasking strategy that produces more variety of generations."""
+        self.schedule = "cosine"
+        self.strategy = "random"
+        self.temperature_annealing = True
+
+
+@define
+class InverseFoldingConfig:
+    invalid_ids: Sequence[int] = []
+    temperature: float = 1.0
+
+
+## Low Level Endpoint Types
+@define
+class SamplingTrackConfig:
+    temperature: float = 1.0
+    top_p: float = 1.0
+    only_sample_masked_tokens: bool = True
+    invalid_ids: Sequence[int] = []
+    topk_logprobs: int = 0
+
+
+@define
+class SamplingConfig:
+    sequence: SamplingTrackConfig | None = attr.field(
+        default=None, metadata={"max_topk": C.MAX_TOPK_SEQUENCE}
+    )
+    structure: SamplingTrackConfig | None = attr.field(
+        default=None, metadata={"max_topk": C.MAX_TOPK_STRUCTURE}
+    )
+    secondary_structure: SamplingTrackConfig | None = attr.field(
+        default=None, metadata={"max_topk": C.MAX_TOPK_SECONDARY_STRUCTURE}
+    )
+    sasa: SamplingTrackConfig | None = attr.field(
+        default=None, metadata={"max_topk": C.MAX_TOPK_SASA}
+    )
+    function: SamplingTrackConfig | None = attr.field(
+        default=None, metadata={"max_topk": C.MAX_TOPK_FUNCTION}
+    )
+
+    return_per_residue_embeddings: bool = False
+    return_mean_embedding: bool = False
+
+
+@define
+class ForwardTrackData:
+    sequence: torch.Tensor | None = None
+    structure: torch.Tensor | None = None
+    secondary_structure: torch.Tensor | None = None
+    sasa: torch.Tensor | None = None
+    function: torch.Tensor | None = None
+
+
+@define
+class LogitsConfig:
+    # Logits.
+    sequence: bool = False
+    structure: bool = False
+    secondary_structure: bool = False
+    sasa: bool = False
+    function: bool = False
+    residue_annotations: bool = False
+
+    # Embeddings.
+    return_embeddings: bool = False
+
+
+@define
+class LogitsOutput:
+    logits: ForwardTrackData | None = None
+    embeddings: torch.Tensor | None = None
+
+    # Residue annotations is multi-hot, so deserves special treatment
+    # It's not a categorical distribution, but instead a bernoulli, so
+    # softmax across the last dimension is _wrong_
+    residue_annotation_logits: torch.Tensor | None = None
+
+
+@define
+class ForwardAndSampleOutput(LogitsOutput):
+    protein_tensor: ESMProteinTensor = ESMProteinTensor()
+
+    entropy: ForwardTrackData | None = None
+    # Probability of sampled token
+    prob: ForwardTrackData | None = None
+    logprob: ForwardTrackData | None = None
+    # Top probability at this position
+    top_prob: ForwardTrackData | None = None
+    topk_logprob: ForwardTrackData | None = None
+    # Which tokens correspond to top probability
+    topk_tokens: ForwardTrackData | None = None
+    per_residue_embedding: torch.Tensor | None = None
+    mean_embedding: torch.Tensor | None = None
+
+
+class ESM3InferenceClient(ABC):
+    def generate(self, input: ProteinType, config: GenerationConfig) -> ProteinType:
+        # This is the easiest and most flexible way to run ESM3. Generate will
+        # iteratively sample tokens an provide an output with the track specified
+        # completely filled out, according to the GenerationConfig provided.
+        # It is a local function wrapping calls for encode -> iterative_sampling -> decode.
+        # if a ESMProteinTensor is provided, encode and decode are skipped
+        raise NotImplementedError
+
+    def batch_generate(
+        self, inputs: Sequence[ProteinType], configs: Sequence[GenerationConfig]
+    ) -> Sequence[ProteinType]:
+        # Same as generate(...), but generates a batch of proteins at once.
+        raise NotImplementedError
+
+    def encode(self, input: ESMProtein) -> ESMProteinTensor:
+        # Encode allows for encoding RawRepresentation into TokenizedRepresentation.
+        # This runs the structure_token_encoder, as well as dealing with PDB => atom37 conversion
+        raise NotImplementedError
+
+    def decode(self, input: ESMProteinTensor) -> ESMProtein:
+        # Decode is the inverse of encode, and runs a structure_token_decoder to output coordinates
+        raise NotImplementedError
+
+    def logits(
+        self, input: ESMProteinTensor, config: LogitsConfig = LogitsConfig()
+    ) -> LogitsOutput:
+        # Our API generally discourages using raw forwards.
+        # This is because sending logits can be prohibitively expensive.
+        # Please use forward_and_sample instead.
+        raise NotImplementedError
+
+    def forward_and_sample(
+        self, input: ESMProteinTensor, sampling_configuration: SamplingConfig
+    ) -> ForwardAndSampleOutput:
+        # forward_and_sample runs a single model forward, sampling tokens according to `SamplingConfiguration`.
+        # This is the way for power users to run ESM3. We hope to design this in a way to enable high throughput
+        # inference, as well as arbitrary chain-of-though invocations of ESM3.
+        raise NotImplementedError
+
+    @property
+    def raw_model(self):
+        # Get underlying esm3 model of an inference client.
+        raise NotImplementedError
+
+
+class ESMCInferenceClient(ABC):
+    def encode(self, input: ESMProtein) -> ESMProteinTensor:
+        # Encode allows for encoding RawRepresentation into TokenizedRepresentation.
+        raise NotImplementedError
+
+    def decode(self, input: ESMProteinTensor) -> ESMProtein:
+        # Decode is the inverse of encode
+        raise NotImplementedError
+
+    def logits(
+        self, input: ESMProteinTensor, config: LogitsConfig = LogitsConfig()
+    ) -> LogitsOutput:
+        raise NotImplementedError
+
+    @property
+    def raw_model(self):
+        # Get underlying esmc model of an inference client.
+        raise NotImplementedError

Dyna-1/esm/sdk/forge.py

@@ -0,0 +1,580 @@
+import asyncio
+from functools import wraps
+from typing import Sequence
+from urllib.parse import urljoin
+
+import requests
+import torch
+from tenacity import retry, retry_if_result, stop_after_attempt, wait_exponential
+
+from esm.sdk.api import (
+    ESM3InferenceClient,
+    ESMProtein,
+    ESMProteinError,
+    ESMProteinTensor,
+    ForwardAndSampleOutput,
+    ForwardTrackData,
+    GenerationConfig,
+    InverseFoldingConfig,
+    LogitsConfig,
+    LogitsOutput,
+    ProteinType,
+    SamplingConfig,
+    SamplingTrackConfig,
+)
+from esm.utils.misc import maybe_list, maybe_tensor
+from esm.utils.sampling import validate_sampling_config
+from esm.utils.types import FunctionAnnotation
+
+
+def _list_to_function_annotations(l) -> list[FunctionAnnotation] | None:
+    if l is None or len(l) <= 0:
+        return None
+    return [FunctionAnnotation(*t) for t in l]
+
+
+def retry_if_specific_error(exception):
+    """
+    We only retry on specific errors.
+    Currently we retry for 502 (bad gateway) and 429 (rate limit)
+    """
+    return isinstance(exception, ESMProteinError) and exception.error_code in {
+        429,
+        502,
+        504,
+    }
+
+
+def log_retry_attempt(retry_state):
+    print(
+        f"Retrying... Attempt {retry_state.attempt_number} after {retry_state.next_action.sleep}s due to: {retry_state.outcome.result()}"
+    )
+
+
+def _validate_protein_tensor_input(input):
+    if not isinstance(input, ESMProteinTensor):
+        raise ValueError(
+            "Input must be an ESMProteinTensor instance. "
+            "Use encode() API to encode an ESMProtein into ESMProteinTensor."
+        )
+
+
+class SequenceStructureForgeInferenceClient:
+    def __init__(
+        self,
+        url: str = "https://forge.evolutionaryscale.ai",
+        token: str = "",
+        request_timeout: int | None = None,
+    ):
+        if token == "":
+            raise RuntimeError(
+                "Please provide a token to connect to Forge via token=YOUR_API_TOKEN_HERE"
+            )
+        self.url = url
+        self.token = token
+        self.headers = {"Authorization": f"Bearer {self.token}"}
+        self.request_timeout = request_timeout
+
+    def fold(
+        self,
+        sequence: str,
+        potential_sequence_of_concern: bool,
+        model_name: str | None = None,
+    ) -> ESMProtein | ESMProteinError:
+        request = {"sequence": sequence}
+        if model_name is not None:
+            request["model"] = model_name
+        try:
+            data = self._post("fold", request, potential_sequence_of_concern)
+        except ESMProteinError as e:
+            return e
+
+        return ESMProtein(
+            coordinates=maybe_tensor(data["coordinates"], convert_none_to_nan=True)
+        )
+
+    def inverse_fold(
+        self,
+        coordinates: torch.Tensor,
+        config: InverseFoldingConfig,
+        potential_sequence_of_concern: bool,
+        model_name: str | None = None,
+    ) -> ESMProtein | ESMProteinError:
+        inverse_folding_config = {
+            "invalid_ids": config.invalid_ids,
+            "temperature": config.temperature,
+        }
+        request = {
+            "coordinates": maybe_list(coordinates, convert_nan_to_none=True),
+            "inverse_folding_config": inverse_folding_config,
+        }
+        if model_name is not None:
+            request["model"] = model_name
+        try:
+            data = self._post("inverse_fold", request, potential_sequence_of_concern)
+        except ESMProteinError as e:
+            return e
+
+        return ESMProtein(sequence=data["sequence"])
+
+    def _post(self, endpoint, request, potential_sequence_of_concern):
+        request["potential_sequence_of_concern"] = potential_sequence_of_concern
+
+        response = requests.post(
+            urljoin(self.url, f"/api/v1/{endpoint}"),
+            json=request,
+            headers=self.headers,
+            timeout=self.request_timeout,
+        )
+
+        if not response.ok:
+            raise ESMProteinError(
+                error_code=response.status_code,
+                error_msg=f"Failure in {endpoint}: {response.text}",
+            )
+
+        data = response.json()
+        # Nextjs puts outputs dict under "data" key.
+        # Lift it up for easier downstream processing.
+        if "outputs" not in data and "data" in data:
+            data = data["data"]
+
+        # Print warning message if there is any.
+        if "warning_messages" in data and data["warning_messages"] is not None:
+            for msg in data["warning_messages"]:
+                print("\033[31m", msg, "\033[0m")
+
+        return data
+
+
+class ESM3ForgeInferenceClient(ESM3InferenceClient):
+    def __init__(
+        self,
+        model: str,
+        url: str = "https://forge.evolutionaryscale.ai",
+        token: str = "",
+        request_timeout: int | None = None,
+        min_retry_wait: int = 1,
+        max_retry_wait: int = 10,
+        max_retry_attempts: int = 5,
+    ):
+        if token == "":
+            raise RuntimeError(
+                "Please provide a token to connect to Forge via token=YOUR_API_TOKEN_HERE"
+            )
+        self.model = model  # Name of the model to run.
+        self.url = url
+        self.token = token
+        self.headers = {"Authorization": f"Bearer {self.token}"}
+        self.request_timeout = request_timeout
+        self.min_retry_wait = min_retry_wait
+        self.max_retry_wait = max_retry_wait
+        self.max_retry_attempts = max_retry_attempts
+
+    @staticmethod
+    def retry_decorator(func):
+        """
+        A static method that returns a retry decorator. This decorator uses the
+        instance's retry settings.
+        """
+
+        @wraps(func)
+        def wrapper(instance, *args, **kwargs):
+            retry_decorator = retry(
+                retry=retry_if_result(retry_if_specific_error),
+                wait=wait_exponential(
+                    multiplier=1,
+                    min=instance.min_retry_wait,
+                    max=instance.max_retry_wait,
+                ),
+                stop=stop_after_attempt(instance.max_retry_attempts),
+                before_sleep=log_retry_attempt,
+            )
+            # Apply the retry decorator to the function
+            return retry_decorator(func)(instance, *args, **kwargs)
+
+        return wrapper
+
+    @retry_decorator
+    def generate(self, input: ProteinType, config: GenerationConfig) -> ProteinType:
+        if isinstance(input, ESMProtein):
+            output = self.__generate_protein(input, config)
+        elif isinstance(input, ESMProteinTensor):
+            output = self.__generate_protein_tensor(input, config)
+        else:
+            return ESMProteinError(
+                error_code=500, error_msg=f"Unknown input type {type(input)}"
+            )
+
+        if (
+            isinstance(output, ESMProtein)
+            and isinstance(input, ESMProtein)
+            and config.track not in ["function", "residue_annotations"]
+        ):
+            # Function and residue annotation encoding/decoding is lossy
+            # There is no guarantee that decoding encoded tokens will yield the same input
+            output.function_annotations = input.function_annotations
+
+        return output
+
+    def batch_generate(
+        self, inputs: Sequence[ProteinType], configs: Sequence[GenerationConfig]
+    ) -> Sequence[ProteinType]:
+        """Forge supports auto-batching. So batch_generate() for the Forge client
+        is as simple as running a collection of generate() in parallel using asyncio.
+        """
+        loop = asyncio.get_event_loop()
+
+        async def _async_generate():
+            futures = [
+                loop.run_in_executor(None, self.generate, protein, config)
+                for protein, config in zip(inputs, configs)
+            ]
+            return await asyncio.gather(*futures, return_exceptions=True)
+
+        results = loop.run_until_complete(_async_generate())
+
+        def _capture_exception(r):
+            if isinstance(r, BaseException) and not isinstance(r, ESMProteinError):
+                return ESMProteinError(500, str(r))
+            return r
+
+        return [_capture_exception(r) for r in results]
+
+    def __generate_protein(
+        self, input: ESMProtein, config: GenerationConfig
+    ) -> ESMProtein | ESMProteinError:
+        req = {}
+        req["sequence"] = input.sequence
+        req["secondary_structure"] = input.secondary_structure
+        req["sasa"] = input.sasa
+        if input.function_annotations is not None:
+            req["function"] = [x.to_tuple() for x in input.function_annotations]
+        req["coordinates"] = maybe_list(input.coordinates, convert_nan_to_none=True)
+
+        request = {
+            "model": self.model,
+            "inputs": req,
+            "track": config.track,
+            "invalid_ids": config.invalid_ids,
+            "schedule": config.schedule,
+            "num_steps": config.num_steps,
+            "temperature": config.temperature,
+            "top_p": config.top_p,
+            "condition_on_coordinates_only": config.condition_on_coordinates_only,
+        }
+        try:
+            data = self._post("generate", request, input.potential_sequence_of_concern)
+        except ESMProteinError as e:
+            return e
+
+        return ESMProtein(
+            sequence=data["outputs"]["sequence"],
+            secondary_structure=data["outputs"]["secondary_structure"],
+            sasa=data["outputs"]["sasa"],
+            function_annotations=_list_to_function_annotations(
+                data["outputs"]["function"]
+            ),
+            coordinates=maybe_tensor(
+                data["outputs"]["coordinates"], convert_none_to_nan=True
+            ),
+            plddt=maybe_tensor(data["outputs"]["plddt"]),
+            ptm=maybe_tensor(data["outputs"]["ptm"]),
+        )
+
+    def __generate_protein_tensor(
+        self, input: ESMProteinTensor, config: GenerationConfig
+    ) -> ESMProteinTensor | ESMProteinError:
+        req = {}
+        req["sequence"] = maybe_list(input.sequence)
+        req["structure"] = maybe_list(input.structure)
+        req["secondary_structure"] = maybe_list(input.secondary_structure)
+        req["sasa"] = maybe_list(input.sasa)
+        req["function"] = maybe_list(input.function)
+        req["coordinates"] = maybe_list(input.coordinates, convert_nan_to_none=True)
+        req["residue_annotation"] = maybe_list(input.residue_annotations)
+
+        request = {
+            "model": self.model,
+            "inputs": req,
+            "track": config.track,
+            "invalid_ids": config.invalid_ids,
+            "schedule": config.schedule,
+            "num_steps": config.num_steps,
+            "temperature": config.temperature,
+            "top_p": config.top_p,
+            "condition_on_coordinates_only": config.condition_on_coordinates_only,
+        }
+
+        try:
+            data = self._post(
+                "generate_tensor", request, input.potential_sequence_of_concern
+            )
+        except ESMProteinError as e:
+            return e
+
+        def _field_to_tensor(field, convert_none_to_nan: bool = False):
+            if field not in data["outputs"]:
+                return None
+            return maybe_tensor(
+                data["outputs"][field], convert_none_to_nan=convert_none_to_nan
+            )
+
+        output = ESMProteinTensor(
+            sequence=_field_to_tensor("sequence"),
+            structure=_field_to_tensor("structure"),
+            secondary_structure=_field_to_tensor("secondary_structure"),
+            sasa=_field_to_tensor("sasa"),
+            function=_field_to_tensor("function"),
+            residue_annotations=_field_to_tensor("residue_annotation"),
+            coordinates=_field_to_tensor("coordinates", convert_none_to_nan=True),
+        )
+
+        return output
+
+    @retry_decorator
+    def forward_and_sample(
+        self, input: ESMProteinTensor, sampling_configuration: SamplingConfig
+    ) -> ForwardAndSampleOutput | ESMProteinError:
+        _validate_protein_tensor_input(input)
+        validate_sampling_config(sampling_configuration, on_invalid="raise")
+
+        req = {}
+        sampling_config = {}
+        embedding_config = {
+            "sequence": sampling_configuration.return_mean_embedding,
+            "per_residue": sampling_configuration.return_per_residue_embeddings,
+        }
+
+        req["sequence"] = maybe_list(input.sequence)
+        req["structure"] = maybe_list(input.structure)
+        req["secondary_structure"] = maybe_list(input.secondary_structure)
+        req["sasa"] = maybe_list(input.sasa)
+        req["function"] = maybe_list(input.function)
+        req["coordinates"] = maybe_list(input.coordinates, convert_nan_to_none=True)
+        req["residue_annotation"] = maybe_list(input.residue_annotations)
+
+        def do_track(t: str):
+            track: SamplingTrackConfig | None
+            if (track := getattr(sampling_configuration, t, None)) is None:
+                sampling_config[t] = None
+            else:
+                sampling_config[t] = {
+                    "temperature": track.temperature,
+                    "top_p": track.top_p,
+                    "only_sample_masked_tokens": track.only_sample_masked_tokens,
+                    "invalid_ids": track.invalid_ids,
+                    "topk_logprobs": track.topk_logprobs,
+                }
+
+        do_track("sequence")
+        do_track("structure")
+        do_track("secondary_structure")
+        do_track("sasa")
+        do_track("function")
+
+        request = {
+            "model": self.model,
+            "inputs": req,
+            "sampling_config": sampling_config,
+            "embedding_config": embedding_config,
+        }
+        try:
+            data = self._post(
+                "forward_and_sample", request, input.potential_sequence_of_concern
+            )
+        except ESMProteinError as e:
+            return e
+
+        def get(k, field):
+            if data[k] is None:
+                return None
+            v = data[k][field]
+            return torch.tensor(v) if v is not None else None
+
+        tokens = ESMProteinTensor(
+            sequence=get("sequence", "tokens"),
+            structure=get("structure", "tokens"),
+            secondary_structure=get("secondary_structure", "tokens"),
+            sasa=get("sasa", "tokens"),
+            function=get("function", "tokens"),
+        )
+
+        def get_track(field):
+            return ForwardTrackData(
+                sequence=get("sequence", field),
+                structure=get("structure", field),
+                secondary_structure=get("secondary_structure", field),
+                sasa=get("sasa", field),
+                function=get("function", field),
+            )
+
+        def operate_on_track(track: ForwardTrackData, fn):
+            apply = lambda x: fn(x) if x is not None else None
+            return ForwardTrackData(
+                sequence=apply(track.sequence),
+                structure=apply(track.structure),
+                secondary_structure=apply(track.secondary_structure),
+                sasa=apply(track.sasa),
+                function=apply(track.function),
+            )
+
+        logprob = get_track("logprobs")
+        output = ForwardAndSampleOutput(
+            protein_tensor=tokens,
+            logprob=logprob,
+            prob=operate_on_track(logprob, torch.exp),
+            entropy=get_track("entropy"),
+            topk_logprob=get_track("topk_logprobs"),
+            topk_tokens=get_track("topk_tokens"),
+            per_residue_embedding=data["embeddings"]["per_residue"],
+            mean_embedding=data["embeddings"]["sequence"],
+        )
+        return output
+
+    @retry_decorator
+    def encode(self, input: ESMProtein) -> ESMProteinTensor | ESMProteinError:
+        tracks = {}
+        tracks["sequence"] = input.sequence
+        tracks["secondary_structure"] = input.secondary_structure
+        tracks["sasa"] = input.sasa
+        if input.function_annotations is not None:
+            tracks["function"] = [x.to_tuple() for x in input.function_annotations]
+        tracks["coordinates"] = maybe_list(input.coordinates, convert_nan_to_none=True)
+
+        request = {"inputs": tracks, "model": self.model}
+
+        try:
+            data = self._post("encode", request, input.potential_sequence_of_concern)
+        except ESMProteinError as e:
+            return e
+
+        return ESMProteinTensor(
+            sequence=maybe_tensor(data["outputs"]["sequence"]),
+            structure=maybe_tensor(data["outputs"]["structure"]),
+            coordinates=maybe_tensor(
+                data["outputs"]["coordinates"], convert_none_to_nan=True
+            ),
+            secondary_structure=maybe_tensor(data["outputs"]["secondary_structure"]),
+            sasa=maybe_tensor(data["outputs"]["sasa"]),
+            function=maybe_tensor(data["outputs"]["function"]),
+            residue_annotations=maybe_tensor(data["outputs"]["residue_annotation"]),
+        )
+
+    @retry_decorator
+    def decode(self, input: ESMProteinTensor) -> ESMProtein | ESMProteinError:
+        _validate_protein_tensor_input(input)
+
+        tokens = {}
+        tokens["sequence"] = maybe_list(input.sequence)
+        tokens["structure"] = maybe_list(input.structure)
+        tokens["secondary_structure"] = maybe_list(input.secondary_structure)
+        tokens["sasa"] = maybe_list(input.sasa)
+        tokens["function"] = maybe_list(input.function)
+        tokens["residue_annotation"] = maybe_list(input.residue_annotations)
+        tokens["coordinates"] = maybe_list(input.coordinates, convert_nan_to_none=True)
+
+        request = {"model": self.model, "inputs": tokens}
+
+        try:
+            data = self._post("decode", request, input.potential_sequence_of_concern)
+        except ESMProteinError as e:
+            return e
+
+        return ESMProtein(
+            sequence=data["outputs"]["sequence"],
+            secondary_structure=data["outputs"]["secondary_structure"],
+            sasa=data["outputs"]["sasa"],
+            function_annotations=_list_to_function_annotations(
+                data["outputs"]["function"]
+            ),
+            coordinates=maybe_tensor(
+                data["outputs"]["coordinates"], convert_none_to_nan=True
+            ),
+            plddt=maybe_tensor(data["outputs"]["plddt"]),
+            ptm=maybe_tensor(data["outputs"]["ptm"]),
+        )
+
+    @retry_decorator
+    def logits(
+        self, input: ESMProteinTensor, config: LogitsConfig = LogitsConfig()
+    ) -> LogitsOutput | ESMProteinError:
+        _validate_protein_tensor_input(input)
+
+        # Note: using raw model forwards is discouraged because of the byte size
+        # of the logits.
+        # Please use forward_and_sample instead.
+        req = {}
+        req["sequence"] = maybe_list(input.sequence)
+        req["structure"] = maybe_list(input.structure)
+        req["secondary_structure"] = maybe_list(input.secondary_structure)
+        req["sasa"] = maybe_list(input.sasa)
+        req["function"] = maybe_list(input.function)
+        req["coordinates"] = maybe_list(input.coordinates, convert_nan_to_none=True)
+        req["residue_annotation"] = maybe_list(input.residue_annotations)
+
+        logits_config = {
+            "sequence": config.sequence,
+            "structure": config.structure,
+            "secondary_structure": config.secondary_structure,
+            "sasa": config.sasa,
+            "function": config.function,
+            "residue_annotations": config.residue_annotations,
+            "return_embeddings": config.return_embeddings,
+        }
+
+        request = {"model": self.model, "inputs": req, "logits_config": logits_config}
+
+        try:
+            data = self._post("logits", request, input.potential_sequence_of_concern)
+        except ESMProteinError as e:
+            return e
+
+        def _maybe_logits(track: str):
+            if "logits" in data and track in data["logits"]:
+                return maybe_tensor(data["logits"][track])
+            return None
+
+        output = LogitsOutput(
+            logits=ForwardTrackData(
+                sequence=_maybe_logits("sequence"),
+                structure=_maybe_logits("structure"),
+                secondary_structure=_maybe_logits("secondary_structure"),
+                sasa=_maybe_logits("sasa"),
+                function=_maybe_logits("function"),
+            ),
+            embeddings=maybe_tensor(data["embeddings"]),
+            residue_annotation_logits=_maybe_logits("residue_annotation"),
+        )
+
+        return output
+
+    def _post(self, endpoint, request, potential_sequence_of_concern):
+        request["potential_sequence_of_concern"] = potential_sequence_of_concern
+
+        response = requests.post(
+            urljoin(self.url, f"/api/v1/{endpoint}"),
+            json=request,
+            headers=self.headers,
+            timeout=self.request_timeout,
+        )
+
+        if not response.ok:
+            raise ESMProteinError(
+                error_code=response.status_code,
+                error_msg=f"Failure in {endpoint}: {response.text}",
+            )
+
+        data = response.json()
+        # Nextjs puts outputs dict under "data" key.
+        # Lift it up for easier downstream processing.
+        if "outputs" not in data and "data" in data:
+            data = data["data"]
+
+        return data
+
+    @property
+    def raw_model(self):
+        raise NotImplementedError(
+            f"Can not get underlying remote model {self.model} from a Forge client."
+        )

Dyna-1/esm/sdk/sagemaker.py

@@ -0,0 +1,110 @@
+import json
+
+import boto3
+
+from esm.sdk.forge import (
+    ESM3ForgeInferenceClient,
+    SequenceStructureForgeInferenceClient,
+)
+
+
+class SequenceStructureSageMakerClient(SequenceStructureForgeInferenceClient):
+    def __init__(self, endpoint_name: str):
+        """SequenceStructure (folding and inverse folding) client that talks to a SageMaker endpoint.
+
+        Args:
+            endpoint_name: Name of the SageMaker endpoint.
+        """
+        # Dummy URL and token to make SequenceStructureForgeInferenceClient happy.
+        super().__init__(url="", token="dummy")
+
+        self._endpoint_name = endpoint_name
+
+        self._client = boto3.client(service_name="sagemaker-runtime")
+
+    def _post(self, endpoint, request, potential_sequence_of_concern):
+        request["potential_sequence_of_concern"] = potential_sequence_of_concern
+        request["model"] = request.get("model", None)
+        invocations_request = {
+            # Duplicate these fields at the top level to make Forge requests consistent.
+            "model": request["model"],
+            "request_id": "",  # Forge specific field.
+            "user_id": "",  # Forge specific field.
+            # Invocation data bits.
+            "api_ver": "v1",  # Must be v1 right now.
+            "endpoint": endpoint,
+            # Wrapped request.
+            endpoint: request,
+        }
+
+        try:
+            response = self._client.invoke_endpoint(
+                EndpointName=self._endpoint_name,
+                ContentType="application/json",
+                Body=json.dumps(invocations_request),
+            )
+        except Exception as e:
+            raise RuntimeError(f"Failure in {endpoint}: {e}") from e
+
+        data = json.loads(response["Body"].read().decode())
+
+        # Response must match request.
+        assert (
+            data["endpoint"] == endpoint
+        ), f"Response endpoint is {data['endpoint']} but request is {endpoint}"
+
+        # Get the actual responses under the endpoint key.
+        data = data[endpoint]
+
+        return data
+
+
+class ESM3SageMakerClient(ESM3ForgeInferenceClient):
+    def __init__(self, endpoint_name: str, model: str):
+        """ESM3 client that talks to a SageMaker endpoint.
+
+        Args:
+            endpoint_name: Name of the SageMaker endpoint.
+            model: Name of the ESM3 model.
+        """
+        # Dummy URL and token to make ESM3ForgeInferenceClient happy.
+        super().__init__(model=model, url="", token="dummy")
+
+        self._endpoint_name = endpoint_name
+        self._model = model
+
+        self._client = boto3.client(service_name="sagemaker-runtime")
+
+    def _post(self, endpoint, request, potential_sequence_of_concern):
+        request["potential_sequence_of_concern"] = potential_sequence_of_concern
+
+        invocations_request = {
+            # Duplicate these fields at the top level to make Forge requests consistent.
+            "model": request["model"],
+            "request_id": "",  # Forge specific field.
+            "user_id": "",  # Forge specific field.
+            # Invocation data bits.
+            "api_ver": "v1",  # Must be v1 right now.
+            "endpoint": endpoint,
+            # Wrapped request.
+            endpoint: request,
+        }
+
+        try:
+            response = self._client.invoke_endpoint(
+                EndpointName=self._endpoint_name,
+                ContentType="application/json",
+                Body=json.dumps(invocations_request),
+            )
+        except Exception as e:
+            raise RuntimeError(f"Failure in {endpoint}: {e}")
+
+        data = json.loads(response["Body"].read().decode())
+
+        # Response must match request.
+        assert data["endpoint"] == endpoint
+
+        # Get the actual responses under the endpoint key.
+        data = data[endpoint]
+
+        return data

Dyna-1/esm/tokenization/__init__.py

@@ -0,0 +1,69 @@
+from dataclasses import dataclass
+from typing import Protocol
+
+from esm.utils.constants.models import (
+    ESM3_OPEN_SMALL,
+    normalize_model_name,
+)
+
+from .function_tokenizer import InterProQuantizedTokenizer
+from .residue_tokenizer import ResidueAnnotationsTokenizer
+from .sasa_tokenizer import SASADiscretizingTokenizer
+from .sequence_tokenizer import EsmSequenceTokenizer
+from .ss_tokenizer import SecondaryStructureTokenizer
+from .structure_tokenizer import StructureTokenizer
+from .tokenizer_base import EsmTokenizerBase
+
+
+class TokenizerCollectionProtocol(Protocol):
+    sequence: EsmSequenceTokenizer
+    structure: StructureTokenizer
+    secondary_structure: SecondaryStructureTokenizer
+    sasa: SASADiscretizingTokenizer
+    function: InterProQuantizedTokenizer
+    residue_annotations: ResidueAnnotationsTokenizer
+
+
+@dataclass
+class TokenizerCollection:
+    sequence: EsmSequenceTokenizer
+    structure: StructureTokenizer
+    secondary_structure: SecondaryStructureTokenizer
+    sasa: SASADiscretizingTokenizer
+    function: InterProQuantizedTokenizer
+    residue_annotations: ResidueAnnotationsTokenizer
+
+
+def get_esm3_model_tokenizers(model: str = ESM3_OPEN_SMALL) -> TokenizerCollection:
+    if normalize_model_name(model) == ESM3_OPEN_SMALL:
+        return TokenizerCollection(
+            sequence=EsmSequenceTokenizer(),
+            structure=StructureTokenizer(),
+            secondary_structure=SecondaryStructureTokenizer(kind="ss8"),
+            sasa=SASADiscretizingTokenizer(),
+            function=InterProQuantizedTokenizer(),
+            residue_annotations=ResidueAnnotationsTokenizer(),
+        )
+    else:
+        raise ValueError(f"Unknown model: {model}")
+
+
+def get_esmc_model_tokenizers() -> EsmSequenceTokenizer:
+    return EsmSequenceTokenizer()
+
+
+def get_invalid_tokenizer_ids(tokenizer: EsmTokenizerBase) -> list[int]:
+    if isinstance(tokenizer, EsmSequenceTokenizer):
+        return [
+            tokenizer.mask_token_id,  # type: ignore
+            tokenizer.pad_token_id,  # type: ignore
+            tokenizer.cls_token_id,  # type: ignore
+            tokenizer.eos_token_id,  # type: ignore
+        ]
+    else:
+        return [
+            tokenizer.mask_token_id,
+            tokenizer.pad_token_id,
+            tokenizer.bos_token_id,
+            tokenizer.eos_token_id,
+        ]

Dyna-1/esm/tokenization/function_tokenizer.py

@@ -0,0 +1,429 @@
+"""Tokenizes annotations of protein function."""
+
+import re
+import string
+from functools import cache, cached_property, partial
+from typing import Collection
+
+import numpy as np
+import pandas as pd
+import scipy.sparse as sp
+import torch
+import torch.nn.functional as F
+
+from esm.tokenization.tokenizer_base import EsmTokenizerBase
+from esm.utils.constants import esm3 as C
+from esm.utils.function import interpro, lsh, tfidf
+from esm.utils.misc import stack_variable_length_tensors
+from esm.utils.types import FunctionAnnotation, PathLike
+
+
+def _default_data_path(x: PathLike | None, d: PathLike) -> PathLike:
+    return x if x is not None else C.data_root("esm3") / d
+
+
+def _default_local_data_path(x: PathLike | None, d: PathLike) -> PathLike:
+    return x if x is not None else d
+
+
+class InterProQuantizedTokenizer(EsmTokenizerBase):
+    """Tokenizer for functional annotations.
+
+    This tokenizer converts InterPro and/or function keywords into a multi-token
+    representation by hashing TF-IDF vector representations of the text associated with
+    the fuction and then applying a locality sensitive hash (LSH).
+    """
+
+    def __init__(
+        self,
+        depth: int = 8,
+        lsh_bits_per_token: int = 8,
+        lsh_path: PathLike | None = None,
+        keyword_vocabulary_path: PathLike | None = None,
+        keyword_idf_path: PathLike | None = None,
+        interpro_entry_path: PathLike | None = None,
+        interpro2keywords_path: PathLike | None = None,
+    ):
+        """Constructs function tokenizer.
+
+        Args:
+            depth: number of tokens emitted in each position.
+            lsh_bits_per_token: Number of LSH bits per token. Determines the vocabulary
+                  size.
+            lsh_path: path to locality sensitive hash (LSH) hyperplanes.
+            keyword_vocabulary_path: path to csv containing function keyword vocabulary.
+            keyword_idf_path: path to IDF values for each keyword.
+            interpro_entry_csv_path: path to list of InterPro entries in CSV format.
+            interpro2keywords_path: path to CSV mapping InterPro IDs to function keywords.
+        """
+        self.depth = depth
+
+        self.keyword_vocabulary_path = _default_local_data_path(
+            keyword_vocabulary_path, C.KEYWORDS_VOCABULARY
+        )
+        self.keyword_idf_path = _default_local_data_path(
+            keyword_idf_path, C.KEYWORDS_IDF
+        )
+
+        self._interpro2keywords_path = _default_local_data_path(
+            interpro2keywords_path, C.INTERPRO2KEYWORDS
+        )
+        self.interpro_ = interpro.InterPro(
+            entries_path=_default_local_data_path(interpro_entry_path, C.INTERPRO_ENTRY)
+        )
+
+        self.lsh_path = lsh_path
+        self.lsh_bits_per_token = lsh_bits_per_token
+        self.lsh_vocab_size = 1 << lsh_bits_per_token
+
+        # This is the offset into the vocabulary where LSH tokens start.
+        self._lsh_token_vocab_offset = len(self.special_tokens) + 1  # +1 for <none>
+
+    @cached_property
+    def _lsh(self) -> lsh.LSHTokenized:
+        """Locality sensitive hash for function annotations."""
+        return lsh.LSHTokenized(
+            self.lsh_bits_per_token,
+            len(self.keyword_vocabulary),
+            self.depth,
+            _default_data_path(self.lsh_path, C.LSH_TABLE_PATHS["8bit"]),
+        )
+
+    @cached_property
+    def interpro2keywords(self) -> dict[str, list[str]]:
+        """Mapping from InterPro ID to function keywords."""
+        df = pd.read_csv(self._interpro2keywords_path)
+        assert "interpro_id" in df.columns and "keywords" in df.columns, df.columns
+        return dict(zip(df.interpro_id, df.keywords.str.split(",")))
+
+    @cached_property
+    def interpro_labels(self) -> list[str]:
+        """The set of supported InterPro labels."""
+        return sorted(self.interpro2keywords.keys())
+
+    @cached_property
+    def interpro_to_index(self) -> dict[str, int]:
+        """Mapping from InterPro id to index."""
+        return {id: i for i, id in enumerate(self.interpro_labels)}
+
+    @property
+    def keyword_vocabulary(self) -> list[str]:
+        """Set of supported keywords."""
+        return self._tfidf.vocabulary
+
+    @property
+    def keyword_to_index(self) -> dict[str, int]:
+        """Mapping from keywords to index."""
+        return self._tfidf.vocab_to_index
+
+    @cached_property
+    def _tfidf(self) -> tfidf.TFIDFModel:
+        """Creates TF-IDF model for encoding function keywords."""
+        return tfidf.TFIDFModel(
+            vocabulary_path=self.keyword_vocabulary_path, idf_path=self.keyword_idf_path
+        )
+
+    @cached_property
+    def special_tokens(self) -> list[str]:
+        """List of special tokens which come before cluster tokens in vocab."""
+        return ["<pad>", "<motif>", "<unk>"]
+
+    @cached_property
+    def vocab(self) -> list[str]:
+        """Vocabulary of function tokens."""
+        lsh_tokens = [f"<lsh:{i}>" for i in range(self.lsh_vocab_size)]
+        return self.special_tokens + ["<none>"] + lsh_tokens
+
+    @cached_property
+    def vocab_to_index(self) -> dict[str, int]:
+        return {token: token_id for token_id, token in enumerate(self.vocab)}
+
+    def get_special_tokens_mask(self, encoded: torch.Tensor) -> torch.Tensor:
+        """Determines where in the sequence are special tokens."""
+        where = encoded < len(self.special_tokens)
+        assert torch.all(torch.all(where, dim=1) | torch.all(~where, dim=1))
+        return where[:, 0]
+
+    def tokenize(
+        self,
+        annotations: list[FunctionAnnotation],
+        seqlen: int,
+        p_keyword_dropout: float = 0.0,
+    ) -> list[str]:
+        """Encodes range-annotations of protein function as tokens.
+
+        Args:
+            features: Annotated function ranges, either as InterPro ids or keywords.
+            seqlen: length of sequence.
+            p_keyword_dropout: Optional probability of dropping out keywords from the
+              input annotations.
+        Returns:
+            Tokenized representation of function annotations as a list of string tokens
+            of size seqlen.
+        """
+        assert seqlen >= 0
+
+        if not annotations:
+            return ["<pad>"] * seqlen
+
+        # Expand the range annotations into positional annotaiton sets.
+        positional_labels: list[set[str]] = [set() for _ in range(seqlen)]
+        for annotation in annotations:
+            assert 1 <= annotation.start <= annotation.end <= seqlen, (
+                f"Invalid annotation range [{annotation.start}, {annotation.end}] for "
+                f"sequence length {seqlen}."
+            )
+            for i in range(annotation.start - 1, annotation.end):
+                positional_labels[i].add(annotation.label)
+
+        if p_keyword_dropout > 0:
+            keyword_mask = (
+                np.random.random(len(self._tfidf.vocabulary)) < p_keyword_dropout
+            )
+        else:
+            keyword_mask = None
+
+        # Annotations tend to be repetitive over the length of the sequence - cache their
+        # hashes to speed up tokenization.
+        hash_fn = cache(partial(self._function_text_hash, keyword_mask=keyword_mask))
+
+        tokens: list[str] = []
+        for labels in positional_labels:
+            if not labels:
+                token = "<none>"
+            else:
+                lsh_hash = hash_fn(frozenset(labels))
+                if lsh_hash is not None:
+                    assert len(lsh_hash) == self.depth
+                    token = "<lsh:" + ",".join(map(str, lsh_hash)) + ">"
+                else:
+                    token = "<unk>"
+
+            tokens.append(token)
+
+        return tokens
+
+    def _function_text_hash(
+        self, labels: Collection[str], keyword_mask: np.ndarray | None = None
+    ) -> np.ndarray | None:
+        """Applies a locality sensitive hash (LSH) to function text.
+
+        Args:
+            labels: InterPro ids and/or keywords.
+            keyword_mask: optional boolean array shaped (keyword_vocab_size,) indicating
+                which keywords to drop before hashing.
+        Returns:
+            LSH shaped (depth,) or None if there is no text or keywords to hash.
+        """
+        # Split labels into either InterPro ids or keywords.
+        interpro_ids = []
+        keywords = []
+        for label in labels:
+            match = re.search(r"IPR\d+", label)
+            if match and match.group() in self.interpro_to_index:
+                interpro_ids.append(match.group())
+            elif label in self._tfidf.vocab_to_index:
+                keywords.append(label)
+            else:
+                raise ValueError(f"Unsupported: {label}")
+
+        vec: sp.csr_matrix = self._tfidf.encode(keywords)
+
+        # Perform an element-wise maximum over TF-IDF vectors from distinct tags to
+        # avoid tags getting "washed out" by eg. 4 very similar tags. Keywords are
+        # incorporated as another TF-IDF vector
+        vec: sp.csr_matrix = self._tfidf.encode(keywords)
+        for interpro_id in interpro_ids:
+            interpro_keywords = self.interpro2keywords.get(interpro_id, [])
+            vec_ = self._tfidf.encode(interpro_keywords)
+            vec = vec.maximum(vec_)
+
+        if keyword_mask is not None:
+            vec.data *= 1 - np.take(keyword_mask, vec.indices)
+
+        if vec.sum() == 0:
+            return None
+
+        return self._lsh(vec)[0, :]
+
+    def encode(
+        self, tokens: list[str], add_special_tokens: bool = True
+    ) -> torch.Tensor:
+        """Encodes string tokens as token-id tensor.
+
+        Args:
+            tokens: list of individual tokens. e.g. ["<none>", "<pq:1,2,3,4>"]
+            add_special_tokens: whether to add a single pad token at the start and end
+              of the sequence to act as <cls> and <eos> tokens.
+        Returns:
+            <int>[length, depth] function tokens. Length will be +2 of input tokens
+            length when add_special_tokens is True.
+        """
+        token_ids = torch.zeros(size=(len(tokens), self.depth), dtype=torch.int64)
+        for i, token in enumerate(tokens):
+            token_ids[i, :] = torch.tensor(self._token2ids(token))
+        if add_special_tokens:
+            token_ids = F.pad(
+                token_ids, (0, 0, 1, 1), value=self.vocab_to_index["<pad>"]
+            )
+        return token_ids
+
+    def lookup_annotation_name(self, annotation: FunctionAnnotation) -> str | None:
+        return self.interpro_.lookup_name(annotation.label)
+
+    def format_annotation(self, annotation: FunctionAnnotation) -> str:
+        annotation_name = self.lookup_annotation_name(annotation)
+        if annotation_name is not None:
+            return f"{annotation_name} ({annotation.label})"
+        else:
+            return annotation.label
+
+    def _token2ids(self, token: str) -> list[int]:
+        """Converts token into token_id set of length depth."""
+        if re.match(r"<lsh:[\d+,]+>", token):
+            lsh_ids = [int(lsh_id) for lsh_id in re.findall(r"\d+", token)]
+            assert (
+                len(lsh_ids) == self.depth
+            ), f"Expected token to have {self.depth} ids found {lsh_ids}"
+            return [self._lsh_token_vocab_offset + lsh_id for lsh_id in lsh_ids]
+        elif token == "<none>" or token in self.special_tokens:
+            return [self.vocab_to_index[token]] * self.depth
+        else:
+            raise ValueError(f"Unknown token: {token}")
+
+    def batch_encode(
+        self, token_batch: list[list[str]], add_special_tokens: bool = True
+    ) -> torch.Tensor:
+        """Encodes batch of function tokens.
+
+        Args:
+            token_batch: batch of function tokens.
+            add_special_tokens: whether to add special tokens.
+        Returns:
+            <int>[batch_size, max_length, depth] batch of encoded tokens.
+        """
+        encoded = [
+            self.encode(tokens, add_special_tokens=add_special_tokens)
+            for tokens in token_batch
+        ]
+        return stack_variable_length_tensors(
+            encoded, constant_value=self.vocab_to_index["<pad>"]
+        )
+
+    def decode(self, encoded: torch.Tensor):
+        raise NotImplementedError(
+            "Function token decoding should be handled with "
+            "util.decoding.decode_function_annotations"
+        )
+
+    @property
+    def mask_token(self) -> str:
+        return "<pad>"
+
+    @property
+    def mask_token_id(self) -> int:
+        return self.vocab_to_index[self.mask_token]
+
+    @property
+    def bos_token(self) -> str:
+        return "<pad>"
+
+    @property
+    def bos_token_id(self) -> int:
+        return self.vocab_to_index[self.bos_token]
+
+    @property
+    def eos_token(self) -> str:
+        return "<pad>"
+
+    @property
+    def eos_token_id(self) -> int:
+        return self.vocab_to_index[self.eos_token]
+
+    @property
+    def pad_token(self) -> str:
+        return "<pad>"
+
+    @property
+    def pad_token_id(self) -> int:
+        return self.vocab_to_index[self.pad_token]
+
+    @property
+    def chain_break_token(self) -> str:
+        return "<pad>"
+
+    @property
+    def chain_break_token_id(self) -> int:
+        return self.vocab_to_index[self.chain_break_token]
+
+    @property
+    def all_token_ids(self):
+        return list(range(len(self.vocab)))
+
+    @property
+    def special_token_ids(self):
+        return [self.vocab_to_index[token] for token in self.special_tokens]
+
+
+def _texts_to_keywords(texts: list[str]) -> list[str]:
+    """Breaks InterPro/GO free-text description set into bag-of-n-grams for n={1,2}.
+
+    Args:
+        texts: collection of text descriptions, i.e. InterPro/GO names.
+    Returns:
+        Collection of terms/n-grams
+    """
+    keywords = []
+    for text in texts:
+        keywords.extend(_keywords_from_text(text))
+    return keywords
+
+
+def _keywords_from_text(text: str) -> list[str]:
+    """Splits text into unigrams and bigrams."""
+    elements = text.split(", ")
+
+    terms = []
+    for element in elements:
+        element = _sanitize(element)
+        words = element.split()
+
+        # Add 1-mers
+        terms.extend(words)
+
+        # Add 2-mers
+        for i in range(len(words) - 1):
+            bigram = words[i] + " " + words[i + 1]
+            terms.append(bigram)
+
+    return [term for term in terms if len(term) > 1 and term not in _EXCLUDED_TERMS]
+
+
+def _sanitize(text: str) -> str:
+    text = text.replace("-", " ")
+    text = text.translate(str.maketrans("", "", string.punctuation))
+    text = text.lower()
+    return text
+
+
+# These terms are omitted from textual representations since they are pervasive and
+# unspecific to particular protein function.
+_EXCLUDED_TERMS = {
+    "binding domain",
+    "biological_process",
+    "biological process",
+    "biologicalprocess",
+    "c",
+    "cellular_component",
+    "cellular component",
+    "cellularcomponent",
+    "cellular_process",
+    "cellularprocess",
+    "cellular process",
+    "cellularprocess",
+    "like domain",
+    "molecular function",
+    "molecular_function",
+    "molecularfunction",
+    "n",
+}

Dyna-1/esm/tokenization/residue_tokenizer.py

@@ -0,0 +1,236 @@
+from functools import cached_property
+from typing import Any
+
+import pandas as pd
+import torch
+import torch.nn.functional as F
+from cloudpathlib import AnyPath
+
+from esm.tokenization.tokenizer_base import EsmTokenizerBase
+from esm.utils.constants import esm3 as C
+
+Sample = dict[str, Any]
+
+
+class ResidueAnnotationsTokenizer(EsmTokenizerBase):
+    def __init__(self, csv_path: str | None = None, max_annotations: int = 16):
+        if csv_path is None:
+            csv_path = str(C.data_root("esm3") / C.RESID_CSV)
+        self.csv_path = csv_path
+        self.max_annotations = max_annotations
+
+    @cached_property
+    def _description2label(self) -> dict[str, str]:
+        with AnyPath(self.csv_path).open() as f:  # type: ignore
+            df = pd.read_csv(f)
+        return dict(zip(df.label, df.label_clean))
+
+    @cached_property
+    def _labels(self) -> list[str]:
+        with AnyPath(self.csv_path).open() as f:  # type: ignore
+            df = pd.read_csv(f)
+        labels = (
+            df.groupby("label_clean")["count"]
+            .sum()
+            .sort_values(ascending=False, kind="stable")  # type: ignore
+            .index.tolist()
+        )
+        assert isinstance(labels, list)
+        return labels  # type: ignore
+
+    def _description2id(self, description: str) -> int | None:
+        label = self._description2label.get(description)
+        return self._label2id.get(label)  # type: ignore
+
+    @cached_property
+    def _label2id(self) -> dict[str, int]:
+        offset = len(self.special_tokens) + 1  # +1 for "<none>"
+        return {label: offset + i for i, label in enumerate(self._labels)}
+
+    @cached_property
+    def special_tokens(self) -> list[str]:
+        """List of special tokens which come before cluster toknes in vocab."""
+        return ["<pad>", "<motif>", "<unk>"]
+
+    @cached_property
+    def vocab(self):
+        annotation_tokens = [f"<ra:{id}>" for _, id in self._label2id.items()]
+        return self.special_tokens + ["<none>"] + annotation_tokens
+
+    @cached_property
+    def vocab_to_index(self) -> dict[str, int]:
+        return {token: token_id for token_id, token in enumerate(self.vocab)}
+
+    @cached_property
+    def vocabulary(self) -> list[str]:
+        """Full vocabulary."""
+        return [*self.special_tokens, "<none>", *self._labels]
+
+    def get_special_tokens_mask(self, encoded: torch.Tensor) -> torch.Tensor:
+        """Determines where in the sequence are special tokens."""
+        return encoded[:, 0] < len(self.special_tokens)
+
+    def tokenize(
+        self, sample: Sample | None, sequence: str, fail_on_mismatch: bool = False
+    ) -> list[str]:
+        """
+        # interpro_site_starts
+        # interpro_site_ends  # should always == interpro_site_starts.  but I haven't checked overall.
+        # interpro_site_residues  # the residue identity of the specfic residue that is annotated.  good for a sanity check that parsing occurred correctly.
+        # interpro_site_descriptions
+        # ASSERT (i.e. drop if bad)
+        # interpro_site_residues matches the residue at that position
+        # all these lists ^ above are the same length
+        """
+        seqlen = len(sequence)
+        assert seqlen >= 0
+        # None mean sequence is *not annotated* - so use full <pad>
+        if sample is None:
+            return ["<pad>"] * seqlen
+
+        if any(
+            sample.get(field) is None
+            for field in [
+                "interpro_site_descriptions",
+                "interpro_site_starts",
+                "interpro_site_ends",
+                "interpro_site_residues",
+            ]
+        ):
+            return ["<pad>"] * seqlen
+
+        num_annotations = len(sample["interpro_site_descriptions"])
+        if any(
+            len(sample[field]) != num_annotations
+            for field in [
+                "interpro_site_starts",
+                "interpro_site_ends",
+                "interpro_site_residues",
+            ]
+        ):
+            # mismatched length.
+            return ["<pad>"] * seqlen
+
+        positional_ids = [set() for _ in range(seqlen)]
+        for description, start, end, residues in zip(
+            sample["interpro_site_descriptions"],
+            sample["interpro_site_starts"],
+            sample["interpro_site_ends"],
+            sample["interpro_site_residues"],
+        ):
+            try:
+                start = int(start)
+                end = int(end)
+            except (TypeError, ValueError):
+                continue
+
+            # Start / End are 1-indexed [inclusive, inclusive].
+            if start <= 0 or end > seqlen or start > end:
+                print(f"invalid start/end: ({start}, {end}), len: {seqlen}")
+                continue
+
+            if len(residues) != (end - start) + 1:
+                print(f"bad reference residue: {residues}")
+                continue
+
+            token_id = self._description2id(description)
+            if token_id is None:
+                token_id = self.vocab_to_index["<unk>"]
+
+            for i, residue in zip(range(start - 1, end), residues):
+                # If there are any mismatching residues, skip the entire sample.
+                if sequence[i] != residue:
+                    if fail_on_mismatch:
+                        raise ValueError(
+                            f"Residue mismatch at position {i} (1-indexed): {sequence[i]} != {residue}"
+                        )
+                    return ["<pad>"] * seqlen
+
+                positional_ids[i].add(token_id)
+
+        tokens = []
+        for token_ids in positional_ids:
+            if token_ids:
+                token = "<ra:" + ",".join(str(token_id) for token_id in token_ids) + ">"
+            else:
+                token = "<none>"
+            tokens.append(token)
+        return tokens
+
+    def _token2ids(self, token: str) -> list[int]:
+        if token.startswith("<ra:") and token.endswith(">"):
+            return [int(token_id) for token_id in token[4:-1].split(",")]
+        else:
+            token_id = self.vocab_to_index[token]
+            return [token_id]
+
+    def encode(
+        self, tokens: list[str], add_special_tokens: bool = True
+    ) -> torch.Tensor:
+        token_ids = torch.full(
+            size=(len(tokens), self.max_annotations),
+            dtype=torch.int64,
+            fill_value=self.vocab_to_index["<pad>"],
+        )
+        for i, token in enumerate(tokens):
+            ids = self._token2ids(token)[: self.max_annotations]
+            token_ids[i, : len(ids)] = torch.tensor(ids)
+
+        if add_special_tokens:
+            token_ids = F.pad(
+                token_ids, (0, 0, 1, 1), value=self.vocab_to_index["<pad>"]
+            )
+        return token_ids
+
+    def decode(self, encoded: torch.Tensor) -> list[str]:
+        raise NotImplementedError(
+            "Residue annotation decoding should be handled with util.decoding.decode_residue_annotations"
+        )
+
+    @property
+    def mask_token(self) -> str:
+        return "<pad>"
+
+    @property
+    def mask_token_id(self) -> int:
+        return self.vocab_to_index[self.mask_token]
+
+    @property
+    def bos_token(self) -> str:
+        return "<pad>"
+
+    @property
+    def bos_token_id(self) -> int:
+        return self.vocab_to_index[self.bos_token]
+
+    @property
+    def eos_token(self) -> str:
+        return "<pad>"
+
+    @property
+    def eos_token_id(self) -> int:
+        return self.vocab_to_index[self.eos_token]
+
+    @property
+    def pad_token(self) -> str:
+        return "<pad>"
+
+    @property
+    def pad_token_id(self) -> int:
+        return self.vocab_to_index[self.pad_token]
+
+    @property
+    def chain_break_token(self) -> str:
+        return "<pad>"
+
+    @property
+    def chain_break_token_id(self) -> int:
+        return self.vocab_to_index[self.chain_break_token]
+
+    @property
+    def all_token_ids(self):
+        return list(range(len(self.vocab)))
+
+    @property
+    def special_token_ids(self):
+        return [self.vocab_to_index[token] for token in self.special_tokens]

Dyna-1/esm/tokenization/sasa_tokenizer.py

@@ -0,0 +1,153 @@
+from functools import cached_property
+
+import torch
+
+from esm.tokenization.tokenizer_base import EsmTokenizerBase
+from esm.utils.constants import esm3 as C
+
+
+class SASADiscretizingTokenizer(EsmTokenizerBase):
+    """Tokenizer for Solvent Accessible Surface Area (SASA)."""
+
+    def __init__(self, boundaries: list[float] = C.SASA_DISCRETIZATION_BOUNDARIES):
+        self._boundaries = sorted(boundaries)
+
+    @cached_property
+    def special_tokens(self) -> list[str]:
+        return ["<pad>", "<motif>", "<unk>"]
+
+    @cached_property
+    def vocab(self) -> list[str]:
+        """Discrete token vocabulary.
+
+        Returns:
+            token vocabulary with ranges represented as "<low-high>".
+        """
+        boundary_strs = ["0"] + [str(b) for b in self._boundaries] + ["inf"]
+        range_tokens = [
+            f"<{low}-{high}>"
+            for low, high in zip(boundary_strs[:-1], boundary_strs[1:])
+        ]
+        return self.special_tokens + range_tokens
+
+    @cached_property
+    def midpoints_tensor(self) -> torch.Tensor:
+        """Midpoints of the SASA token ranges."""
+        boundaries = [0] + self._boundaries + [self._boundaries[-1] * 2]
+        midpoint_tokens = [
+            (float(high) + float(low)) / 2
+            for low, high in zip(boundaries[:-1], boundaries[1:])
+        ]
+        midpoint_tokens = [float("nan"), float("nan"), float("nan")] + midpoint_tokens
+        return torch.Tensor(midpoint_tokens)
+
+    def midpoints(self) -> list[float]:
+        """Midpoints of the SASA token ranges."""
+        return self.midpoints_tensor.tolist()
+
+    @cached_property
+    def vocab_to_index(self) -> dict[str, int]:
+        """Constructs token -> token id mapping."""
+        return {word: i for i, word in enumerate(self.vocab)}
+
+    def get_special_tokens_mask(self, tokens: torch.Tensor) -> torch.Tensor:
+        """Determines which positions are special tokens.
+
+        Args:
+            tokens: <int>[length]
+        Returns:
+            <bool>[length] tensor, true where special tokens are located in the input.
+        """
+        return tokens < len(self.special_tokens)
+
+    def encode(
+        self, values: list[float | str], add_special_tokens: bool = True
+    ) -> torch.Tensor:
+        """Encodes SASA values as discrete tokens.
+
+        Args:
+            values: list of either SASA values or individual tokens. For example
+                [1.2, "<pad>", 10.3, <pad>, 0.]
+        Returns:
+            Token ids as tensor. Adds BOS and EOS special tokens.
+        """
+        ids = []
+        if add_special_tokens:
+            ids.append(self.vocab_to_index["<pad>"])  # BOS
+        for value in values:
+            if isinstance(value, (float, int)):
+                bucket = torch.bucketize(value, torch.tensor(self._boundaries))
+                token_id = len(self.special_tokens) + bucket
+            elif isinstance(value, str):
+                token_id = self.vocab_to_index[value]
+            else:
+                raise TypeError(value)
+            ids.append(token_id)
+        if add_special_tokens:
+            ids.append(self.vocab_to_index["<pad>"])  # EOS
+
+        return torch.tensor(ids, dtype=torch.int64)
+
+    def decode_float(self, encoded: torch.Tensor) -> list[float]:
+        """Decodes SASA token ids into float values."""
+        decoded = self.midpoints_tensor[encoded.cpu()]
+        nan_mask = torch.isnan(decoded)
+        np_arr = decoded.numpy()
+        np_arr[nan_mask.numpy()] = None
+        return np_arr.tolist()
+
+    def decode(self, encoded: torch.Tensor) -> str:
+        """Decodes SASA token ids."""
+        return ",".join(self.vocab[i] for i in encoded)
+
+    def decode_list(self, encoded: torch.Tensor) -> list[str]:
+        """Decodes SASA token ids."""
+        return [self.vocab[i] for i in encoded]
+
+    @property
+    def mask_token(self) -> str:
+        return "<pad>"
+
+    @property
+    def mask_token_id(self) -> int:
+        return self.vocab_to_index[self.mask_token]
+
+    @property
+    def bos_token(self) -> str:
+        return "<pad>"
+
+    @property
+    def bos_token_id(self) -> int:
+        return self.vocab_to_index[self.bos_token]
+
+    @property
+    def eos_token(self) -> str:
+        return "<pad>"
+
+    @property
+    def eos_token_id(self) -> int:
+        return self.vocab_to_index[self.eos_token]
+
+    @property
+    def pad_token(self) -> str:
+        return "<pad>"
+
+    @property
+    def pad_token_id(self) -> int:
+        return self.vocab_to_index[self.pad_token]
+
+    @property
+    def chain_break_token(self) -> str:
+        return "<pad>"
+
+    @property
+    def chain_break_token_id(self) -> int:
+        return self.vocab_to_index[self.chain_break_token]
+
+    @property
+    def all_token_ids(self):
+        return list(range(len(self.vocab)))
+
+    @property
+    def special_token_ids(self):
+        return [self.vocab_to_index[token] for token in self.special_tokens]

Dyna-1/esm/tokenization/sequence_tokenizer.py

@@ -0,0 +1,89 @@
+from tokenizers import Tokenizer
+from tokenizers.models import BPE
+from tokenizers.processors import TemplateProcessing
+from transformers import PreTrainedTokenizerFast
+
+from esm.tokenization.tokenizer_base import EsmTokenizerBase
+from esm.utils.constants import esm3 as C
+
+
+class EsmSequenceTokenizer(PreTrainedTokenizerFast, EsmTokenizerBase):
+    """
+    Constructs an ESM tokenizer.
+    """
+
+    model_input_names = ["sequence_tokens", "attention_mask"]
+
+    def __init__(
+        self,
+        unk_token="<unk>",
+        cls_token="<cls>",
+        pad_token="<pad>",
+        mask_token="<mask>",
+        eos_token="<eos>",
+        chain_break_token="|",
+        **kwargs,
+    ):
+        all_tokens = C.SEQUENCE_VOCAB
+        token_to_id = {tok: ind for ind, tok in enumerate(all_tokens)}
+
+        # a character-level tokenizer is the same as BPE with no token merges
+        bpe = BPE(token_to_id, merges=[], unk_token=unk_token)
+        tokenizer = Tokenizer(bpe)
+        special_tokens = [
+            cls_token,
+            pad_token,
+            mask_token,
+            eos_token,
+            chain_break_token,
+        ]
+        self.cb_token = chain_break_token
+        additional_special_tokens = [chain_break_token]
+
+        tokenizer.add_special_tokens(special_tokens)
+
+        # This is where we configure the automatic addition of special tokens when we call
+        # tokenizer(text, add_special_tokens=True). Note that you can also configure how two
+        # sequences are merged if you want.
+        tokenizer.post_processor = TemplateProcessing(  # type: ignore
+            single="<cls> $A <eos>",
+            special_tokens=[
+                ("<cls>", tokenizer.token_to_id("<cls>")),
+                ("<eos>", tokenizer.token_to_id("<eos>")),
+            ],
+        )
+        super().__init__(
+            tokenizer_object=tokenizer,
+            unk_token=unk_token,
+            cls_token=cls_token,
+            pad_token=pad_token,
+            mask_token=mask_token,
+            eos_token=eos_token,
+            additional_special_tokens=additional_special_tokens,
+            **kwargs,
+        )
+
+    # These are a footgun, we never use the `bos` token anywhere so we're just overriding it here.
+    @property
+    def bos_token(self):
+        return self.cls_token
+
+    @property
+    def bos_token_id(self):
+        return self.cls_token_id
+
+    @property
+    def chain_break_token(self):
+        return self.cb_token
+
+    @property
+    def chain_break_token_id(self):
+        return self.convert_tokens_to_ids(self.chain_break_token)
+
+    @property
+    def all_token_ids(self):
+        return list(range(self.vocab_size))
+
+    @property
+    def special_token_ids(self):
+        return self.all_special_ids

Dyna-1/esm/tokenization/ss_tokenizer.py

@@ -0,0 +1,125 @@
+from functools import cached_property
+from typing import Sequence
+
+import torch
+
+from esm.tokenization.tokenizer_base import EsmTokenizerBase
+from esm.utils.constants import esm3 as C
+
+
+class SecondaryStructureTokenizer(EsmTokenizerBase):
+    """Tokenizer for secondary structure strings."""
+
+    def __init__(self, kind: str = "ss8"):
+        assert kind in ("ss8", "ss3")
+        self.kind = kind
+
+    @property
+    def special_tokens(self) -> list[str]:
+        return ["<pad>", "<motif>", "<unk>"]
+
+    @cached_property
+    def vocab(self):
+        """Tokenzier vocabulary list."""
+        match self.kind:
+            case "ss8":
+                nonspecial_tokens = list(C.SSE_8CLASS_VOCAB)  # "GHITEBSC"
+            case "ss3":
+                nonspecial_tokens = list(C.SSE_3CLASS_VOCAB)  # HEC
+            case _:
+                raise ValueError(self.kind)
+
+        # The non-special tokens ids match amino acid tokens ids when possible.
+        return [*self.special_tokens, *nonspecial_tokens]
+
+    @cached_property
+    def vocab_to_index(self) -> dict[str, int]:
+        """Constructs token -> token id mapping."""
+        return {word: i for i, word in enumerate(self.vocab)}
+
+    def get_special_tokens_mask(self, tokens: torch.Tensor) -> torch.Tensor:
+        """Determines which positions are special tokens.
+
+        Args:
+            tokens: <int>[length]
+        Returns:
+            <bool>[length] tensor, true where special tokens are located in the input.
+        """
+        return tokens < len(self.special_tokens)
+
+    def encode(
+        self, sequence: str | Sequence[str], add_special_tokens: bool = True
+    ) -> torch.Tensor:
+        """Encode secondary structure string
+
+        Args:
+            string: secondary structure string e.g. "GHHIT", or as token listk.
+        Returns:
+            <int>[sequence_length] token ids representing. Will add <cls>/<eos>.
+        """
+        ids = []
+        if add_special_tokens:
+            ids.append(self.vocab_to_index["<pad>"])  # cls
+        for char in sequence:
+            ids.append(self.vocab_to_index[char])
+        if add_special_tokens:
+            ids.append(self.vocab_to_index["<pad>"])  # eos
+        return torch.tensor(ids, dtype=torch.int64)
+
+    def decode(self, encoded: torch.Tensor) -> str:
+        """Decodes token ids into secondary structure string.
+
+        Args:
+            encoded: <int>[length] token id array.
+        Returns
+            Decoded secondary structure string.
+        """
+        return "".join(self.vocab[i] for i in encoded)
+
+    @property
+    def mask_token(self) -> str:
+        return "<pad>"
+
+    @property
+    def mask_token_id(self) -> int:
+        return self.vocab_to_index[self.mask_token]
+
+    @property
+    def bos_token(self) -> str:
+        return "<pad>"
+
+    @property
+    def bos_token_id(self) -> int:
+        return self.vocab_to_index[self.bos_token]
+
+    @property
+    def eos_token(self) -> str:
+        return "<pad>"
+
+    @property
+    def eos_token_id(self) -> int:
+        return self.vocab_to_index[self.eos_token]
+
+    @property
+    def pad_token(self) -> str:
+        return "<pad>"
+
+    @property
+    def pad_token_id(self) -> int:
+        return self.vocab_to_index[self.pad_token]
+
+    @property
+    def chain_break_token(self) -> str:
+        return "<pad>"
+
+    @property
+    def chain_break_token_id(self) -> int:
+        return self.vocab_to_index[self.chain_break_token]
+
+    @property
+    def all_token_ids(self):
+        return list(range(len(self.vocab)))
+
+    @property
+    def special_token_ids(self):
+        return [self.vocab_to_index[token] for token in self.special_tokens]

Dyna-1/esm/tokenization/structure_tokenizer.py

@@ -0,0 +1,83 @@
+from esm.tokenization.tokenizer_base import EsmTokenizerBase
+from esm.utils.constants import esm3 as C
+
+
+class StructureTokenizer(EsmTokenizerBase):
+    """A convenince class for accessing special token ids of
+    the StructureTokenEncoder and StructureTokenDecoder."""
+
+    def __init__(self, codebook_size: int = C.VQVAE_CODEBOOK_SIZE):
+        self.vq_vae_special_tokens = {
+            "MASK": codebook_size,
+            "EOS": codebook_size + 1,
+            "BOS": codebook_size + 2,
+            "PAD": codebook_size + 3,
+            "CHAINBREAK": codebook_size + 4,
+        }
+
+    def mask_token(self) -> str:
+        raise NotImplementedError(
+            "Structure tokens are defined on 3D coordinates, not strings."
+        )
+
+    @property
+    def mask_token_id(self) -> int:
+        return self.vq_vae_special_tokens["MASK"]
+
+    def bos_token(self) -> str:
+        raise NotImplementedError(
+            "Structure tokens are defined on 3D coordinates, not strings."
+        )
+
+    @property
+    def bos_token_id(self) -> int:
+        return self.vq_vae_special_tokens["BOS"]
+
+    def eos_token(self) -> str:
+        raise NotImplementedError(
+            "Structure tokens are defined on 3D coordinates, not strings."
+        )
+
+    @property
+    def eos_token_id(self) -> int:
+        return self.vq_vae_special_tokens["EOS"]
+
+    def pad_token(self) -> str:
+        raise NotImplementedError(
+            "Structure tokens are defined on 3D coordinates, not strings."
+        )
+
+    @property
+    def pad_token_id(self) -> int:
+        return self.vq_vae_special_tokens["PAD"]
+
+    def chain_break_token(self) -> str:
+        raise NotImplementedError(
+            "Structure tokens are defined on 3D coordinates, not strings."
+        )
+
+    @property
+    def chain_break_token_id(self) -> int:
+        return self.vq_vae_special_tokens["CHAINBREAK"]
+
+    @property
+    def all_token_ids(self):
+        return list(range(C.VQVAE_CODEBOOK_SIZE + len(self.vq_vae_special_tokens)))
+
+    @property
+    def special_token_ids(self):
+        return self.vq_vae_special_tokens.values()
+
+    def encode(self, *args, **kwargs):
+        raise NotImplementedError(
+            "The StructureTokenizer class is provided as a convenience for "
+            "accessing special token ids of the StructureTokenEncoder and StructureTokenDecoder.\n"
+            "Please use them instead."
+        )
+
+    def decode(self, *args, **kwargs):
+        raise NotImplementedError(
+            "The StructureTokenizer class is provided as a convenience for "
+            "accessing special token ids of the StructureTokenEncoder and StructureTokenDecoder.\n"
+            "Please use them instead."
+        )

Dyna-1/esm/tokenization/tokenizer_base.py

@@ -0,0 +1,44 @@
+from typing import Protocol, runtime_checkable
+
+
+@runtime_checkable
+class EsmTokenizerBase(Protocol):
+    def encode(self, *args, **kwargs): ...
+
+    def decode(self, *args, **kwargs): ...
+
+    @property
+    def mask_token(self) -> str: ...
+
+    @property
+    def mask_token_id(self) -> int: ...
+
+    @property
+    def bos_token(self) -> str: ...
+
+    @property
+    def bos_token_id(self) -> int: ...
+
+    @property
+    def eos_token(self) -> str: ...
+
+    @property
+    def eos_token_id(self) -> int: ...
+
+    @property
+    def pad_token(self) -> str: ...
+
+    @property
+    def pad_token_id(self) -> int: ...
+
+    @property
+    def chain_break_token(self) -> str: ...
+
+    @property
+    def chain_break_token_id(self) -> int: ...
+
+    @property
+    def all_token_ids(self): ...
+
+    @property
+    def special_token_ids(self): ...

Dyna-1/esm/utils/constants/api.py

@@ -0,0 +1,5 @@
+MAX_TOPK_SEQUENCE = 32
+MAX_TOPK_STRUCTURE = MAX_TOPK_SEQUENCE
+MAX_TOPK_SECONDARY_STRUCTURE = MAX_TOPK_SEQUENCE
+MAX_TOPK_SASA = MAX_TOPK_SEQUENCE
+MAX_TOPK_FUNCTION = MAX_TOPK_SEQUENCE

Dyna-1/esm/utils/constants/esm3.py

@@ -0,0 +1,130 @@
+import os
+from functools import cache
+from pathlib import Path
+
+from huggingface_hub import snapshot_download
+
+SEQUENCE_BOS_TOKEN = 0
+SEQUENCE_PAD_TOKEN = 1
+SEQUENCE_EOS_TOKEN = 2
+SEQUENCE_CHAINBREAK_TOKEN = 31
+SEQUENCE_MASK_TOKEN = 32
+
+VQVAE_CODEBOOK_SIZE = 4096
+VQVAE_SPECIAL_TOKENS = {
+    "MASK": VQVAE_CODEBOOK_SIZE,
+    "EOS": VQVAE_CODEBOOK_SIZE + 1,
+    "BOS": VQVAE_CODEBOOK_SIZE + 2,
+    "PAD": VQVAE_CODEBOOK_SIZE + 3,
+    "CHAINBREAK": VQVAE_CODEBOOK_SIZE + 4,
+}
+VQVAE_DIRECTION_LOSS_BINS = 16
+VQVAE_PAE_BINS = 64
+VQVAE_MAX_PAE_BIN = 31.0
+VQVAE_PLDDT_BINS = 50
+
+STRUCTURE_MASK_TOKEN = VQVAE_SPECIAL_TOKENS["MASK"]
+STRUCTURE_BOS_TOKEN = VQVAE_SPECIAL_TOKENS["BOS"]
+STRUCTURE_EOS_TOKEN = VQVAE_SPECIAL_TOKENS["EOS"]
+STRUCTURE_PAD_TOKEN = VQVAE_SPECIAL_TOKENS["PAD"]
+STRUCTURE_CHAINBREAK_TOKEN = VQVAE_SPECIAL_TOKENS["CHAINBREAK"]
+STRUCTURE_UNDEFINED_TOKEN = 955
+
+SASA_PAD_TOKEN = 0
+
+SS8_PAD_TOKEN = 0
+
+INTERPRO_PAD_TOKEN = 0
+
+RESIDUE_PAD_TOKEN = 0
+
+CHAIN_BREAK_STR = "|"
+
+SEQUENCE_BOS_STR = "<cls>"
+SEQUENCE_EOS_STR = "<eos>"
+
+MASK_STR_SHORT = "_"
+SEQUENCE_MASK_STR = "<mask>"
+SASA_MASK_STR = "<unk>"
+SS8_MASK_STR = "<unk>"
+
+# fmt: off
+SEQUENCE_VOCAB = [
+    "<cls>", "<pad>", "<eos>", "<unk>",
+    "L", "A", "G", "V", "S", "E", "R", "T", "I", "D", "P", "K",
+    "Q", "N", "F", "Y", "M", "H", "W", "C", "X", "B", "U", "Z",
+    "O", ".", "-", "|",
+    "<mask>",
+]
+# fmt: on
+
+SSE_8CLASS_VOCAB = "GHITEBSC"
+SSE_3CLASS_VOCAB = "HEC"
+SSE_8CLASS_TO_3CLASS_MAP = {
+    "G": "H",
+    "H": "H",
+    "I": "H",
+    "T": "C",
+    "E": "E",
+    "B": "E",
+    "S": "C",
+    "C": "C",
+}
+
+SASA_DISCRETIZATION_BOUNDARIES = [
+    0.8,
+    4.0,
+    9.6,
+    16.4,
+    24.5,
+    32.9,
+    42.0,
+    51.5,
+    61.2,
+    70.9,
+    81.6,
+    93.3,
+    107.2,
+    125.4,
+    151.4,
+]
+
+MAX_RESIDUE_ANNOTATIONS = 16
+
+
+TFIDF_VECTOR_SIZE = 58641
+
+
+@staticmethod
+@cache
+def data_root(model: str):
+    if "INFRA_PROVIDER" in os.environ:
+        return Path("")
+    # Try to download from hugginface if it doesn't exist
+    if model.startswith("esm3"):
+        path = Path(snapshot_download(repo_id="EvolutionaryScale/esm3-sm-open-v1"))
+    elif model.startswith("esmc-300"):
+        path = Path(snapshot_download(repo_id="EvolutionaryScale/esmc-300m-2024-12"))
+    elif model.startswith("esmc-600"):
+        path = Path(snapshot_download(repo_id="EvolutionaryScale/esmc-600m-2024-12"))
+    else:
+        raise ValueError(f"{model=} is an invalid model name.")
+    return path
+
+
+IN_REPO_DATA_FOLDER = Path(__file__).parents[2] / "data"
+
+INTERPRO_ENTRY = IN_REPO_DATA_FOLDER / "entry_list_safety_29026.list"
+INTERPRO_HIERARCHY = IN_REPO_DATA_FOLDER / "ParentChildTreeFile.txt"
+INTERPRO2GO = IN_REPO_DATA_FOLDER / "ParentChildTreeFile.txt"
+INTERPRO_2ID = "data/tag_dict_4_safety_filtered.json"
+
+LSH_TABLE_PATHS = {"8bit": "data/hyperplanes_8bit_58641.npz"}
+
+KEYWORDS_VOCABULARY = (
+    IN_REPO_DATA_FOLDER / "keyword_vocabulary_safety_filtered_58641.txt"
+)
+KEYWORDS_IDF = IN_REPO_DATA_FOLDER / "keyword_idf_safety_filtered_58641.npy"
+
+RESID_CSV = "data/uniref90_and_mgnify90_residue_annotations_gt_1k_proteins.csv"
+INTERPRO2KEYWORDS = IN_REPO_DATA_FOLDER / "interpro_29026_to_keywords_58641.csv"

Dyna-1/esm/utils/constants/models.py

@@ -0,0 +1,25 @@
+# Model names
+ESM3_OPEN_SMALL = "esm3_sm_open_v1"
+ESM3_OPEN_SMALL_ALIAS_1 = "esm3-open-2024-03"
+ESM3_OPEN_SMALL_ALIAS_2 = "esm3-sm-open-v1"
+ESM3_OPEN_SMALL_ALIAS_3 = "esm3-open"
+ESM3_STRUCTURE_ENCODER_V0 = "esm3_structure_encoder_v0"
+ESM3_STRUCTURE_DECODER_V0 = "esm3_structure_decoder_v0"
+ESM3_FUNCTION_DECODER_V0 = "esm3_function_decoder_v0"
+ESMC_600M = "esmc_600m"
+ESMC_300M = "esmc_300m"
+
+
+def model_is_locally_supported(x: str):
+    return x in {
+        ESM3_OPEN_SMALL,
+        ESM3_OPEN_SMALL_ALIAS_1,
+        ESM3_OPEN_SMALL_ALIAS_2,
+        ESM3_OPEN_SMALL_ALIAS_3,
+    }
+
+
+def normalize_model_name(x: str):
+    if x in {ESM3_OPEN_SMALL_ALIAS_1, ESM3_OPEN_SMALL_ALIAS_2, ESM3_OPEN_SMALL_ALIAS_3}:
+        return ESM3_OPEN_SMALL
+    return x

Dyna-1/esm/utils/constants/physics.py

@@ -0,0 +1,5 @@
+BB_COORDINATES = [
+    [0.5256, 1.3612, 0.0000],
+    [0.0000, 0.0000, 0.0000],
+    [-1.5251, 0.0000, 0.0000],
+]

Dyna-1/esm/utils/decoding.py

@@ -0,0 +1,244 @@
+import warnings
+from typing import cast
+
+import attr
+import torch
+
+from esm.models.function_decoder import FunctionTokenDecoder
+from esm.models.vqvae import StructureTokenDecoder
+from esm.sdk.api import ESMProtein, ESMProteinTensor
+from esm.tokenization import TokenizerCollectionProtocol
+from esm.tokenization.function_tokenizer import (
+    InterProQuantizedTokenizer,
+)
+from esm.tokenization.residue_tokenizer import (
+    ResidueAnnotationsTokenizer,
+)
+from esm.tokenization.sasa_tokenizer import (
+    SASADiscretizingTokenizer,
+)
+from esm.tokenization.sequence_tokenizer import (
+    EsmSequenceTokenizer,
+)
+from esm.tokenization.ss_tokenizer import (
+    SecondaryStructureTokenizer,
+)
+from esm.tokenization.structure_tokenizer import (
+    StructureTokenizer,
+)
+from esm.tokenization.tokenizer_base import EsmTokenizerBase
+from esm.utils.constants import esm3 as C
+from esm.utils.function.encode_decode import (
+    decode_function_tokens,
+    decode_residue_annotation_tokens,
+)
+from esm.utils.misc import maybe_list
+from esm.utils.structure.protein_chain import ProteinChain
+from esm.utils.types import FunctionAnnotation
+
+
+def decode_protein_tensor(
+    input: ESMProteinTensor,
+    tokenizers: TokenizerCollectionProtocol,
+    structure_token_decoder: StructureTokenDecoder,
+    function_token_decoder: FunctionTokenDecoder | None = None,
+) -> ESMProtein:
+    input = attr.evolve(input)  # Make a copy
+
+    sequence = None
+    secondary_structure = None
+    sasa = None
+    function_annotations = []
+
+    coordinates = None
+
+    # If all pad tokens, set to None
+    for track in attr.fields(ESMProteinTensor):
+        tokens: torch.Tensor | None = getattr(input, track.name)
+        if track.name == "coordinates" or track.name == "potential_sequence_of_concern":
+            continue
+        if tokens is not None:
+            tokens = tokens[1:-1]  # Remove BOS and EOS tokens
+            tokens = tokens.flatten()  # For multi-track tensors
+            track_tokenizer = getattr(tokenizers, track.name)
+            if torch.all(tokens == track_tokenizer.pad_token_id):
+                setattr(input, track.name, None)
+            # If structure track has any mask tokens, do not decode.
+            if track.name == "structure" and torch.any(
+                tokens == track_tokenizer.mask_token_id
+            ):
+                setattr(input, track.name, None)
+
+    if input.sequence is not None:
+        sequence = decode_sequence(input.sequence, tokenizers.sequence)
+
+    plddt, ptm = None, None
+    if input.structure is not None:
+        # Note: We give priority to the structure tokens over the coordinates when decoding
+        coordinates, plddt, ptm = decode_structure(
+            structure_tokens=input.structure,
+            structure_decoder=structure_token_decoder,
+            structure_tokenizer=tokenizers.structure,
+            sequence=sequence,
+        )
+    elif input.coordinates is not None:
+        coordinates = input.coordinates[1:-1, ...]
+
+    if input.secondary_structure is not None:
+        secondary_structure = decode_secondary_structure(
+            input.secondary_structure, tokenizers.secondary_structure
+        )
+    if input.sasa is not None:
+        sasa = decode_sasa(input.sasa, tokenizers.sasa)
+    if input.function is not None:
+        if function_token_decoder is None:
+            raise ValueError(
+                "Cannot decode function annotations without a function token decoder"
+            )
+        function_track_annotations = decode_function_annotations(
+            input.function,
+            function_token_decoder=function_token_decoder,
+            function_tokenizer=tokenizers.function,
+        )
+        function_annotations.extend(function_track_annotations)
+    if input.residue_annotations is not None:
+        residue_annotations = decode_residue_annotations(
+            input.residue_annotations, tokenizers.residue_annotations
+        )
+        function_annotations.extend(residue_annotations)
+
+    return ESMProtein(
+        sequence=sequence,
+        secondary_structure=secondary_structure,
+        sasa=sasa,  # type: ignore
+        function_annotations=function_annotations if function_annotations else None,
+        coordinates=coordinates,
+        plddt=plddt,
+        ptm=ptm,
+        potential_sequence_of_concern=input.potential_sequence_of_concern,
+    )
+
+
+def _bos_eos_warn(msg: str, tensor: torch.Tensor, tok: EsmTokenizerBase):
+    if tensor[0] != tok.bos_token_id:
+        warnings.warn(
+            f"{msg} does not start with BOS token, token is ignored. BOS={tok.bos_token_id} vs {tensor}"
+        )
+    if tensor[-1] != tok.eos_token_id:
+        warnings.warn(
+            f"{msg} does not end with EOS token, token is ignored. EOS='{tok.eos_token_id}': {tensor}"
+        )
+
+
+def decode_sequence(
+    sequence_tokens: torch.Tensor, sequence_tokenizer: EsmSequenceTokenizer, **kwargs
+) -> str:
+    _bos_eos_warn("Sequence", sequence_tokens, sequence_tokenizer)
+    sequence = sequence_tokenizer.decode(sequence_tokens, **kwargs)
+    sequence = sequence.replace(" ", "")
+    sequence = sequence.replace(sequence_tokenizer.mask_token, C.MASK_STR_SHORT)
+    sequence = sequence.replace(sequence_tokenizer.cls_token, "")
+    sequence = sequence.replace(sequence_tokenizer.eos_token, "")
+
+    return sequence
+
+
+def decode_structure(
+    structure_tokens: torch.Tensor,
+    structure_decoder: StructureTokenDecoder,
+    structure_tokenizer: StructureTokenizer,
+    sequence: str | None = None,
+) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
+    is_singleton = len(structure_tokens.size()) == 1
+    if is_singleton:
+        structure_tokens = structure_tokens.unsqueeze(0)
+    else:
+        raise ValueError(
+            f"Only one structure can be decoded at a time, got structure tokens of shape {structure_tokens.size()}"
+        )
+    _bos_eos_warn("Structure", structure_tokens[0], structure_tokenizer)
+
+    decoder_output = structure_decoder.decode(structure_tokens)
+    bb_coords: torch.Tensor = decoder_output["bb_pred"][
+        0, 1:-1, ...
+    ]  # Remove BOS and EOS tokens
+    bb_coords = bb_coords.detach().cpu()
+
+    if "plddt" in decoder_output:
+        plddt = decoder_output["plddt"][0, 1:-1]
+        plddt = plddt.detach().cpu()
+    else:
+        plddt = None
+
+    if "ptm" in decoder_output:
+        ptm = decoder_output["ptm"]
+    else:
+        ptm = None
+
+    chain = ProteinChain.from_backbone_atom_coordinates(bb_coords, sequence=sequence)
+    chain = chain.infer_oxygen()
+    return torch.tensor(chain.atom37_positions), plddt, ptm
+
+
+def decode_secondary_structure(
+    secondary_structure_tokens: torch.Tensor, ss_tokenizer: SecondaryStructureTokenizer
+) -> str:
+    _bos_eos_warn("Secondary structure", secondary_structure_tokens, ss_tokenizer)
+    secondary_structure_tokens = secondary_structure_tokens[1:-1]
+    secondary_structure = ss_tokenizer.decode(secondary_structure_tokens)
+    return secondary_structure
+
+
+def decode_sasa(
+    sasa_tokens: torch.Tensor, sasa_tokenizer: SASADiscretizingTokenizer
+) -> list[float]:
+    if sasa_tokens[0] != 0:
+        raise ValueError("SASA does not start with 0 corresponding to BOS token")
+    if sasa_tokens[-1] != 0:
+        raise ValueError("SASA does not end with 0 corresponding to EOS token")
+    sasa_tokens = sasa_tokens[1:-1]
+    if sasa_tokens.dtype in [
+        torch.int8,
+        torch.int16,
+        torch.int32,
+        torch.int64,
+        torch.long,
+    ]:
+        # Decode if int
+        # handles turning NaN's into None's
+        sasa = sasa_tokenizer.decode_float(sasa_tokens)
+    else:
+        # If already float, just convert to list
+        sasa = cast(list[float], maybe_list(sasa_tokens, convert_nan_to_none=True))
+
+    return sasa
+
+
+def decode_function_annotations(
+    function_annotation_tokens: torch.Tensor,
+    function_token_decoder: FunctionTokenDecoder,
+    function_tokenizer: InterProQuantizedTokenizer,
+    **kwargs,
+) -> list[FunctionAnnotation]:
+    # No need to check for BOS/EOS as function annotations are not affected
+
+    function_annotations = decode_function_tokens(
+        function_annotation_tokens,
+        function_token_decoder=function_token_decoder,
+        function_tokens_tokenizer=function_tokenizer,
+        **kwargs,
+    )
+    return function_annotations
+
+
+def decode_residue_annotations(
+    residue_annotation_tokens: torch.Tensor,
+    residue_annotation_decoder: ResidueAnnotationsTokenizer,
+) -> list[FunctionAnnotation]:
+    # No need to check for BOS/EOS as function annotations are not affected
+
+    residue_annotations = decode_residue_annotation_tokens(
+        residue_annotations_token_ids=residue_annotation_tokens,
+        residue_annotations_tokenizer=residue_annotation_decoder,
+    )
+    return residue_annotations

Dyna-1/esm/utils/encoding.py

@@ -0,0 +1,246 @@
+from typing import Sequence
+
+import torch
+import torch.nn.functional as F
+
+from esm.models.vqvae import StructureTokenEncoder
+from esm.tokenization.function_tokenizer import (
+    InterProQuantizedTokenizer as EsmFunctionTokenizer,
+)
+
+from esm.tokenization.residue_tokenizer import (
+    ResidueAnnotationsTokenizer,
+)
+from esm.tokenization.sasa_tokenizer import (
+    SASADiscretizingTokenizer,
+)
+from esm.tokenization.sequence_tokenizer import (
+    EsmSequenceTokenizer,
+)
+from esm.tokenization.ss_tokenizer import (
+    SecondaryStructureTokenizer,
+)
+from esm.tokenization.structure_tokenizer import (
+    StructureTokenizer,
+)
+from esm.utils.constants import esm3 as C
+from esm.utils.function.encode_decode import (
+    encode_function_annotations,
+)
+from esm.utils.structure.protein_chain import ProteinChain
+from esm.utils.types import FunctionAnnotation
+
+
+# Raw Defaults
+def get_default_sequence(sequence_length: int) -> str:
+    return C.MASK_STR_SHORT * sequence_length
+
+
+def get_default_secondary_structure(sequence_length: int) -> str:
+    return C.MASK_STR_SHORT * sequence_length
+
+
+def get_default_sasa(sequence_length: int) -> Sequence[float | str | None]:
+    return [None] * sequence_length
+
+
+# Tokenization
+def tokenize_sequence(
+    sequence: str,
+    sequence_tokenizer: EsmSequenceTokenizer,
+    add_special_tokens: bool = True,
+) -> torch.Tensor:
+    sequence = sequence.replace(C.MASK_STR_SHORT, sequence_tokenizer.mask_token)
+    sequence_tokens = sequence_tokenizer.encode(
+        sequence, add_special_tokens=add_special_tokens
+    )
+    sequence_tokens = torch.tensor(sequence_tokens, dtype=torch.int64)
+    return sequence_tokens
+
+
+def tokenize_structure(
+    coordinates: torch.Tensor,
+    structure_encoder: StructureTokenEncoder,
+    structure_tokenizer: StructureTokenizer,
+    reference_sequence: str = "",
+    add_special_tokens: bool = True,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    device = next(structure_encoder.parameters()).device
+    chain = ProteinChain.from_atom37(
+        coordinates, sequence=reference_sequence if reference_sequence else None
+    )
+
+    # Setup padding
+    if reference_sequence and len(reference_sequence) != coordinates.size(0):
+        raise ValueError(
+            f"Reference sequence length ({len(reference_sequence)}) does not match the number of residues in the coordinates ({coordinates.size(0)})"
+        )
+
+    left_pad = 0
+    right_pad = 0
+
+    if add_special_tokens:
+        left_pad += 1  # Add space for BOS token
+        right_pad += 1  # Add space for EOS token
+
+    coordinates, plddt, residue_index = chain.to_structure_encoder_inputs()
+    coordinates = coordinates.to(device)  # (1, L, 37, 3)
+    plddt = plddt.to(device)  # (1, L)
+    residue_index = residue_index.to(device)  # (1, L)
+    _, structure_tokens = structure_encoder.encode(
+        coordinates, residue_index=residue_index
+    )
+    coordinates = torch.squeeze(coordinates, dim=0)  # (L, 37, 3)  # type: ignore
+    plddt = torch.squeeze(plddt, dim=0)  # (L,)  # type: ignore
+    structure_tokens = torch.squeeze(structure_tokens, dim=0)  # (L,)  # type: ignore
+
+    # Add space for BOS and EOS tokens
+    if add_special_tokens:
+        coordinates = F.pad(
+            coordinates, (0, 0, 0, 0, left_pad, right_pad), value=torch.inf
+        )
+        plddt = F.pad(plddt, (left_pad, right_pad), value=0)
+        structure_tokens = F.pad(
+            structure_tokens,
+            (left_pad, right_pad),
+            value=structure_tokenizer.mask_token_id,
+        )
+        structure_tokens[0] = structure_tokenizer.bos_token_id
+        structure_tokens[-1] = structure_tokenizer.eos_token_id
+    return coordinates, plddt, structure_tokens
+
+
+def tokenize_secondary_structure(
+    secondary_structure: str | Sequence[str],
+    secondary_structure_tokenizer: SecondaryStructureTokenizer,
+    add_special_tokens: bool = True,
+) -> torch.Tensor:
+    if isinstance(secondary_structure, str):
+        # Ensure only one char per token
+        secondary_structure = secondary_structure.replace(
+            secondary_structure_tokenizer.mask_token, C.MASK_STR_SHORT
+        )
+
+    # Input as list of chars
+    secondary_structure = [char for char in secondary_structure]
+
+    # Use tokenizer's mask token
+    secondary_structure = [
+        secondary_structure_tokenizer.mask_token if char == C.MASK_STR_SHORT else char
+        for char in secondary_structure
+    ]
+
+    secondary_structure_tokens = secondary_structure_tokenizer.encode(
+        secondary_structure, add_special_tokens=add_special_tokens
+    )
+    return secondary_structure_tokens
+
+
+def tokenize_sasa(
+    sasa: Sequence[float | str | None],
+    sasa_tokenizer: SASADiscretizingTokenizer,
+    add_special_tokens: bool = True,
+):
+    sasa_tokens = sasa_tokenizer.encode(
+        [sasa_tokenizer.mask_token if value is None else value for value in sasa],
+        add_special_tokens=add_special_tokens,
+    )
+    return sasa_tokens
+
+
+def tokenize_function_annotations(
+    function_annotations: Sequence[FunctionAnnotation],
+    reference_sequence: str,
+    function_tokenizer: EsmFunctionTokenizer,
+    residue_annotation_tokenizer: ResidueAnnotationsTokenizer,
+    add_special_tokens: bool = True,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    function_tokens, residue_annotation_tokens = encode_function_annotations(
+        sequence=reference_sequence,
+        function_annotations=function_annotations,
+        function_tokens_tokenizer=function_tokenizer,
+        residue_annotations_tokenizer=residue_annotation_tokenizer,
+        add_special_tokens=add_special_tokens,
+    )
+    return function_tokens, residue_annotation_tokens
+
+
+
+
+# Tokenized Defaults
+def get_default_sequence_tokens(
+    sequence_length: int, sequence_tokenizer: EsmSequenceTokenizer
+) -> torch.Tensor:
+    assert sequence_tokenizer.mask_token_id is not None
+    assert sequence_tokenizer.bos_token_id is not None
+    assert sequence_tokenizer.eos_token_id is not None
+
+    sequence_tokens = torch.full(
+        (sequence_length + 2,), sequence_tokenizer.mask_token_id
+    )
+    sequence_tokens[0] = sequence_tokenizer.bos_token_id
+    sequence_tokens[-1] = sequence_tokenizer.eos_token_id
+
+    return sequence_tokens
+
+
+def get_default_structure_tokens(
+    sequence_length: int, structure_tokenizer: StructureTokenizer
+) -> torch.Tensor:
+    structure_tokens = (
+        torch.ones((sequence_length + 2,), dtype=torch.int64)
+        * structure_tokenizer.mask_token_id
+    )
+    # Always include BOS and EOS tokens
+    structure_tokens[0] = structure_tokenizer.bos_token_id
+    structure_tokens[-1] = structure_tokenizer.eos_token_id
+    return structure_tokens
+
+
+def get_default_secondary_structure_tokens(
+    sequence_length: int, secondary_structure_tokenizer: SecondaryStructureTokenizer
+) -> torch.Tensor:
+    ss8_tokens = torch.full(
+        (sequence_length + 2,), secondary_structure_tokenizer.mask_token_id
+    )
+    ss8_tokens[0] = secondary_structure_tokenizer.bos_token_id
+    ss8_tokens[-1] = secondary_structure_tokenizer.eos_token_id
+
+    return ss8_tokens
+
+
+def get_default_sasa_tokens(
+    sequence_length: int, sasa_tokenizer: SASADiscretizingTokenizer
+) -> torch.Tensor:
+    sasa_tokens = torch.full((sequence_length + 2,), sasa_tokenizer.mask_token_id)
+    sasa_tokens[0] = sasa_tokenizer.bos_token_id
+    sasa_tokens[-1] = sasa_tokenizer.eos_token_id
+    return sasa_tokens
+
+
+def get_default_function_tokens(
+    sequence_length: int, function_tokenizer: EsmFunctionTokenizer
+) -> torch.Tensor:
+    function_tokens = (
+        torch.ones((sequence_length + 2, function_tokenizer.depth), dtype=torch.int64)
+        * function_tokenizer.pad_token_id
+    )
+    # Always include BOS and EOS tokens
+    function_tokens[0] = function_tokenizer.bos_token_id
+    function_tokens[-1] = function_tokenizer.eos_token_id
+    return function_tokens
+
+
+def get_default_residue_annotation_tokens(
+    sequence_length: int, residue_annotation_tokenizer: ResidueAnnotationsTokenizer
+) -> torch.Tensor:
+    residue_annotation_tokens = (
+        torch.ones((sequence_length + 2, C.MAX_RESIDUE_ANNOTATIONS), dtype=torch.int64)
+        * residue_annotation_tokenizer.pad_token_id
+    )
+    # Always include BOS and EOS tokens
+    residue_annotation_tokens[0] = residue_annotation_tokenizer.bos_token_id
+    residue_annotation_tokens[-1] = residue_annotation_tokenizer.eos_token_id
+    return residue_annotation_tokens
+
+

Dyna-1/esm/utils/function/encode_decode.py

@@ -0,0 +1,187 @@
+import re
+from typing import Sequence
+
+import torch
+
+from esm.models.function_decoder import (
+    FunctionTokenDecoder,
+    merge_annotations,
+)
+from esm.tokenization.function_tokenizer import (
+    InterProQuantizedTokenizer,
+)
+from esm.tokenization.residue_tokenizer import (
+    ResidueAnnotationsTokenizer,
+)
+from esm.utils.constants import esm3 as C
+from esm.utils.types import FunctionAnnotation
+
+
+def encode_function_annotations(
+    sequence: str,
+    function_annotations: Sequence[FunctionAnnotation],
+    function_tokens_tokenizer: InterProQuantizedTokenizer,
+    residue_annotations_tokenizer: ResidueAnnotationsTokenizer,
+    add_special_tokens: bool = True,
+) -> tuple[torch.Tensor, torch.Tensor]:
+    assert isinstance(
+        residue_annotations_tokenizer, ResidueAnnotationsTokenizer
+    ), "residue_annotations_tokenizer must be of type ResidueAnnotationsTokenizer"
+
+    # Split the user's annotations by type
+    ft_annotations: list[FunctionAnnotation] = []
+    ra_annotations: list[FunctionAnnotation] = []
+    for fa in function_annotations:
+        assert (
+            1 <= fa.start <= fa.end <= len(sequence)
+        ), f"Invalid (start, end) in function annotation {fa}. Indices 1-indexed and [inclusive, inclusive]"
+
+        supported_label = False
+
+        # Is it an InterPro label?
+        if match := re.search(r"IPR\d+", fa.label):
+            if match.group() in function_tokens_tokenizer.interpro_to_index:
+                ft_annotations.append(fa)
+                supported_label = True
+
+        # Is it a function keyword?
+        if fa.label in function_tokens_tokenizer._tfidf.vocab_to_index:
+            ft_annotations.append(fa)
+            supported_label = True
+
+        # Is it a residue annotation?
+        if fa.label in residue_annotations_tokenizer._labels:
+            ra_annotations.append(fa)
+            supported_label = True
+
+        if not supported_label:
+            raise ValueError(f"Unknown label in FunctionAnnotation: {fa.label}")
+
+    # Convert function token FunctionAnnotations -> Tensor
+    function_tokens = function_tokens_tokenizer.tokenize(
+        annotations=ft_annotations, seqlen=len(sequence)
+    )
+    function_token_ids = function_tokens_tokenizer.encode(
+        function_tokens, add_special_tokens=add_special_tokens
+    )
+
+    # Convert residue annotation FunctionAnnotations -> Tensor
+    if ra_annotations:
+        descriptions, starts, ends = zip(
+            *[(anot.label, anot.start, anot.end) for anot in ra_annotations]
+        )
+    else:
+        descriptions = starts = ends = None
+    ra_tokens = residue_annotations_tokenizer.tokenize(
+        {
+            "interpro_site_descriptions": descriptions,
+            "interpro_site_starts": starts,
+            "interpro_site_ends": ends,
+        },
+        sequence=sequence,
+        fail_on_mismatch=True,
+    )
+    residue_annotation_ids = residue_annotations_tokenizer.encode(
+        ra_tokens, add_special_tokens=add_special_tokens
+    )
+
+    return function_token_ids, residue_annotation_ids
+
+
+def decode_function_tokens(
+    function_token_ids: torch.Tensor,
+    function_token_decoder: FunctionTokenDecoder,
+    function_tokens_tokenizer: InterProQuantizedTokenizer,
+    decoder_annotation_threshold: float = 0.1,
+    annotation_min_length: int | None = 5,
+    annotation_gap_merge_max: int | None = 3,
+) -> list[FunctionAnnotation]:
+    """Decodes model prediction logits into function predictions.
+
+    Merges function token and residue annotation predictions into a single
+    set of FunctionAnnotation predictions.
+
+    Args:
+        function_token_ids: Tensor <float>[length, depth] of
+            function token ids.
+        residue_annotation_logits: Tensor  <float>[length, RA-vocab] of residue
+            annotation binary classification logits.
+        function_tokens_tokenizer: InterPro annotation tokenizer.
+        residue_annotation_threshold: tokenizer of residue annotations.
+        residue_annotation_threshold: predicted probability threshold for emitting
+            a predicted residue annotation.
+    Returns:
+        Predicted function annotations merged from both predictions.
+    """
+    assert (
+        function_token_ids.ndim == 2
+    ), "function_token_ids must be of shape (length, depth)"
+
+    annotations: list[FunctionAnnotation] = []
+
+    # Function Annotations from predicted function tokens.
+    decoded = function_token_decoder.decode(
+        function_token_ids,
+        tokenizer=function_tokens_tokenizer,
+        annotation_threshold=decoder_annotation_threshold,
+        annotation_min_length=annotation_min_length,
+        annotation_gap_merge_max=annotation_gap_merge_max,
+    )
+
+    # Convert predicted InterPro annotation to FunctionAnnotation.
+    annotations.extend(decoded["function_keywords"])
+    for annotation in decoded["interpro_annotations"]:
+        annotation: FunctionAnnotation
+        label = function_tokens_tokenizer.format_annotation(annotation)
+        annotations.append(
+            FunctionAnnotation(label=label, start=annotation.start, end=annotation.end)
+        )
+
+    return annotations
+
+
+def decode_residue_annotation_tokens(
+    residue_annotations_token_ids: torch.Tensor,
+    residue_annotations_tokenizer: ResidueAnnotationsTokenizer,
+    annotation_min_length: int | None = 5,
+    annotation_gap_merge_max: int | None = 3,
+) -> list[FunctionAnnotation]:
+    """Decodes residue annotation tokens into FunctionAnnotations.
+
+    Args:
+        tokens: Tensor <int>[length, MAX_RESIDUE_ANNOTATIONS] of residue annotation tokens.
+        residue_annotations_tokenizer: Tokenizer of residue annotations.
+        threshold: predicted probability threshold for emitting a predicted residue
+            annotation.
+    Returns:
+        Predicted residue annotations.
+    """
+    assert (
+        residue_annotations_token_ids.ndim == 2
+    ), "logits must be of shape (length, MAX_RESIDUE_ANNOTATIONS)"
+
+    annotations: list[FunctionAnnotation] = []
+
+    for depth in range(0, C.MAX_RESIDUE_ANNOTATIONS):
+        token_ids = residue_annotations_token_ids[:, depth]
+        nonzero_indices = torch.nonzero(token_ids).squeeze(dim=1).cpu().numpy()
+        if len(nonzero_indices) == 0:
+            continue
+        for loc in nonzero_indices:
+            vocab_index: int = token_ids[loc].item()  # type: ignore
+            label = residue_annotations_tokenizer.vocabulary[vocab_index]
+            if label not in [*residue_annotations_tokenizer.special_tokens, "<none>"]:
+                annotation = FunctionAnnotation(label=label, start=loc, end=loc)
+                annotations.append(annotation)
+
+    annotations = merge_annotations(annotations, merge_gap_max=annotation_gap_merge_max)
+
+    # Drop very small annotations.
+    if annotation_min_length is not None:
+        annotations = [
+            annotation
+            for annotation in annotations
+            if annotation.end - annotation.start + 1 >= annotation_min_length
+        ]
+
+    return annotations

Dyna-1/esm/utils/function/interpro.py

@@ -0,0 +1,178 @@
+"""Utilities for interacting with InterPro."""
+
+import itertools
+import re
+from dataclasses import dataclass
+from enum import IntEnum, auto
+from functools import cached_property
+
+import networkx as nx
+import pandas as pd
+from cloudpathlib import AnyPath
+
+from esm.utils.constants import esm3 as C
+from esm.utils.types import PathLike
+
+
+def parse_go_terms(text: str) -> list[str]:
+    """Parses GO terms from a string.
+
+    Args:
+        text: String containing GO terms. Example: "GO:0008309, GO:1902267" Note that GO
+          terms have exactly 7 digits.
+    Returns:
+        All GO terms found in the string. Example: ['GO:0008309', 'GO:1902267']
+    """
+    return re.findall(r"GO:(?:\d{7,})", text)
+
+
+def _parse_interpro2go(path: PathLike) -> dict[str, list[str]]:
+    """Parses InterPro2GO file into map.
+
+    NOTE: this file has a very strange, non-standard format.
+
+    Args:
+        path: path to InterPro2GO file from: https://www.ebi.ac.uk/GOA/InterPro2GO
+    Returns:
+        Mapping from InterPro to list of associated GO terms.
+    """
+    with AnyPath(path).open("r") as f:
+        text = f.read()
+    df = pd.Series(text.split("\n"), name="line").to_frame()
+    df = df[~df.line.str.startswith("!")]
+    df["interpro_id"] = df.line.apply(lambda line: re.findall(r"IPR\d+", line))
+    df["go_ids"] = df.line.apply(parse_go_terms)
+    df = df[df.go_ids.apply(len).gt(0) & df.interpro_id.apply(len).eq(1)]
+    df["interpro_id"] = df["interpro_id"].apply(lambda xs: xs[0])  # type: ignore
+
+    # Group all mappints together into a single map.
+    df = (
+        df.groupby("interpro_id")["go_ids"]  # type: ignore
+        .apply(lambda group: list(itertools.chain.from_iterable(group)))
+        .reset_index()
+    )
+    return dict(zip(df.interpro_id, df.go_ids))  # type: ignore
+
+
+class InterProEntryType(IntEnum):
+    """InterPro types and representation counts:
+
+    Family                    21,942
+    Domain                    14,053
+    Homologous_superfamily     3,446
+    Conserved_site               728
+    Repeat                       374
+    Active_site                  133
+    Binding_site                  75
+    PTM                           17
+    """
+
+    ACTIVE_SITE = 0
+    BINDING_SITE = auto()
+    CONSERVED_SITE = auto()
+    DOMAIN = auto()
+    FAMILY = auto()
+    HOMOLOGOUS_SUPERFAMILY = auto()
+    PTM = auto()
+    REPEAT = auto()
+    UNKNOWN = auto()
+
+
+@dataclass
+class InterProEntry:
+    """Represents an InterPro entry."""
+
+    id: str  # Example: IPR000006
+    type: InterProEntryType
+    name: str  # Example: "Metallothionein, vertebrate"
+    description: str | None = None
+
+
+class InterPro:
+    """Convenience class interacting with InterPro ontology/data."""
+
+    def __init__(
+        self,
+        entries_path: PathLike | None = None,
+        hierarchy_path: PathLike | None = None,
+        interpro2go_path: PathLike | None = None,
+    ):
+        """Constructs interface to query InterPro entries."""
+
+        def default(x, d):
+            return x if x is not None else d
+
+        self.entries_path = default(entries_path, C.INTERPRO_ENTRY)
+        self.hierarchy_graph_path = default(hierarchy_path, C.INTERPRO_HIERARCHY)
+        self.interpro2go_path = default(interpro2go_path, C.INTERPRO2GO)
+
+    @cached_property
+    def interpro2go(self) -> dict[str, list[str]]:
+        """Reads the InterPro to GO term mapping."""
+        assert self.interpro2go_path is not None
+        return _parse_interpro2go(self.interpro2go_path)
+
+    @cached_property
+    def entries_frame(self) -> pd.DataFrame:
+        """Loads full InterPro entry set as a DataFrame.
+
+        Colums are
+            - "id": str interpro accession /id as
+            - "type": InterProEntryType representing the type of annotation.
+            - "name": Short name of the entry.
+        """
+        with AnyPath(self.entries_path).open("r") as f:
+            df = pd.read_csv(f, sep="\t")
+        assert all(
+            col in df.columns for col in ["ENTRY_AC", "ENTRY_TYPE", "ENTRY_NAME"]
+        )
+        df.rename(
+            columns={"ENTRY_AC": "id", "ENTRY_TYPE": "type", "ENTRY_NAME": "name"},
+            inplace=True,
+        )
+        df["type"] = df.type.str.upper().apply(
+            lambda type_name: InterProEntryType[type_name]
+        )
+        return df
+
+    @cached_property
+    def entries(self) -> dict[str, InterProEntry]:
+        """Returns all InterPro entries."""
+        return {
+            row.id: InterProEntry(  # type: ignore
+                id=row.id,  # type: ignore
+                type=row.type,  # type: ignore
+                name=row.name,  # type: ignore
+            )
+            for row in self.entries_frame.itertuples()
+        }
+
+    def lookup_name(self, interpro_id: str) -> str | None:
+        """Short name / title for an interpro id."""
+        if interpro_id not in self.entries:
+            return None
+        return self.entries[interpro_id].name
+
+    def lookup_entry_type(self, interpro_id: str) -> InterProEntryType:
+        """Looks up entry-type for an interpro id."""
+        if interpro_id in self.entries:
+            return self.entries[interpro_id].type
+        else:
+            return InterProEntryType.UNKNOWN
+
+    @cached_property
+    def graph(self) -> nx.DiGraph:
+        """Reads the InterPro hierarchy of InterPro."""
+        graph = nx.DiGraph()
+        with AnyPath(self.hierarchy_graph_path).open("r") as f:
+            parents = []
+            for line in f:
+                ipr = line.split("::", maxsplit=1)[0]
+                ipr_strip = ipr.lstrip("-")
+                level = (len(ipr) - len(ipr_strip)) // 2
+                parents = parents[:level]
+                graph.add_node(ipr_strip)
+                if parents:
+                    graph.add_edge(ipr_strip, parents[-1])
+                parents.append(ipr_strip)
+        return graph

Dyna-1/esm/utils/function/lsh.py

@@ -0,0 +1,102 @@
+import numpy as np
+from cloudpathlib import AnyPath
+
+from esm.utils.types import PathLike
+
+
+class LSHTable:
+    def __init__(self, n_bits: int, dim: int, hyperplanes: np.ndarray | None = None):
+        if hyperplanes is None:
+            hyperplanes = np.random.randn(n_bits, dim)
+            hyperplanes = hyperplanes / np.linalg.norm(
+                hyperplanes, axis=-1, keepdims=True
+            )
+        else:
+            assert hyperplanes.shape == (n_bits, dim), (
+                hyperplanes.shape,
+                (n_bits, dim),
+            )
+        assert hyperplanes is not None
+        self.hyperplanes: np.ndarray = hyperplanes
+        self.values = 1 << np.arange(n_bits)
+
+    def __call__(self, array, tokenize: bool = True):
+        similarity = self.hyperplanes @ array.T
+        bits = np.where(similarity >= 0, 1, 0)
+        if tokenize:
+            tokens = bits.T @ self.values
+            return tokens
+        else:
+            return bits.T
+
+
+class LSHTokenized:
+    def __init__(
+        self,
+        n_bits: int,
+        dim: int,
+        num_tables: int = 1,
+        filepath: PathLike | None = None,
+        allow_create_hyperplanes: bool = False,  # set this if you want the lsh to allow creation of hyperplanes
+    ):
+        table_hyperplanes = None
+        if filepath is not None:
+            filepath = AnyPath(filepath)
+            if not filepath.exists():
+                raise FileNotFoundError(filepath)
+            table_hyperplanes = np.load(filepath)  # type: ignore
+            for i in range(num_tables):
+                assert str(i) in table_hyperplanes, f"Missing hyperplane for table {i}"
+        elif not allow_create_hyperplanes:
+            raise RuntimeError(
+                "Not allowed to create hyperplanes but no filepath provided"
+            )
+
+        self.tables = [
+            LSHTable(
+                n_bits,
+                dim,
+                table_hyperplanes[str(i)] if table_hyperplanes is not None else None,
+            )
+            for i in range(num_tables)
+        ]
+
+    def write_hyperplanes(self, filepath: PathLike):
+        hyperplanes: dict[str, np.ndarray] = {  # type: ignore
+            str(i): table.hyperplanes for i, table in enumerate(self.tables)
+        }
+        np.savez(filepath, **hyperplanes)
+
+    def __call__(self, array):
+        tokens = np.stack([table(array) for table in self.tables], 1)
+        return tokens
+
+
+class LSHBitstream:
+    def __init__(
+        self,
+        n_bits: int,
+        dim: int,
+        filepath: PathLike | None = None,
+        allow_create_hyperplanes: bool = False,  # set this if you want the lsh to allow creation of hyperplanes
+    ):
+        table_hyperplanes = None
+        if filepath is not None:
+            filepath = AnyPath(filepath)
+            if not filepath.exists():
+                raise FileNotFoundError(filepath)
+            table_hyperplanes = np.load(filepath)
+        elif not allow_create_hyperplanes:
+            raise RuntimeError(
+                "Not allowed to create hyperplanes but no filepath provided"
+            )
+
+        self.table = LSHTable(
+            n_bits, dim, table_hyperplanes if table_hyperplanes is not None else None
+        )
+
+    def write_hyperplanes(self, filepath: PathLike):
+        np.save(filepath, self.table.hyperplanes)
+
+    def __call__(self, array):
+        return self.table(array, tokenize=False)