| 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 MetallonBinding(Dataset): |
| def __init__(self, data_path, prompt='', return_prompt=False): |
| super(MetallonBinding, 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: |
| name = str(row['name']).strip() |
| prot_seq = str(row['aa_seq']).strip() |
| result = str(int(row['label'])).strip() |
|
|
| text_seq = f"<answer>{result}</answer>\n" |
| |
| prompt = f""" |
| Task: Determine whether this protein is a metalloprotein based on the provided sequence and protein name {name}. |
| Background: Metalloproteins are proteins that bind metal ions, often through specific amino acid residues such as histidine (H), cysteine (C), aspartate (D), or glutamate (E). |
| Question: Does this protein bind metal ions? Please choose one of the following options: |
| 0: Non-metalloprotein — This protein does **not** bind to any metal ions. |
| 1: Metalloprotein — This protein **binds** to one or more metal ions. |
| """ |
| 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 |
