app.py

import gradio as gr
import os
import subprocess
from pathlib import Path

def check_training_status():
    results_dir = Path("/app/results")
    if not results_dir.exists():
        return "Training hasn't started yet."
    
    iterations = len(list(results_dir.glob("iter_*")))
    return f"Completed {iterations} training iterations."

def start_training(model_path, instruct_count, max_iter):
    os.environ["MODEL_PATH"] = model_path
    os.environ["INSTRUCT_COUNT"] = str(instruct_count)
    os.environ["MAX_ITER"] = str(max_iter)
    
    try:
        subprocess.run(["bash", "run.sh"], 
                      check=True,
                      cwd="/app/qwen")
        return "Training completed successfully!"
    except subprocess.CalledProcessError as e:
        return f"Error during training: {str(e)}"

# Create the interface
with gr.Blocks() as iface:
    gr.Markdown("# Self-Lengthen Training Interface")
    
    with gr.Row():
        with gr.Column():
            model_path = gr.Textbox(
                label="Model Path",
                value="/app/models/base_model",
                info="Path to the base model"
            )
            instruct_count = gr.Number(
                label="Instruction Count",
                value=5000,
                minimum=100,
                info="Number of instructions to generate"
            )
            max_iter = gr.Number(
                label="Max Iterations",
                value=3,
                minimum=1,
                info="Number of training iterations"
            )
            train_btn = gr.Button("Start Training")
        
        with gr.Column():
            status_output = gr.Textbox(
                label="Status",
                value="Ready to start training...",
                interactive=False
            )
            refresh_btn = gr.Button("Refresh Status")
    
    train_btn.click(
        fn=start_training,
        inputs=[model_path, instruct_count, max_iter],
        outputs=status_output
    )
    
    refresh_btn.click(
        fn=check_training_status,
        inputs=None,
        outputs=status_output
    )

if __name__ == "__main__":
    iface.launch(server_name="0.0.0.0", server_port=7860)