Update utils.py

utils.py

@@ -129,7 +129,8 @@ class ReactionPredictionModel():
                     token = os.environ.get("TOKEN")
                     )
                 self.load_forward_model(candidate_models[model])
-        
+
+        print(self.forward_model.device, self.retro_model.device)
         string_template_path = hf_hub_download(candidate_models[list(candidate_models.keys())[0]], filename="string_template.json", token = os.environ.get("TOKEN"))
         string_template = json.load(open(string_template_path, 'r'))
         reactant_start_str = string_template['REACTANTS_START_STRING']
@@ -205,33 +206,9 @@ class ReactionPredictionModel():
         )
         self.forward_model.config.pad_token_id = self.tokenizer.pad_token_id
         self.forward_model.to("cuda")
-    
-    @spaces.GPU(duration=20)
-    def predict_single_smiles(self, smiles, task_type):
-        if task_type == "full_retro":
-            if "." in smiles:
-                return None
-        
-        task_type = "retrosynthesis" if task_type == "full_retro" else "synthesis"
-        # canonicalize the smiles
-        mol = Chem.MolFromSmiles(smiles)
-        if mol is None:
-            return None
-        smiles = Chem.MolToSmiles(mol)
-
-        smiles_list = [smiles]
-        task_type_list = [task_type]
-
-
-        df = pd.DataFrame({"src": smiles_list, "task_type": task_type_list})
-        test_dataset = Dataset.from_pandas(df)
-        # construct the dataloader
-        test_loader = torch.utils.data.DataLoader(
-            test_dataset,
-            batch_size=1,
-            collate_fn=self.data_collator,
-        )
 
+    @spaces.GPU(duration=20)
+    def predict(self, test_loader):
         predictions = []
         for i, batch in tqdm(enumerate(test_loader), total=len(test_loader), desc="Evaluating"):
             with torch.no_grad():
@@ -276,6 +253,36 @@ class ReactionPredictionModel():
                 predictions.append(canonized_smiles_list)
         
         rank, invalid_rate = compute_rank(predictions)
+
+        return rank
+    
+    def predict_single_smiles(self, smiles, task_type):
+        if task_type == "full_retro":
+            if "." in smiles:
+                return None
+        
+        task_type = "retrosynthesis" if task_type == "full_retro" else "synthesis"
+        # canonicalize the smiles
+        mol = Chem.MolFromSmiles(smiles)
+        if mol is None:
+            return None
+        smiles = Chem.MolToSmiles(mol)
+
+        smiles_list = [smiles]
+        task_type_list = [task_type]
+
+
+        df = pd.DataFrame({"src": smiles_list, "task_type": task_type_list})
+        test_dataset = Dataset.from_pandas(df)
+        # construct the dataloader
+        test_loader = torch.utils.data.DataLoader(
+            test_dataset,
+            batch_size=1,
+            collate_fn=self.data_collator,
+        )
+
+        rank = self.predict(test_loader)
+
         return rank