import gradio as gr
import subprocess
import os

def run_self_lengthen(model_path, instruct_count, max_iter):
    # Set environment variables
    os.environ["MODEL_PATH"] = model_path
    os.environ["INSTRUCT_COUNT"] = str(instruct_count)
    os.environ["MAX_ITER"] = str(max_iter)
    
    # Run the self-lengthen process
    try:
        subprocess.run(["/app/start.sh"], check=True)
        return "Training completed successfully! Check the results in the output directory."
    except subprocess.CalledProcessError as e:
        return f"Error during training: {str(e)}"

# Create Gradio interface
iface = gr.Interface(
    fn=run_self_lengthen,
    inputs=[
        gr.Textbox(label="Model Path", value="models/base_model"),
        gr.Number(label="Instruction Count", value=5000, minimum=100),
        gr.Number(label="Max Iterations", value=3, minimum=1)
    ],
    outputs=gr.Textbox(label="Status"),
    title="Self-Lengthen Training Interface",
    description="Train language models to generate longer texts using the Self-Lengthen approach."
)

iface.launch(server_name="0.0.0.0", server_port=7860)