Update app.py

app.py

@@ -1,18 +1,28 @@
-import gradio as gr
-import subprocess
-
-def run_training():
-    # Execute train.py and capture its output
-    result = subprocess.run(["python", "train.py"], capture_output=True, text=True)
-    return result.stdout + "\n" + result.stderr
-
-# Create a Gradio interface with no inputs and a text output for logs
-iface = gr.Interface(
-    fn=run_training,
-    inputs=[],
-    outputs="text",
-    title="LLaMA LoRA Fine-Tuning",
-    description="Click the button below to start fine-tuning the LLaMA 3.2 3B Instruct model using PEFT/LoRA."
-)
-
-iface.launch()
+import gradio as gr
+import subprocess
+import os
+
+def run_training():
+    # Run train.py 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"
+    else:
+        result_msg = "\n❌ model.zip not found. Check logs below for issues.\n"
+
+    return result.stdout + "\n" + result.stderr + result_msg
+
+iface = gr.Interface(
+    fn=run_training,
+    inputs=[],
+    outputs="text",
+    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()