import gradio as gr

def generate_chatml(messages):
    messages = messages.split("\n")
    chatml = ""
    for i, message in enumerate(messages):
        if message.strip():
            role = "<|im_start|>user" if i % 2 == 0 else "<|im_start|>assistant"
            chatml += f"{role}\n{message.strip()}\n<|im_end|>\n"
    return chatml.strip()

with gr.Blocks() as demo:
    gr.Markdown("# ChatML Template Generator\n## This will generate a prompt for you to enter to ChatGPT.\n### Enter messages alternately for user and assistant.")
    
    with gr.Row():
        with gr.Column():
            messages_input = gr.Textbox(label="Messages (alternate user and assistant, one per line)", lines=10, value="I love you!\nI love my car.\nI love that for you.")
    
    generate_btn = gr.Button("Generate ChatML")
    output = gr.Textbox(label="Generated ChatML", lines=20)
    
    generate_btn.click(
        generate_chatml, 
        inputs=messages_input, 
        outputs=output
    )

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