INNER CODE UNIT · Python
FABindDataSet
QizhiPei/FABind · FABind/fabind/data.py:9
class FABindDataSet(Dataset):
def __init__(self, root, data=None, protein_dict=None, compound_dict=None, proteinMode=0, compoundMode=1,
add_noise_to_com=None, pocket_radius=20, contactCutoff=8.0, predDis=True, args=None,
use_whole_protein=False, compound_coords_init_mode=None, seed=42, pre=None,
transform=None, pre_transform=None, pre_filter=None, noise_for_predicted_pocket=5.0, test_random_rotation=False, pocket_idx_no_noise=True, use_esm2_feat=False):
self.data = data
self.protein_dict = protein_dict
self.compound_dict = compound_dict
# this will call the process function to save the data, protein_dict and compound_dict
super().__init__(root, transform, pre_transform, pre_filter)
print(self.processed_paths)
self.data = torch.load(self.processed_paths[0])
self.compound_rdkit_coords = torch.load(self.processed_paths[3])
self.protein_dict = lmdb.open(self.processed_paths[1], readonly=True, max_readers=1, lock=False, readahead=False, meminit=False)
self.compound_dict = lmdb.open(self.processed_paths[2], readonly=True, max_readers=1, lock=False, readahead=False, meminit=False)
if use_esm2_feat:
self.protein_esm2_feat = lmdb.open(self.processed_paths[4], readonly=True, max_readers=1, lock=False, readahead=False, meminit=False)
self.compound_coords_init_mode = compound_coords_init_mode