import torch import torch.nn as nn import torch.nn.functional as F from src.data.protein_dataset import dynamic_pad from src.data.esm.sdk.api import LogitsConfig import os import pickle from tqdm import tqdm from rdkit import Chem from rdkit.Chem import AllChem import sys; sys.path.append('/nfs_beijing/kubeflow-user/zhangyang_2024/workspace/protein_benchmark/model_zoom') from model_zoom.esm.utils.sampling import _BatchedESMProteinTensor MODEL_ZOOM_PATH = '/nfs_beijing/kubeflow-user/zhangyang_2024/workspace/protein_benchmark/model_zoom' os.environ["TOKENIZERS_PARALLELISM"] = "true" class PretrainModelInterface(nn.Module): """ setup_model: 注册预训练模型 construct_batch: 构建模型的输入batch, 需要进行padding。如果序列加上EOS, BOS, 结构也需要在这里进行padding forward: 调用不同模型提取蛋白质氨基酸级别的embedding, 在后处理阶段去除EOS, BOS embedding，只返回氨基酸embedding """ def __init__(self, pretrain_model_name, batch_size = 64, max_length = 1022, device = 'cuda', sequence_only=False, task_type=None): super(PretrainModelInterface, self).__init__() self.pretrain_model_name = pretrain_model_name self.sequence_only = sequence_only self.batch_size = batch_size self.max_length = max_length self.task_type = task_type self.device = device self.setup_model() def setup_model(self): """ Setup the pre-trained model based on the specified name. """ if self.pretrain_model_name == 'esm2_650m': from transformers import AutoTokenizer, AutoModelForMaskedLM self.pretrain_model = AutoModelForMaskedLM.from_pretrained(f"{MODEL_ZOOM_PATH}/esm2_650m").to(self.device) self.tokenizer = AutoTokenizer.from_pretrained(f"{MODEL_ZOOM_PATH}/esm2_650m") elif self.pretrain_model_name == 'esm3_1.4b': from model_zoom.esm.models.esm3 import ESM3 self.pretrain_model = ESM3.from_pretrained("esm3-sm-open-v1", ).to(self.device) elif self.pretrain_model_name == 'esmc_600m': from model_zoom.esm.models.esmc import ESMC self.pretrain_model = ESMC.from_pretrained("esmc_600m").to(self.device) elif self.pretrain_model_name == 'procyon': from transformers import AutoTokenizer, AutoModelForMaskedLM from model_zoom.procyon.model.model_unified import UnifiedProCyon from model_zoom.GearNet.data.protein import Protein from model_zoom.GearNet.data.transform import ProteinView from model_zoom.GearNet.data.transform import Compose from model_zoom.GearNet.data.geo_graph import GraphConstruction from model_zoom.GearNet.data.function import AlphaCarbonNode, SpatialEdge, KNNEdge, SequentialEdge from model_zoom.GearNet.gearnet import GeometryAwareRelationalGraphNeuralNetwork protein_view_transform = ProteinView(view='residue') self.transform = Compose([protein_view_transform]) self.graph_construction_model = GraphConstruction( node_layers=[AlphaCarbonNode()], edge_layers=[SpatialEdge(radius=10.0, min_distance=5), KNNEdge(k=10, min_distance=5), SequentialEdge(max_distance=2)], edge_feature="gearnet" ) self.gearnet_edge = GeometryAwareRelationalGraphNeuralNetwork(input_dim=21, hidden_dims=[512, 512, 512, 512, 512, 512], num_relation=7, edge_input_dim=59, num_angle_bin=8, batch_norm=True, concat_hidden=True, short_cut=True, readout="sum" ) ckpt = torch.load("/nfs_beijing/kubeflow-user/zhangyang_2024/workspace/protein_benchmark/model_zoom/GearNet/mc_gearnet_edge.pth") self.gearnet_edge.load_state_dict(ckpt) self.gearnet_edge = self.gearnet_edge.to(self.device) self.gearnet_edge.eval() os.environ["HOME_DIR"] = MODEL_ZOOM_PATH os.environ["DATA_DIR"] = "/nfs_beijing/wanghao/2025-onesystem/vllm/ProCyon-Instruct" os.environ["LLAMA3_PATH"] = "/nfs_beijing/wanghao/2025-onesystem/vllm/Meta-Llama-3-8B" procyon_ckpt = '/nfs_beijing/kubeflow-user/zhangyang_2024/workspace/protein_benchmark/model_zoom/procyon/model_weights/ProCyon-Full' self.esm_pretrain_model = AutoModelForMaskedLM.from_pretrained(f"{MODEL_ZOOM_PATH}/esm2_3b").to(self.device) self.esm_pretrain_model.eval() self.esm_tokenizer = AutoTokenizer.from_pretrained(f"{MODEL_ZOOM_PATH}/esm2_3b") # procyon initialization self.pretrain_model, _ = UnifiedProCyon.from_pretrained( pretrained_weights_dir=procyon_ckpt, checkpoint_dir=procyon_ckpt ) self.pretrain_model = self.pretrain_model.to(self.device) elif self.pretrain_model_name == 'gearnet': from model_zoom.GearNet.data.transform import ProteinView from model_zoom.GearNet.data.transform import Compose from model_zoom.GearNet.data.geo_graph import GraphConstruction from model_zoom.GearNet.data.function import AlphaCarbonNode, SpatialEdge, KNNEdge, SequentialEdge from model_zoom.GearNet.gearnet import GeometryAwareRelationalGraphNeuralNetwork protein_view_transform = ProteinView(view='residue') self.transform = Compose([protein_view_transform]) self.graph_construction_model = GraphConstruction( node_layers=[AlphaCarbonNode()], edge_layers=[SpatialEdge(radius=10.0, min_distance=5), KNNEdge(k=10, min_distance=5), SequentialEdge(max_distance=2)], edge_feature="gearnet" ) self.gearnet_edge = GeometryAwareRelationalGraphNeuralNetwork(input_dim=21, hidden_dims=[512, 512, 512, 512, 512, 512], num_relation=7, edge_input_dim=59, num_angle_bin=8, batch_norm=True, concat_hidden=True, short_cut=True, readout="sum" ) ckpt = torch.load("/nfs_beijing/kubeflow-user/zhangyang_2024/workspace/protein_benchmark/model_zoom/GearNet/mc_gearnet_edge.pth") self.gearnet_edge.load_state_dict(ckpt) self.pretrain_model = self.gearnet_edge.to(self.device) self.pretrain_model.eval() elif self.pretrain_model_name == 'prollama': from transformers import LlamaForCausalLM, LlamaTokenizer llama_path = "/nfs_beijing/kubeflow-user/wanghao/workspace/ai4sci/protein_benchmark_project/data/ProLLaMA" self.pretrain_model = LlamaForCausalLM.from_pretrained( llama_path, quantization_config=None ).to(self.device) self.tokenizer = LlamaTokenizer.from_pretrained(llama_path) elif self.pretrain_model_name == 'prost': from transformers import AutoModel, AutoTokenizer, AutoConfig prost_weights = "/nfs_beijing/kubeflow-user/zhangyang_2024/workspace/protein_benchmark/model_zoom/protst" prost_tokenizer_weights = "/nfs_beijing/kubeflow-user/zhangyang_2024/workspace/protein_benchmark/model_zoom/esm1b_650m" protst_model = AutoModel.from_pretrained( prost_weights, trust_remote_code=True, torch_dtype=torch.bfloat16 ) self.pretrain_model = protst_model.protein_model.to(self.device) self.tokenizer = AutoTokenizer.from_pretrained(prost_tokenizer_weights) elif self.pretrain_model_name == 'progen2': from model_zoom.progen2.modeling_progen import ProGenForCausalLM from tokenizers import Tokenizer def create_tokenizer_custom(file): with open(file, 'r') as f: return Tokenizer.from_str(f.read()) self.pretrain_model = ProGenForCausalLM.from_pretrained(f'{MODEL_ZOOM_PATH}/progen2').to(self.device) self.tokenizer = create_tokenizer_custom(file=f'{MODEL_ZOOM_PATH}/progen2/tokenizer.json') elif self.pretrain_model_name == 'prostt5': from transformers import T5Tokenizer, T5EncoderModel import mini3di self.tokenizer = T5Tokenizer.from_pretrained(f'{MODEL_ZOOM_PATH}/ProstT5', do_lower_case=False) self.pretrain_model = T5EncoderModel.from_pretrained(f"{MODEL_ZOOM_PATH}/ProstT5").to(self.device) self.encoder_3di = mini3di.Encoder() elif self.pretrain_model_name == 'protgpt2': from transformers import AutoTokenizer, AutoModelForCausalLM self.tokenizer = AutoTokenizer.from_pretrained(f"{MODEL_ZOOM_PATH}/ProtGPT2") self.pretrain_model = AutoModelForCausalLM.from_pretrained(f"{MODEL_ZOOM_PATH}/ProtGPT2").to(self.device) elif self.pretrain_model_name == 'protrek': from model_zoom.ProTrek.model.ProTrek.protrek_trimodal_model import ProTrekTrimodalModel import mini3di self.encoder_3di = mini3di.Encoder() config = { "protein_config": f"{MODEL_ZOOM_PATH}/ProTrek/ProTrek_650M_UniRef50/esm2_t33_650M_UR50D", "text_config": f"{MODEL_ZOOM_PATH}/ProTrek/ProTrek_650M_UniRef50/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext", "structure_config": f"{MODEL_ZOOM_PATH}/ProTrek/ProTrek_650M_UniRef50/foldseek_t30_150M", "load_protein_pretrained": False, "load_text_pretrained": False, "from_checkpoint": f"{MODEL_ZOOM_PATH}/ProTrek/ProTrek_650M_UniRef50/ProTrek_650M_UniRef50.pt" } self.pretrain_model = ProTrekTrimodalModel(**config).eval().to(self.device) elif self.pretrain_model_name == 'saport': from transformers import EsmTokenizer, EsmForMaskedLM import mini3di self.encoder_3di = mini3di.Encoder() self.tokenizer = EsmTokenizer.from_pretrained(f'{MODEL_ZOOM_PATH}/SaPort/ckpt') self.pretrain_model = EsmForMaskedLM.from_pretrained(f'{MODEL_ZOOM_PATH}/SaPort/ckpt').to(self.device) elif self.pretrain_model_name == 'venusplm': from vplm import TransformerForMaskedLM, TransformerConfig from vplm import VPLMTokenizer venusplm_weight_path = "/nfs_beijing/kubeflow-user/zhangyang_2024/workspace/protein_benchmark/model_zoom/venusplm" config = TransformerConfig.from_pretrained(venusplm_weight_path, attn_impl="sdpa") # or "flash_attn" if you have installed flash-attn self.pretrain_model = TransformerForMaskedLM.from_pretrained(venusplm_weight_path, config=config).to(self.device) self.pretrain_model.eval() self.tokenizer = VPLMTokenizer.from_pretrained(venusplm_weight_path) elif self.pretrain_model_name == 'prosst2048': from transformers import AutoTokenizer, AutoModelForMaskedLM from model_zoom.ProSST.prosst.structure.quantizer import PdbQuantizer prosst_2048_weight_path = "/nfs_beijing/kubeflow-user/zhangyang_2024/workspace/protein_benchmark/model_zoom/ProSST/prosst_2048_weight" self.tokenizer = tokenizer = AutoTokenizer.from_pretrained(prosst_2048_weight_path, trust_remote_code=True) self.ss_processor = PdbQuantizer(structure_vocab_size=2048) self.pretrain_model = AutoModelForMaskedLM.from_pretrained(prosst_2048_weight_path, trust_remote_code=True).to(self.device) self.pretrain_model.eval() else: raise ValueError(f"Unknown pretrain model name: {self.pretrain_model_name}") def construct_batch(self, data, batch_size, task_name=None): """ Constructs batches of data. Args: data (list): List of data samples. batch_size (int): Size of each batch. Yields: dict: A batch of data. """ MAXLEN = 1022 for i in range(0, len(data), batch_size): if self.pretrain_model_name == 'esm2_650m': add_BOS = 1 name_batch, label_batch, smiles_batch = [], [], [] if "pair" in self.task_type: current_batch = [sample["seq"] for sample in data[i:i + batch_size]] max_length_batch = [max(len(s) for s in column)+2 for column in zip(*current_batch)] else: max_length_batch = [max([len(sample['seq']) for sample in data[i:i + batch_size]]) + 2] X_batch, S_batch, mask_batch, t_lengths = [[] for _ in range(len(max_length_batch))], \ [[] for _ in range(len(max_length_batch))], \ [[] for _ in range(len(max_length_batch))], \ [[] for _ in range(len(max_length_batch))] for sample in data[i:i + batch_size]: seq = sample['seq'] label = sample['label'] smiles = sample['smiles'] if 'smiles' in sample else None if smiles: mol = Chem.MolFromSmiles(smiles) fp = AllChem.GetMorganFingerprintAsBitVect(mol, radius=2, nBits=2048) smiles = torch.tensor([int(ele) for ele in list(fp.ToBitString())]).float() if not isinstance(seq, list): seq = [seq] for j, _seq in enumerate(seq): _seq = _seq[:MAXLEN] t_lengths[j].append(len(_seq)) attention_mask = torch.zeros(max_length_batch[j]) seq_token = torch.tensor(self.tokenizer.encode(_seq)) attention_mask[:len(seq_token)] = 1 seq_token = self.pad_data(seq_token, dim=0, max_length=max_length_batch[j]) # _X = dynamic_pad(_X, [add_BOS, add_BOS], dim=0, pad_value=0) # 坐标需要和seq一样加上BOS, EOS # _X = self.pad_data(_X, dim=0, max_length=max_length_batch[j]) S_batch[j].append(seq_token) # X_batch[j].append(_X) mask_batch[j].append(attention_mask) if task_name == 'contact_map': label = F.pad(label, [add_BOS, self.max_length-label.shape[0]-add_BOS, add_BOS, self.max_length-label.shape[0]-add_BOS]) else: label = torch.tensor(label) label_batch.append(label) name_batch.append(sample['name']) if smiles is not None: smiles_batch.append(smiles) S_batch = [torch.stack(ele).to(self.device) for ele in S_batch] # X_batch = [torch.stack(ele).to(self.device) for ele in X_batch] mask_batch = [torch.stack(ele).to(self.device)==1 for ele in mask_batch] t_lengths = [torch.tensor(ele).to(self.device) for ele in t_lengths] smiles_batch = None if len(smiles_batch) == 0 else torch.stack(smiles_batch) yield { 'name': name_batch, 'seq': S_batch, # 'X': X_batch, 't_lengths': t_lengths, 'smiles': smiles_batch, 'attention_mask': mask_batch, 'label': torch.stack(label_batch), } if self.pretrain_model_name == 'esm3_1.4b': addBOS = 1 from model_zoom.esm.utils.misc import stack_variable_length_tensors from model_zoom.esm.utils.sampling import _BatchedESMProteinTensor from model_zoom.esm.utils import encoding model = self.pretrain_model name_batch, sequence_list, coordinates_list, label_batch, structure_tokens_batch, mask_batch = [], [], [], [], [], [] seq_tokenizer = model.tokenizers.sequence struct_tokenizer = model.tokenizers.structure pad = model.tokenizers.sequence.pad_token_id if "pair" in self.task_type: max_length_batch = max( [max(len(seq) for seq in sample['seq']) for sample in data[i:i + batch_size]] ) + addBOS*2 else: max_length_batch = max([len(sample['seq']) for sample in data[i:i + batch_size]])+addBOS*2 for sample in data[i:i + batch_size]: name, label, seq, X = sample['name'], sample['label'], sample['seq'], sample['X'] sequence_tokens, structure_tokens, coordinates, masks = [], [], [], [] for _seq, _X in zip(seq, X): _seq, _X = _seq[:MAXLEN], _X[:MAXLEN] _seq_token = encoding.tokenize_sequence(_seq, seq_tokenizer, add_special_tokens=True) _coordinates, _plddt, _structure_token = encoding.tokenize_structure( _X, model.get_structure_encoder(), struct_tokenizer, add_special_tokens=True ) mask = torch.zeros(max_length_batch) mask[:_coordinates.shape[0]] = 1 # pad _seq_token = self.pad_data(_seq_token, dim=0, pad_value=pad, max_length=max_length_batch) _structure_token = self.pad_data(_structure_token, dim=0, pad_value=pad, max_length=max_length_batch) _coordinates = dynamic_pad(_coordinates, [addBOS, addBOS], dim=0, pad_value=0) # 坐标需要和seq一样加上BOS, EOS _coordinates = self.pad_data(_coordinates, dim=0, max_length=max_length_batch) sequence_tokens.append(_seq_token) structure_tokens.append(_structure_token) coordinates.append(_coordinates) masks.append(mask) sequence_tokens = torch.hstack(sequence_tokens) structure_tokens = torch.hstack(structure_tokens) coordinates = torch.vstack(coordinates) masks = torch.hstack(masks) sequence_list.append(sequence_tokens) coordinates_list.append(coordinates) structure_tokens_batch.append(structure_tokens) mask_batch.append(masks) name_batch.append(sample['name']) label = sample['label'] if task_name == 'contact_map': label = F.pad(label, [add_BOS, self.max_length-label.shape[0]-add_BOS, add_BOS, self.max_length-label.shape[0]-add_BOS]) else: label = torch.tensor(label) label_batch.append(label) sequence_tokens = stack_variable_length_tensors( sequence_list, constant_value=pad, ).to(self.device) structure_tokens_batch = stack_variable_length_tensors( structure_tokens_batch, constant_value=pad, ).to(self.device) coordinates_batch = stack_variable_length_tensors( coordinates_list, constant_value=pad, ).to(self.device) if self.sequence_only: protein_tensor = _BatchedESMProteinTensor(sequence=sequence_tokens, coordinates=coordinates_batch).to(self.device) else: protein_tensor = _BatchedESMProteinTensor(sequence=sequence_tokens, structure=structure_tokens_batch, coordinates=coordinates_batch).to(self.device) yield { 'name': name_batch, 'protein_tensor': protein_tensor, 'label': torch.stack(label_batch).to(self.device), 'attention_mask': torch.stack(mask_batch).to(self.device)==1 } if self.pretrain_model_name == 'esmc_600m': addBOS = 1 from model_zoom.esm.utils.sampling import _BatchedESMProteinTensor from model_zoom.esm.utils.misc import stack_variable_length_tensors name_batch, seq_batch, mask_batch, label_batch = [], [], [], [] if "pair" in self.task_type: max_length_batch = max( [max(len(seq) for seq in sample['seq']) for sample in data[i:i + batch_size]] ) + addBOS*2 else: max_length_batch = max([len(sample['seq']) for sample in data[i:i + batch_size]])+addBOS*2 for sample in data[i:i + batch_size]: seq = sample['seq'] label = sample['label'] if not isinstance(seq, list): seq = [seq] seq_tokens, attention_masks = [], [] for _seq in seq: attention_mask = torch.zeros(max_length_batch) attention_mask[:len(_seq)] = 1 _seq = _seq[:max_length_batch] _seq_token = self.pretrain_model._tokenize([_seq]).flatten() pad = self.pretrain_model.tokenizer.pad_token_id _seq_token = self.pad_data(_seq_token, dim=0, pad_value=pad, max_length=max_length_batch) seq_tokens.append(_seq_token) attention_masks.append(attention_mask) seq_tokens = torch.hstack(seq_tokens) attention_masks = torch.hstack(attention_masks) seq_batch.append(seq_tokens) mask_batch.append(attention_masks) if task_name == 'contact_map': label = F.pad(label, [add_BOS, self.max_length-label.shape[0]-add_BOS, add_BOS, self.max_length-label.shape[0]-add_BOS]) else: label = torch.tensor(label) label_batch.append(label) name_batch.append(sample['name']) sequence_tokens = stack_variable_length_tensors( seq_batch, constant_value=self.pretrain_model.tokenizer.pad_token_id, ) # sequence_tokens = self.pretrain_model._tokenize(seq_batch) protein_tensor = _BatchedESMProteinTensor(sequence=sequence_tokens).to(self.device) yield { 'name': name_batch, 'protein_tensor': protein_tensor, 'attention_mask': torch.stack(mask_batch).to(self.device)==1, 'label': torch.stack(label_batch).to(self.device), } if self.pretrain_model_name == 'procyon': add_BOS = 0 from model_zoom.GearNet.data.protein import Protein name_batch, X_batch, S_batch, label_batch = [], [], [], [] if "pair" in self.task_type: max_length_batch = max( [max(len(seq) for seq in sample['seq']) for sample in data[i:i + batch_size]] ) + 2 else: max_length_batch = max([len(sample['seq']) for sample in data[i:i + batch_size]])+2 for sample in data[i:i + batch_size]: try: seq = sample['seq'] label = sample['label'] pdb_path = sample['pdb_path'] if not isinstance(seq, list) and not isinstance(pdb_path, list): seq, pdb_path = [seq], [pdb_path] seq_embeddings, struct_embeddings = [], [] for _seq, _pdb_path in zip(seq, pdb_path): # seq_token = torch.tensor(self.esm_tokenizer.encode(_seq)) seq_token = self.esm_tokenizer( [_seq], return_tensors="pt", padding=True, max_length=max_length_batch, truncation=True ) seq_embedding = self.esm_pretrain_model.esm( seq_token['input_ids'].to(self.device), attention_mask=seq_token['attention_mask'].to(self.device), return_dict=True, ).last_hidden_state.squeeze(0).mean(0).flatten() seq_embeddings.append(seq_embedding) # strcuture protein = Protein.from_pdb(_pdb_path, bond_feature="length", residue_feature="symbol") protein = self.transform({"graph": protein})["graph"] _protein = Protein.pack([protein]) protein_ = self.graph_construction_model(_protein).to(self.device) with torch.no_grad(): out = self.gearnet_edge(protein_, protein_.node_feature.float()) struct_embeddings.append(out["graph_feature"].flatten()) seq_embeddings = torch.cat(seq_embeddings, dim=-1) struct_embeddings = torch.cat(struct_embeddings, dim=-1) if task_name == 'contact_map': label = F.pad(label, [add_BOS, self.max_length-label.shape[0]-add_BOS, add_BOS, self.max_length-label.shape[0]-add_BOS]) else: label = torch.tensor(label) S_batch.append(seq_embeddings) X_batch.append(struct_embeddings) label_batch.append(label) name_batch.append(sample['name']) except: print(f"Error processing sample {sample['name']}") continue yield { 'name': name_batch, 'seq': torch.stack(S_batch).to(self.device), 'X': torch.stack(X_batch).unsqueeze(1).to(self.device), 'label': torch.stack(label_batch).to(self.device), } if self.pretrain_model_name == 'gearnet': add_BOS = 0 from model_zoom.GearNet.data.protein import Protein name_batch, X_batch, S_batch, label_batch = [], [], [], [] # proteins = [] proteins_pair_1, proteins_pair_2 = [], [] for sample in data[i:i + batch_size]: try: label = sample['label'] pdb_path = sample['pdb_path'] if not isinstance(pdb_path, list): pdb_path = [pdb_path] temp_proteins = [] for j, _pdb_path in enumerate(pdb_path): protein = Protein.from_pdb(_pdb_path, bond_feature="length", residue_feature="symbol") protein = self.transform({"graph": protein})["graph"] temp_proteins.append(protein) # sub_proteins.append(protein) proteins_pair_1.append(temp_proteins[0]) proteins_pair_2.append(temp_proteins[1]) if task_name == 'contact_map': label = F.pad(label, [add_BOS, self.max_length-label.shape[0]-add_BOS, add_BOS, self.max_length-label.shape[0]-add_BOS]) else: label = torch.tensor(label) label_batch.append(label) name_batch.append(sample['name']) except: print(f"Error processing sample {sample['name']}") continue proteins_pairs = [proteins_pair_1, proteins_pair_2] embeddings, attention_masks = [], [] if len(proteins_pair_2) == 0: max_length_batch = Protein.pack(proteins_pair_1).num_residues.max() else: max_length_batch = max([Protein.pack(proteins).num_residues.max() for proteins in proteins_pairs]) for proteins in proteins_pairs: X_batch = [] if len(proteins_pairs) == 0: continue _protein = Protein.pack(proteins) protein_ = self.graph_construction_model(_protein).to(self.device) with torch.no_grad(): out = self.gearnet_edge(protein_, protein_.node_feature.float()) node_feature = out["node_feature"] # max_length_batch = _protein.num_residues.max() split = torch.cumsum(F.pad(_protein.num_residues, (1,0)), dim=0) attention_mask = torch.zeros(len(split)-1, max_length_batch).to(self.device) for i in range(len(split)-1): start, end = split[i], split[i+1] embedding = node_feature[start:end] attention_mask[i, :embedding.shape[0]] = 1 embedding = self.pad_data(embedding, dim=0, max_length=max_length_batch) X_batch.append(embedding) X_batch = torch.stack(X_batch) embeddings.append(X_batch) attention_masks.append(attention_mask) embeddings = torch.cat(embeddings, dim=-1) if "pair" in self.task_type: attention_mask = (attention_masks[0].int() + attention_masks[1].int() > 0) else: attention_mask = attention_masks[0] yield { 'name': name_batch, 'attention_mask': attention_mask, 'X': embeddings.to(self.device), 'label': torch.stack(label_batch).to(self.device), } if self.pretrain_model_name == 'prollama': name_batch, S_batch, mask_batch, label_batch = [], [], [], [] if "pair" in self.task_type: max_length_batch = max( [max(len(seq) for seq in sample['seq']) for sample in data[i:i + batch_size]] ) + 2 else: max_length_batch = max([len(sample['seq']) for sample in data[i:i + batch_size]])+2 for sample in data[i:i + batch_size]: seq = sample['seq'] label = sample['label'] if not isinstance(seq, list): seq = [seq] seq = [f"[Determine superfamily] Seq=<{_seq}>" for _seq in seq] seq_tokens, attention_masks = [], [] for _seq in seq: attention_mask = torch.zeros(max_length_batch) seq_token = torch.tensor(self.tokenizer.encode(_seq)) attention_mask[:len(seq_token)] = 1 seq_token = self.pad_data(seq_token, dim=0, max_length=max_length_batch) seq_tokens.append(seq_token) attention_masks.append(attention_mask) seq_tokens = torch.hstack(seq_tokens) attention_masks = torch.hstack(attention_masks) S_batch.append(seq_tokens) mask_batch.append(attention_masks) if task_name == 'contact_map': label = F.pad(label, [add_BOS, self.max_length-label.shape[0]-add_BOS, add_BOS, self.max_length-label.shape[0]-add_BOS]) else: label = torch.tensor(label) label_batch.append(label) name_batch.append(sample['name']) yield { 'name': name_batch, 'seq': torch.stack(S_batch).to(self.device), 'attention_mask': torch.stack(mask_batch).to(self.device)==1, 'label': torch.stack(label_batch).to(self.device), } if self.pretrain_model_name == "venusplm": name_batch, X_batch, S_batch, mask_batch, label_batch = [], [], [], [], [] if "pair" in self.task_type: max_length_batch = max( [max(len(seq) for seq in sample['seq']) for sample in data[i:i + batch_size]] ) + 2 else: max_length_batch = max([len(sample['seq']) for sample in data[i:i + batch_size]]) + 2 for sample in data[i:i + batch_size]: seq = sample['seq'] label = sample['label'] if not isinstance(seq, list): seq = [seq] seq_tokens, attention_masks = [], [] for _seq in seq: attention_mask = torch.zeros(max_length_batch) seq_token = torch.tensor(self.tokenizer.encode(_seq)) attention_mask[:len(seq_token)] = 1 seq_token = self.pad_data(seq_token, dim=0, max_length=max_length_batch) seq_tokens.append(seq_token) attention_masks.append(attention_mask) seq_tokens = torch.hstack(seq_tokens) attention_masks = torch.hstack(attention_masks) S_batch.append(seq_tokens) mask_batch.append(attention_masks) if task_name == 'contact_map': label = F.pad(label, [add_BOS, self.max_length-label.shape[0]-add_BOS, add_BOS, self.max_length-label.shape[0]-add_BOS]) else: label = torch.tensor(label) label_batch.append(label) name_batch.append(sample['name']) yield { 'name': name_batch, 'seq': torch.stack(S_batch).to(self.device), 'attention_mask': torch.stack(mask_batch).to(self.device)==1, 'label': torch.stack(label_batch).to(self.device), } if self.pretrain_model_name == "prosst2048": def tokenize_structure_sequence(structure_sequence): shift_structure_sequence = [i + 3 for i in structure_sequence] shift_structure_sequence = [1, *shift_structure_sequence, 2] return torch.tensor( # [ shift_structure_sequence, # ], dtype=torch.long, ) name_batch, X_batch, S_batch, mask_batch, label_batch = [], [], [], [], [] max_length_batch = max([len(sample['seq']) for sample in data[i:i + batch_size]])+2 for sample in data[i:i + batch_size]: seq = sample['seq'] label = sample['label'] X = sample['pdb_path'] attention_mask = torch.zeros(max_length_batch) seq_token = torch.tensor(self.tokenizer.encode(seq)) attention_mask[:len(seq_token)] = 1 seq_token = self.pad_data(seq_token, dim=0, max_length=max_length_batch) X = self.ss_processor(X, return_residue_seq=False)['2048']['ranked_unrelax_0.pdb']["struct"] X = tokenize_structure_sequence(X) X = self.pad_data(X, dim=0, max_length=max_length_batch) S_batch.append(seq_token) X_batch.append(X) mask_batch.append(attention_mask) if task_name == 'contact_map': label = F.pad(label, [add_BOS, self.max_length-label.shape[0]-add_BOS, add_BOS, self.max_length-label.shape[0]-add_BOS]) else: label = torch.tensor(label) label_batch.append(label) name_batch.append(sample['name']) yield { 'name': name_batch, 'seq': torch.stack(S_batch).to(self.device), 'X': torch.stack(X_batch).to(self.device), 'attention_mask': torch.stack(mask_batch).to(self.device)==1, 'label': torch.stack(label_batch).to(self.device), } if self.pretrain_model_name == "prost": name_batch, X_batch, S_batch, mask_batch, label_batch = [], [], [], [], [] if "pair" in self.task_type: max_length_batch = max( [max(len(seq) for seq in sample['seq']) for sample in data[i:i + batch_size]] ) + 2 else: max_length_batch = max([len(sample['seq']) for sample in data[i:i + batch_size]]) + 2 max_length_batch = self.max_length for sample in data[i:i + batch_size]: seq = sample['seq'] label = sample['label'] if not isinstance(seq, list): seq = [seq] seq_tokens, attention_masks = [], [] for _seq in seq: _seq = _seq[:1022] attention_mask = torch.zeros(max_length_batch) seq_token = torch.tensor(self.tokenizer.encode(_seq)) attention_mask[:len(seq_token)] = 1 seq_token = self.pad_data(seq_token, dim=0, max_length=max_length_batch) seq_tokens.append(seq_token) attention_masks.append(attention_mask) seq_tokens = torch.hstack(seq_tokens) attention_masks = torch.hstack(attention_masks) S_batch.append(seq_tokens) mask_batch.append(attention_masks) if task_name == 'contact_map': label = F.pad(label, [add_BOS, self.max_length-label.shape[0]-add_BOS, add_BOS, self.max_length-label.shape[0]-add_BOS]) else: label = torch.tensor(label) label_batch.append(label) name_batch.append(sample['name']) yield { 'name': name_batch, 'seq': torch.stack(S_batch).to(self.device), 'attention_mask': torch.stack(mask_batch).to(self.device)==1, 'label': torch.stack(label_batch).to(self.device), } if self.pretrain_model_name == 'progen2': # 没有BOS, EOS, 需要限制长度在1024以内 if "pair" in self.task_type: max_length_batch = max( [max(len(seq) for seq in sample['seq']) for sample in data[i:i + batch_size]] ) else: max_length_batch = max([len(sample['seq']) for sample in data[i:i + batch_size]]) name_batch, S_batch, mask_batch, label_batch = [], [], [], [] for sample in data[i:i + batch_size]: seq = sample['seq'] label = sample['label'] if not isinstance(seq, list): seq = [seq] seq_tokens, attention_masks = [], [] for _seq in seq: attention_mask = torch.zeros(self.max_length) seq_token = torch.tensor(self.tokenizer.encode(_seq).ids) attention_mask[:len(seq_token)] = 1 seq_token = self.pad_data(seq_token, dim=0, max_length=max_length_batch) attention_mask = self.pad_data(attention_mask, dim=0, max_length=max_length_batch) seq_tokens.append(seq_token) attention_masks.append(attention_mask) seq_tokens = torch.hstack(seq_tokens) attention_masks = torch.hstack(attention_masks) S_batch.append(seq_tokens[:1024]) mask_batch.append(attention_masks[:1024]) if task_name == 'contact_map': label = F.pad(label, [add_BOS, self.max_length-label.shape[0]-add_BOS, add_BOS, self.max_length-label.shape[0]-add_BOS]) else: label = torch.tensor(label) label_batch.append(label) name_batch.append(sample['name']) yield { 'name': name_batch, 'seq': torch.stack(S_batch).to(self.device), 'attention_mask': torch.stack(mask_batch).to(self.device)==1, 'label': torch.stack(label_batch).to(self.device), } if self.pretrain_model_name == 'prostt5': import re if "pair" in self.task_type: max_length_batch = max( [max(len(seq) for seq in sample['seq']) for sample in data[i:i + batch_size]] ) + 2 else: max_length_batch = max([len(sample['seq']) for sample in data[i:i + batch_size]]) + 2 name_batch, struct_batch, S_batch, mask_batch, label_batch = [], [], [], [], [] for sample in data[i:i + batch_size]: seq = sample['seq'] X = sample['X'] label = sample['label'] if not isinstance(seq, list) and not isinstance(X, list): seq, X = [seq], [X] seq_tokens, attention_masks = [], [] for _seq, _X in zip(seq, X): N, CA, C, CB, O = _X[:, 0], _X[:, 1], _X[:, 2], _X[:, 3], _X[:, 4] attention_mask = torch.zeros(2, max_length_batch) states = self.encoder_3di.encode_atoms(ca = CA.numpy(), cb = CB.numpy(), n = N.numpy(), c = C.numpy()) struct_sequence = self.encoder_3di.build_sequence(states).lower() if self.sequence_only: sequence_examples = [_seq, _seq] sequence_examples = [" ".join(list(re.sub(r"[UZOB]", "X", sequence))) for sequence in sequence_examples] sequence_examples = [ "" + " " + s if s.isupper() else "" + " " + s # this expects 3Di sequences to be already lower-case for s in sequence_examples ] else: sequence_examples = [_seq, struct_sequence] sequence_examples = [" ".join(list(re.sub(r"[UZOB]", "X", sequence))) for sequence in sequence_examples] sequence_examples = [ "" + " " + s if s.isupper() else "" + " " + s # this expects 3Di sequences to be already lower-case for s in sequence_examples ] seq_token = self.tokenizer.batch_encode_plus(sequence_examples, add_special_tokens=True, padding="longest", return_tensors='pt').to(self.device) attention_mask[:, :seq_token.input_ids.shape[1]] = 1 seq_token = self.pad_data(seq_token.input_ids, dim=1, max_length=max_length_batch) attention_mask = self.pad_data(attention_mask, dim=1, max_length=max_length_batch) seq_tokens.append(seq_token) attention_masks.append(attention_mask) seq_tokens = torch.hstack(seq_tokens) attention_masks = torch.hstack(attention_masks) S_batch.append(seq_tokens) mask_batch.append(attention_masks) if task_name == 'contact_map': label = F.pad(label, [add_BOS, self.max_length-label.shape[0]-add_BOS, add_BOS, self.max_length-label.shape[0]-add_BOS]) else: label = torch.tensor(label) label_batch.append(label) name_batch.append(sample['name']) yield { 'name': name_batch, 'seq': torch.cat(S_batch, dim=0).to(self.device), 'attention_mask': torch.cat(mask_batch, dim=0).to(self.device)==1, 'label': torch.stack(label_batch).to(self.device), } if self.pretrain_model_name == 'protgpt2': # 没有BOS, EOS, 并且长度与氨基酸数量不一致 max_length_batch = 0 for sample in data[i:i + batch_size]: if "pair" in self.task_type: submax = max([len(torch.tensor(self.tokenizer.encode(_seq))) for _seq in sample['seq']]) max_length_batch = max(max_length_batch, submax) else: seq_token = torch.tensor(self.tokenizer.encode(sample['seq'])) max_length_batch = max(max_length_batch, len(seq_token)) name_batch, S_batch, mask_batch, label_batch = [], [], [], [] for sample in data[i:i + batch_size]: seq = sample['seq'] label = sample['label'] if not isinstance(seq, list): seq = [seq] seq_tokens, attention_masks = [], [] for _seq in seq: attention_mask = torch.zeros(self.max_length) seq_token = torch.tensor(self.tokenizer.encode(_seq)) attention_mask[:len(seq_token)] = 1 seq_token = self.pad_data(seq_token, dim=0, max_length=max_length_batch) attention_mask = self.pad_data(attention_mask, dim=0, max_length=max_length_batch) seq_tokens.append(seq_token) attention_masks.append(attention_mask) seq_tokens = torch.hstack(seq_tokens) attention_masks = torch.hstack(attention_masks) S_batch.append(seq_tokens) mask_batch.append(attention_masks) if task_name == 'contact_map': label = F.pad(label, [add_BOS, self.max_length-label.shape[0]-add_BOS, add_BOS, self.max_length-label.shape[0]-add_BOS]) else: label = torch.tensor(label) label_batch.append(label) name_batch.append(sample['name']) yield { 'name': name_batch, 'seq': torch.stack(S_batch).to(self.device), 'attention_mask': torch.stack(mask_batch).to(self.device)==1, 'label': torch.stack(label_batch).to(self.device), } if self.pretrain_model_name == 'protrek': if "pair" in self.task_type: max_length_batch = max( [max(len(seq) for seq in sample['seq']) for sample in data[i:i + batch_size]] ) + 2 else: max_length_batch = max([len(sample['seq']) for sample in data[i:i + batch_size]]) + 2 name_batch, struct_batch, seq_batch, label_batch, mask_batch = [], [], [], [], [] # attention_mask = torch.zeros(len(data[i:i + batch_size]), max_length_batch) for idx, sample in enumerate(data[i:i + batch_size]): seq = sample['seq'] X = sample['X'] label = sample['label'] if not isinstance(seq, list) and not isinstance(X, list): seq, X = [seq], [X] seq_tokens, attention_masks, struct_tokens = [], [], [] for _seq, _X in zip(seq, X): N, CA, C, CB, O = _X[:, 0], _X[:, 1], _X[:, 2], _X[:, 3], _X[:, 4] states = self.encoder_3di.encode_atoms(ca = CA.numpy(), cb = CB.numpy(), n = N.numpy(), c = C.numpy()) struct_sequence = self.encoder_3di.build_sequence(states).lower() if self.sequence_only: struct_sequence = ''.join(['#' for one in struct_sequence]) attention_mask = torch.zeros(max_length_batch) attention_mask[:len(_seq)+2] = 1 attention_masks.append(attention_mask) seq_tokens.append(_seq) struct_tokens.append(struct_sequence) attention_masks = torch.hstack(attention_masks) struct_batch.append(struct_tokens) seq_batch.append(seq_tokens) if task_name == 'contact_map': label = F.pad(label, [add_BOS, self.max_length-label.shape[0]-add_BOS, add_BOS, self.max_length-label.shape[0]-add_BOS]) else: label = torch.tensor(label) label_batch.append(label) name_batch.append(sample['name']) mask_batch.append(attention_masks) yield { 'name': name_batch, 'seq_batch': seq_batch, 'struct_batch': struct_batch, 'attention_mask': torch.stack(mask_batch).to(self.device)==1, 'label': torch.stack(label_batch).to(self.device) } if self.pretrain_model_name == 'saport': if "pair" in self.task_type: max_length_batch = max( [max(len(seq) for seq in sample['seq']) for sample in data[i:i + batch_size]] ) + 2 else: max_length_batch = max([len(sample['seq']) for sample in data[i:i + batch_size]]) + 2 name_batch, struct_batch, S_batch, mask_batch, label_batch = [], [], [], [], [] for sample in data[i:i + batch_size]: seq = sample['seq'] X = sample['X'] label = sample['label'] if not isinstance(seq, list) and not isinstance(X, list): seq, X = [seq], [X] seq_tokens, attention_masks, struct_tokens = [], [], [] for _seq, _X in zip(seq, X): N, CA, C, CB, O = _X[:, 0], _X[:, 1], _X[:, 2], _X[:, 3], _X[:, 4] states = self.encoder_3di.encode_atoms(ca = CA.numpy(), cb = CB.numpy(), n = N.numpy(), c = C.numpy()) struct_sequence = self.encoder_3di.build_sequence(states).lower() attention_mask = torch.zeros(self.max_length) if self.sequence_only: struct_sequence = ''.join(['#' for one in struct_sequence]) merged_seq = ''.join(a + b.lower() for a, b in zip(_seq, struct_sequence)) seq_token = self.tokenizer(merged_seq, return_tensors="pt").input_ids[0] attention_mask[:seq_token.shape[0]] = 1 seq_token = self.pad_data(seq_token, dim=0, max_length=max_length_batch) attention_mask = self.pad_data(attention_mask, dim=0, max_length=max_length_batch) seq_tokens.append(seq_token) attention_masks.append(attention_mask) seq_tokens = torch.hstack(seq_tokens) attention_masks = torch.hstack(attention_masks) S_batch.append(seq_tokens) mask_batch.append(attention_masks) if task_name == 'contact_map': label = F.pad(label, [add_BOS, self.max_length-label.shape[0]-add_BOS, add_BOS, self.max_length-label.shape[0]-add_BOS]) else: label = torch.tensor(label) label_batch.append(label) name_batch.append(sample['name']) yield { 'name': name_batch, 'seq': torch.stack(S_batch, dim=0).to(self.device), 'attention_mask': torch.stack(mask_batch, dim=0).to(self.device)==1, 'label': torch.stack(label_batch).to(self.device), } def forward(self, x): names, labels = x['name'], x['label'] if self.pretrain_model_name == 'esm2_650m': # Forward pass through the pre-trained model seq, attention_mask, t_lengths, smiles = x['seq'], x['attention_mask'], x['t_lengths'], x['smiles'] embeddings, attention_masks = [], [] for i, (_seq, _attention_mask, _t_length) in enumerate(zip(seq, attention_mask, t_lengths)): outputs = self.pretrain_model.esm( _seq, attention_mask=_attention_mask, return_dict=True, ) _embeddings = outputs.last_hidden_state _t_length = torch.where(_t_length==_t_length.max(), _t_length-1, _t_length) if i != len(seq) - 1: _attention_mask.scatter_(1, _t_length.unsqueeze(1), True) embeddings.append(_embeddings) attention_masks.append(_attention_mask) embeddings = torch.cat(embeddings, dim=1) attention_mask = torch.cat(attention_masks, dim=-1) if "pair" in self.task_type: ends = torch.tensor([attention_mask.shape[1]]*attention_mask.shape[0])-1 starts = torch.ones_like(ends) else: ends = attention_mask.sum(dim=-1)-1 starts = torch.ones_like(ends) if self.pretrain_model_name == 'esm3_1.4b': # WORKING protein_tensor, attention_mask = x['protein_tensor'], x['attention_mask'] embeddings = [] if "pair" in self.task_type: split_idx = protein_tensor.sequence.shape[1] // 2 protein_tensor = [ _BatchedESMProteinTensor( sequence=protein_tensor.sequence[:, :split_idx], structure=protein_tensor.structure[:, :split_idx], coordinates=protein_tensor.coordinates[:, :split_idx], ), _BatchedESMProteinTensor( sequence=protein_tensor.sequence[:, split_idx:], structure=protein_tensor.structure[:, split_idx:], coordinates=protein_tensor.coordinates[:, split_idx:], ) ] attention_mask = [attention_mask[:, :split_idx], attention_mask[:, split_idx:]] else: protein_tensor, attention_mask = [protein_tensor], [attention_mask] for _protein_tensor in protein_tensor: output = self.pretrain_model.logits( _protein_tensor, LogitsConfig( sequence=True, structure=True, secondary_structure=True, sasa=True, function=True, residue_annotations=True, return_embeddings=True, ), ) embeddings.append(output.embeddings) embeddings = torch.cat(embeddings, dim=-1) if "pair" in self.task_type: attention_mask = (attention_mask[0].int() + attention_mask[1].int() > 0) ends = attention_mask.sum(dim=-1)-1 starts = torch.ones_like(ends) attention_mask = attention_mask else: ends = attention_mask[0].sum(dim=-1)-1 starts = torch.ones_like(ends) attention_mask = attention_mask[0] # embeddings = F.pad(embeddings, (0, 0, 0, self.max_length-embeddings.shape[1], 0, 0), value=0) if self.pretrain_model_name == 'esmc_600m': protein_tensor, attention_mask = x['protein_tensor'], x['attention_mask'] embeddings = [] if "pair" in self.task_type: split_idx = protein_tensor.sequence.shape[1] // 2 protein_tensor = [ _BatchedESMProteinTensor( sequence=protein_tensor.sequence[:, :split_idx] ), _BatchedESMProteinTensor( sequence=protein_tensor.sequence[:, split_idx:] ) ] attention_mask = [attention_mask[:, :split_idx], attention_mask[:, split_idx:]] else: protein_tensor, attention_mask = [protein_tensor], [attention_mask] for _protein_tensor in protein_tensor: logits_output = self.pretrain_model.logits( _protein_tensor, LogitsConfig(sequence=True, return_embeddings=True) ) embeddings.append(logits_output.embeddings) embeddings = torch.cat(embeddings, dim=-1) if "pair" in self.task_type: attention_mask = (attention_mask[0].int() + attention_mask[1].int() > 0) ends = attention_mask.sum(dim=-1)-1 starts = torch.ones_like(ends) attention_mask = attention_mask else: ends = attention_mask[0].sum(dim=-1)-1 starts = torch.ones_like(ends) attention_mask = attention_mask[0] if self.pretrain_model_name == 'procyon': # 长度和氨基酸数量不同 seq_embedding, struct_embedding = x['seq'], x['X'] embeddings, attention_masks = [], [] if "pair" in self.task_type: seq_split_idx = seq_embedding.shape[-1] // 2 struct_split_idx = struct_embedding.shape[-1] // 2 seq_embedding = [ seq_embedding[:, :seq_split_idx], seq_embedding[:, seq_split_idx:] ] struct_embedding = [ struct_embedding[:, :, :struct_split_idx], struct_embedding[:, :, struct_split_idx:] ] else: seq_embedding, struct_embedding = [seq_embedding], [struct_embedding] for _seq_embedding, _strcture_embedding in zip(seq_embedding, struct_embedding): seq_embeddings = self.pretrain_model.token_projectors["aaseq"]( _seq_embedding ) struct_embeddings = self.pretrain_model.token_projectors["prot_structure"]( _strcture_embedding ) instructions = [ "Describe the following protein with features: <|protein|> <|struct|>" ] * seq_embeddings.shape[0] input_ids, attn_masks = self.pretrain_model._prepare_text_inputs_and_tokenize(instructions, [[]] * seq_embeddings.shape[0], no_pad=True) input_ids, attn_masks = input_ids.to(self.device), attn_masks.to(self.device) if self.sequence_only: input_embeds, ret_output_indices = self.pretrain_model._prepare_input_embeddings( input_ids, protein_soft_tokens=seq_embeddings ) else: input_embeds, ret_output_indices = self.pretrain_model._prepare_input_embeddings( input_ids, protein_soft_tokens=seq_embeddings, protein_struct_tokens=struct_embeddings ) attention_mask = ~(input_ids == self.pretrain_model.tokenizer.pad_token_id) outputs = self.pretrain_model.text_encoder( input_embeds = input_embeds, attn_masks = attn_masks, ) embeddings.append(outputs.hidden_states[-1]) # shape(b, 2048, 4096) attention_masks.append(attention_mask) embeddings = torch.cat(embeddings, dim=-1) if "pair" in self.task_type: attention_mask = (attention_masks[0].int() + attention_masks[1].int() > 0) ends = attention_mask.sum(dim=-1)-1 starts = torch.ones_like(ends) else: ends = attention_masks[0].sum(dim=-1)-1 starts = torch.ones_like(ends) attention_mask = attention_masks[0] if self.pretrain_model_name == 'gearnet': embeddings = x['X'] attention_mask = x['attention_mask'] ends = attention_mask.sum(dim=-1) starts = torch.zeros_like(ends) if self.pretrain_model_name == 'prollama': # 长度和氨基酸数量不同 seq, attention_mask = x['seq'], x['attention_mask'] if "pair" in self.task_type: split_idx = seq.shape[1] // 2 seq = [ seq[:, :split_idx], seq[:, split_idx:] ] attention_mask = [attention_mask[:, :split_idx], attention_mask[:, split_idx:]] else: seq, attention_mask = [seq], [attention_mask] embeddings = [] for _seq, _attention_mask in zip(seq, attention_mask): out = self.pretrain_model( input_ids = _seq, attention_mask = _attention_mask, output_hidden_states=True ) embeddings.append(out.hidden_states[-1].float()) embeddings = torch.cat(embeddings, dim=-1) if "pair" in self.task_type: attention_mask = (attention_mask[0].int() + attention_mask[1].int() > 0) ends = attention_mask.sum(dim=-1)-1 starts = torch.ones_like(ends) else: ends = attention_mask[0].sum(dim=-1)-1 starts = torch.ones_like(ends) attention_mask = attention_mask[0] if self.pretrain_model_name == 'venusplm': seq, attention_mask = x['seq'], x['attention_mask'] if "pair" in self.task_type: split_idx = seq.shape[1] // 2 seq = [ seq[:, :split_idx], seq[:, split_idx:] ] attention_mask = [attention_mask[:, :split_idx], attention_mask[:, split_idx:]] else: seq, attention_mask = [seq], [attention_mask] embeddings = [] for _seq, _attention_mask in zip(seq, attention_mask): outputs = self.pretrain_model( input_ids=_seq, attention_mask=_attention_mask, output_hidden_states=True ) embeddings.append(outputs.hidden_states[-1]) embeddings = torch.cat(embeddings, dim=-1) if "pair" in self.task_type: attention_mask = (attention_mask[0].int() + attention_mask[1].int() > 0) ends = attention_mask.sum(dim=-1)-1 starts = torch.ones_like(ends) else: ends = attention_mask[0].sum(dim=-1)-1 starts = torch.ones_like(ends) attention_mask = attention_mask[0] if self.pretrain_model_name == 'prost': seq, attention_mask = x['seq'], x['attention_mask'] if "pair" in self.task_type: split_idx = seq.shape[1] // 2 seq = [ seq[:, :split_idx], seq[:, split_idx:] ] attention_mask = [attention_mask[:, :split_idx], attention_mask[:, split_idx:]] else: seq, attention_mask = [seq], [attention_mask] embeddings = [] for _seq, _attention_mask in zip(seq, attention_mask): outputs = self.pretrain_model( input_ids=_seq, attention_mask=_attention_mask, return_dict=True ) embeddings.append(outputs.residue_feature) embeddings = torch.cat(embeddings, dim=-1) if "pair" in self.task_type: attention_mask = (attention_mask[0].int() + attention_mask[1].int() > 0) ends = attention_mask.sum(dim=-1)-1 starts = torch.ones_like(ends) else: ends = attention_mask[0].sum(dim=-1)-1 starts = torch.ones_like(ends) attention_mask = attention_mask[0] if self.pretrain_model_name == 'progen2': seq, attention_mask = x['seq'], x['attention_mask'] if "pair" in self.task_type: split_idx = seq.shape[1] // 2 seq = [ seq[:, :split_idx], seq[:, split_idx:] ] attention_mask = [attention_mask[:, :split_idx], attention_mask[:, split_idx:]] else: seq, attention_mask = [seq], [attention_mask] embeddings = [] for _seq, _attention_mask in zip(seq, attention_mask): outputs = self.pretrain_model.transformer( _seq, return_dict=True ) embeddings.append(outputs[0]) embeddings = torch.cat(embeddings, dim=-1) if "pair" in self.task_type: attention_mask = (attention_mask[0].int() + attention_mask[1].int() > 0) ends = attention_mask.sum(dim=-1)-1 starts = torch.ones_like(ends) else: ends = attention_mask[0].sum(dim=-1)-1 starts = torch.ones_like(ends) attention_mask = attention_mask[0] if self.pretrain_model_name == 'prostt5': seq, attention_mask = x['seq'], x['attention_mask'] if "pair" in self.task_type: split_idx = seq.shape[1] // 2 seq = [ seq[:, :split_idx], seq[:, split_idx:] ] attention_mask = [attention_mask[:, :split_idx], attention_mask[:, split_idx:]] else: seq, attention_mask = [seq], [attention_mask] embeddings = [] for _seq, _attention_mask in zip(seq, attention_mask): embedding_repr = self.pretrain_model( _seq, attention_mask=_attention_mask ) _embeddings = embedding_repr.last_hidden_state _embeddings = _embeddings.reshape(_embeddings.shape[0]//2, 2, _embeddings.shape[1], _embeddings.shape[2]) _embeddings = torch.cat([_embeddings[:,0], _embeddings[:,1]], dim=-1) embeddings.append(_embeddings) embeddings = torch.cat(embeddings, dim=-1) if "pair" in self.task_type: attention_mask = (attention_mask[0].int() + attention_mask[1].int() > 0) attention_mask = attention_mask[::2] ends = attention_mask.sum(dim=-1)-1 starts = torch.ones_like(ends) else: attention_mask = attention_mask[0][::2] ends = attention_mask.sum(dim=-1)-1 starts = torch.ones_like(ends) if self.pretrain_model_name == 'protgpt2': seq, attention_mask = x['seq'], x['attention_mask'] if "pair" in self.task_type: split_idx = seq.shape[1] // 2 seq = [ seq[:, :split_idx], seq[:, split_idx:] ] attention_mask = [attention_mask[:, :split_idx], attention_mask[:, split_idx:]] else: seq, attention_mask = [seq], [attention_mask] embeddings = [] for _seq, _attention_mask in zip(seq, attention_mask): outputs = self.pretrain_model.transformer(_seq) embeddings.append(outputs.last_hidden_state) embeddings = torch.cat(embeddings, dim=-1) if "pair" in self.task_type: attention_mask = (attention_mask[0].int() + attention_mask[1].int() > 0) ends = attention_mask.sum(dim=-1)-1 starts = torch.ones_like(ends) else: ends = attention_mask[0].sum(dim=-1)-1 starts = torch.ones_like(ends) attention_mask = attention_mask[0] if self.pretrain_model_name == 'protrek': seq, attention_mask, struct = x['seq_batch'], x['attention_mask'], x['struct_batch'] if "pair" in self.task_type: split_idx = attention_mask.shape[1] // 2 seq = [ [ele[0] for ele in seq], [ele[1] for ele in seq] ] struct = [ [ele[0] for ele in struct], [ele[1] for ele in struct] ] attention_mask = [attention_mask[:, :split_idx], attention_mask[:, split_idx:]] else: seq, attention_mask, struct = [[ele[0] for ele in seq]], [attention_mask], [[ele[0] for ele in struct]] embeddings = [] for _seq, _attention_mask, _struct in zip(seq, attention_mask, struct): seq_embedding = self.pretrain_model.get_protein_repr(_seq) struc_embedding = self.pretrain_model.get_structure_repr(_struct) _embeddings = torch.cat([seq_embedding, struc_embedding], dim=-1) if "pair" in self.task_type: _embeddings = self.pad_data(_embeddings, dim=1, max_length=split_idx) embeddings.append(_embeddings) embeddings = torch.cat(embeddings, dim=-1) if "pair" in self.task_type: attention_mask = (attention_mask[0].int() + attention_mask[1].int() > 0) ends = attention_mask.sum(dim=-1)-1 starts = torch.ones_like(ends) else: ends = attention_mask[0].sum(dim=-1)-1 starts = torch.ones_like(ends) attention_mask = attention_mask[0] if self.pretrain_model_name == 'prosst2048': # Forward pass through the pre-trained model outputs = self.pretrain_model( input_ids=x['seq'], attention_mask=x['attention_mask'], output_hidden_states=True, ss_input_ids=x['X'] ) embeddings = outputs.hidden_states[-1] ends = x['attention_mask'].sum(dim=-1)-1 starts = torch.ones_like(ends) attention_mask = x['attention_mask'] if self.pretrain_model_name == 'saport': seq, attention_mask = x['seq'], x['attention_mask'] if "pair" in self.task_type: split_idx = seq.shape[1] // 2 seq = [ seq[:, :split_idx], seq[:, split_idx:] ] attention_mask = [attention_mask[:, :split_idx], attention_mask[:, split_idx:]] else: seq, attention_mask = [seq], [attention_mask] embeddings = [] for _seq, _attention_mask in zip(seq, attention_mask): output = self.pretrain_model.esm( _seq, attention_mask=_attention_mask, return_dict=True, ) embeddings.append(output.last_hidden_state) embeddings = torch.cat(embeddings, dim=-1) if "pair" in self.task_type: attention_mask = (attention_mask[0].int() + attention_mask[1].int() > 0) attention_mask = attention_mask[::2] ends = attention_mask.sum(dim=-1)-1 starts = torch.ones_like(ends) else: attention_mask = attention_mask[0][::2] ends = attention_mask.sum(dim=-1)-1 starts = torch.ones_like(ends) # 这个embedding 是[B,L,D]的形式, L是氨基酸数量, 但对于语言模型老师, L这个维度可能不刚好等于氨基酸数量，无法用于氨基酸级别的预测任务。另外，有些方法会加上[EOS, BOS]，有些又不需要，最好是在最后的embedding上统一去除[EOS, BOS] return self.post_process(names, labels, embeddings, attention_mask, starts, ends, smiles) def post_process(self, names, labels, embeddings, attention_mask, starts, ends, smiles)->list: results = [] for i, end in enumerate(ends): start = starts[i] label = labels[i].cpu() if self.task_type == 'contact': label = labels[i,start:end, start:end].cpu() results.append( { 'name': names[i], 'embedding': embeddings[i,start:end].cpu(), 'attention_mask': attention_mask[i,start:end].cpu(), 'label': label, 'smiles': smiles[i] if smiles is not None else None } ) return results def pad_data(self, data, dim=0, pad_value=0, max_length=1022): if data.shape[dim] < max_length: data = dynamic_pad(data, [0, max_length-data.shape[dim]], dim=dim, pad_value=pad_value) else: # start = torch.randint(0, data.shape[0]-self.max_length+1, (1,)).item() start = 0 data = data[start:start+max_length] return data def inference_datasets(self, data, task_name=None): self.pretrain_model.eval() with torch.no_grad(): proccessed_data = [] for i, batch in enumerate(tqdm(self.construct_batch(data, self.batch_size, task_name), desc='Extracting embeddings')): # print('batch size:', len(batch['name'])) try: results = self.forward(batch) proccessed_data.extend(results) except: print(f"Error processing batch {i}") return proccessed_data