Merge pull request #12 from gradio-app/ema

Add a checkbox for EMA-based smoothing

Author: abidlabs

trackio/ui.py

@@ -31,14 +31,18 @@ def get_runs(project):
     return storage.get_runs(project)
 
 
-def load_run_data(project, run):
+def load_run_data(project: str | None, run: str | None, smoothing: bool):
     if not project or not run:
         return None
     storage = SQLiteStorage("", "", {})
     metrics = storage.get_metrics(project, run)
     if not metrics:
         return None
     df = pd.DataFrame(metrics)
+    if smoothing:
+        numeric_cols = df.select_dtypes(include="number").columns
+        numeric_cols = [c for c in numeric_cols if c not in RESERVED_KEYS]
+        df[numeric_cols] = df[numeric_cols].ewm(alpha=0.1).mean()
     if "step" not in df.columns:
         df["step"] = range(len(df))
     return df
@@ -64,14 +68,15 @@ def log(project: str, run: str, metrics: dict[str, Any]) -> None:
     storage.log(metrics)
 
 
-with gr.Blocks(theme="citrus") as demo:
+with gr.Blocks(theme="citrus", title="Trackio Dashboard") as demo:
     with gr.Sidebar() as sidebar:
         gr.Markdown(
             f"<div style='display: flex; align-items: center; gap: 8px;'><img src='/gradio_api/file={TRACKIO_LOGO_PATH}' width='32' height='32'><span style='font-size: 2em; font-weight: bold;'>Trackio</span></div>"
         )
         project_dd = gr.Dropdown(label="Project", allow_custom_value=True)
         gr.Markdown("### ⚙️ Settings")
         realtime_cb = gr.Checkbox(label="Refresh realtime", value=True)
+        smoothing_cb = gr.Checkbox(label="Smoothing", value=True)
     with gr.Row():
         run_dd = gr.Dropdown(label="Run", choices=[], multiselect=True)
 
@@ -103,13 +108,13 @@ def log(project: str, run: str, metrics: dict[str, Any]) -> None:
     )
 
     @gr.render(
-        triggers=[run_dd.change, timer.tick],
-        inputs=[project_dd, run_dd],
+        triggers=[run_dd.change, timer.tick, smoothing_cb.change],
+        inputs=[project_dd, run_dd, smoothing_cb],
     )
-    def update_dashboard(project, runs):
+    def update_dashboard(project, runs, smoothing):
         dfs = []
         for run in runs:
-            df = load_run_data(project, run)
+            df = load_run_data(project, run, smoothing)
             if df is not None:
                 df["run"] = run
                 dfs.append(df)