Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,47 @@
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin 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
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*.tfevents* filter=lfs diff=lfs merge=lfs -text
+*.db* filter=lfs diff=lfs merge=lfs -text
+*.ark* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*data* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.index filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.gguf* filter=lfs diff=lfs merge=lfs -text
+*.ggml filter=lfs diff=lfs merge=lfs -text
+*.llamafile* filter=lfs diff=lfs merge=lfs -text
+*.pt2 filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.idea/SaProt.iml

@@ -0,0 +1,7 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module version="4">
+  <component name="PyDocumentationSettings">
+    <option name="format" value="GOOGLE" />
+    <option name="myDocStringFormat" value="Google" />
+  </component>
+</module>

LICENSE

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

README.md

@@ -0,0 +1,605 @@
+---
+license: mit
+language:
+- en
+- zh
+tags:
+- OneScience
+- life-science
+- protein
+- protein-language-model
+- structure-aware
+- mutation-effect
+- embedding
+- inverse-folding
+- SaProt
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">SaProt</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+SaProt (Protein Language Modeling with Structure-aware Vocabulary) is a protein language model that jointly models protein amino acid sequences and structural information. Its central idea is to combine amino acids (AA) with the 3Di structural alphabet generated by Foldseek into structure-aware tokens, allowing the model to learn representations from both protein sequences and structural context.
+
+SaProt can be used for protein representation extraction, zero-shot mutation effect prediction, protein inverse folding, and downstream task fine-tuning.
+
+Paper:
+
+> **SaProt: Protein Language Modeling with Structure-aware Vocabulary**  
+> ICLR 2024 Spotlight  
+> Follow-up work was published in Nature Biotechnology (2025)
+
+# Model Description
+
+SaProt models proteins using a structure-aware vocabulary formed by combining amino acids (AA) with the Foldseek 3Di structural alphabet.
+
+For example, a structure-aware sequence can be represented as:
+
+```text
+M#EvVpQpL#VyQdYaKv
+```
+
+Every two characters form a structure-aware token: the first character represents the amino acid, and the second represents the corresponding 3Di structural state. `#` can be used to mask low-confidence structural regions.
+
+The official release provides pretrained models at multiple scales:
+
+| Model | Parameter scale | Training data |
+| --- | ---: | --- |
+| `SaProt_35M_AF2` | 35M | 40M AF2 structures |
+| `SaProt_650M_PDB` | 650M | 40M AF2 structures + 60K PDB structures |
+| `SaProt_650M_AF2` | 650M | 40M AF2 structures |
+| `SaProt_1.3B_AF2` | 1.3B | 40M AF2 structures |
+| `SaProt_1.3B_AFDB_OMG_NCBI` | 1.3B | AFDB + OMG_prot50 + NCBI |
+
+For the 35M and 650M SaProt models, the official recommendation is to use SA-token inputs containing structural information for the best results. The 1.3B version can handle both structure-aware sequences and amino-acid-only sequences relatively well.
+
+# Use Cases
+
+| Use case | Description |
+| --- | --- |
+| Protein representation extraction | Extract residue-level or protein-level embeddings |
+| Zero-shot mutation effect prediction | Evaluate single or multiple mutations directly without task-specific fine-tuning |
+| Structure-aware protein modeling | Jointly use amino acid and 3Di structural tokens |
+| Protein inverse folding | Design sequences from structural information |
+| Downstream task fine-tuning | Apply to tasks such as EC, GO, stability, PPI, Contact, and DeepLoc |
+
+# Usage
+
+## 1. Using OneCode
+
+Experience intelligent one-click AI4S programming in the OneCode online environment:
+
+[Try intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware Requirements**
+
+- SaProt supports inference on CPUs and GPUs/DCUs.
+- The 35M model can be used for lightweight testing; GPUs/DCUs are recommended for the 650M and 1.3B models.
+- Batch embedding, mutation scanning, pretraining, and fine-tuning substantially increase GPU memory and host memory requirements.
+
+### Set Up the Runtime Environment
+
+#### DCU Environment
+
+```bash
+# Activate DTK and CONDA first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+
+pip install onescience[bio] \
+  -i http://mirrors.onescience.ai:3141/pypi/simple/ \
+  --trusted-host mirrors.onescience.ai
+```
+
+#### Environment Notes
+
+- If you encounter missing dependencies or version incompatibilities during execution, refer to the dependency versions specified in `requirements.txt` and install or adjust the relevant dependencies as needed.
+
+### Prepare Models and Data
+
+#### 1) SaProt Model Weights
+
+The official models are primarily released on Hugging Face:
+
+```text
+SaProt_35M_AF2
+https://huggingface.co/westlake-repl/SaProt_35M_AF2
+
+SaProt_650M_PDB
+https://huggingface.co/westlake-repl/SaProt_650M_PDB
+
+SaProt_650M_AF2
+https://huggingface.co/westlake-repl/SaProt_650M_AF2
+
+SaProt_1.3B_AF2
+https://huggingface.co/westlake-repl/SaProt_1.3B_AF2
+
+SaProt_1.3B_AFDB_OMG_NCBI
+https://huggingface.co/westlake-repl/SaProt_1.3B_AFDB_OMG_NCBI
+```
+
+For example, download `SaProt_650M_AF2` in advance for offline use:
+
+```bash
+huggingface-cli download \
+  westlake-repl/SaProt_650M_AF2 \
+  --local-dir ./weight/PLMs/SaProt_650M_AF2
+```
+
+It is recommended to download all model weights under `weight/PLMs/`. The current SaProt configuration reads from:
+
+```text
+weight/PLMs/SaProt_650M_AF2
+```
+
+To run the ESM2 comparison experiment, also prepare:
+
+```bash
+huggingface-cli download \
+  facebook/esm2_t33_650M_UR50D \
+  --local-dir ./weight/PLMs/esm2_t33_650M_UR50D
+```
+
+The corresponding configuration reads from:
+
+```text
+weight/PLMs/esm2_t33_650M_UR50D
+```
+
+#### 2) Foldseek
+
+SaProt structure-aware inputs require PDB/CIF structures to be encoded as Foldseek 3Di sequences first. The official README provides the following download link:
+
+```text
+https://drive.google.com/file/d/1B_9t3n_nlj8Y3Kpc_mMjtMdY0OPYa7Re/view
+```
+
+You can also use the official Foldseek Linux prebuilt package or a Foldseek installation already available on the system or platform.
+
+For this adapted version, place Foldseek at:
+
+```text
+SaProt/
+└── scripts/
+    └── bin/
+        └── foldseek
+```
+
+Then grant it execute permission:
+
+```bash
+chmod +x scripts/bin/foldseek
+```
+
+The Foldseek path used by the current configuration is:
+
+```text
+scripts/bin/foldseek
+```
+
+#### 3) Downstream Task Datasets
+
+The official downstream task datasets are available at:
+
+```text
+https://drive.google.com/drive/folders/11dNGqPYfLE3M-Mbh4U7IQpuHxJpuRr4g?usp=sharing
+```
+
+For this adapted version, extract the downstream task data to:
+
+```text
+scripts/LMDB/
+```
+
+Typical paths used by the configuration include:
+
+```text
+scripts/LMDB/Thermostability/foldseek/train
+scripts/LMDB/Thermostability/foldseek/valid
+scripts/LMDB/Thermostability/foldseek/test
+scripts/LMDB/ProteinGym/substitutions
+scripts/LMDB/ClinVar
+```
+
+#### 4) Pretraining Dataset
+
+To pretrain SaProt from scratch or continue pretraining, prepare the official pretraining data:
+
+```text
+westlake-repl/AF2_UniRef50
+https://huggingface.co/datasets/westlake-repl/AF2_UniRef50
+```
+
+The official pretraining configuration uses LMDB data directories such as:
+
+```text
+scripts/LMDB/AF2_Uniref50/foldseek/train
+scripts/LMDB/AF2_Uniref50/foldseek/valid
+```
+
+The pretraining dataset is large and is only needed when pretraining from scratch or continuing pretraining.
+
+## 3. Quick Start
+
+### Download the Model Package
+
+```bash
+hf download OneScience-Group/SaProt \
+  --local-dir ./SaProt
+
+cd SaProt
+```
+
+### Quick Verification
+
+Check the dependencies:
+
+```bash
+python - <<'PY'
+import torch
+import transformers
+import esm
+import pytorch_lightning as pl
+
+print("torch:", torch.__version__)
+print("transformers:", transformers.__version__)
+print("pytorch_lightning:", pl.__version__)
+print("SaProt dependencies OK")
+PY
+```
+
+Check Foldseek:
+
+```bash
+./scripts/bin/foldseek version
+```
+
+Test model loading:
+
+```bash
+python - <<'PY'
+from transformers import EsmTokenizer, EsmForMaskedLM
+
+model_path = "./weight/PLMs/SaProt_650M_AF2"
+
+tokenizer = EsmTokenizer.from_pretrained(model_path)
+model = EsmForMaskedLM.from_pretrained(model_path)
+
+print("SaProt load OK")
+PY
+```
+
+# Example Data
+
+The official repository provides:
+
+```text
+scripts/example/8ac8.cif
+```
+
+This can be used to demonstrate conversion from a protein structure to a structure-aware sequence.
+
+Your own structure input can be:
+
+```text
+*.pdb
+*.cif
+```
+
+If you already have a Foldseek-encoded structure-aware sequence, you can pass it directly to SaProt without processing the structure file again.
+
+# Inference Examples
+
+## Load SaProt for Forward Inference
+
+```bash
+python - <<'PY'
+import torch
+from transformers import EsmTokenizer, EsmForMaskedLM
+
+model_path = "weight/PLMs/SaProt_650M_AF2"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+tokenizer = EsmTokenizer.from_pretrained(model_path)
+model = EsmForMaskedLM.from_pretrained(model_path)
+model.to(device)
+model.eval()
+
+seq = "M#EvVpQpL#VyQdYaKv"
+tokens = tokenizer.tokenize(seq)
+print(tokens)
+
+inputs = tokenizer(seq, return_tensors="pt")
+inputs = {k: v.to(device) for k, v in inputs.items()}
+
+with torch.no_grad():
+    outputs = model(**inputs)
+
+print(outputs.logits.shape)
+PY
+```
+
+## Load SaProt with the ESM Interface
+
+If the model directory contains `SaProt_650M_AF2.pt`, you can use the ESM loading function provided by the project:
+
+```bash
+python - <<'PY'
+from scripts.utils.esm_loader import load_esm_saprot
+
+model_path = "weight/PLMs/SaProt_650M_AF2/SaProt_650M_AF2.pt"
+model, alphabet = load_esm_saprot(model_path)
+
+print("ESM SaProt load OK")
+PY
+```
+
+## Convert a Structure File to a Structure-Aware Sequence
+
+```bash
+python - <<'PY'
+from scripts.utils.foldseek_util import get_struc_seq
+
+pdb_path = "scripts/example/8ac8.cif"
+
+parsed_seqs = get_struc_seq("scripts/bin/foldseek", pdb_path, ["A"], plddt_mask=False)["A"]
+seq, foldseek_seq, combined_seq = parsed_seqs
+
+print(f"seq: {seq}")
+print(f"foldseek_seq: {foldseek_seq}")
+print(f"combined_seq: {combined_seq}")
+PY
+```
+
+The `A` chain selection extracts only chain A from the structure file. The `combined_seq` in the returned result is a structure-aware sequence that can be used directly by SaProt.
+
+## Mutation Effect Prediction
+
+```bash
+python - <<'PY'
+import torch
+from model.saprot.saprot_foldseek_mutation_model import SaprotFoldseekMutationModel
+
+config = {
+    "foldseek_path": None,
+    "config_path": "weight/PLMs/SaProt_650M_AF2",
+    "load_pretrained": True,
+}
+model = SaprotFoldseekMutationModel(**config)
+tokenizer = model.tokenizer
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.eval()
+model.to(device)
+
+seq = "M#EvVpQpL#VyQdYaKv"
+
+mut_info = "V3A"
+mut_value = model.predict_mut(seq, mut_info)
+print(mut_value)
+
+mut_info = "V3A:Q4M"
+mut_value = model.predict_mut(seq, mut_info)
+print(mut_value)
+
+mut_pos = 3
+mut_dict = model.predict_pos_mut(seq, mut_pos)
+print(mut_dict)
+
+mut_pos = 3
+mut_dict = model.predict_pos_prob(seq, mut_pos)
+print(mut_dict)
+PY
+```
+
+## Extract Protein Embeddings
+
+```bash
+python - <<'PY'
+import torch
+from model.saprot.base import SaprotBaseModel
+from transformers import EsmTokenizer
+
+config = {
+    "task": "base",
+    "config_path": "weight/PLMs/SaProt_650M_AF2",
+    "load_pretrained": True,
+}
+
+model = SaprotBaseModel(**config)
+tokenizer = EsmTokenizer.from_pretrained(config["config_path"])
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+model.eval()
+
+seq = "M#EvVpQpL#VyQdYaKv"
+tokens = tokenizer.tokenize(seq)
+print(tokens)
+
+inputs = tokenizer(seq, return_tensors="pt")
+inputs = {k: v.to(device) for k, v in inputs.items()}
+
+with torch.no_grad():
+    embeddings = model.get_hidden_states(inputs, reduction="mean")
+
+print(embeddings[0].shape)
+PY
+```
+
+## Protein Inverse Folding
+
+Inverse folding requires additional inverse folding model weights:
+
+```text
+https://huggingface.co/westlake-repl/SaProt_650M_AF2_inverse_folding
+```
+
+After downloading, place them at:
+
+```text
+weight/PLMs/SaProt_650M_AF2_inverse_folding
+```
+
+Example:
+
+```bash
+python - <<'PY'
+import torch
+from model.saprot.saprot_if_model import SaProtIFModel
+
+config = {
+    "config_path": "weight/PLMs/SaProt_650M_AF2_inverse_folding",
+    "load_pretrained": True,
+}
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = SaProtIFModel(**config)
+model = model.to(device)
+
+aa_seq = "##########"
+struc_seq = "dddddddddd"
+
+pred_aa_seq = model.predict(aa_seq, struc_seq)
+print(pred_aa_seq)
+PY
+```
+
+# Training
+
+This repository uses a unified training entry point:
+
+```bash
+python scripts/training.py -c <config_path>
+```
+
+Configuration files are located in `conf/`, model code is located in `model/`, and data processing and utility code is located in `scripts/`. The current configuration uses the pretrained model weights in `weight/PLMs/SaProt_650M_AF2` by default.
+
+## Pretraining
+
+To pretrain SaProt from scratch or continue pretraining, first prepare the pretraining LMDB dataset, then run:
+
+```bash
+python scripts/training.py -c conf/pretrain/saprot.yaml
+```
+
+This configuration reads from the following paths by default:
+
+```text
+scripts/LMDB/AF2_Uniref50/foldseek/train
+scripts/LMDB/AF2_Uniref50/foldseek/valid
+```
+
+## Downstream Fine-Tuning
+
+Use the following commands to fine-tune SaProt on downstream tasks:
+
+```bash
+# Thermostability
+python scripts/training.py -c conf/Thermostability/saprot.yaml
+
+# EC
+python scripts/training.py -c conf/EC/saprot.yaml
+
+# GO
+python scripts/training.py -c conf/GO/MF/saprot.yaml
+python scripts/training.py -c conf/GO/BP/saprot.yaml
+python scripts/training.py -c conf/GO/CC/saprot.yaml
+
+# Metal ion binding
+python scripts/training.py -c conf/MetalIonBinding/saprot.yaml
+
+# Human PPI
+python scripts/training.py -c conf/HumanPPI/saprot.yaml
+
+# Contact prediction
+python scripts/training.py -c conf/Contact/saprot.yaml
+
+# DeepLoc
+python scripts/training.py -c conf/DeepLoc/cls2/saprot.yaml
+python scripts/training.py -c conf/DeepLoc/cls10/saprot.yaml
+```
+
+For single-GPU or limited-memory environments, use `conf/scnet/Thermostability_saprot_1gpu.yaml` as a starting point:
+
+```bash
+python scripts/training.py -c conf/scnet/Thermostability_saprot_1gpu.yaml
+```
+
+## Zero-Shot Mutation Effect Evaluation
+
+ProteinGym evaluation:
+
+```bash
+python scripts/mutation_zeroshot.py -c conf/ProteinGym/saprot.yaml
+```
+
+The output file is saved by default to:
+
+```text
+output/ProteinGym/SaProt_650M_AF2.tsv
+```
+
+ClinVar evaluation:
+
+```bash
+python scripts/mutation_zeroshot.py -c conf/ClinVar/saprot.yaml
+python scripts/compute_clinvar_auc.py -c conf/ClinVar/saprot.yaml
+```
+
+ClinVar prediction results are saved by default to:
+
+```text
+output/ClinVar/SaProt_650M_AF2
+```
+
+Single-GPU environments can also use the adapted configuration:
+
+```bash
+python scripts/mutation_zeroshot.py -c conf/scnet/ClinVar_saprot.yaml
+python scripts/compute_clinvar_auc.py -c conf/scnet/ClinVar_saprot.yaml
+```
+
+## ESM2 Comparison Experiment
+
+To run the ESM2 baseline, additionally prepare the weights in `weight/PLMs/esm2_t33_650M_UR50D` and the corresponding normal LMDB data. Example commands:
+
+```bash
+python scripts/training.py -c conf/Thermostability/esm2.yaml
+python scripts/mutation_zeroshot.py -c conf/ProteinGym/esm2.yaml
+```
+
+# Output Description
+
+| Task | Main output |
+| --- | --- |
+| Structure encoding | AA sequence, 3Di sequence, and structure-aware sequence |
+| Model forward pass | Token-level logits |
+| Protein representation | Residue-level/protein-level embeddings |
+| Mutation effect prediction | Mutation score |
+| Zero-shot evaluation | ProteinGym Spearman results or ClinVar AUC results |
+| Inverse folding | Protein sequences generated or evaluated under structural conditions |
+| Downstream fine-tuning | Prediction results and model checkpoint for the corresponding task |
+
+# Official OneScience Information
+
+| Platform | Main OneScience repository | Skills repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation and License
+
+- The official SaProt source repository is released under the **MIT License**, which permits use, modification, distribution, sublicensing, and commercial use. When copying or distributing it, retain the original copyright notice and the MIT License text.
+
+- SaProt model weights are released independently through Hugging Face. For commercial use, redistribution, or other purposes, check and comply with the license on each corresponding model page. The relevant pretraining and downstream datasets are also subject to the licenses and terms of use on their respective dataset pages.
+
+- This repository is a **DCU-adapted version** of SaProt. Use of the repository code, model weights, and related data remains subject to the licenses and terms of use of their respective original projects.
+

conf/ClinVar/esm2.yaml

@@ -0,0 +1,36 @@
+setting:
+  seed: 20000812
+  os_environ:
+    CUDA_VISIBLE_DEVICES: 1,2,3,4,5,6,7
+    MASTER_ADDR: 127.0.0.1
+    MASTER_PORT: 12316
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+
+  out_path: ~
+  dataset_dir: scripts/LMDB/ClinVar
+
+model:
+  model_py_path: saprot/saprot_mutation_model
+  load_prev_scheduler: True
+  kwargs:
+    config_path: weight/PLMs/esm2_t33_650M_UR50D
+    load_pretrained: True
+    log_clinvar: True
+    log_dir: output/ClinVar/esm2_t33_650M_UR50D
+
+dataset:
+  dataset_py_path: mutation_zeroshot_dataset
+  dataloader_kwargs:
+    batch_size: 16
+    num_workers: 8
+  kwargs: {}
+
+Trainer:
+  log_every_n_steps: 1
+  devices: 1
+  accelerator: gpu
+  num_nodes: 1
+  strategy:
+    find_unused_parameters: False
+  logger: False

conf/ClinVar/saprot.yaml

@@ -0,0 +1,37 @@
+setting:
+  seed: 20000812
+  os_environ:
+    CUDA_VISIBLE_DEVICES: 5
+    MASTER_ADDR: 127.0.0.1
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+
+  out_path: ~
+  dataset_dir: scripts/LMDB/ClinVar
+
+model:
+  model_py_path: saprot/saprot_foldseek_mutation_model
+  kwargs:
+    foldseek_path: scripts/bin/foldseek
+    config_path: weight/PLMs/SaProt_650M_AF2
+    load_pretrained: True
+    log_clinvar: True
+    log_dir: output/ClinVar/SaProt_650M_AF2
+    plddt_threshold: 70.
+
+dataset:
+  dataset_py_path: mutation_zeroshot_dataset
+  dataloader_kwargs:
+    batch_size: 16
+    num_workers: 8
+  kwargs: {}
+
+Trainer:
+  log_every_n_steps: 1
+  devices: 1
+  accelerator: gpu
+  num_nodes: 1
+  strategy:
+    find_unused_parameters: False
+  logger: False

conf/Contact/esm2.yaml

@@ -0,0 +1,67 @@
+setting:
+  seed: 20000812
+  use_wandb: true
+  os_environ:
+    WANDB_API_KEY: ~
+    WANDB_RUN_ID: ~
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    MASTER_ADDR: 127.0.0.1
+    MASTER_PORT: 12316
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+  wandb_config:
+    project: TAPE_Contact
+    name: esm2_t33_650M_UR50D
+
+model:
+#    Which model to use
+  model_py_path: saprot/saprot_contact_model
+  kwargs:
+#    Arguments to initialize the specific class
+    config_path: weight/PLMs/esm2_t33_650M_UR50D
+    load_pretrained: True
+
+#    Arguments to initialize the basic class AbstractModel
+  lr_scheduler_kwargs:
+    last_epoch: -1
+    init_lr: 1.0e-2
+    #    Weather to use this scheduler or not
+    on_use: false
+
+  optimizer_kwargs:
+    betas: [0.9, 0.98]
+    weight_decay: 0.01
+
+  save_path: weight/Contact/esm2_t33_650M_UR50D.pt
+  load_prev_scheduler: false
+  save_weights_only: true
+
+
+dataset:
+#    Arguments to initialize the basic class LMDBDataset
+  dataset_py_path: saprot/saprot_contact_dataset
+  dataloader_kwargs:
+    batch_size: 2
+    num_workers: 8
+  train_lmdb: scripts/LMDB/Contact/normal/train
+  valid_lmdb: scripts/LMDB/Contact/normal/valid
+  test_lmdb: scripts/LMDB/Contact/normal/test
+#    Arguments to initialize the specific class
+  kwargs:
+    tokenizer: weight/PLMs/esm2_t33_650M_UR50D
+
+#  Arguments to initialize Pytorch Lightning Trainer
+Trainer:
+  max_epochs: 30
+  log_every_n_steps: 1
+  strategy:
+    find_unused_parameters: False
+  logger: False
+  enable_checkpointing: false
+  val_check_interval: 0.1
+  accelerator: gpu
+  devices: 1
+  num_nodes: 1
+  accumulate_grad_batches: 1
+  precision: 16
+  num_sanity_val_steps: 0

conf/Contact/saprot.yaml

@@ -0,0 +1,66 @@
+setting:
+  seed: 20000812
+  os_environ:
+    WANDB_API_KEY: ~
+    WANDB_RUN_ID: ~
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    MASTER_ADDR: 127.0.0.1
+    MASTER_PORT: 12316
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+  wandb_config:
+    project: TAPE_Contact
+    name: SaProt_650M_AF2
+
+model:
+#    Which model to use
+  model_py_path: saprot/saprot_contact_model
+  kwargs:
+#    Arguments to initialize the specific class
+    config_path: weight/PLMs/SaProt_650M_AF2
+    load_pretrained: True
+
+#    Arguments to initialize the basic class AbstractModel
+  lr_scheduler_kwargs:
+    last_epoch: -1
+    init_lr: 1.0e-2
+    #    Weather to use this scheduler or not
+    on_use: false
+
+  optimizer_kwargs:
+    betas: [0.9, 0.98]
+    weight_decay: 0.01
+
+  save_path: weight/Contact/SaProt_650M_AF2.pt
+  load_prev_scheduler: false
+  save_weights_only: true
+
+
+dataset:
+#    Arguments to initialize the basic class LMDBDataset
+  dataset_py_path: saprot/saprot_contact_dataset
+  dataloader_kwargs:
+    batch_size: 2
+    num_workers: 8
+  train_lmdb: scripts/LMDB/Contact/foldseek/train
+  valid_lmdb: scripts/LMDB/Contact/foldseek/valid
+  test_lmdb: scripts/LMDB/Contact/foldseek/test
+#    Arguments to initialize the specific class
+  kwargs:
+    tokenizer: weight/PLMs/SaProt_650M_AF2
+
+#  Arguments to initialize Pytorch Lightning Trainer
+Trainer:
+  max_epochs: 30
+  log_every_n_steps: 1
+  strategy:
+    find_unused_parameters: False
+  logger: False
+  enable_checkpointing: false
+  val_check_interval: 0.1
+  accelerator: gpu
+  devices: 1
+  num_nodes: 1
+  accumulate_grad_batches: 1
+  precision: 16
+  num_sanity_val_steps: 0

conf/DeepLoc/cls10/esm2.yaml

@@ -0,0 +1,67 @@
+setting:
+  seed: 20000812
+  os_environ:
+    WANDB_API_KEY: ~
+    WANDB_RUN_ID: ~
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    MASTER_ADDR: localhost
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+  wandb_config:
+    project: DeepLoc_cls10
+    name: esm2_t33_650M_UR50D
+
+model:
+#    Which model to use
+  model_py_path: saprot/saprot_classification_model
+  kwargs:
+#    Arguments to initialize the specific class
+    config_path: weight/PLMs/esm2_t33_650M_UR50D
+    load_pretrained: True
+    num_labels: 10
+
+#    Arguments to initialize the basic class AbstractModel
+  lr_scheduler_kwargs:
+    last_epoch: -1
+    init_lr: 2.0e-5
+#    Weather to use this scheduler or not
+    on_use: false
+
+  optimizer_kwargs:
+    betas: [0.9, 0.98]
+    weight_decay: 0.01
+
+  save_path: weight/DeepLoc/cls10/esm2_t33_650M_UR50D.pt
+
+
+dataset:
+#    Arguments to initialize the basic class LMDBDataset
+  dataset_py_path: saprot/saprot_classification_dataset
+  dataloader_kwargs:
+    batch_size: 8
+    num_workers: 8
+
+  train_lmdb: scripts/LMDB/DeepLoc/cls10/normal/train
+  valid_lmdb: scripts/LMDB/DeepLoc/cls10/normal/valid
+  test_lmdb: scripts/LMDB/DeepLoc/cls10/normal/test
+#    Arguments to initialize the specific class
+  kwargs:
+    tokenizer: weight/PLMs/esm2_t33_650M_UR50D
+
+
+#  Arguments to initialize Pytorch Lightning Trainer
+Trainer:
+  max_epochs: 100
+  log_every_n_steps: 1
+  strategy:
+    find_unused_parameters: True
+  logger: True
+  enable_checkpointing: false
+  val_check_interval: 0.5
+  accelerator: gpu
+  devices: 8
+  num_nodes: 1
+  accumulate_grad_batches: 1
+  precision: 16
+  num_sanity_val_steps: 0

conf/DeepLoc/cls10/saprot.yaml

@@ -0,0 +1,68 @@
+setting:
+  seed: 20000812
+  os_environ:
+    WANDB_API_KEY: ~
+    WANDB_RUN_ID: ~
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    MASTER_ADDR: localhost
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+  wandb_config:
+    project: DeepLoc_cls10
+    name: SaProt_650M_AF2
+
+model:
+#    Which model to use
+  model_py_path: saprot/saprot_classification_model
+  kwargs:
+#    Arguments to initialize the specific class
+    config_path: weight/PLMs/SaProt_650M_AF2
+    load_pretrained: True
+    num_labels: 10
+
+#    Arguments to initialize the basic class AbstractModel
+  lr_scheduler_kwargs:
+    last_epoch: -1
+    init_lr: 2.0e-5
+#    Weather to use this scheduler or not
+    on_use: false
+
+  optimizer_kwargs:
+    betas: [0.9, 0.98]
+    weight_decay: 0.01
+
+  save_path: weight/DeepLoc/cls10/SaProt_650M_AF2.pt
+
+
+dataset:
+#    Arguments to initialize the basic class LMDBDataset
+  dataset_py_path: saprot/saprot_classification_dataset
+  dataloader_kwargs:
+    batch_size: 8
+    num_workers: 8
+
+  train_lmdb: scripts/LMDB/DeepLoc/cls10/foldseek/train
+  valid_lmdb: scripts/LMDB/DeepLoc/cls10/foldseek/valid
+  test_lmdb: scripts/LMDB/DeepLoc/cls10/foldseek/test
+#    Arguments to initialize the specific class
+  kwargs:
+    tokenizer: weight/PLMs/SaProt_650M_AF2
+    plddt_threshold: 70.
+
+
+#  Arguments to initialize Pytorch Lightning Trainer
+Trainer:
+  max_epochs: 100
+  log_every_n_steps: 1
+  strategy:
+    find_unused_parameters: True
+  logger: True
+  enable_checkpointing: false
+  val_check_interval: 0.5
+  accelerator: gpu
+  devices: 8
+  num_nodes: 1
+  accumulate_grad_batches: 1
+  precision: 16
+  num_sanity_val_steps: 0

conf/DeepLoc/cls2/esm2.yaml

@@ -0,0 +1,67 @@
+setting:
+  seed: 20000812
+  os_environ:
+    WANDB_API_KEY: ~
+    WANDB_RUN_ID: ~
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    MASTER_ADDR: localhost
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+  wandb_config:
+    project: DeepLoc_cls2
+    name: esm2_t33_650M_UR50D
+
+model:
+#    Which model to use
+  model_py_path: saprot/saprot_classification_model
+  kwargs:
+#    Arguments to initialize the specific class
+    config_path: weight/PLMs/esm2_t33_650M_UR50D
+    load_pretrained: True
+    num_labels: 2
+
+#    Arguments to initialize the basic class AbstractModel
+  lr_scheduler_kwargs:
+    last_epoch: -1
+    init_lr: 2.0e-5
+#    Weather to use this scheduler or not
+    on_use: false
+
+  optimizer_kwargs:
+    betas: [0.9, 0.98]
+    weight_decay: 0.01
+
+  save_path: weight/DeepLoc/cls2/esm2_t33_650M_UR50D.pt
+
+
+dataset:
+#    Arguments to initialize the basic class LMDBDataset
+  dataset_py_path: saprot/saprot_classification_dataset
+  dataloader_kwargs:
+    batch_size: 8
+    num_workers: 8
+
+  train_lmdb: scripts/LMDB/DeepLoc/cls2/normal/train
+  valid_lmdb: scripts/LMDB/DeepLoc/cls2/normal/valid
+  test_lmdb: scripts/LMDB/DeepLoc/cls2/normal/test
+#    Arguments to initialize the specific class
+  kwargs:
+    tokenizer: weight/PLMs/esm2_t33_650M_UR50D
+
+
+#  Arguments to initialize Pytorch Lightning Trainer
+Trainer:
+  max_epochs: 100
+  log_every_n_steps: 1
+  strategy:
+    find_unused_parameters: True
+  logger: True
+  enable_checkpointing: false
+  val_check_interval: 0.5
+  accelerator: gpu
+  devices: 8
+  num_nodes: 1
+  accumulate_grad_batches: 1
+  precision: 16
+  num_sanity_val_steps: 0

conf/DeepLoc/cls2/saprot.yaml

@@ -0,0 +1,68 @@
+setting:
+  seed: 20000812
+  os_environ:
+    WANDB_API_KEY: ~
+    WANDB_RUN_ID: ~
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    MASTER_ADDR: localhost
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+  wandb_config:
+    project: DeepLoc_cls2
+    name: SaProt_650M_AF2
+
+model:
+#    Which model to use
+  model_py_path: saprot/saprot_classification_model
+  kwargs:
+#    Arguments to initialize the specific class
+    config_path: weight/PLMs/SaProt_650M_AF2
+    load_pretrained: True
+    num_labels: 2
+
+#    Arguments to initialize the basic class AbstractModel
+  lr_scheduler_kwargs:
+    last_epoch: -1
+    init_lr: 2.0e-5
+#    Weather to use this scheduler or not
+    on_use: false
+
+  optimizer_kwargs:
+    betas: [0.9, 0.98]
+    weight_decay: 0.01
+
+  save_path: weight/DeepLoc/cls2/SaProt_650M_AF2.pt
+
+
+dataset:
+#    Arguments to initialize the basic class LMDBDataset
+  dataset_py_path: saprot/saprot_classification_dataset
+  dataloader_kwargs:
+    batch_size: 8
+    num_workers: 8
+
+  train_lmdb: scripts/LMDB/DeepLoc/cls2/foldseek/train
+  valid_lmdb: scripts/LMDB/DeepLoc/cls2/foldseek/valid
+  test_lmdb: scripts/LMDB/DeepLoc/cls2/foldseek/test
+#    Arguments to initialize the specific class
+  kwargs:
+    tokenizer: weight/PLMs/SaProt_650M_AF2
+    plddt_threshold: 70
+
+
+#  Arguments to initialize Pytorch Lightning Trainer
+Trainer:
+  max_epochs: 100
+  log_every_n_steps: 1
+  strategy:
+    find_unused_parameters: True
+  logger: True
+  enable_checkpointing: false
+  val_check_interval: 0.5
+  accelerator: gpu
+  devices: 8
+  num_nodes: 1
+  accumulate_grad_batches: 1
+  precision: 16
+  num_sanity_val_steps: 0

conf/EC/esm2.yaml

@@ -0,0 +1,67 @@
+setting:
+  seed: 20000812
+  os_environ:
+    WANDB_API_KEY: ~
+    WANDB_RUN_ID: ~
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    MASTER_ADDR: localhost
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+  wandb_config:
+    project: EC
+    name: esm2_t33_650M_UR50D
+
+model:
+#    Which model to use
+  model_py_path: saprot/saprot_annotation_model
+  kwargs:
+#    Arguments to initialize the specific class
+    config_path: weight/PLMs/esm2_t33_650M_UR50D
+    load_pretrained: True
+    anno_type: EC
+
+#    Arguments to initialize the basic class AbstractModel
+  lr_scheduler_kwargs:
+    last_epoch: -1
+    init_lr: 2.0e-5
+#    Weather to use this scheduler or not
+    on_use: false
+
+  optimizer_kwargs:
+    betas: [0.9, 0.98]
+    weight_decay: 0.01
+
+  save_path: weight/EC/esm2_t33_650M_UR50D.pt
+
+
+dataset:
+#    Arguments to initialize the basic class LMDBDataset
+  dataset_py_path: saprot/saprot_annotation_dataset
+  dataloader_kwargs:
+    batch_size: 8
+    num_workers: 8
+
+  train_lmdb: scripts/LMDB/EC/AF2/normal/train
+  valid_lmdb: scripts/LMDB/EC/AF2/normal/valid
+  test_lmdb: scripts/LMDB/EC/AF2/normal/test
+#    Arguments to initialize the specific class
+  kwargs:
+    tokenizer: weight/PLMs/esm2_t33_650M_UR50D
+
+
+#  Arguments to initialize Pytorch Lightning Trainer
+Trainer:
+  max_epochs: 100
+  log_every_n_steps: 1
+  strategy:
+    find_unused_parameters: True
+  logger: True
+  enable_checkpointing: false
+  val_check_interval: 0.1
+  accelerator: gpu
+  devices: 8
+  num_nodes: 1
+  accumulate_grad_batches: 1
+  precision: 16
+  num_sanity_val_steps: 0

conf/EC/saprot.yaml

@@ -0,0 +1,68 @@
+setting:
+  seed: 20000812
+  os_environ:
+    WANDB_API_KEY: ~
+    WANDB_RUN_ID: ~
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    MASTER_ADDR: localhost
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+  wandb_config:
+    project: EC
+    name: SaProt_650M_AF2
+
+model:
+#    Which model to use
+  model_py_path: saprot/saprot_annotation_model
+  kwargs:
+#    Arguments to initialize the specific class
+    config_path: weight/PLMs/SaProt_650M_AF2
+    load_pretrained: True
+    anno_type: EC
+
+#    Arguments to initialize the basic class AbstractModel
+  lr_scheduler_kwargs:
+    last_epoch: -1
+    init_lr: 2.0e-5
+#    Weather to use this scheduler or not
+    on_use: false
+
+  optimizer_kwargs:
+    betas: [0.9, 0.98]
+    weight_decay: 0.01
+
+  save_path: weight/EC/SaProt_650M_AF2.pt
+
+
+dataset:
+#    Arguments to initialize the basic class LMDBDataset
+  dataset_py_path: saprot/saprot_annotation_dataset
+  dataloader_kwargs:
+    batch_size: 8
+    num_workers: 8
+
+  train_lmdb: scripts/LMDB/EC/AF2/foldseek/train
+  valid_lmdb: scripts/LMDB/EC/AF2/foldseek/valid
+  test_lmdb: scripts/LMDB/EC/AF2/foldseek/test
+#    Arguments to initialize the specific class
+  kwargs:
+    tokenizer: weight/PLMs/SaProt_650M_AF2
+    plddt_threshold: 70
+
+
+#  Arguments to initialize Pytorch Lightning Trainer
+Trainer:
+  max_epochs: 100
+  log_every_n_steps: 1
+  strategy:
+    find_unused_parameters: True
+  logger: True
+  enable_checkpointing: false
+  val_check_interval: 0.1
+  accelerator: gpu
+  devices: 8
+  num_nodes: 1
+  accumulate_grad_batches: 1
+  precision: 16
+  num_sanity_val_steps: 0

conf/GO/BP/esm2.yaml

@@ -0,0 +1,67 @@
+setting:
+  seed: 20000812
+  os_environ:
+    WANDB_API_KEY: ~
+    WANDB_RUN_ID: ~
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    MASTER_ADDR: localhost
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+  wandb_config:
+    project: GO_BP
+    name: esm2_t33_650M_UR50D
+
+model:
+#    Which model to use
+  model_py_path: saprot/saprot_annotation_model
+  kwargs:
+#    Arguments to initialize the specific class
+    config_path: weight/PLMs/esm2_t33_650M_UR50D
+    load_pretrained: True
+    anno_type: GO_BP
+
+#    Arguments to initialize the basic class AbstractModel
+  lr_scheduler_kwargs:
+    last_epoch: -1
+    init_lr: 2.0e-5
+#    Weather to use this scheduler or not
+    on_use: false
+
+  optimizer_kwargs:
+    betas: [0.9, 0.98]
+    weight_decay: 0.01
+
+  save_path: weight/GO/BP/esm2_t33_650M_UR50D.pt
+
+
+dataset:
+#    Arguments to initialize the basic class LMDBDataset
+  dataset_py_path: saprot/saprot_annotation_dataset
+  dataloader_kwargs:
+    batch_size: 4
+    num_workers: 8
+
+  train_lmdb: scripts/LMDB/GO/AF2/BP/normal/train
+  valid_lmdb: scripts/LMDB/GO/AF2/BP/normal/valid
+  test_lmdb: scripts/LMDB/GO/AF2/BP/normal/test
+#    Arguments to initialize the specific class
+  kwargs:
+    tokenizer: weight/PLMs/esm2_t33_650M_UR50D
+
+
+#  Arguments to initialize Pytorch Lightning Trainer
+Trainer:
+  max_epochs: 100
+  log_every_n_steps: 1
+  strategy:
+    find_unused_parameters: True
+  logger: True
+  enable_checkpointing: false
+  val_check_interval: 0.5
+  accelerator: gpu
+  devices: 8
+  num_nodes: 1
+  accumulate_grad_batches: 2
+  precision: 16
+  num_sanity_val_steps: 0

conf/GO/BP/saprot.yaml

@@ -0,0 +1,68 @@
+setting:
+  seed: 20000812
+  os_environ:
+    WANDB_API_KEY: ~
+    WANDB_RUN_ID: ~
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    MASTER_ADDR: localhost
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+  wandb_config:
+    project: GO_BP
+    name: SaProt_650M_AF2
+
+model:
+#    Which model to use
+  model_py_path: saprot/saprot_annotation_model
+  kwargs:
+#    Arguments to initialize the specific class
+    config_path: weight/PLMs/SaProt_650M_AF2
+    load_pretrained: True
+    anno_type: GO_BP
+
+#    Arguments to initialize the basic class AbstractModel
+  lr_scheduler_kwargs:
+    last_epoch: -1
+    init_lr: 2.0e-5
+#    Weather to use this scheduler or not
+    on_use: false
+
+  optimizer_kwargs:
+    betas: [0.9, 0.98]
+    weight_decay: 0.01
+
+  save_path: weight/GO/BP/SaProt_650M_AF2.pt
+
+
+dataset:
+#    Arguments to initialize the basic class LMDBDataset
+  dataset_py_path: saprot/saprot_annotation_dataset
+  dataloader_kwargs:
+    batch_size: 4
+    num_workers: 8
+
+  train_lmdb: scripts/LMDB/GO/AF2/BP/foldseek/train
+  valid_lmdb: scripts/LMDB/GO/AF2/BP/foldseek/valid
+  test_lmdb: scripts/LMDB/GO/AF2/BP/foldseek/test
+#    Arguments to initialize the specific class
+  kwargs:
+    tokenizer: weight/PLMs/SaProt_650M_AF2
+    plddt_threshold: 70
+
+
+#  Arguments to initialize Pytorch Lightning Trainer
+Trainer:
+  max_epochs: 100
+  log_every_n_steps: 1
+  strategy:
+    find_unused_parameters: True
+  logger: True
+  enable_checkpointing: false
+  val_check_interval: 0.5
+  accelerator: gpu
+  devices: 8
+  num_nodes: 1
+  accumulate_grad_batches: 2
+  precision: 16
+  num_sanity_val_steps: 0

conf/GO/CC/esm2.yaml

@@ -0,0 +1,67 @@
+setting:
+  seed: 20000812
+  os_environ:
+    WANDB_API_KEY: ~
+    WANDB_RUN_ID: ~
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    MASTER_ADDR: localhost
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+  wandb_config:
+    project: GO_CC
+    name: esm2_t33_650M_UR50D
+
+model:
+#    Which model to use
+  model_py_path: saprot/saprot_annotation_model
+  kwargs:
+#    Arguments to initialize the specific class
+    config_path: weight/PLMs/esm2_t33_650M_UR50D
+    load_pretrained: True
+    anno_type: GO_CC
+
+#    Arguments to initialize the basic class AbstractModel
+  lr_scheduler_kwargs:
+    last_epoch: -1
+    init_lr: 2.0e-5
+#    Weather to use this scheduler or not
+    on_use: false
+
+  optimizer_kwargs:
+    betas: [0.9, 0.98]
+    weight_decay: 0.01
+
+  save_path: weight/GO/CC/esm2_t33_650M_UR50D.pt
+
+
+dataset:
+#    Arguments to initialize the basic class LMDBDataset
+  dataset_py_path: saprot/saprot_annotation_dataset
+  dataloader_kwargs:
+    batch_size: 4
+    num_workers: 8
+
+  train_lmdb: scripts/LMDB/GO/AF2/CC/normal/train
+  valid_lmdb: scripts/LMDB/GO/AF2/CC/normal/valid
+  test_lmdb: scripts/LMDB/GO/AF2/CC/normal/test
+#    Arguments to initialize the specific class
+  kwargs:
+    tokenizer: weight/PLMs/esm2_t33_650M_UR50D
+
+
+#  Arguments to initialize Pytorch Lightning Trainer
+Trainer:
+  max_epochs: 100
+  log_every_n_steps: 1
+  strategy:
+    find_unused_parameters: True
+  logger: True
+  enable_checkpointing: false
+  val_check_interval: 0.5
+  accelerator: gpu
+  devices: 8
+  num_nodes: 1
+  accumulate_grad_batches: 2
+  precision: 16
+  num_sanity_val_steps: 0

conf/GO/CC/saprot.yaml

@@ -0,0 +1,68 @@
+setting:
+  seed: 20000812
+  os_environ:
+    WANDB_API_KEY: ~
+    WANDB_RUN_ID: ~
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    MASTER_ADDR: localhost
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+  wandb_config:
+    project: GO_CC
+    name: SaProt_650M_AF2
+
+model:
+#    Which model to use
+  model_py_path: saprot/saprot_annotation_model
+  kwargs:
+#    Arguments to initialize the specific class
+    config_path: weight/PLMs/SaProt_650M_AF2
+    load_pretrained: True
+    anno_type: GO_CC
+
+#    Arguments to initialize the basic class AbstractModel
+  lr_scheduler_kwargs:
+    last_epoch: -1
+    init_lr: 2.0e-5
+#    Weather to use this scheduler or not
+    on_use: false
+
+  optimizer_kwargs:
+    betas: [0.9, 0.98]
+    weight_decay: 0.01
+
+  save_path: weight/GO/CC/SaProt_650M_AF2.pt
+
+
+dataset:
+#    Arguments to initialize the basic class LMDBDataset
+  dataset_py_path: saprot/saprot_annotation_dataset
+  dataloader_kwargs:
+    batch_size: 4
+    num_workers: 8
+
+  train_lmdb: scripts/LMDB/GO/AF2/CC/foldseek/train
+  valid_lmdb: scripts/LMDB/GO/AF2/CC/foldseek/valid
+  test_lmdb: scripts/LMDB/GO/AF2/CC/foldseek/test
+#    Arguments to initialize the specific class
+  kwargs:
+    tokenizer: weight/PLMs/SaProt_650M_AF2
+    plddt_threshold: 70
+
+
+#  Arguments to initialize Pytorch Lightning Trainer
+Trainer:
+  max_epochs: 100
+  log_every_n_steps: 1
+  strategy:
+    find_unused_parameters: True
+  logger: True
+  enable_checkpointing: false
+  val_check_interval: 0.5
+  accelerator: gpu
+  devices: 8
+  num_nodes: 1
+  accumulate_grad_batches: 2
+  precision: 16
+  num_sanity_val_steps: 0

conf/GO/MF/esm2.yaml

@@ -0,0 +1,67 @@
+setting:
+  seed: 20000812
+  os_environ:
+    WANDB_API_KEY: ~
+    WANDB_RUN_ID: ~
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    MASTER_ADDR: localhost
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+  wandb_config:
+    project: GO_MF
+    name: esm2_t33_650M_UR50D
+
+model:
+#    Which model to use
+  model_py_path: saprot/saprot_annotation_model
+  kwargs:
+#    Arguments to initialize the specific class
+    config_path: weight/PLMs/esm2_t33_650M_UR50D
+    load_pretrained: True
+    anno_type: GO_MF
+
+#    Arguments to initialize the basic class AbstractModel
+  lr_scheduler_kwargs:
+    last_epoch: -1
+    init_lr: 2.0e-5
+#    Weather to use this scheduler or not
+    on_use: false
+
+  optimizer_kwargs:
+    betas: [0.9, 0.98]
+    weight_decay: 0.01
+
+  save_path: weight/GO/MF/esm2_t33_650M_UR50D.pt
+
+
+dataset:
+#    Arguments to initialize the basic class LMDBDataset
+  dataset_py_path: saprot/saprot_annotation_dataset
+  dataloader_kwargs:
+    batch_size: 4
+    num_workers: 8
+
+  train_lmdb: scripts/LMDB/GO/AF2/MF/normal/train
+  valid_lmdb: scripts/LMDB/GO/AF2/MF/normal/valid
+  test_lmdb: scripts/LMDB/GO/AF2/MF/normal/test
+#    Arguments to initialize the specific class
+  kwargs:
+    tokenizer: weight/PLMs/esm2_t33_650M_UR50D
+
+
+#  Arguments to initialize Pytorch Lightning Trainer
+Trainer:
+  max_epochs: 100
+  log_every_n_steps: 1
+  strategy:
+    find_unused_parameters: True
+  logger: True
+  enable_checkpointing: false
+  val_check_interval: 0.5
+  accelerator: gpu
+  devices: 8
+  num_nodes: 1
+  accumulate_grad_batches: 2
+  precision: 16
+  num_sanity_val_steps: 0

conf/GO/MF/saprot.yaml

@@ -0,0 +1,68 @@
+setting:
+  seed: 20000812
+  os_environ:
+    WANDB_API_KEY: ~
+    WANDB_RUN_ID: ~
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    MASTER_ADDR: localhost
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+  wandb_config:
+    project: GO_MF
+    name: SaProt_650M_AF2
+
+model:
+#    Which model to use
+  model_py_path: saprot/saprot_annotation_model
+  kwargs:
+#    Arguments to initialize the specific class
+    config_path: weight/PLMs/SaProt_650M_AF2
+    load_pretrained: True
+    anno_type: GO_MF
+
+#    Arguments to initialize the basic class AbstractModel
+  lr_scheduler_kwargs:
+    last_epoch: -1
+    init_lr: 2.0e-5
+#    Weather to use this scheduler or not
+    on_use: false
+
+  optimizer_kwargs:
+    betas: [0.9, 0.98]
+    weight_decay: 0.01
+
+  save_path: weight/GO/MF/SaProt_650M_AF2.pt
+
+
+dataset:
+#    Arguments to initialize the basic class LMDBDataset
+  dataset_py_path: saprot/saprot_annotation_dataset
+  dataloader_kwargs:
+    batch_size: 4
+    num_workers: 8
+
+  train_lmdb: scripts/LMDB/GO/AF2/MF/foldseek/train
+  valid_lmdb: scripts/LMDB/GO/AF2/MF/foldseek/valid
+  test_lmdb: scripts/LMDB/GO/AF2/MF/foldseek/test
+#    Arguments to initialize the specific class
+  kwargs:
+    tokenizer: weight/PLMs/SaProt_650M_AF2
+    plddt_threshold: 70
+
+
+#  Arguments to initialize Pytorch Lightning Trainer
+Trainer:
+  max_epochs: 100
+  log_every_n_steps: 1
+  strategy:
+    find_unused_parameters: True
+  logger: True
+  enable_checkpointing: false
+  val_check_interval: 0.5
+  accelerator: gpu
+  devices: 8
+  num_nodes: 1
+  accumulate_grad_batches: 2
+  precision: 16
+  num_sanity_val_steps: 0

conf/HumanPPI/esm2.yaml

@@ -0,0 +1,66 @@
+setting:
+  seed: 20000812
+  os_environ:
+    WANDB_API_KEY: ~
+    WANDB_RUN_ID: ~
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    MASTER_ADDR: localhost
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+  wandb_config:
+    project: HumanPPI
+    name: esm2_t33_650M_UR50D
+
+model:
+#    Which model to use
+  model_py_path: saprot/saprot_ppi_model
+  kwargs:
+#    Arguments to initialize the specific class
+    config_path: weight/PLMs/esm2_t33_650M_UR50D
+    load_pretrained: True
+
+#    Arguments to initialize the basic class AbstractModel
+  lr_scheduler_kwargs:
+    last_epoch: -1
+    init_lr: 2.0e-5
+#    Weather to use this scheduler or not
+    on_use: false
+
+  optimizer_kwargs:
+    betas: [0.9, 0.98]
+    weight_decay: 0.01
+
+  save_path: weight/HumanPPI/esm2_t33_650M_UR50D.pt
+
+
+dataset:
+#    Arguments to initialize the basic class LMDBDataset
+  dataset_py_path: saprot/saprot_ppi_dataset
+  dataloader_kwargs:
+    batch_size: 4
+    num_workers: 8
+
+  train_lmdb: scripts/LMDB/HumanPPI/normal/train
+  valid_lmdb: scripts/LMDB/HumanPPI/normal/valid
+  test_lmdb: scripts/LMDB/HumanPPI/normal/test
+#    Arguments to initialize the specific class
+  kwargs:
+    tokenizer: weight/PLMs/esm2_t33_650M_UR50D
+
+
+#  Arguments to initialize Pytorch Lightning Trainer
+Trainer:
+  max_epochs: 20
+  log_every_n_steps: 1
+  strategy:
+    find_unused_parameters: True
+  logger: True
+  enable_checkpointing: false
+  val_check_interval: 0.1
+  accelerator: gpu
+  devices: 8
+  num_nodes: 1
+  accumulate_grad_batches: 2
+  precision: 16
+  num_sanity_val_steps: 0

conf/HumanPPI/saprot.yaml

@@ -0,0 +1,67 @@
+setting:
+  seed: 20000812
+  os_environ:
+    WANDB_API_KEY: ~
+    WANDB_RUN_ID: ~
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    MASTER_ADDR: localhost
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+  wandb_config:
+    project: HumanPPI
+    name: SaProt_650M_AF2
+
+model:
+#    Which model to use
+  model_py_path: saprot/saprot_ppi_model
+  kwargs:
+#    Arguments to initialize the specific class
+    config_path: weight/PLMs/SaProt_650M_AF2
+    load_pretrained: True
+
+#    Arguments to initialize the basic class AbstractModel
+  lr_scheduler_kwargs:
+    last_epoch: -1
+    init_lr: 2.0e-5
+#    Weather to use this scheduler or not
+    on_use: false
+
+  optimizer_kwargs:
+    betas: [0.9, 0.98]
+    weight_decay: 0.01
+
+  save_path: weight/HumanPPI/SaProt_650M_AF2.pt
+
+
+dataset:
+#    Arguments to initialize the basic class LMDBDataset
+  dataset_py_path: saprot/saprot_ppi_dataset
+  dataloader_kwargs:
+    batch_size: 4
+    num_workers: 8
+
+  train_lmdb: scripts/LMDB/HumanPPI/foldseek/train
+  valid_lmdb: scripts/LMDB/HumanPPI/foldseek/valid
+  test_lmdb: scripts/LMDB/HumanPPI/foldseek/test
+#    Arguments to initialize the specific class
+  kwargs:
+    tokenizer: weight/PLMs/SaProt_650M_AF2
+    plddt_threshold: 70
+
+
+#  Arguments to initialize Pytorch Lightning Trainer
+Trainer:
+  max_epochs: 20
+  log_every_n_steps: 1
+  strategy:
+    find_unused_parameters: True
+  logger: True
+  enable_checkpointing: false
+  val_check_interval: 0.1
+  accelerator: gpu
+  devices: 8
+  num_nodes: 1
+  accumulate_grad_batches: 2
+  precision: 16
+  num_sanity_val_steps: 0

conf/MetalIonBinding/esm2.yaml

@@ -0,0 +1,67 @@
+setting:
+  seed: 20000812
+  os_environ:
+    WANDB_API_KEY: ~
+    WANDB_RUN_ID: ~
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    MASTER_ADDR: localhost
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+  wandb_config:
+    project: MetalIonBinding
+    name: esm2_t33_650M_UR50D
+
+model:
+#    Which model to use
+  model_py_path: saprot/saprot_classification_model
+  kwargs:
+#    Arguments to initialize the specific class
+    config_path: weight/PLMs/esm2_t33_650M_UR50D
+    load_pretrained: True
+    num_labels: 2
+
+#    Arguments to initialize the basic class AbstractModel
+  lr_scheduler_kwargs:
+    last_epoch: -1
+    init_lr: 2.0e-5
+#    Weather to use this scheduler or not
+    on_use: false
+
+  optimizer_kwargs:
+    betas: [0.9, 0.98]
+    weight_decay: 0.01
+
+  save_path: weight/MetalIonBinding/esm2_t33_650M_UR50D.pt
+
+
+dataset:
+#    Arguments to initialize the basic class LMDBDataset
+  dataset_py_path: saprot/saprot_classification_dataset
+  dataloader_kwargs:
+    batch_size: 8
+    num_workers: 8
+
+  train_lmdb: scripts/LMDB/MetalIonBinding/AF2/normal/train
+  valid_lmdb: scripts/LMDB/MetalIonBinding/AF2/normal/valid
+  test_lmdb: scripts/LMDB/MetalIonBinding/AF2/normal/test
+#    Arguments to initialize the specific class
+  kwargs:
+    tokenizer: weight/PLMs/esm2_t33_650M_UR50D
+
+
+#  Arguments to initialize Pytorch Lightning Trainer
+Trainer:
+  max_epochs: 200
+  log_every_n_steps: 1
+  strategy:
+    find_unused_parameters: True
+  logger: True
+  enable_checkpointing: false
+  val_check_interval: 0.5
+  accelerator: gpu
+  devices: 4
+  num_nodes: 1
+  accumulate_grad_batches: 2
+  precision: 16
+  num_sanity_val_steps: 0

conf/MetalIonBinding/saprot.yaml

@@ -0,0 +1,68 @@
+setting:
+  seed: 20000812
+  os_environ:
+    WANDB_API_KEY: ~
+    WANDB_RUN_ID: ~
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    MASTER_ADDR: localhost
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+  wandb_config:
+    project: MetalIonBinding
+    name: SaProt_650M_AF2
+
+model:
+#    Which model to use
+  model_py_path: saprot/saprot_classification_model
+  kwargs:
+#    Arguments to initialize the specific class
+    config_path: weight/PLMs/SaProt_650M_AF2
+    load_pretrained: True
+    num_labels: 2
+
+#    Arguments to initialize the basic class AbstractModel
+  lr_scheduler_kwargs:
+    last_epoch: -1
+    init_lr: 2.0e-5
+#    Weather to use this scheduler or not
+    on_use: false
+
+  optimizer_kwargs:
+    betas: [0.9, 0.98]
+    weight_decay: 0.01
+
+  save_path: weight/MetalIonBinding/SaProt_650M_AF2.pt
+
+
+dataset:
+#    Arguments to initialize the basic class LMDBDataset
+  dataset_py_path: saprot/saprot_classification_dataset
+  dataloader_kwargs:
+    batch_size: 8
+    num_workers: 8
+
+  train_lmdb: scripts/LMDB/MetalIonBinding/AF2/foldseek/train
+  valid_lmdb: scripts/LMDB/MetalIonBinding/AF2/foldseek/valid
+  test_lmdb: scripts/LMDB/MetalIonBinding/AF2/foldseek/test
+#    Arguments to initialize the specific class
+  kwargs:
+    tokenizer: weight/PLMs/SaProt_650M_AF2
+    plddt_threshold: 70
+
+
+#  Arguments to initialize Pytorch Lightning Trainer
+Trainer:
+  max_epochs: 200
+  log_every_n_steps: 1
+  strategy:
+    find_unused_parameters: True
+  logger: True
+  enable_checkpointing: false
+  val_check_interval: 0.5
+  accelerator: gpu
+  devices: 4
+  num_nodes: 1
+  accumulate_grad_batches: 2
+  precision: 16
+  num_sanity_val_steps: 0

conf/ProteinGym/esm2.yaml

@@ -0,0 +1,33 @@
+setting:
+  seed: 20000812
+  os_environ:
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    MASTER_ADDR: 127.0.0.1
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+
+  out_path: output/ProteinGym/esm2_t33_650M_UR50D.tsv
+  dataset_dir: scripts/LMDB/ProteinGym/substitutions
+
+model:
+  model_py_path: saprot/esm_mutation_model
+  kwargs:
+    config_path: weight/PLMs/esm2_t33_650M_UR50D
+    load_pretrained: True
+
+dataset:
+  dataset_py_path: mutation_zeroshot_dataset
+  dataloader_kwargs:
+    batch_size: 32
+    num_workers: 8
+  kwargs: {}
+
+Trainer:
+  log_every_n_steps: 1
+  devices: 1
+  accelerator: gpu
+  num_nodes: 1
+  strategy:
+    find_unused_parameters: False
+  logger: False

conf/ProteinGym/saprot.yaml

@@ -0,0 +1,36 @@
+setting:
+  seed: 20000812
+  os_environ:
+    #CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    CUDA_VISIBLE_DEVICES: 6
+    MASTER_ADDR: 127.0.0.1
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+
+  out_path: output/ProteinGym/SaProt_650M_AF2.tsv
+  dataset_dir: scripts/LMDB/ProteinGym/substitutions
+
+model:
+  model_py_path: saprot/saprot_foldseek_mutation_model
+  kwargs:
+    foldseek_path: scripts/bin/foldseek
+    config_path: weight/PLMs/SaProt_650M_AF2
+    load_pretrained: True
+    plddt_threshold: 70
+
+dataset:
+  dataset_py_path: mutation_zeroshot_dataset
+  dataloader_kwargs:
+    batch_size: 32
+    num_workers: 8
+  kwargs: {}
+
+Trainer:
+  log_every_n_steps: 1
+  devices: 1
+  accelerator: gpu
+  num_nodes: 1
+  strategy:
+    find_unused_parameters: False
+  logger: False

conf/Thermostability/esm2.yaml

@@ -0,0 +1,69 @@
+setting:
+  seed: 20000812
+  os_environ:
+    WANDB_API_KEY: ~
+    WANDB_RUN_ID: ~
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    MASTER_ADDR: localhost
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+  wandb_config:
+    project: Thermostability
+    name: esm2_t33_650M_UR50D
+
+model:
+#    Which model to use
+  model_py_path: saprot/saprot_regression_model
+  kwargs:
+#    Arguments to initialize the specific class
+    config_path: weight/PLMs/esm2_t33_650M_UR50D
+    load_pretrained: True
+
+#    Arguments to initialize the basic class AbstractModel
+  lr_scheduler_kwargs:
+    last_epoch: -1
+    init_lr: 2.0e-5
+#    Weather to use this scheduler or not
+    on_use: false
+
+  optimizer_kwargs:
+    betas: [0.9, 0.98]
+    weight_decay: 0.01
+
+  save_path: weight/Thermostability/esm2_t33_650M_UR50D.pt
+
+
+dataset:
+#    Arguments to initialize the basic class LMDBDataset
+  dataset_py_path: saprot/saprot_regression_dataset
+  dataloader_kwargs:
+    batch_size: 8
+    num_workers: 8
+
+  train_lmdb: scripts/LMDB/Thermostability/normal/train
+  valid_lmdb: scripts/LMDB/Thermostability/normal/valid
+  test_lmdb: scripts/LMDB/Thermostability/normal/test
+#    Arguments to initialize the specific class
+  kwargs:
+    tokenizer: weight/PLMs/esm2_t33_650M_UR50D
+    mix_max_norm: [40, 67]
+
+
+#  Arguments to initialize Pytorch Lightning Trainer
+Trainer:
+  max_epochs: 200
+  log_every_n_steps: 1
+  strategy:
+    find_unused_parameters: True
+  logger: True
+  enable_checkpointing: false
+  val_check_interval: 0.5
+  accelerator: gpu
+  devices: 8
+  num_nodes: 1
+  accumulate_grad_batches: 1
+  precision: 16
+  num_sanity_val_steps: 0
+
+

conf/Thermostability/saprot.yaml

@@ -0,0 +1,68 @@
+setting:
+  seed: 20000812
+  os_environ:
+    WANDB_API_KEY: ~
+    WANDB_RUN_ID: ~
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    MASTER_ADDR: localhost
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+  wandb_config:
+    project: Thermostability
+    name: SaProt_650M_AF2
+
+model:
+#    Which model to use
+  model_py_path: saprot/saprot_regression_model
+  kwargs:
+#    Arguments to initialize the specific class
+    config_path: weight/PLMs/SaProt_650M_AF2
+    load_pretrained: True
+
+#    Arguments to initialize the basic class AbstractModel
+  lr_scheduler_kwargs:
+    last_epoch: -1
+    init_lr: 2.0e-5
+#    Weather to use this scheduler or not
+    on_use: false
+
+  optimizer_kwargs:
+    betas: [0.9, 0.98]
+    weight_decay: 0.01
+
+  save_path: weight/Thermostability/SaProt_650M_AF2.pt
+
+
+dataset:
+#    Arguments to initialize the basic class LMDBDataset
+  dataset_py_path: saprot/saprot_regression_dataset
+  dataloader_kwargs:
+    batch_size: 8
+    num_workers: 8
+
+  train_lmdb: scripts/LMDB/Thermostability/foldseek/train
+  valid_lmdb: scripts/LMDB/Thermostability/foldseek/valid
+  test_lmdb: scripts/LMDB/Thermostability/foldseek/test
+#    Arguments to initialize the specific class
+  kwargs:
+    tokenizer: weight/PLMs/SaProt_650M_AF2
+    mix_max_norm: [40, 67]
+    plddt_threshold: 70
+
+
+#  Arguments to initialize Pytorch Lightning Trainer
+Trainer:
+  max_epochs: 200
+  log_every_n_steps: 1
+  strategy:
+    find_unused_parameters: True
+  logger: True
+  enable_checkpointing: false
+  val_check_interval: 0.5
+  accelerator: gpu
+  devices: 4
+  num_nodes: 1
+  accumulate_grad_batches: 2
+  precision: 16
+  num_sanity_val_steps: 0

conf/pretrain/saprot.yaml

@@ -0,0 +1,81 @@
+setting:
+  seed: 20000812
+  os_environ:
+    WANDB_API_KEY: ~
+    WANDB_RUN_ID: ~
+    CUDA_VISIBLE_DEVICES: 0,1,2,3,4,5,6,7
+    MASTER_ADDR: 127.0.0.1
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+#      WANDB_MODE: offline
+  wandb_config:
+    project: pretraining
+    name: SaProt_650M_AF2
+
+model:
+#    Which model to use
+  model_py_path: saprot/saprot_lm_model
+  kwargs:
+#    Arguments to initialize the specific class
+    config_path: weight/PLMs/SaProt_650M_AF2
+    load_pretrained: True
+
+#    Arguments to initialize the basic class AbstractModel
+  lr_scheduler_kwargs:
+    last_epoch: -1
+    init_lr: 0
+    max_lr: 4.0e-4
+    final_lr: 4.0e-5
+    warmup_steps: 2000
+    start_decay_after_n_steps: 150000
+    end_decay_after_n_steps: 1500000
+#    Weather to use this scheduler or not
+    on_use: True
+
+  optimizer_kwargs:
+    betas: [0.9, 0.98]
+    weight_decay: 0.01
+
+  save_path: weight/Pretrain/SaProt_650M_AF2.pt
+  load_prev_scheduler: false
+  save_weights_only: false
+
+
+dataset:
+#    Arguments to initialize the basic class LMDBDataset
+  dataset_py_path: saprot/saprot_foldseek_dataset
+  dataloader_kwargs:
+    batch_size: 8
+    num_workers: 8
+
+  train_lmdb: scripts/LMDB/AF2_Uniref50/foldseek/train
+  valid_lmdb: scripts/LMDB/AF2_Uniref50/foldseek/valid
+  test_lmdb: scripts/LMDB/AF2_Uniref50/foldseek/valid
+#    Arguments to initialize the specific class
+  kwargs:
+    tokenizer: weight/PLMs/SaProt_650M_AF2
+    max_length: 1024
+    mask_ratio: 0.15
+
+
+#  Arguments to initialize Pytorch Lightning Trainer
+Trainer:
+#  max_epochs: 100
+  max_steps: 1000000
+  min_steps: 1000000
+  log_every_n_steps: 1
+  accelerator: gpu
+  devices: 1
+  accumulate_grad_batches: 1
+  num_nodes: 1
+  strategy:
+    find_unused_parameters: True
+  logger: False
+  enable_checkpointing: false
+  num_sanity_val_steps: 0
+  val_check_interval: 100
+  precision: 16
+  limit_val_batches: 1.0
+  limit_test_batches: 1.0
+

conf/scnet/ClinVar_saprot.yaml

@@ -0,0 +1,37 @@
+setting:
+  seed: 20000812
+  os_environ:
+    CUDA_VISIBLE_DEVICES: 0
+    MASTER_ADDR: 127.0.0.1
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+
+  out_path: ~
+  dataset_dir: scripts/LMDB/ClinVar
+
+model:
+  model_py_path: saprot/saprot_foldseek_mutation_model
+  kwargs:
+    foldseek_path: scripts/bin/foldseek
+    config_path: weight/PLMs/SaProt_650M_AF2
+    load_pretrained: True
+    log_clinvar: True
+    log_dir: output/ClinVar/SaProt_650M_AF2_scnet
+    plddt_threshold: 70.
+
+dataset:
+  dataset_py_path: mutation_zeroshot_dataset
+  dataloader_kwargs:
+    batch_size: 8
+    num_workers: 4
+  kwargs: {}
+
+Trainer:
+  log_every_n_steps: 1
+  devices: 1
+  accelerator: gpu
+  num_nodes: 1
+  strategy:
+    find_unused_parameters: False
+  logger: False

conf/scnet/Thermostability_saprot_1gpu.yaml

@@ -0,0 +1,68 @@
+setting:
+  seed: 20000812
+  os_environ:
+    WANDB_API_KEY: ~
+    WANDB_RUN_ID: ~
+    CUDA_VISIBLE_DEVICES: 3
+    MASTER_ADDR: localhost
+    MASTER_PORT: 12315
+    WORLD_SIZE: 1
+    NODE_RANK: 0
+  wandb_config:
+    project: Thermostability
+    name: SaProt_650M_AF2_scnet_1gpu
+
+model:
+#    Which model to use
+  model_py_path: saprot/saprot_regression_model
+  kwargs:
+#    Arguments to initialize the specific class
+    config_path: weight/PLMs/SaProt_650M_AF2
+    load_pretrained: True
+
+#    Arguments to initialize the basic class AbstractModel
+  lr_scheduler_kwargs:
+    last_epoch: -1
+    init_lr: 2.0e-5
+#    Weather to use this scheduler or not
+    on_use: false
+
+  optimizer_kwargs:
+    betas: [0.9, 0.98]
+    weight_decay: 0.01
+
+  save_path: weight/Thermostability/SaProt_650M_AF2_scnet_1gpu.pt
+
+
+dataset:
+#    Arguments to initialize the basic class LMDBDataset
+  dataset_py_path: saprot/saprot_regression_dataset
+  dataloader_kwargs:
+    batch_size: 1
+    num_workers: 4
+
+  train_lmdb: scripts/LMDB/Thermostability/foldseek/train
+  valid_lmdb: scripts/LMDB/Thermostability/foldseek/valid
+  test_lmdb: scripts/LMDB/Thermostability/foldseek/test
+#    Arguments to initialize the specific class
+  kwargs:
+    tokenizer: weight/PLMs/SaProt_650M_AF2
+    mix_max_norm: [40, 67]
+    plddt_threshold: 70
+
+
+#  Arguments to initialize Pytorch Lightning Trainer
+Trainer:
+  max_epochs: 200
+  log_every_n_steps: 1
+  strategy:
+    find_unused_parameters: True
+  logger: False
+  enable_checkpointing: false
+  val_check_interval: 0.5
+  accelerator: gpu
+  devices: 1
+  num_nodes: 1
+  accumulate_grad_batches: 8
+  precision: 16
+  num_sanity_val_steps: 0

config.json

@@ -0,0 +1,102 @@
+{
+    "framework":  "PyTorch",
+    "task":  "protein-language-modeling",
+    "model":  "SaProt",
+    "description":  "Structure-aware protein language modeling with amino-acid and Foldseek 3Di tokens for representation extraction, mutation-effect prediction, inverse folding, pretraining, downstream fine-tuning, and benchmark evaluation.",
+    "subtasks":  [
+                     "structure-aware-sequence-encoding",
+                     "masked-language-model-inference",
+                     "protein-embedding-extraction",
+                     "zero-shot-mutation-effect-prediction",
+                     "protein-inverse-folding",
+                     "saprot-pretraining",
+                     "downstream-finetuning",
+                     "benchmark-evaluation",
+                     "esm2-baseline-evaluation"
+                 ],
+    "entrypoints":  {
+                        "train":  "scripts/training.py",
+                        "pretrain_config":  "conf/pretrain/saprot.yaml",
+                        "thermos_finetune_config":  "conf/Thermostability/saprot.yaml",
+                        "scnet_thermos_finetune_config":  "conf/scnet/Thermostability_saprot_1gpu.yaml",
+                        "mutation_zeroshot":  "scripts/mutation_zeroshot.py",
+                        "proteingym_config":  "conf/ProteinGym/saprot.yaml",
+                        "clinvar_config":  "conf/ClinVar/saprot.yaml",
+                        "scnet_clinvar_config":  "conf/scnet/ClinVar_saprot.yaml",
+                        "clinvar_auc":  "scripts/compute_clinvar_auc.py",
+                        "environment_setup":  "scripts/environment.sh"
+                    },
+    "packages":  [
+                     "model",
+                     "scripts.dataset",
+                     "scripts.utils"
+                 ],
+    "configs":  {
+                    "root":  "conf",
+                    "saprot_tasks":  [
+                                         "conf/Thermostability/saprot.yaml",
+                                         "conf/EC/saprot.yaml",
+                                         "conf/GO/MF/saprot.yaml",
+                                         "conf/GO/BP/saprot.yaml",
+                                         "conf/GO/CC/saprot.yaml",
+                                         "conf/MetalIonBinding/saprot.yaml",
+                                         "conf/HumanPPI/saprot.yaml",
+                                         "conf/Contact/saprot.yaml",
+                                         "conf/DeepLoc/cls2/saprot.yaml",
+                                         "conf/DeepLoc/cls10/saprot.yaml",
+                                         "conf/ProteinGym/saprot.yaml",
+                                         "conf/ClinVar/saprot.yaml"
+                                     ],
+                    "esm2_baselines":  [
+                                           "conf/Thermostability/esm2.yaml",
+                                           "conf/EC/esm2.yaml",
+                                           "conf/GO/MF/esm2.yaml",
+                                           "conf/GO/BP/esm2.yaml",
+                                           "conf/GO/CC/esm2.yaml",
+                                           "conf/MetalIonBinding/esm2.yaml",
+                                           "conf/HumanPPI/esm2.yaml",
+                                           "conf/Contact/esm2.yaml",
+                                           "conf/DeepLoc/cls2/esm2.yaml",
+                                           "conf/DeepLoc/cls10/esm2.yaml",
+                                           "conf/ProteinGym/esm2.yaml",
+                                           "conf/ClinVar/esm2.yaml"
+                                       ],
+                    "scnet":  [
+                                  "conf/scnet/Thermostability_saprot_1gpu.yaml",
+                                  "conf/scnet/ClinVar_saprot.yaml"
+                              ]
+                },
+    "example_data":  [
+                         "scripts/example/8ac8.cif"
+                     ],
+    "data_roots":  {
+                       "downstream_lmdb":  "scripts/LMDB",
+                       "pretrain_lmdb":  "scripts/LMDB/AF2_Uniref50/foldseek",
+                       "proteingym":  "scripts/LMDB/ProteinGym/substitutions",
+                       "clinvar":  "scripts/LMDB/ClinVar",
+                       "thermos":  "scripts/LMDB/Thermostability/foldseek"
+                   },
+    "model_cache":  "weight/PLMs",
+    "default_model":  "weight/PLMs/SaProt_650M_AF2",
+    "optional_models":  {
+                            "esm2_650m":  "weight/PLMs/esm2_t33_650M_UR50D",
+                            "saprot_inverse_folding":  "weight/PLMs/SaProt_650M_AF2_inverse_folding"
+                        },
+    "external_tools":  {
+                           "foldseek":  "scripts/bin/foldseek"
+                       },
+    "outputs":  {
+                    "finetuned_checkpoints":  "weight/{task}/*.pt",
+                    "proteingym_scores":  "output/ProteinGym/SaProt_650M_AF2.tsv",
+                    "clinvar_predictions":  "output/ClinVar/SaProt_650M_AF2",
+                    "scnet_clinvar_predictions":  "output/ClinVar/SaProt_650M_AF2_scnet"
+                },
+    "quick_commands":  {
+                           "setup":  "bash scripts/environment.sh",
+                           "train_thermostability":  "python scripts/training.py -c conf/Thermostability/saprot.yaml",
+                           "train_thermostability_scnet_1gpu":  "python scripts/training.py -c conf/scnet/Thermostability_saprot_1gpu.yaml",
+                           "zeroshot_proteingym":  "python scripts/mutation_zeroshot.py -c conf/ProteinGym/saprot.yaml",
+                           "zeroshot_clinvar":  "python scripts/mutation_zeroshot.py -c conf/ClinVar/saprot.yaml",
+                           "compute_clinvar_auc":  "python scripts/compute_clinvar_auc.py -c conf/ClinVar/saprot.yaml"
+                       }
+}

model/README.md

@@ -0,0 +1,9 @@
+# Introduction
+## abstract_model.py
+The `abstract_model.py` file contains the abstract class for all models. 
+
+## model_interface.py
+The `model_interface.py` file is used to dynamically load a model.
+## saprot
+The `saprot` folder contains various implementations for using SaProt model on downstream tasks. Note that 
+SaProt shares the same architecture with ESM models, so we can use the same code to fine-tune both SaProt and ESM-2.

model/abstract_model.py

@@ -0,0 +1,282 @@
+import torch
+import abc
+import os
+
+import pytorch_lightning as pl
+from scripts.utils.lr_scheduler import Esm2LRScheduler
+from torch import distributed as dist
+
+
+class AbstractModel(pl.LightningModule):
+    def __init__(self,
+                 lr_scheduler_kwargs: dict = None,
+                 optimizer_kwargs: dict = None,
+                 save_path: str = None,
+                 from_checkpoint: str = None,
+                 load_prev_scheduler: bool = False,
+                 save_weights_only: bool = True,):
+        """
+
+        Args:
+            lr_scheduler: Kwargs for lr_scheduler
+            optimizer_kwargs: Kwargs for optimizer_kwargs
+            save_path: Save trained model
+            from_checkpoint: Load model from checkpoint
+            load_prev_scheduler: Whether load previous scheduler from save_path
+            load_strict: Whether load model strictly
+            save_weights_only: Whether save only weights or also optimizer and lr_scheduler
+            
+        """
+        super().__init__()
+        self.initialize_model()
+        
+        self.metrics = {}
+        for stage in ["train", "valid", "test"]:
+            stage_metrics = self.initialize_metrics(stage)
+            # Rigister metrics as attributes
+            for metric_name, metric in stage_metrics.items():
+                setattr(self, metric_name, metric)
+                
+            self.metrics[stage] = stage_metrics
+
+        self.lr_scheduler_kwargs = {"init_lr": 0} if lr_scheduler_kwargs is None else lr_scheduler_kwargs
+        self.optimizer_kwargs = {} if optimizer_kwargs is None else optimizer_kwargs
+        self.init_optimizers()
+
+        self.save_path = save_path
+        self.save_weights_only = save_weights_only
+        
+        self.step = 0
+        self.epoch = 0
+        
+        self.load_prev_scheduler = load_prev_scheduler
+        if from_checkpoint:
+            self.load_checkpoint(from_checkpoint, load_prev_scheduler)
+
+    @abc.abstractmethod
+    def initialize_model(self) -> None:
+        """
+        All model initialization should be done here
+        Note that the whole model must be named as "self.model" for model saving and loading
+        """
+        raise NotImplementedError
+    
+    @abc.abstractmethod
+    def forward(self, *args, **kwargs):
+        """
+        Forward propagation
+        """
+        raise NotImplementedError
+    
+    @abc.abstractmethod
+    def initialize_metrics(self, stage: str) -> dict:
+        """
+        Initialize metrics for each stage
+        Args:
+            stage: "train", "valid" or "test"
+        
+        Returns:
+            A dictionary of metrics for the stage. Keys are metric names and values are metric objects
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def loss_func(self, stage: str, outputs, labels) -> torch.Tensor:
+        """
+
+        Args:
+            stage: "train", "valid" or "test"
+            outputs: model outputs for calculating loss
+            labels: labels for calculating loss
+
+        Returns:
+            loss
+
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def load_weights(model, weights):
+        model_dict = model.state_dict()
+
+        unused_params = []
+        missed_params = list(model_dict.keys())
+
+        for k, v in weights.items():
+            if k in model_dict.keys():
+                model_dict[k] = v
+                missed_params.remove(k)
+
+            else:
+                unused_params.append(k)
+
+        if len(missed_params) > 0:
+            print(f"\033[31mSome weights of {type(model).__name__} were not "
+                  f"initialized from the model checkpoint: {missed_params}\033[0m")
+
+        if len(unused_params) > 0:
+            print(f"\033[31mSome weights of the model checkpoint were not used: {unused_params}\033[0m")
+
+        model.load_state_dict(model_dict)
+    
+    # Add 1 to step after each optimizer step
+    def optimizer_step(
+        self,
+        epoch: int,
+        batch_idx: int,
+        optimizer,
+        optimizer_idx: int = 0,
+        optimizer_closure=None,
+        on_tpu: bool = False,
+        using_native_amp: bool = False,
+        using_lbfgs: bool = False,
+    ) -> None:
+        super().optimizer_step(
+            epoch, batch_idx, optimizer, optimizer_idx, optimizer_closure, on_tpu, using_native_amp, using_lbfgs
+        )
+        self.step += 1
+
+    def on_train_epoch_end(self):
+        self.epoch += 1
+
+    def training_step(self, batch, batch_idx):
+        inputs, labels = batch
+        outputs = self(**inputs)
+        loss = self.loss_func('train', outputs, labels)
+        return loss
+    
+    def validation_step(self, batch, batch_idx):
+        inputs, labels = batch
+        outputs = self(**inputs)
+        return self.loss_func('valid', outputs, labels)
+
+    def test_step(self, batch, batch_idx):
+        inputs, labels = batch
+        outputs = self(**inputs)
+        return self.loss_func('test', outputs, labels)
+
+    def load_checkpoint(self, from_checkpoint, load_prev_scheduler):
+        state_dict = torch.load(from_checkpoint, map_location=self.device)
+        self.load_weights(self.model, state_dict["model"])
+        
+        if load_prev_scheduler:
+            try:
+                self.step = state_dict["global_step"]
+                self.epoch = state_dict["epoch"]
+                self.best_value = state_dict["best_value"]
+                self.optimizer.load_state_dict(state_dict["optimizer"])
+                self.lr_scheduler.load_state_dict(state_dict["lr_scheduler"])
+                print(f"Previous training global step: {self.step}")
+                print(f"Previous training epoch: {self.epoch}")
+                print(f"Previous best value: {self.best_value}")
+                print(f"Previous lr_scheduler: {state_dict['lr_scheduler']}")
+            
+            except Exception as e:
+                print(e)
+                raise KeyError("Wrong in loading previous scheduler, please set load_prev_scheduler=False")
+
+    def save_checkpoint(self, save_info: dict = None) -> None:
+        """
+        Save model to save_path
+        Args:
+            save_info: Other info to save
+        """
+        state_dict = {} if save_info is None else save_info
+        state_dict["model"] = self.model.state_dict()
+
+        if not self.save_weights_only:
+            state_dict["global_step"] = self.step
+            state_dict["epoch"] = self.epoch
+            state_dict["best_value"] = getattr(self, f"best_value", None)
+            state_dict["optimizer"] = self.optimizers().optimizer.state_dict()
+            state_dict["lr_scheduler"] = self.lr_schedulers().state_dict()
+
+        torch.save(state_dict, self.save_path)
+
+    def check_save_condition(self, now_value: float, mode: str, save_info: dict = None) -> None:
+        """
+        Check whether to save model. If save_path is not None and now_value is the best, save model.
+        Args:
+            now_value: Current metric value
+            mode: "min" or "max", meaning whether the lower the better or the higher the better
+            save_info: Other info to save
+        """
+
+        assert mode in ["min", "max"], "mode should be 'min' or 'max'"
+
+        if self.save_path is not None:
+            dir = os.path.dirname(self.save_path)
+            os.makedirs(dir, exist_ok=True)
+            
+            if dist.get_rank() == 0:
+                # save the best checkpoint
+                best_value = getattr(self, f"best_value", None)
+                if best_value:
+                    if mode == "min" and now_value < best_value or mode == "max" and now_value > best_value:
+                        setattr(self, "best_value", now_value)
+                        self.save_checkpoint(save_info)
+
+                else:
+                    setattr(self, "best_value", now_value)
+                    self.save_checkpoint(save_info)
+    
+    def reset_metrics(self, stage) -> None:
+        """
+        Reset metrics for given stage
+        Args:
+            stage: "train", "valid" or "test"
+        """
+        for metric in self.metrics[stage].values():
+            metric.reset()
+    
+    def get_log_dict(self, stage: str) -> dict:
+        """
+        Get log dict for the stage
+        Args:
+            stage: "train", "valid" or "test"
+
+        Returns:
+            A dictionary of metrics for the stage. Keys are metric names and values are metric values
+
+        """
+        return {name: metric.compute() for name, metric in self.metrics[stage].items()}
+    
+    def log_info(self, info: dict) -> None:
+        """
+        Record metrics during training and testing
+        Args:
+            info: dict of metrics
+        """
+        if getattr(self, "logger", None) is not None:
+            info["learning_rate"] = self.lr_scheduler.get_last_lr()[0]
+            info["epoch"] = self.epoch
+            self.logger.log_metrics(info, step=self.step)
+
+    def init_optimizers(self):
+        # No decay for layer norm and bias
+        no_decay = ['LayerNorm.weight', 'bias']
+        
+        if "weight_decay" in self.optimizer_kwargs:
+            weight_decay = self.optimizer_kwargs.pop("weight_decay")
+        else:
+            weight_decay = 0.01
+        
+        optimizer_grouped_parameters = [
+            {'params': [p for n, p in self.model.named_parameters() if not any(nd in n for nd in no_decay)],
+             'weight_decay': weight_decay},
+            {'params': [p for n, p in self.model.named_parameters() if any(nd in n for nd in no_decay)],
+             'weight_decay': 0.0}
+        ]
+        
+        self.optimizer = torch.optim.AdamW(optimizer_grouped_parameters,
+                                           lr=self.lr_scheduler_kwargs['init_lr'],
+                                           **self.optimizer_kwargs)
+
+        self.lr_scheduler = Esm2LRScheduler(self.optimizer, **self.lr_scheduler_kwargs)
+    
+    def configure_optimizers(self):
+        return {"optimizer": self.optimizer,
+                "lr_scheduler": {"scheduler": self.lr_scheduler,
+                                 "interval": "step",
+                                 "frequency": 1}
+                }

model/model_interface.py

@@ -0,0 +1,103 @@
+import yaml
+import glob
+
+
+# register all available models through *_model.py files
+# def construct_model():
+#     model_dir = os.path.dirname(__file__)
+#
+#     # lists all model files
+#     model_list = []
+#     for root, _, names in os.walk(model_dir):
+#         for name in names:
+#             if name.endswith('_model.py'):
+#                 sub_dirs = root.replace(model_dir, '').split(os.sep)
+#                 model_list.append((sub_dirs, name[:-3]))
+#
+#     # load model_config.yaml, controlling which models to be loaded
+#     model_config = yaml.safe_load(open(f"{model_dir}/model_config.yaml", "r"))
+#
+#     if model_config["verbose"]:
+#         print("*" * 30 + f" Loading model " + "*" * 30)
+#
+#     # register models
+#     for sub_dirs, name in model_list:
+#         if name in model_config["models"]:
+#             if len(sub_dirs) > 1:
+#                 cmd = f"from {'.'.join(sub_dirs)} import {name}"
+#             else:
+#                 cmd = f"from . import {name}"
+#
+#             exec(cmd)
+#
+#             if model_config["verbose"]:
+#                 info = f"Loaded model: {name}"
+#                 print(f"\033[32m{info}\033[0m")
+#         else:
+#             if model_config["verbose"]:
+#                 info = f"Skipped model: {name}"
+#                 print(f"\033[31m{info}\033[0m")
+#
+#     if model_config["verbose"]:
+#         print("*" * 75)
+#
+#
+# # register function as a wrapper for all models
+# def register_model(cls):
+#     model_dict[cls.__name__] = cls
+#     return cls
+#
+#
+# model_dict = {}
+# construct_model()
+#
+#
+# class ModelInterface:
+#     @classmethod
+#     def get_available_models(cls):
+#         return model_dict.keys()
+#
+#     @classmethod
+#     def init_model(cls, model: str, **kwargs):
+#         """
+#
+#         Args:
+#            model   : Class name of model you want to use. Must be in model_dict.keys()
+#            **kwargs: Kwargs for model initialization
+#
+#         Returns: Corresponding model
+#
+#         """
+#         assert model in model_dict.keys(), f"class {model} doesn't exist!"
+#         return model_dict[model](**kwargs)
+
+
+########################################################################
+#                             Version 2                                #
+########################################################################
+# register function as a wrapper for all models
+def register_model(cls):
+    global now_cls
+    now_cls = cls
+    return cls
+
+
+now_cls = None
+
+
+class ModelInterface:
+    @classmethod
+    def init_model(cls, model_py_path: str, **kwargs):
+        """
+
+        Args:
+            model_py_path: Py file Path of model you want to use.
+           **kwargs: Kwargs for model initialization
+
+        Returns: Corresponding model
+        """
+        sub_dirs = model_py_path.replace("\\", "/").split("/")
+        cmd = f"from {'.' + '.'.join(sub_dirs[:-1])} import {sub_dirs[-1]}"
+        exec(cmd)
+
+        return now_cls(**kwargs)

model/saprot/base.py

@@ -0,0 +1,210 @@
+import torch
+
+from typing import List, Dict
+# from data.pdb2feature import batch_coords2feature
+from transformers import EsmConfig, EsmTokenizer, EsmForMaskedLM, EsmForSequenceClassification
+# from module.esm.structure_module import (
+#     EsmForMaskedLMWithStructure as EsmForMaskedLM,
+#     EsmForSequenceClassificationWithStructure as EsmForSequenceClassification,
+# )
+from ..abstract_model import AbstractModel
+
+
+class SaprotBaseModel(AbstractModel):
+    """
+    ESM base model. It cannot be used directly but provides model initialization for downstream tasks.
+    """
+
+    def __init__(self,
+                 task: str,
+                 config_path: str,
+                 extra_config: dict = None,
+                 load_pretrained: bool = False,
+                 freeze_backbone: bool = False,
+                 use_lora: bool = False,
+                 lora_config_path: str = None,
+                 **kwargs):
+        """
+        Args:
+            task: Task name. Must be one of ['classification', 'regression', 'lm', 'base']
+
+            config_path: Path to the config file of huggingface esm model
+
+            extra_config: Extra config for the model
+
+            load_pretrained: Whether to load pretrained weights of base model
+
+            freeze_backbone: Whether to freeze the backbone of the model
+
+            use_lora: Whether to use LoRA on downstream tasks
+
+            lora_config_path: Path to the config file of LoRA. If not None, LoRA model is for inference only.
+            Otherwise, LoRA model is for training.
+
+            **kwargs: Other arguments for AbstractModel
+        """
+        assert task in ['classification', 'regression', 'lm', 'base']
+        self.task = task
+        self.config_path = config_path
+        self.extra_config = extra_config
+        self.load_pretrained = load_pretrained
+        self.freeze_backbone = freeze_backbone
+        super().__init__(**kwargs)
+
+        # After all initialization done, lora technique is applied if needed
+        self.use_lora = use_lora
+        if use_lora:
+            self._init_lora(lora_config_path)
+
+    def _init_lora(self, lora_config_path):
+        from peft import (
+            PeftModelForSequenceClassification,
+            get_peft_model,
+            LoraConfig,
+        )
+
+        if lora_config_path:
+            # Note that the model is for inference only
+            self.model = PeftModelForSequenceClassification.from_pretrained(self.model, lora_config_path)
+            self.model.merge_and_unload()
+            print("LoRA model is initialized for inference.")
+
+        else:
+            lora_config = {
+                "task_type": "SEQ_CLS",
+                "target_modules": ["query", "key", "value", "intermediate.dense", "output.dense"],
+                "modules_to_save": ["classifier"],
+                "inference_mode": False,
+                "lora_dropout": 0.1,
+                "lora_alpha": 8,
+            }
+
+            peft_config = LoraConfig(**lora_config)
+            self.model = get_peft_model(self.model, peft_config)
+            # original_module is not needed for training
+            self.model.classifier.original_module = None
+
+            print("LoRA model is initialized for training.")
+            self.model.print_trainable_parameters()
+
+        # After LoRA model is initialized, add trainable parameters to optimizer
+        self.init_optimizers()
+
+    def initialize_model(self):
+        # Initialize tokenizer
+        self.tokenizer = EsmTokenizer.from_pretrained(self.config_path)
+
+        # Initialize different models according to task
+        config = EsmConfig.from_pretrained(self.config_path)
+
+        # Add extra config if needed
+        if self.extra_config is None:
+            self.extra_config = {}
+
+        for k, v in self.extra_config.items():
+            setattr(config, k, v)
+
+        if self.task == 'classification':
+            # Note that self.num_labels should be set in child classes
+            if self.load_pretrained:
+                self.model = EsmForSequenceClassification.from_pretrained(
+                    self.config_path, num_labels=self.num_labels, **self.extra_config)
+
+            else:
+                config.num_labels = self.num_labels
+                self.model = EsmForSequenceClassification(config)
+
+        elif self.task == 'regression':
+            if self.load_pretrained:
+                self.model = EsmForSequenceClassification.from_pretrained(
+                    self.config_path, num_labels=1, **self.extra_config)
+
+            else:
+                config.num_labels = 1
+                self.model = EsmForSequenceClassification(config)
+
+        elif self.task == 'lm':
+            if self.load_pretrained:
+                self.model = EsmForMaskedLM.from_pretrained(self.config_path, **self.extra_config)
+
+            else:
+                self.model = EsmForMaskedLM(config)
+
+        elif self.task == 'base':
+            if self.load_pretrained:
+                self.model = EsmForMaskedLM.from_pretrained(self.config_path, **self.extra_config)
+
+            else:
+                self.model = EsmForMaskedLM(config)
+
+            # Remove lm_head as it is not needed for PPI task
+            self.model.lm_head = None
+
+        # Freeze the backbone of the model
+        if self.freeze_backbone:
+            for param in self.model.esm.parameters():
+                param.requires_grad = False
+
+    def initialize_metrics(self, stage: str) -> dict:
+        return {}
+
+    def get_hidden_states(self, inputs, reduction: str = None) -> list:
+        """
+        Get hidden representations of the model.
+
+        Args:
+            inputs:  A dictionary of inputs. It should contain keys ["input_ids", "attention_mask", "token_type_ids"].
+            reduction: Whether to reduce the hidden states. If None, the hidden states are not reduced. If "mean",
+                        the hidden states are averaged over the sequence length.
+
+        Returns:
+            hidden_states: A list of tensors. Each tensor is of shape [L, D], where L is the sequence length and D is
+                            the hidden dimension.
+        """
+        inputs["output_hidden_states"] = True
+        outputs = self.model.esm(**inputs)
+
+        # Get the index of the first <eos> token
+        input_ids = inputs["input_ids"]
+        eos_id = self.tokenizer.eos_token_id
+        ends = (input_ids == eos_id).int()
+        indices = ends.argmax(dim=-1)
+
+        repr_list = []
+        hidden_states = outputs["hidden_states"][-1]
+        for i, idx in enumerate(indices):
+            if reduction == "mean":
+                repr = hidden_states[i][1:idx].mean(dim=0)
+            else:
+                repr = hidden_states[i][1:idx]
+
+            repr_list.append(repr)
+
+        return repr_list
+
+    # def add_bias_feature(self, inputs, coords: List[Dict]) -> torch.Tensor:
+    #     """
+    #     Add structure information as biases to attention map. This function is used to add structure information
+    #     to the model as Evoformer does.
+    #
+    #     Args:
+    #         inputs: A dictionary of inputs. It should contain keys ["input_ids", "attention_mask", "token_type_ids"].
+    #         coords: Coordinates of backbone atoms. Each element is a dictionary with keys ["N", "CA", "C", "O"].
+    #
+    #     Returns
+    #         pair_feature: A tensor of shape [B, L, L, 407]. Here 407 is the RBF of distance(400) + angle(7).
+    #     """
+    #     inputs["pair_feature"] = batch_coords2feature(coords, self.model.device)
+    #     return inputs
+
+    def save_checkpoint(self, save_info: dict = None) -> None:
+        """
+        Rewrite this function for saving LoRA parameters
+        """
+        if not self.use_lora:
+            return super().save_checkpoint(save_info)
+
+        else:
+            self.model.save_pretrained(self.save_path)
+
+

model/saprot/esm_mutation_model.py

@@ -0,0 +1,288 @@
+import os
+import torch
+import copy
+import json
+import torchmetrics
+import torch.distributed as dist
+
+from scripts.utils.constants import aa_set, aa_list
+from ..model_interface import register_model
+from .base import SaprotBaseModel
+
+
+@register_model
+class EsmMutationModel(SaprotBaseModel):
+    def __init__(self,
+                 use_bias_feature: bool = False,
+                 MSA_log_path: str = None,
+                 log_clinvar: bool = False,
+                 log_dir: str = None,
+                 **kwargs):
+        """
+        Args:
+            use_bias_feature: Whether to use structure information as bias feature
+            
+            MSA_log_path: If not None, the model will load MSA log from this path (following Tranception paper)
+            
+            log_clinvar: If True, the model will log the predicted evolutionary indices for ClinVar variants
+            
+            log_dir: If log_clinvar is True, the model will save the predicted evolutionary indices for ClinVar variants
+            
+            **kwargs: other arguments for SaprotBaseModel
+        """
+        self.use_bias_feature = use_bias_feature
+        self.MSA_log_path = MSA_log_path
+        self.MSA_info_dict = {}
+        if MSA_log_path:
+            with open(MSA_log_path, "r") as r:
+                for line in r:
+                    data = json.loads(line)
+                    data["MSA_log_prior"] = torch.tensor(data["MSA_log_prior"])
+                    self.MSA_info_dict[data["DMS_id"]] = data
+        
+        self.log_clinvar = log_clinvar
+        self.log_dir = log_dir
+        if log_clinvar:
+            self.mut_info_list = []
+                    
+        super().__init__(task="lm", **kwargs)
+        
+    def initialize_metrics(self, stage):
+        return {f"{stage}_spearman": torchmetrics.SpearmanCorrCoef()}
+
+    def forward(self, wild_type, seqs, mut_info, structure_content, structure_type, plddt, struc_seq):
+        if self.use_bias_feature and getattr(self, "coords", None) is None:
+            structure_type = "cif" if structure_type == "mmcif" else structure_type
+            tmp_path = f"EsmMutationModel_{self.global_rank}.{structure_type}"
+            with open(tmp_path, "w") as f:
+                f.write(structure_content)
+            
+            self.coords = parse_structure(tmp_path, ["A"])["A"]["coords"]
+            os.remove(tmp_path)
+            
+        ins_seqs = []
+        ori_seqs = []
+        mut_data = []
+        
+        # The running bottleneck is two forward passes of the model to deal with insertion
+        # Therefore we only forward pass the model twice for sequences with insertion
+        ins_dict = {}
+        
+        for i, (seq, info) in enumerate(zip(seqs, mut_info)):
+            # We adopt the same strategy for esm2 model as in esm2 inverse folding paper
+            ori_seq = [aa for aa in wild_type]
+            ins_seq = copy.deepcopy(ori_seq)
+            tmp_data = []
+            ins_num = 0
+            
+            # To indicate whether there is insertion in the sequence
+            flag = False
+            
+            for single in info.split(":"):
+                # Mask the amino acid where the mutation happens
+                # -1 is added because the index starts from 1 and we need to convert it to 0
+                if single[0] in aa_set:
+                    ori_aa, pos, mut_aa = single[0], int(single[1:-1]), single[-1]
+                    ori_seq[pos - ins_num - 1] = self.tokenizer.mask_token
+                    ins_seq[pos - 1] = self.tokenizer.mask_token
+                    
+                    tmp_data.append((ori_aa, pos - ins_num, mut_aa, pos))
+                
+                # For insertion
+                else:
+                    ins_dict[i] = len(ins_dict)
+                    flag = True
+                    
+                    ins_num += 1
+                    ins_pos = int(single[:-1])
+                    ins_seq = ins_seq[:ins_pos - 1] + [self.tokenizer.mask_token] + ins_seq[ins_pos - 1:]
+            
+            if flag:
+                ins_seqs.append(" ".join(ins_seq))
+                
+            ori_seqs.append(" ".join(ori_seq))
+            mut_data.append(tmp_data)
+        
+        device = self.device
+
+        if len(ins_seqs) > 0:
+            ins_inputs = self.tokenizer.batch_encode_plus(ins_seqs, return_tensors="pt", padding=True)
+            ins_inputs = {k: v.to(device) for k, v in ins_inputs.items()}
+            if self.use_bias_feature:
+                coords = [copy.deepcopy(self.coords) for _ in range(len(seqs))]
+                self.add_bias_feature(ins_inputs, coords)
+                
+            ins_outputs = self.model(**ins_inputs)
+            ins_probs = ins_outputs['logits'].softmax(dim=-1)
+            
+        ori_inputs = self.tokenizer.batch_encode_plus(ori_seqs, return_tensors="pt", padding=True)
+        ori_inputs = {k: v.to(device) for k, v in ori_inputs.items()}
+        if self.use_bias_feature:
+            coords = [copy.deepcopy(self.coords) for _ in range(len(seqs))]
+            self.add_bias_feature(ori_inputs, coords)
+            
+        ori_outputs = self.model(**ori_inputs)
+        ori_probs = ori_outputs['logits'].softmax(dim=-1)
+        
+        if self.MSA_log_path is not None:
+            aa2id = {"A": 5, "C": 6, "D": 7, "E": 8, "F": 9, "G": 10, "H": 11, "I": 12, "K": 13, "L": 14, "M": 15,
+                     "N": 16, "P": 17, "Q": 18, "R": 19, "S": 20, "T": 21, "V": 22, "W": 23, "Y": 24}
+            DMS_id = os.path.basename(self.trainer.datamodule.test_lmdb)
+            MSA_info = self.MSA_info_dict[DMS_id]
+            MSA_log_prior = MSA_info["MSA_log_prior"].to(device)
+            st, ed = MSA_info["MSA_start"], MSA_info["MSA_end"]
+
+        preds = []
+        for i, data_list in enumerate(mut_data):
+            pred = 0
+            for data in data_list:
+                ori_aa, ori_pos, mut_aa, ins_pos = data
+
+                ori_prob = ori_probs[i, ori_pos, self.tokenizer.convert_tokens_to_ids(ori_aa)]
+                
+                if i in ins_dict:
+                    mut_prob = ins_probs[ins_dict[i], ins_pos, self.tokenizer.convert_tokens_to_ids(mut_aa)]
+                else:
+                    mut_prob = ori_probs[i, ins_pos, self.tokenizer.convert_tokens_to_ids(mut_aa)]
+
+                # Add MSA info if available
+                if self.MSA_log_path is not None and st <= ori_pos -1 < ed:
+                    ori_msa_prob = MSA_log_prior[ori_pos - 1 - st, aa2id[ori_aa]]
+                    mut_msa_prob = MSA_log_prior[ori_pos - 1 - st, aa2id[mut_aa]]
+                    pred += 0.4 * torch.log(mut_prob / ori_prob) + 0.6 * (mut_msa_prob - ori_msa_prob)
+                
+                else:
+                    # compute zero-shot score
+                    pred += torch.log(mut_prob / ori_prob)
+
+            preds.append(pred)
+
+        if self.log_clinvar:
+            self.mut_info_list.append((mut_info, -torch.tensor(preds)))
+
+        return torch.tensor(preds).to(ori_probs)
+
+    def loss_func(self, stage, outputs, labels):
+        fitness = labels['labels']
+        self.test_spearman(outputs, fitness)
+    
+    def on_test_epoch_end(self):
+        spearman = self.test_spearman.compute()
+        self.reset_metrics("test")
+        self.log("spearman", spearman)
+        if self.use_bias_feature:
+            self.coords = None
+        
+        if self.log_clinvar:
+            # Get dataset name
+            name = os.path.basename(self.trainer.datamodule.test_lmdb)
+            device_rank = dist.get_rank()
+            log_path = f"{self.log_dir}/{name}_{device_rank}.csv"
+            with open(log_path, "w") as w:
+                w.write("protein_name,mutations,evol_indices\n")
+               
+                for mut_info, preds in self.mut_info_list:
+                    for mut, pred in zip(mut_info, preds):
+                        w.write(f"{name},{mut},{pred}\n")
+                        
+            self.mut_info_list = []
+
+    def predict_mut(self, seq: str, mut_info: str) -> float:
+        """
+        Predict the mutational effect of a given mutation
+        Args:
+            seq: The wild type sequence
+
+            mut_info: The mutation information in the format of "A123B", where A is the original amino acid, 123 is the
+                      position and B is the mutated amino acid. If multiple mutations are provided, they should be
+                      separated by colon, e.g. "A123B:C124D".
+
+        Returns:
+            The predicted mutational effect
+        """
+        tokens = self.tokenizer.tokenize(seq)
+        for single in mut_info.split(":"):
+            pos = int(single[1:-1])
+            tokens[pos - 1] = self.tokenizer.mask_token
+
+        mask_seq = " ".join(tokens)
+        inputs = self.tokenizer(mask_seq, return_tensors="pt")
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            logits = outputs.logits
+            probs = logits.softmax(dim=-1)
+
+            score = 0
+            for single in mut_info.split(":"):
+                ori_aa, pos, mut_aa = single[0], int(single[1:-1]), single[-1]
+                ori_prob = probs[0, pos, self.tokenizer.convert_tokens_to_ids(ori_aa)]
+                mut_prob = probs[0, pos, self.tokenizer.convert_tokens_to_ids(mut_aa)]
+
+                score += torch.log(mut_prob / ori_prob)
+
+        return score
+
+    def predict_pos_mut(self, seq: str, pos: int) -> dict:
+        """
+        Predict the mutational effect of mutations at a given position
+        Args:
+            seq: The wild type sequence
+
+            pos: The position of the mutation
+
+        Returns:
+            The predicted mutational effect
+        """
+        tokens = self.tokenizer.tokenize(seq)
+        ori_aa = tokens[pos - 1][0]
+        tokens[pos - 1] = self.tokenizer.mask_token
+
+        mask_seq = " ".join(tokens)
+        inputs = self.tokenizer(mask_seq, return_tensors="pt")
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            logits = outputs.logits
+            probs = logits.softmax(dim=-1)[0, pos]
+
+            scores = {}
+            ori_prob = probs[self.tokenizer.convert_tokens_to_ids(ori_aa)]
+            for mut_aa in aa_list:
+                mut_prob = probs[self.tokenizer.convert_tokens_to_ids(mut_aa)]
+                score = torch.log(mut_prob / ori_prob).item()
+                scores[f"{ori_aa}{pos}{mut_aa}"] = score
+
+        return scores
+
+    def predict_pos_prob(self, seq: str, pos: int) -> dict:
+        """
+        Predict the probability of all amino acids at a given position
+        Args:
+            seq: The wild type sequence
+
+            pos: The position of the mutation
+
+        Returns:
+            The predicted probability of all amino acids
+        """
+        tokens = self.tokenizer.tokenize(seq)
+        tokens[pos - 1] = self.tokenizer.mask_token
+
+        mask_seq = " ".join(tokens)
+        inputs = self.tokenizer(mask_seq, return_tensors="pt")
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            logits = outputs.logits
+            probs = logits.softmax(dim=-1)[0, pos]
+
+            scores = {}
+            for aa in aa_list:
+                prob = probs[self.tokenizer.convert_tokens_to_ids(aa)]
+                scores[aa] = prob.item()
+
+        return scores

model/saprot/saprot_annotation_model.py

@@ -0,0 +1,89 @@
+import torchmetrics
+import torch
+
+
+from torch.nn.functional import binary_cross_entropy_with_logits
+from scripts.utils.metrics import count_f1_max
+from ..model_interface import register_model
+from .base import SaprotBaseModel
+
+
+@register_model
+class SaprotAnnotationModel(SaprotBaseModel):
+    def __init__(self, anno_type: str, **kwargs):
+        """
+        Args:
+            anno_type: one of EC, GO, GO_MF, GO_CC
+            **kwargs: other parameters for SaprotBaseModel
+        """
+        label2num = {"EC": 585, "GO_BP": 1943, "GO_MF": 489, "GO_CC": 320}
+        self.num_labels = label2num[anno_type]
+        super().__init__(task="classification", **kwargs)
+        
+    def initialize_metrics(self, stage):
+        return {f"{stage}_aupr": torchmetrics.AveragePrecision(pos_label=1, average='micro')}
+
+    def forward(self, inputs, coords=None):
+        if coords is not None:
+            inputs = self.add_bias_feature(inputs, coords)
+
+        # If backbone is frozen, the embedding will be the average of all residues
+        if self.freeze_backbone:
+            repr = torch.stack(self.get_hidden_states(inputs, reduction="mean"))
+            x = self.model.classifier.dropout(repr)
+            x = self.model.classifier.dense(x)
+            x = torch.tanh(x)
+            x = self.model.classifier.dropout(x)
+            logits = self.model.classifier.out_proj(x)
+
+        else:
+            logits = self.model(**inputs).logits
+
+        return logits
+
+    def loss_func(self, stage, logits, labels):
+        label = labels['labels'].to(logits)
+        # add weight to balance positive and negative samples
+        # num_pos = label.sum()
+        # pos_weight = (label.numel() - num_pos) / num_pos
+ 
+        loss = binary_cross_entropy_with_logits(logits, label.float())
+        aupr = getattr(self, f"{stage}_aupr")(logits.sigmoid().detach(), label)
+
+        if stage == "train":
+            log_dict = {"train_loss": loss,
+                        # "train_aupr": aupr
+                        }
+            self.log_info(log_dict)
+            self.reset_metrics("train")
+        
+        return loss
+    
+    def test_epoch_end(self, outputs):
+        preds = self.all_gather(torch.cat(self.test_aupr.preds, dim=-1)).view(-1, self.num_labels)
+        target = self.all_gather(torch.cat(self.test_aupr.target, dim=-1)).long().view(-1, self.num_labels)
+        fmax = count_f1_max(preds, target)
+        
+        log_dict = {"test_f1_max": fmax,
+                    "test_loss": torch.cat(self.all_gather(outputs), dim=-1).mean(),
+                    # "test_aupr": self.test_aupr.compute()
+                    }
+        self.log_info(log_dict)
+        print(log_dict)
+        self.reset_metrics("test")
+
+    def validation_epoch_end(self, outputs):
+        aupr = self.valid_aupr.compute()
+
+        preds = self.all_gather(torch.cat(self.valid_aupr.preds, dim=-1)).view(-1, self.num_labels)
+        target = self.all_gather(torch.cat(self.valid_aupr.target, dim=-1)).long().view(-1, self.num_labels)
+        f1_max = count_f1_max(preds, target)
+        
+        log_dict = {"valid_f1_max": f1_max,
+                    "valid_loss": torch.cat(self.all_gather(outputs), dim=-1).mean(),
+                    # "valid_aupr": aupr
+                    }
+        
+        self.log_info(log_dict)
+        self.reset_metrics("valid")
+        self.check_save_condition(log_dict["valid_f1_max"], mode="max")

model/saprot/saprot_classification_model.py

@@ -0,0 +1,74 @@
+import torchmetrics
+import torch
+
+from torch.nn.functional import cross_entropy
+from ..model_interface import register_model
+from .base import SaprotBaseModel
+
+
+@register_model
+class SaprotClassificationModel(SaprotBaseModel):
+    def __init__(self, num_labels: int, **kwargs):
+        """
+        Args:
+            num_labels: number of labels
+            **kwargs: other arguments for SaprotBaseModel
+        """
+        self.num_labels = num_labels
+        super().__init__(task="classification", **kwargs)
+        
+    def initialize_metrics(self, stage):
+        return {f"{stage}_acc": torchmetrics.Accuracy()}
+
+    def forward(self, inputs, coords=None):
+        if coords is not None:
+            inputs = self.add_bias_feature(inputs, coords)
+
+        # If backbone is frozen, the embedding will be the average of all residues
+        if self.freeze_backbone:
+            repr = torch.stack(self.get_hidden_states(inputs, reduction="mean"))
+            x = self.model.classifier.dropout(repr)
+            x = self.model.classifier.dense(x)
+            x = torch.tanh(x)
+            x = self.model.classifier.dropout(x)
+            logits = self.model.classifier.out_proj(x)
+
+        else:
+            logits = self.model(**inputs).logits
+
+        return logits
+
+    def loss_func(self, stage, logits, labels):
+        label = labels['labels']
+        loss = cross_entropy(logits, label)
+
+        # Update metrics
+        for metric in self.metrics[stage].values():
+            metric.update(logits.detach(), label)
+
+        if stage == "train":
+            log_dict = self.get_log_dict("train")
+            log_dict["train_loss"] = loss
+            self.log_info(log_dict)
+
+            # Reset train metrics
+            self.reset_metrics("train")
+
+        return loss
+
+    def test_epoch_end(self, outputs):
+        log_dict = self.get_log_dict("test")
+        log_dict["test_loss"] = torch.cat(self.all_gather(outputs), dim=-1).mean()
+
+        print(log_dict)
+        self.log_info(log_dict)
+
+        self.reset_metrics("test")
+
+    def validation_epoch_end(self, outputs):
+        log_dict = self.get_log_dict("valid")
+        log_dict["valid_loss"] = torch.cat(self.all_gather(outputs), dim=-1).mean()
+
+        self.log_info(log_dict)
+        self.reset_metrics("valid")
+        self.check_save_condition(log_dict["valid_acc"], mode="max")

model/saprot/saprot_contact_model.py

@@ -0,0 +1,145 @@
+import torchmetrics
+import torch
+import numpy as np
+import math
+
+from torch.nn import Linear, ReLU
+from torch.nn.functional import cross_entropy
+from ..model_interface import register_model
+from .base import SaprotBaseModel
+
+
+@register_model
+class SaprotContactModel(SaprotBaseModel):
+    def __init__(self, **kwargs):
+        """
+        Args:
+            num_labels: number of labels
+            **kwargs: other arguments for SaprotBaseModel
+        """
+        super().__init__(task="base", **kwargs)
+
+    def initialize_model(self):
+        super().initialize_model()
+
+        # hidden_size = self.model.config.hidden_size * 2
+        hidden_size = self.model.config.num_attention_heads
+        # hidden_size = self.model.config.num_hidden_layers * self.model.config.num_attention_heads
+        
+        classifier = torch.nn.Sequential(
+            # Linear(hidden_size, hidden_size),
+            # ReLU(),
+            Linear(hidden_size, 2)
+        )
+        
+        # Freeze all parameters except classifier
+        for param in self.model.parameters():
+            param.requires_grad = False
+
+        setattr(self.model, "classifier", classifier)
+
+    def initialize_metrics(self, stage):
+        metric_dict = {}
+        for length in ["P@L", "P@L/2", "P@L/5"]:
+            for range in ["short_range", "medium_range", "long_range"]:
+                metric_dict[f"{stage}_{range}_{length}"] = torchmetrics.Accuracy(ignore_index=-1)
+
+        return metric_dict
+
+    def forward(self, inputs):
+        # inputs["output_hidden_states"] = True
+        # outputs = self.model.esm(**inputs)
+        #
+        # hidden_states = outputs["hidden_states"][-1]
+        # prod = hidden_states[:, :, None, :] * hidden_states[:, None, :, :]
+        # diff = hidden_states[:, :, None, :] - hidden_states[:, None, :, :]
+        # pairwise_features = torch.cat((prod, diff), -1)
+        # pairwise_features = (pairwise_features + pairwise_features.transpose(1, 2)) / 2
+        #
+        # logits = self.model.classifier(pairwise_features)
+        # logits = logits[:, 1: -1, 1: -1].contiguous()
+        #
+        # return logits
+
+        inputs["output_attentions"] = True
+        outputs = self.model.esm(**inputs)
+
+        attention_maps = torch.cat(outputs["attentions"][-1:], 1).permute(0, 2, 3, 1)
+        # attention_maps = torch.cat(outputs["attentions"], 1).permute(0, 2, 3, 1)
+        attention_maps = (attention_maps + attention_maps.transpose(1, 2)) / 2
+        logits = self.model.classifier(attention_maps)
+        logits = logits[:, 1: -1, 1: -1].contiguous()
+
+        return logits
+
+    def loss_func(self, stage, logits, labels):
+        lengths = labels["lengths"]
+        targets = labels["targets"].to(logits.device)
+        loss = cross_entropy(logits.view(-1, logits.size(-1)), targets.flatten(), ignore_index=-1)
+
+        # Iterate through all proteins and count accuracy
+        length_dict = {"P@L": 1, "P@L/2": 2, "P@L/5": 5}
+        range_dict = ["short_range", "medium_range", "long_range"]
+        for pred_map, label_map, L in zip(logits.detach(), targets, lengths):
+            x_inds, y_inds = np.indices(label_map.shape)
+            for r in range_dict:
+                if r == "short_range":
+                    mask = (np.abs(y_inds - x_inds) < 6) | (np.abs(y_inds - x_inds) > 11)
+
+                elif r == "medium_range":
+                    mask = (np.abs(y_inds - x_inds) < 12) | (np.abs(y_inds - x_inds) > 23)
+
+                else:
+                    mask = np.abs(y_inds - x_inds) < 24
+
+                mask = torch.from_numpy(mask)
+                copy_label_map = label_map.clone()
+                copy_label_map[mask] = -1
+                
+                # Mask the lower triangle
+                mask = torch.triu(torch.ones_like(copy_label_map), diagonal=1)
+                copy_label_map[mask == 0] = -1
+
+                selector = copy_label_map != -1
+                preds = pred_map[selector].float()
+                labels = copy_label_map[selector]
+
+                probs = preds.softmax(dim=-1)[:, 1]
+                for k, v in length_dict.items():
+                    l = min(math.ceil(L / v), (labels == 1).sum().item())
+
+                    top_inds = torch.argsort(probs, descending=True)[:l]
+                    top_labels = labels[top_inds]
+
+                    if top_labels.numel() == 0:
+                        continue
+
+                    metric = f"{stage}_{r}_{k}"
+                    self.metrics[stage][metric].update(top_labels, torch.ones_like(top_labels))
+
+        if stage == "train":
+            # log_dict = self.get_log_dict("train")
+            # log_dict["train_loss"] = loss
+            # self.log_info(log_dict)
+
+            # Reset train metrics
+            self.reset_metrics("train")
+
+        return loss
+
+    def test_epoch_end(self, outputs):
+        log_dict = self.get_log_dict("test")
+        log_dict["test_loss"] = torch.cat(self.all_gather(outputs), dim=-1).mean()
+
+        print(log_dict)
+        self.log_info(log_dict)
+
+        self.reset_metrics("test")
+
+    def validation_epoch_end(self, outputs):
+        log_dict = self.get_log_dict("valid")
+        log_dict["valid_loss"] = torch.cat(self.all_gather(outputs), dim=-1).mean()
+
+        self.log_info(log_dict)
+        self.reset_metrics("valid")
+        self.check_save_condition(log_dict["valid_medium_range_P@L/5"], mode="max")

model/saprot/saprot_foldseek_mutation_model.py

@@ -0,0 +1,383 @@
+import copy
+import os
+import torch
+import json
+import torchmetrics
+import random
+import numpy as np
+
+from scripts.utils.constants import aa_set, foldseek_struc_vocab, aa_list
+from ..model_interface import register_model
+from .base import SaprotBaseModel
+
+
+@register_model
+class SaprotFoldseekMutationModel(SaprotBaseModel):
+    def __init__(self,
+                 foldseek_path: str,
+                 plddt_threshold: float = 0.,
+                 mask_rate: float = None,
+                 substitute_rate: float = None,
+                 MSA_log_path: str = None,
+                 log_clinvar: bool = False,
+                 log_dir: str = None,
+                 **kwargs):
+        """
+        Args:
+            foldseek_path: Foldseek binary file path
+            
+            plddt_threshold: The threshold for plddt to determine whether a structure token should be masked
+
+            mask_rate: If not None, the model will randomly mask structure tokens with this rate
+
+            substitute_rate: If not None, the model will randomly substitute structure tokens with this rate
+            
+            MSA_log_path: If not None, the model will load MSA log from this path (follow Tranception paper)
+            
+            log_clinvar: If True, the model will log the predicted evolutionary indices for ClinVar variants
+            
+            log_dir: If log_clinvar is True, the model will save the predicted evolutionary indices for ClinVar variants
+            
+            **kwargs: Other arguments for SaprotBaseModel
+        """
+        self.foldseek_path = foldseek_path
+        self.plddt_threshold = plddt_threshold
+        self.mask_rate = mask_rate
+        self.substitute_rate = substitute_rate
+        
+        self.MSA_log_path = MSA_log_path
+        self.MSA_info_dict = {}
+        if MSA_log_path:
+            with open(MSA_log_path, "r") as r:
+                for line in r:
+                    data = json.loads(line)
+                    data["MSA_log_prior"] = torch.tensor(data["MSA_log_prior"])
+                    self.MSA_info_dict[data["DMS_id"]] = data
+                    
+        self.log_clinvar = log_clinvar
+        self.log_dir = log_dir
+        if log_clinvar:
+            self.mut_info_list = []
+        
+        super().__init__(task="lm", **kwargs)
+
+    def initialize_metrics(self, stage):
+        return {f"{stage}_spearman": torchmetrics.SpearmanCorrCoef()}
+    
+    def get_struc_seq(self, structure_content, structure_type, plddt):
+        structure_type = "cif" if structure_type == "mmcif" else structure_type
+        
+        # Sample a random rank to avoid file conflict
+        rank = random.randint(0, 1000000)
+        
+        tmp_pdb_path = f"EsmFoldseekMutationModel_{self.global_rank}_{rank}.{structure_type}"
+        tmp_save_path = f"EsmFoldseekMutationModel_{self.global_rank}_{rank}.tsv"
+        
+        # Save structure content to temporary file
+        with open(tmp_pdb_path, "w") as w:
+            w.write(structure_content)
+        
+        # Get foldseek structural sequecne
+        cmd = f"{self.foldseek_path} structureto3didescriptor -v 0 --threads 1 {tmp_pdb_path} {tmp_save_path}"
+        os.system(cmd)
+        
+        with open(tmp_save_path, "r") as r:
+            line = r.readline()
+            struc_seq = line.split("\t")[2]
+        
+        if plddt is not None:
+            plddts = np.array(plddt)
+            
+            # Mask regions with plddt < threshold
+            indices = np.where(plddts < self.plddt_threshold)[0]
+            np_seq = np.array(list(struc_seq))
+            np_seq[indices] = "#"
+            struc_seq = "".join(np_seq)
+
+        if self.mask_rate is not None:
+            # Mask random structure tokens
+            indices = np.random.choice(len(struc_seq), int(len(struc_seq) * self.mask_rate), replace=False)
+            np_seq = np.array(list(struc_seq))
+            np_seq[indices] = "#"
+            struc_seq = "".join(np_seq)
+
+        if self.substitute_rate is not None:
+            # Substitute random structure tokens
+            indices = np.random.choice(len(struc_seq), int(len(struc_seq) * self.substitute_rate), replace=False)
+            np_seq = np.array(list(struc_seq))
+            np_seq[indices] = np.random.choice(list(foldseek_struc_vocab), len(indices))
+            struc_seq = "".join(np_seq)
+
+        os.remove(tmp_pdb_path)
+        os.remove(tmp_save_path)
+        os.remove(tmp_save_path + ".dbtype")
+        return struc_seq
+    
+    def forward(self, wild_type, seqs, mut_info, structure_content, structure_type, plddt):
+        device = self.device
+        
+        if getattr(self, "struc_seq", None) is None:
+            self.struc_seq = self.get_struc_seq(structure_content, structure_type, plddt)
+        
+        ins_seqs = []
+        ori_seqs = []
+        mut_data = []
+        
+        # The running bottleneck is two forward passes of the model to deal with insertion
+        # Therefore we only forward pass the model twice for sequences with insertion
+        ins_dict = {}
+        
+        for i, (seq, info) in enumerate(zip(seqs, mut_info)):
+            # We adopt the same strategy for esm2 model as in esm2 inverse folding paper
+            ori_seq = [a+b.lower() for a, b in zip(wild_type, self.struc_seq)]
+            ins_seq = copy.deepcopy(ori_seq)
+            tmp_data = []
+            ins_num = 0
+            
+            # To indicate whether there is insertion in the sequence
+            flag = False
+            
+            for single in info.split(":"):
+                # Mask the amino acid where the mutation happens
+                # -1 is added because the index starts from 1 and we need to convert it to 0
+                if single[0] in aa_set:
+                    ori_aa, pos, mut_aa = single[0], int(single[1:-1]), single[-1]
+                    
+                    ori_seq[pos-ins_num-1] = "#" + ori_seq[pos-ins_num-1][-1]
+                    ins_seq[pos-1] = "#" + ins_seq[pos-1][-1]
+                    # ori_seq[pos - ins_num - 1] = self.tokenizer.mask_token
+                    # ins_seq[pos - 1] = self.tokenizer.mask_token
+
+                    tmp_data.append((ori_aa, pos-ins_num, mut_aa, pos))
+
+                # For insertion
+                else:
+                    ins_dict[i] = len(ins_dict)
+                    flag = True
+                    
+                    ins_num += 1
+                    ins_pos = int(single[:-1])
+                    ins_seq = ins_seq[:ins_pos-1] + ["##"] + ins_seq[ins_pos-1:]
+                    # ins_seq = ins_seq[:ins_pos-1] + [self.tokenizer.mask_token] + ins_seq[ins_pos-1:]
+
+            if flag:
+                ins_seqs.append(" ".join(ins_seq))
+                
+            ori_seqs.append(" ".join(ori_seq))
+            mut_data.append(tmp_data)
+        
+        if len(ins_seqs) > 0:
+            ins_inputs = self.tokenizer.batch_encode_plus(ins_seqs, return_tensors="pt", padding=True)
+            ins_inputs = {k: v.to(device) for k, v in ins_inputs.items()}
+            ins_outputs = self.model(**ins_inputs)
+            ins_probs = ins_outputs['logits'].softmax(dim=-1)
+        
+        ori_inputs = self.tokenizer.batch_encode_plus(ori_seqs, return_tensors="pt", padding=True)
+        ori_inputs = {k: v.to(device) for k, v in ori_inputs.items()}
+        ori_outputs = self.model(**ori_inputs)
+        ori_probs = ori_outputs['logits'].softmax(dim=-1)
+        
+        if self.MSA_log_path is not None:
+            aa2id = {"A": 5, "C": 6, "D": 7, "E": 8, "F": 9, "G": 10, "H": 11, "I": 12, "K": 13, "L": 14, "M": 15,
+                     "N": 16, "P": 17, "Q": 18, "R": 19, "S": 20, "T": 21, "V": 22, "W": 23, "Y": 24}
+            DMS_id = os.path.basename(self.trainer.datamodule.test_lmdb)
+            MSA_info = self.MSA_info_dict[DMS_id]
+            MSA_log_prior = MSA_info["MSA_log_prior"].to(device)
+            st, ed = MSA_info["MSA_start"], MSA_info["MSA_end"]
+        
+        preds = []
+        for i, data_list in enumerate(mut_data):
+            pred = 0
+            for data in data_list:
+                ori_aa, ori_pos, mut_aa, ins_pos = data
+
+                ori_st = self.tokenizer.get_vocab()[ori_aa + foldseek_struc_vocab[0]]
+                mut_st = self.tokenizer.get_vocab()[mut_aa + foldseek_struc_vocab[0]]
+
+                ori_prob = ori_probs[i, ori_pos, ori_st: ori_st + len(foldseek_struc_vocab)].sum()
+                # ori_sturc_aa = ori_aa + self.struc_seq[ori_pos-1].lower()
+                # ori_struc_id = self.tokenizer.get_vocab()[ori_sturc_aa]
+                # ori_prob = ori_probs[i, ori_pos, ori_struc_id]
+
+                if i in ins_dict:
+                    mut_prob = ins_probs[ins_dict[i], ins_pos, mut_st: mut_st + len(foldseek_struc_vocab)].sum()
+                    # mut_prob = ins_probs[ins_dict[i], ins_pos, mut_st: mut_st + len(foldseek_struc_vocab)].max()
+                else:
+                    mut_prob = ori_probs[i, ori_pos, mut_st: mut_st + len(foldseek_struc_vocab)].sum()
+                    # mut_prob = ori_probs[i, ins_pos, mut_st: mut_st + len(foldseek_struc_vocab)].max()
+                    # mut_struc_aa = mut_aa + self.struc_seq[ori_pos-1].lower()
+                    # mut_struc_id = self.tokenizer.get_vocab()[mut_struc_aa]
+                    # mut_prob = ori_probs[i, ori_pos, mut_struc_id]
+                
+                # print(ori_prob, mut_prob)
+                # struc_logits = []
+                # ori_idx, mut_idx = None, None
+                # for aa in aa_list:
+                #     struct_aa = aa + self.struc_seq[ori_pos - 1].lower()
+                #     aa_logits = ori_outputs['logits'][i, ori_pos, self.tokenizer.get_vocab()[struct_aa]]
+                #     struc_logits.append(aa_logits)
+                #
+                #     if aa == ori_aa:
+                #         ori_idx = len(struc_logits) - 1
+                #
+                #     if aa == mut_aa:
+                #         mut_idx = len(struc_logits) - 1
+                #
+                # struc_probs = torch.softmax(torch.tensor(struc_logits), dim=-1)
+                # ori_prob = struc_probs[ori_idx]
+                # mut_prob = struc_probs[mut_idx]
+                
+                # Add MSA info if available
+                if self.MSA_log_path is not None and st <= ori_pos -1 < ed:
+                    ori_msa_prob = MSA_log_prior[ori_pos - 1 - st, aa2id[ori_aa]]
+                    mut_msa_prob = MSA_log_prior[ori_pos - 1 - st, aa2id[mut_aa]]
+                    pred += 0.4 * torch.log(mut_prob / ori_prob) + 0.6 * (mut_msa_prob - ori_msa_prob)
+                
+                # compute zero-shot score
+                else:
+                    pred += torch.log(mut_prob / ori_prob)
+                    # ori_prob = ori_probs[i, ori_pos, ori_st: ori_st + len(foldseek_struc_vocab)]
+                    # mut_prob = ori_probs[i, ori_pos, mut_st: mut_st + len(foldseek_struc_vocab)]
+                    # pred += torch.log(mut_prob / ori_prob).mean()
+
+            preds.append(pred)
+        
+        if self.log_clinvar:
+            self.mut_info_list.append((mut_info, -torch.tensor(preds)))
+
+        return torch.tensor(preds).to(ori_probs)
+
+    def loss_func(self, stage, outputs, labels):
+        fitness = labels['labels']
+
+        # Update metrics
+        for metric in self.metrics[stage].values():
+            metric.update(outputs.detach().float(), fitness.float())
+
+    def test_epoch_end(self, outputs):
+        spearman = self.test_spearman.compute()
+        self.struc_seq = None
+        self.reset_metrics("test")
+        self.log("spearman", spearman)
+        
+        if self.log_clinvar:
+            # Get dataset name
+            name = os.path.basename(self.trainer.datamodule.test_lmdb)
+            log_path = f"{self.log_dir}/{name}.csv"
+            with open(log_path, "w") as w:
+                w.write("protein_name,mutations,evol_indices\n")
+                
+                for mut_info, preds in self.mut_info_list:
+                    for mut, pred in zip(mut_info, preds):
+                        w.write(f"{name},{mut},{pred}\n")
+            
+            self.mut_info_list = []
+    
+    def predict_mut(self, seq: str, mut_info: str) -> float:
+        """
+               Predict the mutational effect of a given mutation
+               Args:
+                   seq: The wild type sequence
+
+                   mut_info: The mutation information in the format of "A123B", where A is the original amino acid, 123 is the
+                             position and B is the mutated amino acid. If multiple mutations are provided, they should be
+                             separated by colon, e.g. "A123B:C124D".
+
+               Returns:
+                   The predicted mutational effect
+               """
+        tokens = self.tokenizer.tokenize(seq)
+        for single in mut_info.split(":"):
+            pos = int(single[1:-1])
+            tokens[pos - 1] = "#" + tokens[pos - 1][-1]
+        
+        mask_seq = " ".join(tokens)
+        inputs = self.tokenizer(mask_seq, return_tensors="pt")
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+        
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            logits = outputs.logits
+            probs = logits.softmax(dim=-1)
+            
+            score = 0
+            for single in mut_info.split(":"):
+                ori_aa, pos, mut_aa = single[0], int(single[1:-1]), single[-1]
+                ori_st = self.tokenizer.get_vocab()[ori_aa + foldseek_struc_vocab[0]]
+                mut_st = self.tokenizer.get_vocab()[mut_aa + foldseek_struc_vocab[0]]
+                
+                ori_prob = probs[0, pos, ori_st: ori_st + len(foldseek_struc_vocab)].sum()
+                mut_prob = probs[0, pos, mut_st: mut_st + len(foldseek_struc_vocab)].sum()
+                
+                score += torch.log(mut_prob / ori_prob)
+        
+        return score.item()
+    
+    def predict_pos_mut(self, seq: str, pos: int) -> dict:
+        """
+        Predict the mutational effect of mutations at a given position
+        Args:
+            seq: The wild type sequence
+
+            pos: The position of the mutation
+
+        Returns:
+            The predicted mutational effect
+        """
+        tokens = self.tokenizer.tokenize(seq)
+        ori_aa = tokens[pos - 1][0]
+        tokens[pos - 1] = "#" + tokens[pos - 1][-1]
+        
+        mask_seq = " ".join(tokens)
+        inputs = self.tokenizer(mask_seq, return_tensors="pt")
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+        
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            logits = outputs.logits
+            probs = logits.softmax(dim=-1)[0, pos]
+            
+            scores = {}
+            ori_st = self.tokenizer.get_vocab()[ori_aa + foldseek_struc_vocab[0]]
+            for mut_aa in aa_list:
+                mut_st = self.tokenizer.get_vocab()[mut_aa + foldseek_struc_vocab[0]]
+                
+                ori_prob = probs[ori_st: ori_st + len(foldseek_struc_vocab)].sum()
+                mut_prob = probs[mut_st: mut_st + len(foldseek_struc_vocab)].sum()
+                
+                score = torch.log(mut_prob / ori_prob)
+                scores[f"{ori_aa}{pos}{mut_aa}"] = score.item()
+        
+        return scores
+    
+    def predict_pos_prob(self, seq: str, pos: int) -> dict:
+        """
+        Predict the probability of all amino acids at a given position
+        Args:
+            seq: The wild type sequence
+
+            pos: The position of the mutation
+
+        Returns:
+            The predicted probability of all amino acids
+        """
+        tokens = self.tokenizer.tokenize(seq)
+        tokens[pos - 1] = "#" + tokens[pos - 1][-1]
+        
+        mask_seq = " ".join(tokens)
+        inputs = self.tokenizer(mask_seq, return_tensors="pt")
+        inputs = {k: v.to(self.device) for k, v in inputs.items()}
+        
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+            logits = outputs.logits
+            probs = logits.softmax(dim=-1)[0, pos]
+            
+            scores = {}
+            for aa in aa_list:
+                st = self.tokenizer.get_vocab()[aa + foldseek_struc_vocab[0]]
+                prob = probs[st: st + len(foldseek_struc_vocab)].sum()
+                
+                scores[aa] = prob.item()
+        
+        return scores

model/saprot/saprot_if_model.py

@@ -0,0 +1,130 @@
+import torch
+import torch.distributed as dist
+import torchmetrics
+
+from torch.nn.functional import cross_entropy
+
+from ..model_interface import register_model
+from .base import SaprotBaseModel
+
+
+@register_model
+class SaProtIFModel(SaprotBaseModel):
+    """
+    SaProt inverse folding model.
+    """
+    def __init__(self, **kwargs):
+        super().__init__(task='lm', **kwargs)
+    
+    def initialize_metrics(self, stage):
+        return {f"{stage}_acc": torchmetrics.Accuracy(ignore_index=-1)}
+    
+    def forward(self, inputs, coords=None):
+        if coords is not None:
+            inputs = self.add_bias_feature(inputs, coords)
+        
+        outputs = self.model(**inputs)
+        return outputs
+    
+    def loss_func(self, stage, outputs, labels):
+        logits = outputs['logits']
+        # merge the first and second dimension of logits
+        logits = logits.view(-1, logits.size(-1))
+        
+        # flatten labels
+        labels = labels['labels'].flatten().to(logits.device)
+        
+        loss = cross_entropy(logits, labels, ignore_index=-1)
+        getattr(self, f"{stage}_acc").update(logits.detach(), labels)
+        
+        if stage == 'train':
+            log_dict = self.get_log_dict("train")
+            log_dict["train_loss"] = loss
+            self.log_info(log_dict)
+            self.reset_metrics("train")
+        
+        return loss
+    
+    def on_train_epoch_end(self):
+        super().on_train_epoch_end()
+        
+        # Re-sample the subset of the training data
+        if getattr(self.trainer.datamodule.train_dataset, "num_per_cluster", None) is not None:
+            self.trainer.datamodule.train_dataset.sample_subset()
+    
+    def on_test_epoch_end(self):
+        log_dict = self.get_log_dict("test")
+        log_dict["test_loss"] = torch.cat(self.all_gather(self.test_outputs), dim=-1).mean()
+        
+        if dist.get_rank() == 0:
+            print(log_dict)
+        self.log_info(log_dict)
+        
+        self.reset_metrics("test")
+    
+    def on_validation_epoch_end(self):
+        log_dict = self.get_log_dict("valid")
+        log_dict["valid_loss"] = torch.cat(self.all_gather(self.valid_outputs), dim=-1).mean()
+        
+        if dist.get_rank() == 0:
+            print(log_dict)
+        self.log_info(log_dict)
+        self.reset_metrics("valid")
+        
+        valid_acc = log_dict["valid_acc"]
+        self.check_save_condition(valid_acc, mode="max")
+    
+    def predict(self, aa_seq: str, struc_seq: str, method: str = "argmax", num_samples: int = 1) -> str:
+        """
+        Predict all masked amino acids in the sequence.
+        Args:
+            aa_seq: Amino acid sequence (could be all masked or partially masked).
+
+            struc_seq: Foldseek sequence.
+
+            method: Prediction method. It could be either "argmax" or "multinomial". If "argmax", the most probable
+                    amino acid will be selected. If "multinomial", the amino acid will be sampled from the multinomial
+                    distribution.
+            
+            num_samples: Number of predicted sequences. Only works when method is "multinomial".
+
+        Returns:
+            Predicted residue sequence.
+        """
+        assert len(aa_seq) == len(struc_seq), "The length of the amino acid sequence and the foldseek sequence must be the same."
+        assert method in ["argmax", "multinomial"], "The prediction method must be either 'argmax' or 'multinomial'."
+        if method == "argmax":
+            assert num_samples == 1, "The sample number must be 1 when the prediction method is 'argmax'."
+            
+        sa_seq = "".join(f"{aa}{struc}" for aa, struc in zip(aa_seq, struc_seq))
+
+        # Record the index of masked amino acids
+        mask_indices = [i for i, aa in enumerate(aa_seq) if aa == '#']
+
+        with torch.no_grad():
+            inputs = self.tokenizer(sa_seq, return_tensors='pt')
+            inputs = {k: v.to(self.device) for k, v in inputs.items()}
+            
+            outputs = self.model(**inputs)
+            logits = outputs.logits[0, 1: -1]
+            probs = torch.softmax(logits, dim=-1)
+
+            # Predict amino acids
+            if method == "argmax":
+                batch_preds = probs.argmax(dim=-1).unsqueeze(0)
+            else:
+                batch_preds = torch.multinomial(probs, num_samples, replacement=True).permute(1, 0)
+                
+            pred_aa_seqs = []
+            for preds in batch_preds:
+                masked_preds = preds[mask_indices]
+                pred_tokens = self.tokenizer.convert_ids_to_tokens(masked_preds)
+    
+                tokens = list(aa_seq)
+                for i, pred_token in zip(mask_indices, pred_tokens):
+                    tokens[i] = pred_token[0]
+    
+                pred_aa_seq = "".join(tokens)
+                pred_aa_seqs.append(pred_aa_seq)
+                
+            return pred_aa_seqs

model/saprot/saprot_lm_model.py

@@ -0,0 +1,74 @@
+import torch
+import os
+import torchmetrics
+
+from torch.nn.functional import cross_entropy
+from ..model_interface import register_model
+from .base import SaprotBaseModel
+
+
+@register_model
+class SaprotLMModel(SaprotBaseModel):
+    def __init__(self, **kwargs):
+        super().__init__(task='lm', **kwargs)
+    
+    def initialize_metrics(self, stage):
+        return {f"{stage}_acc": torchmetrics.Accuracy(ignore_index=-1)}
+    
+    def forward(self, inputs, coords=None):
+        if coords is not None:
+            inputs = self.add_bias_feature(inputs, coords)
+            
+        outputs = self.model(**inputs)
+        
+        # Get hidden representations
+        if "output_hidden_states" in inputs and inputs["output_hidden_states"]:
+            input_ids = inputs["input_ids"]
+            ends = (input_ids == 2).int()
+            indices = ends.argmax(dim=-1)
+            repr_list = []
+            hidden_states = outputs["hidden_states"][-1]
+            for i, idx in enumerate(indices):
+                repr = hidden_states[i][1:idx].mean(dim=0)
+                repr_list.append(repr)
+            
+            reprs = torch.stack(repr_list, dim=0)
+            outputs["hidden_states"] = reprs
+        
+        return outputs
+    
+    def loss_func(self, stage, outputs, labels):
+        logits = outputs['logits']
+        # merge the first and second dimension of logits
+        logits = logits.view(-1, logits.size(-1))
+        
+        # flatten labels
+        labels = labels['labels'].flatten().to(logits.device)
+        
+        loss = cross_entropy(logits, labels, ignore_index=-1)
+        getattr(self, f"{stage}_acc").update(logits.detach(), labels)
+        
+        if stage == 'train':
+            log_dict = self.get_log_dict("train")
+            log_dict["train_loss"] = loss
+            self.log_info(log_dict)
+            self.reset_metrics("train")
+        
+        return loss
+    
+    def test_epoch_end(self, outputs):
+        log_dict = self.get_log_dict("test")
+        log_dict["test_loss"] = torch.cat(self.all_gather(outputs), dim=-1).mean()
+        
+        print(log_dict)
+        self.log_info(log_dict)
+        
+        self.reset_metrics("test")
+    
+    def validation_epoch_end(self, outputs):
+        log_dict = self.get_log_dict("valid")
+        log_dict["valid_loss"] = torch.cat(self.all_gather(outputs), dim=-1).mean()
+        
+        self.log_info(log_dict)
+        self.reset_metrics("valid")
+        self.check_save_condition(log_dict["valid_loss"], mode="min")

model/saprot/saprot_ppi_model.py

@@ -0,0 +1,78 @@
+import torchmetrics
+import torch
+
+from torch.nn import Linear, ReLU
+from torch.nn.functional import cross_entropy
+from ..model_interface import register_model
+from .base import SaprotBaseModel
+
+
+@register_model
+class SaprotPPIModel(SaprotBaseModel):
+    def __init__(self, **kwargs):
+        """
+        Args:
+            **kwargs: other arguments for SaprotBaseModel
+        """
+        super().__init__(task="base", **kwargs)
+
+    def initialize_model(self):
+        super().initialize_model()
+        
+        hidden_size = self.model.config.hidden_size * 2
+        classifier = torch.nn.Sequential(
+                        Linear(hidden_size, hidden_size),
+                        ReLU(),
+                        Linear(hidden_size, 2)
+                    )
+        
+        setattr(self.model, "classifier", classifier)
+
+    def initialize_metrics(self, stage):
+        return {f"{stage}_acc": torchmetrics.Accuracy()}
+
+    def forward(self, inputs_1, inputs_2):
+        if self.freeze_backbone:
+            hidden_1 = torch.stack(self.get_hidden_states(inputs_1, reduction="mean"))
+            hidden_2 = torch.stack(self.get_hidden_states(inputs_2, reduction="mean"))
+        else:
+            hidden_1 = self.model.esm(**inputs_1)[0][:, 0, :]
+            hidden_2 = self.model.esm(**inputs_2)[0][:, 0, :]
+
+        hidden_concat = torch.cat([hidden_1, hidden_2], dim=-1)
+        return self.model.classifier(hidden_concat)
+    
+    def loss_func(self, stage, logits, labels):
+        label = labels['labels']
+        loss = cross_entropy(logits, label)
+
+        # Update metrics
+        for metric in self.metrics[stage].values():
+            metric.update(logits.detach(), label)
+
+        if stage == "train":
+            log_dict = self.get_log_dict("train")
+            log_dict["train_loss"] = loss
+            self.log_info(log_dict)
+
+            # Reset train metrics
+            self.reset_metrics("train")
+
+        return loss
+
+    def test_epoch_end(self, outputs):
+        log_dict = self.get_log_dict("test")
+        log_dict["test_loss"] = torch.cat(self.all_gather(outputs), dim=-1).mean()
+
+        print(log_dict)
+        self.log_info(log_dict)
+
+        self.reset_metrics("test")
+
+    def validation_epoch_end(self, outputs):
+        log_dict = self.get_log_dict("valid")
+        log_dict["valid_loss"] = torch.cat(self.all_gather(outputs), dim=-1).mean()
+
+        self.log_info(log_dict)
+        self.reset_metrics("valid")
+        self.check_save_condition(log_dict["valid_acc"], mode="max")

model/saprot/saprot_regression_model.py

@@ -0,0 +1,95 @@
+import torch.distributed as dist
+import torchmetrics
+import torch
+
+from ..model_interface import register_model
+from .base import SaprotBaseModel
+
+
+@register_model
+class SaprotRegressionModel(SaprotBaseModel):
+    def __init__(self, test_result_path: str = None, **kwargs):
+        """
+        Args:
+            test_result_path: path to save test result
+            **kwargs: other arguments for SaprotBaseModel
+        """
+        self.test_result_path = test_result_path
+        super().__init__(task="regression", **kwargs)
+    
+    def initialize_metrics(self, stage):
+        return {f"{stage}_loss": torchmetrics.MeanSquaredError(),
+                f"{stage}_spearman": torchmetrics.SpearmanCorrCoef(),
+                f"{stage}_R2": torchmetrics.R2Score(),
+                f"{stage}_pearson": torchmetrics.PearsonCorrCoef()}
+    
+    def forward(self, inputs, structure_info=None):
+        if structure_info:
+            # To be implemented
+            raise NotImplementedError
+
+        # If backbone is frozen, the embedding will be the average of all residues
+        if self.freeze_backbone:
+            repr = torch.stack(self.get_hidden_states(inputs, reduction="mean"))
+            x = self.model.classifier.dropout(repr)
+            x = self.model.classifier.dense(x)
+            x = torch.tanh(x)
+            x = self.model.classifier.dropout(x)
+            logits = self.model.classifier.out_proj(x).squeeze(dim=-1)
+
+        else:
+            logits = self.model(**inputs).logits.squeeze(dim=-1)
+
+        return logits
+
+    def loss_func(self, stage, outputs, labels):
+        fitness = labels['labels'].to(outputs)
+        loss = torch.nn.functional.mse_loss(outputs, fitness)
+        
+        # Update metrics
+        for metric in self.metrics[stage].values():
+            # Training is on half precision, but metrics expect float to compute correctly.
+            metric.update(outputs.detach().float(), fitness.float())
+        
+        if stage == "train":
+            # Skip calculating metrics if the batch size is 1
+            if fitness.shape[0] > 1:
+                log_dict = self.get_log_dict("train")
+                self.log_info(log_dict)
+            
+            # Reset train metrics
+            self.reset_metrics("train")
+        
+        return loss
+
+    def test_epoch_end(self, outputs):
+        if self.test_result_path is not None:
+            from torchmetrics.utilities.distributed import gather_all_tensors
+            
+            preds = self.test_spearman.preds
+            preds[-1] = preds[-1].unsqueeze(dim=0) if preds[-1].shape == () else preds[-1]
+            preds = torch.cat(gather_all_tensors(torch.cat(preds, dim=0)))
+            
+            targets = self.test_spearman.target
+            targets[-1] = targets[-1].unsqueeze(dim=0) if targets[-1].shape == () else targets[-1]
+            targets = torch.cat(gather_all_tensors(torch.cat(targets, dim=0)))
+
+            if dist.get_rank() == 0:
+                with open(self.test_result_path, 'w') as w:
+                    w.write("pred\ttarget\n")
+                    for pred, target in zip(preds, targets):
+                        w.write(f"{pred.item()}\t{target.item()}\n")
+        
+        log_dict = self.get_log_dict("test")
+        
+        print(log_dict)
+        self.log_info(log_dict)
+        self.reset_metrics("test")
+
+    def validation_epoch_end(self, outputs):
+        log_dict = self.get_log_dict("valid")
+
+        self.log_info(log_dict)
+        self.reset_metrics("valid")
+        self.check_save_condition(log_dict["valid_loss"], mode="min")
+        

requirements.txt

@@ -0,0 +1,16 @@
+torch==1.13.1
+torchvision==0.14.1
+torchaudio==0.13.1
+pytorch-lightning==1.8.3
+wandb==0.12.10
+transformers==4.28.0
+easydict==1.10
+biopython==1.81
+scipy==1.14.1
+peft==0.10.0
+lmdb==1.4.1
+torchmetrics==0.9.3
+pandas==2.1.1
+fair-esm==2.0.0
+numpy==1.25.2
+protobuf<4,>=3.20.0

scripts/LMDB/README.md

@@ -0,0 +1 @@
+# Place the datasets here

scripts/__init__.py

@@ -0,0 +1 @@
+

scripts/bin/README.md

@@ -0,0 +1 @@
+# Place the foldseek binary file here

scripts/compute_clinvar_auc.py

@@ -0,0 +1,57 @@
+import os
+import argparse
+import pandas as pd
+import yaml
+
+from glob import glob
+from sklearn import metrics
+from easydict import EasyDict
+
+
+def get_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-c', '--config', help="running configurations", type=str, required=True)
+    return parser.parse_args()
+
+
+def main(args):
+    with open(args.config, 'r', encoding='utf-8') as r:
+        config = EasyDict(yaml.safe_load(r))
+
+    output_dir = "output/ClinVar"
+    # file_name = "esm2_t33_650M_UR50D_foldseek_plddt70_iter2900448_mask70"
+
+    label_name = "ClinVar_labels.csv"
+    label_path = os.path.join(output_dir, label_name)
+
+    list_variables_to_keep = ["protein_name", "mutations", "evol_indices"]
+    all_evol_indices = pd.concat(
+        [
+            pd.read_csv(path, low_memory=False)[list_variables_to_keep]
+            for path in glob(f"{config.model.kwargs.log_dir}/*.csv")
+        ],
+        ignore_index=True,
+    )
+    all_evol_indices = all_evol_indices.drop_duplicates()
+
+    labels_dataset = pd.read_csv(label_path, low_memory=False)
+    all_evol_indices_with_labels = pd.merge(
+        all_evol_indices,
+        labels_dataset[["protein_name", "mutations", "ClinVar_labels"]],
+        how="right",
+        on=["protein_name", "mutations"],
+    )
+
+    all_evol_indices_with_labels = all_evol_indices_with_labels[
+        all_evol_indices_with_labels.ClinVar_labels != 0.5
+        ]
+
+    fpr, tpr, threshold = metrics.roc_curve(
+        all_evol_indices_with_labels["ClinVar_labels"], all_evol_indices_with_labels["evol_indices"]
+    )
+    roc_auc = metrics.auc(fpr, tpr)
+    print(roc_auc)
+
+
+if __name__ == '__main__':
+    main(get_args())