Update app.py

app.py

@@ -1,32 +1,16 @@
 import streamlit as st
 
-def number_field(label, columns=None, **input_params):
-    c1, c2 = st.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.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)
+col1, col2 = st.columns([2, 4])
+
+col1.write('Batch size: ')
+bs = col2.number_input(value=10)
+col1.write('Num heads:')
+h = col2.number_input(value=16)
+col1.write('Dimension:')
+d = col2.number_input(value=768)
+col1.write('Seq length:')
+n = col2.number_input(value=1024)
 
 st.header('Query, Key, Value projection')
 
@@ -46,8 +30,12 @@ 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))
+c1, c2 = st.columns([2, 3])
+c1.write("FLOP:")
+c2.write(str(mqa_flop))
+c1.write("Bytes: ")
+c2.write(str(mqa_bytes))
+c1.write("Arithm. intensity:")
+c2.write(str(mqa_flop/mqa_bytes))
 
 st.header('Attention')