import os
import gradio as gr

from train_seallm_khm_sum import main as train_main

HAS_TRAINED_FLAG = "TRAIN_DONE.flag"

def maybe_run_training():
    # Avoid retraining every restart
    if os.path.exists(HAS_TRAINED_FLAG):
        return "Training already finished."

    train_main()
    with open(HAS_TRAINED_FLAG, "w") as f:
        f.write("ok")
    return "Training finished. LoRA saved to ./seallm-khm-sum-lora"

def ui_train():
    msg = maybe_run_training()
    return msg

with gr.Blocks() as demo:
    gr.Markdown("## SeaLLM Khmer Summarization Fine-tune")
    btn = gr.Button("Run training (QLoRA on khm subset)")
    out = gr.Textbox(lines=4)
    btn.click(ui_train, outputs=out)

if __name__ == "__main__":
    demo.launch()