import streamlit as st

def number_field(label, columns=None, **input_params):
    c1, c2 = st.beta_columns(columns or [1, 4])

    # Display field name with some alignment
    c1.markdown("##")
    c1.markdown(label)

    # Sets a default key parameter to avoid duplicate key errors
    input_params.setdefault("key", label)

    # Forward text input parameters
    return c2.number_input("", **input_params)

def key_value(key, value, columns=None):
    c1, c2 = st.beta_columns(columns or [2, 3])

    # Display field name with some alignment
    c1.markdown("##")
    c1.markdown(key)
    c2.markdown("##")
    c2.markdown(value)

st.header("Transformer parameters")
bs = number_field('Batch size: ', value=10)
h = number_field('Num heads: ', value=16)
d = number_field('Dimension: ', value=768)
n = number_field('Seq length: ', value=1024)

st.header('Query, Key, Value projection')

mha_flop = 2*bs*n*d*3*d
mha_bytes = 2*bs*n*d + 2*3*d*d + 2*bs*n*3*d

st.subheader("Multi-query Attention")
key_value("FLOP: ", str(mha_flop))
key_value("bytes: ", str(mha_bytes))
key_value("Arithm. intensity:", str(mha_flop/mha_bytes))

mqa_flop = 2*bs*n*d*(1+2/h)*d
mqa_bytes = 2*bs*n*d + 2*(2/h)*d*d + 2*bs*n*(2/h)*d

st.subheader("Multi-query Attention")
key_value("FLOP: ", str(mqa_flop))
key_value("bytes: ", str(mqa_bytes))
key_value("Arithm. intensity:", str(mqa_flop/mqa_bytes))

st.header('Attention')