Update app.py

app.py

@@ -1,16 +1,18 @@
 import streamlit as st
 
+def number_field(label, **args):
+  c1, c2 = st.columns([2, 4])
+  c1.write(label)
+  
+  return c2.number_input('', **args)
+
 st.header("Transformer parameters")
 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)
+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')