load the adapters firstly

app.py

@@ -19,7 +19,7 @@ def get_models():
 
 candidate_models = get_models()
 properties = list(candidate_models.keys())
-model = MolecularPropertyPredictionModel()
+model = MolecularPropertyPredictionModel(candidate_models)
 
 def get_description(property_name):
     return dataset_descriptions[property_name]

utils.py

@@ -14,7 +14,7 @@ import os
 import pickle
 from sklearn import preprocessing
 import json
-import spaces
+#import spaces
 
 from rdkit import RDLogger, Chem
 # Suppress RDKit INFO messages
@@ -145,7 +145,7 @@ class DataCollator(object):
         return data_dict
 
 class MolecularPropertyPredictionModel():
-    def __init__(self):
+    def __init__(self, candidate_models):
         self.adapter_name = None
 
         # we need to keep track of the paths of adapter scalers
@@ -166,10 +166,17 @@ class MolecularPropertyPredictionModel():
         self.base_model = AutoModelForSequenceClassification.from_pretrained(
             "ChemFM/ChemFM-3B",
             config=config,
-            device_map="cuda",
+            device_map="cpu",
             trust_remote_code=True,
             token = os.environ.get("TOKEN")
         )
+        #self.base_model.to("cuda")
+        # load the adapters firstly
+        for adapter_name in candidate_models:
+            adapter_id = candidate_models[adapter_name]
+            self.base_model.load_adapter(adapter_id, adapter_name=adapter_name)
+            self.apapter_scaler_path[adapter_name] = hf_hub_download(adapter_id, filename="scaler.pkl", token = os.environ.get("TOKEN"))
+
 
         # load the tokenizer
         self.tokenizer = AutoTokenizer.from_pretrained(
@@ -203,30 +210,35 @@ class MolecularPropertyPredictionModel():
             return "keep"
         # switch adapter
         try:
-            self.adapter_name = adapter_name
-            print(self.adapter_name, adapter_id)
-            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)
-            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"))
+            #self.adapter_name = adapter_name
+            #print(self.adapter_name, adapter_id)
+            #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)
+
+            #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"))
             if os.path.exists(self.apapter_scaler_path[adapter_name]):
                 self.scaler = pickle.load(open(self.apapter_scaler_path[adapter_name], "rb"))
             else:
                 self.scaler = None
+            
+            self.adapter_name = adapter_name
 
             return "switched"
         except Exception as e:
             # handle error
             return "error"
     
-    @spaces.GPU
+    #@spaces.GPU
     def predict(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=4,
+            batch_size=32,
             collate_fn=self.data_collator,
         )
         # predict
@@ -234,8 +246,8 @@ class MolecularPropertyPredictionModel():
         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.lora_model.device) for k, v in batch.items()}
-                outputs = self.lora_model(**batch)
+                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: