File size: 1,194 Bytes
3b99abb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 | import json
import torch
from .data_interface import register_dataset
from .lmdb_dataset import LMDBDataset
@register_dataset
class MutationZeroShotDataset(LMDBDataset):
"""
Dataset that deals with mutation data for zero-shot prediction
"""
def __init__(self, **kwargs):
"""
Args: **kwargs: other arguments for LMDBDataset
"""
super().__init__(**kwargs)
def __getitem__(self, index):
data = json.loads(self._get(index))
return data["seq"], data["mut_info"], data["fitness"]
def __len__(self):
return int(self._get("length"))
def collate_fn(self, batch):
seqs, mut_info, fitness = zip(*batch)
plddt = self._get("plddt")
if plddt is not None:
plddt = json.loads(plddt)
inputs = {"wild_type": self._get("wild_type"),
"seqs": seqs,
"mut_info": mut_info,
"structure_content": self._get("structure_content"),
"structure_type": self._get("structure_type"),
"plddt": plddt}
labels = {"labels": torch.Tensor(fitness)}
return inputs, labels |