Update app.py

app.py

@@ -47,10 +47,9 @@ def plot_curve(kn, kd):
     plt.xlabel("Fraction of compute optimal model size")
     plt.ylabel("Compute overhead (%)")
 
-with gr.Blocks() as demo:
-    N = gr.Number(value=1, label="Model size (in B parameters)")
-    D = gr.Number(value=100, label="Dataset size (in B tokens")
 
+
+def compute(N, D):
     C = to_flops(N * Bn, D * Bn)
     N_opt = n_opt(C)
     D_opt = d_opt(C)
@@ -59,7 +58,17 @@ with gr.Blocks() as demo:
 
     plot_curve(kn, 100*overhead(kn, get_kd(kn)))
 
+    text = f"""Compute budget (TFLOPs): {C:.2E}\nTraining compute overhead (%): {100*overhead(kn, get_kd(kn)):.2f}\nInference cost fraction (%): {kn*100:.2f}"""
+    return text
+
+with gr.Blocks() as demo:
+    N = gr.Number(value=1, label="Model size (in B parameters)")
+    D = gr.Number(value=100, label="Dataset size (in B tokens")
+    button = gr.Button("Compute!")
+    
     gr.Plot(value=plt)
-    gr.Markdown(f"""Compute budget (TFLOPs): {C:.2E}\nTraining compute overhead (%): {100*overhead(kn, get_kd(kn)):.2f}\nInference cost fraction (%): {kn*100:.2f}""")
+    md = gr.Markdown(f"""{}""")
+
+    button.click(fn=, inputs=[N, D], ouptus=[md])
 
 demo.launch()