Update utils.py

utils.py

@@ -218,9 +218,10 @@ class MolecularPropertyPredictionModel():
             #self.lora_model = PeftModel.from_pretrained(self.base_model, adapter_id, token = os.environ.get("TOKEN"))
             #self.lora_model.to("cuda")
             #print(self.lora_model)
-
+        
         self.base_model.set_adapter(adapter_name)
         self.base_model.eval()
+        print(f"switch to {adapter_name} successfully")
 
         #if adapter_name not in self.apapter_scaler_path:
         #    self.apapter_scaler_path[adapter_name] = hf_hub_download(adapter_id, filename="scaler.pkl", token = os.environ.get("TOKEN"))
@@ -237,7 +238,7 @@ class MolecularPropertyPredictionModel():
         #    # handle error
         #    return "error"
     
-    @spaces.GPU(duration=10)
+    @spaces.GPU(duration=5)
     def predict(self, valid_df, task_type):
         test_dataset = Dataset.from_pandas(valid_df)
         # construct the dataloader
@@ -263,6 +264,34 @@ class MolecularPropertyPredictionModel():
             y_pred = self.scaler.inverse_transform(y_pred)
 
 
+        return y_pred
+
+    @spaces.GPU(duration=20)
+    def predict_long(self, valid_df, task_type):
+        test_dataset = Dataset.from_pandas(valid_df)
+        # construct the dataloader
+        test_loader = torch.utils.data.DataLoader(
+            test_dataset,
+            batch_size=16,
+            collate_fn=self.data_collator,
+        )
+
+        # predict
+        y_pred = []
+        for i, batch in tqdm(enumerate(test_loader), total=len(test_loader), desc="Evaluating"):
+            with torch.no_grad():
+                batch = {k: v.to(self.base_model.device) for k, v in batch.items()}
+                outputs = self.base_model(**batch)
+            if task_type == "regression": # TODO: check if the model is regression or classification
+                y_pred.append(outputs.logits.cpu().detach().numpy())
+            else:
+                y_pred.append((torch.sigmoid(outputs.logits) > 0.5).cpu().detach().numpy())
+        
+        y_pred = np.concatenate(y_pred, axis=0)
+        if task_type=="regression" and self.scaler is not None:
+            y_pred = self.scaler.inverse_transform(y_pred)
+
+
         return y_pred
     
     def predict_single_smiles(self, smiles, task_type):
@@ -293,7 +322,7 @@ class MolecularPropertyPredictionModel():
         valid_df_smiles = valid_df['smiles'].tolist()
 
         input_df = pd.DataFrame(valid_df_smiles, columns=['smiles'])
-        results = self.predict(input_df, task_type)
+        results = self.predict_long(input_df, task_type)
 
         # add the results to the dataframe
         df.loc[valid_idx, 'prediction'] = results