Update app.py

app.py

@@ -42,15 +42,16 @@ for kn in np.linspace(0.2, 2, 100):
     overheads.append(compute_overhead(kn, kd)*100)
 
 def plot_curve(kn, kd):
+    fig = plt.figure()
     plt.plot(kns, overheads)
     plt.scatter([kn], [kd])
     plt.xlabel("Fraction of compute optimal model size")
     plt.ylabel("Compute overhead (%)")
-
+    return fig
 
 
 def compute(N, D):
-    print(N, D)
+    
     C = to_flops(N * Bn, D * Bn)
     N_opt = n_opt(C)
     D_opt = d_opt(C)
@@ -58,19 +59,21 @@ def compute(N, D):
     kn = N/N_opt
     kd = compute_kd(kn)
 
-    plot_curve(kn, kd)
+    print(N, D, N_opt, D_opt, kn, kd)
+    
+    fig = plot_curve(kn, kd)
 
     text = f"""Compute budget (TFLOPs): {C:.2E}\n\nTraining compute overhead (%): {100*compute_overhead(kn, kd):.2f}\n\nInference cost fraction (%): {kn*100:.2f}"""
-    return text
+    return text, fig
 
 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)
+    plot = gr.Plot(value=plt)
     md = gr.Markdown("")
 
-    button.click(fn=compute, inputs=[N, D], outputs=[md])
+    button.click(fn=compute, inputs=[N, D], outputs=[md, plot])
 
 demo.launch()