Update app.py

app.py

@@ -2,27 +2,22 @@ import gradio as gr
 import subprocess
 import os
 
-def run_training():
-    # Run train.py and capture logs
+def run_training_and_return_zip():
+    # Run training and capture logs
     result = subprocess.run(["python", "train.py"], capture_output=True, text=True)
-    
-    # Check for zipped model
+
     zip_path = "/home/user/app/model.zip"
     if os.path.exists(zip_path):
-        result_msg = "\n✅ Model was successfully zipped to model.zip.\n"
-        result_msg += "📥 You can download it from the Files tab or this URL:\n"
-        result_msg += "https://huggingface.co/spaces/darrenphodgson76/SmolLM2-1.7B-Instruct-Bussiness-Analysis/+/blob/main/model.zip\n"
+        logs = result.stdout + "\n" + result.stderr + "\n✅ Training complete. Download your model below."
+        return logs, zip_path
     else:
-        result_msg = "\n❌ model.zip not found. Check logs below for issues.\n"
-
-    return result.stdout + "\n" + result.stderr + result_msg
+        logs = result.stdout + "\n" + result.stderr + "\n❌ model.zip not found. Check logs for issues."
+        return logs, None
 
-iface = gr.Interface(
-    fn=run_training,
+gr.Interface(
+    fn=run_training_and_return_zip,
     inputs=[],
-    outputs="text",
+    outputs=["text", "file"],
     title="LLaMA LoRA Fine-Tuning",
-    description="Click the button below to start fine-tuning the model using PEFT/LoRA. The fine-tuned model will be saved as a ZIP for download."
-)
-
-iface.launch()
+    description="Click the button to fine-tune. After training, download the zipped model directly."
+).launch()