Upload dti.py

deepscreen/data/dti.py

@@ -192,17 +192,20 @@ class DTIDataset(Dataset):
 
     def __getitem__(self, i):
         sample = self.df.loc[i]
-        return {
+        sample_dict = {
             'N': i,
             'X1': sample['X1'],
             'X1^': self.drug_featurizer(sample['X1^']),
-            'ID1': sample.get('ID1'),
+            # 'ID1': sample.get('ID1'),
             'X2': sample['X2'],
             'X2^': self.protein_featurizer(sample['X2']),
-            'ID2': sample.get('ID2'),
-            'Y': sample.get('Y'),
-            'ID^': sample.get('ID^'),
+            # 'ID2': sample.get('ID2'),
+            # 'Y': sample.get('Y'),
+            # 'ID^': sample.get('ID^'),
         }
+        optional_keys = ['ID1', 'ID2', 'ID^', 'Y']
+        sample_dict.update({key: sample[key] for key in optional_keys if sample.get(key) is not None})
+        return sample_dict
 
 
 class DTIDataModule(LightningDataModule):