changes

Author: abidlabs

trackio/__init__.py

@@ -1,39 +1,46 @@
 from pathlib import Path
-import os
+
 from .storage import TrackioStorage
 from .ui import launch_ui
 
 __version__ = Path(__file__).parent.joinpath("version.txt").read_text().strip()
 
 _current_run = None
 
+
 class Run:
     def __init__(self, project, name, config):
         self.project = project
         self.name = name
         self.config = config
         self.storage = TrackioStorage(project, name, config)
+
     def log(self, metrics):
         self.storage.log(metrics)
+
     def finish(self):
         self.storage.finish()
 
+
 def init(project, name=None, config=None):
     global _current_run
     _current_run = Run(project, name, config)
     return _current_run
 
+
 def log(metrics):
     if _current_run is None:
         raise RuntimeError("Call trackio.init() before log().")
     _current_run.log(metrics)
 
+
 def finish():
     if _current_run is None:
         raise RuntimeError("Call trackio.init() before finish().")
     _current_run.finish()
     global _current_run
     _current_run = None
 
+
 def ui():
     launch_ui()

trackio/storage.py

@@ -2,6 +2,7 @@
 from datasets import Dataset
 import pandas as pd
 
+
 class TrackioStorage:
     def __init__(self, project, name, config):
         self.project = project
@@ -14,12 +15,15 @@ def __init__(self, project, name, config):
         self.config_path = os.path.join(self.dir, "config.json")
         # Save config immediately
         import json
+
         with open(self.config_path, "w") as f:
             json.dump(self.config, f, indent=2)
+
     def log(self, metrics):
         self.logs.append(metrics)
+
     def finish(self):
         if self.logs:
             df = pd.DataFrame(self.logs)
             ds = Dataset.from_pandas(df)
-            ds.to_parquet(self.run_path) 
+            ds.to_parquet(self.run_path)

trackio/ui.py

@@ -1,21 +1,33 @@
-import gradio as gr
-import os
 import json
+import os
+
+import gradio as gr
 import pandas as pd
 import plotly.express as px
 
 TRACKIO_DIR = "trackio"
 
+
 def get_projects():
     if not os.path.exists(TRACKIO_DIR):
         return []
-    return [d for d in os.listdir(TRACKIO_DIR) if os.path.isdir(os.path.join(TRACKIO_DIR, d))]
+    return [
+        d
+        for d in os.listdir(TRACKIO_DIR)
+        if os.path.isdir(os.path.join(TRACKIO_DIR, d))
+    ]
+
 
 def get_runs(project):
     project_dir = os.path.join(TRACKIO_DIR, project)
     if not os.path.exists(project_dir):
         return []
-    return [d for d in os.listdir(project_dir) if os.path.isdir(os.path.join(project_dir, d))]
+    return [
+        d
+        for d in os.listdir(project_dir)
+        if os.path.isdir(os.path.join(project_dir, d))
+    ]
+
 
 def load_run_data(project, run):
     run_dir = os.path.join(TRACKIO_DIR, project, run)
@@ -30,37 +42,58 @@ def load_run_data(project, run):
             config = json.load(f)
     return df, config
 
+
 def plot_metrics(df):
     if df is None or df.empty:
         return None
     plots = []
-    numeric_cols = df.select_dtypes(include='number').columns
+    numeric_cols = df.select_dtypes(include="number").columns
     for col in numeric_cols:
         fig = px.line(df, y=col, title=col)
         plots.append(fig)
     return plots
 
+
 def update_runs(project):
     return gr.Dropdown.update(choices=get_runs(project), value=None)
 
+
 def update_dashboard(project, run):
     if not project or not run:
-        return gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+        return (
+            gr.update(visible=False),
+            gr.update(visible=False),
+            gr.update(visible=False),
+        )
     df, config = load_run_data(project, run)
     plots = plot_metrics(df)
     config_str = json.dumps(config, indent=2)
     return plots, gr.JSON.update(value=config), gr.update(visible=True)
 
+
 def launch_ui():
     with gr.Blocks() as demo:
         gr.Markdown("# Trackio Dashboard")
         with gr.Row():
             project_dd = gr.Dropdown(label="Project", choices=get_projects())
             run_dd = gr.Dropdown(label="Run", choices=[])
         with gr.Row():
-            plot_output = gr.Plot(label="Metrics", visible=False, show_label=True, elem_id="metrics-plot", interactive=True, scale=2)
+            plot_output = gr.Plot(
+                label="Metrics",
+                visible=False,
+                show_label=True,
+                elem_id="metrics-plot",
+                interactive=True,
+                scale=2,
+            )
             config_output = gr.JSON(label="Config", visible=False)
         # Events
-        project_dd.change(fn=lambda p: update_runs(p), inputs=project_dd, outputs=run_dd)
-        run_dd.change(fn=lambda p, r: update_dashboard(p, r), inputs=[project_dd, run_dd], outputs=[plot_output, config_output, plot_output])
-    demo.launch() 
+        project_dd.change(
+            fn=lambda p: update_runs(p), inputs=project_dd, outputs=run_dd
+        )
+        run_dd.change(
+            fn=lambda p, r: update_dashboard(p, r),
+            inputs=[project_dd, run_dd],
+            outputs=[plot_output, config_output, plot_output],
+        )
+    demo.launch()