gradio fix

app.py

@@ -1,9 +1,7 @@
 import gradio as gr
-from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForSeq2SeqLM
-from datasets import load_dataset
 import traceback
 
-def fine_tune_model(model_name, dataset_name, hub_id): #, num_epochs, batch_size, lr, grad):
+def fine_tune_model(model_name, dataset_name, hub_id, num_epochs, batch_size, lr, grad):
     '''
     try:
         
@@ -72,7 +70,10 @@ try:
             gr.inputs.Textbox(label="Model Name (e.g., 'google/t5-efficient-tiny-nh8')"),
             gr.inputs.Textbox(label="Dataset Name (e.g., 'imdb')"),
             gr.inputs.Textbox(label="HF hub to push to after training"),
-
+            gr.inputs.Slider(minimum=1, maximum=10, default=3, label="Number of Epochs"),
+            gr.inputs.Slider(minimum=1, maximum=16, default=4, label="Batch Size"),
+            gr.inputs.Slider(minimum=1, maximum=100, default=50, label="Learning Rate (e-5)"),
+            gr.inputs.Slider(minimum=1, maximum=100, default=1, label="Gradient accumulation (e-1)"), 
         ],
         outputs="text",
         title="Fine-Tune Hugging Face Model",