Update app.py

app.py

@@ -82,7 +82,7 @@ class WeightedTrainer(Trainer):
         return (loss, outputs) if return_outputs else loss
         
 # fine-tuning function
-def train_function_no_sweeps(base_model_path, train_dataset, test_dataset):
+def train_function_no_sweeps(base_model_path):   #, train_dataset, test_dataset):
     
     # Set the LoRA config
     config = {
@@ -170,7 +170,14 @@ def train_function_no_sweeps(base_model_path, train_dataset, test_dataset):
     tokenizer.save_pretrained(save_path)
 
     return save_path
-    
+
+# Constants & Globals
+MODEL_OPTIONS = [
+    "facebook/esm2_t6_8M_UR50D",
+    "facebook/esm2_t12_35M_UR50D",
+    "facebook/esm2_t33_650M_UR50D",
+]  # models users can choose from
+
 # Load the data from pickle files (replace with your local paths)
 with open("./datasets/train_sequences_chunked_by_family.pkl", "rb") as f:
     train_sequences = pickle.load(f)
@@ -198,6 +205,7 @@ test_labels = truncate_labels(test_labels, max_sequence_length)
 train_dataset = Dataset.from_dict({k: v for k, v in train_tokenized.items()}).add_column("labels", train_labels)
 test_dataset = Dataset.from_dict({k: v for k, v in test_tokenized.items()}).add_column("labels", test_labels)
 
+'''
 # Compute Class Weights
 classes = [0, 1]  
 flat_train_labels = [label for sublist in train_labels for label in sublist]
@@ -248,10 +256,46 @@ saved_path = train_function_no_sweeps(base_model_path,train_dataset, test_datase
 
 # debug result
 dubug_result = saved_path  #predictions  #class_weights
+'''
 
 demo = gr.Blocks(title="DEMO FOR ESM2Bind")
 
 with demo:
     gr.Markdown("# DEMO FOR ESM2Bind")
-    gr.Textbox(dubug_result)
+    #gr.Textbox(dubug_result)
+    
+    with gr.Tab("Finetune Pre-trained Model"):
+        gr.Markdown("## Finetune Pre-trained Model")
+        with gr.Column():
+            gr.Markdown("## Load Inputs & Select Parameters")
+            gr.Markdown(
+                """ Pick a dataset, a model & adjust params (_optional_), and press **Finetune Pre-trained Model!"""
+            )
+            with gr.Row():   
+                with gr.Column(scale=0.5, variant="compact"):
+                    base_model_name = gr.Dropdown(
+                        choices=MODEL_OPTIONS,
+                        value=MODEL_OPTIONS[0],
+                        label="Base Model Name",
+                        interactive = True,
+                    )   
+                    finetune_button = gr.Button(
+                        value="Finetune Pre-trained Model", 
+                        interactive=True, 
+                        variant="primary",
+                    )
+                    finetune_output_text = gr.Textbox(
+                        lines=1,
+                        max_lines=12,
+                        label="Finetune Status",
+                        placeholder="Finetune Status Shown Here",
+                    )    
+    # Tab "Finetune Pre-trained Model" actions
+    finetune_button.click(
+        fn = train_function_no_sweeps,
+        inputs=[base_model_name], #finetune_dataset_name],
+        outputs = [finetune_output_text],
+    )
+           
+
 demo.launch()