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()