import gradio as gr

def generate_chatml(user_messages, assistant_messages):
    user_messages = user_messages.split("\n")
    assistant_messages = assistant_messages.split("\n")
    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[17:-11]

with gr.Blocks() as demo:
    gr.Markdown("# ChatML Template Generator\n## This will generate a prompt for you to enter to ChatGPT.\n### There should be one more item in the user field for it to work.")
    
    with gr.Row():
        with gr.Column():
            user_input = gr.Textbox(label="User Messages (one per line)", lines=10, value="I love you!\nI love that for you.")
        with gr.Column():
            assistant_input = gr.Textbox(label="Assistant Messages (one per line)", lines=10, value="I love my car.")
    
    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()