fix: table slider crash (#434)

Co-authored-by: gradio-pr-bot <gradio-pr-bot@users.noreply.github.com>
Co-authored-by: Abubakar Abid <islamrealm@gmail.com>

Author: 3 people

.changeset/tender-dryers-give.md

@@ -0,0 +1,5 @@
+---
+"trackio": minor
+---
+
+feat:fix: table slider crash

trackio/ui/media_page.py

@@ -46,6 +46,10 @@ def extract_media(logs: list[dict]) -> dict[str, list[MediaData]]:
     return media_by_key
 
 
+def is_table(value) -> bool:
+    return isinstance(value, dict) and value.get("_type") == Table.TYPE
+
+
 def refresh_runs_dropdown(project: str | None):
     if project is None:
         runs: list[str] = []
@@ -102,13 +106,7 @@ def display_media_and_tables(project: str | None, selected_run: str | None):
 
         table_cols = df.select_dtypes(include="object").columns
         table_cols = [c for c in table_cols if c not in utils.RESERVED_KEYS]
-        table_cols = [
-            c
-            for c in table_cols
-            if not (metric_df := df.dropna(subset=[c])).empty
-            and isinstance(first_value := metric_df[c].iloc[0], dict)
-            and first_value.get("_type") == Table.TYPE
-        ]
+        table_cols = [c for c in table_cols if any(is_table(x) for x in df[c])]
         has_tables = len(table_cols) > 0
 
         if not has_media and not has_tables:
@@ -169,7 +167,7 @@ def display_media_and_tables(project: str | None, selected_run: str | None):
             with gr.Accordion(f"Tables ({len(table_cols)})", open=True):
                 with gr.Row(key="row"):
                     for metric_idx, metric_name in enumerate(table_cols):
-                        metric_df = df.dropna(subset=[metric_name])
+                        metric_df = df[df[metric_name].apply(is_table)]
                         if not metric_df.empty:
                             value = metric_df[metric_name]
                             first_value = value.iloc[0]
@@ -202,16 +200,27 @@ def display_media_and_tables(project: str | None, selected_run: str | None):
                                             preserved_by_key=None,
                                         )
 
-                                        def get_table_at_index(index: int):
-                                            value = metric_df[metric_name]
-                                            processed_data = Table.to_display_format(
-                                                value.iloc[index - 1]["_value"]
-                                            )
-                                            df_ = pd.DataFrame(processed_data)
-                                            return gr.DataFrame(
-                                                df_,
-                                                label=f"{metric_name} (index {index})",
-                                            )
+                                        def make_table_renderer(
+                                            capture_df, capture_name
+                                        ):
+                                            def get_table_at_index(index: int):
+                                                value = capture_df[capture_name]
+                                                processed_data = (
+                                                    Table.to_display_format(
+                                                        value.iloc[index - 1]["_value"]
+                                                    )
+                                                )
+                                                df_ = pd.DataFrame(processed_data)
+                                                return gr.DataFrame(
+                                                    df_,
+                                                    label=f"{capture_name} (index {index})",
+                                                )
+
+                                            return get_table_at_index
+
+                                        get_table_at_index = make_table_renderer(
+                                            metric_df, metric_name
+                                        )
 
                                         s.input(
                                             get_table_at_index,