dev

app.py

@@ -28,8 +28,9 @@ model_save_path = '/data/lora_finetuned_model'  # Specify your desired save path
 model.save_pretrained(model_save_path)
 '''
 
-def fine_tune_model(model, dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
+def fine_tune_model(dataset_name, hub_id, api_key, num_epochs, batch_size, lr, grad):
     try:
+        model = AutoModelForSeq2SeqLM.from_pretrained("google/t5-efficient-tiny-nh8")
         torch.cuda.empty_cache()
         torch.nn.CrossEntropyLoss()
         #rouge_metric = evaluate.load("rouge", cache_dir='/data/cache')
@@ -335,7 +336,7 @@ except Exception as e:
 # Create Gradio interface
 try:
     iface = gr.Interface(
-        fn=run_train,
+        fn=fine_tune_model,
         inputs=[
             gr.Textbox(label="Dataset Name (e.g., 'imdb')"),
             gr.Textbox(label="HF hub to push to after training"),