| import pandas as pd |
| import json |
| import numpy as np |
|
|
| from torch.utils.data import Subset |
| from transformers import EsmTokenizer |
| from ..lmdb_dataset import * |
| from ..data_interface import register_dataset |
|
|
|
|
| @register_dataset |
| class SaprotAnnotationDataset(LMDBDataset): |
| def __init__(self, |
| tokenizer: str, |
| bias_feature: bool = False, |
| max_length: int = 1024, |
| mask_struc_ratio: float = None, |
| plddt_threshold: float = None, |
| **kwargs): |
| """ |
| |
| Args: |
| tokenizer: EsmTokenizer config path |
| |
| bias_feature: If True, structure information will be used |
| |
| max_length: Max length of sequence |
| |
| mask_struc_ratio: Ratio of masked structure tokens, replace structure tokens with "#" |
| |
| plddt_threshold: If not None, mask structure tokens with pLDDT < threshold |
| |
| **kwargs: other arguments for LMDBDataset |
| |
| """ |
| super().__init__(**kwargs) |
| self.tokenizer = EsmTokenizer.from_pretrained(tokenizer) |
| self.bias_feature = bias_feature |
| self.max_length = max_length |
| self.mask_struc_ratio = mask_struc_ratio |
| self.plddt_threshold = plddt_threshold |
| |
| def __getitem__(self, index): |
| data = json.loads(self._get(index)) |
| seq = data['seq'] |
| |
| |
| if self.mask_struc_ratio is not None: |
| tokens = self.tokenizer.tokenize(seq) |
| mask_candi = [i for i, t in enumerate(tokens) if t[-1] != "#"] |
| |
| |
| mask_num = int(len(mask_candi) * self.mask_struc_ratio) |
| mask_idx = np.random.choice(mask_candi, mask_num, replace=False) |
| for i in mask_idx: |
| tokens[i] = tokens[i][:-1] + "#" |
| |
| seq = "".join(tokens) |
| |
| |
| if self.plddt_threshold is not None: |
| plddt = data["plddt"] |
| tokens = self.tokenizer.tokenize(seq) |
| seq = "" |
| for token, score in zip(tokens, plddt): |
| if score < self.plddt_threshold: |
| seq += token[:-1] + "#" |
| else: |
| seq += token |
| |
| tokens = self.tokenizer.tokenize(seq)[:self.max_length] |
| seq = " ".join(tokens) |
| |
| coords = data['coords'][:self.max_length] if self.bias_feature else None |
| |
| label = data['label'] |
| if isinstance(label, str): |
| label = json.loads(label) |
| |
| return seq, label, coords |
|
|
| def collate_fn(self, batch): |
| seqs, labels, coords = zip(*batch) |
|
|
| model_inputs = self.tokenizer.batch_encode_plus(seqs, return_tensors='pt', padding=True) |
| inputs = {"inputs": model_inputs} |
| |
| if self.bias_feature: |
| inputs['structure_info'] = (coords,) |
|
|
| labels = {"labels": torch.tensor(labels, dtype=torch.long)} |
| |
| return inputs, labels |
| |
| def __len__(self): |
| return int(self._get('length')) |
|
|
