add utils file

app.py

@@ -1,10 +1,82 @@
 import streamlit as st
 from utils import memory_moe_mlp, memory_mlp_layer, memory_for_attention_layer
 
-st.title("Model Memory Usage Calculator")
 
-x = st.slider('Select a value')
-hidden_size = st.slider("The Hidden size (d_model | d)", min_value=128, max_value=2**20, step=128)
+def main():
+    st.title("LLM Model Memory Usage Calculator")
 
-st.write(f"The memory usage in case of Dense model is: {x}")
-st.write(f"The memory usage in case of MOE model with the same argument{hidden_size}")
+    st.sidebar.header("Model Parameters")
+    precession = st.sidebar.number_input("precession in Byte", min_value=1, max_value=4, value=2, step=2)
+    hidden_size = st.sidebar.number_input("Hidden Size", min_value=512, max_value=2 ** 16, value=4096, step=512)
+    num_heads = st.sidebar.number_input("Number of Attention Heads", min_value=4, max_value=128, value=32, step=4)
+    batch_size = st.sidebar.number_input("Batch Size", min_value=1, max_value=256, value=64, step=4)
+    seq_len = st.sidebar.number_input("Sequence Length", min_value=512, max_value=128000, value=2048, step=512)
+    intermediate_size = st.sidebar.number_input("Intermediate Size", min_value=1024, max_value=2 ** 18, value=11008,
+                                                step=128)
+    layers = st.sidebar.number_input("Number of Layers", min_value=6, max_value=48, value=30, step=1)
+    moe = st.sidebar.checkbox("Use Mixture of Experts (MOE)", value=False)
+
+    # Conditional rendering for MOE parameters
+    if moe:
+        top_k = st.sidebar.number_input("Number of Experts to use (Top K)", min_value=1, max_value=16, value=2, step=1)
+        num_experts = st.sidebar.number_input("Total Number of Experts", min_value=2, max_value=32, value=4, step=2)
+    else:
+        top_k = 2  # Default values if MOE is not used
+        num_experts = 4
+
+    attention_memory = memory_for_attention_layer(precession,
+                                                  seq_len,
+                                                  batch_size,
+                                                  hidden_size,
+                                                  num_heads)
+
+    dense_mlp_memory = memory_mlp_layer(precession,
+                                        seq_len,
+                                        batch_size,
+                                        hidden_size,
+                                        intermediate_size)
+
+    dense_model_memory = layers * (attention_memory + dense_mlp_memory) // (1024 ** 3)
+    space = st.empty()
+    space.markdown('<div style="height: 20px;"></div>', unsafe_allow_html=True)
+
+    st.markdown(
+        f'<div style="background-color: #b3f0ff; padding: 30px; border-radius: 5px;">'
+        f'<p style="font-weight: bold;">The memory requirement for this model is ~ {dense_model_memory} GB</p>'
+        f'</div>',
+        unsafe_allow_html=True
+    )
+    space = st.empty()
+    space.markdown('<div style="height: 40px;"></div>', unsafe_allow_html=True)
+
+    if moe:
+        moe_memory = memory_moe_mlp(precession,
+                                    seq_len,
+                                    batch_size,
+                                    hidden_size,
+                                    intermediate_size,
+                                    num_experts,
+                                    top_k)
+        moe_model = layers * (attention_memory + moe_memory) // (1024 ** 3)
+
+        st.markdown(
+            f'<div style="background-color: #99ff99; padding: 30px; border-radius: 5px;">'
+            f'<p style="font-weight: bold;">The memory requirement for the MOE model is ~ {moe_model} GB</p>'
+            f'</div>',
+            unsafe_allow_html=True
+        )
+
+    space = st.empty()
+    space.markdown('<div style="height: 40px;"></div>', unsafe_allow_html=True)
+
+    st.markdown(
+        f'<div style="background-color: #f0f0f0; padding: 30px; border-radius: 5px;">'
+        f'<p style="font-weight: bold;">For more information please read this article</p>'
+        f'<a href="https://medium.com/@khalil.hennara.247/llm-memory-usage-f62a007a509c">Article</a>'
+        f'</div>',
+        unsafe_allow_html=True
+    )
+
+
+if __name__ == "__main__":
+    main()