Update app.py

app.py

@@ -1,20 +1,17 @@
-# roman_estimator_with_presets.py
 import gradio as gr
 import math
 
 # ------------------------
 # GPU presets: TFLOPs (units: TFLOPs)
-# Each GPU: dict of dtype -> TFLOPs (tera-FLOPs)
-# These are presets / approximations — editable by user via the TFLOPs field.
 # ------------------------
 GPUS = {
     "T4":        {"FP32":  8.1,  "FP16":  65.0,  "INT4": 260.0},
     "RTX 3060":  {"FP32": 13.0,  "FP16": 26.0,   "INT4": 52.0},
     "RTX 3090":  {"FP32": 35.0,  "FP16": 70.0,   "INT4": 140.0},
     "RTX 4090":  {"FP32": 83.0,  "FP16": 166.0,  "INT4": 332.0},
-    "A100 80GB": {"FP32": 19.5,  "FP16": 150.0,  "INT4": 600.0},  # FP32 architecture throughput is lower than FP16
+    "A100 80GB": {"FP32": 19.5,  "FP16": 150.0,  "INT4": 600.0},
     "H100 SXM":  {"FP32": 30.0,  "FP16": 300.0,  "INT4": 1200.0},
-    "Custom":    {"FP32": 1.0,   "FP16": 1.0,    "INT4": 1.0},  # placeholder for manual GPUs
+    "Custom":    {"FP32": 1.0,   "FP16": 1.0,    "INT4": 1.0},
 }
 
 # ------------------------
@@ -56,48 +53,32 @@ def estimate_time(params_m: float,
                   dtype: str,
                   tf_override: float,
                   utilization_pct: float):
-    """
-    params_m: model params in millions (e.g., 100 -> 100M)
-    tokens_b: training tokens in billions (e.g., 1.5 -> 1.5B)
-    selected_gpu: key in GPUS
-    dtype: "FP32" / "FP16" / "INT4"
-    tf_override: numeric TFLOPs (if >0 will override preset)
-    utilization_pct: fraction 0..100 representing real-world utilization
-    """
-    # validation
     if params_m <= 0 or tokens_b <= 0:
         return "Enter positive values for parameters and tokens."
 
     params = params_m * 1e6
     tokens = tokens_b * 1e9
 
-    # choose TFLOPs: override if user entered > 0, else use preset
     if tf_override is not None and tf_override > 0:
         chosen_tf = float(tf_override)
         source = "manual override"
     else:
-        # safe fallback
         try:
             chosen_tf = float(GPUS[selected_gpu].get(dtype, 0.0))
             source = f"preset ({selected_gpu} / {dtype})"
         except Exception:
-            chosen_tf = 0.0
-            source = "preset-missing"
+            return "Couldn't determine GPU TFLOPs. Pick a GPU or enter TFLOPs manually."
 
     if chosen_tf <= 0:
         return "Couldn't determine GPU TFLOPs. Pick a GPU or enter TFLOPs manually."
 
-    # convert to FLOPs/sec
-    dtype_mul = 1.0  # GPUS already store per-dtype TFLOPs, so no extra multiplier
     gpu_flops_per_sec = chosen_tf * 1e12 * (max(0.001, utilization_pct / 100.0))
 
-    # FLOPs estimate (industry rule of thumb)
-    flops_total = 6 * params * tokens  # total training FLOPs (approx)
+    flops_total = 6 * params * tokens
     seconds = flops_total / gpu_flops_per_sec
     hours = seconds / 3600.0
     days = hours / 24.0
 
-    # extras
     seq_len = 2048.0
     steps = max(1.0, tokens / seq_len)
     flops_per_step = flops_total / steps if steps > 0 else 0.0
@@ -117,39 +98,45 @@ def estimate_time(params_m: float,
         f"FLOPs / step (avg): {flops_per_step:.3e}",
     ]
 
-    # warning for suspicious override values
     if tf_override and tf_override > 0 and selected_gpu != "Custom":
         out.append("")
-        out.append("⚠️ Note: you overrode the preset TFLOPs. Make sure the value is in TFLOPs (e.g., 150 for A100 FP16-like).")
+        out.append("⚠️ Note: you overrode the preset TFLOPs. Ensure the value is in TFLOPs (e.g., 150 for A100 FP16-like).")
 
     return "\n".join(out)
 
-# ------------------------
-# Helper to return preset TFLOPs for UI update
-# ------------------------
 def preset_tf_for_ui(selected_gpu: str, dtype: str):
-    """Return the preset TFLOPs number for the selected GPU+dtype (or 0 if missing)."""
-    val = 0.0
     if selected_gpu in GPUS:
