Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import matplotlib.pyplot as plt
 
 
-def plot_forecast(num_param, precision, batch_size, seq_len):
+def plot_forecast(num_param, precision, grad_ckpt, batch_size, seq_len):
     # Convert number (input as B)
     num_param = float(num_param) * 1e9
 
@@ -19,6 +19,9 @@ def plot_forecast(num_param, precision, batch_size, seq_len):
     K = 4.6894e-4 * num_param + 1.8494e6
     y3 = batch_size * seq_len * K * precision / 1e9
 
+    if grad_ckpt:
+        y3 /= 5
+
     # Gradients: N×precision
     y4 = num_param * precision / 1e9
 
@@ -86,16 +89,26 @@ def plot_forecast(num_param, precision, batch_size, seq_len):
     return fig
 
 
-demo = gr.Interface(
-    plot_forecast,
-    [
-        gr.Number(3, label="Number of parameters (B)"),
-        gr.Radio(["float32", "float16", "bfloat16"], value="float32", label="Precision"),
-        gr.Slider(1, 128, label="Batch size", step=1, value=8),
-        gr.Slider(1, 1000, label="Sequence Length", step=1, value=256),
-    ],
-    gr.Plot(label="forecast", format="png"),
-)
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            with gr.Accordion("Model"):
+                num_param = gr.Number(3, label="Number of parameters (B)")
+                precision = gr.Radio(["float32", "float16", "bfloat16"], value="float32", label="Precision")
+            with gr.Accordion("Data"):
+                batch_size = gr.Slider(1, 128, label="Batch size", step=1, value=8)
+                seq_len = gr.Slider(1, 1000, label="Sequence Length", step=1, value=256)
+
+            with gr.Accordion("Advanced", open=False):
+                with gr.Accordion("Data"):
+                    grad_ckpt = gr.Checkbox(False, label="Gradient Checkpointing")
+
+            submit = gr.Button("Submit")
+
+        with gr.Column():
+            plot = gr.Plot(label="forecast", format="png")
+
+    submit.click(plot_forecast, [num_param, precision, grad_ckpt, batch_size, seq_len], plot)
 
 if __name__ == "__main__":
     demo.launch()