add lenselink retrieval dataset

src/dataset.py

@@ -0,0 +1,92 @@
+
+import os
+import pickle
+import numpy as np
+from sklearn.preprocessing import StandardScaler
+
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset
+
+
+def collate_target(elem_dicts):
+    """ Data loading for interactions based on protein target. """
+    batch = {'pids': [], 'targets': torch.Tensor(), 'mids': [], 'drugs': torch.Tensor()}
+    labels = torch.Tensor()
+
+    for i, elem_dict in enumerate(elem_dicts):
+
+        labels = torch.cat((labels, torch.tensor(elem_dict['label'])), 0)
+
+        batch['mids'].append(elem_dict['mid'])
+        drug = torch.tensor(elem_dict['drug']).float().unsqueeze(0)
+        batch['drugs'] = drug if len(batch['drugs']) == 0 else torch.cat((batch['drugs'], drug), 0)
+                            
+        batch['pids'].append(elem_dict['pid'])
+        if i == 0:
+            batch['targets'] = torch.tensor(elem_dict['target']).float()
+
+    batch['drugs'] = batch['drugs'].unsqueeze(0)
+
+    return batch, labels
+
+
+class DrugRetrieval(Dataset):
+
+    def __init__(self, data_path, query_target, query_embedding, drug_encoder='CDDD', target_encoder='SeqVec'):
+        super(DrugRetrieval, self).__init__()
+        self.data_path = data_path
+        self.remove_batch = True
+
+        assert os.path.exists(os.path.join(self.data_path, f'processed/{drug_encoder}_encoding.pickle')), 'Drug embeddings not available.'
+        assert os.path.exists(os.path.join(self.data_path, f'processed/{target_encoder}_encoding_train.pickle')), 'Context target embeddings not available.'
+
+        # Drugs
+        emb_dict = self.get_embeddings(encoder_name=drug_encoder)
+        self.drug_ids = list(emb_dict.keys())
+        self.drug_embeddings = list(emb_dict.values())
+        
+        # Context
+        self.target_scaler = StandardScaler()
+        context = self.get_embeddings(encoder_name=target_encoder)
+        self.context = self.standardize(embeddings=context)
+
+        # Query target
+        self.query_target = query_target
+        self.query_embedding = self.target_scaler.transform([query_embedding.tolist()])
+
+    def __getitem__(self, item):
+        return {
+            'pid': self.query_target,
+            'target': self.query_embedding,
+            'mid': self.drug_ids[item],
+            'drug': self.drug_embeddings[item],
+            'label': [0],
+        }
+
+    def get_target_memory(self, exclude_pids=None):
+        memory = list(self.context.values())
+        return torch.tensor(np.stack(memory), dtype=torch.float32)
+
+    def __len__(self):
+        return len(self.drug_ids)
+
+    def get_embeddings(self, encoder_name):
+        with open(os.path.join(self.data_path, f'processed/{encoder_name}_encoding{"_train" if encoder_name == "SeqVec" else ""}.pickle'), 'rb') as handle:
+            embeddings = pickle.load(handle)
+        return embeddings
+
+    def standardize(self, embeddings):
+        split_embeddings = []
+        unique_ids = embeddings.keys()
+        for unique_id in unique_ids:
+            split_embeddings.append(embeddings[unique_id].tolist())
+
+        self.target_scaler.fit(split_embeddings)
+        scaled_embeddings = self.target_scaler.transform(split_embeddings)
+
+        new_dict = {}
+        for unique_id, emb in zip(unique_ids, scaled_embeddings):
+            new_dict[unique_id] = emb
+        return new_dict
+