Update app.py

app.py

@@ -1,10 +1,32 @@
 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 = st.number_input('Batch size: ', value=10)
-h = st.number_input('Num heads: ', value=16)
-d = st.number_input('Dimension: ', value=768)
-n = st.number_input('Seq length: ', value=1024)
+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')
 
@@ -12,16 +34,16 @@ 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")
-st.text("FLOP: " + str(mha_flop))
-st.text("bytes: " + str(mha_bytes))
-st.text("Arithm. intensity:" + str(mha_flop/mha_bytes))
+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")
-st.write("FLOP: " + str(mqa_flop))
-st.write("bytes: " + str(mqa_bytes))
-st.write("Arithm. intensity:" + str(mqa_flop/mqa_bytes))
+key_value("FLOP: ", str(mqa_flop))
+key_value("bytes: ", str(mqa_bytes))
+key_value("Arithm. intensity:", str(mqa_flop/mqa_bytes))
 
 st.header('Attention')