Spaces:
Sleeping
Sleeping
File size: 1,842 Bytes
85ab89d |
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 47 48 49 50 51 52 53 54 55 56 57 |
import os
import sys
import torch
from torch.utils.data import Dataset
import json
import numpy as np
from torch.utils.data.dataloader import default_collate
import time
class ESMDataset(Dataset):
def __init__(self, pdb_root, ann_paths, chain="A"):
"""
protein (string): Root directory of protein (e.g. coco/images/)
ann_root (string): directory to store the annotation file
"""
self.pdb_root = pdb_root
self.annotation = json.load(open(ann_paths, "r"))
self.pdb_ids = {}
self.chain = chain
def __len__(self):
return len(self.annotation)
def __getitem__(self, index):
ann = self.annotation[index]
protein_embedding = '{}.pt'.format(ann["pdb_id"])
protein_embedding_path = os.path.join(self.pdb_root, protein_embedding)
protein_embedding = torch.load(protein_embedding_path, map_location=torch.device('cpu'))
protein_embedding.requires_grad = False
caption = ann["caption"]
return {
"text_input": caption,
"encoder_out": protein_embedding,
"chain": self.chain,
"pdb_id": ann["pdb_id"]
}
def collater(self, samples):
max_len_protein_dim0 = -1
for pdb_json in samples:
pdb_embeddings = pdb_json["encoder_out"]
shape_dim0 = pdb_embeddings.shape[0]
max_len_protein_dim0 = max(max_len_protein_dim0, shape_dim0)
for pdb_json in samples:
pdb_embeddings = pdb_json["encoder_out"]
shape_dim0 = pdb_embeddings.shape[0]
pad1 = ((0, max_len_protein_dim0 - shape_dim0), (0, 0), (0, 0))
arr1_padded = np.pad(pdb_embeddings, pad1, mode='constant', )
pdb_json["encoder_out"] = arr1_padded
return default_collate(samples) |