Update trainer.py

trainer.py

@@ -732,11 +732,12 @@ class Trainer(object):
             G2_path = os.path.join(self.inference_model, '{}-G2.ckpt'.format(self.submodel))
             self.G2.load_state_dict(torch.load(G2_path, map_location=lambda storage, loc: storage))           
         
-        
+
+        smiles_test = [line for line in open("data/chembl_test.smi", 'r').read().splitlines()]
         if self.submodel == "NoTarget":
-            drug_smiles = [line for line in open("data/chembl_test.smi", 'r').read().splitlines()]
+            smiles_train = [line for line in open("data/chembl_train.smi", 'r').read().splitlines()]
         else:
-            drug_smiles = [line for line in open("data/akt_train.smi", 'r').read().splitlines()]
+            smiles_train = [line for line in open("data/akt_train.smi", 'r').read().splitlines()]
 
         if self.submodel == "RL":
             drug_mols = [Chem.MolFromSmiles(smi) for smi in drug_smiles]
@@ -916,7 +917,8 @@ class Trainer(object):
             "Runtime (seconds)": round(et, 2),
             "Validity": f"{fraction_valid(metric_calc_dr)*100:.2f}%",
             "Uniqueness": f"{fraction_unique(metric_calc_dr)*100:.2f}%",
-            "Novelty": f"{novelty(metric_calc_dr, drug_smiles)*100:.2f}%",
+            "Novelty Train": f"{novelty(metric_calc_dr, smiles_train)*100:.2f}%",
+            "Novelty Test": f"{novelty(metric_calc_dr, smiles_test)*100:.2f}%"
         }
         # print("Validity: ", fraction_valid(metric_calc_dr), "\n")
         # print("Uniqueness: ", fraction_unique(metric_calc_dr), "\n")