Update app.py

app.py

@@ -43,10 +43,11 @@ for kn in np.linspace(0.2, 2, 100):
 
 def plot_curve(kn, kd):
     fig = plt.figure()
-    plt.plot(kns, overheads)
-    plt.scatter([kn], [kd])
+    plt.plot(kns, overheads, color="black")
+    plt.scatter([kn], [compute_overhead(kn, kd)*100], marker="D", markerfacecolor="red", markeredgecolor="black", label="You are here!")
     plt.xlabel("Fraction of compute optimal model size")
     plt.ylabel("Compute overhead (%)")
+    plt.legend(loc="best")
     return fig
 
 
@@ -63,7 +64,17 @@ def compute(N, D):
     
     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}"""
+    text = f"""\
+## Compute:
+Compute budget (TFLOPs): {C:.2E}
+
+## Chinchilla optimal:
+Optimal model size:\t\t {N_opt/Bn:.2f}B
+Optimal datset size (tokens):\t {D_opt/Bn:.2f}
+
+## Your setting trade-off:
+Training compute overhead (%):\t {100*compute_overhead(kn, kd):.2f}
+Inference cost fraction (%):\t {kn*100:.2f}"""
     return text, fig
 
 with gr.Blocks() as demo: