"""Data loading utilities for fetching experiment results from HuggingFace.""" import pandas as pd import streamlit as st from huggingface_hub import hf_hub_download, list_repo_files import json from pathlib import Path from typing import Any, Optional REPO_ID = "withcomment/mad" REPO_TYPE = "dataset" TO_NORMALIZE = ('tilde_mu', 'tilde_pi') def normalize(d: dict[str, float]) -> dict[str, float]: """Normalize a distribution""" total = sum(d.values()) if total == 0: return {k: 0.0 for k in d} return {k: v / total for k, v in d.items()} @st.cache_data(ttl=3600) def get_available_runs() -> dict[str, Any]: """List all available run directories in the HuggingFace repo. Returns: A list of records: {"run_path": str, "config": Optional[dict]} """ try: files = list_repo_files(repo_id=REPO_ID, repo_type=REPO_TYPE) # Find directories containing manifest.parquet runs = set() for f in files: if f.endswith("manifest.parquet"): # Extract run path (everything before the dataset name) parts = f.split("/") if len(parts) >= 2: run_path = "/".join(parts[:-2]) # e.g., "mcmc/2026-01-10/12-30-00" runs.add(run_path) run_paths = sorted(runs) return {p: download_config(p) for p in run_paths} except Exception as e: st.warning(f"Failed to list runs: {e}") return {} def _get_nested_value(obj: Any, path: list[str]) -> Any: """Get a nested value from dict/list structures using a path. Supports dict keys and list indices (when the path segment is an int). Returns None when the path cannot be fully resolved. """ current: Any = obj for segment in path: if isinstance(current, dict): if segment not in current: return None current = current[segment] continue if isinstance(current, list): if not segment.isdigit(): return None idx = int(segment) if idx < 0 or idx >= len(current): return None current = current[idx] continue return None return current def _parse_filter_value(raw: str) -> Any: """Parse a user-provided filter value. Tries JSON first to support numbers/bools/null/lists/dicts; falls back to string. """ raw = raw.strip() if raw == "": return "" try: import json return json.loads(raw) except Exception: return raw def filter_runs_by_config_query( runs: list[dict[str, Any]], query: str, ) -> tuple[list[dict[str, Any]], Optional[str]]: """Filter run records by a config query of the form key1.key2=value. Args: runs: Output of list_available_runs(). query: A string like "a.b.c=123". Returns: (filtered_runs, error_message). error_message is None when successful. """ query = (query or "").strip() if not query: return runs, None if "=" not in query: return [], "Invalid filter: expected 'path.to.key=value'" left, right = query.split("=", 1) path = [p for p in left.strip().split(".") if p] if not path: return [], "Invalid filter: empty key path" target_value = _parse_filter_value(right) filtered: list[dict[str, Any]] = [] for record in runs: config = record.get("config") if not isinstance(config, dict): continue actual_value = _get_nested_value(config, path) if actual_value == target_value: filtered.append(record) return filtered, None @st.cache_data(ttl=3600) def list_datasets_for_run(run_path: str) -> list[str]: """List available datasets for a specific run.""" try: files = list_repo_files(repo_id=REPO_ID, repo_type=REPO_TYPE) datasets = set() for f in files: if f.startswith(run_path) and f.endswith("manifest.parquet"): parts = f.split("/") # Dataset name is the second-to-last part if len(parts) >= 2: datasets.add(parts[-2]) return sorted(datasets) except Exception as e: st.warning(f"Failed to list datasets: {e}") return [] @st.cache_data(ttl=3600) def hf_download(run_path: str, dataset: str, file_type: str, suffix: str = 'parquet') -> Optional[pd.DataFrame]: """ Download a parquet file from HuggingFace. Args: run_path: Path to the run (e.g., "mcmc/2026-01-10/12-30-00") dataset: Dataset name (e.g., "csqa") file_type: Either "manifest" or "traces" Returns: DataFrame or None if download fails """ try: file_path = f"{run_path}/{dataset}/{file_type}.{suffix}" local_path = hf_hub_download( repo_id=REPO_ID, filename=file_path, repo_type=REPO_TYPE, ) if suffix == 'json': with open(local_path, 'r') as f: data = json.load(f) elif suffix == 'parquet': data = pd.read_parquet(local_path) return data except Exception as e: st.warning(f"Failed to download {file_type}.{suffix}: {e}") return None @st.cache_data(ttl=3600) def download_config(run_path: str) -> Optional[dict]: """Download config.json for a run.""" try: file_path = f"{run_path}/config.json" local_path = hf_hub_download( repo_id=REPO_ID, filename=file_path, repo_type=REPO_TYPE, ) import json with open(local_path, 'r') as f: return json.load(f) except Exception as e: st.warning(f"Failed to download config: {e}") return None def load_manifest(run_path: str, dataset: str) -> Optional[pd.DataFrame]: """Load manifest dataframe for a run/dataset.""" df = hf_download(run_path, dataset, "manifest") if df is None: return None df['correct'] = df['argmax'] == df['label'] return df def load_metadata(run_path: str, dataset: str) -> Optional[dict]: return hf_download(run_path, dataset, "metadata", suffix="json") def load_traces(run_path: str, dataset: str) -> Optional[pd.DataFrame]: """Load traces dataframe for a run/dataset.""" traces = hf_download(run_path, dataset, "traces") if traces is None: return None for name in TO_NORMALIZE: if name not in traces.columns: continue traces[name] = traces[name].apply(lambda d: normalize(d) if isinstance(d, dict) else d) return traces def load_merged(run_path: str, dataset: str) -> Optional[pd.DataFrame]: """Load and merge manifest and traces dataframes for a run/dataset.""" manifest_df = load_manifest(run_path, dataset) traces_df = load_traces(run_path, dataset) return traces_df.merge(manifest_df, on='chain_idx') def get_traces_for_chain(traces_df: pd.DataFrame, chain_idx: int) -> pd.DataFrame: """Get all traces for a specific chain, sorted by step.""" chain_traces = traces_df[traces_df['chain_idx'] == chain_idx].copy() return chain_traces.sort_values('step') def extract_run_metadata(runs: list[str]) -> pd.DataFrame: """ Extract metadata from run paths for filtering. Returns DataFrame with columns for filtering. """ records = [] for run in runs: parts = run.split("/") records.append({ 'run_path': run, 'date': parts[1] if len(parts) > 1 else '', 'time': parts[2] if len(parts) > 2 else '', }) return pd.DataFrame(records) def get_run_type(run_path: str) -> str: return run_path.split("/")[0] def get_run_timestamp(run_path: str) -> str: parts = run_path.split("/") if len(parts) >= 3: return f"{parts[1]}/{parts[2]}" return "" def get_unique_config_values(runs: dict[str, Any], config_path: list[str]) -> list: """Extract unique values for a config path across all runs.""" values = set() for config in runs.values(): if not config: continue val = _get_nested_value(config, config_path) if val is not None: # Convert lists/dicts to string for display if isinstance(val, (list, dict)): import json val = json.dumps(val) values.add(val) return sorted(values, key=lambda x: (isinstance(x, str), x)) def filter_runs_by_config_value(runs: dict[str, Any], key: list[str] | str, selected_values: list) -> dict[str, Any]: """Filter runs by selected config values.""" if not selected_values: return runs filtered = {} for r, config in runs.items(): if not config: continue val = _get_nested_value(config, key) if val is not None: # Convert lists/dicts to string for comparison if isinstance(val, (list, dict)): import json val = json.dumps(val) if val in selected_values: filtered[r] = config return filtered @st.cache_data(ttl=3600) def load_data(run_path: str, dataset: str) -> Optional[pd.DataFrame]: """Load merged data (manifest + traces) for a run and dataset.""" return load_merged(run_path, dataset)