INNER CODE UNIT · Python
processed_file_names
QizhiPei/FABind · FABind/fabind/data.py:43
def processed_file_names(self):
return ['data.pt', 'protein_1d_3d.lmdb', 'compound_LAS_edge_index.lmdb', 'compound_rdkit_coords.pt', 'esm2_t33_650M_UR50D.lmdb']
def len(self):
return len(self.data)
def get(self, idx):
line = self.data.iloc[idx]
pocket_com = line['pocket_com']
use_compound_com = line['use_compound_com']
use_whole_protein = line['use_whole_protein'] if "use_whole_protein" in line.index else self.use_whole_protein
group = line['group'] if "group" in line.index else 'train'
if group == 'train' and use_compound_com:
add_noise_to_com = self.add_noise_to_com
elif group == 'train' and not use_compound_com:
add_noise_to_com = self.noise_for_predicted_pocket
else:
add_noise_to_com = None