freeze other components when the prediction is processing

app.py

@@ -25,65 +25,75 @@ def get_description(property_name):
     return dataset_descriptions[property_name]
 
 def predict_single_label(smiles, property_name):
-    adapter_id = candidate_models[property_name]
-    info = model.swith_adapter(property_name, adapter_id)
-
-    running_status = None
-    if info == "keep":
-        running_status = "Adapter is the same as the current one"
-        #print("Adapter is the same as the current one")
-    elif info == "switched":
-        running_status = "Adapter is switched successfully"
-        #print("Adapter is switched successfully")
-    elif info == "error":
-        running_status = "Adapter is not found"
-        #print("Adapter is not found")
-        return "NA", running_status
-    else:
-        running_status = "Unknown error"
-        return "NA", running_status
+    try:
+        adapter_id = candidate_models[property_name]
+        info = model.swith_adapter(property_name, adapter_id)
+
+        running_status = None
+        if info == "keep":
+            running_status = "Adapter is the same as the current one"
+            #print("Adapter is the same as the current one")
+        elif info == "switched":
+            running_status = "Adapter is switched successfully"
+            #print("Adapter is switched successfully")
+        elif info == "error":
+            running_status = "Adapter is not found"
+            #print("Adapter is not found")
+            return "NA", running_status
+        else:
+            running_status = "Unknown error"
+            return "NA", running_status
     
-    #prediction = model.predict(smiles, property_name, adapter_id)
-    prediction = model.predict_single_smiles(smiles, task_types[property_name])
-    if prediction is None:
-        return "NA", "Invalid SMILES string"
+        #prediction = model.predict(smiles, property_name, adapter_id)
+        prediction = model.predict_single_smiles(smiles, task_types[property_name])
+        if prediction is None:
+            return "NA", "Invalid SMILES string"
     
-    # if the prediction is a float, round it to 3 decimal places
-    if isinstance(prediction, float):
-        prediction = round(prediction, 3)
+        # if the prediction is a float, round it to 3 decimal places
+        if isinstance(prediction, float):
+            prediction = round(prediction, 3)
+    except Exception as e:
+        # no matter what the error is, we should return
+        print(e)
+        return "NA", "Prediction failed"
 
     return prediction, "Prediction is done"
 
 def predict_file(file, property_name):
-    adapter_id = candidate_models[property_name]
-    info = model.swith_adapter(property_name, adapter_id)
-
-    running_status = None
-    if info == "keep":
-        running_status = "Adapter is the same as the current one"
-        #print("Adapter is the same as the current one")
-    elif info == "switched":
-        running_status = "Adapter is switched successfully"
-        #print("Adapter is switched successfully")
-    elif info == "error":
-        running_status = "Adapter is not found"
-        #print("Adapter is not found")
-        return None, None, file, running_status
-    else:
-        running_status = "Unknown error"
-        return None, None, file, running_status
+    try:
+        adapter_id = candidate_models[property_name]
+        info = model.swith_adapter(property_name, adapter_id)
+
+        running_status = None
+        if info == "keep":
+            running_status = "Adapter is the same as the current one"
+            #print("Adapter is the same as the current one")
+        elif info == "switched":
+            running_status = "Adapter is switched successfully"
+            #print("Adapter is switched successfully")
+        elif info == "error":
+            running_status = "Adapter is not found"
+            #print("Adapter is not found")
+            return None, None, file, running_status
+        else:
+            running_status = "Unknown error"
+            return None, None, file, running_status
     
-    df = pd.read_csv(file)
-    # we have already checked the file contains the "smiles" column
-    df = model.predict_file(df, task_types[property_name])
-    # we should save this file to the disk to be downloaded
-    # rename the file to have "_prediction" suffix
-    prediction_file = file.replace(".csv", "_prediction.csv") if file.endswith(".csv") else file.replace(".smi", "_prediction.csv")
-    print(file, prediction_file)
-    # save the file to the disk
-    df.to_csv(prediction_file, index=False)
+        df = pd.read_csv(file)
+        # we have already checked the file contains the "smiles" column
+        df = model.predict_file(df, task_types[property_name])
+        # we should save this file to the disk to be downloaded
+        # rename the file to have "_prediction" suffix
+        prediction_file = file.replace(".csv", "_prediction.csv") if file.endswith(".csv") else file.replace(".smi", "_prediction.csv")
+        print(file, prediction_file)
+        # save the file to the disk
+        df.to_csv(prediction_file, index=False)
+    except Exception as e:
+        # no matter what the error is, we should return
+        print(e)
+        return None, None, gr.update(visible=False), file, "Prediction failed"
     
-    return gr.update(visible=False), gr.DownloadButton(label="Download", value=prediction_file, visible=True), prediction_file, "Prediction is done"
+    return gr.update(visible=False), gr.DownloadButton(label="Download", value=prediction_file, visible=True), gr.update(visible=False), prediction_file, "Prediction is done"
 
 def validate_file(file):
     try:
