import gradio as gr

def calculate_effective_batch_size(ctx_len, num_gpus, bsz, accum):
    eff_bsz_tokens = ctx_len * num_gpus * bsz * accum
    return "{:,}".format(eff_bsz_tokens)  # Format with commas for readability
with gr.Blocks() as demo:
    gr.Markdown("## Effective Batch Size Calculator")
    with gr.Row():
        ctx_len = gr.Number(label="Context Length", value=16000)
        num_gpus = gr.Number(label="Number of GPUs", value=2)
        bsz = gr.Number(label="Batch Size", value=2)
        accum = gr.Number(label="Gradient Accumulation Steps", value=32)
    output = gr.Textbox(label="Effective Batch Size in Tokens")  # Changed to Textbox for formatted output
    btn = gr.Button("Calculate")
    btn.click(
        fn=calculate_effective_batch_size,
        inputs=[ctx_len, num_gpus, bsz, accum],
        outputs=output
    )

demo.launch()