import gradio as gr

def generate_chatml(user_messages, assistant_messages):
    max_messages = max(len(user_messages), len(assistant_messages))

    chatml = ""
    
    for i in range(max_messages):
        if i < len(user_messages) and user_messages[i].strip():
            chatml += f"<im_start>user\n{user_messages[i].strip()}\n<|im_end>\n"
        
        if i < len(assistant_messages) and assistant_messages[i].strip():
            chatml += f"<im_start>assistant\n{assistant_messages[i].strip()}\n<|im_end>\n"
    return chatml[15:-11]

with gr.Blocks() as demo:
    gr.Markdown("# ChatML Template Generator")
    gr.Markdown("## This will generate a prompt for you to enter to ChatGPT.")
    
    with gr.Row():
        with gr.Column():
            user_input = gr.Textbox(label="User Messages (one per line)", lines=10)
        with gr.Column():
            assistant_input = gr.Textbox(label="Assistant Messages (one per line)", lines=10)
    
    generate_btn = gr.Button("Generate ChatML")
    output = gr.Textbox(label="Generated ChatML", lines=20)
    
    generate_btn.click(
        generate_chatml, 
        inputs=[user_input, assistant_input], 
        outputs=output
    )

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