import gradio as gr
import subprocess
import threading
import time

def run_training(model_name, progress=gr.Progress()):
    progress(0, desc="Starting Training...")
    
    command = ["python", "train.py", "--model_name", model_name]
    
    # Start the training process
    process = subprocess.Popen(
        command, 
        stdout=subprocess.PIPE, 
        stderr=subprocess.STDOUT, 
        text=True, 
        bufsize=1, 
        universal_newlines=True
    )

    # Stream the output
    output = ""
    for line in process.stdout:
        output += line
        # Update progress based on keywords in the output
        if "Starting training..." in line:
            progress(0.1, desc="Training Started...")
        elif "loss" in line:
             progress(0.5, desc=line.strip())
        
        yield output
    
    process.wait()
    
    if process.returncode == 0:
        progress(1, desc="Training Complete!")
        yield output + "\nTraining finished successfully!"
    else:
        yield output + f"\nTraining failed with error code {process.returncode}."


with gr.Blocks() as demo:
    gr.Markdown("# Llama Model Fine-Tuning")
    gr.Markdown("Select a model and click 'Start Fine-Tuning' to begin.")

    models = [
        'meta-llama/Meta-Llama-3-8B-Instruct',
        'meta-llama/Meta-Llama-3-70B-Instruct',
        'meta-llama/Llama-3.3-70B-Instruct'
    ]
    model_dropdown = gr.Dropdown(models, label="Select Model to Fine-Tune", value='meta-llama/Meta-Llama-3-8B-Instruct')
    
    start_button = gr.Button("Start Fine-Tuning")
    
    output_textbox = gr.Textbox(label="Training Log", lines=20, interactive=False)
    
    start_button.click(
        fn=run_training,
        inputs=[model_dropdown],
        outputs=[output_textbox]
    )

demo.launch()