-        val = GPUS[selected_gpu].get(dtype, 0.0)
-    return val
+        return GPUS[selected_gpu].get(dtype, 0.0)
+    return 0.0
+
+# ------------------------
+# Theme setter (returns HTML snippet to run client-side JS)
+# ------------------------
+def set_theme(theme_name: str):
+    # map button label -> class name used in CSS
+    cls = {
+        "Blue": "theme-blue",
+        "Green": "theme-green",
+        "Purple": "theme-purple",
+    }.get(theme_name, "theme-blue")
+    # return script that sets the root class
+    return f"<script>document.documentElement.className='{cls}';</script>"
 
 # ------------------------
 # Build UI
 # ------------------------
-with gr.Blocks(css=CSS, title="Roman's Training Time Estimator") as demo:
-    # set a default theme class
+with gr.Blocks() as demo:
+    # initial theme set (runs immediately on load)
     gr.HTML("<script>document.documentElement.className='theme-blue';</script>")
 
     with gr.Column(elem_classes="card"):
-        with gr.Row(elem_classes="header-row"):
+        with gr.Row():
             gr.Markdown("## 🧠 Roman’s Training Time Estimator")
-            with gr.Row(elem_classes="theme-buttons"):
+            with gr.Row():
                 btn_blue = gr.Button("Blue", elem_classes="btn-theme")
                 btn_green = gr.Button("Green", elem_classes="btn-theme")
                 btn_purple = gr.Button("Purple", elem_classes="btn-theme")
 
-    # Model & hardware card
     with gr.Column(elem_classes="card"):
         gr.Markdown("### Model & Hardware")
         with gr.Row():
@@ -162,31 +149,32 @@ with gr.Blocks(css=CSS, title="Roman's Training Time Estimator") as demo:
             tf_override = gr.Number(value=preset_tf_for_ui("A100 80GB", "FP16"), label="GPU TFLOPs (teraFLOPs) — editable", precision=3)
             utilization = gr.Slider(minimum=1, maximum=100, value=80, step=1, label="Hardware Utilization (%) — realistic throughput")
 
-    # Result card
     with gr.Column(elem_classes="card"):
         gr.Markdown("### Estimate")
         result = gr.Textbox(lines=12, interactive=False, elem_classes="result-box", label="Result")
         run_btn = gr.Button("Estimate Training Time", elem_classes="btn-theme")
 
-    # Wire interactions
-    # When GPU or dtype changes, update tf_override value to the preset for that combo
+    # invisible HTML target used to inject theme-changing script
+    theme_script = gr.HTML(value="")
+
+    # update TF override when gpu/dtype change
     def _update_tf(selected_gpu, dtype):
         return gr.update(value=preset_tf_for_ui(selected_gpu, dtype))
-
     gpu_dropdown.change(_update_tf, inputs=[gpu_dropdown, dtype_dropdown], outputs=[tf_override])
     dtype_dropdown.change(_update_tf, inputs=[gpu_dropdown, dtype_dropdown], outputs=[tf_override])
 
-    # Run button computes estimate
+    # button clicks
     run_btn.click(estimate_time,
                   inputs=[params, tokens, gpu_dropdown, dtype_dropdown, tf_override, utilization],
                   outputs=[result])
 
-    # Theme buttons (JS simply toggles class on documentElement)
-    btn_blue.click(None, None, None, _js="() => { document.documentElement.className='theme-blue'; return []; }")
-    btn_green.click(None, None, None, _js="() => { document.documentElement.className='theme-green'; return []; }")
-    btn_purple.click(None, None, None, _js="() => { document.documentElement.className='theme-purple'; return []; }")
+    # theme buttons now call the Python set_theme and return HTML that runs client-side
+    btn_blue.click(set_theme, inputs=["Blue"], outputs=[theme_script])
+    btn_green.click(set_theme, inputs=["Green"], outputs=[theme_script])
+    btn_purple.click(set_theme, inputs=["Purple"], outputs=[theme_script])
 
-    gr.HTML("<div class='small-muted'>Tip: GPU preset values are TFLOPs per dtype. You can edit the TFLOPs number to override. Utilization reduces theoretical peak to realistic throughput.</div>")
+    gr.HTML("<div class='small-muted'>Tip: GPU presets are TFLOPs per dtype. You can edit the TFLOPs number to override. Utilization reduces theoretical peak to realistic throughput.</div>")
 
+# pass CSS to launch (Gradio 6.0+ API)
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(css=CSS, title="Roman's Training Time Estimator")