Update utils.py

utils.py

@@ -180,6 +180,7 @@ class ReactionPredictionModel():
         )
 
         self.retro_model.to("cuda")
+        self.retro_model.eval()
     
     def load_forward_model(self, model_path):
         config = AutoConfig.from_pretrained(
@@ -206,21 +207,24 @@ class ReactionPredictionModel():
         )
         self.forward_model.config.pad_token_id = self.tokenizer.pad_token_id
         self.forward_model.to("cuda")
+        self.forward_model.eval()
 
     @spaces.GPU(duration=30)
     def predict(self, test_loader, task_type):
         predictions = []
         for i, batch in tqdm(enumerate(test_loader), total=len(test_loader), desc="Evaluating"):
-            with torch.no_grad():
-                generation_prompts = batch['generation_prompts'][0]
-                inputs = self.tokenizer(generation_prompts, return_tensors="pt", padding=True, truncation=True)
-                del inputs['token_type_ids']
-
-                if task_type == "retrosynthesis":
-                    self.retro_model.to("cuda")
-                    inputs = {k: v.to(self.retro_model.device) for k, v in inputs.items()}
-                    print(inputs)
-                    print(self.retro_model.device)
+            
+            generation_prompts = batch['generation_prompts'][0]
+            inputs = self.tokenizer(generation_prompts, return_tensors="pt", padding=True, truncation=True)
+            del inputs['token_type_ids']
+
+            if task_type == "retrosynthesis":
+                self.retro_model.to("cuda")
+                self.retro_model.eval()
+                inputs = {k: v.to(self.retro_model.device) for k, v in inputs.items()}
+                print(inputs)
+                print(self.retro_model.device)
+                with torch.no_grad():
                     outputs = self.retro_model.generate(**inputs, max_length=512, num_return_sequences=10,
                                    do_sample=False, num_beams=10,
                                    eos_token_id=self.tokenizer.eos_token_id,
@@ -228,11 +232,13 @@ class ReactionPredictionModel():
                                    pad_token_id=self.tokenizer.pad_token_id,
                                    length_penalty=0.0,
                                    )
-                else:
-                    self.forward_model.to("cuda")
-                    inputs = {k: v.to(self.forward_model.device) for k, v in inputs.items()}
-                    print(inputs)
-                    print(self.forward_model.device)
+            else:
+                self.forward_model.to("cuda")
+                self.forward_model.eval()
+                inputs = {k: v.to(self.forward_model.device) for k, v in inputs.items()}
+                print(inputs)
+                print(self.forward_model.device)
+                with torch.no_grad():
                     outputs = self.forward_model.generate(**inputs, max_length=512, num_return_sequences=10,
                                    do_sample=False, num_beams=10,
                                    eos_token_id=self.tokenizer.eos_token_id,
@@ -241,22 +247,22 @@ class ReactionPredictionModel():
                                    length_penalty=0.0,
                                    )
 
-                print(outputs)
-                original_smiles_list = self.tokenizer.batch_decode(outputs[:, len(inputs['input_ids'][0]):],
-                                                              skip_special_tokens=True) 
-                original_smiles_list = map(lambda x: x.replace(" ", ""), original_smiles_list)
-                # canonize the SMILES
-                canonized_smiles_list = []
-                temp = []
-                for original_smiles in original_smiles_list:
-                    temp.append(original_smiles)
-                    try:
-                        canonized_smiles_list.append(Chem.MolToSmiles(Chem.MolFromSmiles(original_smiles)))
-                    except:
-                        canonized_smiles_list.append("")
-                #canonized_smiles_list = \
-                #['N#Cc1ccsc1Nc1cc(F)c(F)cc1[N+](=O)[O-]', 'N#Cc1ccsc1Nc1cc(F)c([N+](=O)[O-])cc1F', 'N#Cc1ccsc1Nc1cc(Cl)c(F)cc1[N+](=O)[O-]', 'N#Cc1cnsc1Nc1cc(F)c(F)cc1[N+](=O)[O-]', 'N#Cc1cc(F)c(F)cc1Nc1sccc1C#N', 'N#Cc1ccsc1Nc1cc(F)c(F)cc1[N+](=N)[O-]', 'N#Cc1cc(C#N)c(Nc2cc(F)c(F)cc2[N+](=O)[O-])s1', 'N#Cc1ccsc1Nc1c(F)c(F)cc(F)c1[N+](=O)[O-]', 'Nc1sccc1CNc1cc(F)c(F)cc1[N+](=O)[O-]', 'N#Cc1ccsc1Nc1ccc(F)cc1[N+](=O)[O-]']
-                predictions.append(canonized_smiles_list)
+            print(outputs)
+            original_smiles_list = self.tokenizer.batch_decode(outputs[:, len(inputs['input_ids'][0]):],
+                                                          skip_special_tokens=True) 
+            original_smiles_list = map(lambda x: x.replace(" ", ""), original_smiles_list)
+            # canonize the SMILES
+            canonized_smiles_list = []
+            temp = []
+            for original_smiles in original_smiles_list:
+                temp.append(original_smiles)
+                try:
+                    canonized_smiles_list.append(Chem.MolToSmiles(Chem.MolFromSmiles(original_smiles)))
+                except:
+                    canonized_smiles_list.append("")
+            #canonized_smiles_list = \
+            #['N#Cc1ccsc1Nc1cc(F)c(F)cc1[N+](=O)[O-]', 'N#Cc1ccsc1Nc1cc(F)c([N+](=O)[O-])cc1F', 'N#Cc1ccsc1Nc1cc(Cl)c(F)cc1[N+](=O)[O-]', 'N#Cc1cnsc1Nc1cc(F)c(F)cc1[N+](=O)[O-]', 'N#Cc1cc(F)c(F)cc1Nc1sccc1C#N', 'N#Cc1ccsc1Nc1cc(F)c(F)cc1[N+](=N)[O-]', 'N#Cc1cc(C#N)c(Nc2cc(F)c(F)cc2[N+](=O)[O-])s1', 'N#Cc1ccsc1Nc1c(F)c(F)cc(F)c1[N+](=O)[O-]', 'Nc1sccc1CNc1cc(F)c(F)cc1[N+](=O)[O-]', 'N#Cc1ccsc1Nc1ccc(F)cc1[N+](=O)[O-]']
+            predictions.append(canonized_smiles_list)
         
         rank, invalid_rate = compute_rank(predictions)
         print(predictions, rank)