: the output dir, default: current dir
+ -c : download a cif.gz file for each PDB id
+ -p : download a pdb.gz file for each PDB id (not available for large structures)
+ -a : download a pdb1.gz file (1st bioassembly) for each PDB id (not available for large structures)
+ -x : download a xml.gz file for each PDB id
+ -s : download a sf.cif.gz file for each PDB id (diffraction only)
+ -m : download a mr.gz file for each PDB id (NMR only)
+ -r : download a mr.str.gz for each PDB id (NMR only)
+EOF
+ exit 1
+}
+
+download() {
+ url="$BASE_URL/$1"
+ out=$2/$1
+ if test -e "$out"; then
+ echo "$out already exists, skipping"
+ return 1
+ fi
+ echo "Downloading $url to $out"
+ curl -s -f $url -o $out || echo "Failed to download $url"
+}
+
+listfile=""
+outdir="."
+cif=false
+pdb=false
+pdb1=false
+xml=false
+sf=false
+mr=false
+mrstr=false
+while getopts f:o:cpaxsmr o
+do
+ case $o in
+ (f) listfile=$OPTARG;;
+ (o) outdir=$OPTARG;;
+ (c) cif=true;;
+ (p) pdb=true;;
+ (a) pdb1=true;;
+ (x) xml=true;;
+ (s) sf=true;;
+ (m) mr=true;;
+ (r) mrstr=true;;
+ (*) usage
+ esac
+done
+shift "$((OPTIND - 1))"
+
+if [ "$listfile" == "" ]
+then
+ echo "Parameter -f must be provided"
+ exit 1
+fi
+contents=$(cat $listfile)
+
+# see https://stackoverflow.com/questions/918886/how-do-i-split-a-string-on-a-delimiter-in-bash#tab-top
+IFS=',' read -ra tokens <<< "$contents"
+
+for token in "${tokens[@]}"
+do
+ if [ "$cif" == true ]
+ then
+ download ${token}.cif.gz $outdir
+ fi
+ if [ "$pdb" == true ]
+ then
+ download ${token}.pdb.gz $outdir
+ fi
+ if [ "$pdb1" == true ]
+ then
+ download ${token}.pdb1.gz $outdir
+ fi
+ if [ "$xml" == true ]
+ then
+ download ${token}.xml.gz $outdir
+ fi
+ if [ "$sf" == true ]
+ then
+ download ${token}-sf.cif.gz $outdir
+ fi
+ if [ "$mr" == true ]
+ then
+ download ${token}.mr.gz $outdir
+ fi
+ if [ "$mrstr" == true ]
+ then
+ download ${token}_mr.str.gz $outdir
+ fi
+
+done
diff --git a/src/crawler/utils/get_uid_from_fasta.py b/src/crawler/utils/get_uid_from_fasta.py
new file mode 100644
index 0000000000000000000000000000000000000000..20149bd82a846f33ae496f84d77e857a7115b057
--- /dev/null
+++ b/src/crawler/utils/get_uid_from_fasta.py
@@ -0,0 +1,20 @@
+import argparse
+
+
+"""
+get uniprot ids from multi fasta file
+"""
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("-f", "--multi_fasta_file", type=str, default=None)
+ parser.add_argument("-u", "--uid_file", type=str, default=None)
+ args = parser.parse_args()
+
+ data = []
+ with open(args.multi_fasta_file, "r") as f:
+ for line in f:
+ if line.startswith(">"):
+ uid = line.split("|")[1]
+ data.append(uid)
+ with open(args.uid_file, "w") as f:
+ f.write("\n".join(data))
\ No newline at end of file
diff --git a/src/crawler/utils/get_unfold.py b/src/crawler/utils/get_unfold.py
new file mode 100644
index 0000000000000000000000000000000000000000..c777702b3e322532cb7831e583f952dc358f2fa1
--- /dev/null
+++ b/src/crawler/utils/get_unfold.py
@@ -0,0 +1,35 @@
+import argparse
+import os
+from utils import read_multi_fasta
+
+if __name__ == '__main__':
+ parser = argparse.ArgumentParser()
+ parser.add_argument("-f", "--fasta_file", type=str, default=None)
+ parser.add_argument("-d", "--af_dir", type=str, default=None)
+ parser.add_argument("-o", "--output_dir", type=str, default=None)
+ parser.add_argument("-c", "--chunk_size", type=int, default=50)
+ args = parser.parse_args()
+
+ for f in os.listdir(args.af_dir):
+ if "model" in f:
+ os.rename(os.path.join(args.af_dir, f), os.path.join(args.af_dir, f"{f.split('-')[1]}.pdb"))
+
+ downloaded_uids = [p[:-4] for p in os.listdir(args.af_dir)]
+ seqs = read_multi_fasta(args.fasta_file)
+ unfold_seqs = {}
+ for head, seq in seqs.items():
+ uid = head.split("|")[1].strip()
+ if uid not in downloaded_uids:
+ unfold_seqs[head] = seq
+ total_seqs = len(unfold_seqs)
+ print(f"Total unfold {total_seqs} sequences")
+ idx = 0
+ for head, seq in unfold_seqs.items():
+ uid = head.split("|")[1]
+ chunk_idx = idx // args.chunk_size
+ os.makedirs(os.path.join(args.output_dir, f"chunk_{chunk_idx}"), exist_ok=True)
+ with open(os.path.join(args.output_dir, f"chunk_{chunk_idx}", f"{uid}.fasta"), "w") as f:
+ f.write(f"{head}\n{seq}")
+ idx += 1
+ print("Done")
+
\ No newline at end of file
diff --git a/src/crawler/utils/metadata_utils.py b/src/crawler/utils/metadata_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..aa2474258a78ff069ec627130458accb1660d984
--- /dev/null
+++ b/src/crawler/utils/metadata_utils.py
@@ -0,0 +1,11 @@
+import json
+
+
+def get_uid_from_rcsb_meata_data(meta_data_file):
+ """
+ get uniprot ids from rcsb meta data file
+ """
+ with open(meta_data_file, "r") as f:
+ data = json.load(f)
+ uniprot_id = data["data"]["entry"]["polymer_entities"]["uniprots"]["rcsb_id"]
+ return uniprot_id
\ No newline at end of file
diff --git a/src/crawler/utils/pdb_filter.py b/src/crawler/utils/pdb_filter.py
new file mode 100644
index 0000000000000000000000000000000000000000..564f61dc5336bc42b32d9410d9a5472e7673e417
--- /dev/null
+++ b/src/crawler/utils/pdb_filter.py
@@ -0,0 +1,68 @@
+import argparse
+import os
+import sys
+import shutil
+import warnings
+warnings.filterwarnings("ignore")
+sys.path.append(os.getcwd())
+from tqdm import tqdm
+from Bio.PDB import PDBParser, PPBuilder
+from src.data_collect.utils import unzip, ungzip, get_seq_from_pdb, get_seqs_from_pdb
+
+
+def unzip_files(unzip_dir):
+ files = os.listdir(unzip_dir)
+ bar = tqdm(files)
+ for file in bar:
+ bar.set_postfix({"current": file})
+ ungzip(os.path.join(unzip_dir, file), unzip_dir)
+
+
+def is_apo(pdb_path):
+ parser = PDBParser()
+ structure = parser.get_structure('pdb', pdb_path)
+
+ for model in structure:
+ for chain in model:
+ for residue in chain:
+ hetero_flag = residue.id[0].strip()
+ if hetero_flag != '':
+ return False
+ return True
+
+
+def process(args):
+ if args.is_zip:
+ assert args.raw_dir, "no raw_dir"
+ unzip_files(args.raw_dir)
+
+ pdbs = sorted(os.listdir(args.raw_dir))
+ seq_pdb_dic = {}
+ bar = tqdm(pdbs)
+ for pdb in bar:
+ bar.set_postfix_str(f"{pdb}")
+ try:
+ seq = get_seq_from_pdb(os.path.join(args.raw_dir, pdb))
+ except:
+ continue
+ if seq in seq_pdb_dic.keys():
+ if is_apo(os.path.join(args.raw_dir, pdb)):
+ seq_pdb_dic[seq] = pdb
+ else:
+ continue
+ seq_pdb_dic[seq] = pdb
+ for seq, pdb in seq_pdb_dic.items():
+ shutil.copyfile(os.path.join(args.raw_dir, pdb), os.path.join(args.unique_dir, pdb))
+
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--is_zip", action="store_true")
+ parser.add_argument("--raw_dir", type=str, default="data/MDH/af/raw")
+ parser.add_argument("--unique_dir", type=str, default="data/MDH/af/unique")
+
+ args = parser.parse_args()
+ process(args)
+ # get_seqs_from_pdb("data/MDH/pdb/unique", "data/MDH/pdb/unique_seqs.fasta")
+
+
\ No newline at end of file
diff --git a/src/crawler/utils/secondary_calculate.py b/src/crawler/utils/secondary_calculate.py
new file mode 100644
index 0000000000000000000000000000000000000000..5d353f2d0a993d6bee35816a290b42de884dc4d2
--- /dev/null
+++ b/src/crawler/utils/secondary_calculate.py
@@ -0,0 +1,43 @@
+import argparse
+import os
+import sys
+sys.path.append(os.getcwd())
+import warnings
+warnings.filterwarnings("ignore")
+from tqdm import tqdm
+from Bio.PDB import PDBParser
+from Bio.PDB.DSSP import DSSP
+
+
+def process(args):
+ pdbs = sorted(os.listdir(args.pdb_dir))
+ bar = tqdm(pdbs)
+ wrong_pdb = []
+ for pdb in bar:
+ if os.path.exists(os.path.join(args.out_dir, pdb[:-4])):
+ continue
+ bar.set_postfix_str(f"{pdb}")
+ file = os.path.join(args.pdb_dir, pdb)
+ parser = PDBParser(QUIET=True)
+ structure = parser.get_structure("pdb", file)
+ model = structure[0]
+ try:
+ dssp = DSSP(model, file)
+ with open(os.path.join(args.out_dir, pdb[:-4]), "w") as f:
+ ss = ""
+ for residue in dssp:
+ ss += residue[2]
+ ss = ss.replace("-", "L")
+ f.write(ss)
+ except:
+ wrong_pdb.append(pdb)
+ print(wrong_pdb)
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--pdb_dir", type=str, default="data/MDH/pdb/process/PDB")
+ parser.add_argument("--out_dir", type=str, default="data/MDH/pdb/process/SS")
+
+ args = parser.parse_args()
+
+ process(args)
\ No newline at end of file
diff --git a/src/crawler/utils/utils.py b/src/crawler/utils/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..e68ff6cb3feddd3284d1ddbc5c42b537b198a0f2
--- /dev/null
+++ b/src/crawler/utils/utils.py
@@ -0,0 +1,77 @@
+import zipfile
+import gzip
+import shutil
+import os
+from tqdm import tqdm
+from Bio.PDB import PDBParser, PPBuilder
+
+def unzip(zipath, savefolder):
+ zf = zipfile.ZipFile(zipath)
+ zf.extractall(savefolder)
+ zf.close()
+
+def ungzip(file, out_dir):
+ with gzip.open(file, 'rb') as f_in:
+ with open(os.path.join(out_dir, file.split('/')[-1][:-3]), 'wb') as f_out:
+ shutil.copyfileobj(f_in, f_out)
+
+
+def get_seq_from_pdb(pdb_file):
+ parser = PDBParser()
+ structure = parser.get_structure(pdb_file[-8:-4], pdb_file)
+ ppb = PPBuilder()
+ chain = structure[0]['A']
+ seq = ""
+ for pp in ppb.build_peptides(chain):
+ seq += pp.get_sequence()
+ return seq
+
+
+def get_seqs_from_pdb(pdb_dir, out_file_path):
+ pdbs = os.listdir(pdb_dir)
+ with open(out_file_path, "w") as f:
+ for pdb in tqdm(pdbs):
+ seq = get_seq_from_pdb(os.path.join(pdb_dir, pdb))
+ f.write(f"> {pdb}\n{seq}\n")
+
+def read_multi_fasta(file_path):
+ """
+ params:
+ file_path: path to a fasta file
+ return:
+ a dictionary of sequences
+ """
+ sequences = {}
+ current_sequence = ''
+ with open(file_path, 'r') as file:
+ for line in file:
+ line = line.strip()
+ if line.startswith('>'):
+ if current_sequence:
+ sequences[header] = current_sequence
+ current_sequence = ''
+ header = line
+ else:
+ current_sequence += line
+ if current_sequence:
+ sequences[header] = current_sequence
+ return sequences
+
+def make_uid_chunks(uid_file, chunk_dir=None, chunk_size=10000):
+ """
+ params:
+ uid_file: path to a file containing a list of uniprot ids
+ chunk_size: size of each chunk
+ return:
+ files containing chunks of uniprot ids
+ """
+ uids = [f.strip() for f in open(uid_file, "r").readlines()]
+ uid_path = os.path.dirname(uid_file)
+ if chunk_dir is None:
+ chunk_dir = uid_path + "/chunks"
+ os.makedirs(chunk_dir, exist_ok=True)
+ chunk_num = len(uids) // chunk_size + 1
+ chunk_name = uid_file.split('/')[-1].split(".")[0]
+ for i in range(chunk_num):
+ with open(os.path.join(chunk_dir, f"{chunk_name}_{i}.txt"), "w") as f:
+ f.write("\n".join(uids[i*chunk_size:(i+1)*chunk_size]))
\ No newline at end of file
diff --git a/src/data/add_noise_to_backbone.py b/src/data/add_noise_to_backbone.py
new file mode 100644
index 0000000000000000000000000000000000000000..26771f937a29e4193dcffe25b6e279f98cd77188
--- /dev/null
+++ b/src/data/add_noise_to_backbone.py
@@ -0,0 +1,35 @@
+import numpy as np
+import os
+import json
+from tqdm import tqdm
+
+pdbs = os.listdir('alphafold_pdb')
+for pdb in tqdm(pdbs):
+ pdb_lines = open(f"alphafold_pdb/{pdb}").read().splitlines()
+
+ def add_noise_and_save(variance, file_name):
+ with open(file_name, "w") as file:
+ for line in pdb_lines:
+ if line.startswith("ATOM"):
+ parts = line.split()
+ try:
+ coords = np.array([float(parts[6]), float(parts[7]), float(parts[8])])
+ noise = np.random.normal(0, variance, coords.shape)
+ new_coords = coords + noise
+ new_line = f"{line[:30]}{new_coords[0]:8.3f}{new_coords[1]:8.3f}{new_coords[2]:8.3f}{line[54:]}"
+ file.write(new_line + "\n")
+ except:
+ file.write(line + "\n")
+ else:
+ file.write(line + "\n")
+
+ variances = [0.5]
+
+ for variance in variances:
+ file_name = f"alphafold_pdb_noise_{variance}/{pdb}"
+ try:
+ add_noise_and_save(variance, file_name)
+ except Exception as e:
+ print(e)
+ print(pdb)
+
diff --git a/src/data/batch_sampler.py b/src/data/batch_sampler.py
new file mode 100644
index 0000000000000000000000000000000000000000..4ed7fe17373804a7ae1d284c03efe05eec10247a
--- /dev/null
+++ b/src/data/batch_sampler.py
@@ -0,0 +1,45 @@
+import random
+from torch.utils.data import Sampler
+
+class BatchSampler(Sampler):
+ '''
+ 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_batch_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_batch_nodes=10000, shuffle=True):
+
+ self.node_counts = node_counts
+ self.idx = [i for i in range(len(node_counts)) if node_counts[i] <= max_batch_nodes]
+ self.shuffle = shuffle
+ self.max_batch_nodes = max_batch_nodes
+ self._form_batches()
+
+ def _form_batches(self):
+ self.batches = []
+ if self.shuffle: random.shuffle(self.idx)
+ idx = self.idx
+ while idx:
+ batch = []
+ max_n_node = 0
+ while idx:
+ if max(self.node_counts[idx[0]], max_n_node) * (len(batch) + 1) > self.max_batch_nodes:
+ break
+ next_idx, idx = idx[0], idx[1:]
+ current_n_node = self.node_counts[next_idx]
+ if current_n_node > max_n_node:
+ max_n_node = current_n_node
+ 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
diff --git a/src/data/collator.py b/src/data/collator.py
new file mode 100644
index 0000000000000000000000000000000000000000..06b7d28e7f68ddc4f6e82469b3eaa323793ff3d5
--- /dev/null
+++ b/src/data/collator.py
@@ -0,0 +1,168 @@
+import torch
+import re
+from typing import Dict, List, Any
+from transformers import PreTrainedTokenizer
+from dataclasses import dataclass
+
+VQVAE_CODEBOOK_SIZE = 4096
+VQVAE_SPECIAL_TOKENS = {
+ "MASK": VQVAE_CODEBOOK_SIZE,
+ "EOS": VQVAE_CODEBOOK_SIZE + 1,
+ "BOS": VQVAE_CODEBOOK_SIZE + 2,
+ "PAD": VQVAE_CODEBOOK_SIZE + 3,
+ "CHAINBREAK": VQVAE_CODEBOOK_SIZE + 4,
+}
+
+@dataclass
+class Collator:
+ """Data collator class for protein sequences."""
+ tokenizer: PreTrainedTokenizer
+ max_length: int = None
+ structure_seq: List[str] = None
+ problem_type: str = 'classification'
+ plm_model: str = None
+ num_labels: int = None
+
+ def __call__(self, examples: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
+ """Collate function for batching examples."""
+ # Initialize lists to store sequences and labels
+ if "ProSST" in self.plm_model:
+ aa_seqs, labels, str_tokens = [], [], []
+ else:
+ aa_seqs, labels = [], []
+ structure_seqs = {
+ seq_type: [] for seq_type in (self.structure_seq or [])
+ }
+
+ # Process each example
+ for e in examples:
+ # Process sequences
+ aa_seq = self.process_sequence(e["aa_seq"])
+ aa_seqs.append(aa_seq)
+ if "ProSST" in self.plm_model:
+ stru_token = self.process_stru_tokens(e["prosst_stru_token"])
+ str_tokens.append(stru_token)
+
+ # Process structure sequences if needed
+ for seq_type in structure_seqs:
+ if seq_type == 'esm3_structure_seq':
+ processed_seq = self.process_esm3_structure_seq(e[seq_type])
+ else:
+ processed_seq = self.process_sequence(e[seq_type])
+ structure_seqs[seq_type].append(processed_seq)
+
+ if self.problem_type == 'multi_label_classification':
+ label_list = e['label'].split(',')
+ e['label'] = [int(l) for l in label_list]
+ binary_list = [0] * self.num_labels
+ for index in e['label']:
+ binary_list[index] = 1
+ e['label'] = binary_list
+ # Process labels
+ labels.append(e["label"])
+
+ # Tokenize sequences
+ if "ProSST" in self.plm_model:
+ batch = self.tokenize_sequences(aa_seqs, structure_seqs, str_tokens)
+ else:
+ batch = self.tokenize_sequences(aa_seqs, structure_seqs)
+
+ # Add labels to batch
+ batch["label"] = torch.as_tensor(
+ labels,
+ dtype=torch.float if self.problem_type == 'regression' else torch.long
+ )
+
+ return batch
+
+ def process_sequence(self, seq: str) -> str:
+ """Process sequence based on model type."""
+ if 'prot_bert' in self.plm_model or "prot_t5" in self.plm_model:
+ seq = " ".join(list(seq))
+ seq = re.sub(r"[UZOB]", "X", seq)
+ return seq
+
+ def process_esm3_structure_seq(self, seq: List[int]) -> torch.Tensor:
+ """Process ESM3 structure sequence."""
+ return torch.tensor([VQVAE_SPECIAL_TOKENS["BOS"]] + seq + [VQVAE_SPECIAL_TOKENS["EOS"]])
+
+ def process_stru_tokens(self, seq:List[int]) -> torch.Tensor:
+ """Process ProSST structure token."""
+ if isinstance(seq, str):
+ seq_clean = seq.strip("[]").replace(" ","")
+ tokens = list(map(int, seq_clean.split(','))) if seq_clean else []
+ elif isinstance(seq, (list, tuple)):
+ tokens = [int(x) for x in seq]
+ stru_tokens = [int(num) for num in tokens]
+ return torch.tensor(stru_tokens)
+
+ def tokenize_sequences(
+ self,
+ aa_seqs: List[str],
+ structure_seqs: Dict[str, List[str]],
+ str_tokens: List[str] = None,
+ ) -> Dict[str, torch.Tensor]:
+ """Tokenize all sequences."""
+ # Process amino acid sequences
+ if "esm1b" in self.plm_model or "esm1v" in self.plm_model:
+ self.max_length = 1022
+ aa_encodings = self.tokenizer(
+ aa_seqs,
+ padding=True,
+ truncation=True if self.max_length else False,
+ max_length=self.max_length,
+ return_tensors="pt"
+ )
+
+ aa_max_length = len(aa_encodings["input_ids"][0])
+ padded_tokens = []
+ if str_tokens:
+ for tokens in str_tokens:
+ struct_sequence = [int(num) for num in tokens]
+ padded_tokens.append(struct_sequence + [0] * (aa_max_length - len(struct_sequence) - 2))
+ batch = {
+ "aa_seq_input_ids": aa_encodings["input_ids"],
+ "aa_seq_attention_mask": aa_encodings["attention_mask"],
+ "aa_seq_stru_tokens": torch.tensor(padded_tokens, dtype=torch.long)
+ }
+ else:
+ batch = {
+ "aa_seq_input_ids": aa_encodings["input_ids"],
+ "aa_seq_attention_mask": aa_encodings["attention_mask"]
+ }
+
+ # Process structure sequences if provided
+ for seq_type, seqs in structure_seqs.items():
+ if not seqs:
+ continue
+
+ if seq_type == 'esm3_structure_seq':
+ # ESM3 structure sequences are already tokenized
+ structure_tokens = torch.stack(seqs)
+ # Pad sequences to max length
+ max_len = max(len(seq) for seq in seqs)
+ padded_tokens = torch.zeros(len(seqs), max_len, dtype=torch.long)
+ attention_mask = torch.zeros(len(seqs), max_len, dtype=torch.long)
+
+ for i, seq in enumerate(seqs):
+ seq_len = len(seq)
+ padded_tokens[i, :seq_len] = seq
+ attention_mask[i, :seq_len] = 1
+
+ batch[f"{seq_type}_input_ids"] = padded_tokens
+ batch[f"{seq_type}_attention_mask"] = attention_mask
+
+ else:
+ # Tokenize other structure sequences
+ structure_encodings = self.tokenizer(
+ seqs,
+ padding=True,
+ truncation=True if self.max_length else False,
+ max_length=self.max_length,
+ return_tensors="pt"
+ )
+
+ batch[f"{seq_type}_input_ids"] = structure_encodings["input_ids"]
+ batch[f"{seq_type}_attention_mask"] = structure_encodings["attention_mask"]
+
+ return batch
\ No newline at end of file
diff --git a/src/data/dataloader.py b/src/data/dataloader.py
new file mode 100644
index 0000000000000000000000000000000000000000..e9591a8fdc65172bbadd8196ff7f0cbfe5ef4e98
--- /dev/null
+++ b/src/data/dataloader.py
@@ -0,0 +1,78 @@
+import json
+import torch
+import datasets
+from torch.utils.data import DataLoader
+from .collator import Collator
+from .batch_sampler import BatchSampler
+from .norm import normalize_dataset
+from torch.utils.data import Dataset
+from typing import Dict, Any, List, Union
+import pandas as pd
+
+def prepare_dataloaders(args, tokenizer, logger):
+ """Prepare train, validation and test dataloaders."""
+ # Process datasets
+ train_dataset = datasets.load_dataset(args.dataset)['train']
+ train_dataset_token_lengths = [len(item['aa_seq']) for item in train_dataset]
+ val_dataset = datasets.load_dataset(args.dataset)['validation']
+ val_dataset_token_lengths = [len(item['aa_seq']) for item in val_dataset]
+ test_dataset = datasets.load_dataset(args.dataset)['test']
+ test_dataset_token_lengths = [len(item['aa_seq']) for item in test_dataset]
+
+ if args.normalize is not None:
+ train_dataset, val_dataset, test_dataset = normalize_dataset(train_dataset, val_dataset, test_dataset, args.normalize)
+
+ # log dataset info
+ logger.info("Dataset Statistics:")
+ logger.info("------------------------")
+ logger.info(f"Dataset: {args.dataset}")
+ logger.info(f" Number of train samples: {len(train_dataset)}")
+ logger.info(f" Number of val samples: {len(val_dataset)}")
+ logger.info(f" Number of test samples: {len(test_dataset)}")
+
+ # log 3 data points from train_dataset
+ logger.info("Sample 3 data points from train dataset:")
+ logger.info(f" Train data point 1: {train_dataset[0]}")
+ logger.info(f" Train data point 2: {train_dataset[1]}")
+ logger.info(f" Train data point 3: {train_dataset[2]}")
+ logger.info("------------------------")
+
+ collator = Collator(
+ tokenizer=tokenizer,
+ max_length=args.max_seq_len if args.max_seq_len > 0 else None,
+ structure_seq=args.structure_seq,
+ problem_type=args.problem_type,
+ plm_model=args.plm_model,
+ num_labels=args.num_labels
+ )
+
+ # Common dataloader parameters
+ dataloader_params = {
+ 'num_workers': args.num_workers,
+ 'collate_fn': collator,
+ 'pin_memory': True,
+ 'persistent_workers': True if args.num_workers > 0 else False,
+ 'prefetch_factor': 2,
+ }
+
+ # Create dataloaders based on batching strategy
+ if args.batch_token is not None:
+ train_loader = create_token_based_loader(train_dataset, train_dataset_token_lengths, args.batch_token, True, **dataloader_params)
+ val_loader = create_token_based_loader(val_dataset, val_dataset_token_lengths, args.batch_token, False, **dataloader_params)
+ test_loader = create_token_based_loader(test_dataset, test_dataset_token_lengths, args.batch_token, False, **dataloader_params)
+ else:
+ train_loader = create_size_based_loader(train_dataset, args.batch_size, True, **dataloader_params)
+ val_loader = create_size_based_loader(val_dataset, args.batch_size, False, **dataloader_params)
+ test_loader = create_size_based_loader(test_dataset, args.batch_size, False, **dataloader_params)
+
+ return train_loader, val_loader, test_loader
+
+def create_token_based_loader(dataset, token_lengths, batch_token, shuffle, **kwargs):
+ """Create dataloader with token-based batching."""
+ sampler = BatchSampler(token_lengths, batch_token, shuffle=shuffle)
+ return DataLoader(dataset, batch_sampler=sampler, **kwargs)
+
+def create_size_based_loader(dataset, batch_size, shuffle, **kwargs):
+ """Create dataloader with size-based batching."""
+ return DataLoader(dataset, batch_size=batch_size, shuffle=shuffle, **kwargs)
+
diff --git a/src/data/get_esm3_structure_seq.py b/src/data/get_esm3_structure_seq.py
new file mode 100644
index 0000000000000000000000000000000000000000..a9973d0afdfc733c13b891ad1f6fb0c334569bd4
--- /dev/null
+++ b/src/data/get_esm3_structure_seq.py
@@ -0,0 +1,85 @@
+import torch
+import os
+import sys
+sys.path.append(os.getcwd())
+import json
+import argparse
+import numpy as np
+import biotite.structure.io as bsio
+from tqdm import tqdm
+from biotite.structure.io.pdb import PDBFile
+from esm.utils.structure.protein_chain import ProteinChain
+from esm.models.vqvae import StructureTokenEncoder
+
+VQVAE_CODEBOOK_SIZE = 4096
+VQVAE_SPECIAL_TOKENS = {
+ "MASK": VQVAE_CODEBOOK_SIZE,
+ "EOS": VQVAE_CODEBOOK_SIZE + 1,
+ "BOS": VQVAE_CODEBOOK_SIZE + 2,
+ "PAD": VQVAE_CODEBOOK_SIZE + 3,
+ "CHAINBREAK": VQVAE_CODEBOOK_SIZE + 4,
+}
+
+def ESM3_structure_encoder_v0(device: torch.device | str = "cpu"):
+ model = (
+ StructureTokenEncoder(
+ d_model=1024, n_heads=1, v_heads=128, n_layers=2, d_out=128, n_codes=4096
+ )
+ .to(device)
+ .eval()
+ )
+ state_dict = torch.load(
+ "./src/data/weight/esm3_structure_encoder_v0.pth", map_location=device
+ )
+ model.load_state_dict(state_dict)
+ return model
+
+def get_esm3_structure_seq(pdb_file, encoder, device="cuda:0"):
+ # Extract Unique Chain IDs
+ chain_ids = np.unique(PDBFile.read(pdb_file).get_structure().chain_id)
+ # print(chain_ids)
+ # ['L', 'H']
+
+ # By Default, ProteinChain takes first one
+ chain = ProteinChain.from_pdb(pdb_file, chain_id=chain_ids[0])
+
+ # Encoder
+ coords, plddt, residue_index = chain.to_structure_encoder_inputs()
+ coords = coords.to(device)
+ struct = bsio.load_structure(pdb_file, extra_fields=["b_factor"])
+ plddt = struct.b_factor.mean()
+ residue_index = residue_index.to(device)
+ _, structure_tokens = encoder.encode(coords, residue_index=residue_index)
+
+ result = {
+ 'name': pdb_file.split('/')[-1].split('.')[0],
+ 'esm3_structure_seq':structure_tokens.cpu().numpy().tolist()[0],
+ 'plddt':plddt
+ }
+ return result
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--pdb_file", type=str, default=None)
+ parser.add_argument("--pdb_dir", type=str, default=None)
+ parser.add_argument("--out_file", type=str, default='esm3_structure_seq.json')
+ args = parser.parse_args()
+
+ device="cuda:0"
+ results = []
+ # result_dict = {'name':[], 'aa_seq':[], 'esm3_structure_seq':[], 'plddt':[], 'residue_index':[]}
+
+ encoder = ESM3_structure_encoder_v0(device)
+
+ if args.pdb_file is not None:
+ result = get_esm3_structure_seq(args.pdb_file, encoder, device)
+ results.append(result)
+
+ elif args.pdb_dir is not None:
+ pdb_files = os.listdir(args.pdb_dir)
+ for pdb_file in tqdm(pdb_files):
+ result = get_esm3_structure_seq(os.path.join(args.pdb_dir, pdb_file), encoder, device)
+ results.append(result)
+
+ with open(args.out_file, "w") as f:
+ f.write("\n".join([json.dumps(r) for r in results]))
diff --git a/src/data/get_foldseek_structure_seq.py b/src/data/get_foldseek_structure_seq.py
new file mode 100644
index 0000000000000000000000000000000000000000..6f765e2c5b1f93e1263ac0512c2f427530a41448
--- /dev/null
+++ b/src/data/get_foldseek_structure_seq.py
@@ -0,0 +1,41 @@
+import os
+import argparse
+import json
+from tqdm import tqdm
+
+# conda install -c conda-forge -c bioconda foldseek
+def get_foldseek_structure_seq(pdb_dir, rm_tmp=True):
+ # foldseek createdb INPUT_dir_with_structures tmp_db
+ # foldseek lndb tmp_db_h tmp_db_ss_h
+ # foldseek convert2fasta tmp_db_ss OUTPUT_3di.fasta
+ # use command to generate foldseek structure seq
+ os.makedirs("tmp_db", exist_ok=True)
+ os.system(f"foldseek createdb {pdb_dir} tmp_db/tmp_db")
+ os.system(f"foldseek lndb tmp_db/tmp_db_h tmp_db/tmp_db_ss_h")
+ os.system(f"foldseek convert2fasta tmp_db/tmp_db_ss tmp_db/tmp_db_ss.fasta")
+
+ results = []
+ # read fasta file
+ with open("tmp_db/tmp_db_ss.fasta", "r") as f:
+ for line in tqdm(f):
+ if line.startswith(">"):
+ name = line.split()[0][1:]
+ seq = next(f).strip()
+ results.append({"name":name.split('.')[0], "foldseek_seq":seq})
+
+ if rm_tmp:
+ os.system("rm -rf tmp_db")
+
+ return results
+
+
+if __name__ == '__main__':
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--pdb_dir", type=str, default=None)
+ parser.add_argument("--out_file", type=str, default=None)
+ parser.add_argument("--rm_tmp", type=bool, default=True)
+ args = parser.parse_args()
+
+ results = get_foldseek_structure_seq(args.pdb_dir, args.rm_tmp)
+ with open(args.out_file, "w") as f:
+ f.write("\n".join([json.dumps(r) for r in results]))
\ No newline at end of file
diff --git a/src/data/get_prosst_str_token.py b/src/data/get_prosst_str_token.py
new file mode 100644
index 0000000000000000000000000000000000000000..5eb050c8ea1d18c205fbbc9163d0aff24c0fd513
--- /dev/null
+++ b/src/data/get_prosst_str_token.py
@@ -0,0 +1,106 @@
+import os
+import sys
+import argparse
+import json
+import pandas as pd
+import torch
+from tqdm import tqdm
+from Bio import SeqIO
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from src.data.prosst.structure.quantizer import PdbQuantizer
+from src.utils.data_utils import extract_seq_from_pdb
+import warnings
+warnings.filterwarnings("ignore", category=Warning)
+structure_vocab_size = 20
+processor = PdbQuantizer(structure_vocab_size = structure_vocab_size)
+
+def get_prosst_token(pdb_file):
+ """Generate ProSST structure tokens for a PDB file"""
+ try:
+ # 提取氨基酸序列
+ aa_seq = extract_seq_from_pdb(pdb_file)
+
+ # 处理结构序列
+ structure_result = processor(pdb_file)
+ pdb_name = os.path.basename(pdb_file)
+ # 验证数据结构
+ if structure_vocab_size not in structure_result:
+ raise ValueError(f"Missing structure key: {structure_vocab_size}")
+ if pdb_name not in structure_result[structure_vocab_size]:
+ raise ValueError(f"Missing PDB entry: {pdb_name}")
+
+ struct_sequence = structure_result[structure_vocab_size][pdb_name]['struct']
+ struct_sequence = [int(num) for num in struct_sequence]
+
+ # 添加特殊标记 [1] + sequence + [2]
+ structure_sequence_offset = [3 + num for num in struct_sequence]
+ structure_input_ids = torch.tensor(
+ [[1] + structure_sequence_offset + [2]],
+ dtype=torch.long
+ )
+
+ return {
+ "name": os.path.basename(pdb_file).split('.')[0],
+ "aa_seq": aa_seq,
+ "struct_tokens": structure_input_ids[0].tolist()
+ }, None
+
+ except Exception as e:
+ return pdb_file, f"{str(e)}"
+
+if __name__ == '__main__':
+ parser = argparse.ArgumentParser(description='ProSST structure token generator')
+ parser.add_argument('--pdb_dir', type=str, help='Directory containing PDB files')
+ parser.add_argument('--pdb_file', type=str, help='Single PDB file path')
+ parser.add_argument('--num_workers', type=int, default=16, help='Number of parallel workers')
+ parser.add_argument('--pdb_index_file', type=str, default=None, help='PDB index file for sharding')
+ parser.add_argument('--pdb_index_level', type=int, default=1, help='Directory hierarchy depth')
+ parser.add_argument('--error_file', type=str, help='Error log output path')
+ parser.add_argument('--out_file', type=str, required=True, help='Output JSON file path')
+ args = parser.parse_args()
+
+ if args.pdb_dir is not None:
+ # load pdb index file
+ if args.pdb_index_file:
+ pdbs = open(args.pdb_index_file).read().splitlines()
+ pdb_files = []
+ for pdb in pdbs:
+ pdb_relative_dir = args.pdb_dir
+ for i in range(1, args.pdb_index_level+1):
+ pdb_relative_dir = os.path.join(pdb_relative_dir, pdb[:i])
+ pdb_files.append(os.path.join(pdb_relative_dir, pdb+".pdb"))
+
+ # regular pdb dir
+ else:
+ pdb_files = sorted([os.path.join(args.pdb_dir, p) for p in os.listdir(args.pdb_dir)])
+
+
+ # 并行处理
+ results, errors = [], []
+ with ThreadPoolExecutor(max_workers=args.num_workers) as executor:
+ futures = {executor.submit(get_prosst_token, f): f for f in pdb_files}
+ with tqdm(total=len(futures), desc="Processing PDBs") as progress:
+ for future in as_completed(futures):
+ result, error = future.result()
+ if error:
+ errors.append({"file": result, "error": error})
+ else:
+ results.append(result)
+ progress.update(1)
+
+ if errors:
+ error_path = args.error_file or args.out_file.replace('.json', '_errors.csv')
+ pd.DataFrame(errors).to_csv(error_path, index=False)
+ print(f"Encountered {len(errors)} errors. Saved to {error_path}")
+
+
+ with open(args.out_file, 'w') as f:
+ f.write('\n'.join(json.dumps(r) for r in results))
+
+
+ elif args.pdb_file:
+ result, error = get_prosst_token(args.pdb_file)
+ if error:
+ raise RuntimeError(f"Error processing {args.pdb_file}: {error}")
+ with open(args.out_file, 'w') as f:
+ json.dump(result, f)
\ No newline at end of file
diff --git a/src/data/get_secondary_structure_seq.py b/src/data/get_secondary_structure_seq.py
new file mode 100644
index 0000000000000000000000000000000000000000..74385569d38ec38e4dab5ac2ba54c609d74b9dca
--- /dev/null
+++ b/src/data/get_secondary_structure_seq.py
@@ -0,0 +1,118 @@
+import os
+import sys
+sys.path.append(os.getcwd())
+import argparse
+import json
+import pandas as pd
+import biotite.structure.io as bsio
+from tqdm import tqdm
+from Bio import PDB
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from src.utils.data_utils import extract_seq_from_pdb
+
+
+ss_alphabet = ['H', 'E', 'C']
+ss_alphabet_dic = {
+ "H": "H", "G": "H", "E": "E",
+ "B": "E", "I": "C", "T": "C",
+ "S": "C", "L": "C", "-": "C",
+ "P": "C"
+}
+
+def get_secondary_structure_seq(pdb_file):
+ try:
+ # extract amino acid sequence
+ aa_seq = extract_seq_from_pdb(pdb_file)
+ pdb_parser = PDB.PDBParser(QUIET=True)
+ structure = pdb_parser.get_structure("protein", pdb_file)
+ model = structure[0]
+ dssp = PDB.DSSP(model, pdb_file)
+ # extract secondary structure sequence
+ sec_structures = []
+ for i, dssp_res in enumerate(dssp):
+ sec_structures.append(dssp_res[2])
+
+ except Exception as e:
+ return pdb_file, e
+
+ sec_structure_str_8 = ''.join(sec_structures)
+ sec_structure_str_8 = sec_structure_str_8.replace('-', 'L')
+ if len(aa_seq) != len(sec_structure_str_8):
+ return pdb_file, f"aa_seq {len(aa_seq)} and sec_structure_str_8 {len(sec_structure_str_8)} length mismatch"
+
+ sec_structure_str_3 = ''.join([ss_alphabet_dic[ss] for ss in sec_structures])
+
+ final_dict = {}
+ final_dict["name"] = pdb_file.split('/')[-1].split('.')[0]
+ final_dict["aa_seq"] = aa_seq
+ final_dict["ss8_seq"] = sec_structure_str_8
+ final_dict["ss3_seq"] = sec_structure_str_3
+
+ return final_dict, None
+
+if __name__ == '__main__':
+ parser = argparse.ArgumentParser()
+ parser.add_argument('--pdb_dir', type=str, help='pdb dir')
+ parser.add_argument('--pdb_file', type=str, help='pdb file')
+
+ # multi processing
+ parser.add_argument('--num_workers', type=int, default=4, help='number of workers')
+
+ # index pdb for large scale inference
+ parser.add_argument("--pdb_index_file", default=None, type=str, help="pdb index file")
+ parser.add_argument("--pdb_index_level", default=1, type=int, help="pdb index level")
+
+ # save file
+ parser.add_argument('--error_file', type=str, help='save error file')
+ parser.add_argument('--out_file', type=str, help='save file')
+ args = parser.parse_args()
+
+ out_dir = os.path.dirname(args.out_file)
+ os.makedirs(out_dir, exist_ok=True)
+
+ if args.pdb_dir is not None:
+ # load pdb index file
+ if args.pdb_index_file:
+ pdbs = open(args.pdb_index_file).read().splitlines()
+ pdb_files = []
+ for pdb in pdbs:
+ pdb_relative_dir = args.pdb_dir
+ for i in range(1, args.pdb_index_level+1):
+ pdb_relative_dir = os.path.join(pdb_relative_dir, pdb[:i])
+ pdb_files.append(os.path.join(pdb_relative_dir, pdb+".pdb"))
+
+ # regular pdb dir
+ else:
+ pdb_files = sorted([os.path.join(args.pdb_dir, p) for p in os.listdir(args.pdb_dir)])
+
+ results, error_pdbs, error_messages = [], [], []
+ with ThreadPoolExecutor(max_workers=args.num_workers) as executor:
+ futures = [executor.submit(get_secondary_structure_seq, pdb_file) for pdb_file in pdb_files]
+
+ with tqdm(total=len(pdb_files), desc="Processing pdb") as progress:
+ for future in as_completed(futures):
+ result, message = future.result()
+ if message is None:
+ results.append(result)
+ else:
+ error_pdbs.append(result)
+ error_messages.append(message)
+ progress.update(1)
+ progress.close()
+
+ if error_pdbs:
+ if args.error_file is None:
+ args.error_file = args.out_file.split(".")[0]+"_error.csv"
+ error_dir = os.path.dirname(args.error_file)
+ os.makedirs(error_dir, exist_ok=True)
+ error_info = {"error_pdbs": error_pdbs, "error_messages": error_messages}
+ pd.DataFrame(error_info).to_csv(args.error_file, index=False)
+
+ with open(args.out_file, "w") as f:
+ f.write("\n".join([json.dumps(r) for r in results]))
+
+ elif args.pdb_file is not None:
+ result, message = get_secondary_structure_seq(args.pdb_file)
+ with open(args.out_file, "w") as f:
+ json.dump(result, f)
+
\ No newline at end of file
diff --git a/src/data/get_ss_fs_esm3_seq.py b/src/data/get_ss_fs_esm3_seq.py
new file mode 100644
index 0000000000000000000000000000000000000000..8d795b505082faa0da00f51ae48611c58b2a529c
--- /dev/null
+++ b/src/data/get_ss_fs_esm3_seq.py
@@ -0,0 +1,93 @@
+import os
+import json
+import argparse
+import torch
+import pandas as pd
+from tqdm import tqdm
+from esm.models.vqvae import StructureTokenEncoder
+from get_esm3_structure_seq import get_esm3_structure_seq
+from get_foldseek_structure_seq import get_foldseek_structure_seq
+from get_secondary_structure_seq import get_secondary_structure_seq
+from get_prosst_str_token import get_prosst_token
+
+# ignore the warning
+import warnings
+warnings.filterwarnings("ignore")
+
+def ESM3_structure_encoder_v0(device: torch.device | str = "cpu"):
+ model = (
+ StructureTokenEncoder(
+ d_model=1024, n_heads=1, v_heads=128, n_layers=2, d_out=128, n_codes=4096
+ )
+ .to(device)
+ .eval()
+ )
+ state_dict = torch.load(
+ "./src/data/weight/esm3_structure_encoder_v0.pth", map_location=device
+ )
+ model.load_state_dict(state_dict)
+ return model
+
+if __name__ == "__main__":
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--pdb_dir", type=str, default='dataset/sesadapter/DeepET/esmfold_pdb')
+ parser.add_argument("--pdb_file", type=str, default=None)
+ parser.add_argument("--out_dir", type=str, default='dataset/sesadapter/DeepET')
+ parser.add_argument("--merge_into", type=str, default='csv', choices=['json', 'csv'])
+ parser.add_argument("--save_intermediate", action='store_true')
+ args = parser.parse_args()
+
+ device = "cuda:0"
+ esm3_encoder = ESM3_structure_encoder_v0(device)
+
+ if args.pdb_dir is not None:
+ dir_name = os.path.basename(args.pdb_dir)
+ pdb_files = os.listdir(args.pdb_dir)
+ ss_results, esm3_results = [], []
+ for pdb_file in tqdm(pdb_files):
+ ss_result, error = get_secondary_structure_seq(os.path.join(args.pdb_dir, pdb_file))
+ if error is not None:
+ print(error)
+ continue
+ ss_results.append(ss_result)
+ esm3_result = get_esm3_structure_seq(os.path.join(args.pdb_dir, pdb_file), esm3_encoder, device)
+ esm3_results.append(esm3_result)
+ # clear cuda cache
+ torch.cuda.empty_cache()
+ with open(os.path.join(args.out_dir, f"{dir_name}_ss.json"), "w") as f:
+ f.write("\n".join([json.dumps(r) for r in ss_results]))
+ with open(os.path.join(args.out_dir, f"{dir_name}_esm3.json"), "w") as f:
+ f.write("\n".join([json.dumps(r) for r in esm3_results]))
+
+ fs_results = get_foldseek_structure_seq(args.pdb_dir)
+ with open(os.path.join(args.out_dir, f"{dir_name}_fs.json"), "w") as f:
+ f.write("\n".join([json.dumps(r) for r in fs_results]))
+ prosst_tokens = get_prosst_token(args.pdb_dir)
+ with open(os.path.join(args.out_dir, f"{dir_name}_prosst.json"), "r") as f:
+ f.write("\n".join([json.dumps(r) for r in prosst_tokens]))
+
+ if args.merge_into == 'csv':
+ # read json files and merge to a single csv according to the same 'name' column
+ ss_json = os.path.join(args.out_dir, f"{dir_name}_ss.json")
+ esm3_json = os.path.join(args.out_dir, f"{dir_name}_esm3.json")
+ fs_json = os.path.join(args.out_dir, f"{dir_name}_fs.json")
+ prosst_json = os.path.join(args.out_dir, f"{dir_name}_prosst.json")
+ # load json line files
+ ss_df = pd.read_json(ss_json, lines=True)
+ esm3_df = pd.read_json(esm3_json, lines=True)
+ fs_df = pd.read_json(fs_json, lines=True)
+ prosst_json = os.path_join(prosst_json, lines=True)
+ # merge the three dataframes by the 'name' column
+ df = pd.merge(ss_df, fs_df, on='name', how='inner')
+ df = pd.merge(df, esm3_df, on='name', how='inner')
+ df = pd.merge(df, prosst_json, on='name', how='inner')
+ # sort by name
+ df = df.sort_values(by='name')
+ df.to_csv(os.path.join(args.out_dir, f"{dir_name}.csv"), index=False)
+
+ if not args.save_intermediate:
+ # remove intermediate files
+ os.remove(ss_json)
+ os.remove(esm3_json)
+ os.remove(fs_json)
+ os.remove(prosst_json)
diff --git a/src/data/norm.py b/src/data/norm.py
new file mode 100644
index 0000000000000000000000000000000000000000..f539c7c8a717d999c7b99c6f9500293d497a0e18
--- /dev/null
+++ b/src/data/norm.py
@@ -0,0 +1,127 @@
+import numpy as np
+from typing import List, Tuple, Any
+from sklearn.preprocessing import StandardScaler, RobustScaler
+
+def min_max_normalize_dataset(train_dataset, val_dataset, test_dataset):
+ """Min-max normalization (0-1 scaling)."""
+ labels = [e["label"] for e in train_dataset]
+ min_label, max_label = min(labels), max(labels)
+ normalized_train_dataset = []
+ normalized_val_dataset = []
+ normalized_test_dataset = []
+ for e in train_dataset:
+ e["label"] = (e["label"] - min_label) / (max_label - min_label)
+ normalized_train_dataset.append(e)
+ for e in val_dataset:
+ e["label"] = (e["label"] - min_label) / (max_label - min_label)
+ normalized_val_dataset.append(e)
+ for e in test_dataset:
+ e["label"] = (e["label"] - min_label) / (max_label - min_label)
+ normalized_test_dataset.append(e)
+ print(normalized_train_dataset[0])
+ return normalized_train_dataset, normalized_val_dataset, normalized_test_dataset
+
+def standard_normalize_dataset(train_dataset, val_dataset, test_dataset):
+ """Z-score normalization (standardization)."""
+ train_labels = np.array([e["label"] for e in train_dataset])
+ mean_label = np.mean(train_labels)
+ std_label = np.std(train_labels)
+ normalized_train_dataset = []
+ normalized_val_dataset = []
+ normalized_test_dataset = []
+ for e in train_dataset:
+ e["label"] = (e["label"] - mean_label) / std_label
+ normalized_train_dataset.append(e)
+ for e in val_dataset:
+ e["label"] = (e["label"] - mean_label) / std_label
+ normalized_val_dataset.append(e)
+ for e in test_dataset:
+ e["label"] = (e["label"] - mean_label) / std_label
+ normalized_test_dataset.append(e)
+ return normalized_train_dataset, normalized_val_dataset, normalized_test_dataset
+
+def robust_normalize_dataset(train_dataset, val_dataset, test_dataset):
+ """Robust scaling using statistics that are robust to outliers."""
+ scaler = RobustScaler()
+ train_labels = np.array([e["label"] for e in train_dataset]).reshape(-1, 1)
+ scaler.fit(train_labels)
+ normalized_train_dataset = []
+ normalized_val_dataset = []
+ normalized_test_dataset = []
+ for e in train_dataset:
+ e["label"] = scaler.transform([[e["label"]]])[0][0]
+ normalized_train_dataset.append(e)
+ for e in val_dataset:
+ e["label"] = scaler.transform([[e["label"]]])[0][0]
+ normalized_val_dataset.append(e)
+ for e in test_dataset:
+ e["label"] = scaler.transform([[e["label"]]])[0][0]
+ normalized_test_dataset.append(e)
+ return normalized_train_dataset, normalized_val_dataset, normalized_test_dataset
+
+def log_normalize_dataset(train_dataset, val_dataset, test_dataset, offset=1.0):
+ """Log normalization, useful for skewed data."""
+ normalized_train_dataset = []
+ normalized_val_dataset = []
+ normalized_test_dataset = []
+ for e in train_dataset:
+ e["label"] = np.log(e["label"] + offset)
+ normalized_train_dataset.append(e)
+ for e in val_dataset:
+ e["label"] = np.log(e["label"] + offset)
+ normalized_val_dataset.append(e)
+ for e in test_dataset:
+ e["label"] = np.log(e["label"] + offset)
+ normalized_test_dataset.append(e)
+ return normalized_train_dataset, normalized_val_dataset, normalized_test_dataset
+
+def quantile_normalize_dataset(train_dataset, val_dataset, test_dataset, n_quantiles=1000):
+ """Quantile normalization to achieve a uniform distribution."""
+ from sklearn.preprocessing import QuantileTransformer
+
+ transformer = QuantileTransformer(n_quantiles=n_quantiles, output_distribution='uniform')
+ train_labels = np.array([e["label"] for e in train_dataset]).reshape(-1, 1)
+ transformer.fit(train_labels)
+
+ normalized_train_dataset = []
+ normalized_val_dataset = []
+ normalized_test_dataset = []
+ for e in train_dataset:
+ e["label"] = transformer.transform([[e["label"]]])[0][0]
+ normalized_train_dataset.append(e)
+ for e in val_dataset:
+ e["label"] = transformer.transform([[e["label"]]])[0][0]
+ normalized_val_dataset.append(e)
+ for e in test_dataset:
+ e["label"] = transformer.transform([[e["label"]]])[0][0]
+ normalized_test_dataset.append(e)
+ return normalized_train_dataset, normalized_val_dataset, normalized_test_dataset
+
+def normalize_dataset(train_dataset, val_dataset, test_dataset, method='min_max', **kwargs):
+ """
+ Unified interface for different normalization methods.
+
+ Args:
+ train_dataset: Training dataset
+ val_dataset: Validation dataset
+ test_dataset: Test dataset
+ method: Normalization method ('min_max', 'standard', 'robust', 'log', 'quantile')
+ **kwargs: Additional arguments for specific normalization methods
+
+ Returns:
+ Normalized datasets (train, val, test)
+ """
+ normalization_methods = {
+ 'min_max': min_max_normalize_dataset,
+ 'standard': standard_normalize_dataset,
+ 'robust': robust_normalize_dataset,
+ 'log': log_normalize_dataset,
+ 'quantile': quantile_normalize_dataset
+ }
+
+ if method not in normalization_methods:
+ raise ValueError(f"Unsupported normalization method: {method}. "
+ f"Available methods: {list(normalization_methods.keys())}")
+
+ return normalization_methods[method](train_dataset, val_dataset, test_dataset, **kwargs)
+
diff --git a/src/data/prosst/readme.txt b/src/data/prosst/readme.txt
new file mode 100644
index 0000000000000000000000000000000000000000..54f8b0914198ca65cbd7664edc3272aad1237862
--- /dev/null
+++ b/src/data/prosst/readme.txt
@@ -0,0 +1 @@
+这个文件夹存放prosst的结构权重
\ No newline at end of file
diff --git a/src/data/prosst/structure/__init__.py b/src/data/prosst/structure/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391
diff --git a/src/data/prosst/structure/encoder/__init__.py b/src/data/prosst/structure/encoder/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..2a9678d4bdecd18efc6b4707caf9ce4d69e70232
--- /dev/null
+++ b/src/data/prosst/structure/encoder/__init__.py
@@ -0,0 +1 @@
+from .gvp import AutoGraphEncoder
\ No newline at end of file
diff --git a/src/data/prosst/structure/encoder/gvp.py b/src/data/prosst/structure/encoder/gvp.py
new file mode 100644
index 0000000000000000000000000000000000000000..c1165e2bffd4d1c77028ebcf26164009f6d56441
--- /dev/null
+++ b/src/data/prosst/structure/encoder/gvp.py
@@ -0,0 +1,163 @@
+import torch
+import torch.nn as nn
+from .layer import GVP, GVPConvLayer, LayerNorm
+from torch_scatter import scatter_mean
+
+class AttentionPooling(nn.Module):
+ def __init__(self, input_dim, attention_dim):
+ super(AttentionPooling, self).__init__()
+ self.attention_dim = attention_dim
+ self.query_layer = nn.Linear(input_dim, attention_dim, bias=True)
+ self.key_layer = nn.Linear(input_dim, attention_dim, bias=True)
+ self.value_layer = nn.Linear(input_dim, 1, bias=True) # value layer outputs one score
+ self.softmax = nn.Softmax(dim=1)
+
+ def forward(self, nodes_features1, nodes_features2):
+ # Assuming nodes_features1 and nodes_features2 are both of shape [node_num, 128]
+ nodes_features = nodes_features1 + nodes_features2 # This can also be concatenation or another operation
+
+ query = self.query_layer(nodes_features)
+ key = self.key_layer(nodes_features)
+ value = self.value_layer(nodes_features)
+
+ attention_scores = torch.matmul(query, key.transpose(-2, -1)) # [node_num, node_num]
+ attention_scores = self.softmax(attention_scores)
+
+ pooled_features = torch.matmul(attention_scores, value) # [node_num, 1]
+ return pooled_features
+
+class AutoGraphEncoder(nn.Module):
+ def __init__(self, node_in_dim, node_h_dim,
+ edge_in_dim, edge_h_dim, attention_dim=64,
+ num_layers=4, drop_rate=0.1) -> None:
+ super().__init__()
+ self.W_v = nn.Sequential(
+ LayerNorm(node_in_dim),
+ GVP(node_in_dim, node_h_dim, activations=(None, None))
+ )
+ self.W_e = nn.Sequential(
+ LayerNorm(edge_in_dim),
+ GVP(edge_in_dim, edge_h_dim, activations=(None, None))
+ )
+
+ self.layers = nn.ModuleList(
+ GVPConvLayer(node_h_dim, edge_h_dim, drop_rate=drop_rate)
+ for _ in range(num_layers))
+
+ ns, _ = node_h_dim
+ self.W_out = nn.Sequential(
+ LayerNorm(node_h_dim),
+ GVP(node_h_dim, (ns, 0)))
+
+ self.dense = nn.Sequential(
+ nn.Linear(ns, 2*ns),
+ nn.ReLU(inplace=True),
+ nn.Dropout(p=drop_rate),
+ nn.Linear(2*ns, node_in_dim[0]) # label num
+ )
+
+ self.loss_fn = nn.CrossEntropyLoss()
+
+ def forward(self, h_V, edge_index, h_E, node_s_labels):
+ h_V = self.W_v(h_V)
+ h_E = self.W_e(h_E)
+ for layer in self.layers:
+ h_V = layer(h_V, edge_index, h_E)
+ out = self.W_out(h_V)
+ logits = self.dense(out)
+ loss = self.loss_fn(logits, node_s_labels)
+
+ return loss, logits
+
+ def get_embedding(self, h_V, edge_index, h_E):
+ h_V = self.W_v(h_V)
+ h_E = self.W_e(h_E)
+ for layer in self.layers:
+ h_V = layer(h_V, edge_index, h_E)
+ out = self.W_out(h_V)
+ return out
+
+
+
+class SubgraphClassficationModel(nn.Module):
+ '''
+ :param node_in_dim: node dimensions in input graph, should be
+ (6, 3) if using original features
+ :param node_h_dim: node dimensions to use in GVP-GNN layers
+ :param edge_in_dim: edge dimensions in input graph, should be
+ (32, 1) if using original features
+ :param edge_h_dim: edge dimensions to embed to before use
+ in GVP-GNN layers
+ :param num_layers: number of GVP-GNN layers
+ :param drop_rate: rate to use in all dropout layers
+ '''
+ def __init__(self, node_in_dim, node_h_dim,
+ edge_in_dim, edge_h_dim, attention_dim=64,
+ num_layers=4, drop_rate=0.1):
+
+ super(SubgraphClassficationModel, self).__init__()
+ self.W_v = nn.Sequential(
+ LayerNorm(node_in_dim),
+ GVP(node_in_dim, node_h_dim, activations=(None, None))
+ )
+ self.W_e = nn.Sequential(
+ LayerNorm(edge_in_dim),
+ GVP(edge_in_dim, edge_h_dim, activations=(None, None))
+ )
+
+ self.layers = nn.ModuleList(
+ GVPConvLayer(node_h_dim, edge_h_dim, drop_rate=drop_rate)
+ for _ in range(num_layers))
+
+ ns, _ = node_h_dim
+ self.W_out = nn.Sequential(
+ LayerNorm(node_h_dim),
+ GVP(node_h_dim, (ns, 0)))
+
+ self.attention_classifier = AttentionPooling(ns, attention_dim)
+ self.dense = nn.Sequential(
+ nn.Linear(2*ns, 2*ns),
+ nn.ReLU(inplace=True),
+ nn.Dropout(p=drop_rate),
+ nn.Linear(2*ns, 1)
+ )
+
+ self.loss_fn = nn.BCEWithLogitsLoss()
+
+ def forward(self, h_V_parent, edge_index_parent, h_E_parent, batch_parent,
+ h_V_subgraph, edge_index_subgraph, h_E_subgraph, batch_subgraph,
+ labels):
+ '''
+ :param h_V: tuple (s, V) of node embeddings
+ :param edge_index: `torch.Tensor` of shape [2, num_edges]
+ :param h_E: tuple (s, V) of edge embeddings
+ '''
+ h_V_parent = self.W_v(h_V_parent)
+ h_E_parent = self.W_e(h_E_parent)
+ for layer in self.layers:
+ h_V_parent = layer(h_V_parent, edge_index_parent, h_E_parent)
+ out_parent = self.W_out(h_V_parent)
+ out_parent = scatter_mean(out_parent, batch_parent, dim=0)
+
+ h_V_subgraph = self.W_v(h_V_subgraph)
+ h_E_subgraph = self.W_e(h_E_subgraph)
+ for layer in self.layers:
+ h_V_subgraph = layer(h_V_subgraph, edge_index_subgraph, h_E_subgraph)
+ out_subgraph = self.W_out(h_V_subgraph)
+ out_subgraph = scatter_mean(out_subgraph, batch_subgraph, dim=0)
+
+ labels = labels.float()
+ out = torch.cat([out_parent, out_subgraph], dim=1)
+ logits = self.dense(out)
+ # logits = self.attention_classifier(out_parent, out_subgraph)
+ loss = self.loss_fn(logits.squeeze(-1), labels)
+ return loss, logits
+
+ def get_embedding(self, h_V, edge_index, h_E, batch):
+ h_V = self.W_v(h_V)
+ h_E = self.W_e(h_E)
+ for layer in self.layers:
+ h_V = layer(h_V, edge_index, h_E)
+ out = self.W_out(h_V)
+ out = scatter_mean(out, batch, dim=0)
+ return out
\ No newline at end of file
diff --git a/src/data/prosst/structure/encoder/layer.py b/src/data/prosst/structure/encoder/layer.py
new file mode 100644
index 0000000000000000000000000000000000000000..6b659e20b1e5f331a9e1ec0543453010439ddf83
--- /dev/null
+++ b/src/data/prosst/structure/encoder/layer.py
@@ -0,0 +1,374 @@
+import torch, functools
+from torch import nn
+import torch.nn.functional as F
+from torch_geometric.nn import MessagePassing
+from torch_scatter import scatter_add
+
+def tuple_sum(*args):
+ '''
+ Sums any number of tuples (s, V) elementwise.
+ '''
+ return tuple(map(sum, zip(*args)))
+
+def tuple_cat(*args, dim=-1):
+ '''
+ Concatenates any number of tuples (s, V) elementwise.
+
+ :param dim: dimension along which to concatenate when viewed
+ as the `dim` index for the scalar-channel tensors.
+ This means that `dim=-1` will be applied as
+ `dim=-2` for the vector-channel tensors.
+ '''
+ dim %= len(args[0][0].shape)
+ s_args, v_args = list(zip(*args))
+ return torch.cat(s_args, dim=dim), torch.cat(v_args, dim=dim)
+
+def tuple_index(x, idx):
+ '''
+ Indexes into a tuple (s, V) along the first dimension.
+
+ :param idx: any object which can be used to index into a `torch.Tensor`
+ '''
+ return x[0][idx], x[1][idx]
+
+def randn(n, dims, device="cpu"):
+ '''
+ Returns random tuples (s, V) drawn elementwise from a normal distribution.
+
+ :param n: number of data points
+ :param dims: tuple of dimensions (n_scalar, n_vector)
+
+ :return: (s, V) with s.shape = (n, n_scalar) and
+ V.shape = (n, n_vector, 3)
+ '''
+ return torch.randn(n, dims[0], device=device), \
+ torch.randn(n, dims[1], 3, device=device)
+
+def _norm_no_nan(x, axis=-1, keepdims=False, eps=1e-8, sqrt=True):
+ '''
+ L2 norm of tensor clamped above a minimum value `eps`.
+
+ :param sqrt: if `False`, returns the square of the L2 norm
+ '''
+ out = torch.clamp(torch.sum(torch.square(x), axis, keepdims), min=eps)
+ return torch.sqrt(out) if sqrt else out
+
+def _split(x, nv):
+ '''
+ Splits a merged representation of (s, V) back into a tuple.
+ Should be used only with `_merge(s, V)` and only if the tuple
+ representation cannot be used.
+
+ :param x: the `torch.Tensor` returned from `_merge`
+ :param nv: the number of vector channels in the input to `_merge`
+ '''
+ v = torch.reshape(x[..., -3*nv:], x.shape[:-1] + (nv, 3))
+ s = x[..., :-3*nv]
+ return s, v
+
+def _merge(s, v):
+ '''
+ Merges a tuple (s, V) into a single `torch.Tensor`, where the
+ vector channels are flattened and appended to the scalar channels.
+ Should be used only if the tuple representation cannot be used.
+ Use `_split(x, nv)` to reverse.
+ '''
+ v = torch.reshape(v, v.shape[:-2] + (3*v.shape[-2],))
+ return torch.cat([s, v], -1)
+
+class GVP(nn.Module):
+ '''
+ Geometric Vector Perceptron. See manuscript and README.md
+ for more details.
+
+ :param in_dims: tuple (n_scalar, n_vector)
+ :param out_dims: tuple (n_scalar, n_vector)
+ :param h_dim: intermediate number of vector channels, optional
+ :param activations: tuple of functions (scalar_act, vector_act)
+ :param vector_gate: whether to use vector gating.
+ (vector_act will be used as sigma^+ in vector gating if `True`)
+ '''
+ def __init__(self, in_dims, out_dims, h_dim=None,
+ activations=(F.relu, torch.sigmoid), vector_gate=False):
+ super(GVP, self).__init__()
+ self.si, self.vi = in_dims
+ self.so, self.vo = out_dims
+ self.vector_gate = vector_gate
+ if self.vi:
+ self.h_dim = h_dim or max(self.vi, self.vo)
+ self.wh = nn.Linear(self.vi, self.h_dim, bias=False)
+ self.ws = nn.Linear(self.h_dim + self.si, self.so)
+ if self.vo:
+ self.wv = nn.Linear(self.h_dim, self.vo, bias=False)
+ if self.vector_gate: self.wsv = nn.Linear(self.so, self.vo)
+ else:
+ self.ws = nn.Linear(self.si, self.so)
+
+ self.scalar_act, self.vector_act = activations
+ self.dummy_param = nn.Parameter(torch.empty(0))
+
+ def forward(self, x):
+ '''
+ :param x: tuple (s, V) of `torch.Tensor`,
+ or (if vectors_in is 0), a single `torch.Tensor`
+ :return: tuple (s, V) of `torch.Tensor`,
+ or (if vectors_out is 0), a single `torch.Tensor`
+ '''
+ if self.vi:
+ s, v = x
+ v = torch.transpose(v, -1, -2)
+ vh = self.wh(v)
+ vn = _norm_no_nan(vh, axis=-2)
+ s = self.ws(torch.cat([s, vn], -1))
+ if self.vo:
+ v = self.wv(vh)
+ v = torch.transpose(v, -1, -2)
+ if self.vector_gate:
+ if self.vector_act:
+ gate = self.wsv(self.vector_act(s))
+ else:
+ gate = self.wsv(s)
+ v = v * torch.sigmoid(gate).unsqueeze(-1)
+ elif self.vector_act:
+ v = v * self.vector_act(
+ _norm_no_nan(v, axis=-1, keepdims=True))
+ else:
+ s = self.ws(x)
+ if self.vo:
+ v = torch.zeros(s.shape[0], self.vo, 3,
+ device=self.dummy_param.device)
+ if self.scalar_act:
+ s = self.scalar_act(s)
+
+ return (s, v) if self.vo else s
+
+class _VDropout(nn.Module):
+ '''
+ Vector channel dropout where the elements of each
+ vector channel are dropped together.
+ '''
+ def __init__(self, drop_rate):
+ super(_VDropout, self).__init__()
+ self.drop_rate = drop_rate
+ self.dummy_param = nn.Parameter(torch.empty(0))
+
+ def forward(self, x):
+ '''
+ :param x: `torch.Tensor` corresponding to vector channels
+ '''
+ device = self.dummy_param.device
+ if not self.training:
+ return x
+ mask = torch.bernoulli(
+ (1 - self.drop_rate) * torch.ones(x.shape[:-1], device=device)
+ ).unsqueeze(-1)
+ x = mask * x / (1 - self.drop_rate)
+ return x
+
+class Dropout(nn.Module):
+ '''
+ Combined dropout for tuples (s, V).
+ Takes tuples (s, V) as input and as output.
+ '''
+ def __init__(self, drop_rate):
+ super(Dropout, self).__init__()
+ self.sdropout = nn.Dropout(drop_rate)
+ self.vdropout = _VDropout(drop_rate)
+
+ def forward(self, x):
+ '''
+ :param x: tuple (s, V) of `torch.Tensor`,
+ or single `torch.Tensor`
+ (will be assumed to be scalar channels)
+ '''
+ if type(x) is torch.Tensor:
+ return self.sdropout(x)
+ s, v = x
+ return self.sdropout(s), self.vdropout(v)
+
+class LayerNorm(nn.Module):
+ '''
+ Combined LayerNorm for tuples (s, V).
+ Takes tuples (s, V) as input and as output.
+ '''
+ def __init__(self, dims):
+ super(LayerNorm, self).__init__()
+ self.s, self.v = dims
+ self.scalar_norm = nn.LayerNorm(self.s)
+
+ def forward(self, x):
+ '''
+ :param x: tuple (s, V) of `torch.Tensor`,
+ or single `torch.Tensor`
+ (will be assumed to be scalar channels)
+ '''
+ if not self.v:
+ return self.scalar_norm(x)
+ s, v = x
+ vn = _norm_no_nan(v, axis=-1, keepdims=True, sqrt=False)
+ vn = torch.sqrt(torch.mean(vn, dim=-2, keepdim=True))
+ return self.scalar_norm(s), v / vn
+
+class GVPConv(MessagePassing):
+ '''
+ Graph convolution / message passing with Geometric Vector Perceptrons.
+ Takes in a graph with node and edge embeddings,
+ and returns new node embeddings.
+
+ This does NOT do residual updates and pointwise feedforward layers
+ ---see `GVPConvLayer`.
+
+ :param in_dims: input node embedding dimensions (n_scalar, n_vector)
+ :param out_dims: output node embedding dimensions (n_scalar, n_vector)
+ :param edge_dims: input edge embedding dimensions (n_scalar, n_vector)
+ :param n_layers: number of GVPs in the message function
+ :param module_list: preconstructed message function, overrides n_layers
+ :param aggr: should be "add" if some incoming edges are masked, as in
+ a masked autoregressive decoder architecture, otherwise "mean"
+ :param activations: tuple of functions (scalar_act, vector_act) to use in GVPs
+ :param vector_gate: whether to use vector gating.
+ (vector_act will be used as sigma^+ in vector gating if `True`)
+ '''
+ def __init__(self, in_dims, out_dims, edge_dims,
+ n_layers=3, module_list=None, aggr="mean",
+ activations=(F.relu, torch.sigmoid), vector_gate=False):
+ super(GVPConv, self).__init__(aggr=aggr)
+ self.si, self.vi = in_dims
+ self.so, self.vo = out_dims
+ self.se, self.ve = edge_dims
+
+ GVP_ = functools.partial(GVP,
+ activations=activations, vector_gate=vector_gate)
+
+ module_list = module_list or []
+ if not module_list:
+ if n_layers == 1:
+ module_list.append(
+ GVP_((2*self.si + self.se, 2*self.vi + self.ve),
+ (self.so, self.vo), activations=(None, None)))
+ else:
+ module_list.append(
+ GVP_((2*self.si + self.se, 2*self.vi + self.ve), out_dims)
+ )
+ for i in range(n_layers - 2):
+ module_list.append(GVP_(out_dims, out_dims))
+ module_list.append(GVP_(out_dims, out_dims,
+ activations=(None, None)))
+ self.message_func = nn.Sequential(*module_list)
+
+ def forward(self, x, edge_index, edge_attr):
+ '''
+ :param x: tuple (s, V) of `torch.Tensor`
+ :param edge_index: array of shape [2, n_edges]
+ :param edge_attr: tuple (s, V) of `torch.Tensor`
+ '''
+ x_s, x_v = x
+ message = self.propagate(edge_index,
+ s=x_s, v=x_v.reshape(x_v.shape[0], 3*x_v.shape[1]),
+ edge_attr=edge_attr)
+ return _split(message, self.vo)
+
+ def message(self, s_i, v_i, s_j, v_j, edge_attr):
+ v_j = v_j.view(v_j.shape[0], v_j.shape[1]//3, 3)
+ v_i = v_i.view(v_i.shape[0], v_i.shape[1]//3, 3)
+ message = tuple_cat((s_j, v_j), edge_attr, (s_i, v_i))
+ message = self.message_func(message)
+ return _merge(*message)
+
+
+class GVPConvLayer(nn.Module):
+ '''
+ Full graph convolution / message passing layer with
+ Geometric Vector Perceptrons. Residually updates node embeddings with
+ aggregated incoming messages, applies a pointwise feedforward
+ network to node embeddings, and returns updated node embeddings.
+
+ To only compute the aggregated messages, see `GVPConv`.
+
+ :param node_dims: node embedding dimensions (n_scalar, n_vector)
+ :param edge_dims: input edge embedding dimensions (n_scalar, n_vector)
+ :param n_message: number of GVPs to use in message function
+ :param n_feedforward: number of GVPs to use in feedforward function
+ :param drop_rate: drop probability in all dropout layers
+ :param autoregressive: if `True`, this `GVPConvLayer` will be used
+ with a different set of input node embeddings for messages
+ where src >= dst
+ :param activations: tuple of functions (scalar_act, vector_act) to use in GVPs
+ :param vector_gate: whether to use vector gating.
+ (vector_act will be used as sigma^+ in vector gating if `True`)
+ '''
+ def __init__(self, node_dims, edge_dims,
+ n_message=3, n_feedforward=2, drop_rate=.1,
+ autoregressive=False,
+ activations=(F.relu, torch.sigmoid), vector_gate=False):
+
+ super(GVPConvLayer, self).__init__()
+ self.conv = GVPConv(node_dims, node_dims, edge_dims, n_message,
+ aggr="add" if autoregressive else "mean",
+ activations=activations, vector_gate=vector_gate)
+ GVP_ = functools.partial(GVP,
+ activations=activations, vector_gate=vector_gate)
+ self.norm = nn.ModuleList([LayerNorm(node_dims) for _ in range(2)])
+ self.dropout = nn.ModuleList([Dropout(drop_rate) for _ in range(2)])
+
+ ff_func = []
+ if n_feedforward == 1:
+ ff_func.append(GVP_(node_dims, node_dims, activations=(None, None)))
+ else:
+ hid_dims = 4*node_dims[0], 2*node_dims[1]
+ ff_func.append(GVP_(node_dims, hid_dims))
+ for i in range(n_feedforward-2):
+ ff_func.append(GVP_(hid_dims, hid_dims))
+ ff_func.append(GVP_(hid_dims, node_dims, activations=(None, None)))
+ self.ff_func = nn.Sequential(*ff_func)
+
+ def forward(self, x, edge_index, edge_attr,
+ autoregressive_x=None, node_mask=None):
+ '''
+ :param x: tuple (s, V) of `torch.Tensor`
+ :param edge_index: array of shape [2, n_edges]
+ :param edge_attr: tuple (s, V) of `torch.Tensor`
+ :param autoregressive_x: tuple (s, V) of `torch.Tensor`.
+ If not `None`, will be used as src node embeddings
+ for forming messages where src >= dst. The corrent node
+ embeddings `x` will still be the base of the update and the
+ pointwise feedforward.
+ :param node_mask: array of type `bool` to index into the first
+ dim of node embeddings (s, V). If not `None`, only
+ these nodes will be updated.
+ '''
+
+ if autoregressive_x is not None:
+ src, dst = edge_index
+ mask = src < dst
+ edge_index_forward = edge_index[:, mask]
+ edge_index_backward = edge_index[:, ~mask]
+ edge_attr_forward = tuple_index(edge_attr, mask)
+ edge_attr_backward = tuple_index(edge_attr, ~mask)
+
+ dh = tuple_sum(
+ self.conv(x, edge_index_forward, edge_attr_forward),
+ self.conv(autoregressive_x, edge_index_backward, edge_attr_backward)
+ )
+
+ count = scatter_add(torch.ones_like(dst), dst,
+ dim_size=dh[0].size(0)).clamp(min=1).unsqueeze(-1)
+
+ dh = dh[0] / count, dh[1] / count.unsqueeze(-1)
+
+ else:
+ dh = self.conv(x, edge_index, edge_attr)
+
+ if node_mask is not None:
+ x_ = x
+ x, dh = tuple_index(x, node_mask), tuple_index(dh, node_mask)
+
+ x = self.norm[0](tuple_sum(x, self.dropout[0](dh)))
+
+ dh = self.ff_func(x)
+ x = self.norm[1](tuple_sum(x, self.dropout[1](dh)))
+
+ if node_mask is not None:
+ x_[0][node_mask], x_[1][node_mask] = x[0], x[1]
+ x = x_
+ return x
\ No newline at end of file
diff --git a/src/data/prosst/structure/quantizer.py b/src/data/prosst/structure/quantizer.py
new file mode 100644
index 0000000000000000000000000000000000000000..15e670cbf7f48bde782b327f876d5bfd5df38c47
--- /dev/null
+++ b/src/data/prosst/structure/quantizer.py
@@ -0,0 +1,625 @@
+import biotite
+import joblib
+import math
+import numpy as np
+import os
+import scipy.spatial as spa
+import torch
+import torch.nn.functional as F
+from Bio import PDB
+from Bio.SeqUtils import seq1
+from pathlib import Path
+from torch_geometric.data import Batch, Data
+from torch_scatter import scatter_mean, scatter_sum, scatter_max
+from tqdm import tqdm
+from typing import List
+from biotite.sequence import ProteinSequence
+from biotite.structure import filter_backbone, get_chains
+from biotite.structure.io import pdb, pdbx
+from biotite.structure.residues import get_residues
+from .encoder import AutoGraphEncoder
+
+
+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.0, D_max=20.0, 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 _orientations(X_ca):
+ forward = _normalize(X_ca[1:] - X_ca[:-1])
+ backward = _normalize(X_ca[:-1] - X_ca[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(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 _positional_embeddings(edge_index, num_embeddings=16, period_range=[2, 1000]):
+ # From https://github.com/jingraham/neurips19-graph-protein-design
+ 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 generate_graph(pdb_file, max_distance=10):
+ """
+ generate graph data from pdb file
+
+ params:
+ pdb_file: pdb file path
+ node_level: residue or secondary_structure
+ node_s_type: ss3, ss8, foldseek
+ max_distance: cut off
+ foldseek_fasta_file: foldseek fasta file path
+ foldseek_fasta_multi_chain: pdb multi chain for foldseek fasta
+
+ return:
+ graph data
+
+ """
+ pdb_parser = PDB.PDBParser(QUIET=True)
+ structure = pdb_parser.get_structure("protein", pdb_file)
+ model = structure[0]
+
+ # extract amino acid sequence
+ seq = []
+ # extract amino acid coordinates
+ aa_coords = {"N": [], "CA": [], "C": [], "O": []}
+
+ for model in structure:
+ for chain in model:
+ for residue in chain:
+ if residue.get_id()[0] == " ":
+ seq.append(residue.get_resname())
+ for atom_name in aa_coords.keys():
+ atom = residue[atom_name]
+ aa_coords[atom_name].append(atom.get_coord().tolist())
+ one_letter_seq = "".join([seq1(aa) for aa in seq])
+
+ # aa means amino acid
+ coords = list(zip(aa_coords["N"], aa_coords["CA"], aa_coords["C"], aa_coords["O"]))
+ coords = torch.tensor(coords)
+ # mask out the missing coordinates
+ mask = torch.isfinite(coords.sum(dim=(1, 2)))
+ coords[~mask] = np.inf
+ ca_coords = coords[:, 1]
+ node_s = torch.zeros(len(ca_coords), 20)
+ # build graph and max_distance
+ distances = spa.distance_matrix(ca_coords, ca_coords)
+ edge_index = torch.tensor(np.array(np.where(distances < max_distance)))
+ # remove loop
+ mask = edge_index[0] != edge_index[1]
+ edge_index = edge_index[:, mask]
+
+ # node features
+ orientations = _orientations(ca_coords)
+ sidechains = _sidechains(coords)
+ node_v = torch.cat([orientations, sidechains.unsqueeze(-2)], dim=-2)
+
+ # edge features
+ pos_embeddings = _positional_embeddings(edge_index)
+ E_vectors = ca_coords[edge_index[0]] - ca_coords[edge_index[1]]
+ rbf = _rbf(E_vectors.norm(dim=-1), D_count=16)
+ edge_s = torch.cat([rbf, pos_embeddings], dim=-1)
+ edge_v = _normalize(E_vectors).unsqueeze(-2)
+
+ # node_v: [node_num, 3, 3]
+ # edge_index: [2, edge_num]
+ # edge_s: [edge_num, 16+16]
+ # edge_v: [edge_num, 1, 3]
+ node_s, node_v, edge_s, edge_v = map(
+ torch.nan_to_num, (node_s, node_v, edge_s, edge_v)
+ )
+ data = Data(
+ node_s=node_s,
+ node_v=node_v,
+ edge_index=edge_index,
+ edge_s=edge_s,
+ edge_v=edge_v,
+ distances=distances,
+ aa_seq=one_letter_seq,
+ )
+
+ return data
+
+
+def get_atom_coords_residuewise(atoms: List[str], struct: biotite.structure.AtomArray):
+ """
+ Example for atoms argument: ["N", "CA", "C"]
+ """
+
+ def filterfn(s, axis=None):
+ filters = np.stack([s.atom_name == name for name in atoms], axis=1)
+ sum = filters.sum(0)
+ if not np.all(sum <= np.ones(filters.shape[1])):
+ raise RuntimeError("structure has multiple atoms with same name")
+ index = filters.argmax(0)
+ coords = s[index].coord
+ coords[sum == 0] = float("nan")
+ return coords
+
+ return biotite.structure.apply_residue_wise(struct, struct, filterfn)
+
+
+def extract_coords_from_structure(structure: biotite.structure.AtomArray):
+ """
+ Args:
+ structure: An instance of biotite AtomArray
+ Returns:
+ Tuple (coords, seq)
+ - coords is an L x 3 x 3 array for N, CA, C coordinates
+ - seq is the extracted sequence
+ """
+ coords = get_atom_coords_residuewise(["N", "CA", "C"], structure)
+ residue_identities = get_residues(structure)[1]
+ seq = "".join([ProteinSequence.convert_letter_3to1(r) for r in residue_identities])
+ return coords
+
+def extract_seq_from_pdb(pdb_file, chain=None):
+ """
+ Args:
+ structure: An instance of biotite AtomArray
+ Returns:
+ - seq is the extracted sequence
+ """
+ structure = load_structure(pdb_file, chain)
+ residue_identities = get_residues(structure)[1]
+ seq = "".join([ProteinSequence.convert_letter_3to1(r) for r in residue_identities])
+ return seq
+
+
+def generate_pos_subgraph(
+ graph_data,
+ subgraph_depth=None,
+ subgraph_interval=1,
+ max_distance=10,
+ anchor_nodes=None,
+ pure_subgraph=False,
+ device="cuda" if torch.cuda.is_available() else "cpu"
+):
+
+ # move graph_data to GPU
+ graph_data = Data(
+ node_s=graph_data.node_s.to(device) if torch.is_tensor(graph_data.node_s) else torch.tensor(graph_data.node_s, device=device),
+ node_v=graph_data.node_v.to(device) if torch.is_tensor(graph_data.node_v) else torch.tensor(graph_data.node_v, device=device),
+ edge_index=graph_data.edge_index.to(device) if torch.is_tensor(graph_data.edge_index) else torch.tensor(graph_data.edge_index, device=device),
+ edge_s=graph_data.edge_s.to(device) if torch.is_tensor(graph_data.edge_s) else torch.tensor(graph_data.edge_s, device=device),
+ edge_v=graph_data.edge_v.to(device) if torch.is_tensor(graph_data.edge_v) else torch.tensor(graph_data.edge_v, device=device),
+ distances=graph_data.distances.to(device) if torch.is_tensor(graph_data.distances) else torch.tensor(graph_data.distances, device=device),
+ aa_seq=graph_data.aa_seq
+ )
+
+ distances = graph_data.distances
+ if subgraph_depth is None:
+ subgraph_depth = 50
+
+ # Calculate anchor nodes if not provided
+ if anchor_nodes is None:
+ anchor_nodes = list(range(0, len(graph_data.aa_seq), subgraph_interval))
+ anchor_nodes_tensor = torch.tensor(anchor_nodes, device=device) # Move anchor nodes to device
+
+ # Get the k nearest neighbors for ALL anchor nodes (batched)
+ k = 50
+ nearest_indices = torch.argsort(distances, dim=1)[:, :k] # (num_nodes, k)
+ distance_mask = torch.gather(distances, 1, nearest_indices) < max_distance # (num_nodes, k)
+ nearest_indices = torch.where(distance_mask, nearest_indices, torch.tensor(-1, device=device)) # (num_nodes, k)
+
+ subgraph_dict = {}
+
+ for anchor_node in anchor_nodes: #Reverted back to for loop to ensure everything works with batches
+ try:
+
+ #Get neighbors for each anchornode
+ k_neighbors = nearest_indices[anchor_node]
+ k_neighbors = k_neighbors[k_neighbors != -1]
+
+ if len(k_neighbors) == 0: # Skip if no neighbors found
+ continue
+
+ if len(k_neighbors) > 30:
+ k_neighbors = k_neighbors[:40]
+
+ k_neighbors, _ = torch.sort(k_neighbors)
+
+ sub_matrix = distances.index_select(0, k_neighbors).index_select(1, k_neighbors)
+
+ # Create edge indices efficiently
+ sub_edges = torch.nonzero(sub_matrix < max_distance, as_tuple=False)
+ mask = sub_edges[:, 0] != sub_edges[:, 1]
+ sub_edge_index = sub_edges[mask]
+
+ if len(sub_edge_index) == 0: # Skip if no edges found
+ continue
+ # Move edge_index to GPU only when needed
+ edge_index_device = graph_data.edge_index.to(device)
+ original_edge_index = k_neighbors[sub_edge_index]
+
+ # More memory efficient edge matching
+ matches = []
+ for edge in original_edge_index:
+ match = (edge_index_device[0] == edge[0]) & (edge_index_device[1] == edge[1])
+ matches.append(match)
+ matches = torch.stack(matches)
+ edge_to_feature_idx = torch.nonzero(matches, as_tuple=True)[0].to(device)
+ if len(edge_to_feature_idx) == 0: # Skip if no matching edges
+ continue
+ #Create data
+ new_node_s = graph_data.node_s[k_neighbors].to(device)
+ new_node_v = graph_data.node_v[k_neighbors].to(device)
+ new_edge_s = graph_data.edge_s[edge_to_feature_idx].to(device)
+ new_edge_v = graph_data.edge_v[edge_to_feature_idx].to(device)
+
+ result = Data(
+ edge_index=sub_edge_index.T.to(device),
+ edge_s=new_edge_s.to(device),
+ edge_v=new_edge_v.to(device),
+ node_s=new_node_s.to(device),
+ node_v=new_node_v.to(device),
+ )
+ if not pure_subgraph:
+ result.index_map = {
+ int(old_id.to(device).item()): new_id
+ for new_id, old_id in enumerate(k_neighbors)
+ }
+ subgraph_dict[anchor_node] = result
+ except Exception as e:
+ print(f"Error processing anchor node {anchor_node}: {str(e)}")
+ continue
+
+ return subgraph_dict
+
+def load_structure(fpath, chain=None):
+ """
+ Args:
+ fpath: filepath to either pdb or cif file
+ chain: the chain id or list of chain ids to load
+ Returns:
+ biotite.structure.AtomArray
+ """
+ if fpath.endswith("cif"):
+ with open(fpath) as fin:
+ pdbxf = pdbx.PDBxFile.read(fin)
+ structure = pdbx.get_structure(pdbxf, model=1)
+ elif fpath.endswith("pdb"):
+ with open(fpath) as fin:
+ pdbf = pdb.PDBFile.read(fin)
+ structure = pdb.get_structure(pdbf, model=1)
+ bbmask = filter_backbone(structure)
+ structure = structure[bbmask]
+ all_chains = get_chains(structure)
+ if len(all_chains) == 0:
+ raise ValueError("No chains found in the input file.")
+ if chain is None:
+ chain_ids = all_chains
+ elif isinstance(chain, list):
+ chain_ids = chain
+ else:
+ chain_ids = [chain]
+ for chain in chain_ids:
+ if chain not in all_chains:
+ raise ValueError(f"Chain {chain} not found in input file")
+ chain_filter = [a.chain_id in chain_ids for a in structure]
+ structure = structure[chain_filter]
+ return structure
+
+
+def get_atom_coords_residuewise(atoms: List[str], struct: biotite.structure.AtomArray):
+ """
+ Example for atoms argument: ["N", "CA", "C"]
+ """
+
+ def filterfn(s, axis=None):
+ filters = np.stack([s.atom_name == name for name in atoms], axis=1)
+ sum = filters.sum(0)
+ if not np.all(sum <= np.ones(filters.shape[1])):
+ raise RuntimeError("structure has multiple atoms with same name")
+ index = filters.argmax(0)
+ coords = s[index].coord
+ coords[sum == 0] = float("nan")
+ return coords
+
+ return biotite.structure.apply_residue_wise(struct, struct, filterfn)
+
+
+def extract_coords_from_structure(structure: biotite.structure.AtomArray):
+ """
+ Args:
+ structure: An instance of biotite AtomArray
+ Returns:
+ Tuple (coords, seq)
+ - coords is an L x 3 x 3 array for N, CA, C coordinates
+ - seq is the extracted sequence
+ """
+ coords = get_atom_coords_residuewise(["N", "CA", "C"], structure)
+ residue_identities = get_residues(structure)[1]
+ seq = "".join([ProteinSequence.convert_letter_3to1(r) for r in residue_identities])
+ return coords
+
+
+def extract_seq_from_pdb(pdb_file, chain=None):
+ """
+ Args:
+ structure: An instance of biotite AtomArray
+ Returns:
+ - seq is the extracted sequence
+ """
+ structure = load_structure(pdb_file, chain)
+ residue_identities = get_residues(structure)[1]
+ seq = "".join([ProteinSequence.convert_letter_3to1(r) for r in residue_identities])
+ return seq
+
+
+def convert_graph(graph):
+ graph = Data(
+ node_s=graph.node_s.to(torch.float32),
+ node_v=graph.node_v.to(torch.float32),
+ edge_index=graph.edge_index.to(torch.int64),
+ edge_s=graph.edge_s.to(torch.float32),
+ edge_v=graph.edge_v.to(torch.float32),
+ )
+ return graph
+
+def predict_structure(model, cluster_models, dataloader, datalabels, device):
+ epoch_iterator = dataloader
+ struc_label_dict = {}
+ cluster_model_dict = {}
+
+ for cluster_model_path in cluster_models:
+ cluster_model_name = cluster_model_path.split("/")[-1].split(".")[0]
+ struc_label_dict[cluster_model_name] = {}
+ cluster_model_dict[cluster_model_name] = joblib.load(cluster_model_path)
+
+ with torch.no_grad():
+ for batch, label_dict in zip(epoch_iterator, datalabels):
+ batch.to(device)
+ h_V = (batch.node_s, batch.node_v)
+ h_E = (batch.edge_s, batch.edge_v)
+
+ node_emebddings = model.get_embedding(h_V, batch.edge_index, h_E)
+ graph_emebddings = scatter_mean(node_emebddings, batch.batch, dim=0).to(device)
+ norm_graph_emebddings = F.normalize(graph_emebddings, p=2, dim=1)
+ struc_label_dict[cluster_model_name][label_dict['name']]={}
+ for name, cluster_model in cluster_model_dict.items():
+ batch_structure_labels = cluster_model.predict(
+ norm_graph_emebddings.cpu()
+ ).tolist()
+ struc_label_dict[name][label_dict['name']]['seq']=label_dict['aa_seq']
+ struc_label_dict[name][label_dict['name']]['struct']=batch_structure_labels
+
+ return struc_label_dict
+
+
+def get_embeds(model, dataloader, device, pooling="mean"):
+ epoch_iterator = tqdm(dataloader)
+ embeds = []
+ with torch.no_grad():
+ for batch in epoch_iterator:
+ batch.to(device)
+ h_V = (batch.node_s, batch.node_v)
+ h_E = (batch.edge_s, batch.edge_v)
+ node_embeds = model.get_embedding(h_V, batch.edge_index, h_E).cpu()
+ if pooling == "mean":
+ graph_embeds = scatter_mean(node_embeds, batch.batch.cpu(), dim=0)
+ elif pooling == "sum":
+ graph_embeds = scatter_sum(node_embeds, batch.batch.cpu(), dim=0)
+ elif pooling == "max":
+ graph_embeds, _ = scatter_max(node_embeds, batch.batch.cpu(), dim=0)
+ else:
+ raise ValueError("pooling should be mean, sum or max")
+ embeds.append(graph_embeds)
+
+ embeds = torch.cat(embeds, dim=0)
+ norm_embeds = F.normalize(embeds, p=2, dim=1)
+ return norm_embeds
+
+def process_pdb_file(
+ pdb_file,
+ subgraph_depth,
+ subgraph_interval,
+ max_distance,
+ device="cuda" if torch.cuda.is_available() else "cpu"
+):
+ result_dict, subgraph_dict = {}, {}
+ result_dict["name"] = Path(pdb_file).name
+
+ try:
+ graph = generate_graph(pdb_file, max_distance)
+ except Exception as e:
+ print(f"Error in processing {pdb_file}")
+ result_dict["error"] = str(e)
+ return None, result_dict, 0
+
+ result_dict["aa_seq"] = graph.aa_seq
+ anchor_nodes = list(range(0, len(graph.node_s), subgraph_interval)) #Define anchor nodes
+ try: #Run subgraph generation
+ subgraph_dict = generate_pos_subgraph(
+ graph,
+ subgraph_depth,
+ subgraph_interval,
+ max_distance,
+ anchor_nodes=anchor_nodes,
+ pure_subgraph=True,
+ device=device
+ )
+ #Move all subgraphs to GPU
+ for key in subgraph_dict.keys():
+ subgraph_dict[key] = convert_graph(subgraph_dict[key])
+ except Exception as e:
+ print(f"Error processing subgraph {e}")
+ return None, result_dict, 0
+
+
+ subgraph_dict = dict(sorted(subgraph_dict.items(), key=lambda x: x[0]))
+ subgraphs = list(subgraph_dict.values())
+ return subgraphs, result_dict, len(anchor_nodes)
+
+
+def pdb_converter(
+ pdb_files,
+ subgraph_depth,
+ subgraph_interval,
+ max_distance,
+ device="cuda" if torch.cuda.is_available() else "cpu",
+ batch_size=32
+):
+ error_proteins, error_messages = [], []
+ dataset, results, node_counts = [], [], []
+
+ for i in tqdm(range(0, len(pdb_files), batch_size), desc="Processing PDB files"):
+ batch = pdb_files[i:i + batch_size]
+
+ for pdb_file in batch:
+ pdb_subgraphs, result_dict, node_count = process_pdb_file(
+ pdb_file,
+ subgraph_depth,
+ subgraph_interval,
+ max_distance,
+ device=device
+ )
+
+ if pdb_subgraphs is None:
+ error_proteins.append(result_dict["name"])
+ error_messages.append(result_dict["error"])
+ continue
+ dataset.append(pdb_subgraphs)
+ results.append(result_dict)
+ node_counts.append(node_count)
+
+ if error_proteins:
+ print(f"Found {len(error_proteins)} errors:")
+ for name, msg in zip(error_proteins, error_messages):
+ print(f"{name}: {msg}")
+
+ def collate_fn(batch):
+ batch_graphs = []
+ for d in batch:
+ batch_graphs.extend(d)
+ batch_graphs = Batch.from_data_list(batch_graphs)
+ batch_graphs.node_s = torch.zeros_like(batch_graphs.node_s)
+ return batch_graphs
+
+ def data_loader():
+ for item in dataset:
+ yield collate_fn([item])
+
+ return data_loader(), results
+
+
+class PdbQuantizer:
+
+ def __init__(
+ self,
+ structure_vocab_size=2048,
+ max_distance=10,
+ subgraph_depth=None,
+ subgraph_interval=1,
+ anchor_nodes=None,
+ model_path=None,
+ cluster_dir=None,
+ cluster_model=None,
+ device=None,
+ batch_size=16,
+ ) -> None:
+ assert structure_vocab_size in [20, 64, 128, 512, 1024, 2048, 4096]
+ self.batch_size = batch_size
+ self.max_distance = max_distance
+ self.subgraph_depth = subgraph_depth
+ self.subgraph_interval = subgraph_interval
+ self.anchor_nodes = anchor_nodes
+ if model_path is None:
+ self.model_path = str(Path(__file__).parent / "static" / "AE.pt")
+ else:
+ self.model_path = model_path
+ self.structure_vocab_size = structure_vocab_size
+
+ if cluster_dir is None:
+ self.cluster_dir = str(Path(__file__).parent / "static")
+ self.cluster_model = [
+ Path(self.cluster_dir) / f"{structure_vocab_size}.joblib",
+ ]
+ else:
+ self.cluster_dir = cluster_dir
+ self.cluster_model = cluster_model
+
+ if device is None:
+ self.device = "cuda" if torch.cuda.is_available() else "cpu"
+ else:
+ self.device = device
+
+ # Load model
+ node_dim = (256, 32)
+ edge_dim = (64, 2)
+ model = AutoGraphEncoder(
+ node_in_dim=(20, 3),
+ node_h_dim=node_dim,
+ edge_in_dim=(32, 1),
+ edge_h_dim=edge_dim,
+ num_layers=6,
+ )
+ model.load_state_dict(torch.load(self.model_path))
+ model = model.to(self.device)
+ model = model.eval()
+ self.model = model
+ self.cluster_models = [
+ os.path.join(self.cluster_dir, m) for m in self.cluster_model
+ ]
+
+ def __call__(self, pdb_files, return_residue_seq=False):
+ if isinstance(pdb_files, str):
+ pdb_files = [pdb_files]
+ elif isinstance(pdb_files, list):
+ pass
+ else:
+ raise ValueError("pdb_files should be either a string or a list of strings")
+ data_loader, results = pdb_converter(
+ pdb_files,
+ self.subgraph_depth,
+ self.subgraph_interval,
+ self.max_distance,
+ device=self.device,
+ batch_size=self.batch_size
+ )
+ structures = predict_structure(
+ self.model, self.cluster_models, data_loader, results, self.device
+ )
+ if not return_residue_seq:
+ for clusterModelLabels in structures.keys():
+ for structureDict in structures[clusterModelLabels].keys():
+ structures[clusterModelLabels][structureDict].pop('seq', None)
+ return structures
+
diff --git a/src/data/prosst/structure/static/1024.joblib b/src/data/prosst/structure/static/1024.joblib
new file mode 100644
index 0000000000000000000000000000000000000000..e4614963274b7e1114f4c5ae1de2341d0049e472
--- /dev/null
+++ b/src/data/prosst/structure/static/1024.joblib
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a761adc0b83c4031a5b31a89eeb6275c7dba2ccec9f9f6cbfe181a0ab95c82d5
+size 21040620
diff --git a/src/data/prosst/structure/static/128.joblib b/src/data/prosst/structure/static/128.joblib
new file mode 100644
index 0000000000000000000000000000000000000000..71432e5c592bfa3da1a7c2893d3d51421e2d565f
--- /dev/null
+++ b/src/data/prosst/structure/static/128.joblib
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:951840db5917fe112d4d4a2b8fbfd1357deacf5c67e1c27bba8d0d31c386cb60
+size 19205612
diff --git a/src/data/prosst/structure/static/20.joblib b/src/data/prosst/structure/static/20.joblib
new file mode 100644
index 0000000000000000000000000000000000000000..b38aa67f01ac3049623f110cec53474d4c9c629e
--- /dev/null
+++ b/src/data/prosst/structure/static/20.joblib
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f98a1d678c0a040e608d096806cc430d29f1ebb1e82f6405d6d6c4339491fc83
+size 18984427
diff --git a/src/data/prosst/structure/static/2048.joblib b/src/data/prosst/structure/static/2048.joblib
new file mode 100644
index 0000000000000000000000000000000000000000..0ded82db4291e55a5bafe84a8d0b995d29a90ef4
--- /dev/null
+++ b/src/data/prosst/structure/static/2048.joblib
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7e4a47885983b9f9adc1c3c8a9b549fd4b77a32685e1ffbad592ae6fccc5a20b
+size 23137772
diff --git a/src/data/prosst/structure/static/4096.joblib b/src/data/prosst/structure/static/4096.joblib
new file mode 100644
index 0000000000000000000000000000000000000000..322aa5913f2ba210e029d1fa782b198ebe6e2e76
--- /dev/null
+++ b/src/data/prosst/structure/static/4096.joblib
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:391ad839ef5019058486661547e34cea6e94234e3d52526659ee43bd1460d97c
+size 27332076
diff --git a/src/data/prosst/structure/static/512.joblib b/src/data/prosst/structure/static/512.joblib
new file mode 100644
index 0000000000000000000000000000000000000000..920e8a613d6a9ec64d0c6458c6fc577d8a2d5429
--- /dev/null
+++ b/src/data/prosst/structure/static/512.joblib
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:eacad4ba5aa1475a25fe156d81274cb93aaf737aa1a8fad2d2b08580ed2b719a
+size 19992044
diff --git a/src/data/prosst/structure/static/64.joblib b/src/data/prosst/structure/static/64.joblib
new file mode 100644
index 0000000000000000000000000000000000000000..e9415b02dc181070342928964bde89e0c1ae83fe
--- /dev/null
+++ b/src/data/prosst/structure/static/64.joblib
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:571518ae435f1ec3bbd79a89cbfb9d75303008b431a90224e212724d34317cc2
+size 19074540
diff --git a/src/data/prosst/structure/static/AE.pt b/src/data/prosst/structure/static/AE.pt
new file mode 100644
index 0000000000000000000000000000000000000000..d42313a8bc5a59ea802ef28bf438806d6d995126
--- /dev/null
+++ b/src/data/prosst/structure/static/AE.pt
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:57cf97ace00c8e0413e3c2a1407d299d86d99ed22982202c1bc652f7f762b3a3
+size 23669861
diff --git a/src/esmfold.py b/src/esmfold.py
new file mode 100644
index 0000000000000000000000000000000000000000..052fd30677a0735a47f3d3d68c3df2aef94fd662
--- /dev/null
+++ b/src/esmfold.py
@@ -0,0 +1,169 @@
+import torch
+import os
+import gc
+import argparse
+import biotite.structure.io as bsio
+import pandas as pd
+from tqdm import tqdm
+from Bio import SeqIO
+from transformers import AutoTokenizer, EsmForProteinFolding
+
+from transformers.models.esm.openfold_utils.protein import to_pdb, Protein as OFProtein
+from transformers.models.esm.openfold_utils.feats import atom14_to_atom37
+
+def read_fasta(file_path, key):
+ return str(getattr(SeqIO.read(file_path, 'fasta'), key))
+
+def read_multi_fasta(file_path):
+ """
+ params:
+ file_path: path to a fasta file
+ return:
+ a dictionary of sequences
+ """
+ sequences = {}
+ current_sequence = ''
+ with open(file_path, 'r') as file:
+ for line in file:
+ line = line.strip()
+ if line.startswith('>'):
+ if current_sequence:
+ sequences[header] = current_sequence
+ current_sequence = ''
+ header = line
+ else:
+ current_sequence += line
+ if current_sequence:
+ sequences[header] = current_sequence
+ return sequences
+
+def convert_outputs_to_pdb(outputs):
+ final_atom_positions = atom14_to_atom37(outputs["positions"][-1], outputs)
+ outputs = {k: v.to("cpu").numpy() for k, v in outputs.items()}
+ final_atom_positions = final_atom_positions.cpu().numpy()
+ final_atom_mask = outputs["atom37_atom_exists"]
+ pdbs = []
+ for i in range(outputs["aatype"].shape[0]):
+ aa = outputs["aatype"][i]
+ pred_pos = final_atom_positions[i]
+ mask = final_atom_mask[i]
+ resid = outputs["residue_index"][i] + 1
+ pred = OFProtein(
+ aatype=aa,
+ atom_positions=pred_pos,
+ atom_mask=mask,
+ residue_index=resid,
+ b_factors=outputs["plddt"][i],
+ chain_index=outputs["chain_index"][i] if "chain_index" in outputs else None,
+ )
+ pdbs.append(to_pdb(pred))
+ return pdbs
+
+if __name__ == '__main__':
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--sequence", type=str, default=None)
+ parser.add_argument("--fasta_file", type=str, default=None)
+ parser.add_argument("--fasta_chunk_num", type=int, default=None)
+ parser.add_argument("--fasta_chunk_id", type=int, default=None)
+ parser.add_argument("--fasta_dir", type=str, default=None)
+ parser.add_argument("--out_dir", type=str)
+ parser.add_argument("--out_file", type=str, default="result.pdb")
+ parser.add_argument("--out_info_file", type=str, default=None)
+ parser.add_argument("--fold_chunk_size", type=int)
+ args = parser.parse_args()
+
+ # model = esm.pretrained.esmfold_v1()
+ # model = model.eval().cuda()
+
+ tokenizer = AutoTokenizer.from_pretrained("facebook/esmfold_v1")
+ model = EsmForProteinFolding.from_pretrained("facebook/esmfold_v1", low_cpu_mem_usage=True)
+
+ model = model.cuda()
+ # model.esm = model.esm.half()
+ torch.backends.cuda.matmul.allow_tf32 = True
+ # Optionally, uncomment to set a chunk size for axial attention. This can help reduce memory.
+ # Lower sizes will have lower memory requirements at the cost of increased speed.
+ if args.fold_chunk_size is not None:
+ model.trunk.set_chunk_size(args.fold_chunk_size)
+
+ if args.fasta_file is not None:
+ unfold_proteins = []
+ seq_dict = read_multi_fasta(args.fasta_file)
+ os.makedirs(args.out_dir, exist_ok=True)
+ names, sequences = list(seq_dict.keys()), list(seq_dict.values())
+ if args.fasta_chunk_num is not None:
+ chunk_size = len(names) // args.fasta_chunk_num + 1
+ start = args.fasta_chunk_id * chunk_size
+ end = min((args.fasta_chunk_id + 1) * chunk_size, len(names))
+ names, sequences = names[start:end], sequences[start:end]
+
+ out_info_dict = {"name": [], "plddt": []}
+ bar = tqdm(zip(names, sequences))
+ for name, sequence in bar:
+ bar.set_description(name)
+ name = name[1:].split(" ")[0]
+ out_file = os.path.join(args.out_dir, f"{name}.ef.pdb")
+ if os.path.exists(out_file):
+ out_info_dict["name"].append(name)
+ struct = bsio.load_structure(out_file, extra_fields=["b_factor"])
+ out_info_dict["plddt"].append(struct.b_factor.mean())
+ continue
+
+ # Multimer prediction can be done with chains separated by ':'
+ try:
+ tokenized_input = tokenizer([sequence], return_tensors="pt", add_special_tokens=False)['input_ids'].cuda()
+ with torch.no_grad():
+ output = model(tokenized_input)
+ except Exception as e:
+ print(e)
+ print(f"Failed to predict {name}")
+ unfold_proteins.append(name)
+ continue
+ gc.collect()
+ pdb = convert_outputs_to_pdb(output)
+ with open(out_file, "w") as f:
+ f.write("\n".join(pdb))
+
+ out_info_dict["name"].append(name)
+ struct = bsio.load_structure(out_file, extra_fields=["b_factor"])
+ out_info_dict["plddt"].append(struct.b_factor.mean())
+
+ if args.out_info_file is not None:
+ pd.DataFrame(out_info_dict).to_csv(args.out_info_file, index=False)
+
+ if args.fasta_dir is not None:
+ os.makedirs(args.out_dir, exist_ok=True)
+ proteins = sorted(os.listdir(args.fasta_dir))
+ bar = tqdm(proteins)
+ for p in bar:
+ name = p[:-6]
+ bar.set_description(name)
+ out_file = os.path.join(args.out_dir, f"{name}.ef.pdb")
+ if os.path.exists(out_file):
+ continue
+ bar.set_description(p)
+ sequence = read_fasta(os.path.join(args.fasta_dir, p), "seq")
+ tokenized_input = tokenizer([sequence], return_tensors="pt", add_special_tokens=False)['input_ids'].cuda()
+ # Multimer prediction can be done with chains separated by ':'
+
+ with torch.no_grad():
+ output = model(tokenized_input)
+
+ pdb = convert_outputs_to_pdb(output)
+ with open(out_file, "w") as f:
+ f.write("\n".join(pdb))
+
+ struct = bsio.load_structure(out_file, extra_fields=["b_factor"])
+ print(p, struct.b_factor.mean())
+ elif args.sequence is not None:
+ sequence = args.sequence
+ # Multimer prediction can be done with chains separated by ':'
+
+ with torch.no_grad():
+ output = model.infer_pdb(sequence)
+
+ with open(args.out_file, "w") as f:
+ f.write(output)
+
+ struct = bsio.load_structure(args.out_file, extra_fields=["b_factor"])
+ print(struct.b_factor.mean())
\ No newline at end of file
diff --git a/src/eval.py b/src/eval.py
new file mode 100644
index 0000000000000000000000000000000000000000..569bdeaa7c32a0c6b94e0ce70c14c7f4c27b1673
--- /dev/null
+++ b/src/eval.py
@@ -0,0 +1,460 @@
+import sys
+import os
+
+os.environ["HF_ENDPOINT"]="https://hf-mirror.com"
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'src')))
+import argparse
+import torch
+import re
+import json
+import os
+import warnings
+import pandas as pd
+import torch.nn as nn
+from tqdm import tqdm
+from torchmetrics.classification import Accuracy, Recall, Precision, MatthewsCorrCoef, AUROC, F1Score, MatthewsCorrCoef
+from torchmetrics.classification import BinaryAccuracy, BinaryRecall, BinaryAUROC, BinaryF1Score, BinaryPrecision, BinaryMatthewsCorrCoef, BinaryF1Score
+from torchmetrics.regression import SpearmanCorrCoef
+from transformers import EsmTokenizer, EsmModel, BertModel, BertTokenizer
+from transformers import T5Tokenizer, T5EncoderModel, AutoTokenizer, AutoModelForMaskedLM, AutoModel
+from transformers import logging
+from datasets import load_dataset
+from torch.utils.data import DataLoader
+# from utils.data_utils import BatchSampler
+# from utils.metrics import MultilabelF1Max
+# from models.adapter_mdoel import AdapterModel
+from data.batch_sampler import BatchSampler
+from training.metrics import MultilabelF1Max
+from models.adapter_model import AdapterModel
+from models.lora_model import LoraModel
+from peft import PeftModel
+from typing import Dict, Any, Union, Tuple
+from data.dataloader import prepare_dataloaders
+from datetime import datetime
+
+# ignore warning information
+logging.set_verbosity_error()
+warnings.filterwarnings("ignore")
+
+def evaluate(model, plm_model, metrics, dataloader, loss_function, device=None):
+ total_loss = 0
+ total_samples = len(dataloader.dataset)
+ print(f"Total samples: {total_samples}")
+ epoch_iterator = tqdm(dataloader)
+ pred_labels = []
+
+ for i, batch in enumerate(epoch_iterator, 1):
+
+ for k, v in batch.items():
+ batch[k] = v.to(device)
+ label = batch["label"]
+
+ logits = model(plm_model, batch)
+ pred_labels.extend(logits.argmax(dim=1).cpu().numpy())
+
+ for metric_name, metric in metrics_dict.items():
+ if args.problem_type == 'regression' and args.num_labels == 1:
+ loss = loss_function(logits.squeeze(), label.squeeze())
+ metric(logits.squeeze(), label.squeeze())
+ elif args.problem_type == 'multi_label_classification':
+ loss = loss_function(logits, label.float())
+ metric(logits, label)
+ else:
+ loss = loss_function(logits, label)
+ metric(torch.argmax(logits, 1), label)
+
+ total_loss += loss.item() * len(label)
+ epoch_iterator.set_postfix(eval_loss=loss.item())
+
+ epoch_loss = total_loss / len(dataloader.dataset)
+ for k, v in metrics.items():
+ metrics[k] = [v.compute().item()]
+ print(f"{k}: {metrics[k][0]}")
+ metrics['loss'] = [epoch_loss]
+ return metrics, pred_labels
+
+if __name__ == '__main__':
+ parser = argparse.ArgumentParser()
+
+ # model params
+ parser.add_argument('--eval_method', type=str, default=None, help='evaluation method')
+ parser.add_argument('--hidden_size', type=int, default=None, help='embedding hidden size of the model')
+ parser.add_argument('--num_attention_head', type=int, default=8, help='number of attention heads')
+ parser.add_argument('--attention_probs_dropout', type=float, default=0, help='attention probs dropout prob')
+ parser.add_argument('--plm_model', type=str, default='facebook/esm2_t33_650M_UR50D', help='esm model name')
+ parser.add_argument('--num_labels', type=int, default=2, help='number of labels')
+ parser.add_argument('--pooling_method', type=str, default='mean', help='pooling method')
+ parser.add_argument('--pooling_dropout', type=float, default=0.25, help='pooling dropout')
+
+ # dataset
+ parser.add_argument('--dataset', type=str, default=None, help='dataset name')
+ parser.add_argument('--problem_type', type=str, default=None, help='problem type')
+ parser.add_argument('--test_file', type=str, default=None, help='test file')
+ parser.add_argument('--split', type=str, default=None, help='split name in Huggingface')
+ parser.add_argument('--test_result_dir', type=str, default=None, help='test result directory')
+ parser.add_argument('--metrics', type=str, default=None, help='computation metrics')
+ parser.add_argument('--num_workers', type=int, default=4, help='number of workers')
+ parser.add_argument('--max_seq_len', type=int, default=None, help='max sequence length')
+ parser.add_argument('--batch_size', type=int, default=None, help='batch size for fixed batch size')
+ parser.add_argument('--batch_token', type=int, default=10000, help='max number of token per batch')
+ parser.add_argument('--use_foldseek', action='store_true', help='use foldseek')
+ parser.add_argument('--use_ss8', action='store_true', help='use ss8')
+
+ # model path
+ parser.add_argument('--output_model_name', type=str, default=None, help='model name')
+ parser.add_argument('--output_root', default="result", help='root directory to save trained models')
+ parser.add_argument('--output_dir', default=None, help='directory to save trained models')
+ parser.add_argument('--model_path', default=None, help='model path directly')
+ parser.add_argument('--structure_seq', type=str, default="", help='structure sequence')
+ parser.add_argument('--training_method', type=str, default="freeze", help='training method')
+ args = parser.parse_args()
+
+ if 'foldseek_seq' in args.structure_seq:
+ args.use_foldseek = True
+ print("Enabled foldseek_seq based on structure_seq parameter")
+ if 'ss8_seq' in args.structure_seq:
+ args.use_ss8 = True
+ print("Enabled ss8_seq based on structure_seq parameter")
+
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+ os.makedirs(args.test_result_dir, exist_ok=True)
+ # build tokenizer and protein language model
+ if "esm" in args.plm_model:
+ tokenizer = EsmTokenizer.from_pretrained(args.plm_model)
+ plm_model = EsmModel.from_pretrained(args.plm_model)
+ args.hidden_size = plm_model.config.hidden_size
+ elif "bert" in args.plm_model:
+ tokenizer = BertTokenizer.from_pretrained(args.plm_model, do_lower_case=False)
+ plm_model = BertModel.from_pretrained(args.plm_model)
+ args.hidden_size = plm_model.config.hidden_size
+ elif "prot_t5" in args.plm_model:
+ tokenizer = T5Tokenizer.from_pretrained(args.plm_model, do_lower_case=False)
+ plm_model = T5EncoderModel.from_pretrained(args.plm_model)
+ args.hidden_size = plm_model.config.d_model
+ elif "ankh" in args.plm_model:
+ tokenizer = AutoTokenizer.from_pretrained(args.plm_model, do_lower_case=False)
+ plm_model = T5EncoderModel.from_pretrained(args.plm_model)
+ args.hidden_size = plm_model.config.d_model
+ elif "ProSST" in args.plm_model:
+ tokenizer = AutoTokenizer.from_pretrained(args.plm_model, do_lower_case=False)
+ plm_model = AutoModelForMaskedLM.from_pretrained(args.plm_model)
+ args.hidden_size = plm_model.config.hidden_size
+ elif "Prime" in args.plm_model:
+ tokenizer = AutoTokenizer.from_pretrained(args.plm_model, do_lower_case=False)
+ plm_model = AutoModelForMaskedLM.from_pretrained(args.plm_model)
+ args.hidden_size = plm_model.config.hidden_size
+ else:
+ tokenizer = AutoTokenizer.from_pretrained(args.plm_model)
+ plm_model = AutoModel.from_pretrained(args.plm_model).to(device).eval()
+ args.hidden_size = plm_model.config.hidden_size
+
+ args.vocab_size = plm_model.config.vocab_size
+
+ # Define metric configurations
+ metric_configs = {
+ 'accuracy': {
+ 'binary': BinaryAccuracy,
+ 'multi': lambda: Accuracy(task="multiclass", num_classes=args.num_labels)
+ },
+ 'recall': {
+ 'binary': BinaryRecall,
+ 'multi': lambda: Recall(task="multiclass", num_classes=args.num_labels)
+ },
+ 'precision': {
+ 'binary': BinaryPrecision,
+ 'multi': lambda: Precision(task="multiclass", num_classes=args.num_labels)
+ },
+ 'f1': {
+ 'binary': BinaryF1Score,
+ 'multi': lambda: F1Score(task="multiclass", num_classes=args.num_labels)
+ },
+ 'mcc': {
+ 'binary': BinaryMatthewsCorrCoef,
+ 'multi': lambda: MatthewsCorrCoef(task="multiclass", num_classes=args.num_labels)
+ },
+ 'auroc': {
+ 'binary': BinaryAUROC,
+ 'multi': lambda: AUROC(task="multiclass", num_classes=args.num_labels)
+ },
+ 'f1_max': {
+ 'any': lambda: MultilabelF1Max(num_labels=args.num_labels)
+ },
+ 'spearman_corr': {
+ 'any': SpearmanCorrCoef
+ }
+ }
+
+ # Initialize metrics dictionary
+ metrics_dict = {}
+ args.metrics = args.metrics.split(',')
+
+ # Create metrics based on configurations
+ for metric_name in args.metrics:
+ if metric_name not in metric_configs:
+ raise ValueError(f"Invalid metric: {metric_name}")
+
+ config = metric_configs[metric_name]
+ if 'any' in config:
+ metrics_dict[metric_name] = config['any']()
+ else:
+ metrics_dict[metric_name] = (config['binary']() if args.num_labels == 2
+ else config['multi']())
+
+ # Move metric to device
+ metrics_dict[metric_name].to(device)
+
+
+ # load adapter model
+ print("---------- Load Model ----------")
+ # model = AdapterModel(args)
+ # if args.model_path is not None:
+ # model_path = args.model_path
+ # else:
+ # model_path = f"{args.output_root}/{args.output_dir}/{args.output_model_name}"
+ if args.eval_method in ["full", "ses-adapter", "freeze"]:
+ model = AdapterModel(args)
+
+ elif args.eval_method in ['plm-lora', 'plm-qlora', 'plm-dora', 'plm-adalora', 'plm-ia3']:
+ model = LoraModel(args)
+
+ if args.model_path is not None:
+ model_path = args.model_path
+ else:
+ model_path = f"{args.output_root}/{args.output_dir}/{args.output_model_name}"
+ if args.eval_method == "full":
+ model_weights = torch.load(model_path)
+ model.load_state_dict(model_weights['model_state_dict'])
+ plm_model.load_state_dict(model_weights['plm_state_dict'])
+ else:
+ model.load_state_dict(torch.load(model_path))
+ model.to(device).eval()
+
+ if args.eval_method == 'plm-lora':
+ lora_path = model_path.replace(".pt", "_lora")
+ plm_model = PeftModel.from_pretrained(plm_model,lora_path)
+ plm_model = plm_model.merge_and_unload()
+ elif args.eval_method == 'plm-qlora':
+ lora_path = model_path.replace(".pt", "_qlora")
+ plm_model = PeftModel.from_pretrained(plm_model,lora_path)
+ plm_model = plm_model.merge_and_unload()
+ elif args.eval_method == "plm-dora":
+ dora_path = model_path.replace(".pt", "_dora")
+ plm_model = PeftModel.from_pretrained(plm_model, dora_path)
+ plm_model = plm_model.merge_and_unload()
+ elif args.eval_method == "plm-adalora":
+ adalora_path = model_path.replace(".pt", "_adalora")
+ plm_model = PeftModel.from_pretrained(plm_model, adalora_path)
+ plm_model = plm_model.merge_and_unload()
+ elif args.eval_method == "plm-ia3":
+ ia3_path = model_path.replace(".pt", "_ia3")
+ plm_model = PeftModel.from_pretrained(plm_model, ia3_path)
+ plm_model = plm_model.merge_and_unload()
+ plm_model.to(device).eval()
+
+ def param_num(model):
+ total = sum([param.numel() for param in model.parameters() if param.requires_grad])
+ num_M = total/1e6
+ if num_M >= 1000:
+ return "Number of parameter: %.2fB" % (num_M/1e3)
+ else:
+ return "Number of parameter: %.2fM" % (num_M)
+ print(param_num(model))
+
+ def collate_fn(examples):
+ aa_seqs, labels = [], []
+ if args.use_foldseek:
+ foldseek_seqs = []
+ if args.use_ss8:
+ ss8_seqs = []
+ prosst_stru_tokens = [] if "ProSST" in args.plm_model else None
+
+ for e in examples:
+ aa_seq = e["aa_seq"]
+ if args.use_foldseek:
+ foldseek_seq = e["foldseek_seq"]
+ if args.use_ss8:
+ ss8_seq = e["ss8_seq"]
+
+
+ if "ProSST" in args.plm_model and "prosst_stru_token" in e:
+ stru_token = e["prosst_stru_token"]
+ if isinstance(stru_token, str):
+ seq_clean = stru_token.strip("[]").replace(" ","")
+ tokens = list(map(int, seq_clean.split(','))) if seq_clean else []
+ elif isinstance(stru_token, (list, tuple)):
+ tokens = [int(x) for x in stru_token]
+ else:
+ tokens = []
+ prosst_stru_tokens.append(torch.tensor(tokens))
+
+ if 'prot_bert' in args.plm_model or "prot_t5" in args.plm_model:
+ aa_seq = " ".join(list(aa_seq))
+ aa_seq = re.sub(r"[UZOB]", "X", aa_seq)
+ if args.use_foldseek:
+ foldseek_seq = " ".join(list(foldseek_seq))
+ if args.use_ss8:
+ ss8_seq = " ".join(list(ss8_seq))
+ elif 'ankh' in args.plm_model:
+ aa_seq = list(aa_seq)
+ if args.use_foldseek:
+ foldseek_seq = list(foldseek_seq)
+ if args.use_ss8:
+ ss8_seq = list(ss8_seq)
+
+ aa_seqs.append(aa_seq)
+ if args.use_foldseek:
+ foldseek_seqs.append(foldseek_seq)
+ if args.use_ss8:
+ ss8_seqs.append(ss8_seq)
+ labels.append(e["label"])
+
+ if 'ankh' in args.plm_model:
+ aa_inputs = tokenizer.batch_encode_plus(aa_seqs, add_special_tokens=True, padding=True, is_split_into_words=True, return_tensors="pt")
+ if args.use_foldseek:
+ foldseek_input_ids = tokenizer.batch_encode_plus(foldseek_seqs, add_special_tokens=True, padding=True, is_split_into_words=True, return_tensors="pt")["input_ids"]
+ if args.use_ss8:
+ ss8_input_ids = tokenizer.batch_encode_plus(ss8_seqs, add_special_tokens=True, padding=True, is_split_into_words=True, return_tensors="pt")["input_ids"]
+ else:
+ aa_inputs = tokenizer(aa_seqs, return_tensors="pt", padding=True, truncation=True)
+ if args.use_foldseek:
+ foldseek_input_ids = tokenizer(foldseek_seqs, return_tensors="pt", padding=True, truncation=True)["input_ids"]
+ if args.use_ss8:
+ ss8_input_ids = tokenizer(ss8_seqs, return_tensors="pt", padding=True, truncation=True)["input_ids"]
+
+ aa_input_ids = aa_inputs["input_ids"]
+ attention_mask = aa_inputs["attention_mask"]
+
+ if args.problem_type == 'regression':
+ labels = torch.as_tensor(labels, dtype=torch.float)
+ else:
+ labels = torch.as_tensor(labels, dtype=torch.long)
+
+ data_dict = {
+ "aa_seq_input_ids": aa_input_ids,
+ "aa_seq_attention_mask": attention_mask,
+ "label": labels
+ }
+
+ if "ProSST" in args.plm_model and prosst_stru_tokens:
+ aa_max_length = len(aa_input_ids[0])
+ padded_tokens = []
+ for tokens in prosst_stru_tokens:
+ if tokens is None or len(tokens) == 0:
+
+ padded_tokens.append([0] * aa_max_length)
+ else:
+ struct_sequence = tokens.tolist()
+
+ padded_tokens.append(struct_sequence + [0] * (aa_max_length - len(struct_sequence)))
+
+ data_dict["aa_seq_stru_tokens"] = torch.tensor(padded_tokens, dtype=torch.long)
+
+ if args.use_foldseek:
+ data_dict["foldseek_seq_input_ids"] = foldseek_input_ids
+ if args.use_ss8:
+ data_dict["ss8_seq_input_ids"] = ss8_input_ids
+
+ return data_dict
+
+ loss_function = nn.CrossEntropyLoss()
+
+ def process_data_line(data):
+ if args.problem_type == 'multi_label_classification':
+ label_list = data['label'].split(',')
+ data['label'] = [int(l) for l in label_list]
+ binary_list = [0] * args.num_labels
+ for index in data['label']:
+ binary_list[index] = 1
+ data['label'] = binary_list
+ if args.max_seq_len is not None:
+ data["aa_seq"] = data["aa_seq"][:args.max_seq_len]
+ if args.use_foldseek:
+ data["foldseek_seq"] = data["foldseek_seq"][:args.max_seq_len]
+ if args.use_ss8:
+ data["ss8_seq"] = data["ss8_seq"][:args.max_seq_len]
+ # 如果是 ProSST 模型且有结构标记，也需要截断
+ if "ProSST" in args.plm_model and "prosst_stru_token" in data:
+ # 结构标记可能是字符串或列表形式
+ if isinstance(data["prosst_stru_token"], str):
+
+ pass
+ elif isinstance(data["prosst_stru_token"], (list, tuple)):
+ data["prosst_stru_token"] = data["prosst_stru_token"][:args.max_seq_len]
+ token_num = min(len(data["aa_seq"]), args.max_seq_len)
+ else:
+ token_num = len(data["aa_seq"])
+ return data, token_num
+
+ # process dataset from json file
+ def process_dataset_from_json(file):
+ dataset, token_nums = [], []
+ for l in open(file):
+ data = json.loads(l)
+ data, token_num = process_data_line(data)
+ dataset.append(data)
+ token_nums.append(token_num)
+ return dataset, token_nums
+
+
+ # process dataset from list
+ def process_dataset_from_list(data_list):
+ dataset, token_nums = [], []
+ for l in data_list:
+ data, token_num = process_data_line(l)
+ dataset.append(data)
+ token_nums.append(token_num)
+ return dataset, token_nums
+
+
+ if args.test_file.endswith('json'):
+ test_dataset, test_token_num = process_dataset_from_json(args.test_file)
+ elif args.test_file.endswith('csv'):
+ test_dataset, test_token_num = process_dataset_from_list(load_dataset("csv", data_files=args.test_file)['train'])
+ if args.test_result_dir:
+ test_result_df = pd.read_csv(args.test_file)
+ elif '/' in args.test_file: # Huggingface dataset (only csv now)
+ raw_dataset = load_dataset(args.test_file)
+ # Using the chosen split first.
+ if args.split and args.split in raw_dataset:
+ split = args.split
+ elif 'test' in raw_dataset:
+ split = 'test'
+ elif 'validation' in raw_dataset:
+ split = 'validation'
+ elif 'train' in raw_dataset:
+ split = 'train'
+ else:
+ split = list(raw_dataset.keys())[0]
+
+ test_dataset, test_token_num = process_dataset_from_list(raw_dataset[split])
+ if args.test_result_dir:
+ test_result_df = pd.DataFrame(raw_dataset[split])
+ else:
+ raise ValueError("Invalid file format")
+
+
+ if args.batch_size is None:
+ if args.batch_token is None:
+ raise ValueError("batch_size or batch_token must be specified")
+ test_loader = DataLoader(
+ test_dataset,
+ num_workers=args.num_workers,
+ collate_fn=collate_fn,
+ batch_sampler=BatchSampler(test_token_num, args.batch_token, False)
+ )
+ else:
+ test_loader = DataLoader(
+ test_dataset,
+ batch_size=args.batch_size,
+ num_workers=args.num_workers,
+ collate_fn=collate_fn,
+ shuffle=False
+ )
+
+ print("---------- Start Eval ----------")
+ with torch.no_grad():
+ metric, pred_labels = evaluate(model, plm_model, metrics_dict, test_loader, loss_function, device)
+ if args.test_result_dir:
+ pd.DataFrame(metric).to_csv(f"{args.test_result_dir}/evaluation_metrics.csv", index=False)
+ test_result_df["pred_label"] = pred_labels
+ test_result_df.to_csv(f"{args.test_result_dir}/evaluation_result.csv", index=False)
diff --git a/src/models/adapter_model.py b/src/models/adapter_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..dec32ef2671e7bd62f1b6a027983545e38f1fc0e
--- /dev/null
+++ b/src/models/adapter_model.py
@@ -0,0 +1,217 @@
+import torch
+import gc
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Tuple
+from .pooling import Attention1dPoolingHead, MeanPoolingHead, LightAttentionPoolingHead
+from .pooling import MeanPooling, MeanPoolingProjection
+
+def rotate_half(x):
+ x1, x2 = x.chunk(2, dim=-1)
+ return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(x, cos, sin):
+ cos = cos[:, :, : x.shape[-2], :]
+ sin = sin[:, :, : x.shape[-2], :]
+
+ return (x * cos) + (rotate_half(x) * sin)
+
+class RotaryEmbedding(nn.Module):
+ """
+ Rotary position embeddings based on those in
+ [RoFormer](https://huggingface.co/docs/transformers/model_doc/roformer). Query and keys are transformed by rotation
+ matrices which depend on their relative positions.
+ """
+
+ def __init__(self, dim: int):
+ super().__init__()
+ # Generate and save the inverse frequency buffer (non trainable)
+ inv_freq = 1.0 / (10000 ** (torch.arange(0, dim, 2, dtype=torch.int64).float() / dim))
+ inv_freq = inv_freq
+ self.register_buffer("inv_freq", inv_freq)
+
+ self._seq_len_cached = None
+ self._cos_cached = None
+ self._sin_cached = None
+
+ def _update_cos_sin_tables(self, x, seq_dimension=2):
+ seq_len = x.shape[seq_dimension]
+
+ # Reset the tables if the sequence length has changed,
+ # or if we're on a new device (possibly due to tracing for instance)
+ if seq_len != self._seq_len_cached or self._cos_cached.device != x.device:
+ self._seq_len_cached = seq_len
+ t = torch.arange(x.shape[seq_dimension], device=x.device).type_as(self.inv_freq)
+ freqs = torch.outer(t, self.inv_freq)
+ emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
+
+ self._cos_cached = emb.cos()[None, None, :, :]
+ self._sin_cached = emb.sin()[None, None, :, :]
+
+ return self._cos_cached, self._sin_cached
+
+ def forward(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+ self._cos_cached, self._sin_cached = self._update_cos_sin_tables(k, seq_dimension=-2)
+
+ return (
+ apply_rotary_pos_emb(q, self._cos_cached, self._sin_cached),
+ apply_rotary_pos_emb(k, self._cos_cached, self._sin_cached),
+ )
+
+
+class CrossModalAttention(nn.Module):
+ def __init__(self, args):
+ super().__init__()
+ self.attention_head_size = args.hidden_size // args.num_attention_head
+ assert (
+ self.attention_head_size * args.num_attention_head == args.hidden_size
+ ), "Embed size needs to be divisible by num heads"
+ self.num_attention_head = args.num_attention_head
+ self.hidden_size = args.hidden_size
+
+ self.query_proj = nn.Linear(args.hidden_size, args.hidden_size)
+ self.key_proj = nn.Linear(args.hidden_size, args.hidden_size)
+ self.value_proj = nn.Linear(args.hidden_size, args.hidden_size)
+
+ self.dropout = nn.Dropout(args.attention_probs_dropout)
+
+ self.out_proj = nn.Linear(args.hidden_size, args.hidden_size)
+ self.rotary_embeddings = RotaryEmbedding(dim=self.attention_head_size)
+
+ def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
+ new_x_shape = x.size()[:-1] + (self.num_attention_head, self.attention_head_size)
+ x = x.view(new_x_shape)
+ return x.permute(0, 2, 1, 3)
+
+ def forward(self, query, key, value, attention_mask=None, output_attentions=False):
+ key_layer = self.transpose_for_scores(self.key_proj(key))
+ value_layer = self.transpose_for_scores(self.value_proj(value))
+ query_layer = self.transpose_for_scores(self.query_proj(query))
+ query_layer = query_layer * self.attention_head_size**-0.5
+
+ query_layer, key_layer = self.rotary_embeddings(query_layer, key_layer)
+
+ # Take the dot product between "query" and "key" to get the raw attention scores.
+ attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+
+ if attention_mask is not None:
+ attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+ attention_scores = attention_scores.masked_fill(attention_mask == 0, float('-inf'))
+
+ attention_probs = F.softmax(attention_scores, dim=-1)
+ # This is actually dropping out entire tokens to attend to, which might
+ # seem a bit unusual, but is taken from the original Transformer paper.
+ attention_probs = self.dropout(attention_probs)
+ context_layer = torch.matmul(attention_probs, value_layer)
+
+ context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+ new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size,)
+ context_layer = context_layer.view(new_context_layer_shape)
+
+ outputs = (context_layer, attention_probs) if output_attentions else context_layer
+
+ return outputs
+
+class AdapterModel(nn.Module):
+ def __init__(self, args):
+ super().__init__()
+ self.args = args
+
+ if 'foldseek_seq' in args.structure_seq:
+ self.foldseek_embedding = nn.Embedding(args.vocab_size, args.hidden_size)
+ self.cross_attention_foldseek = CrossModalAttention(args)
+ if 'ss8_seq' in args.structure_seq:
+ self.ss_embedding = nn.Embedding(args.vocab_size, args.hidden_size)
+ self.cross_attention_ss = CrossModalAttention(args)
+ if 'esm3_structure_seq' in args.structure_seq:
+ self.esm3_structure_embedding = nn.Embedding(args.vocab_size, args.hidden_size)
+ self.cross_attention_esm3_structure = CrossModalAttention(args)
+
+ self.layer_norm = nn.LayerNorm(args.hidden_size)
+
+ if args.pooling_method == 'attention1d':
+ self.classifier = Attention1dPoolingHead(args.hidden_size, args.num_labels, args.pooling_dropout)
+ elif args.pooling_method == 'mean':
+ if "PPI" in args.dataset:
+ self.pooling = MeanPooling()
+ self.projection = MeanPoolingProjection(args.hidden_size, args.num_labels, args.pooling_dropout)
+ else:
+ self.classifier = MeanPoolingHead(args.hidden_size, args.num_labels, args.pooling_dropout)
+ elif args.pooling_method == 'light_attention':
+ self.classifier = LightAttentionPoolingHead(args.hidden_size, args.num_labels, args.pooling_dropout)
+ else:
+ raise ValueError(f"classifier method {args.pooling_method} not supported")
+
+ def plm_embedding(self, plm_model, aa_seq, attention_mask, structure_tokens=None):
+ with torch.no_grad():
+ if "ProSST" in self.args.plm_model:
+ outputs = plm_model(input_ids=aa_seq, attention_mask=attention_mask, ss_input_ids=structure_tokens, output_hidden_states=True)
+ elif "Prime" in self.args.plm_model or "deep" in self.args.plm_model:
+ outputs = plm_model(input_ids=aa_seq, attention_mask=attention_mask, output_hidden_states=True)
+ elif self.training and hasattr(self, 'args') and self.args.training_method == 'full':
+ outputs = plm_model(input_ids=aa_seq, attention_mask=attention_mask)
+ else:
+ outputs = plm_model(input_ids=aa_seq, attention_mask=attention_mask)
+ if "ProSST" in self.args.plm_model or "Prime" in self.args.plm_model:
+ seq_embeds = outputs.hidden_states[-1]
+ else:
+ seq_embeds = outputs.last_hidden_state
+ gc.collect()
+ torch.cuda.empty_cache()
+ return seq_embeds
+
+ def forward(self, plm_model, batch):
+ if "ProSST" in self.args.plm_model:
+ aa_seq, attention_mask, stru_tokens = batch['aa_seq_input_ids'], batch['aa_seq_attention_mask'], batch['aa_seq_stru_tokens']
+ seq_embeds = self.plm_embedding(plm_model, aa_seq, attention_mask, stru_tokens)
+ else:
+ aa_seq, attention_mask = batch['aa_seq_input_ids'], batch['aa_seq_attention_mask']
+ seq_embeds = self.plm_embedding(plm_model, aa_seq, attention_mask)
+
+ if 'foldseek_seq' in self.args.structure_seq:
+ foldseek_seq = batch['foldseek_seq_input_ids']
+ foldseek_embeds = self.foldseek_embedding(foldseek_seq)
+ foldseek_embeds = self.cross_attention_foldseek(foldseek_embeds, seq_embeds, seq_embeds, attention_mask)
+ embeds = seq_embeds + foldseek_embeds
+ embeds = self.layer_norm(embeds)
+
+ if 'ss8_seq' in self.args.structure_seq:
+ ss_seq = batch['ss8_seq_input_ids']
+ ss_embeds = self.ss_embedding(ss_seq)
+
+ if 'foldseek_seq' in self.args.structure_seq:
+ # cross attention with foldseek
+ ss_embeds = self.cross_attention_ss(ss_embeds, embeds, embeds, attention_mask)
+ embeds = ss_embeds + embeds
+ else:
+ # cross attention with sequence
+ ss_embeds = self.cross_attention_ss(ss_embeds, seq_embeds, seq_embeds, attention_mask)
+ embeds = ss_embeds + seq_embeds
+ embeds = self.layer_norm(embeds)
+
+ if 'esm3_structure_seq' in self.args.structure_seq:
+ esm3_structure_seq = batch['esm3_structure_seq_input_ids']
+ esm3_structure_embeds = self.esm3_structure_embedding(esm3_structure_seq)
+
+ if 'foldseek_seq' in self.args.structure_seq:
+ # cross attention with foldseek
+ esm3_structure_embeds = self.cross_attention_esm3_structure(esm3_structure_embeds, embeds, embeds, attention_mask)
+ embeds = esm3_structure_embeds + embeds
+ elif 'ss8_seq' in self.args.structure_seq:
+ # cross attention with ss8
+ esm3_structure_embeds = self.cross_attention_esm3_structure(esm3_structure_embeds, ss_embeds, ss_embeds, attention_mask)
+ embeds = esm3_structure_embeds + ss_embeds
+ else:
+ # cross attention with sequence
+ esm3_structure_embeds = self.cross_attention_esm3_structure(esm3_structure_embeds, seq_embeds, seq_embeds, attention_mask)
+ embeds = esm3_structure_embeds + seq_embeds
+ embeds = self.layer_norm(embeds)
+
+ if self.args.structure_seq:
+ logits = self.classifier(embeds, attention_mask)
+ else:
+ logits = self.classifier(seq_embeds, attention_mask)
+
+ return logits
+
\ No newline at end of file
diff --git a/src/models/lora_model.py b/src/models/lora_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..1997e5bcd1a0da0d6d19b952a6705ed52659ed9f
--- /dev/null
+++ b/src/models/lora_model.py
@@ -0,0 +1,80 @@
+"""
+use LoRA finetuning model
+"""
+import torch
+import gc
+import torch.nn as nn
+import torch.nn.functional as F
+from typing import Tuple
+from .pooling import Attention1dPoolingHead, MeanPoolingHead, LightAttentionPoolingHead
+from .pooling import MeanPooling, MeanPoolingProjection
+
+
+class LoraModel(nn.Module):
+ """
+ finetuning encoder
+ """
+
+ def __init__(self, args) -> None:
+ super().__init__()
+ self.args = args
+ if args.pooling_method == "attention1d":
+ self.classifier = Attention1dPoolingHead(
+ args.hidden_size, args.num_labels, args.pooling_dropout
+ )
+ elif args.pooling_method == "mean":
+ if "PPI" in args.dataset:
+ self.pooling = MeanPooling()
+ self.projection = MeanPoolingProjection(
+ args.hidden_size, args.num_labels, args.pooling_dropout
+ )
+ else:
+ self.classifier = MeanPoolingHead(
+ args.hidden_size, args.num_labels, args.pooling_dropout
+ )
+ elif args.pooling_method == "light_attention":
+ self.classifier = LightAttentionPoolingHead(
+ args.hidden_size, args.num_labels, args.pooling_dropout
+ )
+ else:
+ raise ValueError(f"classifier method {args.pooling_method} not supported")
+
+ def plm_embedding(self, plm_model, aa_seq, attention_mask, stru_token=None):
+ if (
+ self.training
+ and hasattr(self, "args")
+ and self.args.training_method in ['plm-lora', 'plm-qlora', 'plm-dora', 'plm-adalora', 'plm-ia3']
+ ):
+ if "ProSST" in self.args.plm_model:
+ outputs = plm_model(input_ids=aa_seq, attention_mask=attention_mask, ss_input_ids=stru_token, output_hidden_states=True)
+ elif "Prime" in self.args.plm_model:
+ outputs = plm_model(input_ids=aa_seq, attention_mask=attention_mask, output_hidden_states=True)
+ else:
+ outputs = plm_model(input_ids=aa_seq, attention_mask=attention_mask)
+ else:
+ with torch.no_grad():
+ if "ProSST" in self.args.plm_model:
+ outputs = plm_model(input_ids=aa_seq, attention_mask=attention_mask, ss_input_ids=stru_token, output_hidden_states=True)
+ else:
+ outputs = plm_model(input_ids=aa_seq, attention_mask=attention_mask)
+ seq_embeds = outputs.last_hidden_state
+ gc.collect()
+ torch.cuda.empty_cache()
+ return seq_embeds
+
+ def forward(self, plm_model, batch):
+ if "ProSST" in self.args.plm_model:
+ aa_seq, attention_mask, stru_token = (
+ batch["aa_seq_input_ids"],
+ batch["aa_seq_attention_mask"],
+ batch["aa_seq_stru_tokens"]
+ )
+ seq_embeds = self.plm_embedding(plm_model, aa_seq, attention_mask, stru_token)
+ else:
+ aa_seq, attention_mask = (
+ batch["aa_seq_input_ids"],
+ batch["aa_seq_attention_mask"],
+ )
+ seq_embeds = self.plm_embedding(plm_model, aa_seq, attention_mask)
+ logits = self.classifier(seq_embeds, attention_mask)
+ return logits
diff --git a/src/models/model_factory.py b/src/models/model_factory.py
new file mode 100644
index 0000000000000000000000000000000000000000..b4c700b3158879489e208d8e932368588b1e51f6
--- /dev/null
+++ b/src/models/model_factory.py
@@ -0,0 +1,322 @@
+import torch
+from transformers import (
+ EsmTokenizer, EsmModel,
+ BertTokenizer, BertModel,
+ T5Tokenizer, T5EncoderModel,
+ AutoTokenizer, PreTrainedModel,
+ AutoModelForMaskedLM, AutoModel
+)
+from peft import prepare_model_for_kbit_training
+from .adapter_model import AdapterModel
+from .lora_model import LoraModel
+
+def create_models(args):
+ """Create and initialize models and tokenizer."""
+ # Create tokenizer and PLM
+ tokenizer, plm_model = create_plm_and_tokenizer(args)
+
+ # Update hidden size based on PLM
+ args.hidden_size = get_hidden_size(plm_model, args.plm_model)
+
+ # Handle structure sequence vocabulary
+ if args.training_method == 'ses-adapter':
+ args.vocab_size = get_vocab_size(plm_model, args.structure_seq)
+
+ # Create adapter model
+ model = AdapterModel(args)
+
+ # Handle PLM parameters based on training method
+ if args.training_method != 'full':
+ freeze_plm_parameters(plm_model)
+ # if args.training_method == 'ses-adapter':
+ # plm_model=create_models(plm_model, args)
+ if args.training_method == 'plm-lora':
+ plm_model=setup_lora_plm(plm_model, args)
+ elif args.training_method == 'plm-qlora':
+ plm_model=create_qlora_model(plm_model, args)
+ elif args.training_method == 'plm-adalora':
+ plm_model=create_adalora_model(plm_model, args)
+ elif args.training_method == "plm-dora":
+ plm_model=create_dora_model(plm_model, args)
+ elif args.training_method == "plm-ia3":
+ plm_model=create_ia3_model(plm_model, args)
+
+ # Move models to device
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+ model = model.to(device)
+ plm_model = plm_model.to(device)
+
+ return model, plm_model, tokenizer
+
+def create_lora_model(args):
+ tokenizer, plm_model = create_plm_and_tokenizer(args)
+ # Update hidden size based on PLM
+ args.hidden_size = get_hidden_size(plm_model, args.plm_model)
+ model = LoraModel(args=args)
+ # Enable gradient checkpointing
+ plm_model.gradient_checkpointing_enable()
+ plm_model=setup_lora_plm(plm_model, args)
+ return model, plm_model, tokenizer
+
+def create_qlora_model(args):
+ qlora_config = setup_quantization_config()
+ tokenizer, plm_model = create_plm_and_tokenizer(args, qlora_config=qlora_config)
+ # Update hidden size based on PLM
+ args.hidden_size = get_hidden_size(plm_model, args.plm_model)
+ model = LoraModel(args=args)
+ # Enable gradient checkpointing
+ plm_model.gradient_checkpointing_enable()
+ plm_model = prepare_model_for_kbit_training(plm_model)
+ plm_model=setup_lora_plm(plm_model, args)
+ return model, plm_model, tokenizer
+
+def create_dora_model(args):
+ tokenizer, plm_model = create_plm_and_tokenizer(args)
+ # Update hidden size based on PLM
+ args.hidden_size = get_hidden_size(plm_model, args.plm_model)
+ model = LoraModel(args=args)
+ # Enable gradient checkpointing
+ plm_model.gradient_checkpointing_enable()
+ plm_model=setup_dora_plm(plm_model, args)
+ return model, plm_model, tokenizer
+
+def create_adalora_model(args):
+ tokenizer, plm_model = create_plm_and_tokenizer(args)
+ # Update hidden size based on PLM
+ args.hidden_size = get_hidden_size(plm_model, args.plm_model)
+ model = LoraModel(args=args)
+ # Enable gradient checkpointing
+ plm_model.gradient_checkpointing_enable()
+ plm_model=setup_adalora_plm(plm_model, args)
+ print(" Using plm adalora ")
+ return model, plm_model, tokenizer
+
+def create_ia3_model(args):
+ tokenizer, plm_model = create_plm_and_tokenizer(args)
+ args.hidden_size = get_hidden_size(plm_model, args.plm_model)
+ model = LoraModel(args=args)
+ plm_model.gradient_checkpointing_enable()
+ plm_model = prepare_model_for_kbit_training(plm_model)
+ plm_model=setup_ia3_plm(plm_model, args)
+ print(" Using plm IA3 ")
+ return model, plm_model, tokenizer
+
+def lora_factory(args):
+ if args.training_method in "plm-lora":
+ model, plm_model, tokenizer = create_lora_model(args)
+ elif args.training_method == "plm-qlora":
+ model, plm_model, tokenizer = create_qlora_model(args)
+ elif args.training_method == "plm-dora":
+ model, plm_model, tokenizer = create_dora_model(args)
+ elif args.training_method == "plm-adalora":
+ model, plm_model, tokenizer = create_adalora_model(args)
+ elif args.training_method == "plm-ia3":
+ model, plm_model, tokenizer = create_ia3_model(args)
+ else:
+ raise ValueError(f"Unsupported lora training method: {args.training_method}")
+ # Move models to device
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+ model = model.to(device)
+ plm_model = plm_model.to(device)
+ return model, plm_model, tokenizer
+
+def freeze_plm_parameters(plm_model):
+ """Freeze all parameters in the pre-trained language model."""
+ for param in plm_model.parameters():
+ param.requires_grad = False
+ plm_model.eval() # Set to evaluation mode
+
+def setup_quantization_config():
+ """Setup quantization configuration."""
+ from transformers import BitsAndBytesConfig
+ # https://huggingface.co/docs/peft/v0.14.0/en/developer_guides/quantization#quantize-a-model
+ qlora_config = BitsAndBytesConfig(
+ load_in_4bit=True,
+ bnb_4bit_use_double_quant=True,
+ bnb_4bit_quant_type="nf4",
+ bnb_4bit_compute_dtype=torch.bfloat16
+ )
+ return qlora_config
+
+def setup_lora_plm(plm_model, args):
+ """Setup LoRA for pre-trained language model."""
+ # Import LoRA configurations
+ from peft import get_peft_config, get_peft_model, LoraConfig, TaskType
+
+ if not isinstance(plm_model, PreTrainedModel):
+ raise TypeError("based_model must be a PreTrainedModel instance")
+
+ # validate lora_target_modules exist in model
+ available_modules = [name for name, _ in plm_model.named_modules()]
+ for module in args.lora_target_modules:
+ if not any(module in name for name in available_modules):
+ raise ValueError(f"Target module {module} not found in model")
+ # Configure LoRA
+ peft_config = LoraConfig(
+ task_type=TaskType.FEATURE_EXTRACTION,
+ inference_mode=False,
+ r=args.lora_r,
+ lora_alpha=args.lora_alpha,
+ lora_dropout=args.lora_dropout,
+ target_modules=args.lora_target_modules,
+ )
+ # Apply LoRA to model
+ plm_model = get_peft_model(plm_model, peft_config)
+ plm_model.print_trainable_parameters()
+ return plm_model
+
+def setup_dora_plm(plm_model, args):
+ """Setup DoRA for pre-trained language model."""
+ # Import DoRA configurations
+ from peft import get_peft_config, get_peft_model, LoraConfig, TaskType
+
+ if not isinstance(plm_model, PreTrainedModel):
+ raise TypeError("based_model must be a PreTrainedModel instance")
+
+ # validate Dora_target_modules exist in model
+ available_modules = [name for name, _ in plm_model.named_modules()]
+ for module in args.lora_target_modules:
+ if not any(module in name for name in available_modules):
+ raise ValueError(f"Target module {module} not found in model")
+ # Configure DoRA
+ peft_config = LoraConfig(
+ task_type=TaskType.FEATURE_EXTRACTION,
+ inference_mode=False,
+ r=args.lora_r,
+ lora_alpha=args.lora_alpha,
+ lora_dropout=args.lora_dropout,
+ target_modules=args.lora_target_modules,
+ use_dora=True
+ )
+ # Apply DoRA to model
+ plm_model = get_peft_model(plm_model, peft_config)
+ plm_model.print_trainable_parameters()
+ return plm_model
+
+def setup_adalora_plm(plm_model, args):
+ """Setup AdaLoRA for pre-trained language model."""
+ # Import AdaLoRA configurations
+ from peft import get_peft_config, get_peft_model, AdaLoraConfig, TaskType
+
+ if not isinstance(plm_model, PreTrainedModel):
+ raise TypeError("based_model must be a PreTrainedModel instance")
+
+ # validate lora_target_modules exist in model
+ available_modules = [name for name, _ in plm_model.named_modules()]
+ for module in args.lora_target_modules:
+ if not any(module in name for name in available_modules):
+ raise ValueError(f"Target module {module} not found in model")
+
+ # Configure AdaLoRA
+ peft_config = AdaLoraConfig(
+ task_type=TaskType.FEATURE_EXTRACTION,
+ peft_type="ADALORA",
+ init_r=12,
+ r=args.lora_r,
+ lora_alpha=args.lora_alpha,
+ lora_dropout=args.lora_dropout,
+ target_modules=args.lora_target_modules
+ )
+ # Apply AdaLoRA to model
+ plm_model = get_peft_model(plm_model, peft_config)
+ plm_model.print_trainable_parameters()
+ return plm_model
+
+def setup_ia3_plm(plm_model, args):
+ """Setup IA3 for pre-trained language model."""
+ # Import LoRA configurations
+ from peft import IA3Model, IA3Config, get_peft_model, TaskType
+
+ if not isinstance(plm_model, PreTrainedModel):
+ raise TypeError("based_model must be a PreTrainedModel instance")
+
+ # validate lora_target_modules exist in model
+ available_modules = [name for name, _ in plm_model.named_modules()]
+ print(available_modules)
+ for module in args.lora_target_modules:
+ if not any(module in name for name in available_modules):
+ raise ValueError(f"Target module {module} not found in model")
+ # Configure LoRA
+ peft_config = IA3Config(
+ task_type=TaskType.FEATURE_EXTRACTION,
+ peft_type="IA3",
+ target_modules=args.lora_target_modules,
+ feedforward_modules=args.feedforward_modules
+ )
+ # Apply LoRA to model
+ plm_model = get_peft_model(plm_model, peft_config)
+ plm_model.print_trainable_parameters()
+ return plm_model
+
+def create_plm_and_tokenizer(args, qlora_config=None):
+ """Create pre-trained language model and tokenizer based on model type."""
+ if "esm" in args.plm_model:
+ tokenizer = EsmTokenizer.from_pretrained(args.plm_model)
+ if qlora_config:
+ plm_model = EsmModel.from_pretrained(args.plm_model, quantization_config=qlora_config)
+ else:
+ plm_model = EsmModel.from_pretrained(args.plm_model)
+ elif "bert" in args.plm_model:
+ tokenizer = BertTokenizer.from_pretrained(args.plm_model, do_lower_case=False)
+ if qlora_config:
+ plm_model = BertModel.from_pretrained(args.plm_model, quantization_config=qlora_config)
+ else:
+ plm_model = BertModel.from_pretrained(args.plm_model)
+ elif "prot_t5" in args.plm_model:
+ tokenizer = T5Tokenizer.from_pretrained(args.plm_model, do_lower_case=False)
+ if qlora_config:
+ plm_model = T5EncoderModel.from_pretrained(args.plm_model, quantization_config=qlora_config)
+ else:
+ plm_model = T5EncoderModel.from_pretrained(args.plm_model)
+ elif "ankh" in args.plm_model:
+ tokenizer = AutoTokenizer.from_pretrained(args.plm_model, do_lower_case=False)
+ if qlora_config:
+ plm_model = T5EncoderModel.from_pretrained(args.plm_model, quantization_config=qlora_config)
+ else:
+ plm_model = T5EncoderModel.from_pretrained(args.plm_model)
+ elif "ProSST" in args.plm_model:
+ tokenizer = AutoTokenizer.from_pretrained(args.plm_model, do_lower_case=False)
+ if qlora_config:
+ plm_model = AutoModelForMaskedLM.from_pretrained(args.plm_model, trust_remote_code=True, quantization_config=qlora_config)
+ else:
+ plm_model = AutoModelForMaskedLM.from_pretrained(args.plm_model, trust_remote_code=True)
+ elif "Prime" in args.plm_model:
+ tokenizer = AutoTokenizer.from_pretrained(args.plm_model, trust_remote_code=True, do_lower_case=False)
+ if qlora_config:
+ plm_model = AutoModel.from_pretrained(args.plm_model, trust_remote_code=True, quantization_config=qlora_config)
+ else:
+ plm_model = AutoModel.from_pretrained(args.plm_model, trust_remote_code=True)
+ elif "deep" in args.plm_model:
+ tokenizer = AutoTokenizer.from_pretrained(args.plm_model, do_lower_case=False)
+ if qlora_config:
+ plm_model = AutoModel.from_pretrained(args.plm_model, trust_remote_code=True, quantization_config=qlora_config)
+ else:
+ plm_model = AutoModel.from_pretrained(args.plm_model, trust_remote_code=True)
+
+ else:
+ raise ValueError(f"Unsupported model type: {args.plm_model}")
+
+ return tokenizer, plm_model
+
+def get_hidden_size(plm_model, model_type):
+ """Get hidden size based on model type."""
+ if "esm" in model_type:
+ return plm_model.config.hidden_size
+ elif "bert" in model_type:
+ return plm_model.config.hidden_size
+ elif "prot_t5" in model_type or "ankh" in model_type:
+ return plm_model.config.d_model
+ elif "ProSST" in model_type:
+ return plm_model.config.hidden_size
+ elif "Prime" in model_type:
+ return plm_model.config.hidden_size
+ elif "deep" in model_type:
+ return plm_model.config.hidden_size
+ else:
+ raise ValueError(f"Unsupported model type: {model_type}")
+
+def get_vocab_size(plm_model, structure_seq):
+ """Get vocabulary size for structure sequences."""
+ if 'esm3_structure_seq' in structure_seq:
+ return max(plm_model.config.vocab_size, 4100)
+ return plm_model.config.vocab_size
diff --git a/src/models/pooling.py b/src/models/pooling.py
new file mode 100644
index 0000000000000000000000000000000000000000..b0feb33743c96773926e190a91de601571503a5a
--- /dev/null
+++ b/src/models/pooling.py
@@ -0,0 +1,198 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers.activations import ACT2FN
+
+class MaskedConv1d(nn.Conv1d):
+ """A masked 1-dimensional convolution layer.
+
+ Takes the same arguments as torch.nn.Conv1D, except that the padding is set automatically.
+
+ Shape:
+ Input: (N, L, in_channels)
+ input_mask: (N, L, 1), optional
+ Output: (N, L, out_channels)
+ """
+
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ kernel_size: int,
+ stride: int = 1,
+ dilation: int = 1,
+ groups: int = 1,
+ bias: bool = True,
+ ):
+ """
+ :param in_channels: input channels
+ :param out_channels: output channels
+ :param kernel_size: the kernel width
+ :param stride: filter shift
+ :param dilation: dilation factor
+ :param groups: perform depth-wise convolutions
+ :param bias: adds learnable bias to output
+ """
+ padding = dilation * (kernel_size - 1) // 2
+ super().__init__(
+ in_channels,
+ out_channels,
+ kernel_size,
+ stride=stride,
+ dilation=dilation,
+ groups=groups,
+ bias=bias,
+ padding=padding,
+ )
+
+ def forward(self, x, input_mask=None):
+ if input_mask is not None:
+ x = x * input_mask
+ return super().forward(x.transpose(1, 2)).transpose(1, 2)
+
+
+class Attention1dPooling(nn.Module):
+ def __init__(self, hidden_size):
+ super().__init__()
+ self.layer = MaskedConv1d(hidden_size, 1, 1)
+
+ def forward(self, x, input_mask=None):
+ batch_szie = x.shape[0]
+ attn = self.layer(x)
+ attn = attn.view(batch_szie, -1)
+ if input_mask is not None:
+ attn = attn.masked_fill_(
+ ~input_mask.view(batch_szie, -1).bool(), float("-inf")
+ )
+ attn = F.softmax(attn, dim=-1).view(batch_szie, -1, 1)
+ out = (attn * x).sum(dim=1)
+ return out
+
+class Attention1dPoolingProjection(nn.Module):
+ def __init__(self, hidden_size, num_labels, dropout=0.25) -> None:
+ super(Attention1dPoolingProjection, self).__init__()
+ self.linear = nn.Linear(hidden_size, hidden_size)
+ self.dropout = nn.Dropout(dropout)
+ self.relu = nn.ReLU()
+ self.final = nn.Linear(hidden_size, num_labels)
+
+ def forward(self, x):
+ x = self.linear(x)
+ x = self.dropout(x)
+ x = self.relu(x)
+ x = self.final(x)
+ return x
+
+class Attention1dPoolingHead(nn.Module):
+ """Outputs of the model with the attention1d"""
+
+ def __init__(
+ self, hidden_size: int, num_labels: int, dropout: float = 0.25
+ ): # [batch x sequence(751) x embedding (1280)] --> [batch x embedding] --> [batch x 1]
+ super(Attention1dPoolingHead, self).__init__()
+ self.attention1d = Attention1dPooling(hidden_size)
+ self.attention1d_projection = Attention1dPoolingProjection(hidden_size, num_labels, dropout)
+
+ def forward(self, x, input_mask=None):
+ x = self.attention1d(x, input_mask=input_mask.unsqueeze(-1))
+ x = self.attention1d_projection(x)
+ return x
+
+class MeanPooling(nn.Module):
+ """Mean Pooling for sentence-level classification tasks."""
+
+ def __init__(self):
+ super().__init__()
+
+ def forward(self, features, input_mask=None):
+ if input_mask is not None:
+ # Applying input_mask to zero out masked values
+ masked_features = features * input_mask.unsqueeze(2)
+ sum_features = torch.sum(masked_features, dim=1)
+ mean_pooled_features = sum_features / input_mask.sum(dim=1, keepdim=True)
+ else:
+ mean_pooled_features = torch.mean(features, dim=1)
+ return mean_pooled_features
+
+
+class MeanPoolingProjection(nn.Module):
+ """Mean Pooling with a projection layer for sentence-level classification tasks."""
+
+ def __init__(self, hidden_size, num_labels, dropout=0.25):
+ super().__init__()
+ self.dense = nn.Linear(hidden_size, hidden_size)
+ self.dropout = nn.Dropout(dropout)
+ self.out_proj = nn.Linear(hidden_size, num_labels)
+
+ def forward(self, mean_pooled_features):
+ x = self.dropout(mean_pooled_features)
+ x = self.dense(x)
+ x = ACT2FN['gelu'](x)
+ x = self.dropout(x)
+ x = self.out_proj(x)
+ return x
+
+
+class MeanPoolingHead(nn.Module):
+ """Mean Pooling Head for sentence-level classification tasks."""
+
+ def __init__(self, hidden_size, num_labels, dropout=0.25):
+ super().__init__()
+ self.mean_pooling = MeanPooling()
+ self.mean_pooling_projection = MeanPoolingProjection(hidden_size, num_labels, dropout)
+
+ def forward(self, features, input_mask=None):
+ mean_pooling_features = self.mean_pooling(features, input_mask=input_mask)
+ x = self.mean_pooling_projection(mean_pooling_features)
+ return x
+
+
+class LightAttentionPoolingHead(nn.Module):
+ def __init__(self, hidden_size=1280, num_labels=11, dropout=0.25, kernel_size=9, conv_dropout: float = 0.25):
+ super(LightAttentionPoolingHead, self).__init__()
+
+ self.feature_convolution = nn.Conv1d(hidden_size, hidden_size, kernel_size, stride=1,
+ padding=kernel_size // 2)
+ self.attention_convolution = nn.Conv1d(hidden_size, hidden_size, kernel_size, stride=1,
+ padding=kernel_size // 2)
+
+ self.softmax = nn.Softmax(dim=-1)
+
+ self.dropout = nn.Dropout(conv_dropout)
+
+ self.linear = nn.Sequential(
+ nn.Linear(2 * hidden_size, 32),
+ nn.Dropout(dropout),
+ nn.ReLU(),
+ nn.BatchNorm1d(32)
+ )
+
+ self.output = nn.Linear(32, num_labels)
+
+ def forward(self, x: torch.Tensor, mask, **kwargs) -> torch.Tensor:
+ """
+ Args:
+ x: [batch_size, sequence_length, hidden_size] embedding tensor that should be classified
+ mask: [batch_size, sequence_length] mask corresponding to the zero padding used for the shorter sequecnes in the batch. All values corresponding to padding are False and the rest is True.
+
+ Returns:
+ classification: [batch_size,num_labels] tensor with logits
+ """
+ x = x.permute(0, 2, 1) # [batch_size, hidden_size, sequence_length]
+ o = self.feature_convolution(x) # [batch_size, hidden_size, sequence_length]
+ o = self.dropout(o) # [batch_gsize, hidden_size, sequence_length]
+ attention = self.attention_convolution(x) # [batch_size, hidden_size, sequence_length]
+
+ # mask out the padding to which we do not want to pay any attention (we have the padding because the sequences have different lenghts).
+ # This padding is added by the dataloader when using the padded_permuted_collate function in utils/general.py
+ attention = attention.masked_fill(mask[:, None, :] == False, -1e9)
+
+ # code used for extracting embeddings for UMAP visualizations
+ # extraction = torch.sum(x * self.softmax(attention), dim=-1)
+ # extraction = self.id0(extraction)
+
+ o1 = torch.sum(o * self.softmax(attention), dim=-1) # [batchsize, hidden_size]
+ o2, _ = torch.max(o, dim=-1) # [batchsize, hidden_size]
+ o = torch.cat([o1, o2], dim=-1) # [batchsize, 2*hidden_size]
+ o = self.linear(o) # [batchsize, 32]
+ return self.output(o) # [batchsize, num_labels]
\ No newline at end of file
diff --git a/src/predict.py b/src/predict.py
new file mode 100644
index 0000000000000000000000000000000000000000..5862709892bc88439600bb25322cd0559108d71b
--- /dev/null
+++ b/src/predict.py
@@ -0,0 +1,376 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+import sys
+import os
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'src')))
+import argparse
+import torch
+import re
+import json
+import os
+import warnings
+import numpy as np
+from pathlib import Path
+from transformers import EsmTokenizer, EsmModel, BertModel, BertTokenizer, AutoModelForMaskedLM
+from transformers import T5Tokenizer, T5EncoderModel, AutoTokenizer, AutoModel
+from transformers import logging
+from peft import PeftModel
+# Import project modules
+from models.adapter_model import AdapterModel
+from models.lora_model import LoraModel
+from models.pooling import MeanPooling, Attention1dPoolingHead, LightAttentionPoolingHead
+
+# Ignore warning information
+logging.set_verbosity_error()
+warnings.filterwarnings("ignore")
+
+def parse_args():
+ parser = argparse.ArgumentParser(description="Predict protein function for a single sequence")
+
+ # Model parameters
+ parser.add_argument('--eval_method', type=str, default="freeze", choices=["freeze", "plm-lora", "plm-qlora", "ses-adapter"], help="Evaluation method")
+ parser.add_argument('--model_path', type=str, required=True, help="Path to the trained model")
+ parser.add_argument('--plm_model', type=str, required=True, help="Pretrained language model name or path")
+ parser.add_argument('--pooling_method', type=str, default="mean", choices=["mean", "attention1d", "light_attention"], help="Pooling method")
+ parser.add_argument('--problem_type', type=str, default="single_label_classification",
+ choices=["single_label_classification", "multi_label_classification", "regression"],
+ help="Problem type")
+ parser.add_argument('--num_labels', type=int, default=2, help="Number of labels")
+ parser.add_argument('--hidden_size', type=int, default=None, help="Embedding hidden size of the model")
+ parser.add_argument('--num_attention_head', type=int, default=8, help="Number of attention heads")
+ parser.add_argument('--attention_probs_dropout', type=float, default=0, help="Attention probs dropout prob")
+ parser.add_argument('--pooling_dropout', type=float, default=0.25, help="Pooling dropout")
+
+ # Input sequence parameters
+ parser.add_argument('--aa_seq', type=str, required=True, help="Amino acid sequence")
+ parser.add_argument('--foldseek_seq', type=str, default="", help="Foldseek sequence (optional)")
+ parser.add_argument('--ss8_seq', type=str, default="", help="Secondary structure sequence (optional)")
+ parser.add_argument('--dataset', type=str, default="single", help="Dataset name (optional)")
+ parser.add_argument('--use_foldseek', action='store_true', help="Use foldseek sequence")
+ parser.add_argument('--use_ss8', action='store_true', help="Use secondary structure sequence")
+ parser.add_argument('--structure_seq', type=str, default=None, help="Structure sequence types to use (comma-separated)")
+
+ # Other parameters
+ parser.add_argument('--max_seq_len', type=int, default=1024, help="Maximum sequence length")
+
+ args = parser.parse_args()
+
+ # Automatically determine whether to use structure sequences based on input
+ args.use_foldseek = bool(args.foldseek_seq)
+ args.use_ss8 = bool(args.ss8_seq)
+
+ return args
+
+def load_model_and_tokenizer(args):
+ print("---------- Loading Model and Tokenizer ----------")
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+
+ # Check if model file exists
+ if not os.path.exists(args.model_path):
+ raise FileNotFoundError(f"Model file not found: {args.model_path}")
+
+ # Load model configuration if available
+ config_path = os.path.join(os.path.dirname(args.model_path), "config.json")
+ try:
+ with open(config_path, "r") as f:
+ config = json.load(f)
+ print(f"Loaded configuration from {config_path}")
+
+ # Update args with config values if they exist
+ if "pooling_method" in config:
+ args.pooling_method = config["pooling_method"]
+ if "problem_type" in config:
+ args.problem_type = config["problem_type"]
+ if "num_labels" in config:
+ args.num_labels = config["num_labels"]
+ if "num_attention_head" in config:
+ args.num_attention_head = config["num_attention_head"]
+ if "attention_probs_dropout" in config:
+ args.attention_probs_dropout = config["attention_probs_dropout"]
+ if "pooling_dropout" in config:
+ args.pooling_dropout = config["pooling_dropout"]
+ except FileNotFoundError:
+ print(f"Model config not found at {config_path}. Using command line arguments.")
+
+ # Build tokenizer and protein language model
+ if "esm" in args.plm_model:
+ tokenizer = EsmTokenizer.from_pretrained(args.plm_model)
+ plm_model = EsmModel.from_pretrained(args.plm_model)
+ args.hidden_size = plm_model.config.hidden_size
+ elif "bert" in args.plm_model:
+ tokenizer = BertTokenizer.from_pretrained(args.plm_model, do_lower_case=False)
+ plm_model = BertModel.from_pretrained(args.plm_model)
+ args.hidden_size = plm_model.config.hidden_size
+ elif "prot_t5" in args.plm_model:
+ tokenizer = T5Tokenizer.from_pretrained(args.plm_model, do_lower_case=False)
+ plm_model = T5EncoderModel.from_pretrained(args.plm_model)
+ args.hidden_size = plm_model.config.d_model
+ elif "ankh" in args.plm_model:
+ tokenizer = AutoTokenizer.from_pretrained(args.plm_model, do_lower_case=False)
+ plm_model = T5EncoderModel.from_pretrained(args.plm_model)
+ args.hidden_size = plm_model.config.d_model
+ elif "ProSST" in args.plm_model:
+ tokenizer = AutoTokenizer.from_pretrained(args.plm_model, do_lower_case=False)
+ plm_model = AutoModelForMaskedLM.from_pretrained(args.plm_model).to(device).eval()
+ args.hidden_size = plm_model.config.hidden_size
+ elif "Prime" in args.plm_model:
+ tokenizer = AutoTokenizer.from_pretrained(args.plm_model, do_lower_case=False)
+ plm_model = AutoModelForMaskedLM.from_pretrained(args.plm_model).to(device).eval()
+ args.hidden_size = plm_model.config.hidden_size
+ else:
+ tokenizer = AutoTokenizer.from_pretrained(args.plm_model)
+ plm_model = AutoModel.from_pretrained(args.plm_model)
+ args.hidden_size = plm_model.config.hidden_size
+
+ args.vocab_size = plm_model.config.vocab_size
+
+ # Determine structure sequence types
+ if args.structure_seq is None:
+ args.structure_seq = ""
+ print("Warning: structure_seq was None, setting to empty string")
+
+ # Auto-set structure sequence flags based on structure_seq parameter
+ if 'foldseek_seq' in args.structure_seq:
+ args.use_foldseek = True
+ print("Enabled foldseek_seq based on structure_seq parameter")
+ if 'ss8_seq' in args.structure_seq:
+ args.use_ss8 = True
+ print("Enabled ss8_seq based on structure_seq parameter")
+
+ # If flags are set but structure_seq is not, update structure_seq
+ structure_seq_list = []
+ if args.use_foldseek and 'foldseek_seq' not in args.structure_seq:
+ structure_seq_list.append("foldseek_seq")
+ if args.use_ss8 and 'ss8_seq' not in args.structure_seq:
+ structure_seq_list.append("ss8_seq")
+
+ if structure_seq_list and not args.structure_seq:
+ args.structure_seq = ",".join(structure_seq_list)
+
+ print(f"Training method: {args.eval_method}") # Default for prediction
+ print(f"Structure sequence: {args.structure_seq}")
+ print(f"Use foldseek: {args.use_foldseek}")
+ print(f"Use ss8: {args.use_ss8}")
+ print(f"Problem type: {args.problem_type}")
+ print(f"Number of labels: {args.num_labels}")
+ print(f"Number of attention heads: {args.num_attention_head}")
+
+ # Create and load model
+ try:
+ if args.eval_method in ["full", "ses-adapter", "freeze"]:
+ model = AdapterModel(args)
+ # ! lora/ qlora
+ elif args.eval_method in ['plm-lora', 'plm-qlora', 'plm-dora', 'plm-adalora', 'plm-ia3']:
+ model = LoraModel(args)
+ if args.model_path is not None:
+ model_path = args.model_path
+ else:
+ model_path = f"{args.output_root}/{args.output_dir}/{args.output_model_name}"
+ model.load_state_dict(torch.load(model_path))
+ model.to(device).eval()
+ # ! lora/ qlora
+ if args.eval_method == 'plm-lora':
+ lora_path = model_path.replace(".pt", "_lora")
+ plm_model = PeftModel.from_pretrained(plm_model,lora_path)
+ plm_model = plm_model.merge_and_unload()
+ elif args.eval_method == 'plm-qlora':
+ lora_path = model_path.replace(".pt", "_qlora")
+ plm_model = PeftModel.from_pretrained(plm_model,lora_path)
+ plm_model = plm_model.merge_and_unload()
+ elif args.eval_method == "plm-dora":
+ dora_path = model_path.replace(".pt", "_dora")
+ plm_model = PeftModel.from_pretrained(plm_model, dora_path)
+ plm_model = plm_model.merge_and_unload()
+ elif args.eval_method == "plm-adalora":
+ adalora_path = model_path.replace(".pt", "_adalora")
+ plm_model = PeftModel.from_pretrained(plm_model, adalora_path)
+ plm_model = plm_model.merge_and_unload()
+ elif args.eval_method == "plm-ia3":
+ ia3_path = model_path.replace(".pt", "_ia3")
+ plm_model = PeftModel.from_pretrained(plm_model, ia3_path)
+ plm_model = plm_model.merge_and_unload()
+ plm_model.to(device).eval()
+ plm_model.to(device).eval()
+ return model, plm_model, tokenizer, device
+
+ except Exception as e:
+ print(f"Error: {str(e)}")
+ raise
+
+def process_sequences(args, tokenizer, plm_model_name):
+ """Process and prepare input sequences for prediction"""
+ print("---------- Processing Input Sequences ----------")
+
+ # Process amino acid sequence
+ aa_seq = args.aa_seq.strip()
+ if not aa_seq:
+ raise ValueError("Amino acid sequence is empty")
+
+ # Process structure sequences if needed
+ foldseek_seq = args.foldseek_seq.strip() if args.foldseek_seq else ""
+ ss8_seq = args.ss8_seq.strip() if args.ss8_seq else ""
+
+ # Check if structure sequences are required but not provided
+ if args.use_foldseek and not foldseek_seq:
+ print("Warning: Foldseek sequence is required but not provided.")
+ if args.use_ss8 and not ss8_seq:
+ print("Warning: SS8 sequence is required but not provided.")
+
+ # Format sequences based on model type
+ if 'prot_bert' in plm_model_name or "prot_t5" in plm_model_name:
+ aa_seq = " ".join(list(aa_seq))
+ aa_seq = re.sub(r"[UZOB]", "X", aa_seq)
+ if args.use_foldseek and foldseek_seq:
+ foldseek_seq = " ".join(list(foldseek_seq))
+ if args.use_ss8 and ss8_seq:
+ ss8_seq = " ".join(list(ss8_seq))
+ elif 'ankh' in plm_model_name:
+ aa_seq = list(aa_seq)
+ if args.use_foldseek and foldseek_seq:
+ foldseek_seq = list(foldseek_seq)
+ if args.use_ss8 and ss8_seq:
+ ss8_seq = list(ss8_seq)
+
+ # Truncate sequences if needed
+ if args.max_seq_len:
+ aa_seq = aa_seq[:args.max_seq_len]
+ if args.use_foldseek and foldseek_seq:
+ foldseek_seq = foldseek_seq[:args.max_seq_len]
+ if args.use_ss8 and ss8_seq:
+ ss8_seq = ss8_seq[:args.max_seq_len]
+
+ # Tokenize sequences
+ if 'ankh' in plm_model_name:
+ aa_inputs = tokenizer.batch_encode_plus([aa_seq], add_special_tokens=True, padding=True, is_split_into_words=True, return_tensors="pt")
+ if args.use_foldseek and foldseek_seq:
+ foldseek_inputs = tokenizer.batch_encode_plus([foldseek_seq], add_special_tokens=True, padding=True, is_split_into_words=True, return_tensors="pt")
+ if args.use_ss8 and ss8_seq:
+ ss8_inputs = tokenizer.batch_encode_plus([ss8_seq], add_special_tokens=True, padding=True, is_split_into_words=True, return_tensors="pt")
+ else:
+ aa_inputs = tokenizer([aa_seq], return_tensors="pt", padding=True, truncation=True)
+ if args.use_foldseek and foldseek_seq:
+ foldseek_inputs = tokenizer([foldseek_seq], return_tensors="pt", padding=True, truncation=True)
+ if args.use_ss8 and ss8_seq:
+ ss8_inputs = tokenizer([ss8_seq], return_tensors="pt", padding=True, truncation=True)
+
+ # Prepare data dictionary
+ data_dict = {
+ "aa_seq_input_ids": aa_inputs["input_ids"],
+ "aa_seq_attention_mask": aa_inputs["attention_mask"],
+ }
+
+ # 只有 ProSST 模型需要结构标记
+ if "ProSST" in plm_model_name and hasattr(args, 'prosst_stru_token') and args.prosst_stru_token:
+ try:
+ # 处理 ProSST 结构标记
+ if isinstance(args.prosst_stru_token, str):
+ seq_clean = args.prosst_stru_token.strip("[]").replace(" ","")
+ tokens = list(map(int, seq_clean.split(','))) if seq_clean else []
+ elif isinstance(args.prosst_stru_token, (list, tuple)):
+ tokens = [int(x) for x in args.prosst_stru_token]
+ else:
+ tokens = []
+
+ # 添加到数据字典
+ if tokens:
+ stru_tokens = torch.tensor([tokens], dtype=torch.long)
+ data_dict["aa_seq_stru_tokens"] = stru_tokens
+ else:
+ # 如果没有标记，则使用零填充
+ data_dict["aa_seq_stru_tokens"] = torch.zeros_like(aa_inputs["input_ids"], dtype=torch.long)
+ except Exception as e:
+ print(f"Warning: Failed to process ProSST structure tokens: {e}")
+ # 使用零填充
+ data_dict["aa_seq_stru_tokens"] = torch.zeros_like(aa_inputs["input_ids"], dtype=torch.long)
+
+ if args.use_foldseek and foldseek_seq:
+ data_dict["foldseek_seq_input_ids"] = foldseek_inputs["input_ids"]
+
+ if args.use_ss8 and ss8_seq:
+ data_dict["ss8_seq_input_ids"] = ss8_inputs["input_ids"]
+
+ print("Processed input sequences with keys:", data_dict.keys())
+ return data_dict
+
+def predict(model, data_dict, device, args, plm_model):
+ """Run prediction on the processed input data"""
+ print("---------- Running Prediction ----------")
+
+ # Move data to device
+ for k, v in data_dict.items():
+ data_dict[k] = v.to(device)
+
+ # Run model inference
+ with torch.no_grad():
+ outputs = model(plm_model, data_dict) # Pass the actual plm_model instead of None
+
+ # Process outputs based on problem type
+ if args.problem_type == "regression":
+ predictions = outputs.squeeze().item()
+ print(f"Prediction result: {predictions}")
+ return {"prediction": predictions}
+
+ elif args.problem_type == "single_label_classification":
+ probabilities = torch.nn.functional.softmax(outputs, dim=1)
+ predicted_class = torch.argmax(probabilities, dim=1).item()
+ class_probs = probabilities.squeeze().tolist()
+
+ # Ensure class_probs is a list
+ if not isinstance(class_probs, list):
+ class_probs = [class_probs]
+
+ print(f"Predicted class: {predicted_class}")
+ print(f"Class probabilities: {class_probs}")
+
+ return {
+ "predicted_class": predicted_class,
+ "probabilities": class_probs
+ }
+
+ elif args.problem_type == "multi_label_classification":
+ sigmoid_outputs = torch.sigmoid(outputs)
+ predictions = (sigmoid_outputs > 0.5).int().squeeze().tolist()
+ probabilities = sigmoid_outputs.squeeze().tolist()
+
+ # Ensure predictions and probabilities are lists
+ if not isinstance(predictions, list):
+ predictions = [predictions]
+ if not isinstance(probabilities, list):
+ probabilities = [probabilities]
+
+ print(f"Predicted labels: {predictions}")
+ print(f"Label probabilities: {probabilities}")
+
+ return {
+ "predictions": predictions,
+ "probabilities": probabilities
+ }
+
+def main():
+ try:
+ # Parse arguments
+ args = parse_args()
+
+ # Load model, tokenizer and get device
+ model, plm_model, tokenizer, device = load_model_and_tokenizer(args)
+
+ # Process input sequences
+ data_dict = process_sequences(args, tokenizer, args.plm_model)
+
+ # Run prediction
+ results = predict(model, data_dict, device, args, plm_model)
+
+ # Output results
+ print("\n---------- Prediction Results ----------")
+ print(json.dumps(results, indent=2))
+
+ return 0
+ except Exception as e:
+ print(f"Error: {str(e)}")
+ import traceback
+ traceback.print_exc()
+ return 1
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/src/predict_batch.py b/src/predict_batch.py
new file mode 100644
index 0000000000000000000000000000000000000000..cb9324c0480a2c4c688d567247677c9fb3e9551c
--- /dev/null
+++ b/src/predict_batch.py
@@ -0,0 +1,442 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+import sys
+import os
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'src')))
+import argparse
+import torch
+import re
+import json
+import os
+import warnings
+import numpy as np
+import pandas as pd
+from pathlib import Path
+from tqdm import tqdm
+from transformers import EsmTokenizer, EsmModel, BertModel, BertTokenizer
+from transformers import T5Tokenizer, T5EncoderModel, AutoTokenizer, AutoModel, AutoModelForMaskedLM
+from transformers import logging
+from peft import PeftModel
+
+# Import project modules
+from models.adapter_model import AdapterModel
+from models.lora_model import LoraModel
+from models.pooling import MeanPooling, Attention1dPoolingHead, LightAttentionPoolingHead
+
+# Ignore warning information
+logging.set_verbosity_error()
+warnings.filterwarnings("ignore")
+
+def parse_args():
+ parser = argparse.ArgumentParser(description="Batch predict protein function for multiple sequences")
+
+ # Model parameters
+ parser.add_argument('--eval_method', type=str, default="freeze", choices=["full", "freeze", "plm-lora", "plm-qlora", "ses-adapter", 'plm-dora', 'plm-adalora', 'plm-ia3'], help="Evaluation method")
+ parser.add_argument('--model_path', type=str, required=True, help="Path to the trained model")
+ parser.add_argument('--plm_model', type=str, required=True, help="Pretrained language model name or path")
+ parser.add_argument('--pooling_method', type=str, default="mean", choices=["mean", "attention1d", "light_attention"], help="Pooling method")
+ parser.add_argument('--problem_type', type=str, default="single_label_classification",
+ choices=["single_label_classification", "multi_label_classification", "regression"],
+ help="Problem type")
+ parser.add_argument('--num_labels', type=int, default=2, help="Number of labels")
+ parser.add_argument('--hidden_size', type=int, default=None, help="Embedding hidden size of the model")
+ parser.add_argument('--num_attention_head', type=int, default=8, help="Number of attention heads")
+ parser.add_argument('--attention_probs_dropout', type=float, default=0, help="Attention probs dropout prob")
+ parser.add_argument('--pooling_dropout', type=float, default=0.25, help="Pooling dropout")
+
+ # Input and output parameters
+ parser.add_argument('--input_file', type=str, required=True, help="Path to input CSV file with sequences")
+ parser.add_argument('--output_dir', type=str, required=True, help="Path to output CSV file dir for predictions")
+ parser.add_argument('--output_file', type=str, required=True, help="output CSV file name for predictions")
+ parser.add_argument('--use_foldseek', action='store_true', help="Use foldseek sequence")
+ parser.add_argument('--use_ss8', action='store_true', help="Use secondary structure sequence")
+ parser.add_argument('--structure_seq', type=str, default=None, help="Structure sequence types to use (comma-separated)")
+
+ # Other parameters
+ parser.add_argument('--max_seq_len', type=int, default=1024, help="Maximum sequence length")
+ parser.add_argument('--batch_size', type=int, default=1, help="Batch size for prediction")
+ parser.add_argument('--dataset', type=str, default="Protein-wise", help="Dataset name")
+
+ args = parser.parse_args()
+ return args
+
+def load_model_and_tokenizer(args):
+ print("---------- Loading Model and Tokenizer ----------")
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+
+ # Check if model file exists
+ if not os.path.exists(args.model_path):
+ raise FileNotFoundError(f"Model file not found: {args.model_path}")
+
+ # Load model configuration if available
+ config_path = os.path.join(os.path.dirname(args.model_path), "config.json")
+ try:
+ with open(config_path, "r") as f:
+ config = json.load(f)
+ print(f"Loaded configuration from {config_path}")
+
+ # Update args with config values if they exist
+ if "pooling_method" in config:
+ args.pooling_method = config["pooling_method"]
+ if "problem_type" in config:
+ args.problem_type = config["problem_type"]
+ if "num_labels" in config:
+ args.num_labels = config["num_labels"]
+ if "num_attention_head" in config:
+ args.num_attention_head = config["num_attention_head"]
+ if "attention_probs_dropout" in config:
+ args.attention_probs_dropout = config["attention_probs_dropout"]
+ if "pooling_dropout" in config:
+ args.pooling_dropout = config["pooling_dropout"]
+ except FileNotFoundError:
+ print(f"Model config not found at {config_path}. Using command line arguments.")
+
+ # Build tokenizer and protein language model
+ if "esm" in args.plm_model:
+ tokenizer = EsmTokenizer.from_pretrained(args.plm_model)
+ plm_model = EsmModel.from_pretrained(args.plm_model).to(device).eval()
+ args.hidden_size = plm_model.config.hidden_size
+ elif "bert" in args.plm_model:
+ tokenizer = BertTokenizer.from_pretrained(args.plm_model, do_lower_case=False)
+ plm_model = BertModel.from_pretrained(args.plm_model).to(device).eval()
+ args.hidden_size = plm_model.config.hidden_size
+ elif "prot_t5" in args.plm_model:
+ tokenizer = T5Tokenizer.from_pretrained(args.plm_model, do_lower_case=False)
+ plm_model = T5EncoderModel.from_pretrained(args.plm_model).to(device).eval()
+ args.hidden_size = plm_model.config.d_model
+ elif "ankh" in args.plm_model:
+ tokenizer = AutoTokenizer.from_pretrained(args.plm_model, do_lower_case=False)
+ plm_model = T5EncoderModel.from_pretrained(args.plm_model).to(device).eval()
+ args.hidden_size = plm_model.config.d_model
+ elif "ProSST" in args.plm_model:
+ tokenizer = AutoTokenizer.from_pretrained(args.plm_model, do_lower_case=False)
+ plm_model = AutoModelForMaskedLM.from_pretrained(args.plm_model).to(device).eval()
+ args.hidden_size = plm_model.config.hidden_size
+ elif "Prime" in args.plm_model:
+ tokenizer = AutoTokenizer.from_pretrained(args.plm_model, do_lower_case=False)
+ plm_model = AutoModelForMaskedLM.from_pretrained(args.plm_model).to(device).eval()
+ args.hidden_size = plm_model.config.hidden_size
+ else:
+ tokenizer = AutoTokenizer.from_pretrained(args.plm_model)
+ plm_model = AutoModel.from_pretrained(args.plm_model).to(device).eval()
+ args.hidden_size = plm_model.config.hidden_size
+
+ args.vocab_size = plm_model.config.vocab_size
+
+ # Determine structure sequence types
+ if args.structure_seq is None:
+ args.structure_seq = ""
+ print("Warning: structure_seq was None, setting to empty string")
+
+ # Auto-set structure sequence flags based on structure_seq parameter
+ if 'foldseek_seq' in args.structure_seq:
+ args.use_foldseek = True
+ print("Enabled foldseek_seq based on structure_seq parameter")
+ if 'ss8_seq' in args.structure_seq:
+ args.use_ss8 = True
+ print("Enabled ss8_seq based on structure_seq parameter")
+
+ # If flags are set but structure_seq is not, update structure_seq
+ structure_seq_list = []
+ if args.use_foldseek and 'foldseek_seq' not in args.structure_seq:
+ structure_seq_list.append("foldseek_seq")
+ if args.use_ss8 and 'ss8_seq' not in args.structure_seq:
+ structure_seq_list.append("ss8_seq")
+
+ if structure_seq_list and not args.structure_seq:
+ args.structure_seq = ",".join(structure_seq_list)
+
+ print(f"Training method: {args.eval_method}") # Default for prediction
+ print(f"Structure sequence: {args.structure_seq}")
+ print(f"Use foldseek: {args.use_foldseek}")
+ print(f"Use ss8: {args.use_ss8}")
+ print(f"Problem type: {args.problem_type}")
+ print(f"Number of labels: {args.num_labels}")
+ print(f"Number of attention heads: {args.num_attention_head}")
+
+ # Create and load model
+ try:
+ if args.eval_method in ["full", "ses-adapter", "freeze"]:
+ model = AdapterModel(args)
+ # ! lora/ qlora
+ elif args.eval_method in ['plm-lora', 'plm-qlora', 'plm-dora', 'plm-adalora', 'plm-ia3']:
+ model = LoraModel(args)
+ if args.model_path is not None:
+ model_path = args.model_path
+ else:
+ model_path = f"{args.output_root}/{args.output_dir}/{args.output_model_name}"
+ if args.eval_method == "full":
+ model_weights = torch.load(model_path)
+ model.load_state_dict(model_weights['model_state_dict'])
+ plm_model.load_state_dict(model_weights['plm_state_dict'])
+ else:
+ model.load_state_dict(torch.load(model_path))
+ model.to(device).eval()
+ # ! lora/ qlora
+ if args.eval_method == 'plm-lora':
+ lora_path = model_path.replace(".pt", "_lora")
+ plm_model = PeftModel.from_pretrained(plm_model,lora_path)
+ plm_model = plm_model.merge_and_unload()
+ elif args.eval_method == 'plm-qlora':
+ lora_path = model_path.replace(".pt", "_qlora")
+ plm_model = PeftModel.from_pretrained(plm_model,lora_path)
+ plm_model = plm_model.merge_and_unload()
+ elif args.eval_method == "plm-dora":
+ dora_path = model_path.replace(".pt", "_dora")
+ plm_model = PeftModel.from_pretrained(plm_model, dora_path)
+ plm_model = plm_model.merge_and_unload()
+ elif args.eval_method == "plm-adalora":
+ adalora_path = model_path.replace(".pt", "_adalora")
+ plm_model = PeftModel.from_pretrained(plm_model, adalora_path)
+ plm_model = plm_model.merge_and_unload()
+ elif args.eval_method == "plm-ia3":
+ ia3_path = model_path.replace(".pt", "_ia3")
+ plm_model = PeftModel.from_pretrained(plm_model, ia3_path)
+ plm_model = plm_model.merge_and_unload()
+ plm_model.to(device).eval()
+ return model, plm_model, tokenizer, device
+
+ except Exception as e:
+ print(f"Error: {str(e)}")
+ raise
+
+def process_sequence(args, tokenizer, plm_model_name, aa_seq, foldseek_seq="", ss8_seq="", prosst_stru_token=None):
+ """Process and prepare a single input sequence for prediction"""
+
+ # Process amino acid sequence
+ aa_seq = aa_seq.strip()
+ if not aa_seq:
+ raise ValueError("Amino acid sequence is empty")
+
+ # Process structure sequences if needed
+ foldseek_seq = foldseek_seq.strip() if foldseek_seq else ""
+ ss8_seq = ss8_seq.strip() if ss8_seq else ""
+
+ # Check if structure sequences are required but not provided
+ if args.use_foldseek and not foldseek_seq:
+ print(f"Warning: Foldseek sequence is required but not provided for sequence: {aa_seq[:20]}...")
+ if args.use_ss8 and not ss8_seq:
+ print(f"Warning: SS8 sequence is required but not provided for sequence: {aa_seq[:20]}...")
+
+ # Format sequences based on model type
+ if 'prot_bert' in plm_model_name or "prot_t5" in plm_model_name:
+ aa_seq = " ".join(list(aa_seq))
+ aa_seq = re.sub(r"[UZOB]", "X", aa_seq)
+ if args.use_foldseek and foldseek_seq:
+ foldseek_seq = " ".join(list(foldseek_seq))
+ if args.use_ss8 and ss8_seq:
+ ss8_seq = " ".join(list(ss8_seq))
+ elif 'ankh' in plm_model_name:
+ aa_seq = list(aa_seq)
+ if args.use_foldseek and foldseek_seq:
+ foldseek_seq = list(foldseek_seq)
+ if args.use_ss8 and ss8_seq:
+ ss8_seq = list(ss8_seq)
+
+ # Truncate sequences if needed
+ if args.max_seq_len:
+ aa_seq = aa_seq[:args.max_seq_len]
+ if args.use_foldseek and foldseek_seq:
+ foldseek_seq = foldseek_seq[:args.max_seq_len]
+ if args.use_ss8 and ss8_seq:
+ ss8_seq = ss8_seq[:args.max_seq_len]
+
+ # Tokenize sequences
+ if 'ankh' in plm_model_name:
+ aa_inputs = tokenizer.batch_encode_plus([aa_seq], add_special_tokens=True, padding=True, is_split_into_words=True, return_tensors="pt")
+ if args.use_foldseek and foldseek_seq:
+ foldseek_inputs = tokenizer.batch_encode_plus([foldseek_seq], add_special_tokens=True, padding=True, is_split_into_words=True, return_tensors="pt")
+ if args.use_ss8 and ss8_seq:
+ ss8_inputs = tokenizer.batch_encode_plus([ss8_seq], add_special_tokens=True, padding=True, is_split_into_words=True, return_tensors="pt")
+ else:
+ aa_inputs = tokenizer([aa_seq], return_tensors="pt", padding=True, truncation=True)
+ if args.use_foldseek and foldseek_seq:
+ foldseek_inputs = tokenizer([foldseek_seq], return_tensors="pt", padding=True, truncation=True)
+ if args.use_ss8 and ss8_seq:
+ ss8_inputs = tokenizer([ss8_seq], return_tensors="pt", padding=True, truncation=True)
+
+ # Prepare data dictionary
+ data_dict = {
+ "aa_seq_input_ids": aa_inputs["input_ids"],
+ "aa_seq_attention_mask": aa_inputs["attention_mask"],
+ }
+
+ if "ProSST" in plm_model_name and prosst_stru_token is not None:
+ try:
+ if isinstance(prosst_stru_token, str):
+ seq_clean = prosst_stru_token.strip("[]").replace(" ","")
+ tokens = list(map(int, seq_clean.split(','))) if seq_clean else []
+ elif isinstance(prosst_stru_token, (list, tuple)):
+ tokens = [int(x) for x in prosst_stru_token]
+ else:
+ tokens = []
+
+ if tokens:
+ stru_tokens = torch.tensor([tokens], dtype=torch.long)
+ data_dict["aa_seq_stru_tokens"] = stru_tokens
+ else:
+ data_dict["aa_seq_stru_tokens"] = torch.zeros_like(aa_inputs["input_ids"], dtype=torch.long)
+ except Exception as e:
+ print(f"Warning: Failed to process ProSST structure tokens: {e}")
+ data_dict["aa_seq_stru_tokens"] = torch.zeros_like(aa_inputs["input_ids"], dtype=torch.long)
+
+ if args.use_foldseek and foldseek_seq:
+ data_dict["foldseek_seq_input_ids"] = foldseek_inputs["input_ids"]
+
+ if args.use_ss8 and ss8_seq:
+ data_dict["ss8_seq_input_ids"] = ss8_inputs["input_ids"]
+
+ return data_dict
+
+def predict_batch(model, plm_model, data_dict, device, args):
+ """Run prediction on a batch of processed input data"""
+
+ # Move data to device
+ for k, v in data_dict.items():
+ data_dict[k] = v.to(device)
+
+ # Run model inference
+ with torch.no_grad():
+ outputs = model(plm_model, data_dict)
+
+ # Process outputs based on problem type
+ if args.problem_type == "regression":
+ predictions = outputs.squeeze().cpu().numpy()
+ # 确保返回标量值
+ if np.isscalar(predictions):
+ return {"predictions": predictions}
+ else:
+ # 如果是批处理，返回整个数组
+ return {"predictions": predictions.tolist() if isinstance(predictions, np.ndarray) else predictions}
+
+
+ elif args.problem_type == "single_label_classification":
+ probabilities = torch.nn.functional.softmax(outputs, dim=1)
+ predicted_classes = torch.argmax(probabilities, dim=1).cpu().numpy()
+ class_probs = probabilities.cpu().numpy()
+
+ return {
+ "predicted_classes": predicted_classes.tolist(),
+ "probabilities": class_probs.tolist()
+ }
+
+ elif args.problem_type == "multi_label_classification":
+ sigmoid_outputs = torch.sigmoid(outputs)
+ predictions = (sigmoid_outputs > 0.5).int().cpu().numpy()
+ probabilities = sigmoid_outputs.cpu().numpy()
+
+ return {
+ "predictions": predictions.tolist(),
+ "probabilities": probabilities.tolist()
+ }
+
+def main():
+ # Parse command line arguments
+ args = parse_args()
+
+ try:
+ # Load model and tokenizer
+ model, plm_model, tokenizer, device = load_model_and_tokenizer(args)
+
+ # Read input CSV file
+ print(f"---------- Reading input file: {args.input_file} ----------")
+ try:
+ df = pd.read_csv(args.input_file)
+ print(f"Found {len(df)} sequences in input file")
+ except Exception as e:
+ print(f"Error reading input file: {str(e)}")
+ sys.exit(1)
+
+ # Check required columns
+ required_columns = ["aa_seq"]
+ if args.use_foldseek:
+ required_columns.append("foldseek_seq")
+ if args.use_ss8:
+ required_columns.append("ss8_seq")
+
+ missing_columns = [col for col in required_columns if col not in df.columns]
+ if missing_columns:
+ print(f"Error: Input file is missing required columns: {', '.join(missing_columns)}")
+ sys.exit(1)
+
+ # Initialize results dataframe
+ results = []
+
+ # Process each sequence
+ print("---------- Processing sequences ----------")
+ for idx, row in tqdm(df.iterrows(), total=len(df), desc="Predicting"):
+ try:
+ # Get sequences from row
+ aa_seq = row["aa_seq"]
+ foldseek_seq = row["foldseek_seq"] if "foldseek_seq" in df.columns and args.use_foldseek else ""
+ ss8_seq = row["ss8_seq"] if "ss8_seq" in df.columns and args.use_ss8 else ""
+
+ # Process sequence
+ data_dict = process_sequence(args, tokenizer, args.plm_model, aa_seq, foldseek_seq, ss8_seq)
+
+ # Run prediction
+ prediction_results = predict_batch(model, plm_model, data_dict, device, args)
+
+ # Create result row
+ result_row = {"aa_seq": aa_seq}
+
+ # Add sequence ID if available
+ if "id" in df.columns:
+ result_row["id"] = row["id"]
+
+ # Add prediction results based on problem type
+ if args.problem_type == "regression":
+ # result_row["prediction"] = prediction_results["predictions"][0]
+ if isinstance(prediction_results["predictions"], (list, np.ndarray)):
+ result_row["prediction"] = prediction_results["predictions"][0]
+ else:
+ result_row["prediction"] = prediction_results["predictions"]
+
+ elif args.problem_type == "single_label_classification":
+ result_row["predicted_class"] = prediction_results["predicted_classes"][0]
+
+ # Add class probabilities
+ for i, prob in enumerate(prediction_results["probabilities"][0]):
+ result_row[f"class_{i}_prob"] = prob
+
+ elif args.problem_type == "multi_label_classification":
+ # Add binary predictions
+ for i, pred in enumerate(prediction_results["predictions"][0]):
+ result_row[f"label_{i}"] = pred
+
+ # Add probabilities
+ for i, prob in enumerate(prediction_results["probabilities"][0]):
+ result_row[f"label_{i}_prob"] = prob
+
+ results.append(result_row)
+
+ except Exception as e:
+ print(f"Error processing sequence at index {idx}: {str(e)}")
+ # Add error row
+ error_row = {"aa_seq": aa_seq, "error": str(e)}
+ if "id" in df.columns:
+ error_row["id"] = row["id"]
+ results.append(error_row)
+
+ # Create results dataframe
+ results_df = pd.DataFrame(results)
+
+ # Save results to output file
+ if not os.path.exists(args.output_dir):
+ os.makedirs(args.output_dir)
+ output_file = os.path.join(args.output_dir, args.output_file)
+ print(f"---------- Saving results to {output_file} ----------")
+ results_df.to_csv(output_file, index=False)
+ print(f"Saved {len(results_df)} prediction results")
+
+ print("---------- Batch prediction completed successfully ----------")
+
+ except Exception as e:
+ print(f"Error: {str(e)}")
+ import traceback
+ traceback.print_exc()
+ sys.exit(1)
+
+if __name__ == "__main__":
+ main()
diff --git a/src/train.py b/src/train.py
new file mode 100644
index 0000000000000000000000000000000000000000..902677224a5c8bf3afaada558c2e2023cc04f143
--- /dev/null
+++ b/src/train.py
@@ -0,0 +1,49 @@
+import os
+os.environ["HF_ENDPOINT"]="https://hf-mirror.com"
+import json
+import wandb
+from utils.args import parse_args
+from utils.logger import setup_logging, print_model_parameters
+from data.dataloader import prepare_dataloaders
+from models.model_factory import create_models, lora_factory
+from training.trainer import Trainer
+
+def main():
+ # Parse arguments
+ args = parse_args()
+
+ # Setup logging and wandb
+ logger = setup_logging(args)
+ if args.wandb:
+ wandb.init(
+ project=args.wandb_project,
+ name=args.wandb_run_name,
+ entity=args.wandb_entity,
+ config=vars(args)
+ )
+
+ # Create output directory
+ os.makedirs(args.output_dir, exist_ok=True)
+
+ # Initialize models and tokenizer
+ if args.training_method in ['plm-lora', 'plm-qlora', 'plm-dora', 'plm-adalora', 'plm-ia3']:
+ model, plm_model, tokenizer = lora_factory(args)
+ else:
+ model, plm_model, tokenizer = create_models(args)
+ print_model_parameters(model, plm_model, logger)
+
+ # Prepare data with tokenizer
+ train_loader, val_loader, test_loader = prepare_dataloaders(args, tokenizer, logger)
+
+ # Create trainer
+ trainer = Trainer(args, model, plm_model, logger)
+
+ # Train and evaluate
+ trainer.train(train_loader, val_loader)
+ trainer.test(test_loader)
+
+ if args.wandb:
+ wandb.finish()
+
+if __name__ == "__main__":
+ main()
\ No newline at end of file
diff --git a/src/training/loss_function.py b/src/training/loss_function.py
new file mode 100644
index 0000000000000000000000000000000000000000..4a5ee737069300edabc672c4b54f8586c1cf3740
--- /dev/null
+++ b/src/training/loss_function.py
@@ -0,0 +1,25 @@
+import torch
+import torch.nn as nn
+
+class MultiClassFocalLossWithAlpha(nn.Module):
+ def __init__(self, num_classes, alpha=None, gamma=1, reduction='mean', device="cuda"):
+ super(MultiClassFocalLossWithAlpha, self).__init__()
+ if alpha is None:
+ self.alpha = torch.ones(num_classes, dtype=torch.float32)
+ self.alpha = torch.tensor(alpha).to(device)
+ self.gamma = gamma
+ self.reduction = reduction
+
+ def forward(self, pred, target):
+ alpha = self.alpha[target]
+ log_softmax = torch.log_softmax(pred, dim=1)
+ logpt = torch.gather(log_softmax, dim=1, index=target.view(-1, 1))
+ logpt = logpt.view(-1)
+ ce_loss = -logpt
+ pt = torch.exp(logpt)
+ focal_loss = alpha * (1 - pt) ** self.gamma * ce_loss
+ if self.reduction == "mean":
+ return torch.mean(focal_loss)
+ if self.reduction == "sum":
+ return torch.sum(focal_loss)
+ return focal_loss
\ No newline at end of file
diff --git a/src/training/metrics.py b/src/training/metrics.py
new file mode 100644
index 0000000000000000000000000000000000000000..8cc287e54cadbfa6486a4115cac8f8f68bbf00f5
--- /dev/null
+++ b/src/training/metrics.py
@@ -0,0 +1,146 @@
+import torch
+from torchmetrics.classification import Accuracy, Recall, Precision, MatthewsCorrCoef, AUROC, F1Score, MatthewsCorrCoef
+from torchmetrics.classification import BinaryAccuracy, BinaryRecall, BinaryAUROC, BinaryF1Score, BinaryPrecision, BinaryMatthewsCorrCoef, BinaryF1Score
+from torchmetrics.regression import SpearmanCorrCoef, MeanSquaredError
+from torchmetrics.classification import MultilabelAveragePrecision
+
+
+def count_f1_max(pred, target):
+ """
+ F1 score with the optimal threshold, Copied from TorchDrug.
+
+ This function first enumerates all possible thresholds for deciding positive and negative
+ samples, and then pick the threshold with the maximal F1 score.
+
+ Parameters:
+ pred (Tensor): predictions of shape :math:`(B, N)`
+ target (Tensor): binary targets of shape :math:`(B, N)`
+ """
+
+ order = pred.argsort(descending=True, dim=1)
+ target = target.gather(1, order)
+ precision = target.cumsum(1) / torch.ones_like(target).cumsum(1)
+ recall = target.cumsum(1) / (target.sum(1, keepdim=True) + 1e-10)
+ is_start = torch.zeros_like(target).bool()
+ is_start[:, 0] = 1
+ is_start = torch.scatter(is_start, 1, order, is_start)
+
+ all_order = pred.flatten().argsort(descending=True)
+ order = (
+ order
+ + torch.arange(order.shape[0], device=order.device).unsqueeze(1)
+ * order.shape[1]
+ )
+ order = order.flatten()
+ inv_order = torch.zeros_like(order)
+ inv_order[order] = torch.arange(order.shape[0], device=order.device)
+ is_start = is_start.flatten()[all_order]
+ all_order = inv_order[all_order]
+ precision = precision.flatten()
+ recall = recall.flatten()
+ all_precision = precision[all_order] - torch.where(
+ is_start, torch.zeros_like(precision), precision[all_order - 1]
+ )
+ all_precision = all_precision.cumsum(0) / is_start.cumsum(0)
+ all_recall = recall[all_order] - torch.where(
+ is_start, torch.zeros_like(recall), recall[all_order - 1]
+ )
+ all_recall = all_recall.cumsum(0) / pred.shape[0]
+ all_f1 = 2 * all_precision * all_recall / (all_precision + all_recall + 1e-10)
+ return all_f1.max()
+
+
+class MultilabelF1Max(MultilabelAveragePrecision):
+
+ def compute(self):
+ return count_f1_max(torch.cat(self.preds), torch.cat(self.target))
+
+def setup_metrics(args):
+ """Setup metrics based on problem type and specified metrics list."""
+ metrics_dict = {}
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+
+ for metric_name in args.metrics:
+ if args.problem_type == 'regression':
+ metric_config = _setup_regression_metrics(metric_name, device)
+ elif args.problem_type == 'single_label_classification':
+ if args.num_labels == 2:
+ metric_config = _setup_binary_metrics(metric_name, device)
+ else:
+ metric_config = _setup_multiclass_metrics(metric_name, args.num_labels, device)
+ elif args.problem_type == 'multi_label_classification':
+ metric_config = _setup_multilabel_metrics(metric_name, args.num_labels, device)
+
+ if metric_config:
+ metrics_dict[metric_name] = metric_config['metric']
+
+ # Add loss to metrics if it's the monitor metric
+ if args.monitor == 'loss':
+ metrics_dict['loss'] = 'loss'
+
+ return metrics_dict
+
+def _setup_regression_metrics(metric_name, device):
+ metrics_config = {
+ 'spearman_corr': {
+ 'metric': SpearmanCorrCoef().to(device),
+ },
+ 'mse': {
+ 'metric': MeanSquaredError().to(device),
+ }
+ }
+ return metrics_config.get(metric_name)
+
+def _setup_multiclass_metrics(metric_name, num_classes, device):
+ metrics_config = {
+ 'accuracy': {
+ 'metric': Accuracy(task='multiclass', num_classes=num_classes).to(device),
+ },
+ 'recall': {
+ 'metric': Recall(task='multiclass', num_classes=num_classes).to(device),
+ },
+ 'precision': {
+ 'metric': Precision(task='multiclass', num_classes=num_classes).to(device),
+ },
+ 'f1': {
+ 'metric': F1Score(task='multiclass', num_classes=num_classes).to(device),
+ },
+ 'mcc': {
+ 'metric': MatthewsCorrCoef(task='multiclass', num_classes=num_classes).to(device),
+ },
+ 'auroc': {
+ 'metric': AUROC(task='multiclass', num_classes=num_classes).to(device),
+ }
+ }
+ return metrics_config.get(metric_name)
+
+def _setup_binary_metrics(metric_name, device):
+ metrics_config = {
+ 'accuracy': {
+ 'metric': BinaryAccuracy().to(device),
+ },
+ 'recall': {
+ 'metric': BinaryRecall().to(device),
+ },
+ 'precision': {
+ 'metric': BinaryPrecision().to(device),
+ },
+ 'f1': {
+ 'metric': BinaryF1Score().to(device),
+ },
+ 'mcc': {
+ 'metric': BinaryMatthewsCorrCoef().to(device),
+ },
+ 'auroc': {
+ 'metric': BinaryAUROC().to(device),
+ }
+ }
+ return metrics_config.get(metric_name)
+
+def _setup_multilabel_metrics(metric_name, num_labels, device):
+ metrics_config = {
+ 'f1_max': {
+ 'metric': MultilabelF1Max(num_labels=num_labels).to(device),
+ }
+ }
+ return metrics_config.get(metric_name)
\ No newline at end of file
diff --git a/src/training/scheduler.py b/src/training/scheduler.py
new file mode 100644
index 0000000000000000000000000000000000000000..d83b9f8c05eeaebcb574b05db7b0b48fb9ac3023
--- /dev/null
+++ b/src/training/scheduler.py
@@ -0,0 +1,25 @@
+from torch.optim.lr_scheduler import LambdaLR, CosineAnnealingLR, StepLR
+from transformers import get_linear_schedule_with_warmup, get_cosine_schedule_with_warmup
+
+def create_scheduler(args, optimizer):
+ if not args.scheduler:
+ return None
+
+ num_training_steps = args.num_training_steps
+ num_warmup_steps = args.warmup_steps or num_training_steps // 10
+
+ scheduler_dict = {
+ 'linear': lambda: get_linear_schedule_with_warmup(
+ optimizer,
+ num_warmup_steps=num_warmup_steps,
+ num_training_steps=num_training_steps
+ ),
+ 'cosine': lambda: get_cosine_schedule_with_warmup(
+ optimizer,
+ num_warmup_steps=num_warmup_steps,
+ num_training_steps=num_training_steps
+ ),
+ 'step': lambda: StepLR(optimizer, step_size=30, gamma=0.1)
+ }
+
+ return scheduler_dict[args.scheduler]()
\ No newline at end of file
diff --git a/src/training/trainer.py b/src/training/trainer.py
new file mode 100644
index 0000000000000000000000000000000000000000..939c6f91f58def1a66834d75cb6ee06d20347343
--- /dev/null
+++ b/src/training/trainer.py
@@ -0,0 +1,500 @@
+import os
+import json
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+from accelerate import Accelerator
+from .scheduler import create_scheduler
+from .metrics import setup_metrics
+from .loss_function import MultiClassFocalLossWithAlpha
+import wandb
+from models.model_factory import create_plm_and_tokenizer
+from peft import PeftModel
+
+class Trainer:
+ def __init__(self, args, model, plm_model, logger):
+ self.args = args
+ self.model = model
+ self.plm_model = plm_model
+ self.logger = logger
+ self.device = "cuda" if torch.cuda.is_available() else "cpu"
+
+ # Setup metrics
+ self.metrics_dict = setup_metrics(args)
+
+ # Setup optimizer with different learning rates
+ if self.args.training_method == 'full':
+ # Use a smaller learning rate for PLM
+ optimizer_grouped_parameters = [
+ {
+ "params": self.model.parameters(),
+ "lr": args.learning_rate
+ },
+ {
+ "params": self.plm_model.parameters(),
+ "lr": args.learning_rate
+ }
+ ]
+ self.optimizer = torch.optim.AdamW(optimizer_grouped_parameters)
+ elif self.args.training_method in ['plm-lora', 'plm-qlora', 'plm-dora', 'plm-adalora', 'plm-ia3']:
+ optimizer_grouped_parameters = [
+ {
+ "params": self.model.parameters(),
+ "lr": args.learning_rate
+ },
+ {
+ "params": [param for param in self.plm_model.parameters() if param.requires_grad],
+ "lr": args.learning_rate
+ }
+ ]
+ self.optimizer = torch.optim.AdamW(optimizer_grouped_parameters)
+ else:
+ self.optimizer = torch.optim.AdamW(self.model.parameters(), lr=args.learning_rate)
+
+ # Setup accelerator
+ self.accelerator = Accelerator(gradient_accumulation_steps=args.gradient_accumulation_steps)
+
+ # Setup scheduler
+ self.scheduler = create_scheduler(args, self.optimizer)
+
+ # Setup loss function
+ self.loss_fn = self._setup_loss_function()
+
+ # Prepare for distributed training
+ if self.args.training_method in ['full', 'plm-lora', 'plm-qlora', 'plm-dora', 'plm-adalora', 'plm-ia3']:
+ self.model, self.plm_model, self.optimizer = self.accelerator.prepare(
+ self.model, self.plm_model, self.optimizer
+ )
+ else:
+ self.model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)
+ if self.scheduler:
+ self.scheduler = self.accelerator.prepare(self.scheduler)
+
+ # Training state
+ self.best_val_loss = float("inf")
+ if self.args.monitor_strategy == 'min':
+ self.best_val_metric_score = float("inf")
+ else:
+ self.best_val_metric_score = -float("inf")
+ self.global_steps = 0
+ self.early_stop_counter = 0
+
+ # Save args
+ with open(os.path.join(self.args.output_dir, f'{self.args.output_model_name.split(".")[0]}.json'), 'w') as f:
+ json.dump(self.args.__dict__, f)
+
+ def _setup_loss_function(self):
+ if self.args.problem_type == 'regression':
+ return torch.nn.MSELoss()
+ elif self.args.problem_type == 'multi_label_classification':
+ return torch.nn.BCEWithLogitsLoss()
+ else:
+ return torch.nn.CrossEntropyLoss()
+
+ def train(self, train_loader, val_loader):
+ """Train the model."""
+ for epoch in range(self.args.num_epochs):
+ self.logger.info(f"---------- Epoch {epoch} ----------")
+
+ # Training phase
+ train_loss = self._train_epoch(train_loader)
+ self.logger.info(f'Epoch {epoch} Train Loss: {train_loss:.4f}')
+
+ # Validation phase
+ val_loss, val_metrics = self._validate(val_loader)
+
+ # Handle validation results (model saving, early stopping)
+ self._handle_validation_results(epoch, val_loss, val_metrics)
+
+ # Early stopping check
+ if self._check_early_stopping():
+ self.logger.info(f"Early stop at Epoch {epoch}")
+ break
+
+ def _train_epoch(self, train_loader):
+ self.model.train()
+ if self.args.training_method in ['full', 'plm-lora', 'plm-qlora', 'plm-dora', 'plm-adalora', 'plm-ia3']:
+ self.plm_model.train()
+ total_loss = 0
+ total_samples = 0
+ epoch_iterator = tqdm(train_loader, desc="Training")
+
+ for batch in epoch_iterator:
+ # choose models to accumulate
+ models_to_accumulate = [self.model, self.plm_model] if self.args.training_method in ['full', 'plm-lora', 'plm-qlora', 'plm-dora', 'plm-adalora', 'plm-ia3'] else [self.model]
+
+ with self.accelerator.accumulate(*models_to_accumulate):
+ # Forward and backward
+ loss = self._training_step(batch)
+ self.accelerator.backward(loss)
+
+ # Update statistics
+ batch_size = batch["label"].size(0)
+ total_loss += loss.item() * batch_size
+ total_samples += batch_size
+
+ # Gradient clipping if needed
+ if self.args.max_grad_norm > 0:
+ params_to_clip = (
+ list(self.model.parameters()) + list(self.plm_model.parameters())
+ if self.args.training_method in ['full', 'plm-lora', 'plm-qlora', 'plm-dora', 'plm-adalora', 'plm-ia3']
+ else self.model.parameters()
+ )
+ self.accelerator.clip_grad_norm_(params_to_clip, self.args.max_grad_norm)
+
+ # Optimization step
+ self.optimizer.step()
+ if self.scheduler:
+ self.scheduler.step()
+ self.optimizer.zero_grad()
+
+ # Logging
+ self.global_steps += 1
+ self._log_training_step(loss)
+
+ # Update progress bar
+ epoch_iterator.set_postfix(
+ train_loss=loss.item(),
+ grad_step=self.global_steps // self.args.gradient_accumulation_steps
+ )
+
+ return total_loss / total_samples
+
+ def _training_step(self, batch):
+ # Move batch to device
+ batch = {k: v.to(self.device) for k, v in batch.items()}
+
+ # Forward pass
+ logits = self.model(self.plm_model, batch)
+ loss = self._compute_loss(logits, batch["label"])
+
+ return loss
+
+ def _validate(self, val_loader):
+ """
+ Validate the model.
+
+ Args:
+ val_loader: Validation data loader
+
+ Returns:
+ tuple: (validation_loss, validation_metrics)
+ """
+ self.model.eval()
+ if self.args.training_method in ['full', 'plm-lora', 'plm-qlora', 'plm-dora', 'plm-adalora', 'plm-ia3']:
+ self.plm_model.eval()
+
+ total_loss = 0
+ total_samples = 0
+
+ # Reset all metrics at the start of validation
+ for metric in self.metrics_dict.values():
+ metric.reset()
+
+ with torch.no_grad():
+ for batch in tqdm(val_loader, desc="Validating"):
+ batch = {k: v.to(self.device) for k, v in batch.items()}
+
+ # Forward pass
+ logits = self.model(self.plm_model, batch)
+ loss = self._compute_loss(logits, batch["label"])
+
+ # Update loss statistics
+ batch_size = len(batch["label"])
+ total_loss += loss.item() * batch_size
+ total_samples += batch_size
+
+ # Update metrics
+ self._update_metrics(logits, batch["label"])
+
+ # Compute average loss
+ avg_loss = total_loss / total_samples
+
+ # Compute final metrics
+ metrics_results = {name: metric.compute().item()
+ for name, metric in self.metrics_dict.items()}
+
+ return avg_loss, metrics_results
+
+ def test(self, test_loader):
+ # Load best model
+ self._load_best_model()
+
+ # Add a clear signal that testing is starting
+ self.logger.info("---------- Starting Test Phase ----------")
+
+ # Run evaluation with a custom testing function instead of reusing _validate
+ test_loss, test_metrics = self._test_evaluate(test_loader)
+
+ # Log results
+ self.logger.info("Test Results:")
+ self.logger.info(f"Test Loss: {test_loss:.4f}")
+ for name, value in test_metrics.items():
+ self.logger.info(f"Test {name}: {value:.4f}")
+
+ if self.args.wandb:
+ wandb.log({f"test/{k}": v for k, v in test_metrics.items()})
+ wandb.log({"test/loss": test_loss})
+
+ def _test_evaluate(self, test_loader):
+ """
+ Dedicated evaluation function for test phase with proper labeling.
+ This is almost identical to _validate but with "Testing" progress bar.
+ """
+ self.model.eval()
+ if self.args.training_method in ['full', 'plm-lora', 'plm-qlora', 'plm-dora', 'plm-adalora', 'plm-ia3']:
+ self.plm_model.eval()
+
+ total_loss = 0
+ total_samples = 0
+
+ # Reset all metrics at the start of testing
+ for metric in self.metrics_dict.values():
+ metric.reset()
+
+ with torch.no_grad():
+ # Note the desc is "Testing" instead of "Validating"
+ for batch in tqdm(test_loader, desc="Testing"):
+ batch = {k: v.to(self.device) for k, v in batch.items()}
+
+ # Forward pass
+ logits = self.model(self.plm_model, batch)
+ loss = self._compute_loss(logits, batch["label"])
+
+ # Update loss statistics
+ batch_size = len(batch["label"])
+ total_loss += loss.item() * batch_size
+ total_samples += batch_size
+
+ # Update metrics
+ self._update_metrics(logits, batch["label"])
+
+ # Compute average loss
+ avg_loss = total_loss / total_samples
+
+ # Compute final metrics
+ metrics_results = {name: metric.compute().item()
+ for name, metric in self.metrics_dict.items()}
+
+ return avg_loss, metrics_results
+
+ def _compute_loss(self, logits, labels):
+ if self.args.problem_type == 'regression' and self.args.num_labels == 1:
+ return self.loss_fn(logits.squeeze(), labels.squeeze())
+ elif self.args.problem_type == 'multi_label_classification':
+ return self.loss_fn(logits, labels.float())
+ else:
+ return self.loss_fn(logits, labels)
+
+ def _update_metrics(self, logits, labels):
+ """Update metrics with current batch predictions."""
+ for metric_name, metric in self.metrics_dict.items():
+ if self.args.problem_type == 'regression' and self.args.num_labels == 1:
+ logits = logits.view(-1, 1)
+ labels = labels.view(-1, 1)
+ metric(logits, labels)
+ elif self.args.problem_type == 'multi_label_classification':
+ metric(torch.sigmoid(logits), labels)
+ else:
+ if self.args.num_labels == 2:
+ if metric_name == 'auroc':
+ metric(torch.sigmoid(logits[:, 1]), labels)
+ else:
+ metric(torch.argmax(logits, 1), labels)
+ else:
+ if metric_name == 'auroc':
+ metric(F.softmax(logits, dim=1), labels)
+ else:
+ metric(torch.argmax(logits, 1), labels)
+
+ def _log_training_step(self, loss):
+ if self.args.wandb:
+ wandb.log({
+ "train/loss": loss.item(),
+ "train/learning_rate": self.optimizer.param_groups[0]['lr']
+ }, step=self.global_steps)
+
+ # def _save_model(self, path):
+ # if self.args.training_method in ['full', 'plm-lora']:
+ # torch.save({
+ # 'model_state_dict': self.model.state_dict(),
+ # 'plm_state_dict': self.plm_model.state_dict()
+ # }, path)
+ # else:
+ # torch.save(self.model.state_dict(), path)
+
+ # def _load_best_model(self):
+ # path = os.path.join(self.args.output_dir, self.args.output_model_name)
+ # if self.args.training_method in ['full', 'plm-lora']:
+ # checkpoint = torch.load(path, weights_only=True)
+ # self.model.load_state_dict(checkpoint['model_state_dict'])
+ # self.plm_model.load_state_dict(checkpoint['plm_state_dict'])
+ # else:
+ # self.model.load_state_dict(torch.load(path, weights_only=True))
+ def _save_model(self, path):
+ if self.args.training_method in ['full', 'lora']:
+ model_state = {k: v.cpu() for k, v in self.model.state_dict().items()}
+ plm_state = {k: v.cpu() for k, v in self.plm_model.state_dict().items()}
+ torch.save({
+ 'model_state_dict': model_state,
+ 'plm_state_dict': plm_state
+ }, path)
+ elif self.args.training_method == "plm-lora":
+ model_state = {k: v.cpu() for k, v in self.model.state_dict().items()}
+ torch.save(model_state, path)
+ plm_lora_path = path.replace('.pt', '_lora')
+ self.plm_model.save_pretrained(plm_lora_path)
+ elif self.args.training_method == "plm-qlora":
+ # save model state dict
+ model_state = {k: v.cpu() for k, v in self.model.state_dict().items()}
+ torch.save(model_state, path)
+ plm_qlora_path = path.replace('.pt', '_qlora')
+ # save plm model lora weights
+ self.plm_model.save_pretrained(plm_qlora_path)
+ elif self.args.training_method == "plm-dora":
+ # save model state dict
+ model_state = {k: v.cpu() for k, v in self.model.state_dict().items()}
+ torch.save(model_state, path)
+ plm_dora_path = path.replace('.pt', '_dora')
+ # save plm model lora weights
+ self.plm_model.save_pretrained(plm_dora_path)
+ elif self.args.training_method == "plm-adalora":
+ # save model state dict
+ model_state = {k: v.cpu() for k, v in self.model.state_dict().items()}
+ torch.save(model_state, path)
+ plm_adalora_path = path.replace('.pt', '_adalora')
+ self.plm_model.save_pretrained(plm_adalora_path)
+ elif self.args.training_method == "plm-ia3":
+ model_state = {k: v.cpu() for k, v in self.model.state_dict().items()}
+ torch.save(model_state, path)
+ plm_ia3_path = path.replace('.pt', '_ia3')
+ self.plm_model.save_pretrained(plm_ia3_path)
+ else:
+ model_state = {k: v.cpu() for k, v in self.model.state_dict().items()}
+ torch.save(model_state, path)
+
+ def _load_best_model(self):
+ path = os.path.join(self.args.output_dir, self.args.output_model_name)
+ if self.args.training_method in ['full', 'lora']:
+ checkpoint = torch.load(path, map_location="cpu")
+ self.model.load_state_dict(checkpoint['model_state_dict'])
+ self.plm_model.load_state_dict(checkpoint['plm_state_dict'])
+ self.model.to(self.device)
+ self.plm_model.to(self.device)
+ elif self.args.training_method == "plm-lora":
+ checkpoint = torch.load(path, map_location="cpu")
+ self.model.load_state_dict(checkpoint)
+ plm_lora_path = path.replace('.pt', '_lora')
+ _, self.plm_model = create_plm_and_tokenizer(self.args)
+ self.plm_model = PeftModel.from_pretrained(self.plm_model, plm_lora_path)
+ self.plm_model = self.plm_model.merge_and_unload()
+ self.model.to(self.device)
+ self.plm_model.to(self.device)
+ elif self.args.training_method == "plm-qlora":
+ # load model state dict
+ checkpoint = torch.load(path, map_location="cpu")
+ self.model.load_state_dict(checkpoint)
+ plm_qlora_path = path.replace('.pt', '_qlora')
+ # reload plm model and apply qlora weights
+ _, self.plm_model = create_plm_and_tokenizer(self.args)
+ self.plm_model = PeftModel.from_pretrained(self.plm_model, plm_qlora_path)
+ self.plm_model = self.plm_model.merge_and_unload()
+ self.model.to(self.device)
+ self.plm_model.to(self.device)
+ elif self.args.training_method == "plm-dora":
+ # load model state dict
+ checkpoint = torch.load(path, map_location="cpu")
+ self.model.load_state_dict(checkpoint)
+ plm_dora_path = path.replace('.pt', '_dora')
+ # reload plm model and apply dora weights
+ _, self.plm_model = create_plm_and_tokenizer(self.args)
+ self.plm_model = PeftModel.from_pretrained(self.plm_model, plm_dora_path)
+ self.plm_model = self.plm_model.merge_and_unload()
+ self.model.to(self.device)
+ self.plm_model.to(self.device)
+ elif self.args.training_method == "plm-adalora":
+ # load model state dict
+ checkpoint = torch.load(path, map_location="cpu")
+ self.model.load_state_dict(checkpoint)
+ plm_adalora_path = path.replace('.pt', '_adalora')
+ # reload plm model and apply adalora weights
+ _, self.plm_model = create_plm_and_tokenizer(self.args)
+ self.plm_model = PeftModel.from_pretrained(self.plm_model, plm_adalora_path)
+ self.plm_model = self.plm_model.merge_and_unload()
+ self.model.to(self.device)
+ self.plm_model.to(self.device)
+ elif self.args.training_method == "plm-ia3":
+ checkpoint = torch.load(path, map_location="cpu")
+ self.model.load_state_dict(checkpoint)
+ plm_ia3_path = path.replace('.pt', '_ia3')
+ _, self.plm_model = create_plm_and_tokenizer(self.args)
+ self.plm_model = PeftModel.from_pretrained(self.plm_model, plm_ia3_path)
+ self.plm_model = self.plm_model.merge_and_unload()
+ self.model.to(self.device)
+ self.plm_model.to(self.device)
+ else:
+ checkpoint = torch.load(path, map_location="cpu")
+ self.model.load_state_dict(checkpoint)
+ self.model.to(self.device)
+
+ def _handle_validation_results(self, epoch: int, val_loss: float, val_metrics: dict):
+ """
+ Handle validation results, including model saving and early stopping checks.
+
+ Args:
+ epoch: Current epoch number
+ val_loss: Validation loss
+ val_metrics: Dictionary of validation metrics
+ """
+ # Log validation results
+ self.logger.info(f'Epoch {epoch} Val Loss: {val_loss:.4f}')
+ for metric_name, metric_value in val_metrics.items():
+ self.logger.info(f'Epoch {epoch} Val {metric_name}: {metric_value:.4f}')
+
+ if self.args.wandb:
+ wandb.log({
+ "val/loss": val_loss,
+ **{f"val/{k}": v for k, v in val_metrics.items()}
+ }, step=self.global_steps)
+
+ # Check if we should save the model
+ should_save = False
+ monitor_value = val_loss
+
+ # If monitoring a specific metric
+ if self.args.monitor != 'loss' and self.args.monitor in val_metrics:
+ monitor_value = val_metrics[self.args.monitor]
+
+ # Check if current result is better
+ if self.args.monitor_strategy == 'min':
+ if monitor_value < self.best_val_metric_score:
+ should_save = True
+ self.best_val_metric_score = monitor_value
+ self.early_stop_counter = 0
+ else:
+ self.early_stop_counter += 1
+ else: # strategy == 'max'
+ if monitor_value > self.best_val_metric_score:
+ should_save = True
+ self.best_val_metric_score = monitor_value
+ self.early_stop_counter = 0
+ else:
+ self.early_stop_counter += 1
+
+ # Save model if improved
+ if should_save:
+ self.logger.info(f"Saving model with best val {self.args.monitor}: {monitor_value:.4f}")
+ save_path = os.path.join(self.args.output_dir, self.args.output_model_name)
+ self._save_model(save_path)
+
+ def _check_early_stopping(self) -> bool:
+ """
+ Check if training should be stopped early.
+
+ Returns:
+ bool: True if training should stop, False otherwise
+ """
+ if self.args.patience > 0 and self.early_stop_counter >= self.args.patience:
+ self.logger.info(f"Early stopping triggered after {self.early_stop_counter} epochs without improvement")
+ return True
+ return False
diff --git a/src/utils/args.py b/src/utils/args.py
new file mode 100644
index 0000000000000000000000000000000000000000..a9e8a421ae9b1016c6b7b384c5bd2acf5240bb28
--- /dev/null
+++ b/src/utils/args.py
@@ -0,0 +1,162 @@
+import argparse
+import json
+import os
+import warnings
+from typing import Dict, Any
+from datetime import datetime
+
+def parse_args() -> Dict[str, Any]:
+ """Parse and validate command line arguments."""
+ parser = create_argument_parser()
+ args = parser.parse_args()
+
+ # Validate and process arguments
+ validate_args(args)
+ process_dataset_config(args)
+ setup_output_dirs(args)
+ setup_wandb_config(args)
+
+ return args
+
+def create_argument_parser() -> argparse.ArgumentParser:
+ """Create argument parser with all training arguments."""
+ parser = argparse.ArgumentParser()
+
+ # Model parameters
+ add_model_args(parser)
+
+ # Dataset parameters
+ add_dataset_args(parser)
+
+ # Training parameters
+ add_training_args(parser)
+
+ # Output parameters
+ add_output_args(parser)
+
+ # Wandb parameters
+ add_wandb_args(parser)
+
+ return parser
+
+def add_model_args(parser: argparse.ArgumentParser):
+ """Add model-related arguments."""
+ model_group = parser.add_argument_group('Model Configuration')
+ model_group.add_argument('--hidden_size', type=int, default=None)
+ model_group.add_argument('--num_attention_head', type=int, default=8)
+ model_group.add_argument('--attention_probs_dropout', type=float, default=0.1)
+ model_group.add_argument('--plm_model', type=str, default='facebook/esm2_t33_650M_UR50D')
+ model_group.add_argument('--pooling_method', type=str, default='mean',
+ choices=['mean', 'attention1d', 'light_attention'])
+ model_group.add_argument('--pooling_dropout', type=float, default=0.1)
+
+def add_dataset_args(parser: argparse.ArgumentParser):
+ """Add dataset-related arguments."""
+ data_group = parser.add_argument_group('Dataset Configuration')
+ data_group.add_argument('--dataset', type=str)
+ data_group.add_argument('--dataset_config', type=str)
+ data_group.add_argument('--normalize', type=str)
+ data_group.add_argument('--num_labels', type=int)
+ data_group.add_argument('--problem_type', type=str)
+ data_group.add_argument('--pdb_type', type=str)
+ data_group.add_argument('--train_file', type=str)
+ data_group.add_argument('--valid_file', type=str)
+ data_group.add_argument('--test_file', type=str)
+ data_group.add_argument('--metrics', type=str)
+
+def add_training_args(parser: argparse.ArgumentParser):
+ """Add training-related arguments."""
+ train_group = parser.add_argument_group('Training Configuration')
+ train_group.add_argument('--seed', type=int, default=3407)
+ train_group.add_argument('--learning_rate', type=float, default=1e-3)
+ train_group.add_argument('--scheduler', type=str, choices=['linear', 'cosine', 'step'])
+ train_group.add_argument('--warmup_steps', type=int, default=0)
+ train_group.add_argument('--num_workers', type=int, default=4)
+ train_group.add_argument('--batch_size', type=int)
+ train_group.add_argument('--batch_token', type=int)
+ train_group.add_argument('--num_epochs', type=int, default=100)
+ train_group.add_argument('--max_seq_len', type=int, default=-1)
+ train_group.add_argument('--gradient_accumulation_steps', type=int, default=1)
+ train_group.add_argument('--max_grad_norm', type=float, default=-1)
+ train_group.add_argument('--patience', type=int, default=10)
+ train_group.add_argument('--monitor', type=str)
+ train_group.add_argument('--monitor_strategy', type=str, choices=['max', 'min'])
+ train_group.add_argument('--training_method', type=str, default='freeze',
+ choices=['full', 'freeze', 'lora', 'ses-adapter', 'plm-lora', 'plm-qlora', 'plm-adalora', 'plm-dora', 'plm-ia3'])
+ parser.add_argument("--lora_r", type=int, default=8, help="lora r")
+ parser.add_argument("--lora_alpha", type=int, default=32, help="lora_alpha")
+ parser.add_argument("--lora_dropout", type=float, default=0.1, help="lora_dropout")
+ parser.add_argument("--feedforward_modules", type=str, default="w0")
+
+ parser.add_argument(
+ "--lora_target_modules",
+ nargs="+",
+ default=["query", "key", "value"],
+ help="lora target module",
+ )
+ train_group.add_argument('--structure_seq', type=str, default='')
+
+def add_output_args(parser: argparse.ArgumentParser):
+ """Add output-related arguments."""
+ output_group = parser.add_argument_group('Output Configuration')
+ output_group.add_argument('--output_model_name', type=str)
+ output_group.add_argument('--output_root', default="ckpt")
+ output_group.add_argument('--output_dir', default=None)
+
+def add_wandb_args(parser: argparse.ArgumentParser):
+ """Add wandb-related arguments."""
+ wandb_group = parser.add_argument_group('Wandb Configuration')
+ wandb_group.add_argument('--wandb', action='store_true')
+ wandb_group.add_argument('--wandb_entity', type=str)
+ wandb_group.add_argument('--wandb_project', type=str, default='VenusFactory')
+ wandb_group.add_argument('--wandb_run_name', type=str)
+
+def validate_args(args: argparse.Namespace):
+ """Validate command line arguments."""
+ if args.batch_size is None and args.batch_token is None:
+ raise ValueError("batch_size or batch_token must be provided")
+
+ if args.training_method == 'ses-adapter':
+ if args.structure_seq is None:
+ raise ValueError("structure_seq must be provided for ses-adapter")
+ args.structure_seq = args.structure_seq.split(',')
+ else:
+ args.structure_seq = []
+
+def process_dataset_config(args: argparse.Namespace):
+ """Process dataset configuration file."""
+ if not args.dataset_config:
+ return
+
+ config = json.load(open(args.dataset_config))
+
+ # Update args with dataset config values if not already set
+ for key in ['dataset', 'pdb_type', 'train_file', 'valid_file', 'test_file',
+ 'num_labels', 'problem_type', 'monitor', 'monitor_strategy',
+ 'metrics', 'normalize']:
+ if getattr(args, key) is None and key in config:
+ setattr(args, key, config[key])
+
+ # Handle metrics specially
+ if args.metrics:
+ args.metrics = args.metrics.split(',')
+ if args.metrics == ['None']:
+ args.metrics = ['loss']
+ warnings.warn("No metrics provided, using default metrics: loss")
+
+def setup_output_dirs(args: argparse.Namespace):
+ """Setup output directories."""
+ if args.output_dir is None:
+ current_date = strftime("%Y%m%d", localtime())
+ args.output_dir = os.path.join(args.output_root, current_date)
+ else:
+ args.output_dir = os.path.join(args.output_root, args.output_dir)
+ os.makedirs(args.output_dir, exist_ok=True)
+
+def setup_wandb_config(args: argparse.Namespace):
+ """Setup wandb configuration."""
+ if args.wandb:
+ if args.wandb_run_name is None:
+ args.wandb_run_name = f"VenusFactory-{args.dataset}"
+ if args.output_model_name is None:
+ args.output_model_name = f"{args.wandb_run_name}.pt"
diff --git a/src/utils/common.py b/src/utils/common.py
new file mode 100644
index 0000000000000000000000000000000000000000..8a031c68cfe70241bed7244da13f3cdbfac97b7c
--- /dev/null
+++ b/src/utils/common.py
@@ -0,0 +1,41 @@
+import random
+import numpy as np
+import torch
+from typing import List, Dict, Any
+
+def set_seed(seed: int):
+ """Set random seed for reproducibility."""
+ random.seed(seed)
+ np.random.seed(seed)
+ torch.manual_seed(seed)
+ if torch.cuda.is_available():
+ torch.cuda.manual_seed_all(seed)
+
+def min_max_normalize_dataset(train_dataset: List[Dict[str, Any]],
+ val_dataset: List[Dict[str, Any]],
+ test_dataset: List[Dict[str, Any]]) -> tuple:
+ """Normalize datasets using min-max normalization."""
+ # Get all labels from training set
+ train_labels = [data['label'] for data in train_dataset]
+
+ # Calculate min and max from training set
+ min_val = min(train_labels)
+ max_val = max(train_labels)
+
+ # Normalize all datasets
+ for dataset in [train_dataset, val_dataset, test_dataset]:
+ for data in dataset:
+ data['label'] = (data['label'] - min_val) / (max_val - min_val)
+
+ return train_dataset, val_dataset, test_dataset
+
+def check_early_stopping(val_list: List[float],
+ optimize_func: callable,
+ patience: int = 10) -> bool:
+ """Check if training should stop early."""
+ if len(val_list) < patience:
+ return False
+
+ best_val = optimize_func(val_list[:-patience])
+ return all(optimize_func([best_val, val]) == best_val
+ for val in val_list[-patience:])
\ No newline at end of file
diff --git a/src/utils/data_utils.py b/src/utils/data_utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..ae0207a7508f9978fa5ed9549dada00f53e024c6
--- /dev/null
+++ b/src/utils/data_utils.py
@@ -0,0 +1,87 @@
+import random
+import biotite
+import numpy as np
+import torch.utils.data as data
+from typing import List
+from biotite.structure.residues import get_residues
+from biotite.sequence import ProteinSequence
+from biotite.structure.io import pdbx, pdb
+from biotite.structure import filter_backbone
+from biotite.structure import get_chains
+
+def load_structure(fpath, chain=None):
+ """
+ Args:
+ fpath: filepath to either pdb or cif file
+ chain: the chain id or list of chain ids to load
+ Returns:
+ biotite.structure.AtomArray
+ """
+ if fpath.endswith('cif'):
+ with open(fpath) as fin:
+ pdbxf = pdbx.PDBxFile.read(fin)
+ structure = pdbx.get_structure(pdbxf, model=1)
+ elif fpath.endswith('pdb'):
+ with open(fpath) as fin:
+ pdbf = pdb.PDBFile.read(fin)
+ structure = pdb.get_structure(pdbf, model=1)
+ bbmask = filter_backbone(structure)
+ structure = structure[bbmask]
+ all_chains = get_chains(structure)
+ if len(all_chains) == 0:
+ raise ValueError('No chains found in the input file.')
+ if chain is None:
+ chain_ids = all_chains
+ elif isinstance(chain, list):
+ chain_ids = chain
+ else:
+ chain_ids = [chain]
+ for chain in chain_ids:
+ if chain not in all_chains:
+ raise ValueError(f'Chain {chain} not found in input file')
+ chain_filter = [a.chain_id in chain_ids for a in structure]
+ structure = structure[chain_filter]
+ return structure
+
+def get_atom_coords_residuewise(atoms: List[str], struct: biotite.structure.AtomArray):
+ """
+ Example for atoms argument: ["N", "CA", "C"]
+ """
+ def filterfn(s, axis=None):
+ filters = np.stack([s.atom_name == name for name in atoms], axis=1)
+ sum = filters.sum(0)
+ if not np.all(sum <= np.ones(filters.shape[1])):
+ raise RuntimeError("structure has multiple atoms with same name")
+ index = filters.argmax(0)
+ coords = s[index].coord
+ coords[sum == 0] = float("nan")
+ return coords
+
+ return biotite.structure.apply_residue_wise(struct, struct, filterfn)
+
+def extract_coords_from_structure(structure: biotite.structure.AtomArray):
+ """
+ Args:
+ structure: An instance of biotite AtomArray
+ Returns:
+ Tuple (coords, seq)
+ - coords is an L x 3 x 3 array for N, CA, C coordinates
+ - seq is the extracted sequence
+ """
+ coords = get_atom_coords_residuewise(["N", "CA", "C"], structure)
+ residue_identities = get_residues(structure)[1]
+ seq = ''.join([ProteinSequence.convert_letter_3to1(r) for r in residue_identities])
+ return coords
+
+def extract_seq_from_pdb(pdb_file, chain=None):
+ """
+ Args:
+ structure: An instance of biotite AtomArray
+ Returns:
+ - seq is the extracted sequence
+ """
+ structure = load_structure(pdb_file, chain)
+ residue_identities = get_residues(structure)[1]
+ seq = ''.join([ProteinSequence.convert_letter_3to1(r) for r in residue_identities])
+ return seq
+
diff --git a/src/utils/logger.py b/src/utils/logger.py
new file mode 100644
index 0000000000000000000000000000000000000000..16acf74c31c13bcc7884b2040219ed21b525da70
--- /dev/null
+++ b/src/utils/logger.py
@@ -0,0 +1,102 @@
+import logging
+import os
+from typing import Dict, Any, Union, Tuple
+import torch.nn as nn
+from transformers import PreTrainedModel
+
+def setup_logging(args: Dict[str, Any]) -> logging.Logger:
+ """Setup logging configuration."""
+ # Create logger
+ logger = logging.getLogger('training')
+ logger.setLevel(logging.INFO)
+
+ # Create formatters and handlers
+ formatter = logging.Formatter(
+ '%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+ datefmt='%Y-%m-%d %H:%M:%S'
+ )
+
+ # Console handler
+ console_handler = logging.StreamHandler()
+ console_handler.setLevel(logging.INFO)
+ console_handler.setFormatter(formatter)
+ logger.addHandler(console_handler)
+
+ # File handler
+ log_file = os.path.join(args.output_dir, f'{args.output_model_name.split(".")[0]}_training.log')
+ file_handler = logging.FileHandler(log_file)
+ file_handler.setLevel(logging.INFO)
+ file_handler.setFormatter(formatter)
+ logger.addHandler(file_handler)
+
+ # Log initial info
+ logger.info("Starting training with configuration:")
+ for key, value in vars(args).items():
+ logger.info(f"{key}: {value}")
+
+ return logger
+
+def count_parameters(model: Union[nn.Module, PreTrainedModel]) -> Tuple[int, int]:
+ """
+ Count total and trainable parameters in model.
+
+ Args:
+ model: PyTorch model or Hugging Face model
+
+ Returns:
+ Tuple of (total_params, trainable_params)
+ """
+ total_params = sum(p.numel() for p in model.parameters())
+ trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+ return total_params, trainable_params
+
+def format_parameter_count(count: int) -> str:
+ """Format parameter count with appropriate unit."""
+ if count >= 1e9:
+ return f"{count/1e9:.2f}B"
+ elif count >= 1e6:
+ return f"{count/1e6:.2f}M"
+ elif count >= 1e3:
+ return f"{count/1e3:.2f}K"
+ return str(count)
+
+def print_model_parameters(model: nn.Module, plm_model: PreTrainedModel, logger=None):
+ """
+ Print parameter statistics for both adapter and PLM models.
+
+ Args:
+ model: Adapter model
+ plm_model: Pre-trained language model
+ logger: Optional logger for output
+ """
+ # Count adapter parameters
+ adapter_total, adapter_trainable = count_parameters(model)
+
+ # Count PLM parameters
+ plm_total, plm_trainable = count_parameters(plm_model)
+
+ # Prepare output strings
+ output = [
+ "------------------------",
+ "Model Parameters Statistics:",
+ "------------------------",
+ f"Adapter Model:",
+ f" Total parameters: {format_parameter_count(adapter_total)}",
+ f" Trainable parameters: {format_parameter_count(adapter_trainable)}",
+ f"Pre-trained Model:",
+ f" Total parameters: {format_parameter_count(plm_total)}",
+ f" Trainable parameters: {format_parameter_count(plm_trainable)}",
+ f"Combined:",
+ f" Total parameters: {format_parameter_count(adapter_total + plm_total)}",
+ f" Trainable parameters: {format_parameter_count(adapter_trainable + plm_trainable)}",
+ f" Trainable percentage: {((adapter_trainable + plm_trainable) / (adapter_total + plm_total)) * 100:.2f}%",
+ "------------------------"
+ ]
+
+ # Print output
+ if logger:
+ for line in output:
+ logger.info(line)
+ else:
+ for line in output:
+ print(line)
\ No newline at end of file
diff --git a/src/utils/lora_tools.py b/src/utils/lora_tools.py
new file mode 100644
index 0000000000000000000000000000000000000000..7037e15ebeb16af84a792d727b5d872e4952a7b9
--- /dev/null
+++ b/src/utils/lora_tools.py
@@ -0,0 +1,75 @@
+import torch
+from transformers import AutoTokenizer, EsmModel, T5Tokenizer, T5EncoderModel, BertModel, AutoModelForMaskedLM
+from transformers import BertTokenizer, EsmTokenizer, T5Tokenizer
+from peft import LoraConfig, get_peft_model, PeftModel, PeftConfig
+from typing import List, Dict, Any, Tuple
+from transformers import PreTrainedModel
+
+
+def prepare_for_lora_model(
+ based_model,
+ lora_r: int = 8,
+ lora_alpha: int = 32,
+ lora_dropout: float = 0.1,
+ target_modules: List[str,] = ["key", "query", "value"],
+):
+ if not isinstance(based_model, PreTrainedModel):
+ raise TypeError("based_model must be a PreTrainedModel instance")
+
+ # validate target_modules exist in model
+ available_modules = [name for name, _ in based_model.named_modules()]
+ for module in target_modules:
+ if not any(module in name for name in available_modules):
+ raise ValueError(f"Target module {module} not found in model")
+ # get lora config
+ lora_config = LoraConfig(
+ r=lora_r,
+ lora_alpha=lora_alpha,
+ lora_dropout=lora_dropout,
+ target_modules=target_modules,
+ )
+ # get lora model
+ model = get_peft_model(based_model, lora_config)
+ print("Lora model is ready! num of trainable_parameters: ")
+ model.print_trainable_parameters()
+ return model
+
+
+def load_lora_model(base_model, lora_ckpt_path):
+ model = PeftModel.from_pretrained(base_model, lora_ckpt_path)
+ return model
+
+
+def load_eval_base_model(plm_model):
+ device = "cuda" if torch.cuda.is_available() else "cpu"
+ if "esm" in plm_model:
+ base_model = EsmModel.from_pretrained(plm_model).to(device)
+ elif "bert" in plm_model:
+ base_model = BertModel.from_pretrained(plm_model).to(device)
+ elif "prot_t5" in plm_model:
+ base_model = T5EncoderModel.from_pretrained(plm_model).to(device)
+ elif "ankh" in plm_model:
+ base_model = T5EncoderModel.from_pretrained(plm_model).to(device)
+ elif "ProSST" in plm_model:
+ base_model = AutoModelForMaskedLM.from_pretrained(plm_model).to(device)
+
+ return base_model
+
+
+def check_lora_params(model):
+ lora_params = [
+ (name, param) for name, param in model.named_parameters() if "lora_" in name
+ ]
+ print(f"\n num of lora params: {len(lora_params)}")
+
+ if len(lora_params) == 0:
+ print("warning: no lora params found!")
+ else:
+ print("\n first lora param:")
+ name, param = lora_params[0]
+ print(f"name: {name}")
+ print(f"param.shape: {param.shape}")
+ print(f"param.dtype: {param.dtype}")
+ print(f"param.device: {param.device}")
+ # print(f"param_value:\n{param.data.cpu().numpy()}")
+ print(f"requires_grad: {param.requires_grad}")
diff --git a/src/web/download_tab.py b/src/web/download_tab.py
new file mode 100644
index 0000000000000000000000000000000000000000..d74d551cc1910b474436b2b9c8d7074ddeaa0410
--- /dev/null
+++ b/src/web/download_tab.py
@@ -0,0 +1,676 @@
+import gradio as gr
+import os
+import subprocess
+import numpy as np
+import plotly.graph_objects as go
+from Bio.PDB import PDBParser
+import io
+import base64
+from typing import Dict, Any, Tuple, List, Optional
+
+def create_download_tab(constant: Dict[str, Any]) -> Dict[str, Any]:
+ """
+ Create the download tab with various options for downloading protein data.
+
+ Args:
+ constant: Dictionary containing constant values for the application
+
+ Returns:
+ Dictionary containing any state information
+ """
+ def run_download_script(script_name: str, **kwargs) -> str:
+ """
+ Run a download script with the specified arguments.
+
+ Args:
+ script_name: Name of the script to run
+ **kwargs: Arguments to pass to the script
+
+ Returns:
+ Output of the script as a string
+ """
+ cmd = ["python", f"src/crawler/{script_name}"]
+ for k, v in kwargs.items():
+ if v is None: # Skip None values
+ continue
+ if isinstance(v, bool): # Handle boolean flags
+ if v:
+ cmd.append(f"--{k}")
+ elif v == "--merge": # Handle special merge flag
+ cmd.append(v)
+ else: # Handle regular arguments
+ cmd.extend([f"--{k}", str(v)])
+
+ try:
+ result = subprocess.run(cmd, capture_output=True, text=True)
+ if result.returncode == 0:
+ return f"Download completed successfully\n{result.stdout}"
+ else:
+ return f"Error during download:\n{result.stderr}"
+ except Exception as e:
+ return f"Failed to run download script: {str(e)}"
+
+ # Function to visualize protein structure using Plotly
+ def visualize_protein_structure(pdb_file: str) -> Tuple[str, go.Figure]:
+ """
+ Visualize a protein structure from a PDB file using Plotly for interactive 3D visualization.
+
+ Args:
+ pdb_file: Path to the PDB file
+
+ Returns:
+ Tuple containing status message and Plotly figure
+ """
+ try:
+ if not os.path.exists(pdb_file):
+ return f"File not found: {pdb_file}", None
+
+ # Parse the PDB file
+ parser = PDBParser(QUIET=True)
+ structure = parser.get_structure("protein", pdb_file)
+
+ # Extract atom coordinates and information for all atoms
+ all_atoms_x, all_atoms_y, all_atoms_z = [], [], []
+ all_atoms_text = [] # For hover information
+ all_atoms_color = []
+
+ # Color mapping for different atom types
+ color_map = {
+ 'C': '#333333', # Dark gray for carbon
+ 'N': '#3050F8', # Blue for nitrogen
+ 'O': '#FF2010', # Red for oxygen
+ 'S': '#FFFF30', # Yellow for sulfur
+ 'P': '#FF8000', # Orange for phosphorus
+ 'H': '#E0E0E0', # Light gray for hydrogen
+ 'CA': '#00FF00' # Green for alpha carbon
+ }
+
+ # Extract backbone (CA atoms) for the ribbon representation
+ ca_x, ca_y, ca_z = [], [], []
+ ca_text = []
+
+ # Track chains for coloring
+ chains = {}
+ chain_colors = [
+ '#1F77B4', '#FF7F0E', '#2CA02C', '#D62728', '#9467BD',
+ '#8C564B', '#E377C2', '#7F7F7F', '#BCBD22', '#17BECF'
+ ]
+
+ # Create a Plotly figure
+ fig = go.Figure()
+
+ # Track the number of backbone traces for visibility control
+ backbone_trace_count = 0
+
+ # Extract coordinates and properties
+ for model in structure:
+ for chain in model:
+ chain_id = chain.get_id()
+ if chain_id not in chains:
+ chains[chain_id] = len(chains) % len(chain_colors)
+
+ chain_color = chain_colors[chains[chain_id]]
+
+ # Collect CA atoms for this chain
+ chain_ca_x, chain_ca_y, chain_ca_z = [], [], []
+ chain_ca_text = []
+
+ for residue in chain:
+ res_id = residue.get_id()
+ res_name = residue.get_resname()
+ res_num = res_id[1]
+
+ # Extract CA atoms for backbone trace
+ if 'CA' in residue:
+ ca = residue['CA'].get_coord()
+ chain_ca_x.append(ca[0])
+ chain_ca_y.append(ca[1])
+ chain_ca_z.append(ca[2])
+ chain_ca_text.append(f"Chain {chain_id}, {res_name} {res_num}")
+
+ # Also add to global CA lists
+ ca_x.append(ca[0])
+ ca_y.append(ca[1])
+ ca_z.append(ca[2])
+ ca_text.append(f"Chain {chain_id}, {res_name} {res_num}")
+
+ # Extract all atoms
+ for atom in residue:
+ coord = atom.get_coord()
+ all_atoms_x.append(coord[0])
+ all_atoms_y.append(coord[1])
+ all_atoms_z.append(coord[2])
+
+ atom_name = atom.get_name()
+ atom_element = atom.element
+
+ all_atoms_text.append(f"Chain {chain_id}, {res_name} {res_num}, {atom_name}")
+
+ # Determine atom color
+ if atom_name == 'CA':
+ all_atoms_color.append(color_map.get('CA', '#808080'))
+ else:
+ all_atoms_color.append(color_map.get(atom_element, '#808080'))
+
+ # Add this chain's CA atoms as a separate trace for better visualization
+ if chain_ca_x:
+ fig.add_trace(go.Scatter3d(
+ x=chain_ca_x,
+ y=chain_ca_y,
+ z=chain_ca_z,
+ mode='lines',
+ name=f'Chain {chain_id}',
+ line=dict(color=chain_color, width=8), # Increased line width
+ text=chain_ca_text,
+ hoverinfo='text',
+ showlegend=True
+ ))
+ backbone_trace_count += 1
+
+ # Add backbone trace (CA atoms as markers)
+ fig.add_trace(go.Scatter3d(
+ x=ca_x,
+ y=ca_y,
+ z=ca_z,
+ mode='markers',
+ name='Backbone',
+ marker=dict(
+ size=7, # Increased marker size
+ color='#00FF00',
+ opacity=0.8,
+ symbol='circle'
+ ),
+ text=ca_text,
+ hoverinfo='text',
+ showlegend=True
+ ))
+ backbone_trace_count += 1
+
+ # Add all atoms as small markers
+ fig.add_trace(go.Scatter3d(
+ x=all_atoms_x,
+ y=all_atoms_y,
+ z=all_atoms_z,
+ mode='markers',
+ name='All Atoms',
+ marker=dict(
+ size=2.5,
+ color=all_atoms_color,
+ opacity=0.6
+ ),
+ text=all_atoms_text,
+ hoverinfo='text',
+ showlegend=True,
+ visible='legendonly' # Hide by default, can be toggled in legend
+ ))
+
+ # Set layout properties
+ pdb_id = os.path.basename(pdb_file).split('.')[0]
+ fig.update_layout(
+ title=dict(
+ text=f"Structure: {pdb_id}",
+ font=dict(size=20, family="Arial, sans-serif")
+ ),
+ scene=dict(
+ xaxis=dict(title='X (Å)', showbackground=False, showgrid=True, gridcolor='lightgray'),
+ yaxis=dict(title='Y (Å)', showbackground=False, showgrid=True, gridcolor='lightgray'),
+ zaxis=dict(title='Z (Å)', showbackground=False, showgrid=True, gridcolor='lightgray'),
+ aspectmode='data',
+ camera=dict(
+ eye=dict(x=1.5, y=1.5, z=1.5)
+ )
+ ),
+ margin=dict(l=0, r=0, b=0, t=40),
+ legend=dict(
+ yanchor="top",
+ y=0.99,
+ xanchor="left",
+ x=0.01,
+ bgcolor="rgba(255, 255, 255, 0.8)",
+ bordercolor="lightgray",
+ borderwidth=1
+ ),
+ template="plotly_white",
+ height=600, # Increase height for better visualization
+ width=800 # Set width for better aspect ratio
+ )
+
+ # Create visibility arrays for the buttons
+ # For "Backbone Only": all backbone traces visible, all atoms hidden
+ backbone_only_visibility = [True] * backbone_trace_count + [False]
+ # For "All Atoms": all traces visible
+ all_atoms_visibility = [True] * (backbone_trace_count + 1)
+
+ # Add buttons for different views
+ fig.update_layout(
+ updatemenus=[
+ dict(
+ type="buttons",
+ direction="right",
+ buttons=[
+ dict(
+ args=[{"visible": backbone_only_visibility}],
+ label="Backbone Only",
+ method="update"
+ ),
+ dict(
+ args=[{"visible": all_atoms_visibility}],
+ label="All Atoms",
+ method="update"
+ )
+ ],
+ pad={"r": 10, "t": 10},
+ showactive=True,
+ x=0.1,
+ xanchor="left",
+ y=1.1,
+ yanchor="top",
+ bgcolor="rgba(255, 255, 255, 0.8)",
+ bordercolor="lightgray",
+ borderwidth=1
+ ),
+ ]
+ )
+
+ return f"Successfully visualized structure from {pdb_file}", fig
+
+ except Exception as e:
+ import traceback
+ error_details = traceback.format_exc()
+ print(f"Error visualizing structure: {str(e)}\n{error_details}")
+ return f"Error visualizing structure: {str(e)}", None
+
+ # Create the main download tab
+ with gr.Tab("Download"):
+ with gr.Row():
+ with gr.Column():
+ gr.Markdown("### Download Protein Data (See help for more details)")
+
+ # InterPro Metadata tab
+ with gr.Tab("InterPro Metadata"):
+ with gr.Row():
+ interpro_method = gr.Radio(
+ choices=["Single ID", "From JSON"],
+ label="Download Method",
+ value="Single ID"
+ )
+
+ with gr.Column():
+ interpro_id = gr.Textbox(label="InterPro ID", value="IPR000001")
+ interpro_json = gr.Textbox(label="InterPro JSON Path", value="download/interpro_domain/interpro_json.customization", visible=False)
+ interpro_out = gr.Textbox(label="Output Directory", value="download/interpro_domain")
+ interpro_error = gr.Checkbox(label="Save error file", value=True)
+ interpro_btn = gr.Button("Download InterPro Data")
+ interpro_output = gr.Textbox(label="Output", interactive=False)
+
+ def update_interpro_visibility(method):
+ """Update visibility of InterPro input fields based on selected method"""
+ return {
+ interpro_id: gr.update(visible=(method == "Single ID")),
+ interpro_json: gr.update(visible=(method == "From JSON"))
+ }
+
+ interpro_method.change(
+ fn=update_interpro_visibility,
+ inputs=[interpro_method],
+ outputs=[interpro_id, interpro_json]
+ )
+
+ # RCSB Metadata tab
+ with gr.Tab("RCSB Metadata"):
+ with gr.Row():
+ rcsb_method = gr.Radio(
+ choices=["Single ID", "From File"],
+ label="Download Method",
+ value="Single ID"
+ )
+
+ with gr.Column():
+ rcsb_id = gr.Textbox(label="PDB ID", value="1a0j")
+ rcsb_file = gr.Textbox(label="PDB List File", value="download/rcsb.txt", visible=False)
+ rcsb_out = gr.Textbox(label="Output Directory", value="download/rcsb_metadata")
+ rcsb_error = gr.Checkbox(label="Save error file", value=True)
+ rcsb_btn = gr.Button("Download RCSB Metadata")
+ rcsb_output = gr.Textbox(label="Output", interactive=False)
+
+ def update_rcsb_visibility(method):
+ """Update visibility of RCSB input fields based on selected method"""
+ return {
+ rcsb_id: gr.update(visible=(method == "Single ID")),
+ rcsb_file: gr.update(visible=(method == "From File"))
+ }
+
+ rcsb_method.change(
+ fn=update_rcsb_visibility,
+ inputs=[rcsb_method],
+ outputs=[rcsb_id, rcsb_file]
+ )
+
+ # UniProt Sequences tab
+ with gr.Tab("UniProt Sequences"):
+ with gr.Row():
+ uniprot_method = gr.Radio(
+ choices=["Single ID", "From File"],
+ label="Download Method",
+ value="Single ID"
+ )
+
+ with gr.Column():
+ uniprot_id = gr.Textbox(label="UniProt ID", value="P00734")
+ uniprot_file = gr.Textbox(label="UniProt ID List File", value="download/uniprot.txt", visible=False)
+ uniprot_out = gr.Textbox(label="Output Directory", value="download/uniprot_sequences")
+ uniprot_merge = gr.Checkbox(label="Merge into single FASTA", value=False)
+ uniprot_error = gr.Checkbox(label="Save error file", value=True)
+ uniprot_btn = gr.Button("Download UniProt Sequences")
+ uniprot_output = gr.Textbox(label="Output", interactive=False)
+
+ def update_uniprot_visibility(method):
+ """Update visibility of UniProt input fields based on selected method"""
+ return {
+ uniprot_id: gr.update(visible=(method == "Single ID")),
+ uniprot_file: gr.update(visible=(method == "From File"))
+ }
+
+ uniprot_method.change(
+ fn=update_uniprot_visibility,
+ inputs=[uniprot_method],
+ outputs=[uniprot_id, uniprot_file]
+ )
+
+ # RCSB Structures tab
+ with gr.Tab("RCSB Structures"):
+ with gr.Row():
+ # Left column for inputs
+ with gr.Column(scale=3):
+ with gr.Group(): # Group for better visual separation
+ struct_method = gr.Radio(
+ choices=["Single ID", "From File"],
+ label="Download Method",
+ value="Single ID"
+ )
+
+ # Input parameters section with consistent spacing
+ with gr.Row():
+ struct_id = gr.Textbox(label="PDB ID", value="1a0j")
+
+ with gr.Row():
+ struct_file = gr.Textbox(label="PDB List File", value="download/rcsb.txt", visible=False)
+
+ with gr.Row():
+ struct_out = gr.Textbox(label="Output Directory", value="download/rcsb_structures")
+
+ with gr.Row():
+ struct_type = gr.Dropdown(
+ choices=["cif", "pdb", "pdb1", "xml", "sf", "mr", "mrstr"],
+ value="pdb",
+ label="Structure Type"
+ )
+
+ with gr.Row():
+ with gr.Column(scale=1):
+ struct_unzip = gr.Checkbox(label="Unzip downloaded files", value=True)
+ with gr.Column(scale=1):
+ struct_error = gr.Checkbox(label="Save error file", value=True)
+
+ with gr.Row():
+ struct_btn = gr.Button("Download RCSB Structures", size="lg")
+
+ # Output section
+ struct_output = gr.Textbox(label="Download Output", interactive=False, lines=4)
+ struct_viz_status = gr.Textbox(label="Visualization Status", interactive=False)
+
+ # Right column for visualization
+ with gr.Column(scale=5):
+ # Visualization section with full height
+ struct_viz = gr.Plot(label="Structure Visualization", elem_id="struct_viz_plot")
+
+ def update_struct_visibility(method):
+ """Update visibility of RCSB structure input fields based on selected method"""
+ return {
+ struct_id: gr.update(visible=(method == "Single ID")),
+ struct_file: gr.update(visible=(method == "From File"))
+ }
+
+ struct_method.change(
+ fn=update_struct_visibility,
+ inputs=[struct_method],
+ outputs=[struct_id, struct_file]
+ )
+
+ # AlphaFold2 Structures tab
+ with gr.Tab("AlphaFold2 Structures"):
+ with gr.Row():
+ # Left column for inputs
+ with gr.Column(scale=3):
+ with gr.Group(): # Group for better visual separation
+ af_method = gr.Radio(
+ choices=["Single ID", "From File"],
+ label="Download Method",
+ value="Single ID"
+ )
+
+ # Input parameters section with consistent spacing
+ with gr.Row():
+ af_id = gr.Textbox(label="UniProt ID", value="P00734")
+
+ with gr.Row():
+ af_file = gr.Textbox(label="UniProt ID List File", value="download/uniprot.txt", visible=False)
+
+ with gr.Row():
+ af_out = gr.Textbox(label="Output Directory", value="download/alphafold2_structures")
+
+ with gr.Row():
+ af_index_level = gr.Number(label="Index Level", value=0, precision=0)
+
+ with gr.Row():
+ af_error = gr.Checkbox(label="Save error file", value=True)
+
+ with gr.Row():
+ af_btn = gr.Button("Download AlphaFold Structures", size="lg")
+
+ # Output section
+ af_output = gr.Textbox(label="Download Output", interactive=False, lines=4)
+ af_viz_status = gr.Textbox(label="Visualization Status", interactive=False)
+
+ # Right column for visualization
+ with gr.Column(scale=5):
+ # Visualization section with full height
+ af_viz = gr.Plot(label="Structure Visualization", elem_id="af_viz_plot")
+
+ def update_af_visibility(method):
+ """Update visibility of AlphaFold input fields based on selected method"""
+ return {
+ af_id: gr.update(visible=(method == "Single ID")),
+ af_file: gr.update(visible=(method == "From File"))
+ }
+
+ af_method.change(
+ fn=update_af_visibility,
+ inputs=[af_method],
+ outputs=[af_id, af_file]
+ )
+
+ # Handler functions for download buttons
+ def handle_interpro_download(method, id_val, json_val, out_dir, error):
+ """Handle InterPro data download"""
+ if method == "Single ID":
+ return run_download_script(
+ "metadata/download_interpro.py",
+ interpro_id=id_val,
+ out_dir=out_dir,
+ error_file=f"{out_dir}/failed.txt" if error else None
+ )
+ else:
+ return run_download_script(
+ "metadata/download_interpro.py",
+ interpro_json=json_val,
+ out_dir=out_dir,
+ error_file=f"{out_dir}/failed.txt" if error else None
+ )
+
+ interpro_btn.click(
+ fn=handle_interpro_download,
+ inputs=[interpro_method, interpro_id, interpro_json, interpro_out, interpro_error],
+ outputs=interpro_output
+ )
+
+ def handle_rcsb_download(method, id_val, file_val, out_dir, error):
+ """Handle RCSB metadata download"""
+ if method == "Single ID":
+ return run_download_script(
+ "metadata/download_rcsb.py",
+ pdb_id=id_val,
+ out_dir=out_dir,
+ error_file=f"{out_dir}/failed.txt" if error else None
+ )
+ else:
+ return run_download_script(
+ "metadata/download_rcsb.py",
+ pdb_id_file=file_val,
+ out_dir=out_dir,
+ error_file=f"{out_dir}/failed.txt" if error else None
+ )
+
+ rcsb_btn.click(
+ fn=handle_rcsb_download,
+ inputs=[rcsb_method, rcsb_id, rcsb_file, rcsb_out, rcsb_error],
+ outputs=rcsb_output
+ )
+
+ def handle_uniprot_download(method, id_val, file_val, out_dir, merge, error):
+ """Handle UniProt sequence download"""
+ if method == "Single ID":
+ return run_download_script(
+ "sequence/download_uniprot_seq.py",
+ uniprot_id=id_val,
+ out_dir=out_dir,
+ merge="--merge" if merge else None,
+ error_file=f"{out_dir}/failed.txt" if error else None
+ )
+ else:
+ return run_download_script(
+ "sequence/download_uniprot_seq.py",
+ file=file_val,
+ out_dir=out_dir,
+ merge="--merge" if merge else None,
+ error_file=f"{out_dir}/failed.txt" if error else None
+ )
+
+ uniprot_btn.click(
+ fn=handle_uniprot_download,
+ inputs=[uniprot_method, uniprot_id, uniprot_file, uniprot_out, uniprot_merge, uniprot_error],
+ outputs=uniprot_output
+ )
+
+ def handle_struct_download(method, id_val, file_val, out_dir, type_val, unzip, error):
+ """
+ Handle RCSB structure download and visualization
+
+ Args:
+ method: Download method (Single ID or From File)
+ id_val: PDB ID for single download
+ file_val: File path for batch download
+ out_dir: Output directory
+ type_val: Structure file type
+ unzip: Whether to unzip downloaded files
+ error: Whether to save error file
+
+ Returns:
+ Tuple containing download output, visualization status, and Plotly figure
+ """
+ # Download the structure
+ if method == "Single ID":
+ download_output = run_download_script(
+ "structure/download_rcsb.py",
+ pdb_id=id_val,
+ out_dir=out_dir,
+ type=type_val,
+ unzip="--unzip" if unzip else None,
+ error_file=f"{out_dir}/failed.txt" if error else None
+ )
+
+ # Visualize the downloaded structure
+ if "Download completed successfully" in download_output:
+ pdb_file = f"{out_dir}/{id_val.lower()}.{type_val}"
+ if type_val == "pdb" and os.path.exists(pdb_file):
+ viz_status, viz_fig = visualize_protein_structure(pdb_file)
+ return download_output, viz_status, viz_fig
+ else:
+ return download_output, f"Cannot visualize {type_val} format or file not found", None
+ else:
+ return download_output, "Download failed, cannot visualize", None
+ else:
+ download_output = run_download_script(
+ "structure/download_rcsb.py",
+ pdb_id_file=file_val,
+ out_dir=out_dir,
+ type=type_val,
+ unzip="--unzip" if unzip else None,
+ error_file=f"{out_dir}/failed.txt" if error else None
+ )
+ return download_output, "Batch download completed, select a single ID to visualize", None
+
+ struct_btn.click(
+ fn=handle_struct_download,
+ inputs=[struct_method, struct_id, struct_file, struct_out, struct_type, struct_unzip, struct_error],
+ outputs=[struct_output, struct_viz_status, struct_viz]
+ )
+
+ def handle_af_download(method, id_val, file_val, out_dir, index_level, error):
+ """
+ Handle AlphaFold structure download and visualization
+
+ Args:
+ method: Download method (Single ID or From File)
+ id_val: UniProt ID for single download
+ file_val: File path for batch download
+ out_dir: Output directory
+ index_level: Index level for directory structure
+ error: Whether to save error file
+
+ Returns:
+ Tuple containing download output, visualization status, and Plotly figure
+ """
+ # Download the structure
+ if method == "Single ID":
+ download_output = run_download_script(
+ "structure/download_alphafold.py",
+ uniprot_id=id_val,
+ out_dir=out_dir,
+ index_level=index_level,
+ error_file=f"{out_dir}/failed.txt" if error else None
+ )
+
+ # Visualize the downloaded structure
+ if "Download completed successfully" in download_output:
+ # Try different possible file paths
+ possible_paths = [
+ f"{out_dir}/AF-{id_val}-F1-model_v4.pdb",
+ f"{out_dir}/{id_val}.pdb"
+ ]
+
+ for pdb_file in possible_paths:
+ if os.path.exists(pdb_file):
+ viz_status, viz_fig = visualize_protein_structure(pdb_file)
+ return download_output, viz_status, viz_fig
+
+ return download_output, f"PDB file not found in expected locations", None
+ else:
+ return download_output, "Download failed, cannot visualize", None
+ else:
+ download_output = run_download_script(
+ "structure/download_alphafold.py",
+ uniprot_id_file=file_val,
+ out_dir=out_dir,
+ index_level=index_level,
+ error_file=f"{out_dir}/failed.txt" if error else None
+ )
+ return download_output, "Batch download completed, select a single ID to visualize", None
+
+ af_btn.click(
+ fn=handle_af_download,
+ inputs=[af_method, af_id, af_file, af_out, af_index_level, af_error],
+ outputs=[af_output, af_viz_status, af_viz]
+ )
+
+ return {}
diff --git a/src/web/eval_tab.py b/src/web/eval_tab.py
new file mode 100644
index 0000000000000000000000000000000000000000..583a6c2b7ccaffc29fbf98fdcac06546ca693662
--- /dev/null
+++ b/src/web/eval_tab.py
@@ -0,0 +1,1030 @@
+import gradio as gr
+import json
+import os
+import subprocess
+import sys
+import signal
+import threading
+import queue
+import time
+import pandas as pd
+import numpy as np
+import matplotlib.pyplot as plt
+import re
+from datasets import load_dataset
+from web.utils.command import preview_eval_command
+
+def create_eval_tab(constant):
+ plm_models = constant["plm_models"]
+ dataset_configs = constant["dataset_configs"]
+ is_evaluating = False
+ current_process = None
+ output_queue = queue.Queue()
+ stop_thread = False
+ process_aborted = False # 新增标志，表示进程是否被手动终止
+ plm_models = constant["plm_models"]
+
+ def format_metrics(metrics_file):
+ """Convert metrics to HTML table format for display"""
+ try:
+ df = pd.read_csv(metrics_file)
+ metrics_dict = df.iloc[0].to_dict()
+
+ # 定义指标优先级顺序
+ priority_metrics = ['loss', 'accuracy', 'f1', 'precision', 'recall', 'auroc', 'mcc']
+
+ # 构建优先级排序键
+ def get_priority(item):
+ name = item[0]
+ if name in priority_metrics:
+ return priority_metrics.index(name)
+ return len(priority_metrics)
+
+ # 按优先级排序
+ sorted_metrics = sorted(metrics_dict.items(), key=get_priority)
+
+ # 计算指标数量
+ metrics_count = len(sorted_metrics)
+
+ html = f"""
+
+
{metrics_count} metrics found
+
+
+
+ | Metric |
+ Value |
+
+
+
+ """
+
+ # 添加每个指标行，使用交替行颜色
+ for i, (metric_name, metric_value) in enumerate(sorted_metrics):
+ row_style = 'background-color: #f9f9f9;' if i % 2 == 0 else ''
+
+ # 对优先级指标使用粗体
+ is_priority = metric_name in priority_metrics
+ name_style = 'font-weight: bold;' if is_priority else ''
+
+ # 转换指标名称：缩写用大写，非缩写首字母大写
+ display_name = metric_name
+ if metric_name.lower() in ['f1', 'mcc', 'auroc']:
+ display_name = metric_name.upper()
+ else:
+ display_name = metric_name.capitalize()
+
+
+ html += f"""
+
+ | {display_name} |
+ {metric_value:.4f} |
+
+ """
+
+ html += """
+
+
+
Test completed at: """ + time.strftime("%Y-%m-%d %H:%M:%S") + """
+
+
Error: Invalid model path
+
+
Error: Invalid dataset selection
+
+
Error reading output: {str(e)}
+
+
Evaluation completed successfully!
+
Total evaluation time: {time_str}
+
Evaluation dataset: {dataset_display_name}
+
Total samples: {progress_info.get('total_samples', 0)}
+
Evaluation Results
+ {metrics_html}
+ """
+
+ # 设置下载按钮可见并指向结果文件
+ yield summary_html, gr.update(value=result_file, visible=True)
+ else:
+ error_output = "\n".join(progress_info.get("lines", []))
+ yield f"""
+
+
Evaluation completed, but metrics file not found at: {result_file}
+
+
Evaluation failed:
+
{error_output}
+
+
Error during evaluation process:
+
{str(e)}
+
+
+
+ Evaluation Status:
+ {stage}
+
+
+ {progress:.1f}%
+
+
+
+
+
+
+ {f'
Total samples: {total_samples}
' if total_samples > 0 else ''}
+ {f'
Progress: {current}/{total}
' if current > 0 and total > 0 else ''}
+ {f'
Time: {elapsed_time}
' if elapsed_time else ''}
+
+
+
No evaluation in progress to terminate.
+
+
Evaluation successfully terminated!
+
All evaluation state has been reset.
+
+
Failed to terminate evaluation: {str(e)}
+
Evaluation state has been reset.
+
+
+
+ | Dataset |
+ Train Samples |
+ Val Samples |
+ Test Samples |
+
+
+ | {custom_dataset} |
+ {len(dataset["train"]) if "train" in dataset else 0} |
+ {len(dataset["validation"]) if "validation" in dataset else 0} |
+ {len(dataset["test"]) if "test" in dataset else 0} |
+
+
+
+
Error loading dataset
+
{str(e)}
+
+
+
+ | Dataset |
+ Train Samples |
+ Val Samples |
+ Test Samples |
+
+
+ | {config["dataset"]} |
+ {len(dataset["train"]) if "train" in dataset else 0} |
+ {len(dataset["validation"]) if "validation" in dataset else 0} |
+ {len(dataset["test"]) if "test" in dataset else 0} |
+
+
+
+
Error loading dataset
+
{str(e)}
+
Click the 「Start Evaluation」 button to begin model evaluation
'
+
+ # 处理本地图片路径
+ try:
+ # 去掉开头的/以获取正确的路径
+ if img_path.startswith('/'):
+ img_path = img_path[1:]
+
+ # 获取绝对路径
+ current_dir = os.path.dirname(os.path.abspath(__file__))
+ project_root = os.path.dirname(os.path.dirname(current_dir))
+ abs_img_path = os.path.join(project_root, img_path)
+
+ # 读取图片并转换为base64
+ import base64
+ from pathlib import Path
+
+ image_path = Path(abs_img_path)
+ if image_path.exists():
+ image_type = image_path.suffix.lstrip('.').lower()
+ if image_type == 'jpg':
+ image_type = 'jpeg'
+
+ with open(image_path, "rb") as img_file:
+ encoded_string = base64.b64encode(img_file.read()).decode('utf-8')
+
+ return f'
'
+ else:
+ print(f"图片文件不存在: {abs_img_path}")
+ return f'[图片不存在: {img_path}]'
+
+ except Exception as e:
+ print(f"处理图片时出错: {e}, 路径: {img_path}")
+ return f'[图片加载错误: {img_path}]'
+
+ # 使用正则表达式处理所有图片标记
+ pattern = r'!\[(.*?)\]\((.*?)\)'
+ processed_content = re.sub(pattern, embed_image, markdown_content)
+
+ # 使用Python的markdown库进行转换
+ html = markdown.markdown(
+ processed_content,
+ extensions=[
+ 'tables',
+ 'fenced_code',
+ 'codehilite',
+ 'nl2br',
+ 'extra',
+ 'mdx_truly_sane_lists'
+ ],
+ extension_configs={
+ 'mdx_truly_sane_lists': {
+ 'nested_indent': 2,
+ 'truly_sane': True
+ }
+ }
+ )
+
+ return html
+
+ # 从Markdown内容生成HTML导航栏和处理内容
+ def generate_toc_and_content(markdown_content):
+ """从Markdown内容生成HTML导航栏和处理内容"""
+ # 提取所有标题
+ headers = re.findall(r'^(#{1,3})\s+(.+)$', markdown_content, re.MULTILINE)
+
+ if not headers:
+ return "", markdown_content
+
+ toc_html = ""
+
+ # 为每个标题创建导航项
+ for i, (level, title) in enumerate(headers):
+ level_num = len(level)
+ header_id = f"header-{i}"
+
+ # 根据标题级别添加类
+ css_class = ""
+ if level_num == 2:
+ css_class = "nav-h2"
+ elif level_num == 3:
+ css_class = "nav-h3"
+
+ toc_html += f"- {title}
"
+
+ toc_html += "
+ {toc_html}
+
{html_content}
+
+ {toc_html}
+
{html_content}
+
+ {toc_html}
+
{html_content}
+
+ {toc_html}
+
{html_content}
+
+ {toc_html}
+
{html_content}
+
"
+
+ # 为每个标题创建导航项
+ for i, (level, title) in enumerate(headers):
+ level_num = len(level)
+ header_id = f"header-{i}"
+
+ # 根据标题级别添加类
+ css_class = ""
+ if level_num == 2:
+ css_class = "nav-h2"
+ elif level_num == 3:
+ css_class = "nav-h3"
+
+ toc_html += f"- {title}
"
+
+ toc_html += "
+ {training_toc}
+
{training_html}
+
+ {prediction_toc}
+
{prediction_html}
+
+ {evaluation_toc}
+
{evaluation_html}
+
+ {download_toc}
+
{download_html}
+
+ {faq_toc}
+
{faq_html}
+
+
+ {icon}
+
+
{stage}
+
{status.capitalize()}
+
+
+
A prediction is already running. Please wait or abort it.
+
+
❌
+
Please provide a model path
+
+
❌
+
Invalid model path - directory does not exist
+
+
❌
+
Amino acid sequence is required
+
+
❌
+
Error starting prediction process: {str(e)}
+
+
⚠️
+
Warning reading output: {str(e)}
+
+
Prediction was aborted by user
+
+
Regression Prediction Results
+
+
+ | Output | Value |
|---|
+
+
+ | Predicted Value | {prediction_data['prediction']:.4f} |
+
+
+
| Class {i} | {prob:.4f} |
"
+ for i, prob in enumerate(prediction_data['probabilities'])
+ ])
+ elif isinstance(prediction_data.get('probabilities'), dict):
+ prob_rows = "".join([
+ f"| Class {label} | {prob:.4f} |
"
+ for label, prob in prediction_data['probabilities'].items()
+ ])
+ else:
+ # Handle case where probabilities is not a list or dict
+ prob_value = prediction_data.get('probabilities', 0)
+ prob_rows = f"| Class 0 | {prob_value:.4f} |
"
+
+ html_result = f"""
+
+
Single-Label Classification Results
+
+
+ | Output | Value |
|---|
+
+
+ | Predicted Class | {prediction_data['predicted_class']} |
+
+
+
Class Probabilities
+
+
+ | Class | Probability |
|---|
+
+
+ {prob_rows}
+
+
+
| Label {i} | {pred} | {prob:.4f} |
"
+ for i, (pred, prob) in enumerate(zip(prediction_data['predictions'], prediction_data['probabilities']))
+ ])
+ elif (isinstance(prediction_data['predictions'], dict) and
+ isinstance(prediction_data['probabilities'], dict)):
+ pred_rows = "".join([
+ f"| Label {label} | {pred} | {prediction_data['probabilities'].get(label, 0):.4f} |
"
+ for label, pred in prediction_data['predictions'].items()
+ ])
+ else:
+ # Handle case where predictions or probabilities is not a list or dict
+ pred = prediction_data['predictions'] if 'predictions' in prediction_data else "N/A"
+ prob = prediction_data['probabilities'] if 'probabilities' in prediction_data else 0.0
+ pred_rows = f"| Label 0 | {pred} | {prob:.4f} |
"
+ else:
+ # Handle other prediction data formats
+ for key, value in prediction_data.items():
+ if 'label' in key.lower() or 'class' in key.lower():
+ label_name = key
+ label_value = value
+ prob_value = prediction_data.get(f"{key}_prob", 0.0)
+ pred_rows += f"| {label_name} | {label_value} | {prob_value:.4f} |
"
+
+ html_result = f"""
+
+
Multi-Label Classification Results
+
+
+ | Label | Prediction | Probability |
|---|
+
+
+ {pred_rows}
+
+
+
+
Prediction Completed
+
No prediction results found in output.
+
+
+
Prediction Failed
+
Error code: {current_process.returncode if current_process else 'Unknown'}
+
+
{stderr_output}\n{result_output}
+
+
+
Error
+
{str(e)}
+
+
{traceback.format_exc()}
+
+
+
A prediction is already running. Please wait or abort it.
+
+
+
+
+
+ Initializing prediction environment...
+
+
0%
+
+
+
+
+
+
Error: Model path is required
+
+
Error: Invalid model path - directory does not exist
+
+
Error: Input file is required
+
+
Error: Failed to save input file
+
+
Error reading CSV file:
+
{str(e)}
+
+
Error saving input file:
+
{str(e)}
+
+
Error starting prediction process:
+
{str(e)}
+
+
Warning reading output:
+
{str(e)}
+
+
Prediction was manually terminated.
+
All prediction state has been reset.
+
+
Prediction completed successfully!
+
Results saved to: {output_path}
+
Total sequences processed: {progress_info.get('total', 0)}
+
+
+
{len(df)}
+
Predictions
+
+
+
{df['prediction'].mean():.4f}
+
Mean
+
+
+
{df['prediction'].min():.4f}
+
Min
+
+
+
{df['prediction'].max():.4f}
+
Max
+
+
+
{count}
+
Class {class_label}
+
+
+
{len(df)}
+
Predictions
+
+ {class_stats}
+
+
{df[col].sum()}
+
{col}
+
+
+
{len(df)}
+
Predictions
+
+ {label_stats}
+
+
Batch Prediction Results Preview
+ {summary_html}
+
+
+
+
+ {' '.join([f'| {col} | ' for col in df.columns])}
+
+
+
+ {generate_table_rows(df)}
+
+
+
+
+
You can download the complete prediction results using the button below.
+
+
+
Unable to load preview results: {str(e)}
+
You can still download the complete prediction results file.
+
+
Prediction failed to complete
+
Process return code: {current_process.returncode if current_process else 'Unknown'}
+
{result_output}
+
+
Prediction completed, but output file not found at {output_path}
+
{result_output}
+
+
Error during batch prediction: {str(e)}
+
{error_traceback}
+
+
+
+ Prediction Status:
+ {stage}
+
+
+ {percentage:.1f}%
+
+
+
+
+
+
+ {f'
Total sequences: {total}
' if total > 0 else ''}
+ {f'
Progress: {current}/{total}
' if current > 0 and total > 0 else ''}
+ {f'
Status: {progress_info.get("status", "").capitalize()}
' if "status" in progress_info else ''}
+
+
{value[:30]}... | '
+ # Format numeric values to 4 decimal places
+ elif isinstance(value, (int, float)) and not isinstance(value, bool):
+ formatted_value = f"{value:.4f}" if isinstance(value, float) else value
+ cell = f'{formatted_value} | '
+ else:
+ cell = f'{value} | '
+ cells.append(cell)
+
+ # Add alternating row background color
+ bg_color = "#f9f9f9" if i % 2 == 1 else "white"
+ rows.append(f'{" ".join(cells)}
')
+
+ if len(df) > max_rows:
+ cols_count = len(df.columns)
+ rows.append(f'| Showing {max_rows} of {len(df)} rows |
')
+
+ return '\n'.join(rows)
+
+ def handle_abort():
+ """Handle abortion of the prediction process for both single and batch prediction"""
+ nonlocal is_predicting, current_process, stop_thread, process_aborted
+
+ if not is_predicting or current_process is None:
+ empty_html = """
+
+
No prediction process is currently running.
+
+
Prediction successfully terminated!
+
All prediction state has been reset.
+
+
Failed to terminate prediction: {str(e)}
+
Prediction state has been reset.
+
+
No prediction process is currently running.
+
+
Prediction successfully terminated!
+
All prediction state has been reset.
+
+
No prediction process is currently running.
+
+
Prediction successfully terminated!
+
All prediction state has been reset.
+
No file uploaded
"), gr.update(value=[["No file uploaded", "-"]], headers=["name", "sequence"]), gr.update(open=False)
+ try:
+ df = pd.read_csv(file.name)
+ info_html = f"""
+
+
+
+ | File |
+ Total Sequences |
+ Columns |
+
+
+ | {file.name.split('/')[-1]} |
+ {len(df)} |
+ {', '.join(df.columns.tolist())} |
+
+
+
+
Error reading file
+
{str(e)}
+
+
+
+ Training Status:
+ Click Start to train your model
+
+
+
+
+
+ Training Status:
+ Click Start to train your model
+
+
+
+
+
+ Training Status:
+ Failed
+
+
+
+
+
+ Training Status:
+ Waiting to start...
+
+
+
+
Testing in progress, please wait for results...
+
+
Test results will be displayed after testing phase completes
+
+
+
+ Training Status:
+ Training complete!
+
+
+ 100%
+
+
+
+
+
+
+
+ Training Status:
+ {stage} (Epoch {epoch}/{epoch_total})
+
+
+ {progress_percentage:.1f}%
+
+
+
+
+
+
+
Progress: {current}/{total}
+ {f'
Time: {elapsed_time}<{remaining_time}, {it_per_sec:.2f}it/s>
' if elapsed_time and remaining_time else ''}
+ {f'
Loss: {loss:.4f}
' if stage == 'Training' and loss > 0 else ''}
+ {f'
Grad steps: {grad_step}
' if stage == 'Training' and grad_step > 0 else ''}
+
+
+
+
+ Training Status:
+ Initializing training environment...
+
+
+
+
• Parsing configuration parameters
+
• Preparing training environment
+
• This may take a few moments...
+
+
+
+
+ Training Status:
+ Preparing dataset and model...
+
+
+
+
• Loading dataset
+
• Initializing model architecture
+
• Setting up training environment
+
+
+
+
+ Training Status:
+ Starting training process...
+
+
+
+
• Training process launched
+
• Waiting for first statistics to appear
+
• This may take a moment for large models
+
+
+
Training failed with error code {return_code}:
+
{error_output}
+
+
Error during training:
+
{str(e)}
+
{error_output}
+
+
Training was manually terminated.
+
+
Error in final update:
+
{str(e)}
+
{error_output}
+
+
Training initialization failed:
+
{str(e)}
+
+
No training process is currently running.
+
+
Training successfully terminated!
+
All training state has been reset. You can start a new training session.
+
+
Failed to terminate training: {str(e)}
+
Training state has been reset.
+
+
+
+ Training Status:
+ Preparing to start training...
+
+
+
+
+
+ | Dataset |
+ Train Samples |
+ Val Samples |
+ Test Samples |
+
+
+ | {custom_dataset} |
+ {len(dataset["train"]) if "train" in dataset else 0} |
+ {len(dataset["validation"]) if "validation" in dataset else 0} |
+ {len(dataset["test"]) if "test" in dataset else 0} |
+
+
+
+
Error loading dataset
+
{str(e)}
+
+
+
+ | Dataset |
+ Train Samples |
+ Val Samples |
+ Test Samples |
+
+
+ | {config["dataset"]} |
+ {len(dataset["train"]) if "train" in dataset else 0} |
+ {len(dataset["validation"]) if "validation" in dataset else 0} |
+ {len(dataset["test"]) if "test" in dataset else 0} |
+
+
+
+
Error loading dataset
+
{str(e)}
+
+
Test Results
+
{metrics_count} metrics found
+
+
+
+ | Metric |
+ Value |
+
+
+
+ """
+
+ # Sort by priority and alphabetically to ensure important metrics are displayed first
+ priority_metrics = ['loss', 'accuracy', 'f1', 'precision', 'recall', 'auroc', 'mcc']
+
+ # Build priority sorting key
+ def get_priority(item):
+ name = item[0]
+ if name in priority_metrics:
+ return priority_metrics.index(name)
+ return len(priority_metrics)
+
+ # Sort by priority
+ sorted_metrics = sorted(self.test_results.items(), key=get_priority)
+
+ # Add a row for each metric, using alternating row colors
+ for i, (metric_name, metric_value) in enumerate(sorted_metrics):
+ row_style = 'background-color: #f9f9f9;' if i % 2 == 0 else ''
+
+ # Use bold for priority metrics
+ is_priority = metric_name in priority_metrics
+ name_style = 'font-weight: bold;' if is_priority else ''
+
+ # 转换指标名称：缩写用大写，非缩写首字母大写
+ display_name = metric_name
+ if metric_name.lower() in ['f1', 'mcc', 'auroc']:
+ display_name = metric_name.upper()
+ else:
+ display_name = metric_name.capitalize()
+
+ html_content += f"""
+
+ | {display_name} |
+ {metric_value:.4f} |
+
+ """
+
+ html_content += """
+
+
+
Test completed at: """ + time.strftime("%Y-%m-%d %H:%M:%S") + """
+
+
+