import gradio as gr


def deepmind_flops(n_layer, d_model, d_ff, d_attn, n_ctx, n_vocab, n_heads):
    embeddings = 2 * n_ctx * n_vocab * d_model
    attn_qkv = 2 * n_ctx * 3 * d_model * (d_attn * n_heads)
    attn_logits = 2 * n_ctx * n_ctx * (d_attn * n_heads)
    attn_softmax = 3 * n_heads * n_ctx * n_ctx
    attn_reduce = 2 * n_ctx * n_ctx * (d_attn * n_heads)
    attn_project = 2 * n_ctx * (d_attn * n_heads) * d_model
    ff = 2 * n_ctx * (d_model * d_ff + d_model * d_ff)
    logits = 2 * n_ctx * d_model * n_vocab

    return (
        embeddings,
        attn_qkv * n_layer,
        attn_logits * n_layer,
        attn_softmax * n_layer,
        attn_reduce * n_layer,
        attn_project * n_layer,
        ff * n_layer,
        logits,
    )


def calculator(n_layer, d_model, n_heads, n_vocab, n_ctx, ff_ratio, incl_embed):
    d_attn = d_model // n_heads
    if d_model % n_heads != 0:
        raise gr.Error("d_model must be divisible by n_heads")
    d_ff = d_model * ff_ratio

    flops_terms = deepmind_flops(
        n_layer, d_model, d_ff, d_attn, n_ctx, n_vocab, n_heads
    )

    if incl_embed:
        flops_per_sequence = sum(flops_terms)
    else:
        flops_per_sequence = sum(flops_terms[1:-1])

    return flops_per_sequence, flops_per_sequence / n_ctx


with gr.Blocks() as iface:
    with gr.Row():
        with gr.Column():
            n_layer = gr.Number(label="Number of layers (n_layer)")
            d_model = gr.Number(label="Model dimensions (d_model)")
            n_heads = gr.Number(label="Number of attention heads per layer (n_heads)")
            n_vocab = gr.Number(label="Vocabulary size (n_vocab)")
            n_ctx = gr.Number(label="Sequence length")
            ff_ratio = gr.Number(value=4, label="Feedforward ratio")
            incl_embed = gr.Checkbox(
                value=True, label="Include embedding and logits FLOPs"
            )

            btn = gr.Button(value="Submit", variant="primary")

        with gr.Column():
            flops_per_sequence = gr.Number(label="FLOPs per sequence")
            flops_per_token = gr.Number(label="FLOPs per token")

    btn.click(
        calculator,
        inputs=[n_layer, d_model, n_heads, n_vocab, n_ctx, ff_ratio, incl_embed],
        outputs=[flops_per_sequence, flops_per_token],
    )

    gr.Markdown("### GPT-3 model family examples")
    gr.Markdown(
        "In order are the 125M, 350M, 1.3B, 2.7B, 6.7B, 13B, 30B, 66B, and 175B parameter variants."
    )
    gr.Examples(
        [
            [12, 768, 12, 50257, 4096, 4, True],
            [24, 1024, 16, 50257, 4096, 4, True],
            [24, 2048, 32, 50257, 4096, 4, True],
            [32, 2560, 32, 50257, 4096, 4, True],
            [32, 4096, 32, 50257, 4096, 4, True],
            [40, 5120, 40, 50257, 4096, 4, True],
            [48, 7168, 56, 50257, 4096, 4, True],
            [64, 9216, 72, 50257, 4096, 4, True],
            [96, 12288, 96, 50257, 4096, 4, True],
        ],
        [n_layer, d_model, n_heads, n_vocab, n_ctx, ff_ratio, incl_embed],
        [flops_per_sequence, flops_per_token],
        calculator,
        cache_examples=False,
    )

iface.launch()