| | import sys; sys.path.append('/nfs_beijing/kubeflow-user/zhangyang_2024/workspace/protein_benchmark/model_zoom') |
| | import torch |
| | import torch.nn as nn |
| | import os |
| | import numpy as np |
| | from tqdm import tqdm |
| | from src.model.pretrain_modules import ( |
| | ESM2Model, SmilesModel, ESM3Model, ESMC600MModel, ProCyonModel, |
| | GearNetModel, ProLLAMAModel, ProSTModel, ProtGPT2Model, ProTrekModel, |
| | SaPortModel, VenusPLMModel, ProSST2048Model, ProGen2Model, ProstT5Model, |
| | ProtT5, DPLMModel, OntoProteinModel, ANKHBase, PGLMModel,BIOModel |
| | ) |
| | 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.task_type = task_type |
| | self.device = device |
| | self.start, self.end = 1, -1 |
| | self.setup_model() |
| | |
| | def setup_model(self): |
| | """ |
| | Setup the pre-trained model based on the specified name. |
| | ['esm2_650m', 'esm3_1.4b', 'esmc_600m', 'procyon', 'prollama', 'progen2', 'prostt5', 'protgpt2', 'protrek', 'saport', 'gearnet', 'prost', 'prosst2048', 'venusplm'] |
| | """ |
| | device = 'cuda' |
| | self.smiles_model = SmilesModel(device) |
| | if self.pretrain_model_name == 'esm2_650m': |
| | self.pretrain_model = ESM2Model(device) |
| | elif self.pretrain_model_name == 'esm2_35m': |
| | self.pretrain_model = ESM2Model(device, model_path='esm2_35m') |
| | elif self.pretrain_model_name == 'esm2_150m': |
| | self.pretrain_model = ESM2Model(device, model_path='esm2_150m') |
| | elif self.pretrain_model_name == 'esm2_3b': |
| | self.pretrain_model = ESM2Model(device, model_path='esm2_3b') |
| | elif self.pretrain_model_name == 'esm2_15b': |
| | self.pretrain_model = ESM2Model(device, model_path='esm2_15b') |
| | elif self.pretrain_model_name == 'esm3_1.4b': |
| | self.pretrain_model = ESM3Model(device) |
| | elif self.pretrain_model_name == 'esmc_600m': |
| | self.pretrain_model = ESMC600MModel(device) |
| | elif self.pretrain_model_name == 'procyon': |
| | self.pretrain_model = ProCyonModel(device) |
| | elif self.pretrain_model_name == 'prollama': |
| | self.pretrain_model = ProLLAMAModel(device) |
| | elif self.pretrain_model_name == 'progen2': |
| | self.pretrain_model = ProGen2Model(device) |
| | elif self.pretrain_model_name == 'prostt5': |
| | self.pretrain_model = ProstT5Model(device) |
| | elif self.pretrain_model_name == 'protgpt2': |
| | self.pretrain_model = ProtGPT2Model(device) |
| | self.start, self.end = 0, -1 |
| | elif self.pretrain_model_name == 'protrek_35m': |
| | self.pretrain_model = ProTrekModel(device, model_path='protrek_35m') |
| | elif self.pretrain_model_name == 'protrek': |
| | self.pretrain_model = ProTrekModel(device) |
| | elif self.pretrain_model_name == 'saport': |
| | self.pretrain_model = SaPortModel(device) |
| | elif self.pretrain_model_name == 'saport_1.3b': |
| | self.pretrain_model = SaPortModel(device, model_path='saprot_1.3b') |
| | elif self.pretrain_model_name == 'saport_35m': |
| | self.pretrain_model = SaPortModel(device, model_path='saprot_35m') |
| | elif self.pretrain_model_name == 'gearnet': |
| | self.pretrain_model = GearNetModel(device) |
| | elif self.pretrain_model_name == 'prost': |
| | self.pretrain_model = ProSTModel(device) |
| | elif self.pretrain_model_name == 'prosst2048': |
| | self.pretrain_model = ProSST2048Model(device) |
| | elif self.pretrain_model_name == 'venusplm': |
| | self.pretrain_model = VenusPLMModel(device) |
| | elif self.pretrain_model_name == 'prott5': |
| | self.pretrain_model = ProtT5(device) |
| | elif self.pretrain_model_name == 'dplm': |
| | self.pretrain_model = DPLMModel(device) |
| | elif self.pretrain_model_name == 'dplm_150m': |
| | self.pretrain_model = DPLMModel(device, model_path='dplm_150m') |
| | elif self.pretrain_model_name == 'dplm_3b': |
| | self.pretrain_model = DPLMModel(device, model_path='dplm_3b') |
| | elif self.pretrain_model_name == 'ontoprotein': |
| | self.pretrain_model = OntoProteinModel(device) |
| | elif self.pretrain_model_name == "ankh_base": |
| | self.pretrain_model = ANKHBase(device) |
| | elif self.pretrain_model_name == "pglm": |
| | self.pretrain_model = PGLMModel(device) |
| | elif self.pretrain_model_name == "pglm-3b": |
| | self.pretrain_model = PGLMModel(device, model_path="proteinglm-3b-mlm") |
| | elif self.pretrain_model_name == "pretrain_bio": |
| | self.pretrian_model = BIOModel(device,model_path="/nas/shared/kilab/wangyujia/ProtT3/all_checkpoints/stage2_06301657/last.ckpt/converted.ckpt") |
| | |
| | def setup_peft(self, peft_type="lora", **kwargs): |
| | if self.pretrain_model is None: |
| | raise RuntimeError("pretrained model is not initialized, please initial it first.") |
| | self.pretrain_model.setup_peft( |
| | peft_type=peft_type, |
| | **kwargs |
| | ) |
| |
|
| | @torch.no_grad() |
| | def inference_datasets(self, data, task_name=None): |
| | self.pretrain_model.eval() |
| | proccessed_data = [] |
| | for i in tqdm(range(0, len(data), self.batch_size), desc='Extracting embeddings'): |
| | if "|" in data[0]['seq']: |
| | samples_A, samples_B = [], [] |
| | for sample in data[i:i + self.batch_size]: |
| | sample_A = {key: value for key, value in sample.items() if key != 'seq'} |
| | sample_B = {key: value for key, value in sample.items() if key != 'seq'} |
| | sample_A['seq'] = sample['seq'].split('|')[0] |
| | sample_B['seq'] = sample['seq'].split('|')[1] |
| | if 'pdb_path' in sample: |
| | sample_A['pdb_path'] = sample['pdb_path'].split('|')[0] |
| | sample_B['pdb_path'] = sample['pdb_path'].split('|')[1] |
| | sample_A['X'] = sample['X'][0] |
| | sample_B['X'] = sample['X'][1] |
| | samples_A.append(sample_A) |
| | samples_B.append(sample_B) |
| |
|
| | batch_A = self.pretrain_model.construct_batch(samples_A) |
| | batch_B = self.pretrain_model.construct_batch(samples_B) |
| | results_A = self.pretrain_model(batch_A, task_type=self.task_type) |
| | results_B = self.pretrain_model(batch_B, task_type=self.task_type) |
| | results = [] |
| | for idx in range(len(results_A)): |
| | result = {} |
| | for key in results_A[0].keys(): |
| | PAD = torch.ones_like(results_A[idx]['embedding'])[:1] |
| | if key == 'embedding': |
| | result[key] = torch.cat([results_A[idx]['embedding'], PAD, results_B[idx]['embedding']], dim=0) |
| | elif key == 'attention_mask': |
| | result[key] = torch.cat([results_A[idx]['attention_mask'], torch.tensor([True]), results_B[idx]['attention_mask']], dim=0)==1 |
| | else: |
| | result[key] = results_A[idx][key] |
| | results.append(result) |
| | else: |
| | samples = data[i:i + self.batch_size] |
| | batch = self.pretrain_model.construct_batch(samples) |
| | results = self.pretrain_model(batch, task_type=self.task_type) |
| | |
| | if samples[0].get('smiles') is not None: |
| | batch_smi = self.smiles_model.construct_batch(samples) |
| | for idx in range(len(samples)): |
| | results[idx]['smiles'] = batch_smi['smiles'][idx] |
| | |
| | proccessed_data.extend(results) |
| | |
| | return proccessed_data |
| | |
| | def forward(self, data): |
| | if "|" in data[0]['seq']: |
| | samples_A, samples_B = [], [] |
| | for sample in data: |
| | sample_A = {key: value for key, value in sample.items() if key != 'seq'} |
| | sample_B = {key: value for key, value in sample.items() if key != 'seq'} |
| | sample_A['seq'] = sample['seq'].split('|')[0] |
| | sample_B['seq'] = sample['seq'].split('|')[1] |
| | if 'pdb_path' in sample: |
| | sample_A['pdb_path'] = sample['pdb_path'].split('|')[0] |
| | sample_B['pdb_path'] = sample['pdb_path'].split('|')[1] |
| | sample_A['X'] = sample['X'][0] |
| | sample_B['X'] = sample['X'][1] |
| | samples_A.append(sample_A) |
| | samples_B.append(sample_B) |
| |
|
| | batch_A = self.pretrain_model.construct_batch(samples_A) |
| | batch_B = self.pretrain_model.construct_batch(samples_B) |
| | embedding_A = self.pretrain_model(batch_A, task_type=self.task_type, post_process=False)[:,self.start:self.end,:] |
| | embedding_B = self.pretrain_model(batch_B, task_type=self.task_type, post_process=False)[:,self.start:self.end,:] |
| | bs, hidden_dim = embedding_A.shape[0], embedding_A.shape[-1] |
| | PAD = torch.ones((bs, 1, hidden_dim), device=embedding_A.device) |
| | embedding = torch.cat([embedding_A, PAD, embedding_B], dim=1).contiguous() |
| | labels = torch.stack(batch_A["label"]).to(embedding.device).to(embedding.dtype) |
| | PAD_MASK = torch.ones((bs, 1), device=embedding_A.device) |
| | attention_mask = torch.cat( |
| | [batch_A['attention_mask'][:,self.start:self.end], PAD_MASK, batch_B['attention_mask'][:,self.start:self.end]], dim=1 |
| | )==1 |
| | else: |
| | batch = self.pretrain_model.construct_batch(data) |
| | embedding = self.pretrain_model(batch, task_type=self.task_type, post_process=False)[:,self.start:self.end,:] |
| | |
| | labels = torch.stack(batch["label"]).to(embedding.device).to(embedding.dtype) |
| | |
| | attention_mask = batch["attention_mask"][:,self.start:self.end]==1 |
| |
|
| | batch_smi = None |
| | if data[0].get('smiles') is not None: |
| | batch_smi = self.smiles_model.construct_batch(data) |
| | batch_smi = torch.stack(batch_smi["smiles"]).contiguous().to(embedding.device).to(embedding.dtype) |
| | |
| | if self.task_type == "contact": |
| | true_length = attention_mask.shape[-1] |
| | labels = labels[:, :true_length, :true_length] |
| | elif self.task_type == "residual_classification": |
| | true_length = attention_mask.shape[-1] |
| | labels = labels[:, :true_length] |
| | if self.pretrain_model_name == 'prostt5': |
| | attention_mask = attention_mask[::2] |
| | return embedding, labels, attention_mask, batch_smi |
| |
|
