dev

app.py

@@ -25,7 +25,7 @@ model = get_peft_model(model, lora_config)
 model.gradient_checkpointing_enable()   
 
 @spaces.GPU(duration=120)
-def fine_tune_model(dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
+def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
     try:
         torch.cuda.empty_cache()
         def compute_metrics(eval_pred):
@@ -145,11 +145,15 @@ def predict(text):
     predictions = outputs.logits.argmax(dim=-1)
     return predictions.item()
 '''
+
+def run_train(dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
+    result = fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad)
+    return result
 # Create Gradio interface
 try:
     model = AutoModelForSeq2SeqLM.from_pretrained('google/t5-efficient-tiny-nh8'.strip(), num_labels=2, force_download=True)
     iface = gr.Interface(
-        fn=fine_tune_model,
+        fn=run_train,
         inputs=[
             gr.Textbox(label="Dataset Name (e.g., 'imdb')"),
             gr.Textbox(label="HF hub to push to after training"),