| | import json |
| | from pytorch_lightning import LightningDataModule |
| | from torch.utils.data import DataLoader, ConcatDataset, Dataset |
| | from data_provider.stage1_dm import SwissProtDataset, OntoProteinDataset |
| | import pandas as pd |
| |
|
| | class BindingDB(Dataset): |
| | def __init__(self, data_path, prompt='', return_prompt=False): |
| | super(BindingDB, self).__init__() |
| | self.data_path = data_path |
| | self.user_prompt = prompt |
| | self.return_prompt = return_prompt |
| | |
| | self.data_list = self._load_and_preprocess(self.data_path) |
| | self.text2id = self._build_text_vocab() |
| |
|
| | def _load_and_preprocess(self, data_path): |
| | data_list = [] |
| | df = pd.read_csv(data_path) |
| | for _, row in df.iterrows(): |
| | try: |
| | ligand_smiles = str(row['ligand']).strip() |
| | prot_seq = str(row['protein']).strip() |
| | result = str(row['ic50']).strip() |
| |
|
| | text_seq = f"<answer>{result}</answer>\n" |
| | |
| | prompt = f""" |
| | 【Protein sequence (1-letter amino acid codes)】；{ligand_smiles}【Ligand structure (SMILES)】 |
| | Task: Evaluate the inhibitory effect of the ligand on the given protein. |
| | Note: IC50 (half maximal inhibitory concentration) is the concentration of a substance required to inhibit 50% of the protein's activity. Lower IC50 values indicate stronger inhibition. |
| | Based on the provided protein and ligand, predict the inhibitory strength by classifying the IC50 level: |
| | """ |
| | if self.user_prompt: |
| | prompt += self.user_prompt |
| |
|
| | |
| | |
| | data_list.append((prot_seq, text_seq, prompt)) |
| | except Exception as e: |
| | print(f"警告: 跳过有问题的行: {row}，原因: {e}") |
| | return data_list |
| |
|
| | def _build_text_vocab(self): |
| | text2id = {} |
| | for _, text_seq, _ in self.data_list: |
| | if text_seq not in text2id: |
| | text2id[text_seq] = len(text2id) |
| | return text2id |
| |
|
| | def shuffle(self): |
| | random.shuffle(self.data_list) |
| | return self |
| |
|
| | def __len__(self): |
| | return len(self.data_list) |
| |
|
| | def __getitem__(self, index): |
| | prot_seq, text_seq, prompt = self.data_list[index] |
| | if self.return_prompt: |
| | return prot_seq, prompt, text_seq,index |
| | return prot_seq, text_seq, index |
