import gradio as gr
import subprocess
import os

def run_training_and_return_zip():
    # Run training and capture logs
    result = subprocess.run(["python", "train.py"], capture_output=True, text=True)

    zip_path = "/home/user/app/model.zip"
    if os.path.exists(zip_path):
        logs = result.stdout + "\n" + result.stderr + "\n✅ Training complete. Download your model below."
        return logs, zip_path
    else:
        logs = result.stdout + "\n" + result.stderr + "\n❌ model.zip not found. Check logs for issues."
        return logs, None

gr.Interface(
    fn=run_training_and_return_zip,
    inputs=[],
    outputs=["text", "file"],
    title="LLaMA LoRA Fine-Tuning",
    description="Click the button to fine-tune. After training, download the zipped model directly."
).launch()