@@ -166,18 +176,57 @@ def build_inference():
                        file_types=[".smi", ".csv"], height=300)
         predict_file_button = gr.Button("Predict", size='sm', visible=False)
         download_button = gr.DownloadButton("Download", size='sm', visible=False)
+        stop_button = gr.Button("Stop", size='sm', visible=False)
 
         # dropdown change event
         dropdown.change(get_description, inputs=dropdown, outputs=description_box)
         # predict single button click event
-        predict_single_smiles_button.click(predict_single_label, inputs=[textbox, dropdown], outputs=[prediction, running_terminal_label])
+        predict_single_smiles_button.click(lambda:(gr.update(interactive=False), 
+                                                   gr.update(interactive=False),
+                                                   gr.update(interactive=False),
+                                                   gr.update(interactive=False),
+                                                   gr.update(interactive=False),
+                                                   gr.update(interactive=False),
+                                                   gr.update(interactive=False),
+                                                   gr.update(interactive=False),
+                                                   ) , outputs=[dropdown, textbox, predict_single_smiles_button, predict_file_button, download_button, stop_button, input_file, running_terminal_label])\
+                                                   .then(predict_single_label, inputs=[textbox, dropdown], outputs=[prediction, running_terminal_label])\
+                                                   .then(lambda:(gr.update(interactive=True), 
+                                                                 gr.update(interactive=True),
+                                                                 gr.update(interactive=True),
+                                                                 gr.update(interactive=True),
+                                                                 gr.update(interactive=True),
+                                                                 gr.update(interactive=True),
+                                                                 gr.update(interactive=True),
+                                                                 gr.update(interactive=True),
+                                                                 ) , outputs=[dropdown, textbox, predict_single_smiles_button, predict_file_button, download_button, stop_button, input_file, running_terminal_label])
         # input file upload event
         file_status = gr.State()
         input_file.upload(fn=validate_file, inputs=input_file, outputs=[file_status, input_file, predict_file_button, download_button]).success(raise_error, inputs=file_status, outputs=file_status)
         # input file clear event
         input_file.clear(fn=clear_file, inputs=[download_button], outputs=[predict_file_button, download_button, input_file])
         # predict file button click event
-        predict_file_button.click(predict_file, inputs=[input_file, dropdown], outputs=[predict_file_button, download_button, input_file, running_terminal_label])
+        predict_file_event = predict_file_button.click(lambda:(gr.update(interactive=False), 
+                                                               gr.update(interactive=False),
+                                                               gr.update(interactive=False),
+                                                               gr.update(interactive=False, visible=True),
+                                                               gr.update(interactive=False),
+                                                               gr.update(interactive=True, visible=False),
+                                                               gr.update(interactive=False),
+                                                               gr.update(interactive=False),
+                                                               ) , outputs=[dropdown, textbox, predict_single_smiles_button, predict_file_button, download_button, stop_button, input_file, running_terminal_label])\
+                                                               .then(predict_file, inputs=[input_file, dropdown], outputs=[predict_file_button, download_button, stop_button, input_file, running_terminal_label])\
+                                                               .then(lambda:(gr.update(interactive=True), 
+                                                                             gr.update(interactive=True),
+                                                                             gr.update(interactive=True),
+                                                                             gr.update(interactive=True),
+                                                                             gr.update(interactive=True),
+                                                                             gr.update(interactive=True),
+                                                                             gr.update(interactive=True),
+                                                                             gr.update(interactive=True),
+                                                                             ) , outputs=[dropdown, textbox, predict_single_smiles_button, predict_file_button, download_button, stop_button, input_file, running_terminal_label])
+        # stop button click event
+        #stop_button.click(fn=None, inputs=None, outputs=None, cancels=[predict_file_event])
         
     return demo
 

utils.py

@@ -201,7 +201,7 @@ class MolecularPropertyPredictionModel():
             self.apapter_scaler_path[adapter_name] = hf_hub_download(adapter_id, filename="scaler.pkl", token = os.environ.get("TOKEN"))
 
         #self.base_model.to("cuda")
-        print(self.base_model)
+        #print(self.base_model)
     
     def swith_adapter(self, adapter_name, adapter_id):
         # return flag:
@@ -220,6 +220,7 @@ class MolecularPropertyPredictionModel():
             #print(self.lora_model)
 
             self.base_model.set_adapter(adapter_name)
+            self.base_model.eval()
 
             #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"))
@@ -244,8 +245,8 @@ class MolecularPropertyPredictionModel():
             batch_size=16,
             collate_fn=self.data_collator,
         )
-        # predict
 
+        # predict
         y_pred = []
         for i, batch in tqdm(enumerate(test_loader), total=len(test_loader), desc="Evaluating"):
             with torch.no_grad():