initial commit

.gitattributes

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

.gitignore

@@ -0,0 +1,3 @@
+models/weights/
+Flexpert-Design/
+data/atlas/

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "gradio_molecule3d"]
+	path = gradio_molecule3d
+	url = https://github.com/Honzus/gradio_molecule3d

.gradio/certificate.pem

@@ -0,0 +1,31 @@
+-----BEGIN CERTIFICATE-----
+MIIFazCCA1OgAwIBAgIRAIIQz7DSQONZRGPgu2OCiwAwDQYJKoZIhvcNAQELBQAw
+TzELMAkGA1UEBhMCVVMxKTAnBgNVBAoTIEludGVybmV0IFNlY3VyaXR5IFJlc2Vh
+cmNoIEdyb3VwMRUwEwYDVQQDEwxJU1JHIFJvb3QgWDEwHhcNMTUwNjA0MTEwNDM4
+WhcNMzUwNjA0MTEwNDM4WjBPMQswCQYDVQQGEwJVUzEpMCcGA1UEChMgSW50ZXJu
+ZXQgU2VjdXJpdHkgUmVzZWFyY2ggR3JvdXAxFTATBgNVBAMTDElTUkcgUm9vdCBY
+MTCCAiIwDQYJKoZIhvcNAQEBBQADggIPADCCAgoCggIBAK3oJHP0FDfzm54rVygc
+h77ct984kIxuPOZXoHj3dcKi/vVqbvYATyjb3miGbESTtrFj/RQSa78f0uoxmyF+
+0TM8ukj13Xnfs7j/EvEhmkvBioZxaUpmZmyPfjxwv60pIgbz5MDmgK7iS4+3mX6U
+A5/TR5d8mUgjU+g4rk8Kb4Mu0UlXjIB0ttov0DiNewNwIRt18jA8+o+u3dpjq+sW
+T8KOEUt+zwvo/7V3LvSye0rgTBIlDHCNAymg4VMk7BPZ7hm/ELNKjD+Jo2FR3qyH
+B5T0Y3HsLuJvW5iB4YlcNHlsdu87kGJ55tukmi8mxdAQ4Q7e2RCOFvu396j3x+UC
+B5iPNgiV5+I3lg02dZ77DnKxHZu8A/lJBdiB3QW0KtZB6awBdpUKD9jf1b0SHzUv
+KBds0pjBqAlkd25HN7rOrFleaJ1/ctaJxQZBKT5ZPt0m9STJEadao0xAH0ahmbWn
+OlFuhjuefXKnEgV4We0+UXgVCwOPjdAvBbI+e0ocS3MFEvzG6uBQE3xDk3SzynTn
+jh8BCNAw1FtxNrQHusEwMFxIt4I7mKZ9YIqioymCzLq9gwQbooMDQaHWBfEbwrbw
+qHyGO0aoSCqI3Haadr8faqU9GY/rOPNk3sgrDQoo//fb4hVC1CLQJ13hef4Y53CI
+rU7m2Ys6xt0nUW7/vGT1M0NPAgMBAAGjQjBAMA4GA1UdDwEB/wQEAwIBBjAPBgNV
+HRMBAf8EBTADAQH/MB0GA1UdDgQWBBR5tFnme7bl5AFzgAiIyBpY9umbbjANBgkq
+hkiG9w0BAQsFAAOCAgEAVR9YqbyyqFDQDLHYGmkgJykIrGF1XIpu+ILlaS/V9lZL
+ubhzEFnTIZd+50xx+7LSYK05qAvqFyFWhfFQDlnrzuBZ6brJFe+GnY+EgPbk6ZGQ
+3BebYhtF8GaV0nxvwuo77x/Py9auJ/GpsMiu/X1+mvoiBOv/2X/qkSsisRcOj/KK
+NFtY2PwByVS5uCbMiogziUwthDyC3+6WVwW6LLv3xLfHTjuCvjHIInNzktHCgKQ5
+ORAzI4JMPJ+GslWYHb4phowim57iaztXOoJwTdwJx4nLCgdNbOhdjsnvzqvHu7Ur
+TkXWStAmzOVyyghqpZXjFaH3pO3JLF+l+/+sKAIuvtd7u+Nxe5AW0wdeRlN8NwdC
+jNPElpzVmbUq4JUagEiuTDkHzsxHpFKVK7q4+63SM1N95R1NbdWhscdCb+ZAJzVc
+oyi3B43njTOQ5yOf+1CceWxG1bQVs5ZufpsMljq4Ui0/1lvh+wjChP4kqKOJ2qxq
+4RgqsahDYVvTH9w7jXbyLeiNdd8XM2w9U/t7y0Ff/9yi0GE44Za4rF2LN9d11TPA
+mRGunUHBcnWEvgJBQl9nJEiU0Zsnvgc/ubhPgXRR4Xq37Z0j4r7g1SgEEzwxA57d
+emyPxgcYxn/eR44/KJ4EBs+lVDR3veyJm+kXQ99b21/+jh5Xos1AnX5iItreGCc=
+-----END CERTIFICATE-----

Flexpert-Design/README.md

@@ -0,0 +1,69 @@
+# Flexpert-Design
+
+In this directory we provide the code to train and run inference with Flexpert-Design. To expedite the release of the codebase, this part of code was not thoroughly curated and contains redundant files and code. The codebase might be revised in the future but probably it will get completely rewritten as part of a future project with an improved model.
+
+## Environment
+
+Tested for Python 3.9. For other versions enviroment might need to be adapted.
+
+Assuming you have already installed the environment for Flexpert-3D and Flexpert-Seq, install the additional dependencies for Flexpert-Design using the `requirements.txt` file in this directory.
+
+```bash
+pip install -r requirements.txt
+```
+
+
+## Inference
+
+In this example we will illustrate how to run inference with the trained model (trained wights are provided inside the train/results directory, you do not need to train the model again necessarily).
+
+Place the PDB files you want to predict in the `predict_example` directory. It is expected that the files are named like `PDBCODE_CHAINID.pdb`, example file '1ahy_A.pdb' is provided. For each PDB file in that folder, add the instructions on flexibility you want to be considered by the ProteinMPNN model in the `PDBCODE_CHAINID_instructions.csv` file - example file '1ah7_A_instructions.csv' is provided. Then run the following command to run inference.
+
+```bash
+python3 predict.py \
+    --infer_path predict_example/
+```
+
+The output will be in the `predict_example/predictions.txt` file.
+
+The origininal sequence and the regenerated sequence can be compared using the following script.
+
+```bash
+python3 predict_example/compare_seqs.py \
+    --pdb_code 1ah7_A
+```
+
+## Training
+
+First make sure you have the Flexpert-3D model weights in the `Flexpert/models/weights` directory. Alternatively run the following script to download the weights.
+
+```bash
+. ../download_flexpert_weights.sh
+```
+
+Download the training data:
+
+```bash
+. ../download-cath-data.sh
+```
+
+Then run the following command to train the model.
+
+```bash
+export HF_HOME=./HF_cache
+python3 train.py \
+    --batch_size 4 \
+    --model_name 'ProteinMPNN' \
+    --stage 'fit' \
+    --dataset FLEX_CATH4.3 \
+    --ex_name training-reproduction \
+    --offline 0 \
+    --gpus 1 \
+    --epoch 11 \
+    --use_dynamics 1 \
+    --flex_loss_coeff 0.8 \
+    --init_flex_features 1 \
+    --grad_normalization 0 \
+    --loss_fn MSE \
+    --use_pmpnn_checkpoint 1
+```

Flexpert-Design/configs/ANMAwareFlexibilityProtTrans.yaml

@@ -0,0 +1,12 @@
+checkpoint_path: ../models/weights/flexpert_3d_weights.bin
+data_jsonl_name: /NEWcath_ANM_gt_flex_annotated.jsonl #/debug_cath_ANM_gt_flex_annotated.jsonl #/cath_ANM_gt_flex_annotated.jsonl
+half_precision: False #mixed_precision
+gumbel_temperature: 0.2
+num_labels: 1
+add_pearson_loss: False
+add_sse_loss: False
+adaptor_architecture: 'conv'
+enm_embed_dim: 128
+enm_att_heads: 8
+num_layers: 3
+kernel_size: 5

Flexpert-Design/configs/Flexpert-Design-inference.yaml

@@ -0,0 +1 @@
+pmpnn_model_path: 'train/results/MSEloss_flex_cath_coeff_0.8/checkpoints/last.ckpt'

Flexpert-Design/configs/ProteinMPNN.py

@@ -0,0 +1,14 @@
+method = 'ProteinMPNN'
+hidden = 128
+k_neighbors=30
+num_letters = 33
+num_encoder_layers = 3
+num_decoder_layers = 3
+vocab = 33
+dropout = 0.1
+smoothing = 0.1
+batch_size = 8
+lr = 0.001
+proteinmpnn_type = 0
+patience = 100
+epoch=100

Flexpert-Design/configs/ProteinMPNN.yaml

@@ -0,0 +1,13 @@
+augment_eps: 0.0
+num_encoder_layers: 3
+hidden_dim: 128
+hidden: 128
+k_neighbors: 30
+num_letters: 33
+num_decoder_layers: 3
+vocab: 33
+dropout: 0.1
+smoothing: 0.1
+proteinmpnn_type: 0
+init_flex_features: 1
+starting_checkpoint_path: 'vanilla_mpnn_weights/best-epoch=99-recovery=0.485.ckpt'

Flexpert-Design/data_interface.py

@@ -0,0 +1,205 @@
+import inspect
+from torch.utils.data import DataLoader
+from src.interface.data_interface import DInterface_base
+import torch
+import os.path as osp
+from src.tools.utils import cuda
+import pdb
+from src.tools.utils import load_yaml_config
+
+class MyDataLoader(DataLoader):
+    def __init__(self, dataset, model_name, batch_size=64, num_workers=8, *args, **kwargs):
+        super().__init__(dataset, batch_size=batch_size, num_workers=num_workers, *args, **kwargs)
+        self.pretrain_device = 'cuda:0'
+        self.model_name = model_name
+    
+    def __iter__(self):
+        for batch in super().__iter__():
+            # 在这里对batch进行处理
+            # ...
+            try:
+                self.pretrain_device = f'cuda:{torch.distributed.get_rank()}'
+            except:
+                self.pretrain_device = 'cuda:0'
+
+            stream = torch.cuda.Stream(
+                self.pretrain_device
+            )
+            with torch.cuda.stream(stream):
+                if self.model_name=='GVP':
+                    batch = batch.cuda(non_blocking=True, device=self.pretrain_device)
+                    yield batch
+                else:
+                    for key, val in batch.items():
+                        if type(val) == torch.Tensor:
+                            batch[key] = batch[key].cuda(non_blocking=True, device=self.pretrain_device)
+
+                    # X = batch['X'].cuda(non_blocking=True, device=self.pretrain_device)
+                    # S = batch['S'].cuda(non_blocking=True, device=self.pretrain_device)
+                    # score = batch['score'].cuda(non_blocking=True, device=self.pretrain_device)
+                    # mask = batch['mask'].cuda(non_blocking=True, device=self.pretrain_device)
+                    # lengths = batch['lengths'].cuda(non_blocking=True, device=self.pretrain_device)
+                    # chain_mask = batch['chain_mask'].cuda(non_blocking=True, device=self.pretrain_device)
+                    # chain_encoding = batch['chain_encoding'].cuda(non_blocking=True, device=self.pretrain_device)
+                
+                    yield batch
+
+
+class DInterface(DInterface_base):
+    def __init__(self,**kwargs):
+        super().__init__(**kwargs)
+        self.save_hyperparameters()
+        self.load_data_module()
+
+    def setup(self, stage=None):
+        from src.datasets.featurizer import (featurize_AF, featurize_GTrans, featurize_GVP,
+                         featurize_ProteinMPNN, featurize_Inversefolding)
+        if self.hparams.model_name in ['AlphaDesign', 'PiFold', 'KWDesign', 'GraphTrans', 'StructGNN', 'GCA', 'E3PiFold']:
+            self.collate_fn = featurize_GTrans
+        elif self.hparams.model_name == 'GVP':
+            featurizer = featurize_GVP()
+            self.collate_fn = featurizer.collate
+        elif self.hparams.model_name == 'ProteinMPNN':
+            self.collate_fn = featurize_ProteinMPNN
+        elif self.hparams.model_name == 'ESMIF':
+            self.collate_fn = featurize_Inversefolding
+            
+        # Assign train/val datasets for use in dataloaders
+        if stage == 'fit' or stage is None:
+            self.trainset = self.instancialize(split = 'train')
+            self.valset = self.instancialize(split='valid')
+
+        # Assign test dataset for use in dataloader(s)
+        if stage == 'test' or stage is None:
+            self.testset = self.instancialize(split='test')
+        
+        if stage in ['predict','eval']:
+            self.predictset = self.instancialize(split='predict')
+
+    def train_dataloader(self):
+        return MyDataLoader(self.trainset, model_name=self.hparams.model_name, batch_size=self.batch_size, num_workers=self.hparams.num_workers, shuffle=True, prefetch_factor=8, pin_memory=True, collate_fn=self.collate_fn)
+
+    def val_dataloader(self):
+        return MyDataLoader(self.valset, model_name=self.hparams.model_name, batch_size=self.batch_size, num_workers=self.hparams.num_workers, shuffle=False, pin_memory=True, collate_fn=self.collate_fn)
+
+    def test_dataloader(self):
+        return MyDataLoader(self.testset, model_name=self.hparams.model_name, batch_size=self.batch_size, num_workers=self.hparams.num_workers, shuffle=False, pin_memory=True, collate_fn=self.collate_fn)
+
+    def predict_dataloader(self):
+        return MyDataLoader(self.predictset, model_name=self.hparams.model_name, batch_size=self.batch_size, num_workers=self.hparams.num_workers, shuffle=False, pin_memory=True, collate_fn=self.collate_fn)
+
+    def load_data_module(self):
+        
+        name = self.hparams.dataset
+        if name == 'AF2DB':
+            from src.datasets.AF2DB_dataset_lmdb import Af2dbDataset
+            self.data_module = Af2dbDataset
+        
+        if name == 'TS':
+            from src.datasets.ts_dataset  import TSDataset
+            self.data_module = TSDataset
+            self.hparams['path'] = osp.join(self.hparams.data_root, 'ts')
+        
+        if name == 'CASP15':
+            from src.datasets.casp_dataset  import CASPDataset
+            self.data_module = CASPDataset
+            self.hparams['path'] = osp.join(self.hparams.data_root, 'casp15')
+        
+        if name == 'CATH4.2':
+            from src.datasets.cath_dataset import CATHDataset
+            self.data_module = CATHDataset
+            self.hparams['version'] = 4.2
+            self.hparams['path'] = osp.join(self.hparams.data_root, 'cath4.2')
+            
+        if name == 'CATH4.3':
+            from src.datasets.cath_dataset import CATHDataset
+            self.data_module = CATHDataset
+            self.hparams['version'] = 4.3
+            self.hparams['path'] = osp.join(self.hparams.data_root, 'cath4.3')
+        
+        if name == 'MPNN':
+            from src.datasets.mpnn_dataset import MPNNDataset
+            self.data_module = MPNNDataset
+
+        if name == 'FOLDSWITCHERS_1':
+            from src.datasets.foldswitchers_dataset import FoldswitchersDataset
+            self.data_module = FoldswitchersDataset
+            self.hparams['path'] = osp.join(self.hparams.data_root, 'fold_switchers/fold_1')
+
+        if name == 'FOLDSWITCHERS_2':
+            from src.datasets.foldswitchers_dataset import FoldswitchersDataset
+            self.data_module = FoldswitchersDataset
+            self.hparams['path'] = osp.join(self.hparams.data_root, 'fold_switchers/fold_2')
+
+        if name == 'PDBInference':
+            from src.datasets.pdb_inference import PDBInference
+            self.data_module = PDBInference
+            self.hparams['path'] = osp.join(self.hparams.infer_path)
+
+        if name == 'ATLAS_DIST_1':
+            from src.datasets.atlas_dataset import AtlasDataset
+            self.data_module = AtlasDataset
+            self.hparams['path'] = osp.join(self.hparams.data_root, 'atlas/distant-frame-pairs_NO_SUPERPOSITION/frames_1')
+        
+        if name == 'ATLAS_DIST_2':
+            from src.datasets.atlas_dataset import AtlasDataset
+            self.data_module = AtlasDataset
+            self.hparams['path'] = osp.join(self.hparams.data_root, 'atlas/distant-frame-pairs_NO_SUPERPOSITION/frames_2')
+        
+        if name == 'ATLAS_CLUSTER_1':
+            from src.datasets.atlas_dataset import AtlasDataset
+            self.data_module = AtlasDataset
+            self.hparams['path'] = osp.join(self.hparams.data_root, 'atlas/cluster-representatives/frames_1')
+
+        if name == 'ATLAS_CLUSTER_2':
+            from src.datasets.atlas_dataset import AtlasDataset
+            self.data_module = AtlasDataset
+            self.hparams['path'] = osp.join(self.hparams.data_root, 'atlas/cluster-representatives/frames_2')
+
+        if name == 'ATLAS_PDB':
+            from src.datasets.atlas_dataset import AtlasDataset
+            self.data_module = AtlasDataset
+            self.hparams['path'] = osp.join(self.hparams.data_root, '../atlas_pdb_inference/')
+
+        if name == 'ATLAS_FULL_MINIMIZED':
+            from src.datasets.atlas_dataset import AtlasDataset
+            self.data_module = AtlasDataset
+            self.hparams['path'] = osp.join(self.hparams.data_root, '../atlas_eval_proteinmpnn/atlas_full/minimized_PDBs/pdbs/')
+        
+        if name == 'ATLAS_FULL_REFOLDED':
+            from src.datasets.atlas_dataset import AtlasDataset
+            self.data_module = AtlasDataset
+            self.hparams['path'] = osp.join(self.hparams.data_root, '../atlas_eval_proteinmpnn/atlas_full/refolded_PDBs/pdbs/')
+        
+        if name == 'ATLAS_FULL_CRYSTAL':
+            from src.datasets.atlas_dataset import AtlasDataset
+            self.data_module = AtlasDataset
+            self.hparams['path'] = osp.join(self.hparams.data_root, '../atlas_eval_proteinmpnn/atlas_full/crystal_PDBs/pdbs/')
+
+        if name == 'FLEX_CATH4.3':
+            from src.datasets.flex_cath_dataset import FlexCATHDataset
+            self.data_module = FlexCATHDataset
+            self.hparams['version'] = 4.3
+            self.hparams['path'] = osp.join(self.hparams.data_root, 'cath4.3')
+
+
+    def instancialize(self, **other_args):
+        """ Instancialize a model using the corresponding parameters
+            from self.hparams dictionary. You can also input any args
+            to overwrite the corresponding value in self.kwargs.
+        """
+        class_args =  list(inspect.signature(self.data_module.__init__).parameters)[1:]
+        inkeys = self.hparams.keys()
+        args1 = {}
+        for arg in class_args:
+            if arg in inkeys:
+                args1[arg] = self.hparams[arg]
+        args1.update(other_args)
+
+        # if self.hparams['test_engineering'] and self.hparams['use_dynamics']:
+        #     args1['data_jsonl_name'] = self.hparams['test_eng_data_path']
+        #elif self.hparams['use_dynamics']:
+        if self.hparams['use_dynamics']:
+            args1['data_jsonl_name'] = load_yaml_config('configs/ANMAwareFlexibilityProtTrans.yaml')['data_jsonl_name']
+        # import pdb; pdb.set_trace()
+        return self.data_module(**args1) #Here this leads to __init__ of the class dataset

Flexpert-Design/data_utils.py

@@ -0,0 +1,535 @@
+#From https://github.com/JoreyYan/zetadesign/blob/master/data/data.py
+import glob
+import json
+import numpy as np
+import gzip
+import re
+import multiprocessing
+import tqdm
+import shutil
+SENTINEL = 1
+import biotite.structure as struc
+import biotite.application.dssp as dssp
+import biotite.structure.io.pdb.file as file
+
+
+def parse_PDB_biounits(x, sse,ssedssp,atoms=['N', 'CA', 'C'], chain=None):
+    '''
+    input:  x = PDB filename
+            atoms = atoms to extract (optional)
+    output: (length, atoms, coords=(x,y,z)), sequence
+    '''
+
+    alpha_1 = list("ARNDCQEGHILKMFPSTWYV-")
+    states = len(alpha_1)
+    alpha_3 = ['ALA', 'ARG', 'ASN', 'ASP', 'CYS', 'GLN', 'GLU', 'GLY', 'HIS', 'ILE',
+               'LEU', 'LYS', 'MET', 'PHE', 'PRO', 'SER', 'THR', 'TRP', 'TYR', 'VAL', 'GAP']
+
+    aa_1_N = {a: n for n, a in enumerate(alpha_1)}
+    aa_3_N = {a: n for n, a in enumerate(alpha_3)}
+    aa_N_1 = {n: a for n, a in enumerate(alpha_1)}
+    aa_1_3 = {a: b for a, b in zip(alpha_1, alpha_3)}
+    aa_3_1 = {b: a for a, b in zip(alpha_1, alpha_3)}
+
+    def AA_to_N(x):
+        # ["ARND"] -> [[0,1,2,3]]
+        x = np.array(x);
+        if x.ndim == 0: x = x[None]
+        return [[aa_1_N.get(a, states - 1) for a in y] for y in x]
+
+    def N_to_AA(x):
+        # [[0,1,2,3]] -> ["ARND"]
+        x = np.array(x);
+        if x.ndim == 1: x = x[None]
+        return ["".join([aa_N_1.get(a, "-") for a in y]) for y in x]
+
+    xyz, seq, plddts, min_resn, max_resn = {}, {}, [],  1e6, -1e6
+
+    pdbcontents = x.split('\n')[0]
+    with open(pdbcontents) as f:
+        pdbcontents = f.readlines()
+    for line in pdbcontents:
+        #line = line.decode("utf-8", "ignore").rstrip()
+
+        if line[:6] == "HETATM" and line[17:17 + 3] == "MSE":
+            line = line.replace("HETATM", "ATOM  ")
+            line = line.replace("MSE", "MET")
+
+        if line[:4] == "ATOM":
+            ch = line[21:22]
+            if ch == chain or chain is None or ch==' ':
+                atom = line[12:12 + 4].strip()
+                resi = line[17:17 + 3]
+                resn = line[22:22 + 5].strip()
+                plddt=line[60:60 + 6].strip()
+
+
+
+                x, y, z = [float(line[i:(i + 8)]) for i in [30, 38, 46]]
+
+                if resn[-1].isalpha():
+                    resa, resn = resn[-1], int(resn[:-1]) - 1 # in same pos ,use last atoms
+                else:
+                    resa, resn = "_", int(resn) - 1
+                #         resn = int(resn)
+                if resn < min_resn:
+                    min_resn = resn
+                if resn > max_resn:
+                    max_resn = resn
+
+
+
+                if resn not in xyz:
+                    xyz[resn] = {}
+                if resa not in xyz[resn]:
+                    xyz[resn][resa] = {}
+                if resn not in seq:
+                    seq[resn] = {}
+
+                if resa not in seq[resn]:
+                    seq[resn][resa] = resi
+
+                if atom not in xyz[resn][resa]:
+                    xyz[resn][resa][atom] = np.array([x, y, z])
+
+
+
+    # convert to numpy arrays, fill in missing values
+    seq_, xyz_ ,sse_,ssedssp_= [], [], [], []
+    dsspidx=0
+    sseidx=0
+    # try:
+    #     for resn in range(min_resn, max_resn + 1):
+    #         if resn in seq:
+    #             for k in sorted(seq[resn]):
+    #                 seq_.append(aa_3_N.get(seq[resn][k], 20))
+    #                 try:
+    #                     if 'CA' in xyz[resn][k]:
+    #                         sse_.append(sse[sseidx])
+    #                         sseidx = sseidx + 1
+    #                     else:
+    #                         sse_.append('-')
+    #                 except:
+    #                     print('error sse')
+
+
+    #         else:
+    #             seq_.append(20)
+    #             sse_.append('-')
+
+    #         misschianatom = False
+    #         if resn in xyz:
+
+
+    #             for k in sorted(xyz[resn]):
+    #                 for atom in atoms:
+    #                     if atom in xyz[resn][k]:
+    #                         xyz_.append(xyz[resn][k][atom])  #some will miss C and O ,but sse is normal,because sse just depend on CA
+    #                     else:
+    #                         xyz_.append(np.full(3, np.nan))
+    #                         misschianatom=True
+    #                 if misschianatom:
+    #                     ssedssp_.append('-')
+    #                     misschianatom = False
+    #                 else:
+    #                     try:
+    #                         ssedssp_.append(ssedssp[dsspidx])         # if miss chain atom,xyz ,seq think is ok , but dssp miss this
+    #                         dsspidx = dsspidx + 1
+    #                     except:
+    #                         print(dsspidx)
+
+
+    #         else:
+    #             for atom in atoms:
+    #                 xyz_.append(np.full(3, np.nan))
+    #             ssedssp_.append('-')
+
+
+    #     return np.array(xyz_).reshape(-1, len(atoms), 3), N_to_AA(np.array(seq_)),np.array(sse_),np.array(ssedssp_)
+    # except TypeError as e:
+    #     print(f"TypeError: {e}")
+    #     return 'no_chain', 'no_chain','no_chain'
+
+    for resn in range(int(min_resn), int(max_resn + 1)):
+        if resn in seq:
+            for k in sorted(seq[resn]):
+                seq_.append(aa_3_N.get(seq[resn][k], 20))
+                try:
+                    if 'CA' in xyz[resn][k]:
+                        sse_.append(sse[sseidx])
+                        sseidx = sseidx + 1
+                    else:
+                        sse_.append('-')
+                except:
+                    print('error sse')
+
+
+        else:
+            seq_.append(20)
+            sse_.append('-')
+
+        misschianatom = False
+        if resn in xyz:
+
+
+            for k in sorted(xyz[resn]):
+                for atom in atoms:
+                    if atom in xyz[resn][k]:
+                        xyz_.append(xyz[resn][k][atom])  #some will miss C and O ,but sse is normal,because sse just depend on CA
+                    else:
+                        xyz_.append(np.full(3, np.nan))
+                        misschianatom=True
+                if misschianatom:
+                    ssedssp_.append('-')
+                    misschianatom = False
+                else:
+                    try:
+                        ssedssp_.append(ssedssp[dsspidx])         # if miss chain atom,xyz ,seq think is ok , but dssp miss this
+                        dsspidx = dsspidx + 1
+                    except:
+                        print(dsspidx)
+
+
+        else:
+            for atom in atoms:
+                xyz_.append(np.full(3, np.nan))
+            ssedssp_.append('-')
+
+
+    return np.array(xyz_).reshape(-1, len(atoms), 3), N_to_AA(np.array(seq_)),np.array(sse_),np.array(ssedssp_)
+
+
+def parse_PDB(path_to_pdb,name, input_chain_list=None):
+    """
+    make sure every time just input 1 line
+    """
+    c = 0
+    pdb_dict_list = []
+
+
+    if input_chain_list:
+        chain_alphabet = input_chain_list
+    else:
+        init_alphabet = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S',
+                         'T',
+                         'U', 'V', 'W', 'X', 'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm',
+                         'n',
+                         'o', 'p', 'q', 'r', 's', 't', 'u', 'v', 'w', 'x', 'y', 'z']
+        extra_alphabet = [str(item) for item in list(np.arange(300))]
+        chain_alphabet = init_alphabet + extra_alphabet
+
+    biounit_names = [path_to_pdb]
+    for biounit in biounit_names:
+        my_dict = {}
+        s = 0
+        concat_seq = ''
+
+
+        for letter in chain_alphabet:
+
+            PDBFile = file.PDBFile.read(biounit)
+            array_stack = PDBFile.get_structure(altloc="all")
+
+
+            sse1 = struc.annotate_sse(array_stack[0], chain_id=letter).tolist()
+            if len(sse1)==0:
+                sse1 = struc.annotate_sse(array_stack[0], chain_id='').tolist()
+            #ssedssp1 = dssp.DsspApp.annotate_sse(array_stack).tolist()
+            ssedssp1 = [] #not annotating dssp for now
+
+
+            xyz, seq, _, _= parse_PDB_biounits(biounit,sse1,ssedssp1,atoms=['N', 'CA', 'C','O'], chain=letter) #TODO: fix the float error
+            #ssedssp = sse  #faking it for now
+            # if len(sse)!=len(seq[0]):
+            #     xxxx=len(seq[0])
+            #     print(name)
+            #assert len(sse)==len(seq[0])
+            #assert len(ssedssp) == len(seq[0])
+
+            if type(xyz) != str:
+                concat_seq += seq[0]
+                my_dict['seq_chain_' + letter] = seq[0]
+
+                coords_dict_chain = {}
+                coords_dict_chain['N'] = xyz[:, 0, :].tolist()
+                coords_dict_chain['CA'] = xyz[:, 1, :].tolist()
+                coords_dict_chain['C'] = xyz[:, 2, :].tolist()
+                coords_dict_chain['O'] = xyz[:, 3, :].tolist()
+                my_dict['coords_chain_' + letter] = coords_dict_chain
+
+                #sse=''.join(sse)
+                #ssedssp=''.join(ssedssp)
+                #my_dict['sse3' ] = sse
+                #my_dict['sse8'] = ssedssp
+                s += 1
+        #fi = biounit.rfind("/")
+        my_dict['name'] = name#biounit[(fi + 1):-4]
+        my_dict['num_of_chains'] = s
+        my_dict['seq'] = concat_seq
+        if s <= len(chain_alphabet):
+            pdb_dict_list.append(my_dict)
+            c += 1
+    return pdb_dict_list
+
+
+
+
+
+def parse_pdb_split_chain(pdbgzFile):
+
+    with open(pdbgzFile) as f:
+
+
+        lines = f.readlines()
+        # pdbcontent = f.decode()
+
+
+        pattern = re.compile('ATOM\s+\d+\s*\w+\s*[A-Z]{3,4}\s*(\w)\s*.+\n', re.MULTILINE)
+        match = list(set(list(pattern.findall(lines[0]))))
+
+
+    name=pdbgzFile.split('/')[-1]
+        #for chain in match:
+        # parse_PDB
+    # match=[name[4]]
+    # match=['A']
+    pdb_data=parse_PDB(pdbgzFile,name,match)
+
+    return pdb_data
+def parse_pdb_split_chain_af(pdbgzFile):
+    with gzip.open(pdbgzFile, 'rb') as pdbF:
+        try:
+            pdbcontent = pdbF.read()
+        except:
+            print(pdbgzFile)
+
+        pdbcontent = pdbcontent.decode()
+
+
+        pattern = re.compile('ATOM\s+\d+\s*\w+\s*[A-Z]{3,4}\s*(\w)\s*.+\n', re.MULTILINE)
+        match = list(set(list(pattern.findall(pdbcontent))))
+
+
+    name=pdbgzFile.split('/')[-1].split('.')[0]
+        #for chain in match:
+        # parse_PDB
+    # match=[name[4]]
+    # match=[1]
+    pdb_data=parse_PDB('/media/junyu/data/perotin/aftest080_1000/'+pdbgzFile.split('/')[-1].split('.')[0]+'.pdb',name,match)
+
+    return pdb_data
+
+def parse_pdb_split_chain_af_3dcnn(pdbgzFile):
+    with gzip.open(pdbgzFile, 'rb') as pdbF:
+        try:
+            pdbcontent = pdbF.read()
+        except:
+            print(pdbgzFile)
+
+        pdbcontent = pdbcontent.decode()
+
+
+        pattern = re.compile('ATOM\s+\d+\s*\w+\s*[A-Z]{3,4}\s*(\w)\s*.+\n', re.MULTILINE)
+        match = list(set(list(pattern.findall(pdbcontent))))
+
+
+    name=pdbgzFile.split('/')[-1].split('.')[0]
+    namelist=[]
+    for chain in match:
+        namelist.append(name+'__'+chain)
+    # match=[name[4]]
+    # match=[1]
+
+
+
+    return namelist
+def run_net(files_path,output_path):
+    """
+    input is pdbgz's dir
+    from pdb to jsonl
+    """
+    list=glob.glob(files_path+'*.pdb')#[:3110]
+    data=[]
+    for i in tqdm.tqdm(list):
+        data_chains=parse_pdb_split_chain(i)
+        #for chian in data_chains:
+        data.append(data_chains[0])
+
+    print('we want to write now')
+    with open(output_path, 'w') as f:
+        for entry in data:
+            f.write(json.dumps(entry) + '\n')
+
+    f.close()
+    print('finished')
+
+def run_netbyondif(filelist,output_path):
+    with open(filelist) as f:
+
+        lines = f.readlines()
+    data=[]
+    data_1=[]
+    # data_2 = []
+    # data_3 = []
+    # data_4 = []
+    # data_5 = []
+    # data_6 = []
+    # data_7 = []
+    # data_8 = []
+    # data_9 = []
+    # data_10 = []
+    nums_dict={1:0,2:0,3:0,4:0,5:0,6:0,7:0,8:0,9:0,10:0,}
+
+    for i in tqdm.tqdm(lines):
+        data_chains,match=parse_pdb_split_chain(i.split('"')[1])
+
+        for chian in data_chains:
+            for i in match:
+                meanplddt = round(float(np.asarray(chian['plddts_chain_' + i]).mean()),2)
+                data.append({'name':chian['name'],'lens':len(chian['seq']),'meanplddt':meanplddt})
+                if int(meanplddt/10)==1:
+                    #data_1.append(chian)
+                    nums_dict[1]=nums_dict[1]+1
+                elif int(meanplddt/10)==2:
+                    #data_2.append(chian)
+                    nums_dict[2] = nums_dict[2] + 1
+                elif int(meanplddt / 10) == 3:
+                    #data_3.append(chian)
+                    nums_dict[3] = nums_dict[3] + 1
+                elif int(meanplddt / 10) == 4:
+                    #data_4.append(chian)
+                    nums_dict[4] = nums_dict[4] + 1
+                elif int(meanplddt / 10) == 5:
+                    #data_5.append(chian)
+                    nums_dict[5] = nums_dict[5] + 1
+                elif int(meanplddt / 10) == 6:
+                    #data_6.append(chian)
+                    nums_dict[6] = nums_dict[6] + 1
+                elif int(meanplddt / 10) == 7:
+                    #data_7.append(chian)
+                    nums_dict[7] = nums_dict[7] + 1
+                elif int(meanplddt / 10) == 8:
+                    #data_8.append(chian)
+                    nums_dict[8] = nums_dict[8] + 1
+                elif int(meanplddt / 10) == 9:
+                    #data_9.append(chian)
+                    nums_dict[9] = nums_dict[9] + 1
+                elif int(meanplddt / 10) == 10:
+                    #data_10.append(chian)
+                    nums_dict[10] = nums_dict[10] + 1
+                else:
+                    print(chian['name'])
+
+
+    #         data.append(chian)
+    #
+    f.close()
+    output_pathindex=output_path+filelist.split('/')[-1].split('.')[0]+'_detail.jsonl'
+    print('we want to write now')
+    with open(output_pathindex, 'w') as f:
+        for entry in data:
+            f.write(json.dumps(entry) + '\n')
+
+    f.close()
+    #print(nums_dict)
+    # count(output_pathindex)
+    print('finished')
+def list_of_groups(list_info, per_list_len):
+    '''
+    :param list_info:   列表
+    :param per_list_len:  每个小列表的长度
+    :return:
+    '''
+    list_of_group = zip(*(iter(list_info),) *per_list_len)
+    end_list = [list(i) for i in list_of_group] # i is a tuple
+    count = len(list_info) % per_list_len
+    end_list.append(list_info[-count:]) if count !=0 else end_list
+    return end_list
+
+def count(filelist):
+    with open(filelist) as f:
+
+        lines = f.readlines()
+    plddts=[]
+
+    for i in tqdm.tqdm(lines):
+        pl=json.loads(i)['meanplddt']
+        plddts.append(int(pl/10))
+
+    for i in range(10):
+        print('counts '+str(i),plddts.count(i))
+
+def run_net_aftest(files_path,output_path):
+    """
+    input is pdbgz's dir
+    """
+    with open(files_path) as f:
+        lines = f.readlines()
+    data=[]
+    for i in tqdm.tqdm(lines):
+
+        data_chains=parse_pdb_split_chain_af('/media/junyu/data/point_cloud/'+i.split('"')[1])
+        for chian in data_chains:
+            data.append(chian)
+
+    # print('we want to write now')
+    # with open(output_path, 'w') as f:
+    #     for entry in data:
+    #         f.write(json.dumps(entry) + '\n')
+    #
+    # f.close()
+    # print('finished')
+
+    output_pathindex = output_path + str(80) + 'bigthanclass_1000.text'
+    print('we want to write now')
+    with open(output_pathindex, 'w') as f:
+        for entry in data:
+            f.write(entry + '\n')
+
+
+    f.close()
+
+# if __name__ == "__main__":
+#     files_path='/media/junyu/data/perotin/chain_set/AFDATA/details/80bigthanclass_1000.jsonl' #'/home/junyu/下载/splits/'#
+#     output_path='/media/junyu/data/perotin/chain_set/'
+
+
+
+#     # run_net_aftest(files_path,output_path)
+
+#     fakedata='//home/oem/pdb-tools/pdbtools/fixed/'
+#     run_net(fakedata,output_path+'tim184.jsonl')
+
+
+
+        #
+        # f.close()
+        # # print(nums_dict)
+        # print('finished ' +str(i))
+
+
+    # alllist=list_of_groups(lists,10000)
+
+    # for i in range(len(alllist)):
+    #     thislist=alllist[i]
+    #     with open(output_path+'_'+str(i)+'.jsonl', 'w') as f:
+    #         for entry in thislist:
+    #             f.write(json.dumps(entry) + '\n')
+    #
+    #     f.close()
+    #     # print(nums_dict)
+    #     print('finished ' +str(i))
+
+    # _processes = []
+
+
+
+    # q = multiprocessing.Queue()
+    #
+    # proc.start()
+    # for eachlist in alllist:
+    #     _process = multiprocessing.Process(target=run_netbyondif, args=(eachlist,))
+    #     _process.start()
+
+
+    # run_netbyondif(lists,output_path)

Flexpert-Design/download-cath-data.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+echo "Downloading CATH data..."
+
+# Create data directory if it doesn't exist
+mkdir -p ../data/
+
+# Set file information
+URL="https://data.ciirc.cvut.cz/public/projects/2025Flexpert/cath4.3/"
+OUTPUT_DIR="../data/cath4.3"
+
+# Download directory recursively
+echo "Downloading CATH data..."
+wget --no-check-certificate -r -np -nH --cut-dirs=3 --reject "index.html*" \
+     --directory-prefix=${OUTPUT_DIR} ${URL}
+
+echo "CATH data download completed."
+

Flexpert-Design/model_interface.py

@@ -0,0 +1,631 @@
+import sys; sys.path.append('/huyuqi/xmyu/DiffSDS')
+import inspect
+import torch
+from src.tools.utils import cuda
+import torch.nn as nn
+import os
+from torcheval.metrics.text import Perplexity
+from src.interface.model_interface import MInterface_base
+import math
+import torch.nn.functional as F
+from omegaconf import OmegaConf
+from src.tools.utils import load_yaml_config
+import torchmetrics
+ 
+class MInterface(MInterface_base):
+    def __init__(self, model_name=None, loss=None, lr=None, **kwargs):
+        super().__init__()
+        self.save_hyperparameters()
+        self.load_model()
+        self.use_dynamics = kwargs.get('use_dynamics', 0)
+        self.flex_loss_coeff = torch.Tensor([kwargs.get('flex_loss_coeff', 0)]).to('cuda:0').to(torch.float)
+        self.flex_loss_coeff.requires_grad = False
+        if self.use_dynamics:
+            self.load_flex_predictor()
+            self.flex_loss_type = kwargs.get('loss_fn', 0)
+            if self.flex_loss_type == 'MSE':
+                self.flex_loss_fn = nn.MSELoss(reduction='none')
+            elif self.flex_loss_type == 'L1':
+                self.flex_loss_fn = nn.L1Loss(reduction='none')
+            elif self.flex_loss_type == 'DPO':
+                self.flex_loss_fn = ...
+            else:
+                raise ValueError(f"Not recognized type of loss function {self.flex_loss_type}")
+        self.cross_entropy = nn.NLLLoss(reduction='none')
+        os.makedirs(os.path.join(self.hparams.res_dir, self.hparams.ex_name), exist_ok=True)
+
+        self.control_sum_recovery = 0
+        self.control_sum_batch_sizes = 0
+
+        self.grad_normalization = kwargs.get('grad_normalization', 0)
+        self.use_pmpnn_checkpoint = kwargs.get('use_pmpnn_checkpoint',0)
+        
+        if self.use_pmpnn_checkpoint:
+            print('Loading pmpnn checkpoint from {}'.format(self.model.pmpnn_init_weights_path))
+            state_dict = torch.load(self.model.pmpnn_init_weights_path)['state_dict'] #['module']
+            state_dict = {key: value for key, value in state_dict.items() if 'model.' in key[:6]}
+            state_dict = {key.replace("model.", ""): value for key, value in state_dict.items()}
+            self.model.load_state_dict(state_dict)
+
+        self.MSE = nn.MSELoss(reduction='none')
+        self.automatic_optimization = False
+
+        if self.hparams.use_dynamics:
+            self.pearson = torchmetrics.PearsonCorrCoef()
+            self.spearman = torchmetrics.SpearmanCorrCoef()
+            self.validation_step_outputs = []
+            self.test_step_outputs = []
+
+        #### setting forward hook
+
+        # def forward_hook(module, input, output):
+        #     def check_nan(tensor):
+        #         if isinstance(tensor, torch.Tensor):
+        #             if torch.isnan(tensor).any():
+        #                 print(f"NaN detected in the output of {type(module).__name__}")
+        #                 print(f"Tensor shape: {tensor.shape}")
+        #                 print(f"Tensor stats: mean={tensor.mean()}, std={tensor.std()}, min={tensor.min()}, max={tensor.max()}, all={torch.isnan(tensor).all()}")
+        #         elif isinstance(tensor, tuple):
+        #             for i, t in enumerate(tensor):
+        #                 if isinstance(t, torch.Tensor):
+        #                     if torch.isnan(t).any():
+        #                         print(f"NaN detected in the output[{i}] of {type(module).__name__}")
+        #                         print(f"Tensor shape: {t.shape}")
+        #                         print(f"Tensor stats: mean={t.mean()}, std={t.std()}, min={t.min()}, max={t.max()}, all={torch.isnan(tensor).all()}")
+
+        #     if isinstance(output, tuple):
+        #         for i, out in enumerate(output):
+        #             check_nan(out)
+        #     else:
+        #         check_nan(output)
+
+        # for name, module in self.model.named_modules():
+        #     module.register_forward_hook(forward_hook)
+        
+        # for name, module in self.flex_model.named_modules():
+        #     module.register_forward_hook(forward_hook)
+
+        ####
+
+    def forward(self, batch, mode='train', temperature=1.0):
+        if self.hparams.augment_eps>0:
+            batch['X'] = batch['X'] + self.hparams.augment_eps * torch.randn_like(batch['X'])
+
+        batch = self.model._get_features(batch)
+        results = self.model(batch)
+        
+        log_probs, mask = results['log_probs'], batch['mask']
+        if len(log_probs.shape) == 3:
+            if self.hparams.use_dynamics:
+                loss = self.combined_flex_aware_loss(batch, pred_log_probs=log_probs)
+                #loss = loss_dict['combined_loss']
+            else:
+                loss = self.cross_entropy(log_probs.permute(0,2,1), batch['S'])
+                loss = (loss*mask).sum()/(mask.sum())
+        elif len(log_probs.shape) == 2:
+            if self.hparams.model_name == 'GVP':
+                loss = self.cross_entropy(log_probs, batch.seq)
+            else:
+                loss = self.cross_entropy(log_probs, batch['S'])
+            
+            if self.hparams.model_name == 'AlphaDesign':
+                loss += self.cross_entropy(results['log_probs0'], batch['S'])
+            loss = (loss*mask).sum()/(mask.sum())
+        
+        cmp = log_probs.argmax(dim=-1)==batch['S']
+        recovery = (cmp*mask).sum()/(mask.sum())
+
+        if mode == 'predict':
+            return {'original_sequence':batch['S'],'correct_positions': cmp, 'mask':mask,'loss':loss, 'recovery':recovery, 'title':batch['title'], 'log_probs': log_probs, 'batch':batch} #, 'gt_bfactors': batch['norm_bfactors'], 'batch':batch}
+        elif mode == 'eval':
+            return {'original_sequence':batch['S'],'correct_positions': cmp, 'mask':mask,'loss':loss, 'recovery':recovery, 'title':batch['title'], 'log_probs': log_probs, 'batch':batch}
+        else:
+            return loss, recovery
+
+    def avgCorrelations(self, preds, gts, masks):
+        pearson_R = 0
+        spearman_R = 0
+        valid_datapoints = 0
+        for pred, gt, mask in zip(preds, gts, masks):
+            dpR = self.pearson(pred[torch.where(mask)], gt[torch.where(mask)])
+            if torch.isnan(dpR):
+                continue
+            else:
+                pearson_R += dpR
+                spearman_R += self.spearman(pred[torch.where(mask)], gt[torch.where(mask)])
+                valid_datapoints += 1
+        return pearson_R/valid_datapoints, spearman_R/valid_datapoints
+
+    def temperature_schedular(self, batch_idx):
+        total_steps = self.hparams.steps_per_epoch*self.hparams.epoch
+        
+        initial_lr = 1.0
+        circle_steps = total_steps//100
+        x = batch_idx / total_steps
+        threshold = 0.48
+        if x<threshold:
+            linear_decay = 1 - 2*x
+        else:
+            K = 1 - 2*threshold
+            linear_decay = K - K*(x-threshold)/(1-threshold)
+        
+        new_lr = (1+math.cos(batch_idx/circle_steps*math.pi))/2*linear_decay*initial_lr
+
+        return new_lr
+    
+    # def get_grad_norm(self):
+    #     total_norm = 0
+    #     parameters = [p for p in self.parameters() if p.grad is not None and p.requires_grad]
+    #     for p in parameters:
+    #         param_norm = p.grad.detach().data.norm(2)
+    #         total_norm += param_norm.item() ** 2
+    #     total_norm = total_norm ** 0.5                
+    #     return total_norm
+
+    #https://lightning.ai/docs/pytorch/1.9.0/notebooks/lightning_examples/basic-gan.html
+    def training_step(self, batch, batch_idx, **kwargs):
+        if self.use_dynamics:
+            raw_loss, recovery = self(batch)
+            if type(raw_loss) == dict:
+                flex_loss = raw_loss['flex_loss']
+                seq_loss = raw_loss['seq_loss']
+                opt = self.optimizers()
+                opt.zero_grad()
+                
+                _params_for_optimization = [p for p in self.model.parameters() if p.requires_grad]
+                _params_for_optimization += [p for p in self.flex_model.parameters() if p.requires_grad]
+                
+                grads_flex = torch.autograd.grad(flex_loss, _params_for_optimization, create_graph=True)
+                grads_seq = torch.autograd.grad(seq_loss, _params_for_optimization, create_graph=True)
+                if self.grad_normalization:
+                    norm_grads_flex = [g / (g.norm() + 1e-10) for g in grads_flex]
+                    norm_grads_seq = [g / (g.norm() + 1e-10) for g in grads_seq]
+                else:
+                    norm_grads_flex = grads_flex
+                    norm_grads_seq = grads_seq
+
+                combined_grads = [self.flex_loss_coeff * gflex + (1-self.flex_loss_coeff) * gseq for gflex, gseq in zip(norm_grads_flex, norm_grads_seq)]
+            
+                #maybe track the angle between the gradients?
+                self.log_dict({'flex_grad_norm':torch.mean(torch.tensor([g.detach().norm() for g in norm_grads_flex])), 'seq_grad_norm': torch.mean(torch.tensor([g.detach().norm() for g in norm_grads_seq])), 'combined_grad_norm': torch.mean(torch.tensor([g.detach().norm() for g in combined_grads]))}, on_step=True, on_epoch=False, prog_bar=True)
+                
+
+                for param, grad in zip(_params_for_optimization, combined_grads):
+                    if param.grad is None:
+                        param.grad = grad.detach()
+                    else:
+                        param.grad += grad.detach()
+
+                self.clip_gradients(opt, gradient_clip_val=1., gradient_clip_algorithm="norm")
+                opt.step()
+                
+                # Update learning rate
+                sch = self.lr_schedulers()
+                if sch is not None:
+                    sch.step()
+                
+                loss = flex_loss + seq_loss
+
+                self.log_dict({'train_flex_loss':flex_loss, 'train_seq_loss':seq_loss}, on_step=True, on_epoch=False, prog_bar=True)
+                
+                # Log the current learning rate
+                if sch is not None:
+                    current_lr = sch.get_last_lr()[0]
+                    self.log('learning_rate', current_lr, on_step=True, on_epoch=False, prog_bar=True)
+            else:
+                loss = raw_loss
+                self.log('loss', loss, on_step=True, on_epoch=True, prog_bar=True)
+            return loss
+        else:
+            raw_loss, recovery = self(batch)
+            if type(raw_loss) == dict:
+                loss = raw_loss['combined_loss']
+                _ = raw_loss.pop('pred_flex')
+                # _ = raw_loss.pop('gt_bfactors')
+                _ = raw_loss.pop('gt_flex')
+                _ = raw_loss.pop('flex_mask')
+
+                self.log_dict(raw_loss, on_step=True, on_epoch=True, prog_bar=True)
+            else:
+                loss = raw_loss
+                self.log('loss', loss, on_step=True, on_epoch=True, prog_bar=True)
+            return loss
+    
+    def validation_step(self, batch, batch_idx):
+        raw_loss, recovery = self(batch)
+        if type(raw_loss) == dict:
+            loss = raw_loss['flex_loss']+raw_loss['seq_loss'] #raw_loss['combined_loss']
+            raw_loss['recovery'] = recovery
+            pred_flex = raw_loss.pop('pred_flex')
+            gt_flex = batch['gt_flex']
+
+            flex_mask = raw_loss.pop('flex_mask')
+            #epoch_metric_ingredients = {'pred_bfactors':pred_bfactors, 'gt_bfactors':gt_bfactors, 'flex_mask':flex_mask}
+            epoch_metric_ingredients = {'pred_flex': pred_flex,'gt_flex':gt_flex, 'flex_mask':flex_mask}
+            self.validation_step_outputs.append(epoch_metric_ingredients)
+            self.log_dict({ "val_combined_loss":loss,
+                            "val_seq_loss":raw_loss['seq_loss'],
+                            "val_flex_loss":raw_loss['flex_loss'],
+                            "recovery": recovery})
+        else:
+            loss = raw_loss
+            self.log_dict({"val_loss":loss,
+                        "recovery": recovery})
+        #if there is issue with validation metrics - see the test_step below
+        return self.log_dict
+
+    def on_validation_epoch_end(self):
+        if self.hparams.use_dynamics:
+            # all_preds = [b['pred_bfactors'] for b in self.validation_step_outputs]
+            # all_gts = [b['gt_bfactors'] for b in self.validation_step_outputs]
+            all_preds = [b['pred_flex'] for b in self.validation_step_outputs]
+            all_gts = [b['gt_flex'] for b in self.validation_step_outputs]
+            all_masks = [b['flex_mask'] for b in self.validation_step_outputs]
+            
+            max_seq_length = max([pred.size()[1] for pred in all_preds])
+
+            for set_of_tensors in [all_preds, all_gts, all_masks]:
+                for i in range(len(set_of_tensors)):
+                    set_of_tensors[i] = F.pad(set_of_tensors[i], (0, max_seq_length - set_of_tensors[i].shape[1],0,0), value=float(0))
+            all_preds = torch.cat(all_preds, dim=0)
+            all_gts = torch.cat(all_gts, dim=0)
+            all_masks = torch.cat(all_masks, dim=0)
+            # print(all_preds.shape, all_gts.shape, all_masks.shape)
+            # do something with all preds
+            
+            # pearson_R = self.pearson(all_preds[torch.where(all_masks)], all_gts[torch.where(all_masks)])
+            pearson_R, spearman_R = self.avgCorrelations(all_preds, all_gts, all_masks)
+            # try:
+            #     spearman_R = self.spearman(all_preds[torch.where(all_masks)], all_gts[torch.where(all_masks)])
+            # except IndexError:
+            #     spearman_R = pearson_R
+            self.log_dict({"val_pearson_R":pearson_R, "val_spearman_R":spearman_R})
+            self.validation_step_outputs.clear()  # free memory
+        return super().on_validation_epoch_end()
+
+    def on_test_epoch_end(self):
+        import pickle            #use pickle to save the self.test_step_outputs to a file
+        with open(f'rebuttal_experiments/test_step_outputs_{self.hparams.starting_checkpoint_path.split("/")[-3]}_initFF{self.hparams.init_flex_features}_{self.hparams.test_eng_data_path.split("/")[-1][:-5]}.pkl', 'wb') as f:
+            pickle.dump(self.test_step_outputs, f)
+        if self.hparams.test_engineering and self.hparams.use_dynamics:
+            all_preds = [b['pred_flex'] for b in self.test_step_outputs]
+            all_eng_gts = [b['gt_flex'] for b in self.test_step_outputs]
+            all_masks = [b['flex_mask'] for b in self.test_step_outputs]
+            all_eng_masks = [b['eng_mask'] for b in self.test_step_outputs]
+            all_original_gt_flex = [b['original_gt_flex'] for b in self.test_step_outputs]
+
+            avg_sequence_recovery = sum([b['sequence_recovery'] for b in self.test_step_outputs]) / len(self.test_step_outputs)
+            avg_sequence_recovery = avg_sequence_recovery.cpu().tolist()
+            max_seq_length = max([pred.size()[1] for pred in all_preds])
+            
+            _pred_flex_pool = []
+            _eng_gt_flex_pool = []
+            _original_gt_flex_pool = []
+            _original_gt_flex_ranks_pool = []
+            _eng_gt_flex_ranks_pool = []
+            _pred_flex_ranks_pool = []
+
+
+            import numpy as np
+            for eng_mask, flex_mask, original_gt_flex, eng_gt_flex, pred_flex in zip(all_eng_masks, all_masks, all_original_gt_flex, all_eng_gts, all_preds):
+                #select only the values where the engineering mask is 1 and flex mask is 1
+                _original_gt_flex = original_gt_flex[eng_mask == 1]
+                _eng_gt_flex = eng_gt_flex[eng_mask == 1]
+                _pred_flex = pred_flex[eng_mask == 1]
+                _pred_flex_pool.append(_pred_flex.cpu().numpy())
+                _eng_gt_flex_pool.append(_eng_gt_flex.cpu().numpy())
+                _original_gt_flex_pool.append(_original_gt_flex.cpu().numpy())
+                
+                _original_gt_flex_ranks = torch.argsort(torch.argsort(torch.nan_to_num(original_gt_flex, nan=0)))[eng_mask == 1].cpu().numpy()
+                _eng_gt_flex_ranks = torch.argsort(torch.argsort(torch.nan_to_num(eng_gt_flex, nan=0)))[eng_mask == 1].cpu().numpy()
+                _pred_flex_ranks = torch.argsort(torch.argsort(torch.nan_to_num(pred_flex, nan=0)))[eng_mask == 1].cpu().numpy()
+            
+                _original_gt_flex_ranks_pool.append(_original_gt_flex_ranks)
+                _eng_gt_flex_ranks_pool.append(_eng_gt_flex_ranks)
+                _pred_flex_ranks_pool.append(_pred_flex_ranks)
+
+            
+            import matplotlib.pyplot as plt
+            import os
+
+            # # Create 'paper_figures' folder if it doesn't exist
+            # if not os.path.exists('paper_figures'):
+            #     os.makedirs('paper_figures')
+
+            #pool the numpy arrays in the lists into one numpy array
+            _pred_flex_pool = np.concatenate(_pred_flex_pool)
+            _eng_gt_flex_pool = np.concatenate(_eng_gt_flex_pool)
+            _original_gt_flex_pool = np.concatenate(_original_gt_flex_pool)
+            
+            ############################################################################
+            all_gt_seqs = [b['gt_seq'] for b in self.test_step_outputs]
+            all_pred_logprobs = [b['pred_logprobs'] for b in self.test_step_outputs]
+            _gt_seq_pool = []
+            _pred_seq_pool = []
+            _outside_eng_region_pred_seq_pool = []
+            _outside_eng_region_gt_seq_pool = []
+            for eng_mask, gt_seq, pred_logprobs in zip(all_eng_masks, all_gt_seqs, all_pred_logprobs):
+                #select only the values where the engineering mask is 1
+                _outside_eng_region_pred_seq_pool.append(torch.argmax(pred_logprobs[(eng_mask == 0) & (flex_mask == 1)], dim=1).cpu().numpy())
+                _outside_eng_region_gt_seq_pool.append(gt_seq[(eng_mask == 0) & (flex_mask == 1)].cpu().numpy())
+
+                _pred_seq = torch.argmax(pred_logprobs[eng_mask == 1], dim=1)
+                _gt_seq = gt_seq[eng_mask == 1]
+                
+                # create and add to the pools the numpy arrays
+                _gt_seq_pool.append(_gt_seq.cpu().numpy())
+                _pred_seq_pool.append(_pred_seq.cpu().numpy())
+            _gt_seq_pool = np.concatenate(_gt_seq_pool)
+            _pred_seq_pool = np.concatenate(_pred_seq_pool)
+            _outside_eng_region_pred_seq_pool = np.concatenate(_outside_eng_region_pred_seq_pool)
+            _outside_eng_region_gt_seq_pool = np.concatenate(_outside_eng_region_gt_seq_pool)
+            #output these pools together with the other pools to a json_file
+            import json
+            with open(f'paper_figures/pools_{self.hparams.starting_checkpoint_path.split("/")[-3]}_initFF{self.hparams.init_flex_features}_{self.hparams.test_eng_data_path.split("/")[-1][:-5]}.json', 'w') as f:
+                json.dump({
+                    '_pred_flex_pool': _pred_flex_pool.tolist(),
+                    '_eng_gt_flex_pool': _eng_gt_flex_pool.tolist(),
+                    '_original_gt_flex_pool': _original_gt_flex_pool.tolist(),
+                    '_pred_seq_pool': _pred_seq_pool.tolist(),
+                    '_gt_seq_pool': _gt_seq_pool.tolist(),
+                    '_sequence_recovery': avg_sequence_recovery,
+                    '_outside_eng_region_pred_seq_pool': _outside_eng_region_pred_seq_pool.tolist(),
+                    '_outside_eng_region_gt_seq_pool': _outside_eng_region_gt_seq_pool.tolist()
+                }, f)
+
+
+
+            ############################################################################
+
+
+            self.test_step_outputs.clear()
+        else:
+            # all_preds = [b['pred_bfactors'] for b in self.test_step_outputs]
+            # all_gts = [b['gt_bfactors'] for b in self.test_step_outputs]
+            all_preds = [b['pred_flex'] for b in self.test_step_outputs]
+            all_gts = [b['gt_flex'] for b in self.test_step_outputs]
+            all_masks = [b['flex_mask'] for b in self.test_step_outputs]
+            
+            max_seq_length = max([pred.size()[1] for pred in all_preds])
+
+            for set_of_tensors in [all_preds, all_gts, all_masks]:
+                for i in range(len(set_of_tensors)):
+                    set_of_tensors[i] = F.pad(set_of_tensors[i], (0, max_seq_length - set_of_tensors[i].shape[1],0,0), value=float(0))
+
+            all_preds = torch.cat(all_preds, dim=0)
+            all_gts = torch.cat(all_gts, dim=0)
+            all_masks = torch.cat(all_masks, dim=0)
+            # print(all_preds.shape, all_gts.shape, all_masks.shape)
+            # do something with all preds
+            # pearson_R = self.pearson(all_preds[torch.where(all_masks)], all_gts[torch.where(all_masks)])
+            pearson_R, spearman_R = self.avgCorrelations(all_preds, all_gts, all_masks)
+            try:
+                spearman_R = self.spearman(all_preds[torch.where(all_masks)], all_gts[torch.where(all_masks)])
+            except IndexError:
+                spearman_R = pearson_R
+            self.log_dict({"test_pearson_R":pearson_R, "test_spearman_R":spearman_R})
+            self.test_step_outputs.clear()  # free memory
+        return super().on_test_epoch_end()
+
+    def test_step(self, batch, batch_idx):
+        # Here we just reuse the validation_step for testing
+        #return self.validation_step(batch, batch_idx)
+        
+        raw_loss, recovery = self(batch)
+        if type(raw_loss) == dict:
+            #loss = raw_loss['combined_loss']
+            loss = raw_loss['flex_loss']+raw_loss['seq_loss'] #raw_loss['combined_loss']
+            raw_loss['recovery'] = recovery
+            # pred_bfactors = raw_loss.pop('pred_bfactors')
+            pred_flex = raw_loss.pop('pred_flex')
+            # gt_bfactors = raw_loss.pop('gt_bfactors')
+            gt_flex = raw_loss.pop('gt_flex')
+            flex_mask = raw_loss.pop('flex_mask')
+            epoch_metric_ingredients = {'pred_flex':pred_flex, 'gt_flex':gt_flex, 'flex_mask':flex_mask}
+
+            if self.hparams.test_engineering and self.hparams.use_dynamics:
+                eng_mask = raw_loss.pop('eng_mask')
+                original_gt_flex = raw_loss.pop('original_gt_flex')
+                epoch_metric_ingredients['eng_mask'] = eng_mask
+                epoch_metric_ingredients['original_gt_flex'] = original_gt_flex
+                epoch_metric_ingredients['gt_seq'] = raw_loss['gt_seq']
+                epoch_metric_ingredients['pred_logprobs'] = raw_loss['pred_logprobs']
+                epoch_metric_ingredients['sequence_recovery'] = raw_loss['recovery']
+                epoch_metric_ingredients['id'] = batch['title']
+
+            self.test_step_outputs.append(epoch_metric_ingredients)
+            out_dict = {"val_combined_loss":loss,
+                        "val_seq_loss":raw_loss['seq_loss'],
+                        "val_flex_loss":raw_loss['flex_loss'],
+                        "recovery": recovery}
+        else:
+            out_dict = {"val_loss":raw_loss, "recovery": recovery}
+        self.log_dict(out_dict,on_step=True,on_epoch=True, sync_dist=True)
+        #print(out_dict) #This print statement is fixing it - ultimately fixed by setting 'n_step=True' above
+        #Below validation of the correctness of the above loging
+        self.control_sum_batch_sizes += len(batch['X'])
+        self.control_sum_recovery += len(batch['X'])*recovery
+        return out_dict
+
+    def predict_step(self, batch, batch_idx):
+        predict_out = self(batch, mode=self.hparams.stage)
+        return predict_out
+    
+    def combined_flex_aware_loss(self, batch, pred_log_probs):
+
+        _mask = batch['mask']
+    
+        gt_seq = batch['S']
+        gt_flex = batch['gt_flex']
+        anm_input = batch['enm_vals'] #TODO: manage the loading of the anm input
+
+        trail_idcs = torch.argmax((batch['S'] == 0).int(), dim=1)
+        trail_idcs[trail_idcs == 0] = batch['S'].shape[1]  # For sequences without padding
+
+        # # #TODO: test on one example - remove later
+        # # trail_idcs = trail_idcs[0].unsqueeze(0)
+
+        # # # ###########################################################################
+        # # # #### TODO: change back to precomputed GT_FLEX once debugged ###############
+        # dl_gtseq = batch['S']
+        # dl_anm = batch['enm_vals']
+
+
+        # attention_mask = torch.zeros_like(batch['mask'])
+        # for i in range(attention_mask.size(0)):
+        #     attention_mask[i, :trail_idcs[i]] = 1
+
+        # dl_predflex_bs4 = self.flex_model(None, dl_anm, trail_idcs, attention_mask = attention_mask, sampled_pmpnn_sequence = dl_gtseq, alphabet='pmpnn') #['predicted_flex'][:,:-1,0]
+        # dl_predflex_bs1 = self.flex_model(None, dl_anm[0].unsqueeze(0), trail_idcs[0].unsqueeze(0) , attention_mask = attention_mask[0].unsqueeze(0), sampled_pmpnn_sequence = dl_gtseq[0].unsqueeze(0), alphabet='pmpnn') #['predicted_flex'][:,:-1,0]
+        
+        # testseq = 'MKKAVINGEQIRSISDLHQTLKKELALPEYYGENLDALWDCLTGWVEYPLVLEWRQFEQSKQLTENGAESVLQVFREAKAEGADITIILS'
+        # tokenizer_predflex_bs4 = self.flex_model(None, dl_anm[0,:90].unsqueeze(0), trail_idcs[0].unsqueeze(0) , attention_mask = attention_mask[0,:90].unsqueeze(0), sampled_pmpnn_sequence = testseq, alphabet='aa') #['predicted_flex'][:,:-1,0] #['predicted_flex'][:,:-1,0]
+        # import pdb; pdb.set_trace()
+        # input_ids_predflex_bs4 = self.flex_model(dl_gtseq, dl_anm, trail_idcs, attention_mask = attention_mask, sampled_pmpnn_sequence = None, alphabet='aa') #['predicted_flex'][:,:-1,0]
+        # gt_flex = batch['gt_flex']
+        # # ####
+        # import pdb; pdb.set_trace() #check the mask and the gt_flex vs. onthefly computed gt_flex
+        # #TODO: here fix the mask for the prottrans and clean this,
+        # #      the mask should have all 1s where there is sequence or eos token
+
+        # attention_mask = ...
+        # if self.hparams.get_gt_flex_onthefly:
+            
+        #     cache_keys = list(batch['title'])
+
+        #     # Check if all cache_keys are in self.gt_flex_cache
+        #     all_keys_in_cache = all(cache_key in self.model.gt_flex_cache for cache_key in cache_keys)
+            
+        #     if not all_keys_in_cache:
+        #         gt_flex = self.flex_model(None, anm_input, trail_idcs, attention_mask=attention_mask, sampled_pmpnn_sequence=gt_seq, alphabet='pmpnn')['predicted_flex'][:,:-1,0]
+        #         for key, val in zip(cache_keys, gt_flex):
+        #             #TODO: iteruje to spravne?
+        #             self.model.gt_flex_cache[key] = val
+        #     else:
+        #         retrieved_gt_flexs = []
+        #         for key in cache_keys:
+        #             _gt_flex = self.model.gt_flex_cache[key]
+        #             retrieved_gt_flexs.append(_gt_flex)
+        #         gt_flex = torch.cat(retrieved_gt_flexs, dim=0) #TODO: concat spravne?
+        # else:
+        #     raise NotImplementedError('The precomputed data were not realiable.')
+        #     gt_flex = batch['gt_flex']
+        # ###########################################################################
+
+
+        attention_mask = torch.zeros_like(batch['mask'])
+        for i in range(attention_mask.size(0)):
+            attention_mask[i, :trail_idcs[i]] = 1
+
+        #Original sequence loss
+        seq_loss = self.cross_entropy(pred_log_probs.permute(0,2,1), gt_seq)
+        seq_loss = (seq_loss*_mask).sum()/(_mask.sum())
+        #New Dynamics-aware loss
+        flex_model_input = pred_log_probs.permute(0,2,1)
+        pred_flex = self.flex_model(flex_model_input, anm_input, trail_idcs, attention_mask=attention_mask)['predicted_flex'][:,:-1,0]
+        #check here that the loss function is working properly (with the masking and all)
+        # import pdb; pdb.set_trace()
+        _filter_nans_mask = ~torch.isnan(pred_flex) & ~torch.isnan(gt_flex)
+        flex_loss = self.flex_loss_fn(pred_flex[_filter_nans_mask]*_mask[_filter_nans_mask], gt_flex[_filter_nans_mask]*_mask[_filter_nans_mask])
+        _flex_mask = _mask*_filter_nans_mask
+        _flex_mask = _flex_mask.int()
+        flex_loss = flex_loss.sum()/_flex_mask.sum()
+        
+        retVal ={'seq_loss':seq_loss, 'flex_loss':flex_loss, 'pred_flex':pred_flex, 'flex_mask':_flex_mask, 'gt_flex':gt_flex}
+        if self.hparams.test_engineering and self.hparams.use_dynamics:
+            retVal['eng_mask'] = batch['eng_mask']
+            retVal['original_gt_flex'] = batch['original_gt_flex']
+            retVal['gt_seq'] = batch['S']
+            retVal['pred_logprobs'] = pred_log_probs
+        return retVal
+
+
+    def configure_loss(self):
+        def loss_function(pred_angle, angles, pred_seq, seqs, seq_loss_mask, angle_loss_mask):
+            angle_loss = self.MSE(torch.cat([angles[...,:1],torch.sin(angles[...,1:3]), torch.cos(angles[...,1:3])],dim=-1),
+            torch.cat([pred_angle[...,:1],torch.sin(pred_angle[...,1:3]), torch.cos(pred_angle[...,1:3])],dim=-1))
+            
+            angle_loss = angle_loss[angle_loss_mask].sum(dim=-1).mean()
+            logits = pred_seq.permute(0,2,1)
+            seq_loss = self.cross_entropy(logits, seqs)
+            seq_loss = seq_loss[seq_loss_mask].mean()
+
+            metric=Perplexity()
+            metric.update(pred_seq[seq_loss_mask][None,...].cpu(), seqs[seq_loss_mask][None,...].cpu())
+            perp = metric.compute()
+            
+            return {"angle_loss": angle_loss, "seq_loss": seq_loss, "perp":perp}
+
+        self.loss_function = loss_function
+        
+    def load_model(self):
+        params = OmegaConf.load(f'configs/{self.hparams.model_name}.yaml')
+        params.update(self.hparams)
+
+        if self.hparams.model_name == 'GraphTrans':
+            from src.models.graphtrans_model import GraphTrans_Model
+            self.model = GraphTrans_Model(params)
+        
+        if self.hparams.model_name == 'StructGNN':
+            from src.models.structgnn_model import StructGNN_Model
+            self.model = StructGNN_Model(params)
+            
+        if self.hparams.model_name == 'GVP':
+            from src.models.gvp_model import GVP_Model
+            self.model = GVP_Model(params)
+
+        if self.hparams.model_name == 'GCA':
+            from src.models.gca_model import GCA_Model
+            self.model = GCA_Model(params)
+
+        if self.hparams.model_name == 'AlphaDesign':
+            from src.models.alphadesign_model import AlphaDesign_Model
+            self.model = AlphaDesign_Model(params)
+
+        if self.hparams.model_name == 'ProteinMPNN':
+            from src.models.proteinmpnn_model import ProteinMPNN_Model
+            self.model = ProteinMPNN_Model(params)
+
+        if self.hparams.model_name == 'ESMIF':
+            pass
+
+        if self.hparams.model_name == 'PiFold':
+            from src.models.pifold_model import PiFold_Model
+            self.model = PiFold_Model(params)
+
+        if self.hparams.model_name == 'KWDesign':
+            from src.models.kwdesign_model import KWDesign_model#Design_Model
+            self.model = KWDesign_model(params) #Design_Model(params) - this required to significantly change the constructor of Design_Model
+        
+        if self.hparams.model_name == 'E3PiFold':
+            from src.models.E3PiFold_model import E3PiFold
+            self.model = E3PiFold(params)
+    
+    def load_flex_predictor(self):
+        from src.models.anm_prottrans import ANMAwareFlexibilityProtTrans
+        flex_params = load_yaml_config(f'configs/ANMAwareFlexibilityProtTrans.yaml')
+        # flex_params_dict = OmegaConf.to_container(flex_params, resolve=True)
+        self.flex_model = ANMAwareFlexibilityProtTrans(**flex_params)
+        
+        # consider turning on the gradients for debug purposes
+        self.flex_model.eval()
+        for params in self.flex_model.parameters():
+            params.requires_grad = False
+        
+        #also pass it to proteinmpnn:
+        # self.model.flex_model = self.flex_model
+
+
+    def instancialize(self, Model, **other_args):
+        """ Instancialize a model using the corresponding parameters
+            from self.hparams dictionary. You can also input any args
+            to overwrite the corresponding value in self.hparams.
+        """
+        class_args = inspect.getargspec(Model.__init__).args[1:]
+        inkeys = self.hparams.keys()
+        args1 = {}
+        for arg in class_args:
+            if arg in inkeys:
+                args1[arg] = getattr(self.hparams, arg)
+        args1.update(other_args)
+        return Model(**args1)

Flexpert-Design/predict.py

@@ -0,0 +1,148 @@
+import os, sys, warnings, argparse, math, tqdm, datetime
+import pytorch_lightning as pl
+import torch
+from pytorch_lightning.trainer import Trainer
+import pytorch_lightning.callbacks as plc
+import pytorch_lightning.loggers as plog
+from model_interface import MInterface
+from data_interface import DInterface
+from src.tools.logger import SetupCallback, BackupCodeCallback
+from shutil import ignore_patterns
+from transformers import AutoTokenizer
+import numpy as np
+import yaml
+import wandb
+warnings.filterwarnings("ignore")
+
+def create_parser():
+    parser = argparse.ArgumentParser()
+
+
+    parser.add_argument('--infer_path', type=str, help='Path where to read the data to be predicted and where to save the predictions.')
+
+    # Set-up parameters
+    parser.add_argument('--res_dir', default='./train/results', type=str)
+    parser.add_argument('--ex_name', default='debug', type=str)
+    parser.add_argument('--check_val_every_n_epoch', default=1, type=int)
+    parser.add_argument('--stage', default='predict', type=str) #'fit', 'test' or 'predict'
+    parser.add_argument('--val_check_interval', default=0.5, type=float, help='Validation check interval')
+    
+    parser.add_argument('--dataset', default='PDBInference') # AF2DB_dataset, CATH_dataset
+    parser.add_argument('--model_name', default='ProteinMPNN', choices=['StructGNN', 'GraphTrans', 'GVP', 'GCA', 'AlphaDesign', 'ESMIF', 'PiFold', 'ProteinMPNN', 'KWDesign', 'E3PiFold'])
+    # parser.add_argument('--lr', default=4e-4, type=float, help='Learning rate')
+    # parser.add_argument('--lr_scheduler', default='onecycle')
+    # parser.add_argument('--offline', default=1, type=int)
+    parser.add_argument('--seed', default=111, type=int)
+    
+    parser.add_argument('--num_workers', default=12, type=int)
+    parser.add_argument('--pad', default=1024, type=int)
+    parser.add_argument('--min_length', default=40, type=int)
+    parser.add_argument('--data_root', default='./data/')
+    
+    # Training parameters
+    # parser.add_argument('--epoch', default=10, type=int, help='end epoch')
+    parser.add_argument('--augment_eps', default=0.0, type=float, help='noise level')
+    # parser.add_argument('--gpus', default=1, type=int, help='how many GPUs to train on')
+    # parser.add_argument('--weight_decay', default=0.0, type=float, help='Weight decay for optimizer')
+
+    # # Eval parameters
+    # parser.add_argument('--eval_sequences_sampled', default=1, type=int, help='How many sequences to sample in evaluation.')
+    # parser.add_argument('--eval_sequences_temperature', default=0, type=float, help='What temperature to use for the sampling in evaluation.')
+    # parser.add_argument('--eval_output_dir', default=None, type=str, help='Where to save the evaluation output.')
+
+    # Model parameters
+    parser.add_argument('--use_dist', default=1, type=int)
+    parser.add_argument('--use_product', default=0, type=int)
+    parser.add_argument('--use_pmpnn_checkpoint', default=1, type=int, help='By 1 or 0 decide whether to start with pretrained ProteinMPNN.')
+    parser.add_argument('--checkpoint_path', type=str, default=None, help='Path to the model checkpoint to load weights from')
+
+    # Dynamics aware parameters
+    parser.add_argument('--use_dynamics', default=0, type=int)
+    # parser.add_argument('--flex_loss_coeff', default=0.5, type=float)
+    # parser.add_argument('--get_gt_flex_onthefly', default=0, type=int, help='Flag to get ground truth flexibility on-the-fly (with subsequent caching)')
+    parser.add_argument('--init_flex_features', default=1, type=int, help="Set to 0 if no flexibility information should be passed on input to the node features h_V")
+    # parser.add_argument('--loss_fn', default='MSE', type=str, help= 'Define what loss to use. Choose MSE, L1 or DPO.')
+    # parser.add_argument('--grad_normalization', default=1, type=int, help="Set to 0 if the gradients of the seq and flex losses should not be normalized.")
+    # parser.add_argument('--test_engineering', default=0, type=int, help="In this main.py should be set to 0 to not overwrite the training dataset.")
+    
+    args = parser.parse_args()
+    return args
+
+
+if __name__ == "__main__":
+    
+    args = create_parser()
+    args.batch_size = 1
+    print('In the predict stage, defaulting batch size to 1.')
+
+    assert args.use_dynamics == 0, "In the inference script this should be set to 0."
+
+    if not os.path.exists(args.infer_path):
+        os.makedirs(args.infer_path)
+        
+    if (len(args.infer_path) > 0 or args.dataset=='PDBInference') and (len(args.infer_path) == 0 or args.dataset!='PDBInference'):
+        raise ValueError("You should only use --infer_path with --dataset 'PDBInference' and vice versa.")
+
+
+    # Load model weights from checkpoint if provided
+    if args.checkpoint_path is not None:
+        trained_model_path = args.checkpoint_path
+        print(f"Loading model weights from checkpoint passed by argument: {trained_model_path}")
+    else:
+        with open('configs/Flexpert-Design-inference.yaml', 'r') as f:
+            config = yaml.load(f, Loader=yaml.FullLoader)
+        trained_model_path = config['pmpnn_model_path']
+        print(f"Loading model weights from checkpoint specified in Flexpert-Design-inference.yaml: {trained_model_path}")
+
+    if os.path.exists(trained_model_path):
+        print(f"Rewriting the path to the Flexpert-Design trained ProteinMPNN weights in the model interface.")
+        args.starting_checkpoint_path = trained_model_path
+    else:
+        raise FileNotFoundError(f"Checkpoint file not found at {trained_model_path}")
+
+    pl.seed_everything(args.seed)
+
+    data_module = DInterface(**vars(args))
+
+    data_module.setup(stage='predict')
+
+    model = MInterface(**vars(args))
+
+
+    trainer_config = {
+        'devices': 1,
+        'max_epochs': 1,
+        'num_nodes': 1,
+        "strategy": 'ddp',
+        "precision": '32',
+        'accelerator': 'gpu',
+        'val_check_interval': args.val_check_interval,
+        'check_val_every_n_epoch': args.check_val_every_n_epoch
+    }
+
+    trainer = Trainer(**trainer_config)
+
+    predictions = trainer.predict(model, data_module)
+
+    tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D", cache_dir='./cache_dir/') # mask token: 32
+
+
+    serializable_predictions = []
+    for pred_idx, pred in enumerate(predictions):
+        logprobs = pred['log_probs'].cpu().numpy()[0]  # [L, 21]
+        pmpnn_alphabet_tokens_argmax = logprobs.argmax(axis=-1)  # [L]
+        
+        aa_sequence = ''.join(tokenizer.decode(pmpnn_alphabet_tokens_argmax, skip_special_tokens=True).split())
+
+        # Get probability of the predicted sequence
+        seq_probs = np.exp(logprobs.max(axis=-1))  # [L]
+        avg_prob = float(np.mean(seq_probs))
+        
+        serializable_predictions.append({
+            'prediction_id': pred['batch']['title'][0],
+            'amino_acid_sequence': aa_sequence
+        })
+
+    with open(f'{args.infer_path}/predictions.txt', 'w') as f:
+        for pred in serializable_predictions:
+            f.write(f'>{pred["prediction_id"]}\n{pred["amino_acid_sequence"]}\n')

Flexpert-Design/predict_example/1ah7_A_instructions.csv

@@ -0,0 +1 @@
+0.4214228391647339, 0.37651416659355164, 0.1882496476173401, 0.13774731755256653, 0.11560429632663727, 0.12345632910728455, 0.11075370758771896, 0.09350624680519104, 0.06162628158926964, 0.08504123985767365, 0.05511573329567909, 0.03457929939031601, 0.018956221640110016, 0.05267956107854843, 0.021582268178462982, 0.019682325422763824, 0.005200381390750408, 0.01862833835184574, 0.037708550691604614, 0.02962341532111168, 0.0414130762219429, 0.032966580241918564, 0.04219468683004379, 0.043324653059244156, 0.038419052958488464, 0.06062019616365433, 0.0754077360033989, 0.09575366973876953, 0.09765047580003738, 0.1067374125123024, 0.08417803794145584, 0.09050130844116211, 0.07099245488643646, 0.06242087855935097, 0.046906158328056335, 0.024977944791316986, 0.04039282351732254, 0.04056069627404213, 0.02624698355793953, 0.014836937189102173, 0.033674903213977814, 0.03443623706698418, 0.04525991156697273, 0.05213414505124092, 0.02986733242869377, 0.01742064766585827, 0.03752005845308304, 0.02649688348174095, 0.02672465145587921, 0.03430021554231644, 0.011848143301904202, 0.03361964225769043, 0.027863629162311554, 0.03575271740555763, 0.041042227298021317, 0.08238421380519867, 0.08222152292728424, 0.10173829644918442, 0.11664807796478271, 0.13249793648719788, 0.14384658634662628, 0.1345130354166031, 0.13609081506729126, 0.12496572732925415, 0.10717709362506866, 0.08230947703123093, 0.07971317321062088, 0.07025592774152756, 0.06319792568683624, 0.06464961171150208, 0.04482023045420647, 0.051742203533649445, 0.07986844331026077, 0.09591078013181686, 0.11425718665122986, 0.11205209791660309, 0.10624780505895615, 0.10313349217176437, 0.13002970814704895, 0.13183605670928955, 0.15288424491882324, 0.14854931831359863, 0.13990001380443573, 0.09912189096212387, 0.09130637347698212, 0.07575594633817673, 0.061887726187705994, 0.06014912575483322, 0.0577777624130249, 0.051302842795848846, 0.03530939295887947, 0.040248580276966095, 0.0013590790331363678, 0.015310827642679214, 0.03272499516606331, 0.02609187364578247, 0.0054176910780370235, 0.05427498370409012, 0.051064278930425644, 0.06116481125354767, 0.06309916824102402, 0.08470715582370758, 0.1002785935997963, 0.10495362430810928, 0.09807638078927994, 0.0662725567817688, 0.06513857841491699, 0.048988863825798035, 0.029838263988494873, 0.025865966454148293, 0.02097484841942787, 0.014891650527715683, 0.024081528186798096, 0.045654937624931335, 0.052093200385570526, 0.017663143575191498, 0.02189275622367859, 0.08543915301561356, 0.03505314886569977, 0.019413039088249207, 0.045589953660964966, 0.06793230772018433, 0.041016142815351486, 0.05003933981060982, 0.053235944360494614, 0.05916681885719299, 0.058036819100379944, 0.060588233172893524, 0.07040168344974518, 0.07345925271511078, 0.08298061043024063, 0.09419634193181992, 0.11146273463964462, 0.1405934989452362, 0.15075145661830902, 0.13765454292297363, 0.13978315889835358, 0.1482282280921936, 0.1423584669828415, 0.10484395921230316, 0.07584157586097717, 0.06757079809904099, 0.10134144872426987, 0.08083963394165039, 0.07369125634431839, 0.05454648658633232, 0.08305331319570541, 0.07765821367502213, 0.06511223316192627, 0.056034114211797714, 0.08081831783056259, 0.08526752144098282, 0.07231731712818146, 0.07028429210186005, 0.08094073832035065, 0.06563611328601837, 0.07806897908449173, 0.0859430730342865, 0.08600828796625137, 0.08605027943849564, 0.08578154444694519, 0.07862624526023865, 0.07963275909423828, 0.06170313060283661, 0.05005127564072609, 0.05146761238574982, 0.05499078333377838, 0.059220947325229645, 0.06969373673200607, 0.05268307402729988, 0.06721088290214539, 0.04827176779508591, 0.029251907020807266, 0.04153018817305565, 0.03697451949119568, 0.025836892426013947, 0.04521643742918968, 0.05554598197340965, 0.06007472425699234, 0.04923863708972931, 0.06502534449100494, 0.04392743483185768, 0.036296453326940536, 0.0436030775308609, 0.05658774450421333, 0.034551363438367844, 0.049478061497211456, 0.059964731335639954, 0.07313965260982513, 0.062442418187856674, 0.06896451860666275, 0.08025174587965012, 0.08270157873630524, 0.09252781420946121, 0.09688305854797363, 0.11343701928853989, 0.11080081015825272, 0.12969090044498444, 0.10972093790769577, 0.14756494760513306, 0.1637764275074005, 0.18948377668857574, 0.18522633612155914, 0.18577809631824493, 0.20364972949028015, 0.17919236421585083, 0.1657918244600296, 0.1515847146511078, 0.11915461719036102, 0.10438721626996994, 0.10422837734222412, 0.08969437330961227, 0.07429874688386917, 0.07801567018032074, 0.06531910598278046, 0.05813637375831604, 0.04699350893497467, 0.05086237192153931, 0.060301560908555984, 0.04986414313316345, 0.050366174429655075, 0.05464963987469673, 0.05319518595933914, 0.04274186119437218, 0.047863349318504333, 0.036163944751024246, 0.03777360916137695, 0.04280579090118408, 0.04440606012940407, 0.04717888683080673, 0.02609632909297943, 0.05858827754855156, 0.050790246576070786, 0.03004802018404007, 0.04584358632564545, 0.05146845430135727, 0.039567168802022934, 0.03470978885889053, 0.045542243868112564, 0.05142106115818024, 0.05224252864718437, 0.07936417311429977, 0.11145134270191193, 0.14930342137813568, 0.21797531843185425

Flexpert-Design/predict_example/compare_seqs.py

@@ -0,0 +1,59 @@
+# read in the predictions.txt file
+# take the sequence from there
+import argparse
+import os
+import biotite.structure.io.pdb as pdb
+from biotite.structure import get_residues
+
+def compare_sequences(pdb_code):
+    # Read the predicted sequence from predictions.txt
+    with open('predict_example/predictions.txt', 'r') as f:
+        predictions = f.readlines()
+        # Extract the sequence (skip the header line that starts with '>')
+        predicted_seqs = {}
+        current_pdb = None
+        
+        for line in predictions:
+            if line.startswith('>'):
+                current_pdb = line.strip()[1:]  # Remove the '>' character
+            elif current_pdb and line.strip():
+                predicted_seqs[current_pdb] = line.strip()
+        
+        # Use the provided pdb_code to get the corresponding sequence
+        predicted_seq = predicted_seqs.get(pdb_code, "")
+
+    # Read the PDB file
+    pdb_file = f'predict_example/{pdb_code}.pdb'
+    with open(pdb_file, 'r') as f:
+        structure = pdb.PDBFile.read(f)
+        atoms = pdb.get_structure(structure)
+        
+        # Get residue names from the structure
+        residues = get_residues(atoms)[1]
+        # Convert three-letter codes to one-letter codes
+        aa_dict = {
+            'ALA': 'A', 'CYS': 'C', 'ASP': 'D', 'GLU': 'E', 'PHE': 'F',
+            'GLY': 'G', 'HIS': 'H', 'ILE': 'I', 'LYS': 'K', 'LEU': 'L',
+            'MET': 'M', 'ASN': 'N', 'PRO': 'P', 'GLN': 'Q', 'ARG': 'R',
+            'SER': 'S', 'THR': 'T', 'VAL': 'V', 'TRP': 'W', 'TYR': 'Y'
+        }
+        pdb_seq = ''.join([aa_dict.get(res, 'X') for res in residues])
+
+    # Compare the two sequences
+    match_count = sum(1 for a, b in zip(predicted_seq, pdb_seq) if a == b)
+    total_length = max(len(predicted_seq), len(pdb_seq))
+    percent_identity = (match_count / min(len(predicted_seq), len(pdb_seq))) * 100
+
+    # Print the result
+    print(f"Predicted sequence: {predicted_seq}")
+    print(f"PDB sequence:       {pdb_seq}")
+    print(f"Sequence length - Predicted: {len(predicted_seq)}, PDB: {len(pdb_seq)}")
+    print(f"Matching residues: {match_count}/{min(len(predicted_seq), len(pdb_seq))}")
+    print(f"Percent identity: {percent_identity:.2f}%")
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='Compare predicted sequence with PDB sequence')
+    parser.add_argument('--pdb_code', type=str, help='PDB code (e.g., 1ah7_A)')
+    args = parser.parse_args()
+    
+    compare_sequences(args.pdb_code)

Flexpert-Design/predict_example/predictions.txt

@@ -0,0 +1,2 @@
+>1ah7_A
+GSSLDKTEVEESTGLRLVNQAIDILKNDKTRVDKEYLDLIEKYKPELQEGIYKAYHSEPYNDNGKFSRHYYNPVVHTSRIPDAVTAAETGSHYYNKAGEYYKKGDYEEAYFYLGIALAYLSDACNPMNASGYTNESFPEGFYEALQKYVCTIAKKYENTTGEPYYNLTGKNPKDHIRGAATKARELFSGIYHERVKEDFEKGKTSEEARLKWRERIEPQLGKLLLFAQRVMAGAIERFFDTAGGL

Flexpert-Design/requirements.txt

@@ -0,0 +1,23 @@
+Bio
+biotite
+fair-esm
+evaluate
+joblib
+matplotlib
+numpy
+omegaconf
+packaging
+pandas
+python_dateutil
+pytorch_lightning
+PyYAML
+requests
+safetensors
+scikit_learn
+scipy
+torch
+torch_geometric
+torcheval
+torchmetrics
+tqdm
+transformers

Flexpert-Design/src/__init__.py

@@ -0,0 +1,49 @@
+# Copyright (c) CAIRI AI Lab. All rights reserved
+
+import warnings
+from packaging.version import parse
+
+from .version import __version__
+
+
+def digit_version(version_str: str, length: int = 4):
+    """Convert a version string into a tuple of integers.
+
+    This method is usually used for comparing two versions. For pre-release
+    versions: alpha < beta < rc.
+
+    Args:
+        version_str (str): The version string.
+        length (int): The maximum number of version levels. Default: 4.
+
+    Returns:
+        tuple[int]: The version info in digits (integers).
+    """
+    version = parse(version_str)
+    assert version.release, f'failed to parse version {version_str}'
+    release = list(version.release)
+    release = release[:length]
+    if len(release) < length:
+        release = release + [0] * (length - len(release))
+    if version.is_prerelease:
+        mapping = {'a': -3, 'b': -2, 'rc': -1}
+        val = -4
+        # version.pre can be None
+        if version.pre:
+            if version.pre[0] not in mapping:
+                warnings.warn(f'unknown prerelease version {version.pre[0]}, '
+                              'version checking may go wrong')
+            else:
+                val = mapping[version.pre[0]]
+            release.extend([val, version.pre[-1]])
+        else:
+            release.extend([val, 0])
+
+    elif version.is_postrelease:
+        release.extend([1, version.post])
+    else:
+        release.extend([0, 0])
+    return tuple(release)
+
+
+__all__ = ['__version__', 'digit_version']

Flexpert-Design/src/datasets/__init__.py

@@ -0,0 +1,15 @@
+# Copyright (c) CAIRI AI Lab. All rights reserved
+
+from .alphafold_dataset import AlphaFoldDataset
+from .cath_dataset import CATHDataset
+from .dataloader import load_data
+from .featurizer import (featurize_AF, featurize_GTrans, featurize_GVP,
+                         featurize_ProteinMPNN, featurize_Inversefolding)
+from .ts_dataset import TSDataset
+
+__all__ = [
+    'AlphaFoldDataset', 'CATHDataset', 'TSDataset',
+    'load_data',
+    'featurize_AF', 'featurize_GTrans', 'featurize_GVP',
+    'featurize_ProteinMPNN', 'featurize_Inversefolding'
+]

Flexpert-Design/src/datasets/alphafold_dataset.py

@@ -0,0 +1,112 @@
+import os
+import os.path as osp
+import json
+import numpy as np
+import pickle as cPickle
+
+import torch.utils.data as data
+from src.datasets.utils import cached_property
+
+
+class AlphaFoldDataset(data.Dataset):
+    def __init__(self, path='./', upid='', mode='train', max_length=500, limit_length=1, joint_data=0):
+        
+        self.path = path
+        self.upid = upid
+        self.max_length = max_length
+        self.limit_length = limit_length
+        self.joint_data = joint_data
+        
+        if mode in ['train', 'valid', 'test']:
+            self.data = self.cache_data[mode]
+        
+        if mode == 'all':
+            self.data = self.cache_data['train'] + self.cache_data['valid'] + self.cache_data['test']
+        
+        self.lengths = np.array([ len(sample['seq']) for sample in self.data])
+        self.max_len = np.max(self.lengths)
+        self.min_len = np.min(self.lengths)
+
+    def _raw_data(self, path, upid):
+        if not os.path.exists(path):
+            raise "no such file:{} !!!".format(path)
+        else:
+            path = osp.join(path, upid)
+            data_ = cPickle.load(open(path+'/data_{}.pkl'.format(upid),'rb'))
+            score_ = cPickle.load(open(path+'/data_{}_score.pkl'.format(upid),'rb'))
+            for i in range(len(data_)):
+                data_[i]['score'] = score_[i]['res_score']
+        return data_
+
+    def _data_info(self, data):
+        len_inds = []
+        seq2ind = {}
+        for ind, temp in enumerate(data):
+            if self.limit_length:
+                if 30 < len(temp['seq']) and len(temp['seq']) < self.max_length:
+                    # 'title', 'seq', 'CA', 'C', 'O', 'N'
+                    len_inds.append(ind)
+                    seq2ind[temp['seq']] = ind
+            else:
+                len_inds.append(ind)
+                seq2ind[temp['seq']] = ind
+        return len_inds, seq2ind
+        
+    def get_data(self, path, upid, **kwargs):
+        data_ = self._raw_data(path, upid)
+        path = osp.join(path, upid)
+
+        file_name = 'split_clu_l.json' if self.limit_length else 'split_clu.json'
+
+        assert os.path.exists(osp.join(path, file_name))
+        split = json.load(open(osp.join(path, file_name),'r'))
+        data_dict = {'train':[data_[i] for i in split['train']],
+                     'valid':[data_[i] for i in split['valid']],
+                     'test':[data_[i] for i in split['test']]}
+        return data_dict
+
+    def get_full_data(self, path, **kwargs):
+        datanames = [dataname for dataname in os.listdir(path) if ('_v2' in dataname)]
+        file_name = 'split_clu_l.json' if self.limit_length else 'split_clu.json'
+        assert os.path.exists(osp.join(path, 'full', file_name))
+        split = json.load(open(osp.join(path, 'full', file_name),'r'))
+        return split 
+    
+    @cached_property
+    def cache_data(self): # TODO: joint_data
+        path = self.path
+        upid = self.upid
+        if self.joint_data:
+            datanames = [dataname for dataname in os.listdir(path) if ('_v2' in dataname)]
+            data_dict = {'train':[], 'valid':[], 'test':[]}
+            full_inds = self.get_full_data(path)
+
+            for dataname in datanames:
+                temp = self._raw_data(path, dataname)
+                train_idx, valid_idx, test_idx = map(lambda fold: full_inds[dataname][fold], ['train', 'valid', 'test'])
+                data_dict['train'] += [temp[i] for i in train_idx]
+                data_dict['valid'] += [temp[i] for i in valid_idx]
+
+                data_test = []
+                for i in test_idx:
+                    item = temp[i]
+                    item['category'] = dataname
+                    data_test.append(temp[i])
+                
+                data_dict['test'] += data_test
+
+        else:
+            data_dict = self.get_data(path, upid)
+            for item in data_dict['test']:
+                item['category'] = upid
+
+        return data_dict
+    
+    def change_mode(self, mode):
+        self.data = self.cache_data[mode]
+    
+    def __len__(self):
+        return len(self.data)
+
+    def __getitem__(self, index):
+        return self.data[index]

Flexpert-Design/src/datasets/atlas_dataset.py

@@ -0,0 +1,133 @@
+import os
+import json
+import numpy as np
+from tqdm import tqdm
+import random
+import pdb
+import torch.utils.data as data
+from .utils import cached_property
+from transformers import AutoTokenizer
+
+class AtlasDataset(data.Dataset):
+    def __init__(self, path='./',  split='train', max_length=500, test_name='All', data = None, removeTS=0):
+        self.path = path
+        self.mode = split
+        self.max_length = max_length
+        self.test_name = test_name
+        self.removeTS = removeTS
+        if self.removeTS:
+            self.remove = json.load(open(self.path+'/remove.json', 'r'))['remove']
+        
+        if data is None:
+            if self.mode in ['eval','predict']:
+                self.data = self.cache_data['test'] #This calls the cache_data property
+            else:
+                self.data = self.cache_data[split] #This calls the cache_data property
+        else:
+            self.data = data
+        
+        self.tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D", cache_dir="./cache_dir/")
+    
+    @cached_property
+    def cache_data(self):
+        alphabet='ACDEFGHIKLMNPQRSTVWY'
+        alphabet_set = set([a for a in alphabet])
+        print("path is: ", self.path)
+
+        if not os.path.exists(self.path):
+            raise "no such file:{} !!!".format(self.path)
+        else:
+
+            with open(self.path+'/chain_set.jsonl') as f:
+                lines = f.readlines()
+            data_list = []
+            
+            for line in tqdm(lines):
+                entry = json.loads(line)
+
+                if self.removeTS and entry['name'] in self.remove:
+                    continue
+                seq = entry['seq']
+
+                for key, val in entry['coords'].items():
+                    entry['coords'][key] = np.asarray(val)
+                
+                bad_chars = set([s for s in seq]).difference(alphabet_set)
+
+                if len(bad_chars) == 0:
+                    if len(entry['seq']) <= self.max_length: 
+                        chain_length = len(entry['seq'])
+                        chain_mask = np.ones(chain_length)
+                        data_list.append({
+                            'title':entry['name'],
+                            'seq':entry['seq'],
+                            'CA':entry['coords']['CA'],
+                            'C':entry['coords']['C'],
+                            'O':entry['coords']['O'],
+                            'N':entry['coords']['N'],
+                            'chain_mask': chain_mask,
+                            'chain_encoding': 1*chain_mask
+                        })
+            
+            with open(self.path+'/chain_set_splits.json') as f:
+                dataset_splits = json.load(f)
+            
+            if self.test_name == 'L100':
+                with open(self.path+'/test_split_L100.json') as f:
+                    test_splits = json.load(f)
+                dataset_splits['test'] = test_splits['test']
+
+            if self.test_name == 'sc':
+                with open(self.path+'/test_split_sc.json') as f:
+                    test_splits = json.load(f)
+                dataset_splits['test'] = test_splits['test']
+            
+            name2set = {}
+            name2set.update({name:'train' for name in dataset_splits['train']})
+            name2set.update({name:'valid' for name in dataset_splits['validation']})
+            name2set.update({name:'test' for name in dataset_splits['test']})
+
+            data_dict = {'train':[],'valid':[],'test':[]}
+            for data in data_list:
+                #raise ValueError("only 1015 sequences get loaded to the predict set!!! why not whole 1390??? FIX!")
+                if name2set.get(data['title']): #This was causing the trouble with empty datasets - missmatch of names in the chain_set and chain_set_split
+                    if name2set[data['title']] == 'train':
+                        data_dict['train'].append(data)
+                    
+                    if name2set[data['title']] == 'valid':
+                        data_dict['valid'].append(data)
+                    
+                    if name2set[data['title']] == 'test':
+                        data['category'] = 'Unkown'
+                        data['score'] = 100.0
+                        data_dict['test'].append(data)
+                else:
+                    import pdb; pdb.set_trace()
+            return data_dict
+
+    def change_mode(self, mode):
+        self.data = self.cache_data[mode]
+    
+    def __len__(self):
+        return len(self.data)
+    
+    def get_item(self, index):
+        return self.data[index]
+    
+    def __getitem__(self, index):
+        item = self.data[index]
+        L = len(item['seq'])
+        if L>self.max_length:
+            # 计算截断的最大索引
+            max_index = L - self.max_length
+            # 生成随机的截断索引
+            truncate_index = random.randint(0, max_index)
+            # 进行截断
+            item['seq'] = item['seq'][truncate_index:truncate_index+self.max_length]
+            item['CA'] = item['CA'][truncate_index:truncate_index+self.max_length]
+            item['C'] = item['C'][truncate_index:truncate_index+self.max_length]
+            item['O'] = item['O'][truncate_index:truncate_index+self.max_length]
+            item['N'] = item['N'][truncate_index:truncate_index+self.max_length]
+            item['chain_mask'] = item['chain_mask'][truncate_index:truncate_index+self.max_length]
+            item['chain_encoding'] = item['chain_encoding'][truncate_index:truncate_index+self.max_length]
+        return item

Flexpert-Design/src/datasets/casp_dataset.py

@@ -0,0 +1,57 @@
+import os
+import json
+import numpy as np
+import torch.utils.data as data
+
+
+class CASPDataset(data.Dataset):
+    def __init__(self, path = './', split='test'):
+        if not os.path.exists(path):
+            raise "no such file:{} !!!".format(path)
+        else:
+            with open(os.path.join(path,'casp15.jsonl')) as f:
+                lines = f.readlines()
+                
+            # casp15_data = json.load(open(path+'casp15.json', 'r'))
+
+        alphabet='ACDEFGHIKLMNPQRSTVWY'
+        alphabet_set = set([a for a in alphabet])
+        
+        self.data = []
+        for line in lines:
+            entry = json.loads(line)
+            seq = entry['seq']
+
+            for key, val in entry['coords'].items():
+                entry['coords'][key] = np.asarray(val)
+            
+            bad_chars = set([s for s in seq]).difference(alphabet_set)
+
+            if len(bad_chars) == 0:
+                chain_length = len(entry['seq'])
+                chain_mask = np.ones(chain_length)
+                self.data.append({
+                    'title':entry['name'],
+                    'seq':entry['seq'],
+                    'CA':entry['coords']['CA'],
+                    'C':entry['coords']['C'],
+                    'O':entry['coords']['O'],
+                    'N':entry['coords']['N'],
+                    'chain_mask': chain_mask,
+                    'chain_encoding': 1*chain_mask,
+                    'classification': entry['classification']
+                })
+
+    def __len__(self):
+        return len(self.data)
+    
+    def get_item(self, index):
+        return self.data[index]
+
+    def __getitem__(self, index):
+        return self.data[index]
+    
+if __name__ == '__main__':
+    dataset = CASPDataset('/gaozhangyang/experiments/OpenCPD/data/casp15/')
+    for data in dataset:
+        print(data)

Flexpert-Design/src/datasets/cath_dataset.py

@@ -0,0 +1,141 @@
+import os
+import json
+import numpy as np
+from tqdm import tqdm
+import random
+import torch.utils.data as data
+from .utils import cached_property
+from transformers import AutoTokenizer
+from src.tools.utils import load_yaml_config
+
+class CATHDataset(data.Dataset):
+    def __init__(self, path='./',  split='train', max_length=500, test_name='All', data = None, removeTS=0, version=4.2, data_jsonl_name='/chain_set.jsonl'):
+        self.version = version
+        self.path = path
+        self.mode = split
+        self.max_length = max_length
+        self.test_name = test_name
+        self.removeTS = removeTS
+        self.data_jsonl_name = data_jsonl_name
+
+        self.using_dynamics = data_jsonl_name == load_yaml_config('/scratch/project/fta-24-31/koubapet/ProteinInvBench/src/models/configs/FlexibilityProtTrans.yaml')['data_jsonl_name']
+
+        print(self.data_jsonl_name)
+        if self.removeTS:
+            self.remove = json.load(open(self.path+'/remove.json', 'r'))['remove']
+        
+        if data is None:
+            if split == 'predict':
+                _split = 'valid'
+                print('In predict mode for CATH4.3 using VALIDATION split as the data. Consider switching to TEST set.')
+            else:
+                _split = split
+            self.data = self.cache_data[_split]
+        else:
+            self.data = data
+        
+        self.tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D", cache_dir="./cache_dir/")
+    
+    @cached_property
+    def cache_data(self):
+        alphabet='ACDEFGHIKLMNPQRSTVWY'
+        alphabet_set = set([a for a in alphabet])
+        print("path is: ", self.path)
+        if not os.path.exists(self.path):
+            raise "no such file:{} !!!".format(self.path)
+        else:
+            with open(self.path+'/'+self.data_jsonl_name) as f:
+                lines = f.readlines()
+            data_list = []
+            for line in tqdm(lines):
+                entry = json.loads(line)
+                if self.removeTS and entry['name'] in self.remove:
+                    continue
+                seq = entry['seq']
+
+                for key, val in entry['coords'].items():
+                    entry['coords'][key] = np.asarray(val)
+                
+                bad_chars = set([s for s in seq]).difference(alphabet_set)
+
+                if len(bad_chars) == 0:
+                    if len(entry['seq']) <= self.max_length: 
+                        chain_length = len(entry['seq'])
+                        chain_mask = np.ones(chain_length)
+                        data_list.append({
+                            'title':entry['name'],
+                            'seq':entry['seq'],
+                            'CA':entry['coords']['CA'],
+                            'C':entry['coords']['C'],
+                            'O':entry['coords']['O'],
+                            'N':entry['coords']['N'],
+                            'chain_mask': chain_mask,
+                            'chain_encoding': 1*chain_mask
+                        })
+                        if self.using_dynamics: #TODO: pass this bool properly
+                            data_list[-1]['norm_bfactors'] = entry['bfactor']
+                        
+            if self.version==4.2:
+                with open(self.path+'/chain_set_splits.json') as f:
+                    dataset_splits = json.load(f)
+            
+            if self.version==4.3:
+                with open(self.path+'/chain_set_splits.json') as f:
+                    dataset_splits = json.load(f)
+            
+            if self.test_name == 'L100':
+                with open(self.path+'/test_split_L100.json') as f:
+                    test_splits = json.load(f)
+                dataset_splits['test'] = test_splits['test']
+
+            if self.test_name == 'sc':
+                with open(self.path+'/test_split_sc.json') as f:
+                    test_splits = json.load(f)
+                dataset_splits['test'] = test_splits['test']
+            
+            name2set = {}
+            name2set.update({name:'train' for name in dataset_splits['train']})
+            name2set.update({name:'valid' for name in dataset_splits['validation']})
+            name2set.update({name:'test' for name in dataset_splits['test']})
+
+            data_dict = {'train':[],'valid':[],'test':[]}
+            for data in data_list:
+                if name2set.get(data['title']):
+                    if name2set[data['title']] == 'train':
+                        data_dict['train'].append(data)
+                    
+                    if name2set[data['title']] == 'valid':
+                        data_dict['valid'].append(data)
+                    
+                    if name2set[data['title']] == 'test':
+                        data['category'] = 'Unkown'
+                        data['score'] = 100.0
+                        data_dict['test'].append(data)
+            return data_dict
+
+    def change_mode(self, mode):
+        self.data = self.cache_data[mode]
+    
+    def __len__(self):
+        return len(self.data)
+    
+    def get_item(self, index):
+        return self.data[index]
+    
+    def __getitem__(self, index):
+        item = self.data[index]
+        L = len(item['seq'])
+        if L>self.max_length:
+            # 计算截断的最大索引
+            max_index = L - self.max_length
+            # 生成随机的截断索引
+            truncate_index = random.randint(0, max_index)
+            # 进行截断
+            item['seq'] = item['seq'][truncate_index:truncate_index+self.max_length]
+            item['CA'] = item['CA'][truncate_index:truncate_index+self.max_length]
+            item['C'] = item['C'][truncate_index:truncate_index+self.max_length]
+            item['O'] = item['O'][truncate_index:truncate_index+self.max_length]
+            item['N'] = item['N'][truncate_index:truncate_index+self.max_length]
+            item['chain_mask'] = item['chain_mask'][truncate_index:truncate_index+self.max_length]
+            item['chain_encoding'] = item['chain_encoding'][truncate_index:truncate_index+self.max_length]
+        return item

Flexpert-Design/src/datasets/dataloader.py

@@ -0,0 +1,161 @@
+import copy
+import random
+import os.path as osp
+
+import torch
+import torch.utils.data as data
+import pdb
+
+from .cath_dataset import CATHDataset
+from .alphafold_dataset import AlphaFoldDataset
+from .ts_dataset import TSDataset
+from .casp_dataset import CASPDataset
+from .mpnn_dataset import MPNNDataset
+from .featurizer import (featurize_AF, featurize_GTrans, featurize_GVP,
+                         featurize_ProteinMPNN, featurize_Inversefolding)
+from .fast_dataloader import DataLoaderX
+
+class GTransDataLoader(torch.utils.data.DataLoader):
+    def __init__(self, dataset, batch_size=1, shuffle=False, sampler=None, batch_sampler=None, num_workers=0,
+                 collate_fn=None, **kwargs):
+        super(GTransDataLoader, self).__init__(dataset, batch_size, shuffle, sampler, batch_sampler, num_workers, collate_fn,**kwargs)
+        self.featurizer = collate_fn
+
+
+class BatchSampler(data.Sampler):
+    '''
+    From https://github.com/jingraham/neurips19-graph-protein-design.
+    
+    A `torch.utils.data.Sampler` which samples batches according to a
+    maximum number of graph nodes.
+    
+    :param node_counts: array of node counts in the dataset to sample from
+    :param max_nodes: the maximum number of nodes in any batch,
+                      including batches of a single element
+    :param shuffle: if `True`, batches in shuffled order
+    '''
+    def __init__(self, node_counts, max_nodes=3000, shuffle=True):
+        self.node_counts = node_counts
+        self.idx = [i for i in range(len(node_counts))  
+                        if node_counts[i] <= max_nodes]
+        self.shuffle = shuffle
+        self.max_nodes = max_nodes
+        self._form_batches()
+    
+    def _form_batches(self):
+        self.batches = []
+        if self.shuffle: random.shuffle(self.idx)
+        idx = self.idx
+        while idx:
+            batch = []
+            n_nodes = 0
+            while idx and n_nodes + self.node_counts[idx[0]] <= self.max_nodes:
+                next_idx, idx = idx[0], idx[1:]
+                n_nodes += self.node_counts[next_idx]
+                batch.append(next_idx)
+            self.batches.append(batch)
+    
+    def __len__(self): 
+        if not self.batches: self._form_batches()
+        return len(self.batches)
+    
+    def __iter__(self):
+        if not self.batches: 
+            self._form_batches()
+        for batch in self.batches: 
+            yield batch
+
+
+class GVPDataLoader(torch.utils.data.DataLoader):
+    def __init__(self, dataset, num_workers=0,
+                 featurizer=None, max_nodes=3000, **kwargs):
+        super(GVPDataLoader, self).__init__(dataset, 
+                                            batch_sampler = BatchSampler(node_counts = [ len(data['seq']) for data in dataset], max_nodes=max_nodes), 
+                                            num_workers = num_workers, 
+                                            collate_fn = featurizer.collate,
+                                            **kwargs)
+        self.featurizer = featurizer
+
+
+def load_data(data_name, method, batch_size, data_root, pdb_path, split_csv, max_nodes=3000, num_workers=8, removeTS=0, test_casp=False, **kwargs):
+    if data_name == 'CATH4.2' or data_name == 'TS':
+        cath_set = CATHDataset(osp.join(data_root, 'cath4.2'), mode='train', test_name='All', removeTS=removeTS)
+        train_set, valid_set, test_set = map(lambda x: copy.copy(x), [cath_set] * 3)
+        valid_set.change_mode('valid')
+        test_set.change_mode('test')
+        if data_name == 'TS':
+            test_set = TSDataset(osp.join(data_root, 'ts'))
+            
+        collate_fn = featurize_GTrans
+    elif data_name == 'CATH4.3':
+        cath_set = CATHDataset(osp.join(data_root, 'cath4.3'), mode='train', test_name='All', removeTS=removeTS, version=4.3)
+        train_set, valid_set, test_set = map(lambda x: copy.copy(x), [cath_set] * 3)
+        valid_set.change_mode('valid')
+        test_set.change_mode('test')
+        
+        collate_fn = featurize_GTrans
+    elif data_name == 'AlphaFold':
+        af_set = AlphaFoldDataset(osp.join(data_root, 'af2db'), upid=upid, mode='train', limit_length=limit_length, joint_data=joint_data)
+        train_set, valid_set, test_set = map(lambda x: copy.copy(x), [af_set] * 3)
+        valid_set.change_mode('valid')
+        test_set.change_mode('test')
+        collate_fn = featurize_AF
+    elif data_name=='MPNN':
+        train_set = MPNNDataset(mode='train')
+        valid_set = MPNNDataset(mode='valid')
+        test_set = MPNNDataset(mode='test')
+        collate_fn = featurize_GTrans
+    
+    elif data_name == 'S350':
+        cath_set = CATHDataset(osp.join(data_root, 's350'), mode='train', test_name='All', removeTS=removeTS, version=4.3)
+        train_set, valid_set, test_set = map(lambda x: copy.copy(x), [cath_set] * 3)
+        valid_set.change_mode('train')
+        test_set.change_mode('train')
+        
+        collate_fn = featurize_GTrans
+        
+    elif data_name == 'Protherm':
+        cath_set = CATHDataset(osp.join(data_root, 'protherm'), mode='train', test_name='All', removeTS=removeTS, version=4.3)
+        train_set, valid_set, test_set = map(lambda x: copy.copy(x), [cath_set] * 3)
+        valid_set.change_mode('valid')
+        test_set.change_mode('test')
+        
+        collate_fn = featurize_GTrans
+    if test_casp:
+        test_set = CASPDataset(osp.join(data_root, 'casp15'))
+
+    if method in ['AlphaDesign', 'PiFold', 'KWDesign', 'GraphTrans', 'StructGNN']:
+        pass
+    elif method == 'GVP':
+        featurizer = featurize_GVP()
+        collate_fn = featurizer.collate
+    elif method == 'ProteinMPNN':
+        collate_fn = featurize_ProteinMPNN
+    elif method == 'ESMIF':
+        collate_fn = featurize_Inversefolding
+        
+    # train_set.data = train_set.data[:100]
+    # valid_set.data = valid_set.data[:100]
+    # test_set.data = test_set.data[:100]
+    pdb.set_trace()
+    train_loader = DataLoaderX(local_rank=0, dataset=train_set, batch_size=batch_size, shuffle=True, num_workers=num_workers, collate_fn=collate_fn, prefetch_factor=8)
+    valid_loader = DataLoaderX(local_rank=0,dataset=valid_set, batch_size=batch_size, shuffle=False, num_workers=num_workers, collate_fn=collate_fn, prefetch_factor=8)
+    test_loader = DataLoaderX(local_rank=0,dataset=test_set, batch_size=batch_size, shuffle=False, num_workers=num_workers, collate_fn=collate_fn, prefetch_factor=8)
+        
+    return train_loader, valid_loader, test_loader
+
+
+def make_cath_loader(test_set, method, batch_size, max_nodes=3000, num_workers=8):
+    if method in ['pifold','adesign', 'graphtrans', 'structgnn', 'gca']:
+        collate_fn = featurize_GTrans
+        test_loader = GTransDataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=num_workers, collate_fn=collate_fn)
+    elif method == 'gvp':
+        featurizer = featurize_GVP()
+        test_loader = GVPDataLoader(test_set, num_workers=num_workers, featurizer=featurizer, max_nodes=max_nodes)
+    elif method == 'proteinmpnn':
+        collate_fn = featurize_ProteinMPNN
+        test_loader = GTransDataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=num_workers, collate_fn=collate_fn)
+    elif method == 'esmif':
+        collate_fn = featurize_Inversefolding
+        test_loader = GTransDataLoader(test_set, batch_size=batch_size, shuffle=False, num_workers=num_workers, collate_fn=collate_fn)
+    return test_loader

Flexpert-Design/src/datasets/fast_dataloader.py

@@ -0,0 +1,52 @@
+import threading
+import torch
+import queue
+from torch.utils.data import DataLoader
+
+
+class DataLoaderX(DataLoader):
+    def __init__(self, local_rank, **kwargs):
+        super().__init__(**kwargs)
+        self.stream = torch.cuda.Stream(
+            local_rank
+        )  # create a new cuda stream in each process
+        self.local_rank = local_rank
+        # self.custom_collect_fn = custom_collect_fn
+        
+    def __iter__(self):
+        self.iter = super().__iter__()
+        self.preload()
+        return self
+
+
+    def preload(self):
+        while True:
+            #获取下一个值
+            self.batch = next(self.iter, None)
+            if self.batch is not None:
+                break
+            if self.iter._send_idx==len(self.iter):
+                break
+        
+        if (self.batch is None):
+            return None
+
+        with torch.cuda.stream(self.stream): # 将数据预先放进gpu
+            for key, val in self.batch.items():
+                if type(val) == torch.Tensor:
+                    self.batch[key] = val.to(
+                            device=self.local_rank, non_blocking=True
+                        )
+
+
+    def __next__(self):
+        torch.cuda.current_stream().wait_stream(
+            self.stream
+        )  # wait tensor to put on GPU
+        batch = self.batch
+        # batch = self.custom_collect_fn(self.batch)
+        if batch is None:
+            raise StopIteration
+        self.preload()
+        return batch
+

Flexpert-Design/src/datasets/featurizer.py

@@ -0,0 +1,743 @@
+import torch
+import numpy as np
+import itertools
+import torch.nn.functional as F
+import math
+import torch_geometric
+# import torch_cluster
+from collections.abc import Mapping, Sequence
+from torch_geometric.data import Data, Batch
+from torch.utils.data.dataloader import default_collate
+from transformers import AutoTokenizer
+import pdb
+tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D", cache_dir='./cache_dir/') # mask token: 32
+
+def _normalize(tensor, dim=-1):
+    '''
+    Normalizes a `torch.Tensor` along dimension `dim` without `nan`s.
+    '''
+    return torch.nan_to_num(
+        torch.div(tensor, torch.norm(tensor, dim=dim, keepdim=True)))
+
+
+def _rbf(D, D_min=0., D_max=20., D_count=16, device='cpu'):
+    '''
+    From https://github.com/jingraham/neurips19-graph-protein-design
+    
+    Returns an RBF embedding of `torch.Tensor` `D` along a new axis=-1.
+    That is, if `D` has shape [...dims], then the returned tensor will have
+    shape [...dims, D_count].
+    '''
+    D_mu = torch.linspace(D_min, D_max, D_count, device=device)
+    D_mu = D_mu.view([1, -1])
+    D_sigma = (D_max - D_min) / D_count
+    D_expand = torch.unsqueeze(D, -1)
+
+    RBF = torch.exp(-((D_expand - D_mu) / D_sigma) ** 2)
+    return RBF
+
+def shuffle_subset(n, p):
+    n_shuffle = np.random.binomial(n, p)
+    ix = np.arange(n)
+    ix_subset = np.random.choice(ix, size=n_shuffle, replace=False)
+    ix_subset_shuffled = np.copy(ix_subset)
+    np.random.shuffle(ix_subset_shuffled)
+    ix[ix_subset] = ix_subset_shuffled
+    return ix
+
+
+def featurize_AF(batch, shuffle_fraction=0.):
+    """ Pack and pad batch into torch tensors """
+    alphabet = 'ACDEFGHIKLMNPQRSTVWY'
+    B = len(batch)
+    lengths = np.array([len(b['seq']) for b in batch], dtype=np.int32)
+    L_max = max([len(b['seq']) for b in batch])
+    X = np.zeros([B, L_max, 4, 3])
+    S = np.zeros([B, L_max], dtype=np.int32)
+    score = np.zeros([B, L_max])
+
+    # Build the batch
+    for i, b in enumerate(batch):
+        x = np.stack([b[c] for c in ['N', 'CA', 'C', 'O']], 1) # [#atom, 4, 3]
+        
+        l = len(b['seq'])
+        x_pad = np.pad(x, [[0,L_max-l], [0,0], [0,0]], 'constant', constant_values=(np.nan, )) # [#atom, 4, 3]
+        X[i,:,:,:] = x_pad
+
+        # Convert to labels
+        indices = np.asarray([alphabet.index(a) for a in b['seq']], dtype=np.int32)
+        if shuffle_fraction > 0.:
+            idx_shuffle = shuffle_subset(l, shuffle_fraction)
+            S[i, :l] = indices[idx_shuffle]
+            score[i,:l] = b['score'][idx_shuffle]
+        else:
+            S[i, :l] = indices
+            score[i,:l] = b['score']
+
+    mask = np.isfinite(np.sum(X,(2,3))).astype(np.float32) # atom mask
+    numbers = np.sum(mask, axis=1).astype(np.int)
+    S_new = np.zeros_like(S)
+    score_new = np.zeros_like(score)
+    X_new = np.zeros_like(X)+np.nan
+    for i, n in enumerate(numbers):
+        X_new[i,:n,::] = X[i][mask[i]==1]
+        S_new[i,:n] = S[i][mask[i]==1]
+        score_new[i,:n] = score[i][mask[i]==1]
+
+    X = X_new
+    S = S_new
+    score = score_new
+    isnan = np.isnan(X)
+    mask = np.isfinite(np.sum(X,(2,3))).astype(np.float32)
+    X[isnan] = 0.
+    # Conversion
+    S = torch.from_numpy(S).to(dtype=torch.long)
+    score = torch.from_numpy(score).float()
+    X = torch.from_numpy(X).to(dtype=torch.float32)
+    mask = torch.from_numpy(mask).to(dtype=torch.float32)
+    return X, S, score, mask, lengths
+
+
+def featurize_GTrans(batch):
+    """ Pack and pad batch into torch tensors """
+    # alphabet = 'ACDEFGHIKLMNPQRSTVWYX'
+    batch = [one for one in batch if one is not None]
+    B = len(batch)
+    if B==0:
+        return None
+    lengths = np.array([len(b['seq']) for b in batch], dtype=np.int32)
+    L_max = max([len(b['seq']) for b in batch])
+    X = np.zeros([B, L_max, 4, 3])
+    S = np.zeros([B, L_max], dtype=np.int32)
+    score = np.ones([B, L_max]) * 100.0
+    chain_mask = np.zeros([B, L_max])-1 # 1:需要被预测的掩码部分 0:可见部分
+    chain_encoding = np.zeros([B, L_max])-1
+    
+
+    # Build the batch
+    for i, b in enumerate(batch):
+        x = np.stack([b[c] for c in ['N', 'CA', 'C', 'O']], 1) # [#atom, 4, 3]
+        
+        l = len(b['seq'])
+        x_pad = np.pad(x, [[0,L_max-l], [0,0], [0,0]], 'constant', constant_values=(np.nan, )) # [#atom, 4, 3]
+        X[i,:,:,:] = x_pad
+
+        # Convert to labels
+        indices = np.array(tokenizer.encode(b['seq'], add_special_tokens=False))
+        # indices = np.asarray([alphabet.index(a) for a in b['seq']], dtype=np.int32)
+
+
+        S[i, :l] = indices
+        chain_mask[i,:l] = b['chain_mask']
+        chain_encoding[i,:l] = b['chain_encoding']
+
+    mask = np.isfinite(np.sum(X,(2,3))).astype(np.float32) # atom mask
+    numbers = np.sum(mask, axis=1).astype(np.int32)
+    S_new = np.zeros_like(S)
+    X_new = np.zeros_like(X)+np.nan
+    for i, n in enumerate(numbers):
+        X_new[i,:n,::] = X[i][mask[i]==1]
+        S_new[i,:n] = S[i][mask[i]==1]
+
+    X = X_new
+    S = S_new
+    isnan = np.isnan(X)
+    mask = np.isfinite(np.sum(X,(2,3))).astype(np.float32)
+    X[isnan] = 0.
+    # Conversion
+    S = torch.from_numpy(S).to(dtype=torch.long)
+    score = torch.from_numpy(score).float()
+    X = torch.from_numpy(X).to(dtype=torch.float32)
+    mask = torch.from_numpy(mask).to(dtype=torch.float32)
+    lengths = torch.from_numpy(lengths)
+    chain_mask = torch.from_numpy(chain_mask)
+    chain_encoding = torch.from_numpy(chain_encoding)
+    
+    return {"title": [b['title'] for b in batch],
+            "X":X,
+            "S":S,
+            "score": score,
+            "mask":mask,
+            "lengths":lengths,
+            "chain_mask":chain_mask,
+            "chain_encoding":chain_encoding}
+
+
+class featurize_GVP:
+    def __init__(self, num_positional_embeddings=16, top_k=30, num_rbf=16):
+        self.top_k = top_k
+        self.num_rbf = num_rbf
+        self.num_positional_embeddings = num_positional_embeddings
+        # self.letter_to_num = {'C': 4, 'D': 3, 'S': 15, 'Q': 5, 'K': 11, 'I': 9,
+        #                'P': 14, 'T': 16, 'F': 13, 'A': 0, 'G': 7, 'H': 8,
+        #                'E': 6, 'L': 10, 'R': 1, 'W': 17, 'V': 19, 
+        #                'N': 2, 'Y': 18, 'M': 12}
+        # self.num_to_letter = {v:k for k, v in self.letter_to_num.items()}
+    
+    def featurize(self, batch):
+        data_all = []
+        for b in batch:
+            if b is None:
+                continue
+            coords = torch.tensor(np.stack([b[c] for c in ['N', 'CA', 'C', 'O']], 1))
+            seq = torch.tensor(np.array(tokenizer.encode(b['seq'], add_special_tokens=False)))
+        
+            mask = torch.isfinite(coords.sum(dim=(1,2)))
+            coords[~mask] = np.inf
+            
+            X_ca = coords[:, 1].float()
+            edge_index = torch_geometric.nn.knn_graph(X_ca, k=self.top_k)#torch_cluster.knn_graph(X_ca, k=self.top_k)
+            
+            pos_embeddings = self._positional_embeddings(edge_index) # [E, 16]
+            E_vectors = X_ca[edge_index[0]] - X_ca[edge_index[1]] # [E, 3]
+            rbf = _rbf(E_vectors.norm(dim=-1), D_count=self.num_rbf) # [E, 16]
+            
+            dihedrals = self._dihedrals(coords)  # [n,6]
+            orientations = self._orientations(X_ca) # [n,2,3]
+            sidechains = self._sidechains(coords) # [n,3]
+            
+            node_s = dihedrals.float() # [n,6]
+
+            node_v = torch.cat([orientations, sidechains.unsqueeze(-2)], dim=-2).float() # [n, 3, 3]
+
+            edge_s = torch.cat([rbf, pos_embeddings], dim=-1).float() # [E, 32]
+            edge_v = _normalize(E_vectors).unsqueeze(-2).float() # [E, 1, 3]
+            
+            node_s, node_v, edge_s, edge_v = map(torch.nan_to_num,(node_s, node_v, edge_s, edge_v))
+            
+            data = torch_geometric.data.Data(x=X_ca, seq=seq,
+                                            node_s=node_s, node_v=node_v,
+                                            edge_s=edge_s, edge_v=edge_v,
+                                            edge_index=edge_index, mask=mask)
+            data_all.append(data)
+        return data_all
+    
+    def _positional_embeddings(self, edge_index, 
+                               num_embeddings=None,
+                               period_range=[2, 1000]):
+        # From https://github.com/jingraham/neurips19-graph-protein-design
+        num_embeddings = num_embeddings or self.num_positional_embeddings
+        d = edge_index[0] - edge_index[1]
+     
+        frequency = torch.exp(
+            torch.arange(0, num_embeddings, 2, dtype=torch.float32)
+            * -(np.log(10000.0) / num_embeddings)
+        )
+        angles = d.unsqueeze(-1) * frequency
+        E = torch.cat((torch.cos(angles), torch.sin(angles)), -1)
+        return E
+
+    def _dihedrals(self, X, eps=1e-7):
+        # From https://github.com/jingraham/neurips19-graph-protein-design
+        
+        X = torch.reshape(X[:, :3], [3*X.shape[0], 3])
+        dX = X[1:] - X[:-1]
+        U = _normalize(dX, dim=-1)
+        u_2 = U[:-2]
+        u_1 = U[1:-1]
+        u_0 = U[2:]
+
+        # Backbone normals
+        n_2 = _normalize(torch.cross(u_2, u_1), dim=-1)
+        n_1 = _normalize(torch.cross(u_1, u_0), dim=-1)
+
+        # Angle between normals
+        cosD = torch.sum(n_2 * n_1, -1)
+        cosD = torch.clamp(cosD, -1 + eps, 1 - eps)
+        D = torch.sign(torch.sum(u_2 * n_1, -1)) * torch.acos(cosD)
+
+        # This scheme will remove phi[0], psi[-1], omega[-1]
+        D = F.pad(D, [1, 2]) 
+        D = torch.reshape(D, [-1, 3])
+        # Lift angle representations to the circle
+        D_features = torch.cat([torch.cos(D), torch.sin(D)], 1)
+        return D_features
+    
+    def _orientations(self, X):
+        forward = _normalize(X[1:] - X[:-1])
+        backward = _normalize(X[:-1] - X[1:])
+        forward = F.pad(forward, [0, 0, 0, 1])
+        backward = F.pad(backward, [0, 0, 1, 0])
+        return torch.cat([forward.unsqueeze(-2), backward.unsqueeze(-2)], -2)
+
+    def _sidechains(self, X):
+        n, origin, c = X[:, 0], X[:, 1], X[:, 2]
+        c, n = _normalize(c - origin), _normalize(n - origin)
+        bisector = _normalize(c + n)
+        perp = _normalize(torch.cross(c, n))
+        vec = -bisector * math.sqrt(1 / 3) - perp * math.sqrt(2 / 3)
+        return vec 
+    
+    def collate(self, batch):
+        batch = self.featurize(batch)
+        if (batch is None) or (len(batch)==0):
+            return None
+        
+        elem = batch[0]
+        if isinstance(elem, Data):
+            return Batch.from_data_list(batch)
+        elif isinstance(elem, torch.Tensor):
+            return default_collate(batch)
+        elif isinstance(elem, float):
+            return torch.tensor(batch, dtype=torch.float)
+        elif isinstance(elem, int):
+            return torch.tensor(batch)
+        elif isinstance(elem, str):
+            return batch
+        elif isinstance(elem, Mapping):
+            return {key: self.collate([d[key] for d in batch]) for key in elem}
+        elif isinstance(elem, tuple) and hasattr(elem, '_fields'):
+            return type(elem)(*(self.collate(s) for s in zip(*batch)))
+        elif isinstance(elem, Sequence) and not isinstance(elem, str):
+            return [self.collate(s) for s in zip(*batch)]
+
+        raise TypeError('DataLoader found invalid type: {}'.format(type(elem)))
+
+
+def featurize_ProteinMPNN(batch, is_testing=False, chain_dict=None, fixed_position_dict=None, omit_AA_dict=None, tied_positions_dict=None, pssm_dict=None, bias_by_res_dict=None):
+    """ Pack and pad batch into torch tensors """
+    
+    batch = [one for one in batch if one is not None]
+    # print('______________________________________________________')
+    # print('______________________________________________________')
+    # print('______________________________________________________')
+    # print('______________________________________________________')
+    # print(batch[0].keys())
+    USING_DYNAMICS = True if ('norm_bfactors' in batch[0].keys()) or ('gt_flex' in batch[0].keys()) or  ('enm_vals'  in batch[0].keys()) or ('original_gt_flex' in batch[0].keys()) or ('eng_mask' in batch[0].keys()) else False
+
+    alphabet = 'ACDEFGHIKLMNPQRSTVWYX'
+    B = len(batch)
+    if B==0:
+        return None
+    lengths = np.array([len(b['seq']) for b in batch], dtype=np.int32) #sum of chain seq lengths
+    L_max = max([len(b['seq']) for b in batch])
+    X = np.zeros([B, L_max, 4, 3])
+    residue_idx = -100*np.ones([B, L_max], dtype=np.int32)
+    chain_M = np.zeros([B, L_max], dtype=np.int32) #1.0 for the bits that need to be predicted
+    pssm_coef_all = np.zeros([B, L_max], dtype=np.float32) #1.0 for the bits that need to be predicted
+    pssm_bias_all = np.zeros([B, L_max, 21], dtype=np.float32) #1.0 for the bits that need to be predicted
+    pssm_log_odds_all = 10000.0*np.ones([B, L_max, 21], dtype=np.float32) #1.0 for the bits that need to be predicted
+    chain_M_pos = np.zeros([B, L_max], dtype=np.int32) #1.0 for the bits that need to be predicted
+    bias_by_res_all = np.zeros([B, L_max, 21], dtype=np.float32)
+    chain_encoding_all = np.zeros([B, L_max], dtype=np.int32) #1.0 for the bits that need to be predicted
+    S = np.zeros([B, L_max], dtype=np.int32)
+    score = np.zeros([B, L_max])
+    omit_AA_mask = np.zeros([B, L_max, len(alphabet)], dtype=np.int32)
+    # Build the batch
+    letter_list_list = []
+    visible_list_list = []
+    masked_list_list = []
+    masked_chain_length_list_list = []
+    tied_pos_list_of_lists_list = []
+    # shuffle all chains before the main loop
+    if USING_DYNAMICS:
+        if ('norm_bfactors' in batch[0].keys()):
+            b_factors = np.zeros([B, L_max])
+        if ('gt_flex' in batch[0].keys()):
+            gt_flex = np.zeros([B, L_max])
+        if ('enm_vals' in batch[0].keys()):
+            enm_vals = np.zeros([B, L_max])
+        if ('original_gt_flex' in batch[0].keys()):
+            original_gt_flex = np.zeros([B, L_max])
+        if ('eng_mask' in batch[0].keys()):
+            eng_mask = np.zeros([B, L_max])
+
+    for i, b in enumerate(batch):
+        if chain_dict != None:
+            masked_chains, visible_chains = chain_dict[b['name']] #masked_chains a list of chain letters to predict [A, D, F]
+        else:
+            # masked_chains = [item[-1:] for item in list(b) if item[:10]=='seq_chain_']
+            masked_chains = ['']
+            visible_chains = []
+        # num_chains = b['num_of_chains']
+        all_chains = masked_chains + visible_chains
+        #random.shuffle(all_chains)
+    for i, b in enumerate(batch):
+        mask_dict = {}
+        a = 0
+        x_chain_list = []
+        chain_mask_list = []
+        chain_seq_list = []
+        chain_encoding_list = []
+        c = 1
+        letter_list = []
+        global_idx_start_list = [0]
+        visible_list = []
+        masked_list = []
+        masked_chain_length_list = []
+        fixed_position_mask_list = []
+        omit_AA_mask_list = []
+        pssm_coef_list = []
+        pssm_bias_list = []
+        pssm_log_odds_list = []
+        bias_by_res_list = []
+
+        if USING_DYNAMICS:
+            if ('norm_bfactors' in batch[0].keys()):
+                b_factors_list = []
+            if ('gt_flex' in batch[0].keys()):
+                gt_flex_list = []
+            if ('enm_vals' in batch[0].keys()):
+                enm_vals_list = []
+            if ('original_gt_flex' in batch[0].keys()):
+                original_gt_flex_list = []
+            if ('eng_mask' in batch[0].keys()):
+                eng_mask_list = []
+        l0 = 0
+        l1 = 0
+        for step, letter in enumerate(all_chains):
+            if letter in visible_chains:
+                letter_list.append(letter)
+                visible_list.append(letter)
+                chain_seq = b[f'seq_chain_{letter}']
+                chain_seq = ''.join([a if a!='-' else 'X' for a in chain_seq])
+                chain_length = len(chain_seq)
+                global_idx_start_list.append(global_idx_start_list[-1]+chain_length)
+                chain_coords = b[f'coords_chain_{letter}'] #this is a dictionary
+                chain_mask = np.zeros(chain_length) #0.0 for visible chains
+                x_chain = np.stack([chain_coords[c] for c in [f'N_chain_{letter}', f'CA_chain_{letter}', f'C_chain_{letter}', f'O_chain_{letter}']], 1) #[chain_lenght,4,3]
+                x_chain_list.append(x_chain)
+                chain_mask_list.append(chain_mask)
+                chain_seq_list.append(chain_seq)
+                chain_encoding_list.append(c*np.ones(np.array(chain_mask).shape[0]))
+                l1 += chain_length
+                residue_idx[i, l0:l1] = 100*(c-1)+np.arange(l0, l1)
+                l0 += chain_length
+                c+=1
+                fixed_position_mask = np.ones(chain_length)
+                fixed_position_mask_list.append(fixed_position_mask)
+                omit_AA_mask_temp = np.zeros([chain_length, len(alphabet)], np.int32)
+                omit_AA_mask_list.append(omit_AA_mask_temp)
+                pssm_coef = np.zeros(chain_length)
+                pssm_bias = np.zeros([chain_length, 21])
+                pssm_log_odds = 10000.0*np.ones([chain_length, 21])
+                pssm_coef_list.append(pssm_coef)
+                pssm_bias_list.append(pssm_bias)
+                pssm_log_odds_list.append(pssm_log_odds)
+                bias_by_res_list.append(np.zeros([chain_length, 21]))
+            if letter in masked_chains:
+                masked_list.append(letter)
+                letter_list.append(letter)
+                
+                if USING_DYNAMICS:
+                    if ('norm_bfactors' in batch[0].keys()):
+                        chain_b_factors = b['norm_bfactors']
+                        b_factors_list.append(chain_b_factors)
+                    if ('gt_flex' in batch[0].keys()):
+                        chain_gt_flex = b['gt_flex']
+                        gt_flex_list.append(chain_gt_flex)
+                    if ('enm_vals' in batch[0].keys()):
+                        chain_enm_vals = b['enm_vals']
+                        enm_vals_list.append(chain_enm_vals)
+                    if ('original_gt_flex' in batch[0].keys()):
+                        chain_original_gt_flex = b['original_gt_flex']
+                        original_gt_flex_list.append(chain_original_gt_flex)
+                    if ('eng_mask' in batch[0].keys()):
+                        chain_eng_mask = b['eng_mask']
+                        eng_mask_list.append(chain_eng_mask)
+
+                # chain_seq = b[f'seq_chain_{letter}']
+                chain_seq = b[f'seq{letter}']
+                chain_seq = ''.join([a if a!='-' else 'X' for a in chain_seq])
+                chain_length = len(chain_seq)
+                global_idx_start_list.append(global_idx_start_list[-1]+chain_length)
+                masked_chain_length_list.append(chain_length)
+                # chain_coords = b[f'coords_chain_{letter}'] #this is a dictionary
+                chain_coords = b
+                chain_mask = np.ones(chain_length) #1.0 for masked
+                # x_chain = np.stack([chain_coords[c] for c in [f'N_chain_{letter}', f'CA_chain_{letter}', f'C_chain_{letter}', f'O_chain_{letter}']], 1) #[chain_lenght,4,3]
+                x_chain = np.stack([chain_coords[c] for c in [f'N', f'CA', f'C', f'O']], 1) #[chain_lenght,4,3]
+                x_chain_list.append(x_chain)
+                chain_mask_list.append(chain_mask)
+                chain_seq_list.append(chain_seq)
+                chain_encoding_list.append(c*np.ones(np.array(chain_mask).shape[0]))
+                l1 += chain_length
+                residue_idx[i, l0:l1] = 100*(c-1)+np.arange(l0, l1)
+                l0 += chain_length
+                c+=1
+                fixed_position_mask = np.ones(chain_length)
+                if fixed_position_dict!=None:
+                    fixed_pos_list = fixed_position_dict[b['name']][letter]
+                    if fixed_pos_list:
+                        fixed_position_mask[np.array(fixed_pos_list)-1] = 0.0
+                fixed_position_mask_list.append(fixed_position_mask)
+                omit_AA_mask_temp = np.zeros([chain_length, len(alphabet)], np.int32)
+                if omit_AA_dict!=None:
+                    for item in omit_AA_dict[b['name']][letter]:
+                        idx_AA = np.array(item[0])-1
+                        AA_idx = np.array([np.argwhere(np.array(list(alphabet))== AA)[0][0] for AA in item[1]]).repeat(idx_AA.shape[0])
+                        idx_ = np.array([[a, b] for a in idx_AA for b in AA_idx])
+                        omit_AA_mask_temp[idx_[:,0], idx_[:,1]] = 1
+                omit_AA_mask_list.append(omit_AA_mask_temp)
+                pssm_coef = np.zeros(chain_length)
+                pssm_bias = np.zeros([chain_length, 21])
+                pssm_log_odds = 10000.0*np.ones([chain_length, 21])
+                if pssm_dict:
+                    if pssm_dict[b['name']][letter]:
+                        pssm_coef = pssm_dict[b['name']][letter]['pssm_coef']
+                        pssm_bias = pssm_dict[b['name']][letter]['pssm_bias']
+                        pssm_log_odds = pssm_dict[b['name']][letter]['pssm_log_odds']
+                pssm_coef_list.append(pssm_coef)
+                pssm_bias_list.append(pssm_bias)
+                pssm_log_odds_list.append(pssm_log_odds)
+                if bias_by_res_dict:
+                    bias_by_res_list.append(bias_by_res_dict[b['name']][letter])
+                else:
+                    bias_by_res_list.append(np.zeros([chain_length, 21]))
+
+       
+        letter_list_np = np.array(letter_list)
+        tied_pos_list_of_lists = []
+        tied_beta = np.ones(L_max)
+        if tied_positions_dict!=None:
+            tied_pos_list = tied_positions_dict[b['name']]
+            if tied_pos_list:
+                set_chains_tied = set(list(itertools.chain(*[list(item) for item in tied_pos_list])))
+                for tied_item in tied_pos_list:
+                    one_list = []
+                    for k, v in tied_item.items():
+                        start_idx = global_idx_start_list[np.argwhere(letter_list_np == k)[0][0]]
+                        if isinstance(v[0], list):
+                            for v_count in range(len(v[0])):
+                                one_list.append(start_idx+v[0][v_count]-1)#make 0 to be the first
+                                tied_beta[start_idx+v[0][v_count]-1] = v[1][v_count]
+                        else:
+                            for v_ in v:
+                                one_list.append(start_idx+v_-1)#make 0 to be the first
+                    tied_pos_list_of_lists.append(one_list)
+        tied_pos_list_of_lists_list.append(tied_pos_list_of_lists)
+ 
+        x = np.concatenate(x_chain_list,0) #[L, 4, 3]
+
+        if USING_DYNAMICS:
+            if ('norm_bfactors' in batch[0].keys()):
+                bf = np.concatenate(b_factors_list,0) #[L,]
+            if ('gt_flex' in batch[0].keys()):
+                gt = np.concatenate(gt_flex_list,0) #[L,]
+            if ('enm_vals' in batch[0].keys()):
+                enm = np.concatenate(enm_vals_list,0)
+            if ('original_gt_flex' in batch[0].keys()):
+                orig_gt = np.concatenate(original_gt_flex_list,0)
+            if ('eng_mask' in batch[0].keys()):
+                eng = np.concatenate(eng_mask_list,0)
+
+        all_sequence = "".join(chain_seq_list)
+        m = np.concatenate(chain_mask_list,0) #[L,], 1.0 for places that need to be predicted
+        chain_encoding = np.concatenate(chain_encoding_list,0)
+        m_pos = np.concatenate(fixed_position_mask_list,0) #[L,], 1.0 for places that need to be predicted
+
+        pssm_coef_ = np.concatenate(pssm_coef_list,0) #[L,], 1.0 for places that need to be predicted
+        pssm_bias_ = np.concatenate(pssm_bias_list,0) #[L,], 1.0 for places that need to be predicted
+        pssm_log_odds_ = np.concatenate(pssm_log_odds_list,0) #[L,], 1.0 for places that need to be predicted
+
+        bias_by_res_ = np.concatenate(bias_by_res_list, 0)  #[L,21], 0.0 for places where AA frequencies don't need to be tweaked
+
+        l = len(all_sequence)
+        x_pad = np.pad(x, [[0, L_max-l], [0,0], [0,0]], 'constant', constant_values=(np.nan, ))
+        if USING_DYNAMICS:
+            if ('norm_bfactors' in batch[0].keys()):
+                bf_pad = np.pad(bf, [[0, L_max-l]], 'constant', constant_values=(np.nan, ))
+            if ('gt_flex' in batch[0].keys()):
+                gt_pad = np.pad(gt, [[0, L_max-l]], 'constant', constant_values=(np.nan, ))
+            if ('enm_vals' in batch[0].keys()):
+                enm_pad = np.pad(enm, [[0, L_max-l]], 'constant', constant_values=(np.nan, ))
+            if ('original_gt_flex' in batch[0].keys()):
+                orig_gt_pad = np.pad(orig_gt, [[0, L_max-l]], 'constant', constant_values=(0, ))
+            if ('eng_mask' in batch[0].keys()):
+                eng_pad = np.pad(eng, [[0, L_max-l]], 'constant', constant_values=(0, ))
+            
+        X[i,:,:,:] = x_pad
+        if USING_DYNAMICS:
+            if ('norm_bfactors' in batch[0].keys()):
+                b_factors[i, :] = bf_pad
+            if ('gt_flex' in batch[0].keys()):
+                gt_flex[i, :] = gt_pad[:-1]
+            if ('enm_vals' in batch[0].keys()):
+                enm_vals[i, :] = enm_pad
+            if ('original_gt_flex' in batch[0].keys()):
+                original_gt_flex[i, :] = orig_gt_pad[:-1]
+            if ('eng_mask' in batch[0].keys()):
+                eng_mask[i, :] = eng_pad[:-1]
+                
+        if 'score' in b.keys():
+            score[i, :l] = b['score']
+        else:
+            score[i, :l] = 100.0
+ 
+        m_pad = np.pad(m, [[0, L_max-l]], 'constant', constant_values=(0.0, ))
+        m_pos_pad = np.pad(m_pos, [[0,L_max-l]], 'constant', constant_values=(0.0, ))
+        omit_AA_mask_pad = np.pad(np.concatenate(omit_AA_mask_list,0), [[0,L_max-l], [0, 0]], 'constant', constant_values=(0.0, ))
+        chain_M[i,:] = m_pad
+        chain_M_pos[i,:] = m_pos_pad
+        omit_AA_mask[i,] = omit_AA_mask_pad
+
+        chain_encoding_pad = np.pad(chain_encoding, [[0,L_max-l]], 'constant', constant_values=(0.0, ))
+        chain_encoding_all[i,:] = chain_encoding_pad
+
+        pssm_coef_pad = np.pad(pssm_coef_, [[0, L_max-l]], 'constant', constant_values=(0.0, ))
+        pssm_bias_pad = np.pad(pssm_bias_, [[0, L_max-l], [0,0]], 'constant', constant_values=(0.0, ))
+        pssm_log_odds_pad = np.pad(pssm_log_odds_, [[0,L_max-l], [0,0]], 'constant', constant_values=(0.0, ))
+
+        pssm_coef_all[i,:] = pssm_coef_pad
+        pssm_bias_all[i,:] = pssm_bias_pad
+        pssm_log_odds_all[i,:] = pssm_log_odds_pad
+
+        bias_by_res_pad = np.pad(bias_by_res_, [[0,L_max-l], [0,0]], 'constant', constant_values=(0.0, ))
+        bias_by_res_all[i,:] = bias_by_res_pad
+
+        # Convert to labels
+        indices = np.array(tokenizer.encode(b['seq'], add_special_tokens=False))
+        S[i, :l] = indices
+        letter_list_list.append(letter_list)
+        visible_list_list.append(visible_list)
+        masked_list_list.append(masked_list)
+        masked_chain_length_list_list.append(masked_chain_length_list)
+
+    
+    isnan = np.isnan(X)
+    mask = np.isfinite(np.sum(X,(2,3))).astype(np.float32)
+    X[isnan] = 0.
+
+    # Conversion
+    pssm_coef_all = torch.from_numpy(pssm_coef_all).to(dtype=torch.float32)
+    pssm_bias_all = torch.from_numpy(pssm_bias_all).to(dtype=torch.float32)
+    pssm_log_odds_all = torch.from_numpy(pssm_log_odds_all).to(dtype=torch.float32)
+
+    tied_beta = torch.from_numpy(tied_beta).to(dtype=torch.float32)
+
+    jumps = ((residue_idx[:,1:]-residue_idx[:,:-1])==1).astype(np.float32)
+    bias_by_res_all = torch.from_numpy(bias_by_res_all).to(dtype=torch.float32)
+    phi_mask = np.pad(jumps, [[0,0],[1,0]])
+    psi_mask = np.pad(jumps, [[0,0],[0,1]])
+    omega_mask = np.pad(jumps, [[0,0],[0,1]])
+    dihedral_mask = np.concatenate([phi_mask[:,:,None], psi_mask[:,:,None], omega_mask[:,:,None]], -1) #[B,L,3]
+    dihedral_mask = torch.from_numpy(dihedral_mask).to(dtype=torch.float32)
+    residue_idx = torch.from_numpy(residue_idx).to(dtype=torch.long)
+    S = torch.from_numpy(S).to(dtype=torch.long)
+    X = torch.from_numpy(X).to(dtype=torch.float32)
+    if USING_DYNAMICS:
+        if ('norm_bfactors' in batch[0].keys()):
+            b_factors = torch.from_numpy(b_factors).to(dtype=torch.float32)
+        if ('gt_flex' in batch[0].keys()):
+            gt_flex = torch.from_numpy(gt_flex).to(dtype=torch.float32)
+        if ('enm_vals' in batch[0].keys()):
+            enm_vals = torch.from_numpy(enm_vals).to(dtype=torch.float32)
+        if ('original_gt_flex' in batch[0].keys()):
+            original_gt_flex = torch.from_numpy(original_gt_flex).to(dtype=torch.float32)
+        if ('eng_mask' in batch[0].keys()):
+            eng_mask = torch.from_numpy(eng_mask).to(dtype=torch.float32)
+    score = torch.from_numpy(score).float()
+    mask = torch.from_numpy(mask).to(dtype=torch.float32)
+    chain_M = torch.from_numpy(chain_M).to(dtype=torch.float32)
+    chain_M_pos = torch.from_numpy(chain_M_pos).to(dtype=torch.float32)
+    omit_AA_mask = torch.from_numpy(omit_AA_mask).to(dtype=torch.float32)
+    chain_encoding_all = torch.from_numpy(chain_encoding_all).to(dtype=torch.long)
+
+    if is_testing is False:
+        retVal = {"title": [b['title'] for b in batch],
+                "X":X,
+                "S":S,
+                "score": score,
+                "mask":mask,
+                "lengths":lengths,
+                "chain_M":chain_M,
+                "chain_M_pos":chain_M_pos,
+                "residue_idx":residue_idx,
+                "chain_encoding_all":chain_encoding_all}
+        if USING_DYNAMICS:
+            if ('norm_bfactors' in batch[0].keys()):
+                retVal['norm_bfactors'] = b_factors
+            if ('gt_flex' in batch[0].keys()):
+                retVal['gt_flex'] = gt_flex
+            if ('enm_vals' in batch[0].keys()):
+                retVal['enm_vals'] = enm_vals
+            if ('original_gt_flex' in batch[0].keys()):
+                retVal['original_gt_flex'] = original_gt_flex
+            if ('eng_mask' in batch[0].keys()):
+                retVal['eng_mask'] = eng_mask
+            
+        return retVal
+    else:
+        retVal = {"title": [b['title'] for b in batch],
+                "X":X,
+                "S":S,
+                "score": score,
+                "mask":mask,
+                "lengths":lengths,
+                "chain_M":chain_M,
+                "chain_M_pos":chain_M_pos,
+                "residue_idx":residue_idx,
+                "chain_encoding_all":chain_encoding_all}
+        if USING_DYNAMICS:
+            if ('norm_bfactors' in batch[0].keys()):
+                retVal['norm_bfactors'] = b_factors
+            if ('gt_flex' in batch[0].keys()):
+                retVal['gt_flex'] = gt_flex
+            if ('enm_vals' in batch[0].keys()):
+                retVal['enm_vals'] = enm_vals
+            if ('original_gt_flex' in batch[0].keys()):
+                retVal['original_gt_flex'] = original_gt_flex
+            if ('eng_mask' in batch[0].keys()):
+                retVal['eng_mask'] = eng_mask
+        return retVal
+
+
+def featurize_Inversefolding(batch, shuffle_fraction=0.):
+    """ Pack and pad batch into torch tensors """
+    alphabet = 'ACDEFGHIKLMNPQRSTVWY'
+    B = len(batch)
+    lengths = np.array([len(b['seq']) for b in batch], dtype=np.int32)
+    L_max = max([len(b['seq']) for b in batch])
+    X = np.zeros([B, L_max, 3, 3])
+    S = np.zeros([B, L_max], dtype=np.int32)
+    score = np.ones([B, L_max]) * 100.0
+    chain_mask = np.zeros([B, L_max])-1 # 1:需要被预测的掩码部分 0:可见部分
+    chain_encoding = np.zeros([B, L_max])-1
+
+    # Build the batch
+    for i, b in enumerate(batch):
+        x = np.stack([b[c] for c in ['N', 'CA', 'C']], 1) # [#atom, 4, 3]
+        
+        l = len(b['seq'])
+        x_pad = np.pad(x, [[0,L_max-l], [0,0], [0,0]], 'constant', constant_values=(np.nan, )) # [#atom, 3, 3]
+        X[i,:,:,:] = x_pad
+
+        # Convert to labels
+        indices = np.array(tokenizer.encode(b['seq'], add_special_tokens=False))
+        if shuffle_fraction > 0.:
+            idx_shuffle = shuffle_subset(l, shuffle_fraction)
+            S[i, :l] = indices[idx_shuffle]
+        else:
+            S[i, :l] = indices
+        
+        chain_mask[i,:l] = b['chain_mask']
+        chain_encoding[i,:l] = b['chain_encoding']
+
+    mask = np.isfinite(np.sum(X,(2,3))).astype(np.float32) # atom mask
+    numbers = np.sum(mask, axis=1).astype(np.int)
+    S_new = np.zeros_like(S)
+    X_new = np.zeros_like(X)+np.nan
+    for i, n in enumerate(numbers):
+        X_new[i,:n,::] = X[i][mask[i]==1]
+        S_new[i,:n] = S[i][mask[i]==1]
+
+    X = X_new
+    S = S_new
+    isnan = np.isnan(X)
+    mask = np.isfinite(np.sum(X,(2,3))).astype(np.float32)
+    X[isnan] = 0.
+    # Conversion
+    S = torch.from_numpy(S).to(dtype=torch.long)
+    score = torch.from_numpy(score).float()
+    X = torch.from_numpy(X).to(dtype=torch.float32)
+    mask = torch.from_numpy(mask).to(dtype=torch.float32)
+    chain_mask = torch.from_numpy(chain_mask)
+    chain_encoding = torch.from_numpy(chain_encoding)
+    return {"title": [b['title'] for b in batch],
+            "X":X,
+            "S":S,
+            "score": score,
+            "mask":mask,
+            "lengths":lengths,
+            "chain_mask":chain_mask,
+            "chain_encoding":chain_encoding}

Flexpert-Design/src/datasets/flex_cath_dataset.py

@@ -0,0 +1,155 @@
+import os
+import json
+import numpy as np
+from tqdm import tqdm
+import random
+import torch.utils.data as data
+from .utils import cached_property
+from transformers import AutoTokenizer
+from src.tools.utils import load_yaml_config
+
+class FlexCATHDataset(data.Dataset):
+    def __init__(self, path='./',  split='train', max_length=500, test_name='All', data = None, removeTS=0, version=4.3, data_jsonl_name='/chain_set.jsonl', use_dynamics=True):
+        self.version = version
+        self.path = path
+        self.mode = split
+        self.max_length = max_length
+        self.test_name = test_name
+        self.removeTS = removeTS
+        self.data_jsonl_name = data_jsonl_name
+
+        self.using_dynamics = use_dynamics
+
+        print(self.data_jsonl_name)
+        if self.removeTS:
+            self.remove = json.load(open(self.path+'/remove.json', 'r'))['remove']
+        
+        if data is None:
+            if split == 'predict':
+                _split = 'valid'
+                print('In predict mode for CATH4.3 using VALIDATION split as the data. Consider switching to TEST set.')
+            else:
+                _split = split
+            self.data = self.cache_data[_split]
+        else:
+            self.data = data
+        
+        self.tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D", cache_dir="./cache_dir/")
+    
+    @cached_property
+    def cache_data(self):
+        alphabet='ACDEFGHIKLMNPQRSTVWY'
+        alphabet_set = set([a for a in alphabet])
+        print("path is: ", self.path)
+        if not os.path.exists(self.path):
+            raise "no such file:{} !!!".format(self.path)
+        else:
+            with open(self.path+'/'+self.data_jsonl_name) as f:
+                lines = f.readlines()
+            data_list = []
+            for line in tqdm(lines):
+                entry = json.loads(line)
+                if self.removeTS and entry['name'] in self.remove:
+                    continue
+                seq = entry['seq']
+
+                for key, val in entry['coords'].items():
+                    entry['coords'][key] = np.asarray(val)
+                
+                bad_chars = set([s for s in seq]).difference(alphabet_set)
+
+                if len(bad_chars) == 0:
+                    if len(entry['seq']) <= self.max_length: 
+                        chain_length = len(entry['seq'])
+                        chain_mask = np.ones(chain_length)
+                        data_list.append({
+                            'title':entry['name'],
+                            'seq':entry['seq'],
+                            'CA':entry['coords']['CA'],
+                            'C':entry['coords']['C'],
+                            'O':entry['coords']['O'],
+                            'N':entry['coords']['N'],
+                            'chain_mask': chain_mask,
+                            'chain_encoding': 1*chain_mask
+                        })
+                        if self.using_dynamics:
+                            data_list[-1]['gt_flex'] = entry['gt_flex']
+                            data_list[-1]['enm_vals'] = entry['enm_vals']
+                            if 'original_gt_flex' in entry:
+                                data_list[-1]['original_gt_flex'] = entry['original_gt_flex']
+                            if 'eng_mask' in entry:
+                                data_list[-1]['eng_mask'] = entry['eng_mask']
+                # else:
+                #     import pdb; pdb.set_trace()
+                #     print("Bad chars found in sequence: ", bad_chars)
+                        
+            if self.version==4.2:
+                with open(self.path+'/chain_set_splits.json') as f:
+                    dataset_splits = json.load(f)
+            
+            if self.version==4.3:
+                with open(self.path+'/chain_set_splits.json') as f:
+                    dataset_splits = json.load(f)
+                    # _dataset_splits = json.load(f)
+                    # dataset_splits = {k: _dataset_splits['train'] for k,_ in _dataset_splits.items()}
+                    # print("TODO: FIX THIS BACK!!!")
+                    # import pdb; pdb.set_trace()
+            
+            if self.test_name == 'L100':
+                with open(self.path+'/test_split_L100.json') as f:
+                    test_splits = json.load(f)
+                dataset_splits['test'] = test_splits['test']
+
+            if self.test_name == 'sc':
+                with open(self.path+'/test_split_sc.json') as f:
+                    test_splits = json.load(f)
+                dataset_splits['test'] = test_splits['test']
+            
+            name2set = {}
+            name2set.update({name:'train' for name in dataset_splits['train']})
+            name2set.update({name:'valid' for name in dataset_splits['validation']})
+            name2set.update({name:'test' for name in dataset_splits['test']})
+
+            data_dict = {'train':[],'valid':[],'test':[]}
+            for data in data_list:
+                if name2set.get(data['title']):
+                    if name2set[data['title']] == 'train':
+                        data_dict['train'].append(data)
+                    
+                    if name2set[data['title']] == 'valid':
+                        data_dict['valid'].append(data)
+                    
+                    if name2set[data['title']] == 'test':
+                        data['category'] = 'Unkown'
+                        data['score'] = 100.0
+                        data_dict['test'].append(data)
+            return data_dict
+
+    def change_mode(self, mode):
+        self.data = self.cache_data[mode]
+    
+    def __len__(self):
+        return len(self.data)
+    
+    def get_item(self, index):
+        return self.data[index]
+    
+    def __getitem__(self, index):
+        item = self.data[index]
+        L = len(item['seq'])
+        if L>self.max_length:
+            # 计算截断的最大索引
+            max_index = L - self.max_length
+            # 生成随机的截断索引
+            truncate_index = random.randint(0, max_index)
+            # 进行截断
+            item['seq'] = item['seq'][truncate_index:truncate_index+self.max_length]
+            item['CA'] = item['CA'][truncate_index:truncate_index+self.max_length]
+            item['C'] = item['C'][truncate_index:truncate_index+self.max_length]
+            item['O'] = item['O'][truncate_index:truncate_index+self.max_length]
+            item['N'] = item['N'][truncate_index:truncate_index+self.max_length]
+            item['chain_mask'] = item['chain_mask'][truncate_index:truncate_index+self.max_length]
+            item['chain_encoding'] = item['chain_encoding'][truncate_index:truncate_index+self.max_length]
+            item['gt_flex'] = item['gt_flex'][truncate_index:truncate_index+self.max_length]
+            item['enm_vals'] = item['enm_vals'][truncate_index:truncate_index+self.max_length]
+        return item

Flexpert-Design/src/datasets/foldswitchers_dataset.py

@@ -0,0 +1,128 @@
+import os
+import json
+import numpy as np
+from tqdm import tqdm
+import random
+import pdb
+import torch.utils.data as data
+from .utils import cached_property
+from transformers import AutoTokenizer
+
+class FoldswitchersDataset(data.Dataset):
+    def __init__(self, path='./',  split='train', max_length=500, test_name='All', data = None, removeTS=0):
+        self.path = path
+        self.mode = split
+        self.max_length = max_length
+        self.test_name = test_name
+        self.removeTS = removeTS
+        if self.removeTS:
+            self.remove = json.load(open(self.path+'/remove.json', 'r'))['remove']
+        
+        if data is None:
+            self.data = self.cache_data[split] #This calls the cache_data property
+        else:
+            self.data = data
+        
+        self.tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D", cache_dir="./cache_dir/")
+    
+    @cached_property
+    def cache_data(self):
+        alphabet='ACDEFGHIKLMNPQRSTVWY'
+        alphabet_set = set([a for a in alphabet])
+        print("path is: ", self.path)
+
+        if not os.path.exists(self.path):
+            raise "no such file:{} !!!".format(self.path)
+        else:
+
+            with open(self.path+'/chain_set.jsonl') as f:
+                lines = f.readlines()
+            data_list = []
+
+            for line in tqdm(lines):
+                entry = json.loads(line)
+
+                if self.removeTS and entry['name'] in self.remove:
+                    continue
+                seq = entry['seq']
+
+                for key, val in entry['coords'].items():
+                    entry['coords'][key] = np.asarray(val)
+                
+                bad_chars = set([s for s in seq]).difference(alphabet_set)
+
+                if len(bad_chars) == 0:
+                    if len(entry['seq']) <= self.max_length: 
+                        chain_length = len(entry['seq'])
+                        chain_mask = np.ones(chain_length)
+                        data_list.append({
+                            'title':entry['name'],
+                            'seq':entry['seq'],
+                            'CA':entry['coords']['CA'],
+                            'C':entry['coords']['C'],
+                            'O':entry['coords']['O'],
+                            'N':entry['coords']['N'],
+                            'chain_mask': chain_mask,
+                            'chain_encoding': 1*chain_mask
+                        })
+                        
+            with open(self.path+'/chain_set_splits_cleaned.json') as f:
+                dataset_splits = json.load(f)
+            
+            if self.test_name == 'L100':
+                with open(self.path+'/test_split_L100.json') as f:
+                    test_splits = json.load(f)
+                dataset_splits['test'] = test_splits['test']
+
+            if self.test_name == 'sc':
+                with open(self.path+'/test_split_sc.json') as f:
+                    test_splits = json.load(f)
+                dataset_splits['test'] = test_splits['test']
+            
+            name2set = {}
+            name2set.update({name:'train' for name in dataset_splits['train']})
+            name2set.update({name:'valid' for name in dataset_splits['validation']})
+            name2set.update({name:'test' for name in dataset_splits['test']})
+
+            data_dict = {'train':[],'valid':[],'test':[]}
+            for data in data_list:
+                #pdb.set_trace()
+                if name2set.get(data['title']): #This was causing the trouble with empty datasets - missmatch of names in the chain_set and chain_set_split
+                    if name2set[data['title']] == 'train':
+                        data_dict['train'].append(data)
+                    
+                    if name2set[data['title']] == 'valid':
+                        data_dict['valid'].append(data)
+                    
+                    if name2set[data['title']] == 'test':
+                        data['category'] = 'Unkown'
+                        data['score'] = 100.0
+                        data_dict['test'].append(data)
+            return data_dict
+
+    def change_mode(self, mode):
+        self.data = self.cache_data[mode]
+    
+    def __len__(self):
+        return len(self.data)
+    
+    def get_item(self, index):
+        return self.data[index]
+    
+    def __getitem__(self, index):
+        item = self.data[index]
+        L = len(item['seq'])
+        if L>self.max_length:
+            # 计算截断的最大索引
+            max_index = L - self.max_length
+            # 生成随机的截断索引
+            truncate_index = random.randint(0, max_index)
+            # 进行截断
+            item['seq'] = item['seq'][truncate_index:truncate_index+self.max_length]
+            item['CA'] = item['CA'][truncate_index:truncate_index+self.max_length]
+            item['C'] = item['C'][truncate_index:truncate_index+self.max_length]
+            item['O'] = item['O'][truncate_index:truncate_index+self.max_length]
+            item['N'] = item['N'][truncate_index:truncate_index+self.max_length]
+            item['chain_mask'] = item['chain_mask'][truncate_index:truncate_index+self.max_length]
+            item['chain_encoding'] = item['chain_encoding'][truncate_index:truncate_index+self.max_length]
+        return item

Flexpert-Design/src/datasets/mpnn_dataset.py

@@ -0,0 +1,492 @@
+import os
+import json
+import numpy as np
+from tqdm import tqdm
+import pandas as pd
+import torch.utils.data as data
+from Bio.PDB import PDBParser
+import torch
+import random
+import csv
+from dateutil import parser
+from .fast_dataloader import DataLoaderX
+from torch.utils.data import DataLoader
+import time
+
+from joblib import Parallel, delayed, cpu_count
+from tqdm import tqdm
+
+
+def pmap_multi(pickleable_fn, data, n_jobs=None, verbose=1, desc=None, **kwargs):
+    """
+
+    Parallel map using joblib.
+
+    Parameters
+    ----------
+    pickleable_fn : callable
+        Function to map over data.
+    data : iterable
+        Data over which we want to parallelize the function call.
+    n_jobs : int, optional
+        The maximum number of concurrently running jobs. By default, it is one less than
+        the number of CPUs.
+    verbose: int, optional
+        The verbosity level. If nonzero, the function prints the progress messages.
+        The frequency of the messages increases with the verbosity level. If above 10,
+        it reports all iterations. If above 50, it sends the output to stdout.
+    kwargs
+        Additional arguments for :attr:`pickleable_fn`.
+
+    Returns
+    -------
+    list
+        The i-th element of the list corresponds to the output of applying
+        :attr:`pickleable_fn` to :attr:`data[i]`.
+    """
+    if n_jobs is None:
+        n_jobs = cpu_count() - 1
+
+    results = Parallel(n_jobs=n_jobs, verbose=verbose, timeout=None)(
+    delayed(pickleable_fn)(*d, **kwargs) for i, d in tqdm(enumerate(data),desc=desc)
+    )
+
+
+
+def build_training_clusters(params, debug):
+    val_ids = set([int(l) for l in open(params['VAL']).readlines()])
+    test_ids = set([int(l) for l in open(params['TEST']).readlines()])
+   
+    if debug:
+        val_ids = []
+        test_ids = []
+ 
+    # read & clean list.csv
+    with open(params['LIST'], 'r') as f:
+        reader = csv.reader(f)
+        next(reader)
+        rows = [[r[0],r[3],int(r[4])] for r in reader
+                if float(r[2])<=params['RESCUT'] and
+                parser.parse(r[1])<=parser.parse(params['DATCUT'])]
+    
+    # compile training and validation sets
+    train = {}
+    valid = {}
+    test = {}
+
+    if debug:
+        rows = rows[:20]
+    for r in rows:
+        if r[2] in val_ids:
+            if r[2] in valid.keys():
+                valid[r[2]].append(r[:2])
+            else:
+                valid[r[2]] = [r[:2]]
+        elif r[2] in test_ids:
+            if r[2] in test.keys():
+                test[r[2]].append(r[:2])
+            else:
+                test[r[2]] = [r[:2]]
+        else:
+            if r[2] in train.keys():
+                train[r[2]].append(r[:2])
+            else:
+                train[r[2]] = [r[:2]]
+    if debug:
+        valid=train       
+    return train, valid, test
+
+
+def loader_pdb(item,params):
+
+    pdbid,chid = item[0].split('_')
+    PREFIX = "%s/pdb/%s/%s"%(params['DIR'],pdbid[1:3],pdbid)
+    
+    # load metadata
+    if not os.path.isfile(PREFIX+".pt"):
+        return {'seq': np.zeros(5)}
+    meta = torch.load(PREFIX+".pt")
+    asmb_ids = meta['asmb_ids']
+    asmb_chains = meta['asmb_chains']
+    chids = np.array(meta['chains'])
+
+    # find candidate assemblies which contain chid chain
+    asmb_candidates = set([a for a,b in zip(asmb_ids,asmb_chains)
+                           if chid in b.split(',')])
+
+    # if the chains is missing is missing from all the assemblies
+    # then return this chain alone
+    if len(asmb_candidates)<1:
+        chain = torch.load("%s_%s.pt"%(PREFIX,chid))
+        L = len(chain['seq'])
+        return {'seq'    : chain['seq'],
+                'xyz'    : chain['xyz'],
+                'idx'    : torch.zeros(L).int(),
+                'masked' : torch.Tensor([0]).int(),
+                'label'  : item[0]}
+
+    # randomly pick one assembly from candidates
+    asmb_i = random.sample(list(asmb_candidates), 1)
+
+    # indices of selected transforms
+    idx = np.where(np.array(asmb_ids)==asmb_i)[0]
+
+    # load relevant chains
+    chains = {c:torch.load("%s_%s.pt"%(PREFIX,c))
+              for i in idx for c in asmb_chains[i]
+              if c in meta['chains']}
+
+    # generate assembly
+    asmb = {}
+    for k in idx:
+
+        # pick k-th xform
+        xform = meta['asmb_xform%d'%k]
+        u = xform[:,:3,:3]
+        r = xform[:,:3,3]
+
+        # select chains which k-th xform should be applied to
+        s1 = set(meta['chains'])
+        s2 = set(asmb_chains[k].split(','))
+        chains_k = s1&s2
+
+        # transform selected chains 
+        for c in chains_k:
+            try:
+                xyz = chains[c]['xyz']
+                xyz_ru = torch.einsum('bij,raj->brai', u, xyz) + r[:,None,None,:]
+                asmb.update({(c,k,i):xyz_i for i,xyz_i in enumerate(xyz_ru)})
+            except KeyError:
+                return {'seq': np.zeros(5)}
+
+    # select chains which share considerable similarity to chid
+    seqid = meta['tm'][chids==chid][0,:,1]
+    homo = set([ch_j for seqid_j,ch_j in zip(seqid,chids)
+                if seqid_j>params['HOMO']])
+    # stack all chains in the assembly together
+    seq,xyz,idx,masked = "",[],[],[]
+    seq_list = []
+    for counter,(k,v) in enumerate(asmb.items()):
+        seq += chains[k[0]]['seq']
+        seq_list.append(chains[k[0]]['seq'])
+        xyz.append(v)
+        idx.append(torch.full((v.shape[0],),counter))
+        if k[0] in homo:
+            masked.append(counter)
+
+    return {'seq'    : seq,
+            'xyz'    : torch.cat(xyz,dim=0),
+            'idx'    : torch.cat(idx,dim=0),
+            'masked' : torch.Tensor(masked).int(),
+            'label'  : item[0]}
+
+def get_pdbs(data, max_length=10000, num_units=1000000):
+    init_alphabet = ['A', 'B', 'C', 'D', 'E', 'F', 'G','H', 'I', 'J','K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T','U', 'V','W','X', 'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f', 'g','h', 'i', 'j','k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't','u', 'v','w','x', 'y', 'z']
+    extra_alphabet = [str(item) for item in list(np.arange(300))]
+    chain_alphabet = init_alphabet + extra_alphabet
+    c = 0
+    c1 = 0
+
+ 
+    data = {k:v for k,v in data.items()}
+    c1 += 1
+    if 'label' in list(data):
+        my_dict = {}
+        s = 0
+        concat_seq = ''
+        concat_N = []
+        concat_CA = []
+        concat_C = []
+        concat_O = []
+        concat_mask = []
+        coords_dict = {}
+        mask_list = []
+        visible_list = []
+        if len(list(np.unique(data['idx']))) < 352:
+            for idx in list(np.unique(data['idx'])):
+                letter = chain_alphabet[idx]
+                res = np.argwhere(data['idx']==idx)
+                initial_sequence= "".join(list(np.array(list(data['seq']))[res][0,]))
+                if initial_sequence[-6:] == "HHHHHH":
+                    res = res[:,:-6]
+                if initial_sequence[0:6] == "HHHHHH":
+                    res = res[:,6:]
+                if initial_sequence[-7:-1] == "HHHHHH":
+                    res = res[:,:-7]
+                if initial_sequence[-8:-2] == "HHHHHH":
+                    res = res[:,:-8]
+                if initial_sequence[-9:-3] == "HHHHHH":
+                    res = res[:,:-9]
+                if initial_sequence[-10:-4] == "HHHHHH":
+                    res = res[:,:-10]
+                if initial_sequence[1:7] == "HHHHHH":
+                    res = res[:,7:]
+                if initial_sequence[2:8] == "HHHHHH":
+                    res = res[:,8:]
+                if initial_sequence[3:9] == "HHHHHH":
+                    res = res[:,9:]
+                if initial_sequence[4:10] == "HHHHHH":
+                    res = res[:,10:]
+                if res.shape[1] < 4:
+                    pass
+                else:
+                    my_dict['seq_chain_'+letter]= "".join(list(np.array(list(data['seq']))[res][0,]))
+                    concat_seq += my_dict['seq_chain_'+letter]
+                    if idx in data['masked']:
+                        mask_list.append(letter)
+                    else:
+                        visible_list.append(letter)
+                    coords_dict_chain = {}
+                    all_atoms = np.array(data['xyz'][res,])[0,] #[L, 14, 3]
+                    coords_dict_chain['N_chain_'+letter]=all_atoms[:,0,:].tolist()
+                    coords_dict_chain['CA_chain_'+letter]=all_atoms[:,1,:].tolist()
+                    coords_dict_chain['C_chain_'+letter]=all_atoms[:,2,:].tolist()
+                    coords_dict_chain['O_chain_'+letter]=all_atoms[:,3,:].tolist()
+                    my_dict['coords_chain_'+letter]=coords_dict_chain
+            my_dict['name']= data['label']
+            my_dict['masked_list']= mask_list
+            my_dict['visible_list']= visible_list
+            my_dict['num_of_chains'] = len(mask_list) + len(visible_list)
+            my_dict['seq'] = concat_seq
+            if len(concat_seq) <= max_length:
+                return my_dict
+    return None
+
+def safe_iter(ID, split_dict, params, alphabet_set, max_length=1000):
+    sel_idx = np.random.randint(0, len(split_dict[ID]))
+    out = loader_pdb(split_dict[ID][sel_idx], params)
+    entry = get_pdbs(out)
+    if entry is None:
+        return None
+    
+    seq = entry['seq']
+    bad_chars = set([s for s in seq]).difference(alphabet_set)
+    if len(bad_chars) != 0:
+        return None
+    
+    if len(entry['seq']) > max_length:
+        return None
+    
+    masked_chains = entry['masked_list']
+    visible_chains = entry['visible_list']
+    
+    all_chains = masked_chains + visible_chains
+    visible_temp_dict = {}
+    masked_temp_dict = {}
+    
+    for step, letter in enumerate(all_chains):
+        chain_seq = entry[f'seq_chain_{letter}']
+        if letter in visible_chains:
+            visible_temp_dict[letter] = chain_seq
+        elif letter in masked_chains:
+            masked_temp_dict[letter] = chain_seq
+            
+    for km, vm in masked_temp_dict.items():
+        for kv, vv in visible_temp_dict.items():
+            if vm == vv:
+                if kv not in masked_chains:
+                    masked_chains.append(kv)
+                if kv in visible_chains:
+                    visible_chains.remove(kv)
+    
+    all_chains = masked_chains + visible_chains
+    random.shuffle(all_chains)
+    
+    
+    x_chain_list = []
+    chain_mask_list = []
+    chain_seq_list = []
+    chain_encoding_list = []
+    c = 1
+    
+    for step, letter in enumerate(all_chains):
+        if letter in visible_chains:
+            chain_seq = entry[f'seq_chain_{letter}']
+            chain_length = len(chain_seq)
+            chain_coords = entry[f'coords_chain_{letter}'] #this is a dictionary
+            chain_mask = np.zeros(chain_length) #0.0 for visible chains
+            x_chain = np.stack([chain_coords[c] for c in [f'N_chain_{letter}', f'CA_chain_{letter}', f'C_chain_{letter}', f'O_chain_{letter}']], 1) #[chain_length,4,3]
+            x_chain_list.append(x_chain)
+            chain_mask_list.append(chain_mask)
+            chain_seq_list.append(chain_seq)
+            chain_encoding_list.append(c*np.ones(np.array(chain_mask).shape[0]))
+            c+=1
+        elif letter in masked_chains: 
+            chain_seq = entry[f'seq_chain_{letter}']
+            chain_length = len(chain_seq)
+            chain_coords = entry[f'coords_chain_{letter}'] #this is a dictionary
+            chain_mask = np.ones(chain_length) #0.0 for visible chains
+            x_chain = np.stack([chain_coords[c] for c in [f'N_chain_{letter}', f'CA_chain_{letter}', f'C_chain_{letter}', f'O_chain_{letter}']], 1) #[chain_lenght,4,3]
+            x_chain_list.append(x_chain)
+            chain_mask_list.append(chain_mask)
+            chain_seq_list.append(chain_seq)
+            chain_encoding_list.append(c*np.ones(np.array(chain_mask).shape[0]))
+            c+=1
+    
+    chain_mask_all = torch.from_numpy(np.concatenate(chain_mask_list))
+    chain_encoding_all = torch.from_numpy(np.concatenate(chain_encoding_list))
+    x_chain_all = torch.from_numpy(np.concatenate(x_chain_list))
+    
+    data = {
+        "title":entry['name'],
+        "seq":''.join(chain_seq_list), #len(seq)=n
+        "chain_mask":chain_mask_all,
+        "chain_encoding":chain_encoding_all,
+        "CA":x_chain_all[:,1], # [n,3]
+        "C":x_chain_all[:,2],
+        "O":x_chain_all[:,3],
+        "N":x_chain_all[:,0]} # [n,]
+    return data
+
+class MPNNDataset(data.Dataset):
+    def __init__(self, data_path='/gaozhangyang/drug_dataset/proteinmpnn_data/pdb_2021aug02', rescut=3.5, split='train'):
+        self.data_path = data_path
+        self.rescut = rescut
+        self.params = {
+            "LIST"    : f"{self.data_path}/list.csv", 
+            "VAL"     : f"{self.data_path}/valid_clusters.txt",
+            "TEST"    : f"{self.data_path}/test_clusters.txt",
+            "DIR"     : f"{self.data_path}",
+            "DATCUT"  : "2030-Jan-01",
+            "RESCUT"  : self.rescut, #resolution cutoff for PDBs
+            "HOMO"    : 0.70 #min seq.id. to detect homo chains
+        }
+        
+        if not os.path.exists("/gaozhangyang/experiments/OpenCPD/data/mpnn_data/split.pt"):
+            train, valid, test = build_training_clusters(self.params, False)
+            split = {"train": train, "valid":valid, "test":test}
+            torch.save(split, "/gaozhangyang/experiments/OpenCPD/data/mpnn_data/split.pt")
+        else:
+            split = torch.load("/gaozhangyang/experiments/OpenCPD/data/mpnn_data/split.pt")
+        
+        self.split_dict = split[mode]
+        alphabet='ACDEFGHIKLMNPQRSTVWYX'
+        self.alphabet_set = set([a for a in alphabet])
+        self.IDs = list(self.split_dict.keys())
+        # self.data = self.preprocess()
+    
+    def cache_split(self,):
+        train, valid, test = build_training_clusters(self.params, False)
+        
+        return {"train": train, "valid":valid, "test":test}
+    
+    @classmethod
+    def safe_iter(self, ID, split_dict, params, alphabet_set, max_length=1000):
+        # sel_idx = np.random.randint(0, len(split_dict[ID]))
+        sel_idx = 0
+        out = loader_pdb(split_dict[ID][sel_idx], params)
+        entry = get_pdbs(out)
+        if entry is None:
+            return None
+        
+        seq = entry['seq']
+        bad_chars = set([s for s in seq]).difference(alphabet_set)
+        if len(bad_chars) != 0:
+            return None
+        
+        if len(entry['seq']) > max_length:
+            return None
+        
+        masked_chains = entry['masked_list']
+        visible_chains = entry['visible_list']
+        
+        all_chains = masked_chains + visible_chains
+        visible_temp_dict = {}
+        masked_temp_dict = {}
+        
+        for step, letter in enumerate(all_chains):
+            chain_seq = entry[f'seq_chain_{letter}']
+            if letter in visible_chains:
+                visible_temp_dict[letter] = chain_seq
+            elif letter in masked_chains:
+                masked_temp_dict[letter] = chain_seq
+                
+        for km, vm in masked_temp_dict.items():
+            for kv, vv in visible_temp_dict.items():
+                if vm == vv:
+                    if kv not in masked_chains:
+                        masked_chains.append(kv)
+                    if kv in visible_chains:
+                        visible_chains.remove(kv)
+        
+        all_chains = masked_chains + visible_chains
+        random.shuffle(all_chains)
+        
+        
+        x_chain_list = []
+        chain_mask_list = []
+        chain_seq_list = []
+        chain_encoding_list = []
+        c = 1
+        
+        for step, letter in enumerate(all_chains):
+            if letter in visible_chains:
+                chain_seq = entry[f'seq_chain_{letter}']
+                chain_length = len(chain_seq)
+                chain_coords = entry[f'coords_chain_{letter}'] #this is a dictionary
+                chain_mask = np.zeros(chain_length) #0.0 for visible chains
+                x_chain = np.stack([chain_coords[c] for c in [f'N_chain_{letter}', f'CA_chain_{letter}', f'C_chain_{letter}', f'O_chain_{letter}']], 1) #[chain_length,4,3]
+                x_chain_list.append(x_chain)
+                chain_mask_list.append(chain_mask)
+                chain_seq_list.append(chain_seq)
+                chain_encoding_list.append(c*np.ones(np.array(chain_mask).shape[0]))
+                c+=1
+            elif letter in masked_chains: 
+                chain_seq = entry[f'seq_chain_{letter}']
+                chain_length = len(chain_seq)
+                chain_coords = entry[f'coords_chain_{letter}'] #this is a dictionary
+                chain_mask = np.ones(chain_length) #0.0 for visible chains
+                x_chain = np.stack([chain_coords[c] for c in [f'N_chain_{letter}', f'CA_chain_{letter}', f'C_chain_{letter}', f'O_chain_{letter}']], 1) #[chain_lenght,4,3]
+                x_chain_list.append(x_chain)
+                chain_mask_list.append(chain_mask)
+                chain_seq_list.append(chain_seq)
+                chain_encoding_list.append(c*np.ones(np.array(chain_mask).shape[0]))
+                c+=1
+        
+        chain_mask_all = np.concatenate(chain_mask_list)
+        chain_encoding_all = np.concatenate(chain_encoding_list)
+        x_chain_all = np.concatenate(x_chain_list)
+        
+        data = {
+            "title":entry['name']+str(int(chain_mask_all.sum())),
+            "seq":''.join(chain_seq_list), #len(seq)=n
+            "chain_mask":chain_mask_all,
+            "chain_encoding":chain_encoding_all,
+            "CA":x_chain_all[:,1], # [n,3]
+            "C":x_chain_all[:,2],
+            "O":x_chain_all[:,3],
+            "N":x_chain_all[:,0]} # [n,]
+        return data
+
+        
+    
+    def preprocess(self):
+        data = pmap_multi(self.safe_iter, [(ID,) for ID in self.IDs], split_dict=self.split_dict, params=self.params, alphabet_set=self.alphabet_set)
+        return data
+        
+    def __len__(self):
+        # return len(self.data)
+        return len(self.IDs)
+
+    def __getitem__(self, index):
+        ID = self.IDs[index]
+        out = self.safe_iter(ID, split_dict=self.split_dict, params=self.params, alphabet_set=self.alphabet_set)
+        return out
+
+
+def collate_fn(batch):
+    return batch
+
+
+if __name__ == "__main__":
+    MPNNDataset = MPNNDataset()
+    loader = DataLoaderX(local_rank=0, dataset = MPNNDataset, collate_fn=collate_fn, batch_size=4)
+    # loader = DataLoader(dataset = MPNNDataset, collate_fn=collate_fn, batch_size=4, prefetch_factor=4, num_workers=4)
+    for batch in tqdm(loader):
+        for one in batch:
+            if one is not None:
+                for key, val in one.items():
+                    if type(val) == torch.Tensor:
+                        result = val.to('cuda:0')
+        time.sleep(2)
+    print()

Flexpert-Design/src/datasets/pdb_inference.py

@@ -0,0 +1,329 @@
+import os
+import json
+import numpy as np
+import random
+import pdb
+import torch.utils.data as data
+from .utils import cached_property
+from transformers import AutoTokenizer
+
+#Imports for the PDB parser utils
+import glob
+import json
+import numpy as np
+import gzip
+import re
+import multiprocessing
+import tqdm
+import shutil
+SENTINEL = 1
+import biotite.structure as struc
+import biotite.application.dssp as dssp
+import biotite.structure.io.pdb.file as file
+
+class PDBInference(data.Dataset):
+    def __init__(self, path='./',  max_length=500, *args, **kwargs):
+        self.path = path
+        self.max_length = max_length
+
+        self.data = self.cache_data #TODO
+        self.tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D", cache_dir="./cache_dir/")
+    
+    @cached_property
+    def cache_data(self):
+        alphabet='ACDEFGHIKLMNPQRSTVWY'
+        alphabet_set = set([a for a in alphabet])
+        print("path is: ", self.path)
+
+        if not os.path.exists(self.path):
+            raise "no such folder:{} !!!".format(self.path)
+        else:
+
+            #list all PDBs
+            pdb_files = []
+            _files = os.listdir(self.path)
+            for _file in _files:
+                if _file.endswith('.pdb'):
+                    pdb_files.append(_file)
+            print(f'pdb_files size = {len(pdb_files)}')
+            #parse the PDBs into lines like if it was from the chain_set.json
+            lines = []
+            for _pdb in pdb_files:
+                _input_chain = _pdb.split('_')[1].split('.')[0] #ASSUMING NAMING 'PDBCODE_CHAINCODE_XXX'
+                _line = self.parse_PDB(self.path+'/'+_pdb, name=_pdb.split('.')[0], input_chain=_input_chain) #Input chain list can be parsed here as well
+                #pdb.set_trace()
+                lines.append(_line[0])
+            
+            print(f'lines size = {len(lines)}')
+            data_list = []
+
+            flex_instructions = {}
+            flexibility_files = glob.glob(self.path + '/*instructions.csv')
+            for file in flexibility_files:
+                with open(file, 'r') as f:
+                    flexibility_instructions_parsed= f.read().strip().split(',')
+                    flexibility_instructions_parsed = [float(i) for i in flexibility_instructions_parsed] + [0.0] #add the padding here
+                    flex_instructions[file.split('/')[-1].split('_instructions')[0]] = flexibility_instructions_parsed
+
+            for line in tqdm.tqdm(lines):
+                entry = line
+
+                seq = entry['seq']
+
+                for key, val in entry['coords'].items():
+                    entry['coords'][key] = np.asarray(val)
+                
+                bad_chars = set([s for s in seq]).difference(alphabet_set)
+                try:
+                    _flex_instructions = flex_instructions[entry['name']]
+                except KeyError:
+                    _flex_instructions = [0.0] * len(seq)
+                    print(f"No flexibility instructions found for {entry['name']}. Passing zeros.")
+
+                if len(bad_chars) == 0:
+                    if len(entry['seq']) <= self.max_length: 
+                        chain_length = len(entry['seq'])
+                        chain_mask = np.ones(chain_length)
+                        data_list.append({
+                            'title':entry['name'],
+                            'seq':entry['seq'],
+                            'CA':entry['coords']['CA'],
+                            'C':entry['coords']['C'],
+                            'O':entry['coords']['O'],
+                            'N':entry['coords']['N'],
+                            'chain_mask': chain_mask,
+                            'chain_encoding': 1*chain_mask,
+                            'gt_flex': _flex_instructions
+                        })
+                else:
+                    print(f'Skipping PDBs with Bad chars, e.g. gaps in the sequence: {entry["name"]}')
+            
+            #data_dict = {'train':[],'valid':data_list,'test':data_list}
+            print(f'data_list size = {len(data_list)}')
+            return data_list#data_dict
+
+    def change_mode(self, mode):
+        self.data = self.cache_data[mode]
+    
+    def __len__(self):
+        return len(self.data)
+    
+    def get_item(self, index):
+        return self.data[index]
+    
+    def __getitem__(self, index):
+        item = self.data[index]
+        L = len(item['seq'])
+        if L>self.max_length:
+            # 计算截断的最大索引
+            max_index = L - self.max_length
+            # 生成随机的截断索引
+            truncate_index = random.randint(0, max_index)
+            # 进行截断
+            item['seq'] = item['seq'][truncate_index:truncate_index+self.max_length]
+            item['CA'] = item['CA'][truncate_index:truncate_index+self.max_length]
+            item['C'] = item['C'][truncate_index:truncate_index+self.max_length]
+            item['O'] = item['O'][truncate_index:truncate_index+self.max_length]
+            item['N'] = item['N'][truncate_index:truncate_index+self.max_length]
+            item['chain_mask'] = item['chain_mask'][truncate_index:truncate_index+self.max_length]
+            item['chain_encoding'] = item['chain_encoding'][truncate_index:truncate_index+self.max_length]
+            item['gt_flex'] = item['gt_flex'][truncate_index:truncate_index+self.max_length]
+        return item
+
+    #Code from data_utils on local PC, based on: https://github.com/JoreyYan/zetadesign/blob/master/data/data.py
+    def parse_PDB_biounits(self, x, sse,ssedssp,atoms=['N', 'CA', 'C'], chain=None):
+        '''
+        input:  x = PDB filename
+                atoms = atoms to extract (optional)
+        output: (length, atoms, coords=(x,y,z)), sequence
+        '''
+
+        alpha_1 = list("ARNDCQEGHILKMFPSTWYV-")
+        states = len(alpha_1)
+        alpha_3 = ['ALA', 'ARG', 'ASN', 'ASP', 'CYS', 'GLN', 'GLU', 'GLY', 'HIS', 'ILE',
+                'LEU', 'LYS', 'MET', 'PHE', 'PRO', 'SER', 'THR', 'TRP', 'TYR', 'VAL', 'GAP']
+
+        aa_1_N = {a: n for n, a in enumerate(alpha_1)}
+        aa_3_N = {a: n for n, a in enumerate(alpha_3)}
+        aa_N_1 = {n: a for n, a in enumerate(alpha_1)}
+        aa_1_3 = {a: b for a, b in zip(alpha_1, alpha_3)}
+        aa_3_1 = {b: a for a, b in zip(alpha_1, alpha_3)}
+
+        def AA_to_N(x):
+            x = np.array(x)
+            if x.ndim == 0: x = x[None]
+            return [[aa_1_N.get(a, states - 1) for a in y] for y in x]
+
+        def N_to_AA(x):
+            x = np.array(x)
+            if x.ndim == 1: x = x[None]
+            return ["".join([aa_N_1.get(a, "-") for a in y]) for y in x]
+
+        xyz, seq, plddts, min_resn, max_resn = {}, {}, [],  1e6, -1e6
+
+        pdbcontents = x.split('\n')[0]
+        with open(pdbcontents) as f:
+            pdbcontents = f.readlines()
+        for line in pdbcontents:
+
+            if line[:6] == "HETATM" and line[17:17 + 3] == "MSE":
+                line = line.replace("HETATM", "ATOM  ")
+                line = line.replace("MSE", "MET")
+
+            if line[:4] == "ATOM":
+                ch = line[21:22]
+                if ch == chain or chain is None or ch==' ':
+                    atom = line[12:12 + 4].strip()
+                    resi = line[17:17 + 3]
+                    resn = line[22:22 + 5].strip()
+                    plddt=line[60:60 + 6].strip()
+
+
+
+                    x, y, z = [float(line[i:(i + 8)]) for i in [30, 38, 46]]
+
+                    if resn[-1].isalpha():
+                        resa, resn = resn[-1], int(resn[:-1]) - 1 # in same pos ,use last atoms
+                    else:
+                        resa, resn = "_", int(resn) - 1
+                    #         resn = int(resn)
+                    if resn < min_resn:
+                        min_resn = resn
+                    if resn > max_resn:
+                        max_resn = resn
+
+
+
+                    if resn not in xyz:
+                        xyz[resn] = {}
+                    if resa not in xyz[resn]:
+                        xyz[resn][resa] = {}
+                    if resn not in seq:
+                        seq[resn] = {}
+
+                    if resa not in seq[resn]:
+                        seq[resn][resa] = resi
+
+                    if atom not in xyz[resn][resa]:
+                        xyz[resn][resa][atom] = np.array([x, y, z])
+
+        # convert to numpy arrays, fill in missing values
+        seq_, xyz_ ,sse_,ssedssp_= [], [], [], []
+        dsspidx=0
+        sseidx=0
+
+        for resn in range(int(min_resn), int(max_resn + 1)):
+            if resn in seq:
+                for k in sorted(seq[resn]):
+                    seq_.append(aa_3_N.get(seq[resn][k], 20))
+                    try:
+                        if 'CA' in xyz[resn][k]:
+                            sse_.append(sse[sseidx])
+                            sseidx = sseidx + 1
+                        else:
+                            sse_.append('-')
+                    except:
+                        print('error sse')
+
+
+            else:
+                seq_.append(20)
+                sse_.append('-')
+
+            misschianatom = False
+            if resn in xyz:
+
+
+                for k in sorted(xyz[resn]):
+                    for atom in atoms:
+                        if atom in xyz[resn][k]:
+                            xyz_.append(xyz[resn][k][atom])  #some will miss C and O ,but sse is normal,because sse just depend on CA
+                        else:
+                            xyz_.append(np.full(3, np.nan))
+                            misschianatom=True
+                    if misschianatom:
+                        ssedssp_.append('-')
+                        misschianatom = False
+                    else:
+                        try:
+                            ssedssp_.append(ssedssp[dsspidx])         # if miss chain atom,xyz ,seq think is ok , but dssp miss this
+                            dsspidx = dsspidx + 1
+                        except:
+                            pass
+                            #print(dsspidx)
+
+
+            else:
+                for atom in atoms:
+                    xyz_.append(np.full(3, np.nan))
+                ssedssp_.append('-')
+
+
+        return np.array(xyz_).reshape(-1, len(atoms), 3), N_to_AA(np.array(seq_)),np.array(sse_),np.array(ssedssp_)
+
+    def parse_PDB(self, path_to_pdb, name, input_chain):
+        """
+        make sure every time just input 1 line
+        """
+        c = 0
+        pdb_dict_list = []
+
+
+        biounit_names = [path_to_pdb]
+        for biounit in biounit_names:
+            my_dict = {}
+            s = 0
+            concat_seq = ''
+
+
+            letter = input_chain #Assuming single chain!!
+
+            PDBFile = file.PDBFile.read(biounit)
+            array_stack = PDBFile.get_structure(altloc="all")
+
+            #In case the passed letter is unknown, select one chain from the PDB file based on the dominant protein chain
+            if letter not in array_stack.chain_id:
+                is_protein = struc.filter_amino_acids(array_stack)
+                protein_atoms = array_stack[0][is_protein]
+                chain_ids, chain_counts = np.unique(protein_atoms.chain_id, return_counts=True)
+                dominant_chain_id = chain_ids[np.argmax(chain_counts)]
+                letter = dominant_chain_id
+
+
+            sse1 = struc.annotate_sse(array_stack[0], chain_id=letter).tolist()
+            if len(sse1)==0:
+                sse1 = struc.annotate_sse(array_stack[0], chain_id='').tolist()
+
+            ssedssp1 = [] #not annotating dssp for now
+
+
+            xyz, seq, sse, ssedssp = self.parse_PDB_biounits(biounit,sse1,ssedssp1,atoms=['N', 'CA', 'C','O'], chain=letter) #TODO: fix the float error
+            ssedssp = sse  #faking it for now
+
+            assert len(sse)==len(seq[0])
+            assert len(ssedssp) == len(seq[0])
+
+            if type(xyz) != str:
+                concat_seq += seq[0]
+                my_dict['seq_chain_' + letter] = seq[0]
+
+                coords_dict_chain = {}
+                coords_dict_chain['N'] = xyz[:, 0, :].tolist()
+                coords_dict_chain['CA'] = xyz[:, 1, :].tolist()
+                coords_dict_chain['C'] = xyz[:, 2, :].tolist()
+                coords_dict_chain['O'] = xyz[:, 3, :].tolist()
+                my_dict['coords_chain_' + letter] = coords_dict_chain
+                my_dict['coords'] = coords_dict_chain
+                s += 1
+
+            # if s>1:
+            #     raise NotImplementedError('Inference so far implemented only for single chain proteins')
+
+            my_dict['name'] = name
+            my_dict['num_chains'] = s
+            my_dict['seq'] = my_dict[f'seq_chain_{letter}'] #concat_seq
+            # if s <= len(chain_alphabet):
+            #     pdb_dict_list.append(my_dict)
+            #     c += 1
+            pdb_dict_list.append(my_dict)
+        return pdb_dict_list

Flexpert-Design/src/datasets/ts_dataset.py

@@ -0,0 +1,47 @@
+import os
+import json
+import numpy as np
+import torch.utils.data as data
+
+
+class TSDataset(data.Dataset):
+    def __init__(self, path = './', split='test'):
+        if not os.path.exists(path):
+            raise "no such file:{} !!!".format(path)
+        else:
+            ts50_data = json.load(open(path+'/ts50.json'))
+            ts500_data = json.load(open(path+'/ts500.json'))
+
+        # TS500 has proteins with lengths of 500+
+        # TS50 only contains proteins with lengths less than 500
+        self.data = []
+        for temp in ts50_data:
+            coords = np.array(temp['coords'])
+            self.data.append({'title':temp['name'],
+                                'seq':temp['seq'],
+                                'CA':coords[:,1,:],
+                                'C':coords[:,2,:],
+                                'O':coords[:,3,:],
+                                'N':coords[:,0,:],
+                                'category': 'ts50'
+                                })
+
+        for temp in ts500_data:
+            coords = np.array(temp['coords'])
+            self.data.append({'title':temp['name'],
+                                'seq':temp['seq'],
+                                'CA':coords[:,1,:],
+                                'C':coords[:,2,:],
+                                'O':coords[:,3,:],
+                                'N':coords[:,0,:],
+                                'category': 'ts500'
+                                })
+
+    def __len__(self):
+        return len(self.data)
+    
+    def get_item(self, index):
+        return self.data[index]
+
+    def __getitem__(self, index):
+        return self.data[index]

Flexpert-Design/src/datasets/utils.py

@@ -0,0 +1,99 @@
+import random
+
+
+class cached_property(object):
+    """
+    Descriptor (non-data) for building an attribute on-demand on first use.
+    """
+    def __init__(self, factory):
+        """
+        <factory> is called such: factory(instance) to build the attribute.
+        """
+        self._attr_name = factory.__name__
+        self._factory = factory
+
+    def __get__(self, instance, owner):
+        # Build the attribute.
+        attr = self._factory(instance)
+
+        # Cache the value; hide ourselves.
+        setattr(instance, self._attr_name, attr)
+        return attr
+
+
+def get_inds(expected_num, clu_nums, cid2clu, seq2ind):
+    cur_len, cur_idx, query_cids, query_idx = 0, 0, [], []
+    while cur_len < expected_num:
+        cid, l = clu_nums[cur_idx % (len(clu_nums))]
+        cur_idx += 1
+        # check if this cluster has been selected
+        if cid in query_cids:
+            continue
+        if random.random() > 0.5:
+            for seq in cid2clu[cid]:
+                # seq2ind: ensure it is in limited lengths
+                if seq in seq2ind.keys():
+                    query_idx.append(seq2ind[seq])
+                    cur_len += 1
+
+            query_cids.append(cid)
+    return query_cids, query_idx
+
+
+def get_num(N, valid_num=100):
+    train_n, valid_n = int(0.9 * N), min(valid_num, int(0.05 * N))
+    test_n = N - train_n - valid_n
+    return train_n, valid_n, test_n
+
+
+def get_full_inds(expected_num, clu_nums, cid2clu, full_seq2ind):
+    cur_len, cur_idx, query_cids, query_idx = 0, 0, [], {}
+    # build query_idx for each dataset
+    for dataname in full_seq2ind.keys():
+        if dataname not in query_idx.keys():
+            query_idx[dataname] = []
+    cur_idx_lst = list(range(len(clu_nums)))
+    while cur_len < expected_num:
+        cur_idx = random.choice(cur_idx_lst)
+        cid, l = clu_nums[cur_idx]
+        # check if this cluster has been selected
+        if cid in query_cids:
+            continue
+        for seq in set(cid2clu[cid]):
+            # seq2ind: ensure it is in limited lengths
+            for dataname in full_seq2ind.keys():
+                if seq in full_seq2ind[dataname].keys():
+                    query_idx[dataname].append(full_seq2ind[dataname][seq])
+                    cur_len += 1
+        query_cids.append(cid)
+        cur_idx_lst.remove(cur_idx)
+    return query_cids, query_idx
+
+
+def get_inds(expected_num, clu_nums, cid2clu, seq2ind):
+    cur_len, query_cids, query_idx = 0, [], []
+    cur_idx_lst = list(range(len(clu_nums)))
+    while cur_len < expected_num:
+        try:
+            cur_idx = random.choice(cur_idx_lst)
+            cid, l = clu_nums[cur_idx]
+            # check if this cluster has been selected
+            if cid in query_cids:
+                continue
+
+            # check if this cluster is too big
+            pre = abs(expected_num - cur_len)
+            aft = abs(cur_len + l - expected_num)
+            if pre < aft:
+                continue
+
+            for seq in cid2clu[cid]:
+                # seq2ind: ensure it is in limited lengths
+                if seq in seq2ind.keys():
+                    query_idx.append(seq2ind[seq])
+                    cur_len += 1
+            query_cids.append(cid)
+            cur_idx_lst.remove(cur_idx)
+        except:
+            break
+    return query_cids, query_idx

Flexpert-Design/src/interface/data_interface.py

@@ -0,0 +1,66 @@
+import inspect
+import importlib
+import pytorch_lightning as pl
+from torch.utils.data import DataLoader
+
+
+class DInterface_base(pl.LightningDataModule):
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.save_hyperparameters()
+        self.batch_size = self.hparams.batch_size
+        print("batch_size", self.batch_size)
+        self.load_data_module()
+
+    def setup(self, stage=None):
+        # Assign train/val datasets for use in dataloaders
+        if stage == 'fit' or stage is None:
+            self.trainset = self.instancialize(split = 'train')
+            self.valset = self.instancialize(split='valid')
+
+        # Assign test dataset for use in dataloader(s)
+        if stage == 'test' or stage is None:
+            self.testset = self.instancialize(split='test')
+
+    def train_dataloader(self):
+        return DataLoader(self.trainset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=True, prefetch_factor=3)
+
+    def val_dataloader(self):
+        return DataLoader(self.valset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False)
+
+    def test_dataloader(self):
+        return DataLoader(self.testset, batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False)
+
+    def load_data_module(self):
+        name = self.dataset
+        # Change the `snake_case.py` file name to `CamelCase` class name.
+        # Please always name your model file name as `snake_case.py` and
+        # class name corresponding `CamelCase`.
+        camel_name = ''.join([i.capitalize() for i in name.split('_')])
+        try:
+            self.data_module = getattr(importlib.import_module(
+                '.'+name, package=__package__), camel_name)
+        except:
+            raise ValueError(
+                f'Invalid Dataset File Name or Invalid Class Name data.{name}.{camel_name}')
+
+    def instancialize(self, **other_args):
+        """ Instancialize a model using the corresponding parameters
+            from self.hparams dictionary. You can also input any args
+            to overwrite the corresponding value in self.kwargs.
+        """
+        if other_args['split'] == 'train':
+            self.data_module = getattr(importlib.import_module(
+                    '.AF2DB_dataset', package='data'), 'Af2dbDataset')
+        else:
+            self.data_module = getattr(importlib.import_module(
+                    '.CASP15_dataset', package='data'), 'CASP15Dataset')
+        
+        class_args =  list(inspect.signature(self.data_module.__init__).parameters)[1:]
+        inkeys = self.kwargs.keys()
+        args1 = {}
+        for arg in class_args:
+            if arg in inkeys:
+                args1[arg] = self.kwargs[arg]
+        args1.update(other_args)
+        return self.data_module(**args1)

Flexpert-Design/src/interface/model_interface.py

@@ -0,0 +1,89 @@
+import torch
+import pytorch_lightning as pl
+import torch.nn as nn
+import os
+import torch.optim.lr_scheduler as lrs
+import inspect
+
+class MInterface_base(pl.LightningModule):
+    def __init__(self, model_name=None, loss=None, lr=None, **kargs):
+        super().__init__()
+        self.save_hyperparameters()
+        self.load_model()
+        self.configure_loss()
+        os.makedirs(os.path.join(self.hparams.res_dir, self.hparams.ex_name), exist_ok=True)
+        
+    def forward(self, input):
+        pass
+    
+    
+    def training_step(self, batch, batch_idx, **kwargs):
+        pass
+
+
+    def validation_step(self, batch, batch_idx):
+        pass
+
+    def test_step(self, batch, batch_idx):
+        # Here we just reuse the validation_step for testing
+        return self.validation_step(batch, batch_idx)
+
+    def on_validation_epoch_end(self):
+        # Make the Progress Bar leave there
+        self.print('')
+    
+    def get_schedular(self, optimizer, lr_scheduler='onecycle'):
+        if lr_scheduler == 'step':
+            scheduler = lrs.StepLR(optimizer,
+                                    step_size=self.hparams.lr_decay_steps,
+                                    gamma=self.hparams.lr_decay_rate)
+        elif lr_scheduler == 'cosine':
+            scheduler = lrs.CosineAnnealingLR(optimizer,
+                                                T_max=self.hparams.lr_decay_steps,
+                                                eta_min=self.hparams.lr_decay_min_lr)
+        elif lr_scheduler == 'onecycle':
+            scheduler = lrs.OneCycleLR(optimizer, max_lr=self.hparams.lr, steps_per_epoch=self.hparams.steps_per_epoch, epochs=self.hparams.epoch, three_phase=False)
+        else:
+            raise ValueError('Invalid lr_scheduler type!')
+
+        return scheduler
+
+    def configure_optimizers(self):
+        if hasattr(self.hparams, 'weight_decay'):
+            weight_decay = self.hparams.weight_decay
+        else:
+            weight_decay = 0
+    
+        optimizer_g = torch.optim.AdamW(self.model.parameters(), lr=self.hparams.lr, weight_decay=weight_decay, betas=(0.9, 0.98), eps=1e-8)
+
+        schecular_g = self.get_schedular(optimizer_g, self.hparams.lr_scheduler)
+
+        return [optimizer_g], [{"scheduler": schecular_g, "interval": "step"}]
+        
+    def lr_scheduler_step(self, *args, **kwargs):
+        scheduler = self.lr_schedulers()
+        scheduler.step()
+        
+    
+    def configure_devices(self):
+        self.device = torch.device(self.hparams.device)
+
+    def configure_loss(self):
+        self.loss_function = nn.CrossEntropyLoss(reduction='none')
+        
+    def load_model(self):
+        self.model = None
+
+    def instancialize(self, Model, **other_args):
+        """ Instancialize a model using the corresponding parameters
+            from self.hparams dictionary. You can also input any args
+            to overwrite the corresponding value in self.hparams.
+        """
+        class_args = inspect.getargspec(Model.__init__).args[1:]
+        inkeys = self.hparams.keys()
+        args1 = {}
+        for arg in class_args:
+            if arg in inkeys:
+                args1[arg] = getattr(self.hparams, arg)
+        args1.update(other_args)
+        return Model(**args1)

Flexpert-Design/src/interface/pretrain_interface.py

@@ -0,0 +1,405 @@
+import torch
+from omegaconf import OmegaConf
+from transformers import AutoTokenizer, EsmForMaskedLM
+import torch.nn.functional as F
+
+class PretrainInterface(torch.nn.Module):
+    def __init__(self, name):
+        super().__init__()
+        self.name = name
+        if name == "ESM35M":
+            self.esm_dim = 480
+            self.tokenizer = AutoTokenizer.from_pretrained("/huyuqi/model_zoom/transformers/models--facebook--esm2_t12_35M_UR50D")
+            self.pretrain_model = EsmForMaskedLM.from_pretrained("/huyuqi/model_zoom/transformers/models--facebook--esm2_t12_35M_UR50D")
+        if name == "ESM650M":
+            self.esm_dim = 1280
+            self.tokenizer = AutoTokenizer.from_pretrained("/huyuqi/model_zoom/transformers/models--facebook--esm2_t33_650M_UR50D/snapshots/08e4846e537177426273712802403f7ba8261b6c")
+            self.pretrain_model = EsmForMaskedLM.from_pretrained("/huyuqi/model_zoom/transformers/models--facebook--esm2_t33_650M_UR50D/snapshots/08e4846e537177426273712802403f7ba8261b6c")
+        if name == "ESM3B":
+            self.esm_dim = 2560
+            self.tokenizer = AutoTokenizer.from_pretrained("/huyuqi/model_zoom/transformers/models--facebook--esm2_t36_3B_UR50D/snapshots/476b639933c8baad5ad09a60ac1a87f987b656fc")
+            self.pretrain_model = EsmForMaskedLM.from_pretrained("/huyuqi/model_zoom/transformers/models--facebook--esm2_t36_3B_UR50D/snapshots/476b639933c8baad5ad09a60ac1a87f987b656fc")
+        
+        if name == "vanilla":
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/DiffSDS/Pretrain_lightning/results/ESMVQ/base/configs/10-18T01-15-36-project.yaml")
+            pretrain_args.diffusion = False
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load('/huyuqi/xmyu/DiffSDS/Pretrain_lightning/results/ESMVQ/base/checkpoints/best-epoch=14-val_loss=0.314.pth', map_location=torch.device('cpu'))
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict, strict=False)
+        
+        # if name == "LFQ":
+        #     from step1_VQ.model_interface import MInterface
+        #     pretrain_args = OmegaConf.load("/huyuqi/xmyu/DiffSDS/Pretrain_lightning/results/ESMFVQ/LFQ_seg_linear/configs/10-17T15-46-37-project.yaml")
+        #     pretrain_args.diffusion = False
+        #     self.pretrain_model = MInterface(**pretrain_args)
+        #     ckpt = torch.load('/huyuqi/xmyu/DiffSDS/Pretrain_lightning/results/ESMFVQ/LFQ_seg_linear/checkpoints/best-epoch=14-val_loss=0.161.pth', map_location=torch.device('cpu'))
+        #     state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+        #     self.pretrain_model.load_state_dict(state_dict, strict=False)
+            
+
+
+        if name == "softgroup-1":
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/ESMSoftGroup/softgroup-1/configs/12-16T14-57-28-project.yaml")
+            pretrain_args.diffusion = False
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load('/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/ESMSoftGroup/softgroup-1/checkpoints/best-epoch=13-val_loss=0.111.pth', map_location=torch.device('cpu'))
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+        if name == "softgroup-2":
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/DiffSDS/Pretrain_lightning/results/ESMSoftGroup/softgroup-2/configs/10-24T12-51-57-project.yaml")
+            pretrain_args.diffusion = False
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load('/huyuqi/xmyu/DiffSDS/Pretrain_lightning/results/ESMSoftGroup/softgroup-2/checkpoints/best-epoch=14-val_loss=0.067.pth', map_location=torch.device('cpu'))
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+        if name == "softgroup-3":
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/DiffSDS/Pretrain_lightning/results/ESMSoftGroup/softgroup-3/configs/10-25T00-04-15-project.yaml")
+            pretrain_args.diffusion = False
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load('/huyuqi/xmyu/DiffSDS/Pretrain_lightning/results/ESMSoftGroup/softgroup-3/checkpoints/best-epoch=14-val_loss=0.063.pth', map_location=torch.device('cpu'))
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+        if name == "softgroup-4":
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/DiffSDS/Pretrain_lightning/results/ESMSoftGroup/softgroup_32_vectors/configs/10-19T01-03-55-project.yaml")
+            pretrain_args.diffusion = False
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load('/huyuqi/xmyu/DiffSDS/Pretrain_lightning/results/ESMSoftGroup/softgroup_32_vectors/checkpoints/best-epoch=14-val_loss=0.056.pth', map_location=torch.device('cpu'))
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict, strict=False)
+        
+        if name == "softgroup-5":
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/DiffSDS/Pretrain_lightning/results/ESMSoftGroup/softgroup-5-gzy/configs/10-27T17-15-56-project.yaml")
+            pretrain_args.diffusion = False
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load('/huyuqi/xmyu/DiffSDS/Pretrain_lightning/results/ESMSoftGroup/softgroup-5-gzy/checkpoints/best-epoch=14-val_loss=0.039.pth', map_location=torch.device('cpu'))
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+        if name == "softgroup-6":
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/DiffSDS/Pretrain_lightning/results/ESMSoftGroup/softgroup_128_group/configs/10-28T01-28-50-project.yaml")
+            pretrain_args.diffusion = False
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load('/huyuqi/xmyu/DiffSDS/Pretrain_lightning/results/ESMSoftGroup/softgroup_128_group/checkpoints/best-epoch=14-val_loss=0.011.pth', map_location=torch.device('cpu'))
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+
+
+        if name == "softgroup_128_group":
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/DiffSDS/Pretrain_lightning/results/ESMSoftGroup/softgroup_128_group/configs/10-28T01-28-50-project.yaml")
+            pretrain_args.diffusion = False
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load('/huyuqi/xmyu/DiffSDS/Pretrain_lightning/results/ESMSoftGroup/softgroup_128_group/checkpoints/best-epoch=14-val_loss=0.011.pth', map_location=torch.device('cpu'))
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+
+        if name == "diff-softgroup-1":
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/DiffESMSoftGroup/diff-softgroup-rm-dist/configs/12-17T14-19-21-project.yaml")
+            pretrain_args.diffusion = True
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load('/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/DiffESMSoftGroup/diff-softgroup-rm-dist/checkpoints/best-epoch=12-val_loss=0.496.pth', map_location=torch.device('cpu'))
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+        if name == "diff-softgroup-4":
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/DiffESMSoftGroup/diff-softgroup-vq32/configs/12-19T01-54-15-project.yaml")
+            pretrain_args.diffusion = True
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load('/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/DiffESMSoftGroup/diff-softgroup-vq32/checkpoints/best-epoch=13-val_loss=0.184.pth', map_location=torch.device('cpu'))
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+        if name == "diff-softgroup-5":
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/DiffESMSoftGroup/diff-softgroup-vq64/configs/12-19T01-57-07-project.yaml")
+            pretrain_args.diffusion = True
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load('/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/DiffESMSoftGroup/diff-softgroup-vq64/checkpoints/best-epoch=13-val_loss=0.100.pth', map_location=torch.device('cpu'))
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+        if name == "diff-softgroup-6":
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/DiffESMSoftGroup/diff-softgroup-vq128/configs/12-19T10-47-37-project.yaml")
+            pretrain_args.diffusion = True
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load('/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/DiffESMSoftGroup/diff-softgroup-vq128/checkpoints/best-epoch=13-val_loss=0.081.pth', map_location=torch.device('cpu'))
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+        if name == 'vanilla-1':
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/DiffSDS/Pretrain_lightning/results/ESMVQ/base/configs/10-18T01-15-37-project.yaml")
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load("/huyuqi/xmyu/DiffSDS/Pretrain_lightning/results/ESMVQ/base/checkpoints/best-epoch=14-val_loss=0.314.pth")
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+
+        if name == 'soft-1':
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/ESMSoft/soft_rerun/configs/12-10T12-38-16-project.yaml")
+            pretrain_args.diffusion=False
+            pretrain_args.attn_type = 'raw'
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/ESMSoft/soft_rerun/checkpoints/best-epoch=14-val_loss=0.018.pth")
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+
+        if name == 'soft_64_vecs':
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/DiffSDS/Pretrain_lightning/results/ESMSoft/soft_vq_num64/configs/10-19T11-11-58-project.yaml")
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load("/huyuqi/xmyu/DiffSDS/Pretrain_lightning/results/ESMSoft/soft_vq_num64/checkpoints/best-epoch=14-val_loss=8.768.pth")
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+        if name == 'LFQ':
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/DiffSDS/Pretrain_lightning/results/ESMFVQ/vanilla_L1loss/configs/10-24T01-36-37-project.yaml")
+            pretrain_args.diffusion = False
+            pretrain_args.attn_type = 'raw'
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load("/huyuqi/xmyu/DiffSDS/Pretrain_lightning/results/ESMFVQ/vanilla_L1loss/checkpoints/best-epoch=14-val_loss=11.328.pth")
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict, strict=False)
+        
+        if name == 'SCQ-mlp3-vqdim32':
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/ESMSoftBV/SoftBV-vq16-conditional-mlp3-vqdim32/configs/12-22T07-52-47-project.yaml")
+            pretrain_args.diffusion = False
+            pretrain_args.vq_dim, pretrain_args.condition_layer, pretrain_args.sphere = 32, 3, False
+
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/ESMSoftBV/SoftBV-vq16-conditional-mlp3-vqdim32/checkpoints/best-epoch=14-val_loss=0.376.pth")
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+        if name == 'SCQ-mlp3-vqdim32-sphere':
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/ESMSoftBV/SoftBV-mlp3-vqdim32-sphere/configs/12-22T10-44-46-project.yaml")
+            pretrain_args.diffusion = False
+
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/ESMSoftBV/SoftBV-mlp3-vqdim32-sphere/checkpoints/best-epoch=14-val_loss=0.454.pth")
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+        if name == 'SCQ-mlp6-vqdim32-sphere':
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/ESMSoftBV/SoftBV-mlp6BN-vqdim32-sphere/configs/12-22T18-28-04-project.yaml")
+            pretrain_args.diffusion = False
+            pretrain_args.attn_type = 'raw'
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/ESMSoftBV/SoftBV-mlp6BN-vqdim32-sphere/checkpoints/best-epoch=14-val_loss=0.148.pth")
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+        if name == 'SCQ-mlp2-vqdim32':
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/ESMSoftBV/SoftBV-vq16-conditional-mlp2-vqdim32/configs/12-22T00-21-35-project.yaml")
+            pretrain_args.diffusion = False
+            pretrain_args.vq_dim, pretrain_args.condition_layer, pretrain_args.sphere = 32, 2, False
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/ESMSoftBV/SoftBV-vq16-conditional-mlp2-vqdim32/checkpoints/best-epoch=14-val_loss=0.362.pth")
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+        if name == 'SCQ-mlp2-vqdim32-sphere':
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/ESMSoftBV/SoftBV-vq16-conditional-sphere-vqdim32/configs/12-22T00-06-35-project.yaml")
+            pretrain_args.diffusion = False
+            pretrain_args.vq_dim, pretrain_args.condition_layer, pretrain_args.sphere = 32, 2, True
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/ESMSoftBV/SoftBV-vq16-conditional-sphere-vqdim32/checkpoints/best-epoch=14-val_loss=0.338.pth")
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+        if name == 'SCQ-mlp2-vqdim16':
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/ESMSoftBV/SoftBV-vq16-conditional/configs/12-21T13-13-11-project.yaml")
+            pretrain_args.diffusion = False
+            pretrain_args.vq_dim, pretrain_args.condition_layer, pretrain_args.sphere = 16, 2, False
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/ESMSoftBV/SoftBV-vq16-conditional/checkpoints/best-epoch=14-val_loss=0.094.pth")
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+        if name == 'SCQ-mlp2-vqdim16-sphere':
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/ESMSoftBV/SoftBV-vq16-conditional-sphere/configs/12-21T16-38-57-project.yaml")
+            pretrain_args.diffusion = False
+            pretrain_args.vq_dim, pretrain_args.condition_layer, pretrain_args.sphere = 16, 2, True
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/ESMSoftBV/SoftBV-vq16-conditional-sphere/checkpoints/best-epoch=14-val_loss=1.080.pth")
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+        if name == 'SCQ-vq8-mlp6-vqdim16-sphere':
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/ESMSoftBV/SoftBV-vq8-mlp6BN-vqdim32-sphere/configs/12-23T05-15-56-project.yaml")
+            pretrain_args.diffusion = False
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/ESMSoftBV/SoftBV-vq8-mlp6BN-vqdim32-sphere/checkpoints/best-epoch=14-val_loss=0.892.pth")
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+
+        if name == 'SCQ-mlp9-vqdim32-sphere':
+            from step1_VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/ESMSoftBV/SoftBV-mlp9BN-vqdim32-sphere/configs/12-23T16-20-07-project.yaml")
+            pretrain_args.diffusion = False
+
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load("/huyuqi/xmyu/VQProteinFormer/step1_VQ/results/ESMSoftBV/SoftBV-mlp9BN-vqdim32-sphere/checkpoints/best-epoch=14-val_loss=0.151.pth")
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+
+        
+        if name == 'AF2VQ':
+            from step3_AF2VQ.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step3_AF2VQ/results/AF2VQ_softgroup16/configs/12-13T07-59-50-project.yaml")
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load("/huyuqi/xmyu/VQProteinFormer/step3_AF2VQ/results/AF2VQ_softgroup16/checkpoints/best-epoch=11-val_loss=0.812.pth")
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+
+        if name == "ProGLM":
+            self.vq_dim=480
+            from step2_ProGLM.model.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/DiffSDS/Inpainting_representation/results/softgroup_bin_1127/version_4/hparams.yaml")
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load('/huyuqi/xmyu/DiffSDS/Inpainting_representation/results/softgroup_bin_1127/checkpoints/best-epoch=08-valid_acc=0.804.ckpt', map_location=torch.device('cpu'))['state_dict']
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+        if name == 'ProGLM_softgroup_af2db':
+            from step2_ProGLM.model.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/DiffSDS/Inpainting_representation/results/softgroup_bin_2/version_3/hparams.yaml")
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load('/huyuqi/xmyu/DiffSDS/Inpainting_representation/results/softgroup_bin_2/checkpoints/best-epoch=13-valid_acc=0.863.ckpt', map_location=torch.device('cpu'))['state_dict']
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+        if name == 'ProGLM_SoftVQ_cath':
+            from step2_ProGLM.model.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step2_ProGLM/results/ProGLM_SoftVQ_epoch15_pad300/configs/12-25T01-20-35-project.yaml")
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load('/huyuqi/xmyu/VQProteinFormer/step2_ProGLM/results/ProGLM_SoftVQ_epoch15_pad300/checkpoints/best-epoch=27-valid_acc=0.001.pth')
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+
+
+        if name == 'ProGLM_SoftCVQ_cath':
+            from step2_ProGLM.model.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step2_ProGLM/results/ProGLM_SoftCVQ_epoch15_pad300_BCE/configs/12-25T01-42-37-project.yaml")
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load('/huyuqi/xmyu/VQProteinFormer/step2_ProGLM/results/ProGLM_SoftCVQ_epoch15_pad300_BCE/checkpoints/best-epoch=14-valid_acc=0.614.pth')
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+
+        if name == 'ProGLM_SoftCVQ_cath_inpaint':
+            from step2_ProGLM.model.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step2_ProGLM/results/ProGLM_SoftCVQ_epoch15_pad300_BCE_inpaint/configs/12-25T07-47-52-project.yaml")
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load('/huyuqi/xmyu/VQProteinFormer/step2_ProGLM/results/ProGLM_SoftCVQ_epoch15_pad300_BCE_inpaint/checkpoints/best-epoch=14-valid_acc=0.616.pth')
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+        if name == 'ProGLM_SoftCVQ_AF2DB':
+            from step2_ProGLM.model.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step2_ProGLM/results/ProGLM_SoftCVQ_epoch15_AF2DB/configs/12-25T13-01-12-project.yaml")
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load('/huyuqi/xmyu/VQProteinFormer/step2_ProGLM/results/ProGLM_SoftCVQ_epoch15_AF2DB/checkpoints/best-epoch=14-valid_acc=0.631.pth')
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+
+        if name == 'ProGLM_SoftCVQ_ESM1B_CATH':
+            from step2_ProGLM.model.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step2_ProGLM/results/ProGLM_SoftCVQ_ESM1B_CATH_lr5e-5/configs/12-25T16-02-35-project.yaml")
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load('/huyuqi/xmyu/VQProteinFormer/step2_ProGLM/results/ProGLM_SoftCVQ_ESM1B_CATH_lr5e-5/checkpoints/best-epoch=14-valid_acc=0.616.pth')
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+
+        if name == 'ProGLM_SoftCVQ_CATH':
+            from step2_ProGLM.model.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step2_ProGLM/results/ProGPT_SoftCVQ_CATH/configs/12-26T08-13-41-project.yaml")
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load('/huyuqi/xmyu/VQProteinFormer/step2_ProGLM/results/ProGPT_SoftCVQ_CATH/checkpoints/best-epoch=14-gpt_acc=0.758.pth')
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+        if name == 'ProGLM_SoftCVQ_CATH_epoch10k':
+            from step2_ProGLM.model.model_interface import MInterface
+            pretrain_args = OmegaConf.load("/huyuqi/xmyu/VQProteinFormer/step2_ProGLM/results/ProGPT_SoftCVQ_CATH_epoch1000/configs/12-27T02-36-49-project.yaml")
+            self.pretrain_model = MInterface(**pretrain_args)
+            ckpt = torch.load('/huyuqi/xmyu/VQProteinFormer/step2_ProGLM/results/ProGPT_SoftCVQ_CATH_epoch10000_resume/checkpoints/best-epoch=1887-gpt_loss=0.181.pth')
+            state_dict = {k.replace('_forward_module.', ''): v for k, v in ckpt.items()}
+            self.pretrain_model.load_state_dict(state_dict)
+        
+        
+        if name == 'GearNet':
+            from model.PretrainGearNet import PretrainGearNet_Model
+            self.pretrain_model = PretrainGearNet_Model()
+        
+        self.pretrain_model.eval()
+    
+    def get_vq_id(self, seqs, angles, attn_mask):
+        # if ('softgroup' in self.name) or ('LFQ' in self.name):
+        #     h_input = self.pretrain_model.model.input(seqs.squeeze(-1), angles)
+        #     h_enc = self.pretrain_model.model.ProteinEnc(h_input, attn_mask, None).last_hidden_state          
+        #     vq_id, e_enc = self.pretrain_model.model.VQLayer.get_vq(h_enc, attn_mask, temperature=1e-5)
+        #     return F.pad(vq_id, [0,1,0,0])
+
+        h_input = self.pretrain_model.model.input(seqs.squeeze(-1), angles)
+        h_enc = self.pretrain_model.model.ProteinEnc(h_input, attn_mask, None).last_hidden_state
+        vq_id, e_enc = self.pretrain_model.model.VQLayer.get_vq(h_enc, attn_mask, temperature=1e-5)
+        return vq_id
+    
+    def forward(self, batch):
+        if self.name in ["ESM35M", "ESM650M", "ESM3B"]:
+            seqs, attn_mask = batch['seqs'], batch['attn_mask']
+            outputs = self.pretrain_model.model(input_ids=seqs[:,:,0], attention_mask=attn_mask)
+            pretrain_embedding = outputs.hidden_states
+            pretrain_embedding = pretrain_embedding.reshape(-1,self.esm_dim)[attn_mask.view(-1)==1]
+            return pretrain_embedding
+        if self.name in ["softgroup_128_group"]:
+            seqs, angles, attn_mask = batch['seqs'], batch['angles'] , batch['attn_mask']
+            vq_id = self.pretrain_model.model.get_vqid(seqs[...,0], angles, attn_mask)
+            return vq_id
+        if self.name in ["ProGLM"]:
+            vq_id, attn_mask, seg, pos = batch['vq_id'], batch['attn_mask'], batch['seg'], batch['pos']
+            feat = self.pretrain_model.model.get_feat(vq_id, attn_mask, seg, pos)
+            feat = feat.reshape(-1,self.vq_dim)[attn_mask.view(-1)==1]
+            return feat
+        if self.name in ["GearNet"]:
+            seqs = batch['seqs']
+            batch = batch['batch']
+            attn_mask = batch['attn_mask']
+            for idx in range(seqs.shape[0]):
+                seq_str = self.pretrain_featurizer.ESM_tokenizer.decode(seqs[idx,attn_mask[idx,:].bool(),0])
+                seq_strs.append(seq_str.split(" "))
+            seq_strs = sum(seq_strs, [])
+            node_index = torch.arange(batch.batch.shape[0], device=batch.batch.device)
+            node2graph = batch.batch
+            chain_id = torch.ones_like(batch.batch)
+
+            pretrain_embedding = self.pretrain_gearnet_model(seq_strs, node_index, node2graph, chain_id, batch.pos)
+            return pretrain_embedding
+
+
+
+

Flexpert-Design/src/models/E3PiFold_model.py

@@ -0,0 +1,90 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Optional
+from torch import Tensor
+from omegaconf import OmegaConf
+from src.modules.E3PiFold import GaussianEncoder, TransformerEncoderWithPair
+from src.tools import gather_nodes, _dihedrals, _get_rbf, _get_dist, _rbf, _orientations_coarse_gl_tuple
+
+class E3PiFold(nn.Module):
+    def __init__(self, config) -> None:
+        super().__init__()
+        self.node_embed = nn.Linear(21, config.embed_dim)
+        self.protein_embedder = GaussianEncoder(config.kernel_num, config.embed_dim, config.attention_heads, config.use_dist, config.use_product)
+
+        self.encoder = TransformerEncoderWithPair(
+            config.encoder_layers,
+            config.embed_dim,
+            config.ffn_embed_dim,
+            config.attention_heads,
+            config.emb_dropout,
+            config.dropout,
+            config.attention_dropout,
+            config.activation_dropout,
+            config.max_seq_len,
+        )
+        self.predictor = nn.Linear(config.embed_dim, 33)
+    
+    def _full_dist(self, X, mask, top_k=30, eps=1E-6):
+        mask_2D = torch.unsqueeze(mask,1) * torch.unsqueeze(mask,2)
+        dX = torch.unsqueeze(X,1) - torch.unsqueeze(X,2)
+        D = (1. - mask_2D)*10000 + mask_2D* torch.sqrt(torch.sum(dX**2, 3) + eps)
+
+        D_max, _ = torch.max(D, -1, keepdim=True)
+        D_adjust = D + (1. - mask_2D) * (D_max+1)
+        D_neighbors, E_idx = torch.topk(D_adjust, min(top_k, D_adjust.shape[-1]), dim=-1, largest=False)
+        return D_neighbors, E_idx 
+    
+    def _get_features(self, batch):
+        X = batch['X']
+        X_ca = X[:,:,1,:]
+        D_neighbors, E_idx = self._full_dist(X_ca, batch['mask'], 30)
+        V_angles = _dihedrals(X.float())
+        V_direct, E_direct, E_angles = _orientations_coarse_gl_tuple(X.float(), E_idx)
+        h_V = torch.cat([V_angles, V_direct], dim=-1).to(X.dtype)
+        batch['h_V'] = h_V 
+        return batch
+    
+    def forward(self, batch):
+        '''
+        X, H, seq_mask
+        '''
+        X = batch['X'][:,:,1]
+        H = self.node_embed(batch['h_V'])
+        seq_mask = batch['mask']
+        pair_mask = seq_mask[..., None] * seq_mask[..., None, :]
+        padding_mask = 1 - seq_mask
+        x, graph_attn_bias = self.protein_embedder(X, H, pair_mask)
+        (
+            encoder_rep, 
+            encoder_pair_rep,
+            delta_encoder_pair_rep,
+            x_norm,
+            delta_encoder_pair_rep_norm,
+        ) = self.encoder(x, padding_mask=padding_mask, attn_mask=graph_attn_bias, pair_mask=pair_mask)
+        logits = self.predictor(x)
+        log_probs = F.log_softmax(logits, dim=-1)
+
+        return {'log_probs': log_probs}
+
+
+if __name__ == '__main__':
+    B, N, dim = 16, 512, 768
+    X = torch.randn(B, N, 3)
+    H = torch.randn(B, N, dim)
+    seq_mask = (torch.ones(B, N)>0.5).float()
+
+    config = {'encoder_layers': 12,
+              'kernel_num':16,
+              'embed_dim': 768,
+              'ffn_embed_dim': 3072,
+              'attention_heads': 8,
+              'emb_dropout': 0.1,
+              'dropout': 0.1,
+              'attention_dropout': 0.1,
+              'activation_dropout': 0.0,
+              'max_seq_len': 256}
+    config = OmegaConf.create(config)
+    model = E3PiFold(config)
+    feat = model(X, H, seq_mask)

Flexpert-Design/src/models/MemoryESM.py

@@ -0,0 +1,164 @@
+import subprocess
+import os
+from joblib import Parallel, delayed, cpu_count
+from tqdm import tqdm
+import pandas as pd
+import torch
+import torch.nn as nn
+import shutil
+from .PretrainESM_model import PretrainESM_Model
+
+
+class MemoESM(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.PretrainESM = PretrainESM_Model(args)
+        self.tokenizer = self.PretrainESM.tokenizer
+        self.memory = {}
+        # self.fix_memory = False
+    
+    # def save_memory(self, path):
+    #     params = {key:val for key,val in self.state_dict().items() if "GNNTuning" in key}
+    #     torch.save({"params":params,"memory": self.memory}, path)
+
+    # def load_memory(self, path):
+    #     data = torch.load(path)
+    #     self.load_state_dict(data['params'], strict=False)
+    #     self.memory = data['memory']
+        
+    def clean_input(self, batch, score_cut=0.99):
+        '''
+        require: batch['pred_ids'], batch['attention_mask'], batch['confs']
+        '''
+        symbol = "<mask>"
+        replace_dict = {"-":symbol,
+                        ".":symbol,
+                        "<eos>":symbol,
+                        "<unk>":symbol,
+                        "<cls>":symbol,
+                        "<pad>":symbol,
+                        "<null_1>":symbol,
+                        "<mask>":symbol,
+                        "U":symbol,
+                        "O":symbol}
+
+        device = batch['pred_ids'].device
+        query_seqs = []
+        for pred_ids, mask, score in zip(batch['pred_ids'], batch['attention_mask'], batch['confs']):
+            seq = self.tokenizer.decode(pred_ids[mask], clean_up_tokenization_spaces=False)
+            elements = []
+            for idx, x in enumerate(seq.split(" ")):
+                symbol = replace_dict.get(x, x)
+                if score[idx] < score_cut:
+                    symbol = "<mask>"
+                elements.append(symbol)
+            seq = "".join(elements)
+            query_seqs.append(seq)
+        
+        results = self.tokenizer.batch_encode_plus(query_seqs, return_tensors="pt", padding=True)
+        return query_seqs
+    
+    def initoutput(self, B, maxL, device):
+        self.out_pred_ids = torch.zeros(B, maxL, dtype=torch.long, device=device)
+        self.out_confs = torch.zeros(B, maxL, dtype=torch.float, device=device)
+        self.out_embeds = torch.zeros(B, maxL, 1280, dtype=torch.float, device=device)
+        self.titles = [None for i in range(B)]
+        
+    def retrivel(self, titles, num_nodes, device, use_memory):
+        # retrieval
+        unseen = []
+        for idx in range(len(titles)):
+            name = titles[idx]
+            if (name in self.memory) and use_memory:
+                memo_pred_ids = self.memory[name]['pred_ids'].to(device)
+                memo_confs = self.memory[name]['confs'].to(device)
+                memo_embeds = self.memory[name]['embeds'].to(device)
+                
+                self.out_pred_ids[idx, :num_nodes[idx]] = memo_pred_ids
+                self.out_confs[idx, :num_nodes[idx]] = memo_confs
+                self.out_embeds[idx, :num_nodes[idx]] = memo_embeds
+                self.titles[idx] = name
+            else:
+                unseen.append(idx)
+        return unseen
+    
+    def rebatch(self, unseen, batch):
+        unseen_pred_ids = []
+        unseen_attention_mask = []
+        for i in unseen:
+            unseen_pred_ids.append(batch['pred_ids'][i])
+            unseen_attention_mask.append(batch['attention_mask'][i])
+        unseen_pred_ids = torch.stack(unseen_pred_ids)
+        unseen_attention_mask = torch.stack(unseen_attention_mask)
+        return {"pred_ids":unseen_pred_ids, "attention_mask":unseen_attention_mask}
+    
+    def save2memory(self, unseen,outputs, titles, unseen_attention_mask):
+         # save to memory
+        for i in range(len(unseen)):
+            name = titles[unseen[i]]
+            self.titles[unseen[i]] = name
+            mask = unseen_attention_mask[i]
+            self.memory[name] = {"pred_ids":outputs['pred_ids'][i][mask].detach().to('cpu'), 
+                                "confs":outputs['confs'][i][mask].detach().to('cpu'), 
+                                "embeds":outputs['embeds'][i][mask].detach().to('cpu')}
+    
+    def update(self, unseen, unseen_attention_mask, num_nodes, outputs):
+        # update
+        for idx in range(len(unseen)):
+            mask = unseen_attention_mask[idx]==1
+            self.out_pred_ids[unseen[idx], :num_nodes[unseen[idx]]] = outputs['pred_ids'][idx][mask]
+            self.out_confs[unseen[idx], :num_nodes[unseen[idx]]] = outputs['confs'][idx][mask]
+            self.out_embeds[unseen[idx], :num_nodes[unseen[idx]]] = outputs['embeds'][idx][mask]
+    
+    @torch.no_grad()
+    def forward(self, batch, use_memory=False):
+        # debatch
+        # clean_seqs = self.clean_input(batch)
+        device = batch['probs'].device
+        B, maxL, _ = batch['probs'].shape
+        num_nodes = batch['attention_mask'].sum(dim=-1).tolist()
+        self.initoutput(B, maxL, device)
+        unseen = self.retrivel(batch['title'], num_nodes, device, use_memory)
+        
+        
+        if len(unseen)>0:
+            # batch forward
+            new_batch = self.rebatch(unseen, batch)
+            outputs = self.PretrainESM(new_batch)
+            
+            self.save2memory(unseen,outputs, batch['title'], new_batch['attention_mask'])
+            self.update(unseen, new_batch['attention_mask'], num_nodes, outputs)
+        
+        return {'title':self.titles,'pred_ids':self.out_pred_ids, 'confs':self.out_confs, 'embeds':self.out_embeds, 'attention_mask':batch['attention_mask']}
+
+
+
+if __name__ == '__main__': 
+    
+    # work_space = '/gaozhangyang/experiments/PiFoldV2/data/mmseq_workspace2'
+    # target_seqs = ["MVLSPADKTNVKAAWGKVGAHAGEYGAEALERMFLSFPQTKTYFPHFDLSHGSAQVKGHG", "MVHLTPEEKSAVTALWGKVNVDEVGVEALGRLLVVYPWTQRFFESFGDLSTPDAVMGNPKV",
+    #  "MVLSPADKTNVKAAWGKVGAGGAEALERMFLSFPQKTYYTYFPHFDLSHGSAQVKGHG"]
+
+    # query_seqs = ["MVLSPADKTNVKAAWGKVGAHAGEYGAEALERMFLSFPTTKFPHFDLSHGSAQV", "MVHLTPEEKSAVTALWGKVNVDEVGGGRLLVVYPWTQRFFESFGDLSTPDAV",]
+    
+    # results = search_seqs(query_seqs, target_seqs, work_space)
+    # print(results)
+        
+
+    import biotite.sequence as seq
+    import biotite.sequence.align as align
+
+    # Create example query and target protein sequences
+    query_seq1 = seq.ProteinSequence("MSKXXKAFLNKXXL")
+    target_seq1 = seq.ProteinSequence("MSKVKAALNKVLL")
+    target_seq2 = seq.ProteinSequence("MSKVKKALNKVLL")
+    target_seq3 = seq.ProteinSequence("MSTVAAALKMLLL")
+
+    results = search_seqs_biotite(["MSKXXKAFLNKXXL"], ["MSKVKAALNKVLL", "MSKVKKALNKVLL", "MSTVAAALKMLLL"])
+
+    # Print the alignments
+    print("Query alignments:")
+
+
+
+    

Flexpert-Design/src/models/MemoryESMIF.py

@@ -0,0 +1,116 @@
+import subprocess
+import os
+from joblib import Parallel, delayed, cpu_count
+from tqdm import tqdm
+import pandas as pd
+import torch
+import torch.nn as nn
+from .PretrainESMIF_model import PretrainESMIF_Model
+from torch_scatter import scatter_sum
+
+class MemoESMIF(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.PretrainESMIF = PretrainESMIF_Model()
+        self.memory = {}
+        # self.fix_memory = False
+    
+    # def save_memory(self, path):
+    #     params = {key:val for key,val in self.state_dict().items() if "GNNTuning" in key}
+    #     torch.save({"params":params,"memory": self.memory}, path)
+
+    # def load_memory(self, path):
+    #     data = torch.load(path)
+    #     self.load_state_dict(data['params'], strict=False)
+    #     self.memory = data['memory']
+        
+    def initoutput(self, B, maxL, device):
+        self.out_embeds = torch.zeros(B, maxL, 512, dtype=torch.float, device=device)
+        self.titles = [None for i in range(B)]
+        
+    def retrivel(self, titles, num_nodes, device, use_memory):
+        # retrieval
+        unseen = []
+        for idx in range(len(titles)):
+            name = titles[idx]
+            if (name in self.memory) and use_memory:
+                memo_embeds = self.memory[name]['embeds'].to(device)
+                self.out_embeds[idx, :num_nodes[idx]] = memo_embeds
+                self.titles[idx] = name
+            else:
+                unseen.append(idx)
+        return unseen
+    
+    def rebatch(self, unseen, batch):
+        unseen_position = []
+        for i in unseen:
+            mask = batch['batch_id']==i
+            unseen_position.append(batch['position'][mask][:,:3,:])
+        return {"position":unseen_position}
+    
+    def save2memory(self, unseen,outputs, titles, num_nodes):
+         # save to memory
+        for i in range(len(unseen)):
+            name = titles[unseen[i]]
+            self.titles[unseen[i]] = name
+            num = num_nodes[unseen[i]]
+            self.memory[name] = {"embeds":outputs['feat'][i,:num].detach().to('cpu')}
+    
+    def update(self, unseen, num_nodes, outputs):
+        # update
+        for idx in range(len(unseen)):
+            num = num_nodes[unseen[idx]]
+            self.out_embeds[unseen[idx], :num_nodes[unseen[idx]]] = outputs['feat'][idx, :num]
+    
+    @torch.no_grad()
+    def forward(self, batch, use_memory=False):
+        # debatch
+        # clean_seqs = self.clean_input(batch)
+        device = batch['position'].device
+        num_nodes = scatter_sum(torch.ones_like(batch['batch_id']), batch['batch_id'], dim=0)
+        B, maxL = num_nodes.shape[0], num_nodes.max()
+        self.initoutput(B, maxL, device)
+        unseen = self.retrivel(batch['title'], num_nodes, device, use_memory)
+        
+        
+        if len(unseen)>0:
+            # batch forward
+            new_batch = self.rebatch(unseen, batch)
+            outputs = self.PretrainESMIF(new_batch['position'])
+            
+            self.save2memory(unseen,outputs, batch['title'], num_nodes)
+            self.update(unseen, num_nodes, outputs)
+        
+        return {'title':self.titles, 'embeds':self.out_embeds}
+
+
+
+if __name__ == '__main__': 
+    
+    # work_space = '/gaozhangyang/experiments/PiFoldV2/data/mmseq_workspace2'
+    # target_seqs = ["MVLSPADKTNVKAAWGKVGAHAGEYGAEALERMFLSFPQTKTYFPHFDLSHGSAQVKGHG", "MVHLTPEEKSAVTALWGKVNVDEVGVEALGRLLVVYPWTQRFFESFGDLSTPDAVMGNPKV",
+    #  "MVLSPADKTNVKAAWGKVGAGGAEALERMFLSFPQKTYYTYFPHFDLSHGSAQVKGHG"]
+
+    # query_seqs = ["MVLSPADKTNVKAAWGKVGAHAGEYGAEALERMFLSFPTTKFPHFDLSHGSAQV", "MVHLTPEEKSAVTALWGKVNVDEVGGGRLLVVYPWTQRFFESFGDLSTPDAV",]
+    
+    # results = search_seqs(query_seqs, target_seqs, work_space)
+    # print(results)
+        
+
+    import biotite.sequence as seq
+    import biotite.sequence.align as align
+
+    # Create example query and target protein sequences
+    query_seq1 = seq.ProteinSequence("MSKXXKAFLNKXXL")
+    target_seq1 = seq.ProteinSequence("MSKVKAALNKVLL")
+    target_seq2 = seq.ProteinSequence("MSKVKKALNKVLL")
+    target_seq3 = seq.ProteinSequence("MSTVAAALKMLLL")
+
+    results = search_seqs_biotite(["MSKXXKAFLNKXXL"], ["MSKVKAALNKVLL", "MSKVKKALNKVLL", "MSTVAAALKMLLL"])
+
+    # Print the alignments
+    print("Query alignments:")
+
+
+
+    

Flexpert-Design/src/models/MemoryPiFold.py

@@ -0,0 +1,143 @@
+import torch
+import torch.nn as nn
+from .PretrainPiFold_model import PretrainPiFold_Model
+from torch_scatter import scatter_sum
+import torch.nn.functional as F
+
+class MemoPiFold_model(nn.Module):
+    def __init__(self, args):
+        super().__init__()
+        self.PretrainPiFold = PretrainPiFold_Model(args)
+        self.memory = {}
+    
+    def save_memory(self, path):
+        params = {key:val for key,val in self.state_dict().items() if "GNNTuning" in key}
+        torch.save({"params":params,"memory": self.memory}, path)
+
+    def load_memory(self, path):
+        data = torch.load(path)
+        self.load_state_dict(data['params'], strict=False)
+        self.memory = data['memory']
+        
+    def initoutput(self, B, max_L, nums, device):
+        self.confs = torch.ones(B, max_L, device=device)
+        self.embeds = torch.ones(B, max_L, 128, device=device)
+        self.probs = torch.ones(B, max_L, 33, device=device)
+        self.attention_mask = torch.ones_like(self.confs)==0
+        self.titles = [None for i in range(B)]
+        for id, num in enumerate(nums):
+            self.attention_mask[id, :num] = True
+        self.edge_feats = []
+        
+    def retrivel(self, batch, nums, batch_uid, device, use_memory):
+        # retrieval
+        unseen = []
+        
+        for idx, name in enumerate(batch['title']):
+            if (name in self.memory) and use_memory:
+                try:
+                    self.confs[batch_uid[idx],:nums[idx]] = self.memory[name]['conf'].to(device)
+                except:
+                    self.confs[batch_uid[idx],:nums[idx]] = self.memory[name]['conf'].to(device)
+                self.embeds[batch_uid[idx],:nums[idx]] = self.memory[name]['embed'].to(device)
+                self.probs[batch_uid[idx],:nums[idx]] = self.memory[name]['prob'].to(device)
+                self.edge_feats.append((batch_uid[idx], self.memory[name]['h_E'].to(device)))
+                self.titles[batch_uid[idx]] = name
+            else:
+                unseen.append(idx)
+        return unseen
+
+    def rebatch(self, unseen, batch_uid, batch_id, batch, shift, nums, device):
+        h_V2, h_E2, E_idx2, batch_id2 = [], [], [], []
+        shift2 = [0]
+        idx=0
+        for id in batch_uid:
+            if id not in unseen:
+                continue
+            node_mask = batch_id == id
+            edge_mask = batch_id[batch['E_idx'][0]] == id
+            h_V2.append(batch['h_V'][node_mask])
+            h_E2.append(batch['h_E'][edge_mask])
+            new_E_idx = batch['E_idx'][:,edge_mask] 
+            new_E_idx = new_E_idx- shift[batch_id[new_E_idx[0]]]+shift2[-1]
+            E_idx2.append(new_E_idx)
+            new_batch_id = torch.ones(node_mask.sum().long(), device=device)*idx
+            batch_id2.append(new_batch_id)
+            shift2.append(shift2[-1]+nums[id])
+            idx+=1
+        
+        h_V2 = torch.cat(h_V2)
+        h_E2 = torch.cat(h_E2)
+        E_idx2 = torch.cat(E_idx2, dim=-1)
+        batch_id2 = torch.cat(batch_id2).long()
+        return {"h_V":h_V2, 'h_E':h_E2, 'E_idx':E_idx2, 'batch_id':batch_id2}
+    
+    def update_save2memory(self, unseen, batch_id2, E_idx2, batch, pretrain_gnn, max_L):
+        for id in batch_id2.unique():
+            node_mask = batch_id2 == id
+            edge_mask = batch_id2[E_idx2[0]] == id
+            title = batch['title'][unseen[int(id)]]
+            conf = pretrain_gnn['confs'][id]
+            conf = F.pad(conf, (0, max_L-len(conf)))
+            embed = pretrain_gnn['embeds'][id]
+            embed = F.pad(embed, (0,0,0,max_L-len(embed)))
+            prob = pretrain_gnn['probs'][id]
+            prob = F.pad(prob, (0,0,0,max_L-len(prob)))
+            self.edge_feats.append((unseen[int(id)], pretrain_gnn['h_E'][edge_mask]))
+            
+            self.confs[unseen[int(id)]] = conf
+            self.embeds[unseen[int(id)]] = embed
+            self.probs[unseen[int(id)]] = prob
+            self.titles[unseen[int(id)]] = title
+            
+            attn_mask = self.attention_mask[unseen[int(id)]]
+            
+            # save to memory
+            self.memory[title] = {'conf': conf[attn_mask].detach().to('cpu'), 
+                                'embed': embed[attn_mask].detach().to('cpu'), 
+                                'prob': prob[attn_mask].detach().to('cpu'),
+                                'h_E':pretrain_gnn['h_E'][edge_mask].detach().to('cpu')}
+    
+    @torch.no_grad()
+    def forward(self, batch, use_memory=False):
+        batch_id = batch['batch_id']
+        batch_uid = batch_id.unique()
+        device = batch_id.device
+        
+        nums = scatter_sum(torch.ones_like(batch_id), batch_id)
+        shift = torch.cat([torch.zeros(1, device=device), torch.cumsum(nums, dim=0)]).long()
+        max_L, B = nums.max(), batch_uid.shape[0]
+        
+        self.initoutput(B, max_L, nums, device)
+        unseen = self.retrivel(batch, nums, batch_uid, device, use_memory)
+        
+        # organize data
+        if len(unseen)>0:
+            # rebatch
+            new_batch = self.rebatch(unseen, batch_uid, batch_id, batch, shift, nums, device)
+            
+            # forward pass 
+            pretrain_gnn = self.PretrainPiFold(new_batch)
+            
+            self.update_save2memory(unseen, pretrain_gnn['batch_id'], pretrain_gnn['E_idx'], batch, pretrain_gnn, max_L)
+            
+        
+        self.edge_feats = sorted(self.edge_feats, key=lambda x: x[0])
+        self.edge_feats = torch.cat([one[1] for one in self.edge_feats])
+        
+        pred_ids = self.probs.argmax(dim=-1)*self.attention_mask + (~self.attention_mask)*1
+        
+        return {'title': self.titles,
+                'pred_ids': pred_ids,
+                'confs': self.confs, 
+                'embeds': self.embeds, 
+                'probs': self.probs, 
+                'attention_mask': self.attention_mask,
+                'h_E':self.edge_feats,
+                'E_idx': batch['E_idx'],
+                'batch_id': batch['batch_id']}
+            
+            
+    
+    def _get_features(self, S, score, X, mask, chain_mask, chain_encoding):
+        return self.PretrainPiFold._get_features(S, score, X, mask, chain_mask, chain_encoding)

Flexpert-Design/src/models/MemoryTuning.py

@@ -0,0 +1,213 @@
+import torch
+import torch.nn as nn
+from .Tuning import GNNTuning_Model
+
+
+class MemoTuning(nn.Module):
+    def __init__(self, args, tunning_layers_n, tunning_layers_dim, input_design_dim, input_esm_dim, tunning_dropout, tokenizer, fix_memory=False):
+        super().__init__()
+        self.args = args
+        self.tunning_layers_dim = tunning_layers_dim
+        self.GNNTuning = GNNTuning_Model(args, num_encoder_layers=tunning_layers_n, hidden_dim=tunning_layers_dim, input_design_dim=input_design_dim, input_esm_dim=input_esm_dim, dropout = tunning_dropout)
+        self.tokenizer = tokenizer
+        self.memory = {}
+        
+    def save_param_memory(self, path):
+        torch.save({"params":self.state_dict(),"memory": self.memory}, path)
+
+    def load_param_memory(self, path):
+        data = torch.load(path)
+        self.load_state_dict(data['params'])
+        self.memory = data['memory']
+    
+    def get_seqs(self, pred_ids_raw, attention_mask):
+        query_seqs = []
+        for pred_ids, mask in zip(pred_ids_raw, attention_mask):
+            seq = self.tokenizer.decode(pred_ids[mask], clean_up_tokenization_spaces=False)
+            seq = "".join(seq.split(" "))
+            query_seqs.append(seq)
+        return query_seqs
+    
+    def initoutput(self, pretrain_design, B, max_L, device):
+        # initialize output
+        self.out_pred_ids = torch.zeros_like(pretrain_design['pred_ids'])
+        self.out_confs = torch.zeros_like(pretrain_design['confs'])
+        self.out_embeds = torch.zeros(B, max_L, self.tunning_layers_dim, device = device)
+        self.out_attention_mask = torch.zeros_like(pretrain_design['attention_mask'])
+        self.out_probs = torch.zeros_like(pretrain_design['probs'])
+        self.out_log_probs = torch.zeros_like(pretrain_design['probs'])
+        self.titles = [None for i in range(B)]
+        
+    
+    
+    def retrivel(self, keys, num_nodes,device, use_memory):
+        unseen = []
+        for idx in range(len(keys)):
+            key = keys[idx]
+            if (key in self.memory) and use_memory:
+                self.out_pred_ids[idx, :num_nodes[idx]] = self.memory[key]['pred_ids'].to(device)
+                self.out_confs[idx, :num_nodes[idx]] = self.memory[key]['confs'].to(device)
+                self.out_embeds[idx, :num_nodes[idx]] = self.memory[key]['embeds'].to(device)
+                self.out_attention_mask[idx, :num_nodes[idx]] = self.memory[key]['attention_mask'].to(device)
+                self.out_probs[idx, :num_nodes[idx]] = self.memory[key]['probs'].to(device)
+                self.out_log_probs[idx, :num_nodes[idx]] = self.memory[key]['log_probs'].to(device)
+                self.titles[idx] = key
+            else:
+                unseen.append(idx)
+        return unseen
+    
+    def rebatch(self,unseen, batch_id_raw, E_idx_raw, h_E_raw, shift, num_nodes, pretrain_design, pretrain_esm_msa, pretrain_struct, pretrain_esmif, device):
+        unseen_design_pred_ids = []
+        unseen_design_confs = []
+        unseen_design_embeds = []
+        unseen_design_attention_mask = []
+        
+        unseen_esm_pred_ids = []
+        unseen_esm_confs = []
+        unseen_esm_embeds = []
+        unseen_esm_attention_mask = []
+        unseen_struct_embeds = []
+        unseen_esmif_embeds = []
+        h_E = []
+        E_idx = []
+        batch_id = []
+        
+        new_shift = 0
+        for bid, i in enumerate(unseen):
+            edge_mask = batch_id_raw[E_idx_raw[0]] == i
+            h_E.append(h_E_raw[edge_mask])
+            E_idx.append(E_idx_raw[:,edge_mask]-shift[i]+new_shift)
+            batch_id.append(torch.ones(num_nodes[i], device=device).long()*bid)
+            new_shift += num_nodes[i]
+            
+            unseen_design_pred_ids.append(pretrain_design['pred_ids'][i])
+            unseen_design_confs.append(pretrain_design['confs'][i])
+            unseen_design_embeds.append(pretrain_design['embeds'][i])
+            unseen_design_attention_mask.append(pretrain_design['attention_mask'][i])
+            
+            if self.args.use_LM:
+                unseen_esm_pred_ids.append(pretrain_esm_msa['pred_ids'][:,i])
+                unseen_esm_confs.append(pretrain_esm_msa['confs'][:,i])
+                unseen_esm_embeds.append(pretrain_esm_msa['embeds'][:,i])
+                unseen_esm_attention_mask.append(pretrain_esm_msa['attention_mask'][:,i])
+            
+            if self.args.use_gearnet:
+                unseen_struct_embeds.append(pretrain_struct['embeds'][:,i])
+            
+            if self.args.use_esmif:
+                unseen_esmif_embeds.append(pretrain_esmif['embeds'][i])
+            
+            
+        unseen_design_pred_ids = torch.stack(unseen_design_pred_ids)
+        unseen_design_confs = torch.stack(unseen_design_confs)
+        unseen_design_embeds = torch.stack(unseen_design_embeds)
+        unseen_design_attention_mask = torch.stack(unseen_design_attention_mask)
+        
+        if self.args.use_LM:
+            unseen_esm_pred_ids = torch.stack(unseen_esm_pred_ids, dim=1)
+            unseen_esm_confs = torch.stack(unseen_esm_confs, dim=1)
+            unseen_esm_embeds = torch.stack(unseen_esm_embeds, dim=1)
+            unseen_esm_attention_mask = torch.stack(unseen_esm_attention_mask, dim=1)
+        
+        if self.args.use_gearnet:
+            unseen_struct_embeds = torch.stack(unseen_struct_embeds, dim=1)
+        
+        if self.args.use_esmif:
+            unseen_esmif_embeds = torch.stack(unseen_esmif_embeds, dim=0)
+            
+        
+        unseen_batch = {"pretrain_design":
+                            {"pred_ids": unseen_design_pred_ids, 
+                            "confs":unseen_design_confs, 
+                            "embeds": unseen_design_embeds, 
+                            "attention_mask":unseen_design_attention_mask},
+                        "h_E": torch.cat(h_E),
+                        "E_idx": torch.cat(E_idx, dim=1),
+                        "batch_id": torch.cat(batch_id),
+                        "attention_mask":unseen_design_attention_mask
+                        }
+
+        if self.args.use_LM:
+            unseen_batch["pretrain_esm_msa"]={"pred_ids": unseen_esm_pred_ids, 
+                            "confs":unseen_esm_confs, 
+                            "embeds": unseen_esm_embeds, 
+                            "attention_mask":unseen_esm_attention_mask}
+        
+        if self.args.use_gearnet:
+            unseen_batch["pretrain_struct"] = {
+                            "embeds":unseen_struct_embeds}
+        
+        if self.args.use_esmif:
+            unseen_batch["pretrain_esmif"] = {"embeds":unseen_esmif_embeds}
+        return unseen_batch
+    
+    def save2memory(self,keys,unseen,num_nodes, unseen_results):
+        # save to memory
+        for i in range(len(unseen)):
+            key = keys[unseen[i]]
+            num = num_nodes[unseen[i]]
+            self.memory[key] = {"pred_ids":unseen_results['pred_ids'][i][:num].detach().to('cpu'), 
+                                "confs":unseen_results['confs'][i][:num].detach().to('cpu'), 
+                                "embeds":unseen_results['embeds'][i][:num].detach().to('cpu'),
+                                "probs":unseen_results['probs'][i][:num].detach().to('cpu'),
+                                "log_probs":unseen_results['log_probs'][i][:num].detach().to('cpu'),
+                                "attention_mask":unseen_results['attention_mask'][i][:num].detach().to('cpu')}
+    
+    def update(self, unseen, num_nodes, unseen_results, keys):
+        # update
+        for i in range(len(unseen)):
+            num = num_nodes[unseen[i]]
+            self.out_pred_ids[unseen[i], :num] = unseen_results['pred_ids'][i][:num]
+            self.out_confs[unseen[i], :num] = unseen_results['confs'][i][:num]
+            self.out_embeds[unseen[i], :num] = unseen_results['embeds'][i][:num]
+            self.out_probs[unseen[i], :num] = unseen_results['probs'][i][:num]
+            self.out_log_probs[unseen[i], :num] = unseen_results['log_probs'][i][:num]
+            self.titles[unseen[i]] = keys[unseen[i]]
+
+    def forward(self, batch, use_memory=False):
+        self.use_memory = use_memory
+        pretrain_design,  h_E_raw, E_idx_raw, mask_attend, batch_id_raw = batch['pretrain_design'] ,batch['h_E'], batch['E_idx'], batch['attention_mask'], batch['batch_id']
+        device = h_E_raw.device
+        
+        pretrain_esm_msa = None
+        if self.args.use_LM:
+            pretrain_esm_msa = batch['pretrain_esm_msa']
+        
+        pretrain_struct = None
+        if self.args.use_gearnet:
+            pretrain_struct = batch['pretrain_struct']
+        
+        pretrain_esmif = None
+        if self.args.use_esmif:
+            pretrain_esmif = batch['esm_feat']
+        
+        
+        num_nodes = batch['attention_mask'].sum(dim=-1)
+        shift = torch.cat([torch.zeros(1, device=device), torch.cumsum(num_nodes, dim=0)]).long()
+        
+        B, max_L = num_nodes.shape[0], num_nodes.max()
+        
+        self.initoutput(pretrain_design, B, max_L, device)
+        
+        
+        # keys = list(zip(design_seqs, *lm_seqs))
+        keys = batch['title']
+        unseen = self.retrivel(keys, num_nodes,device, use_memory)
+        
+                
+        if len(unseen)>0:
+            unseen_batch = self.rebatch(unseen, batch_id_raw, E_idx_raw, h_E_raw, shift, num_nodes, pretrain_design, pretrain_esm_msa, pretrain_struct, pretrain_esmif, device)
+            unseen_results = self.GNNTuning(unseen_batch)
+            
+            self.save2memory(keys,unseen,num_nodes, unseen_results)
+            self.update(unseen, num_nodes, unseen_results, keys)
+            
+        return {'title':self.titles,'pred_ids':self.out_pred_ids, 'confs':self.out_confs, 'embeds':self.out_embeds, 'probs':self.out_probs, "log_probs":self.out_log_probs, 'attention_mask':pretrain_design['attention_mask']}
+        
+        
+        
+        
+
+
+
+    

Flexpert-Design/src/models/PretrainESMIF_model.py

@@ -0,0 +1,32 @@
+import esm
+import torch.nn as nn
+import torch
+from esm.inverse_folding.util import CoordBatchConverter
+
+class PretrainESMIF_Model(nn.Module):
+    def __init__(self):
+        super(PretrainESMIF_Model, self).__init__()
+        #  /root/.cache/torch/hub/checkpoints
+        model_data = torch.load("./model_zoo/esmif/esm_if1_gvp4_t16_142M_UR50.pt")
+        self.model, self.alphabet = esm.pretrained.load_model_and_alphabet_core("esm_if1_gvp4_t16_142M_UR50", model_data, None)
+    
+    def forward(self, coords_list):
+        self.model.eval()
+        batch_converter = CoordBatchConverter(self.model.decoder.dictionary)
+        batch_coords, confidence, _, _, padding_mask = (
+            batch_converter([(coord, None, None) for coord in coords_list], device=coords_list[0].device)
+        )
+        with torch.no_grad():
+            encoder_out = self.model.encoder(batch_coords, padding_mask, confidence)
+            
+        feat = encoder_out['encoder_out'][0].permute(1,0,2)[:,1:-1] # 2,1046-2,512
+        attention_mask = encoder_out['encoder_padding_mask'][0][:,1:-1]==False # 2,1046-2
+        
+        return {"feat":feat}
+
+if __name__ == '__main__':
+    model = PretrainESMIF_Model(0.1)
+    coords1 = torch.rand(1044,3,3)#N, CA, C
+    coords2 = torch.rand(500,3,3)
+    model([coords1, coords2])
+    print()

Flexpert-Design/src/models/PretrainESM_model.py

@@ -0,0 +1,35 @@
+import time
+import torch
+import math
+import torch.nn as nn
+from transformers import AutoTokenizer, EsmForMaskedLM # EsmForMaskedLM, 1041 line
+import torch
+
+
+class PretrainESM_Model(nn.Module):
+    def __init__(self, args):
+        """ Graph labeling network """
+        super(PretrainESM_Model, self).__init__()
+        self.args=args
+        # {0: '<cls>', 1: '<pad>', 2: '<eos>', 3: '<unk>', 4: 'L', 5: 'A', 6: 'G', 7: 'V', 8: 'S', 9: 'E', 10: 'R', 11: 'T', 12: 'I', 13: 'D', 14: 'P', 15: 'K', 16: 'Q', 17: 'N', 18: 'F', 19: 'Y', 20: 'M', 21: 'H', 22: 'W', 23: 'C', 24: 'X', 25: 'B', 26: 'U', 27: 'Z', 28: 'O', 29: '.', 30: '-', 31: '<null_1>', 32: '<mask>'}
+        self.tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D", cache_dir="./cache_dir/")
+        self.model = EsmForMaskedLM.from_pretrained("facebook/esm2_t33_650M_UR50D", cache_dir="./cache_dir/")
+        
+    
+    def forward(self,batch):
+        outputs = self.model(input_ids=batch['pred_ids'], attention_mask=batch['attention_mask'])
+        logits = outputs.logits
+
+        prop = logits.softmax(dim=-1)
+        confidences, pred_ids = prop.max(dim=-1)
+            
+        ret = {"pred_ids": pred_ids,
+               "confs": confidences,
+               "embeds": outputs.hidden_states,
+               "attention_mask": batch['attention_mask']}
+        return ret
+    
+if __name__ == '__main__':
+    tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D", cache_dir="./cache_dir/")
+    tokenizer.convert_ids_to_tokens
+    print()

Flexpert-Design/src/models/PretrainPiFold_model.py

@@ -0,0 +1,64 @@
+import torch
+import os.path as osp
+from src.models.pifold_model import PiFold_Model
+import torch.nn.functional as F
+
+
+class PretrainPiFold_Model(PiFold_Model):
+    def __init__(self, args, **kwargs):
+        """ Graph labeling network """
+        PiFold_Model.__init__(self, args)
+        if args.augment_eps>0:
+            pretrain_pifold_path = osp.join(self.args.res_dir, self.args.dataset, f"PiFold_{args.augment_eps}", "checkpoint.pth")
+        else:
+            # pretrain_pifold_path = osp.join(self.args.res_dir, self.args.dataset, "PiFold", "checkpoint.pth")
+            pretrain_pifold_path = osp.join('model_zoo', self.args.dataset, "PiFold", "checkpoint.pth")
+        self.load_state_dict(torch.load(pretrain_pifold_path))
+    
+    @torch.no_grad()
+    def forward(self, batch):
+        h_V, h_P, P_idx, batch_id = batch['h_V'], batch['h_E'], batch['E_idx'], batch['batch_id']
+        device = h_V.device
+        h_V = self.W_v(self.norm_nodes(self.node_embedding(h_V)))
+        h_P = self.W_e(self.norm_edges(self.edge_embedding(h_P)))
+        
+        h_V, h_P = self.encoder(h_V, h_P, P_idx, batch_id)
+        log_probs, logits = self.decoder(h_V, batch_id)
+        probs = F.softmax(logits, dim=-1)
+        conf, pred_id = probs.max(dim=-1)
+        
+        maxL = 0
+        for b in batch_id.unique():
+            mask = batch_id==b
+            L = mask.sum()
+            if L>maxL:
+                maxL=L
+        
+        confs = []
+        seqs = []
+        embeds = []
+        probs2 = []
+        for b in batch_id.unique():
+            mask = batch_id==b
+            # elements = [alphabet[int(id)] for id in pred_id[mask]]
+            elements = self.tokenizer.decode(pred_id[mask]).split(" ")
+            seqs.append(elements)
+            confs.append(conf[mask])
+            embeds.append(h_V[mask])
+            probs2.append(probs[mask])
+        
+        seqs = self.tokenizer(["".join(one) for one in seqs], padding=True, truncation=True, return_tensors='pt', add_special_tokens=False)
+        confs = torch.stack([F.pad(one, (0, maxL-len(one))) for one in confs])
+        embeds = torch.stack([F.pad(one, (0,0, 0, maxL-len(one))) for one in embeds])
+        probs2 = torch.stack([F.pad(one, (0,0, 0, maxL-len(one)), value=1/33) for one in probs2])
+        
+        ret = {"pred_ids":seqs['input_ids'].to(device),
+               "confs":confs,
+               "embeds":embeds,
+               "probs":probs2,
+               "attention_mask":seqs['attention_mask'].to(device),
+               "E_idx":P_idx,
+               "batch_id":batch_id,
+               "h_E":h_P}
+        return ret
+

Flexpert-Design/src/models/Tuning.py

@@ -0,0 +1,275 @@
+import torch.nn as nn
+from src.modules.pifold_module import *
+from torch_scatter import scatter_softmax, scatter_log_softmax
+
+def positional_encoding(x):
+    batch_size, seq_len, hidden_size = x.size()
+    pos = torch.arange(0, seq_len).float().unsqueeze(1).repeat(1, hidden_size // 2)
+    div = torch.exp(torch.arange(0, hidden_size, 2).float() * (-torch.log(torch.tensor(10000.0)) / hidden_size))
+    sin = torch.sin(pos * div)
+    cos = torch.cos(pos * div)
+    pos_encoding = torch.cat([sin, cos], dim=-1).unsqueeze(0).repeat(batch_size, 1, 1)
+    return pos_encoding
+
+
+class MSAAttention(nn.Module):
+    def __init__(self, hidden_dim) -> None:
+        super().__init__()
+        self.MSA_Q = nn.Linear(hidden_dim, hidden_dim)
+        self.MSA_K = nn.Linear(hidden_dim, hidden_dim)
+        self.MSA_V = nn.Linear(hidden_dim, hidden_dim)
+    
+    def forward(self, inputs_embeds):
+        pos_enc = positional_encoding(inputs_embeds)
+        inputs_embeds = inputs_embeds + pos_enc
+        
+        query = self.MSA_Q(inputs_embeds)  # shape: [batch, N, 128]
+        key = self.MSA_K(inputs_embeds)  # shape: [batch, N, 128]
+        value = self.MSA_V(inputs_embeds)  # shape: [batch, N, 128]
+        attn_scores = torch.bmm(query, key.transpose(1, 2))
+        attn_weights = nn.functional.softmax(attn_scores, dim=2)
+        
+        attn_output = torch.bmm(attn_weights, value)
+        return attn_output
+
+
+class GNNTuning_Model(nn.Module):
+    def __init__(self, args, num_encoder_layers, hidden_dim, input_design_dim, input_esm_dim, input_struct_dim=3072, input_esmif_dim=512, dropout=0.1):
+        super(GNNTuning_Model, self).__init__()
+        self.args = args
+        encoder_layers = []
+        for i in range(num_encoder_layers):
+            encoder_layers.append(
+                GeneralGNN(hidden_dim, hidden_dim*2, dropout=dropout, node_net = "AttMLP", edge_net = "EdgeMLP", node_context = 1, edge_context = 0),
+            )
+        self.encoder_layers = nn.Sequential(*encoder_layers)
+        
+        from transformers import AutoTokenizer
+        from transformers.models.esm.modeling_esm import EsmModel, EsmEmbeddings
+        from transformers.models.esm.configuration_esm import EsmConfig
+        
+        self.tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t33_650M_UR50D", cache_dir="./cache_dir/")
+        config = EsmConfig(attention_probs_dropout_prob=0,
+                           hidden_size=hidden_dim,
+                           intermediate_size=1280,
+                           mask_token_id=32,
+                           num_attention_heads=12,
+                           num_hidden_layers=3,
+                           pad_token_id=1,
+                           position_embedding_type="rotary",
+                           token_dropout=False,
+                           vocab_size=33
+                           )
+        
+        self.DesignEmbed = EsmEmbeddings(config)
+        self.ESMEmbed = EsmEmbeddings(config)
+        self.EdgeEmbed = nn.Sequential(nn.Linear(416+16+16, 512),
+                                     nn.ReLU(),
+                                     nn.Linear(512, hidden_dim),
+                                     nn.ReLU(),
+                                     nn.Linear(hidden_dim,hidden_dim))
+        
+        self.DesignConf = nn.Sequential(nn.Linear(1, 128),
+                                     nn.ReLU(),
+                                     nn.Linear(128, 128),
+                                     nn.ReLU(),
+                                     nn.Linear(128,1),
+                                     nn.Sigmoid())
+        
+        self.ESMConf = nn.Sequential(nn.Linear(1, 128),
+                                     nn.ReLU(),
+                                     nn.Linear(128, 128),
+                                     nn.ReLU(),
+                                     nn.Linear(128,1))
+        
+        self.DesignProj = nn.Sequential(nn.Linear(input_design_dim, 512),
+                                     nn.ReLU(),
+                                     nn.Linear(512, hidden_dim),
+                                     nn.ReLU(),
+                                     nn.Linear(hidden_dim,hidden_dim))
+        
+        self.ESMProj = nn.Sequential(nn.Linear(input_esm_dim, 512),
+                                     nn.ReLU(),
+                                     nn.Linear(512, hidden_dim),
+                                     nn.ReLU(),
+                                     nn.Linear(hidden_dim,hidden_dim))
+        
+        self.StructProj = nn.Sequential(nn.Linear(input_struct_dim, 512),
+                                     nn.ReLU(),
+                                     nn.Linear(512, hidden_dim),
+                                     nn.ReLU(),
+                                     nn.Linear(hidden_dim,hidden_dim))
+        
+        self.ESMIFProj = nn.Sequential(nn.Linear(input_esmif_dim, 512),
+                                     nn.ReLU(),
+                                     nn.Linear(512, hidden_dim),
+                                     nn.ReLU(),
+                                     nn.Linear(hidden_dim,hidden_dim))
+        
+        self.ReadOut = nn.Linear(hidden_dim,33)
+        # self.TimeEmbed = nn.Embedding(20, hidden_dim)
+        # self.ProbEmbed = nn.Sequential(nn.Linear(33, 512),
+        #                              nn.ReLU(),
+        #                              nn.Linear(512, hidden_dim),
+        #                              nn.ReLU(),
+        #                              nn.Linear(hidden_dim,hidden_dim))
+        
+        self.MLP1 = nn.Sequential(nn.Linear(1, 512),
+                                     nn.ReLU(),
+                                     nn.Linear(512, hidden_dim),
+                                     nn.ReLU(),
+                                     nn.Linear(hidden_dim,1),
+                                     nn.Sigmoid())
+        
+        self.MLP2 = nn.Sequential(nn.Linear(1, 512),
+                                     nn.ReLU(),
+                                     nn.Linear(512, hidden_dim),
+                                     nn.ReLU(),
+                                     nn.Linear(hidden_dim,1),
+                                     nn.Sigmoid())
+
+        
+        
+    
+    # def embed_gnn(self, pretrain_gnn, mask_select_id, mask_select_feat):
+    #     gnn_embed = self.DesignEmbed(mask_select_id(pretrain_gnn['pred_ids'])).squeeze()
+    #     gnn_conf = self.DesignConf(mask_select_id(pretrain_gnn['confs']))
+    #     gnn_proj = self.DesignProj(mask_select_feat(pretrain_gnn['embeds']))
+        
+    #     if self.args.use_confembed:
+    #         return gnn_embed*F.sigmoid(gnn_conf)  + gnn_proj 
+    #     else:
+    #         return gnn_embed  + gnn_proj 
+    
+    # def embed_esm(self, pretrain_esm, mask_select_id, mask_select_feat):
+    #     esm_embed = self.ESMEmbed(mask_select_id(pretrain_esm['pred_ids'])).squeeze()
+    #     esm_conf = self.ESMConf(mask_select_id(pretrain_esm['confs']))
+    #     esm_proj = self.ESMProj(mask_select_feat(pretrain_esm['embeds']))
+    #     if self.args.use_confembed:
+    #         return esm_embed*F.sigmoid(esm_conf) + esm_proj 
+    #     else:
+    #         return esm_embed + esm_proj
+    
+    # def embed_struct(self, pretrain_struct, mask_select_feat):
+    #     struct_proj = self.StructProj(mask_select_feat(pretrain_struct['embeds']))
+    #     return struct_proj
+    
+    # def embed_esmif(self, pretrain_esmif, mask_select_feat):
+    #     struct_proj = self.ESMIFProj(mask_select_feat(pretrain_esmif['embeds']))
+    #     return struct_proj
+    
+    def fuse(self, mask_select_feat, mask_select_id, gnn_embed=None, esm_embed=None, gearnet_embed=None, esmif_embed=None, gnn_pred_id=None, esm_pred_id=None, confidence=None, confidence_esm=None):
+        gnn, esm, gearnet, esmif, conf = 0, 0, 0, 0, 1.0
+        if gnn_embed is not None:
+            gnn = self.DesignProj(mask_select_feat(gnn_embed))
+            gnn += self.DesignEmbed(mask_select_id(gnn_pred_id)).squeeze()
+        
+        if esm_embed is not None:
+            esm = self.ESMProj(mask_select_feat(esm_embed))
+            esm += self.ESMEmbed(mask_select_id(esm_pred_id)).squeeze()
+        
+        if gearnet_embed is not None:
+            gearnet = self.StructProj(mask_select_feat(gearnet_embed))
+        
+        if esmif_embed is not None:
+            esmif = self.ESMIFProj(mask_select_feat(esmif_embed))
+        
+        if conf is not None:
+            conf = self.DesignConf(mask_select_id(confidence))
+            esm_conf = self.ESMConf(mask_select_id(confidence_esm))
+        
+        return (gnn*conf+esm*esm_conf+gearnet+esmif)
+        
+
+    def forward(self, batch):
+        pretrain_design,  h_E_raw, E_idx, mask_attend, batch_id = batch['pretrain_design'], batch['h_E'], batch['E_idx'], batch['attention_mask'], batch['batch_id']
+        
+        if self.args.use_LM:
+            pretrain_esm_msa = batch['pretrain_esm_msa']
+        
+        if self.args.use_gearnet:
+            pretrain_struct = batch['pretrain_struct']
+        
+        if self.args.use_esmif:
+            pretrain_esmif = batch['pretrain_esmif']
+            
+        mask_select_id = lambda x:  torch.masked_select(x, mask_attend.bool()).reshape(-1,1)
+        mask_select_feat = lambda x:  torch.masked_select(x, mask_attend.bool().unsqueeze(-1)).reshape(-1,x.shape[-1])
+        
+        inputs_embeds = 0
+        for i in range(self.args.msa_n):
+            gnn_embed = pretrain_design['embeds']
+            esm_embed = pretrain_esm_msa['embeds'][i] if self.args.use_LM else None
+            gearnet_embed = pretrain_struct['embeds'][i] if self.args.use_gearnet else None
+            esmif_embed = pretrain_esmif['embeds'] if self.args.use_esmif else None
+            confidence = pretrain_design['confs']
+            confidence_esm = pretrain_esm_msa['confs'][i]
+            inputs_embeds += self.fuse(mask_select_feat, mask_select_id, gnn_embed, esm_embed, gearnet_embed, esmif_embed, pretrain_design['pred_ids'], pretrain_esm_msa['pred_ids'][i], confidence, confidence_esm)
+        
+        h_V = inputs_embeds
+        h_E = self.EdgeEmbed(h_E_raw)
+        
+        for layer in self.encoder_layers:
+            h_V, h_E = layer(h_V, h_E, E_idx, batch_id)
+            
+        logits = self.ReadOut(h_V)
+        
+        # confidence update
+        old_confs = mask_select_id(pretrain_design['confs'])
+        confs = torch.softmax(logits, dim=-1).max(dim=-1)[0][:,None]
+        h_V = h_V*self.MLP1(confs-old_confs) + inputs_embeds*self.MLP2(old_confs-confs)
+        logits = self.ReadOut(h_V)
+        
+        B, N = pretrain_design['confs'].shape
+        vocab_size = logits.shape[-1]
+        
+        new_logits = torch.zeros(B,N,vocab_size, device=logits.device).reshape(B*N, vocab_size)
+        new_logits = new_logits.masked_scatter_(mask_attend.bool().view(-1,1), logits)
+        new_logits = new_logits.reshape(B,N,vocab_size)
+        log_probs = torch.log_softmax(new_logits, dim=-1)
+        
+        device = logits.device
+        seqs, confs, embeds, probs2 = self.to_matrix(h_V, logits, batch_id)
+        
+        
+        ret = {"pred_ids":seqs['input_ids'].to(device),
+               "confs":confs,
+               "embeds":embeds,
+               "probs":probs2,
+               "attention_mask":seqs['attention_mask'].to(device),
+               "h_E":h_E_raw,
+               "E_idx":E_idx,
+               "batch_id":batch_id,
+               "log_probs":log_probs}
+        return ret
+
+    def to_matrix(self, h_V, logits, batch_id):
+        
+        probs = F.softmax(logits, dim=-1)
+        conf, pred_id = probs.max(dim=-1)
+        
+        maxL = 0
+        for b in batch_id.unique():
+            mask = batch_id==b
+            L = mask.sum()
+            if L>maxL:
+                maxL=L
+
+        confs = []
+        seqs = []
+        embeds = []
+        probs2 = []
+        for b in batch_id.unique():
+            mask = batch_id==b
+            # elements = [alphabet[int(id)] for id in pred_id[mask]]
+            elements = self.tokenizer.decode(pred_id[mask]).split(" ")
+            seqs.append(elements)
+            confs.append(conf[mask])
+            embeds.append(h_V[mask])
+            probs2.append(probs[mask])
+        
+        seqs = self.tokenizer(["".join(one) for one in seqs], padding=True, truncation=True, return_tensors='pt', add_special_tokens=False)
+        confs = torch.stack([F.pad(one, (0, maxL-len(one))) for one in confs])
+        embeds = torch.stack([F.pad(one, (0,0, 0, maxL-len(one))) for one in embeds])
+        probs2 = torch.stack([F.pad(one, (0,0, 0, maxL-len(one)), value=1/33) for one in probs2])
+        return seqs, confs, embeds, probs2

Flexpert-Design/src/models/__init__.py

@@ -0,0 +1,16 @@
+# Copyright (c) CAIRI AI Lab. All rights reserved
+
+# from .alphadesign_model import AlphaDesign_Model
+# from .esmif_model import GVPTransformerModel as ESMIF_Model
+# from .gca_model import GCA_Model
+# from .graphtrans_model import GraphTrans_Model
+# from .gvp_model import GVP_Model
+# from .pifold_model import PiFold_Model
+from .proteinmpnn_model import ProteinMPNN_Model
+# from .structgnn_model import StructGNN_Model
+# from .kwdesign_model import KWDesign_model
+
+__all__ = [
+    'AlphaDesign_Model', 'ESMIF_Model', 'GCA_Model', 'GraphTrans_Model', 'GVP_Model',
+    'PiFold_Model', 'ProteinMPNN_Model', 'StructGNN_Model', 'KWDesign_model'
+]

Flexpert-Design/src/models/alphadesign_model.py

@@ -0,0 +1,138 @@
+import torch
+import torch.nn as nn
+
+from src.modules.alphadesign_module import ATDecoder, CNNDecoder, CNNDecoder2, StructureEncoder
+from src.tools.design_utils import gather_nodes, _dihedrals, _rbf, _orientations_coarse_gl
+
+
+class AlphaDesign_Model(nn.Module):
+    def __init__(self, args, **kwargs):
+        """ Graph labeling network """
+        super(AlphaDesign_Model, self).__init__()
+        self.args = args
+        node_features = args.node_features
+        edge_features = args.edge_features
+        hidden_dim = args.hidden_dim
+        dropout = args.dropout
+        num_encoder_layers = args.num_encoder_layers
+        self.top_k = args.k_neighbors
+        self.num_rbf = 16
+        self.num_positional_embeddings = 16
+
+        if args.use_new_feat:
+            node_in, edge_in = 12, 16+7
+        else:
+            node_in, edge_in = 6, 16+7
+        self.node_embedding = nn.Linear(node_in, node_features, bias=True)
+        self.edge_embedding = nn.Linear(edge_in, edge_features, bias=True)
+        self.norm_nodes = nn.BatchNorm1d(node_features)
+        self.norm_edges = nn.BatchNorm1d(edge_features)
+
+        self.W_v = nn.Sequential(
+            nn.Linear(node_features, hidden_dim, bias=True),
+            nn.LeakyReLU(),
+            nn.BatchNorm1d(hidden_dim),
+            nn.Linear(hidden_dim, hidden_dim, bias=True),
+            nn.LeakyReLU(),
+            nn.BatchNorm1d(hidden_dim),
+            nn.Linear(hidden_dim, hidden_dim, bias=True)
+        )
+        
+        
+        self.W_e = nn.Linear(edge_features, hidden_dim, bias=True) 
+        self.W_f = nn.Linear(edge_features, hidden_dim, bias=True)
+
+        self.encoder = StructureEncoder(hidden_dim, num_encoder_layers, dropout, use_SGT=self.args.use_SGT)
+
+        if args.autoregressive:
+            self.decoder = ATDecoder(args, hidden_dim, dropout)
+        else:
+            self.decoder = CNNDecoder(hidden_dim, hidden_dim)
+            self.decoder2 = CNNDecoder2(hidden_dim, hidden_dim)
+            
+        # self.chain_embed = nn.Embedding(2,16)
+        self._init_params()
+    
+    def forward(self, batch,  AT_test = False, return_logit=False):
+        h_V, h_P, P_idx, batch_id = batch['_V'], batch['_E'], batch['E_idx'], batch['batch_id']
+        h_V = self.W_v(self.norm_nodes(self.node_embedding(h_V)))
+        h_P = self.W_e(self.norm_edges(self.edge_embedding(h_P)))
+        
+        h_V = self.encoder(h_V, h_P, P_idx, batch_id)
+        log_probs0 = None
+        if AT_test:
+            log_probs = self.decoder.sampling(h_V, h_P, P_idx,  batch_id)
+        else:
+            log_probs0, logits = self.decoder(h_V, batch_id)
+            log_probs, logits = self.decoder2(h_V, logits, batch_id)
+        if return_logit:
+            return {'log_probs': log_probs, 'logits': logits}
+        else:
+            return {'log_probs': log_probs, 'log_probs0': log_probs0}
+        
+    def _init_params(self):
+        for p in self.parameters():
+            if p.dim() > 1:
+                nn.init.xavier_uniform_(p)
+
+    def _full_dist(self, X, mask, top_k=30, eps=1E-6):
+        mask_2D = torch.unsqueeze(mask,1) * torch.unsqueeze(mask,2)
+        dX = torch.unsqueeze(X,1) - torch.unsqueeze(X,2)
+        D = (1. - mask_2D)*10000 + mask_2D* torch.sqrt(torch.sum(dX**2, 3) + eps)
+
+        D_max, _ = torch.max(D, -1, keepdim=True)
+        D_adjust = D + (1. - mask_2D) * (D_max+1)
+        D_neighbors, E_idx = torch.topk(D_adjust, min(top_k, D_adjust.shape[-1]), dim=-1, largest=False)
+        return D_neighbors, E_idx  
+
+    def _get_features(self, batch):
+        S, score, X, mask = batch['S'], batch['score'], batch['X'], batch['mask']
+        mask_bool = (mask==1)
+        
+        B, N, _,_ = X.shape
+        X_ca = X[:,:,1,:]
+        D_neighbors, E_idx = self._full_dist(X_ca, mask, self.top_k)
+
+        # sequence
+        S = torch.masked_select(S, mask_bool)
+        if score is not None:
+            score = torch.masked_select(score, mask_bool)
+
+        # node feature
+        _V = _dihedrals(X) 
+        if not self.args.use_new_feat:
+            _V = _V[...,:6]
+        _V = torch.masked_select(_V, mask_bool.unsqueeze(-1)).reshape(-1,_V.shape[-1])
+
+        # edge feature
+        _E = torch.cat((_rbf(D_neighbors, self.num_rbf), _orientations_coarse_gl(X, E_idx)), -1) # [4,387,387,23]
+        mask_attend = gather_nodes(mask.unsqueeze(-1), E_idx).squeeze(-1) # 一阶邻居节点的mask: 1代表节点存在, 0代表节点不存在
+        mask_attend = (mask.unsqueeze(-1) * mask_attend) == 1 # 自身的mask*邻居节点的mask
+        _E = torch.masked_select(_E, mask_attend.unsqueeze(-1)).reshape(-1,_E.shape[-1])
+        
+        # edge index
+        shift = mask.sum(dim=1).cumsum(dim=0) - mask.sum(dim=1)
+        src = shift.view(B,1,1) + E_idx
+        src = torch.masked_select(src, mask_attend).view(1,-1)
+        dst = shift.view(B,1,1) + torch.arange(0, N, device=src.device).view(1,-1,1).expand_as(mask_attend)
+        dst = torch.masked_select(dst, mask_attend).view(1,-1)
+        E_idx = torch.cat((dst, src), dim=0).long()
+        
+        
+        # 3D point
+        sparse_idx = mask.nonzero()
+        X = X[sparse_idx[:,0],sparse_idx[:,1],:,:]
+        batch_id = sparse_idx[:,0]
+
+        mask = torch.masked_select(mask, mask_bool)
+
+        batch.update({'X':X,
+                'S':S,
+                'score':score,
+                '_V':_V,
+                '_E':_E,
+                'E_idx':E_idx,
+                'batch_id': batch_id,
+                'mask':mask})
+
+        return batch

Flexpert-Design/src/models/anm_prottrans.py

@@ -0,0 +1,677 @@
+#import dependencies
+import os.path
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
+from torch.utils.data import DataLoader
+
+import re
+import numpy as np
+import pandas as pd
+import copy
+import pdb
+
+import transformers, datasets
+from transformers.modeling_outputs import TokenClassifierOutput, BaseModelOutputWithPastAndCrossAttentions
+from transformers.models.t5.modeling_t5 import T5Config, T5PreTrainedModel, T5Stack
+from transformers.utils.model_parallel_utils import assert_device_map, get_device_map
+from transformers import T5EncoderModel, T5Tokenizer
+from transformers import TrainingArguments, Trainer, set_seed
+from safetensors import safe_open
+
+#DataCollator
+from transformers.data.data_collator import DataCollatorMixin
+from transformers.tokenization_utils_base import PreTrainedTokenizerBase
+from transformers.utils import PaddingStrategy
+
+import random
+import warnings
+from collections.abc import Mapping
+from dataclasses import dataclass
+from random import randint
+from typing import Any, Callable, Dict, List, NewType, Optional, Tuple, Union
+
+from evaluate import load
+from datasets import Dataset
+
+from tqdm import tqdm
+import random
+
+from scipy import stats
+from sklearn.metrics import accuracy_score
+
+import matplotlib.pyplot as plt
+
+from Bio import SeqIO
+from io import StringIO
+import requests
+import tempfile
+
+from sklearn.model_selection import train_test_split
+import csv
+
+
+#### UTILS
+
+class LoRAConfig:
+    def __init__(self):
+        self.lora_rank = 4
+        self.lora_init_scale = 0.01
+        self.lora_modules = ".*SelfAttention|.*EncDecAttention"
+        self.lora_layers = "q|k|v|o"
+        self.trainable_param_names = ".*layer_norm.*|.*lora_[ab].*"
+        self.lora_scaling_rank = 1
+        # lora_modules and lora_layers are speicified with regular expressions
+        # see https://www.w3schools.com/python/python_regex.asp for reference
+        
+class LoRALinear(nn.Module):
+    def __init__(self, linear_layer, rank, scaling_rank, init_scale):
+        super().__init__()
+        self.in_features = linear_layer.in_features
+        self.out_features = linear_layer.out_features
+        self.rank = rank
+        self.scaling_rank = scaling_rank
+        self.weight = linear_layer.weight
+        self.bias = linear_layer.bias
+        if self.rank > 0:
+            self.lora_a = nn.Parameter(torch.randn(rank, linear_layer.in_features) * init_scale)
+            if init_scale < 0:
+                self.lora_b = nn.Parameter(torch.randn(linear_layer.out_features, rank) * init_scale)
+            else:
+                self.lora_b = nn.Parameter(torch.zeros(linear_layer.out_features, rank))
+        if self.scaling_rank:
+            self.multi_lora_a = nn.Parameter(
+                torch.ones(self.scaling_rank, linear_layer.in_features)
+                + torch.randn(self.scaling_rank, linear_layer.in_features) * init_scale
+            )
+            if init_scale < 0:
+                self.multi_lora_b = nn.Parameter(
+                    torch.ones(linear_layer.out_features, self.scaling_rank)
+                    + torch.randn(linear_layer.out_features, self.scaling_rank) * init_scale
+                )
+            else:
+                self.multi_lora_b = nn.Parameter(torch.ones(linear_layer.out_features, self.scaling_rank))
+
+    def forward(self, input):
+        if self.scaling_rank == 1 and self.rank == 0:
+            # parsimonious implementation for ia3 and lora scaling
+            if self.multi_lora_a.requires_grad:
+                hidden = F.linear((input * self.multi_lora_a.flatten()), self.weight, self.bias)
+            else:
+                hidden = F.linear(input, self.weight, self.bias)
+            if self.multi_lora_b.requires_grad:
+                hidden = hidden * self.multi_lora_b.flatten()
+            return hidden
+        else:
+            # general implementation for lora (adding and scaling)
+            weight = self.weight
+            if self.scaling_rank:
+                weight = weight * torch.matmul(self.multi_lora_b, self.multi_lora_a) / self.scaling_rank
+            if self.rank:
+                weight = weight + torch.matmul(self.lora_b, self.lora_a) / self.rank
+            return F.linear(input, weight, self.bias)
+
+    def extra_repr(self):
+        return "in_features={}, out_features={}, bias={}, rank={}, scaling_rank={}".format(
+            self.in_features, self.out_features, self.bias is not None, self.rank, self.scaling_rank
+        )
+
+
+def modify_with_lora(transformer, config):
+    for m_name, module in dict(transformer.named_modules()).items():
+        if re.fullmatch(config.lora_modules, m_name):
+            for c_name, layer in dict(module.named_children()).items():
+                if re.fullmatch(config.lora_layers, c_name):
+                    assert isinstance(
+                        layer, nn.Linear
+                    ), f"LoRA can only be applied to torch.nn.Linear, but {layer} is {type(layer)}."
+                    setattr(
+                        module,
+                        c_name,
+                        LoRALinear(layer, config.lora_rank, config.lora_scaling_rank, config.lora_init_scale),
+                    )
+    return transformer
+
+class ClassConfig:
+    def __init__(self, dropout=0.2, num_labels=1, add_pearson_loss=False, add_sse_loss=False, adaptor_architecture = None , enm_embed_dim = 512, enm_att_heads = 8, kernel_size = 3, num_layers = 2, **kwargs):
+        self.dropout_rate = dropout
+        self.num_labels = num_labels
+        self.add_pearson_loss = add_pearson_loss
+        self.add_sse_loss = add_sse_loss
+        self.adaptor_architecture = adaptor_architecture
+        self.enm_embed_dim = enm_embed_dim
+        self.enm_att_heads = enm_att_heads
+        self.kernel_size = kernel_size
+        self.num_layers = num_layers
+
+class ENMAdaptedAttentionClassifier(nn.Module):
+    def __init__(self, seq_embedding_dim, out_dim, enm_embed_dim, num_att_heads):
+        super(ENMAdaptedAttentionClassifier, self).__init__()
+        self.embedding = nn.Linear(1, enm_embed_dim)
+        self.enm_attention = nn.MultiheadAttention(enm_embed_dim, num_att_heads)
+        self.layer_norm = nn.LayerNorm(enm_embed_dim)
+        self.enm_adaptor = nn.Linear(enm_embed_dim, seq_embedding_dim)
+        self.adapted_classifier = nn.Linear(2*seq_embedding_dim, out_dim)
+    
+    def forward(self, seq_embedding, enm_input):
+        enm_input = enm_input.transpose(0, 1)  # Transpose to shape (N, B, E) for MultiheadAttention
+        enm_input = enm_input.unsqueeze(-1)  # Add a dimension for the embedding
+        enm_input_embedded = self.embedding(enm_input)
+        enm_att, _ = self.enm_attention(enm_input_embedded, enm_input_embedded, enm_input_embedded)
+        enm_att = enm_att.transpose(0, 1)  # Transpose back to shape (B, N, E)
+        enm_att = self.layer_norm(enm_att + enm_input.transpose(0, 1))
+        enm_embedding = self.enm_adaptor(enm_att)
+        combined_embedding = torch.cat((seq_embedding, enm_embedding), dim=-1)
+        logits = self.adapted_classifier(combined_embedding)
+        return logits
+    
+class ENMAdaptedConvClassifier(nn.Module):
+    def __init__(self, seq_embedding_dim, out_dim, kernel_size, enm_embedding_dim, num_layers):
+        super(ENMAdaptedConvClassifier, self).__init__()
+        layers = []
+        self.conv1 = nn.Conv1d(1, enm_embedding_dim, kernel_size=kernel_size, padding=(kernel_size-1)//2)
+        layers.append(self.conv1)
+        layers.append(nn.ReLU())
+        for i in range(num_layers-1):
+            layers.append(nn.Conv1d(enm_embedding_dim, enm_embedding_dim, kernel_size=kernel_size, padding=(kernel_size-1)//2))
+            layers.append(nn.ReLU())
+        self.conv_net = nn.Sequential(*layers)
+        self.adapted_classifier = nn.Linear(seq_embedding_dim+1, out_dim)
+
+    def forward(self, seq_embedding, enm_input, attention_mask=None):
+        enm_input = torch.nan_to_num(enm_input, nan=0.0)
+        enm_input = enm_input.unsqueeze(1)
+        conv_out = self.conv_net(enm_input)
+        enm_embedding = conv_out.transpose(1,2)
+        
+        if attention_mask is not None:
+            # Use attention_mask to ignore padded elements
+            mask = attention_mask.unsqueeze(-1).float()
+            enm_embedding = enm_embedding * mask
+            # Compute mean over non-padded elements
+            
+            enm_embedding = enm_embedding.mean(dim=-1).unsqueeze(-1)
+            # enm_embedding = enm_embedding.sum(dim=2)/ mask.sum(dim=2).clamp(min=1e-9)
+        else:
+            raise ValueError('We actually want to provide the mask.')
+            enm_embedding = torch.mean(enm_embedding, dim=1)
+            
+        # enm_embedding = enm_embedding.unsqueeze(1).expand(-1, seq_embedding.size(1), -1)
+        combined_embedding = torch.cat((seq_embedding, enm_embedding), dim=-1)
+        logits = self.adapted_classifier(combined_embedding)
+        return logits
+    
+
+        
+class ENMAdaptedDirectClassifier(nn.Module):
+    def __init__(self, seq_embedding_dim, out_dim):
+        super(ENMAdaptedDirectClassifier, self).__init__()
+        self.adapted_classifier = nn.Linear(seq_embedding_dim+1, out_dim)
+
+    def forward(self, seq_embedding, enm_input):
+            enm_input = enm_input.unsqueeze(-1)
+            combined_embedding = torch.cat((seq_embedding, enm_input), dim=-1)
+            logits = self.adapted_classifier(combined_embedding)
+            return logits
+
+class ENMNoAdaptorClassifier(nn.Module):
+    def __init__(self, seq_embedding_dim, out_dim):
+        super(ENMNoAdaptorClassifier, self).__init__()
+        self.adapted_classifier = nn.Linear(seq_embedding_dim, out_dim)
+
+    def forward(self, seq_embedding, enm_input):
+            _ = enm_input #ignoring enm_input
+            logits = self.adapted_classifier(seq_embedding)
+            return logits
+
+
+class T5EncoderForTokenClassification(T5PreTrainedModel):
+
+    def __init__(self, config: T5Config, class_config):
+        super().__init__(config)
+        self.num_labels = class_config.num_labels
+        self.config = config
+        self.add_pearson_loss = class_config.add_pearson_loss
+        self.add_sse_loss = class_config.add_sse_loss
+        self.shared = nn.Embedding(config.vocab_size, config.d_model)
+
+        encoder_config = copy.deepcopy(config)
+        encoder_config.use_cache = False
+        encoder_config.is_encoder_decoder = False
+        self.encoder = T5Stack(encoder_config, self.shared)
+        # self.encoder = CustomT5Stack(encoder_config, self.shared)
+
+        # import pdb; pdb.set_trace()
+        original_embedding = self.encoder.embed_tokens
+        in_dim, out_dim = tuple(original_embedding.weight.shape)
+        self.new_embedding = nn.Linear(in_dim, out_dim, bias=False).to('cuda:0') #TODO: pass the correct weights!!! And careful! the embedding layer and the linear layer are maybe mutually "transposed"
+        print("Initialized new_embedding layer - without weights yet!")
+        # self.new_embedding.weight = nn.Parameter(original_embedding.weight.T)
+
+        # self.weight = original_embedding.weight
+        # self.weight = nn.Parameter(self.new_embedding.weight.T)
+        # self.encoder.forward = new_forward.__get__(self.encoder, self.encoder.__class__)
+
+        self.dropout = nn.Dropout(class_config.dropout_rate)
+        if class_config.adaptor_architecture == 'attention':
+            self.classifier = ENMAdaptedAttentionClassifier(config.hidden_size, class_config.num_labels, class_config.enm_embed_dim, class_config.enm_att_heads) #nn.Linear(config.hidden_size, class_config.num_labels)
+        elif class_config.adaptor_architecture == 'direct':
+            self.classifier = ENMAdaptedDirectClassifier(config.hidden_size, class_config.num_labels)
+        elif class_config.adaptor_architecture == 'conv':
+            self.classifier = ENMAdaptedConvClassifier(config.hidden_size, class_config.num_labels, class_config.kernel_size, class_config.enm_embed_dim, class_config.num_layers)
+        elif class_config.adaptor_architecture == 'no-adaptor':
+            self.classifier = ENMNoAdaptorClassifier(config.hidden_size, class_config.num_labels)
+        else:
+            raise ValueError('Only attention, direct, conv and no-adaptor architectures are supported for the adaptor.')
+
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+        # Model parallel
+        self.model_parallel = False
+        self.device_map = None
+
+    def parallelize(self, device_map=None):
+        self.device_map = (
+            get_device_map(len(self.encoder.block), range(torch.cuda.device_count()))
+            if device_map is None
+            else device_map
+        )
+        assert_device_map(self.device_map, len(self.encoder.block))
+        self.encoder.parallelize(self.device_map)
+        self.classifier = self.classifier.to(self.encoder.first_device)
+        self.model_parallel = True
+
+    def deparallelize(self):
+        self.encoder.deparallelize()
+        self.encoder = self.encoder.to("cpu")
+        self.model_parallel = False
+        self.device_map = None
+        torch.cuda.empty_cache()
+
+    def get_input_embeddings(self):
+        return self.shared
+
+    def set_input_embeddings(self, new_embeddings):
+        self.shared = new_embeddings
+        self.encoder.set_input_embeddings(new_embeddings)
+
+    def get_encoder(self):
+        return self.encoder
+
+    def _prune_heads(self, heads_to_prune):
+        """
+        Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
+        class PreTrainedModel
+        """
+        for layer, heads in heads_to_prune.items():
+            self.encoder.layer[layer].attention.prune_heads(heads)
+
+    def forward(
+        self,
+        enm_vals = None,
+        input_ids=None,
+        attention_mask=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+        if inputs_embeds is not None:
+            outputs = self.encoder(input_ids=None,attention_mask=attention_mask,inputs_embeds=inputs_embeds,head_mask=head_mask,output_attentions=output_attentions,output_hidden_states=output_hidden_states,return_dict=return_dict,)
+        elif input_ids is not None:
+            outputs = self.encoder(input_ids=input_ids,attention_mask=attention_mask,inputs_embeds=None,head_mask=head_mask,output_attentions=output_attentions,output_hidden_states=output_hidden_states,return_dict=return_dict,)
+        sequence_output = outputs[0]
+        # import pdb; pdb.set_trace() #TODO: CHECK EVERYTHING IS IN EVAL MODE and the dropout below is OFF
+        sequence_output = self.dropout(sequence_output)
+        #TODO: check the enm_vals are padded properly and check that the sequence limit (in the transformer) is indeed 512
+        # logits = self.classifier(sequence_output, enm_vals)
+        
+        logits = self.classifier(sequence_output, enm_vals, attention_mask)
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return TokenClassifierOutput(
+            #loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+class ENMAdaptedTrainer(Trainer):
+    def compute_loss(self, model, inputs, return_outputs=False):
+        labels = inputs.get("labels")
+        #enm_vals = inputs.get("enm_vals")
+        
+        outputs = model(**inputs)
+        logits = outputs.get('logits')
+        mask = inputs.get('attention_mask')
+        loss_fct = MSELoss()
+
+        active_loss = mask.view(-1) == 1
+        active_logits = logits.view(-1)
+        active_labels = torch.where(active_loss, labels.view(-1), torch.tensor(-100).type_as(labels))
+        valid_logits=active_logits[active_labels!=-100]
+        valid_labels=active_labels[active_labels!=-100]
+
+        loss = loss_fct(valid_labels, valid_logits)
+        return (loss, outputs) if return_outputs else loss
+    
+
+
+def PT5_classification_model(half_precision, class_config):
+    # Load PT5 and tokenizer
+    # possible to load the half preciion model (thanks to @pawel-rezo for pointing that out)
+    if not half_precision:
+        model = T5EncoderModel.from_pretrained("Rostlab/prot_t5_xl_uniref50")
+        tokenizer = T5Tokenizer.from_pretrained("Rostlab/prot_t5_xl_uniref50")
+    elif half_precision and torch.cuda.is_available() : 
+        tokenizer = T5Tokenizer.from_pretrained('Rostlab/prot_t5_xl_half_uniref50-enc', do_lower_case=False)
+        model = T5EncoderModel.from_pretrained("Rostlab/prot_t5_xl_half_uniref50-enc", torch_dtype=torch.float16).to(torch.device('cuda'))
+    else:
+          raise ValueError('Half precision can be run on GPU only.')
+    
+
+    
+    # Create new Classifier model with PT5 dimensions
+    class_model=T5EncoderForTokenClassification(model.config,class_config)
+    
+    # Set encoder and embedding weights to checkpoint weights
+    class_model.shared=model.shared
+    class_model.encoder=model.encoder    
+    
+    # Delete the checkpoint model
+    model=class_model
+    del class_model
+    
+    # Print number of trainable parameters
+    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+    params = sum([np.prod(p.size()) for p in model_parameters])
+    print("ProtT5_Classfier\nTrainable Parameter: "+ str(params))    
+ 
+    # Add model modification lora
+    config = LoRAConfig()
+    
+    # Add LoRA layers
+    model = modify_with_lora(model, config)
+    
+    # Freeze Embeddings and Encoder (except LoRA)
+    for (param_name, param) in model.shared.named_parameters():
+                param.requires_grad = False
+    for (param_name, param) in model.encoder.named_parameters():
+                param.requires_grad = False       
+
+    for (param_name, param) in model.named_parameters():
+            if re.fullmatch(config.trainable_param_names, param_name):
+                param.requires_grad = True
+
+    # Print trainable Parameter          
+    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+    params = sum([np.prod(p.size()) for p in model_parameters])
+    print("ProtT5_LoRA_Classfier\nTrainable Parameter: "+ str(params) + "\n")
+    
+    return model, tokenizer
+
+
+@dataclass
+class DataCollatorForTokenRegression(DataCollatorMixin):
+    """
+    Data collator that will dynamically pad the inputs received, as well as the labels.
+    Args:
+        tokenizer ([`PreTrainedTokenizer`] or [`PreTrainedTokenizerFast`]):
+            The tokenizer used for encoding the data.
+        padding (`bool`, `str` or [`~utils.PaddingStrategy`], *optional*, defaults to `True`):
+            Select a strategy to pad the returned sequences (according to the model's padding side and padding index)
+            among:
+            - `True` or `'longest'` (default): Pad to the longest sequence in the batch (or no padding if only a single
+              sequence is provided).
+            - `'max_length'`: Pad to a maximum length specified with the argument `max_length` or to the maximum
+              acceptable input length for the model if that argument is not provided.
+            - `False` or `'do_not_pad'`: No padding (i.e., can output a batch with sequences of different lengths).
+        max_length (`int`, *optional*):
+            Maximum length of the returned list and optionally padding length (see above).
+        pad_to_multiple_of (`int`, *optional*):
+            If set will pad the sequence to a multiple of the provided value.
+            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=
+            7.5 (Volta).
+        label_pad_token_id (`int`, *optional*, defaults to -100):
+            The id to use when padding the labels (-100 will be automatically ignore by PyTorch loss functions).
+        return_tensors (`str`):
+            The type of Tensor to return. Allowable values are "np", "pt" and "tf".
+    """
+
+    tokenizer: PreTrainedTokenizerBase
+    padding: Union[bool, str, PaddingStrategy] = True
+    max_length: Optional[int] = None
+    pad_to_multiple_of: Optional[int] = None
+    label_pad_token_id: int = -100
+    return_tensors: str = "pt"
+
+    def torch_call(self, features):
+        label_name = "label" if "label" in features[0].keys() else "labels"
+        labels = [feature[label_name] for feature in features] if label_name in features[0].keys() else None
+
+        no_labels_features = [{k: v for k, v in feature.items() if k != label_name and k!= 'enm_vals'} for feature in features]
+
+        batch = self.tokenizer.pad(
+            no_labels_features,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors="pt",
+        )
+
+        batch['enm_vals'] = torch.nn.utils.rnn.pad_sequence([torch.tensor(feature['enm_vals'], dtype=torch.float) for feature in features], batch_first=True, padding_value=0.0)
+#batch = self.tokenizer.pad(no_labels_features,padding=self.padding,max_length=self.max_length,pad_to_multiple_of=self.pad_to_multiple_of,return_tensors="pt")
+        if labels is None:
+            return batch
+
+        sequence_length = batch["input_ids"].shape[1]
+        padding_side = self.tokenizer.padding_side
+
+        def to_list(tensor_or_iterable):
+            if isinstance(tensor_or_iterable, torch.Tensor):
+                return tensor_or_iterable.tolist()
+            return list(tensor_or_iterable)
+
+        if padding_side == "right":
+            batch[label_name] = [
+                to_list(label) + [self.label_pad_token_id] * (sequence_length - len(label)) for label in labels
+
+            ]
+        else:
+            batch[label_name] = [
+                [self.label_pad_token_id] * (sequence_length - len(label)) + to_list(label) for label in labels
+            ]
+
+        batch[label_name] = torch.tensor(batch[label_name], dtype=torch.float)
+        return batch
+
+def _torch_collate_batch(examples, tokenizer, pad_to_multiple_of: Optional[int] = None):
+    """Collate `examples` into a batch, using the information in `tokenizer` for padding if necessary."""
+    # Tensorize if necessary.
+    if isinstance(examples[0], (list, tuple, np.ndarray)):
+        examples = [torch.tensor(e, dtype=torch.long) for e in examples]
+
+    length_of_first = examples[0].size(0)
+
+    # Check if padding is necessary.
+
+    are_tensors_same_length = all(x.size(0) == length_of_first for x in examples)
+    if are_tensors_same_length and (pad_to_multiple_of is None or length_of_first % pad_to_multiple_of == 0):
+        return torch.stack(examples, dim=0)
+
+    # If yes, check if we have a `pad_token`.
+    if tokenizer._pad_token is None:
+        raise ValueError(
+            "You are attempting to pad samples but the tokenizer you are using"
+            f" ({tokenizer.__class__.__name__}) does not have a pad token."
+        )
+
+    # Creating the full tensor and filling it with our data.
+    max_length = max(x.size(0) for x in examples)
+    if pad_to_multiple_of is not None and (max_length % pad_to_multiple_of != 0):
+        max_length = ((max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of
+    result = examples[0].new_full([len(examples), max_length], tokenizer.pad_token_id)
+    for i, example in enumerate(examples):
+        if tokenizer.padding_side == "right":
+            result[i, : example.shape[0]] = example
+        else:
+            result[i, -example.shape[0] :] = example
+    return result
+
+def tolist(x):
+    if isinstance(x, list):
+        return x
+    elif hasattr(x, "numpy"):  # Checks for TF tensors without needing the import
+        x = x.numpy()
+    return x.tolist()
+
+#### END OF UTILS
+
+def do_topology_split(df, split_path):
+    import json
+    with open(split_path, 'r') as f:
+        splits = json.load(f)
+    
+    #split the dataframe according to the splits
+    train_df = df[df['name'].isin(splits['train'])]
+    valid_df = df[df['name'].isin(splits['validation'])]
+    test_df = df[df['name'].isin(splits['test'])]
+    return train_df, valid_df, test_df
+    
+
+class ANMAwareFlexibilityProtTrans(nn.Module):
+    def __init__(self, gumbel_temperature, **kwargs):
+        super(ANMAwareFlexibilityProtTrans, self).__init__()
+        
+        model, tokenizer = self.load_finetuned_model(**kwargs)
+        self.model = model
+        self.tokenizer = tokenizer
+        self.device = torch.device('cuda')
+        self.model.to(self.device)
+        self.model.eval()
+        self.gumbel_temperature = gumbel_temperature
+        self.logit_transform = nn.functional.gumbel_softmax #Use the Straight Through Gumbel SoftMax - in forward process it does argmax, 
+
+        # in the backward process it approximates the gradient of argmax by the gradient of the Gumbel Softmax
+        # https://pytorch.org/docs/stable/generated/torch.nn.functional.gumbel_softmax.html set hard=True to do the Straight-Through trick
+
+        self.conversion_tensor = self.construct_pmpnn_t5_conversion_tensor()
+        
+
+    def construct_pmpnn_t5_conversion_tensor(self):
+        """
+        Creates tensor which takes the onehot encodings in the proteinmpnn vocabulary and maps them to ProtTrans vocabulary.
+        """
+        _one_hots = []
+        for idx in [[0,1], 2, [3, 29, 30, 31, 32], 5, 4, 6, 7, 8, 10, 9, 13, 11, 12, 14, 15, 18, 16, 17, 19, 20, 21, 22, 23, 24, 25, 28, 26, 27]:
+            if isinstance(idx, int):
+                _oh = F.one_hot(torch.tensor([idx]), 33)
+            else:
+                _sohs = []
+                for subidx in idx:
+                    _soh = F.one_hot(torch.tensor([subidx]), 33)
+                    _sohs.append(_soh)
+                _oh = torch.sum(torch.stack(_sohs), dim=0)
+            _one_hots.append(_oh)
+        #_one_hots = [F.one_hot(torch.tensor([idx]), 33)[0] if isinstance(idx, int) else torch.sum(torch.stack([F.one_hot(torch.tensor([subidx]), 33)[0] for subidx in idx]), dim=0) for idx in [[0,1], 2, [3, 29, 30, 31, 32], 5, 4, 6, 7, 8, 10, 9, 13, 11, 12, 14, 15, 18, 16, 17, 19, 20, 21, 22, 23, 24, 25, 28, 26, 27]]
+        _one_hots.extend([torch.zeros((1,33)) for _ in range(100)])
+        return torch.cat(_one_hots, dim=0).to(torch.device('cuda')).float()
+
+    def load_finetuned_model(self, checkpoint_path, half_precision, **kwargs):#num_labels, add_pearson_loss, add_sse_loss, adaptor_architecture, enm_embed_dim, enm_att_heads, num_layers, kernel_size):
+        class_config=ClassConfig(**kwargs) #um_labels=num_labels, add_pearson_loss=add_pearson_loss, add_sse_loss=add_sse_loss, adaptor_architecture = adaptor_architecture, enm_embed_dim = enm_embed_dim, enm_att_heads = enm_att_heads, num_layers = num_layers, kernel_size = kernel_size)
+        model, tokenizer = PT5_classification_model(half_precision=half_precision, class_config=class_config) #.from_pretrained(args.model_path)
+
+        # model.load_state_dict(torch.load(args.model_path))
+        # try:
+        #     with safe_open(f"{checkpoint_path}/model.safetensors", framework="pt", device="cuda:0") as f:
+        #         state_dict = {}
+        #         for key in f.keys():
+        #             state_dict[key] = f.get_tensor(key)
+        #         model.load_state_dict(state_dict, strict=False)
+        # except:
+        #     state_dict = torch.load(f"{checkpoint_path}/pytorch_model.bin", map_location='cuda:0')
+        #     model.load_state_dict(state_dict, strict=False)
+        state_dict = torch.load(checkpoint_path, map_location='cuda:0')
+        model.load_state_dict(state_dict, strict=False)
+        model.eval()
+        
+        original_embedding = model.encoder.embed_tokens
+        model.new_embedding.weight = nn.Parameter(original_embedding.weight.T)
+        print('Set the weights for the new embedding layer!')
+        return model, tokenizer
+    
+    def translate_to_model_vocab(self, batch_one_hot, trail_idcs):
+            # Pad the batch_one_hot tensor with zeros along the last dimension
+            batch_one_hot = F.pad(batch_one_hot, (0, 1, 0, 0, 0, 0), 'constant', 0)
+            
+            #TODO: VERIFY THAT THE GRADIENTS ARE OK AFTER THE MASKED_SCATTER OPERATION
+            # Create a mask for the '2' token
+            mask = torch.zeros_like(batch_one_hot, dtype=torch.bool)
+            for i, trail_idx in enumerate(trail_idcs):
+                if trail_idx < batch_one_hot.size(2):  # Ensure index is within bounds
+                    mask[i, :, trail_idx] = True
+            
+            # Create a tensor with '2' in the one-hot encoding
+            token_2 = torch.zeros_like(batch_one_hot)
+            token_2[:, 2, :] = 1  # Assuming '2' corresponds to index 2 in the one-hot encoding
+            
+            # Use masked_scatter_ to modify the tensor in-place while preserving gradients
+            batch_one_hot.masked_scatter_(mask, token_2[mask])
+
+            T5_translation = torch.einsum('ej,ijk->iek', self.conversion_tensor, batch_one_hot)
+            T5_translation = T5_translation.permute(0,2,1)
+            return T5_translation
+    
+    def forward(self, pmpnn_logits, anm_input, trail_idcs, attention_mask, sampled_pmpnn_sequence = None, alphabet = None): #batch example 32x33x395 (batch_size x ProteinMPNN vocab size x seq length)
+    
+        anm_input = F.pad(anm_input, (0, 1, 0, 0), 'constant', 0)
+        attention_mask = F.pad(attention_mask, (0, 1, 0, 0), 'constant', 1)
+        
+        if sampled_pmpnn_sequence is None:
+            if alphabet is None:
+                batch_one_hot = self.logit_transform(pmpnn_logits, tau=self.gumbel_temperature, hard=True, dim=1)
+                batch_token_ids = self.translate_to_model_vocab(batch_one_hot, trail_idcs)
+                inputs = batch_token_ids #.to(torch.int)
+            # elif alphabet == 'aa':
+            #     batch_one_hot = ... #TODO one hot encode the pmpnn tokens
+            #     batch_token_ids = self.translate_to_model_vocab(batch_one_hot, trail_idcs)
+            #     input_ids = ... #TODO: argmax to get the tokens from the one hot encodings
+            #     outputs = self.model(input_ids = input_ids, enm_vals=anm_input, attention_mask = attention_mask) #TODO?: pass the mask as well (take it from the batch, pad it for the end of sequence, convert to Tensor)
+            #     predicted_flex = outputs.logits
+            #     return {'predicted_flex': predicted_flex, 'enm_vals': anm_input, 'input_ids': input_ids}
+                
+        # elif alphabet is None:
+        #     raise ValueError('need to specify what alphabet is used to encode sampled_pmpnn_sequence!')
+        # elif alphabet is 'pmpnn':
+        #     # Convert sampled_pmpnn_sequence to one-hot encoding
+        #     batch_one_hot = F.one_hot(sampled_pmpnn_sequence, num_classes=33).float().permute(0,2,1)
+        #     batch_token_ids = self.translate_to_model_vocab(batch_one_hot, trail_idcs)
+        #     inputs = batch_token_ids
+        # elif alphabet is 'pt5':
+        #     inputs = F.one_hot(sampled_pmpnn_sequence, num_classes=128).float() #.permute(0,2,1)
+        # elif alphabet is 'aa':
+        #     ... #TODO apply tokenizer
+        #     #tokens = self.tokenizer(" ".join(sampled_pmpnn_sequence))
+        #     tokens = self.tokenizer(" ".join(sampled_pmpnn_sequence))
+        #     input_ids = torch.tensor(tokens['input_ids']).cuda().unsqueeze(0)
+            
+        #     outputs = self.model(input_ids = input_ids, enm_vals=anm_input, attention_mask = attention_mask) #TODO?: pass the mask as well (take it from the batch, pad it for the end of sequence, convert to Tensor)
+        #     predicted_flex = outputs.logits
+        #     return {'predicted_flex': predicted_flex, 'enm_vals': anm_input, 'input_ids': input_ids}
+
+        inputs_embeds = self.model.new_embedding(inputs) #TODO pass through embedding
+        outputs = self.model(enm_vals=anm_input, inputs_embeds = inputs_embeds, attention_mask = attention_mask) #TODO?: pass the mask as well (take it from the batch, pad it for the end of sequence, convert to Tensor)
+        #TODO: above it throws RuntimeError: Expected tensor for argument #1 'indices' to have one of the following scalar types: 
+        # Long, Int; but got torch.cuda.FloatTensor instead (while checking arguments for embedding)
+        
+        predicted_flex = outputs.logits
+        return {'predicted_flex': predicted_flex, 'enm_vals': anm_input, 'input_ids': inputs}