fix target context to Lenselink due to trained checkpoint

src/dataset.py

@@ -39,16 +39,16 @@ class DrugRetrieval(Dataset):
         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.'
+        assert os.path.exists(f'data/Lenselink/processed/{target_encoder}_encoding_train.pickle')), 'Context target embeddings not available.'
 
         # Drugs
-        emb_dict = self.get_embeddings(encoder_name=drug_encoder)
+        emb_dict = self.get_drug_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)
+        context = self.get_target_embeddings(encoder_name=target_encoder)
         self.context = self.standardize(embeddings=context)
 
         # Query target
@@ -71,8 +71,13 @@ class DrugRetrieval(Dataset):
     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:
+    def get_drug_embeddings(self, encoder_name):
+        with open(os.path.join(self.data_path, f'processed/{encoder_name}_encoding.pickle'), 'rb') as handle:
+            embeddings = pickle.load(handle)
+        return embeddings
+
+    def get_target_embeddings(self, encoder_name):
+        with open(f'data/Lenselink/processed/{encoder_name}_encoding_train.pickle'), 'rb') as handle:
             embeddings = pickle.load(handle)
         return embeddings