"""Load published Praxa behavioral experiment metrics from Hugging Face Hub. Source of truth: EdyVision/praxa-behavioral-experiments (private dataset). Does not depend on the training harness package. Load strategy: walk ``domains/.../results///metrics.json`` on the Hub so newly published conditions (e.g. RFDT) appear after Refresh without editing this file. Known labels stay curated; unknown dirs get a readable fallback label. Caching: in-process frames plus a disk snapshot under ``PRAXA_CACHE_DIR`` (default ``/tmp/praxa_hub_cache``) shared across gunicorn workers. Ordinary page loads reuse memory/disk; the Refresh button force-clears both and re-lists the Hub. """ from __future__ import annotations import fcntl import json import os import pickle import threading import time from dataclasses import dataclass from pathlib import Path from typing import Any import pandas as pd DEFAULT_EXPERIMENTS_REPO = os.environ.get( "PRAXA_EXPERIMENTS_REPO", "EdyVision/praxa-behavioral-experiments", ) DOMAIN_PREFIX: dict[str, str] = { "medicareplus": "domains/healthcare/medicareplus_v1", "citygov": "domains/government/citygov_v1", } DOMAIN_LABELS: dict[str, str] = { "medicareplus": "MediCarePlus (healthcare)", "citygov": "CityGov (government)", } MODEL_KEYS: tuple[str, ...] = ( "qwen3-8b", "granite41-8b", "ministral3-8b-instruct-bf16", "modernbert-base", ) MATCHED_8B_KEYS: tuple[str, ...] = ( "qwen3-8b", "granite41-8b", "ministral3-8b-instruct-bf16", ) ENCODER_MODEL_KEYS: tuple[str, ...] = ("modernbert-base",) MODEL_LABELS: dict[str, str] = { "qwen3-8b": "Qwen3-8B", "granite41-8b": "Granite4.1-8B", "ministral3-8b-instruct-bf16": "Ministral3-8B", "modernbert-base": "ModernBERT-base", } FROZEN_CONDITIONS: tuple[str, ...] = ( "base", "explicit_full_json", "explicit_retrieved", "1shot", "2shot", "4shot", ) LEARNED_CONDITIONS: tuple[str, ...] = ("lora", "cartridge") RFDT_CONDITIONS: tuple[str, ...] = ("rfdt_base", "rfdt_lora") ENCODER_CONDITIONS: tuple[str, ...] = ("encoder_ft",) ALL_CONDITIONS: tuple[str, ...] = ( FROZEN_CONDITIONS + LEARNED_CONDITIONS + RFDT_CONDITIONS + ENCODER_CONDITIONS ) CONDITION_LABELS: dict[str, str] = { "base": "Base", "explicit_full_json": "Explicit policy", "explicit_retrieved": "Explicit retrieved", "1shot": "1-shot", "2shot": "2-shot", "4shot": "4-shot", "lora": "LoRA", "cartridge": "Behavioral KV", "rfdt_base": "RFDT Base", "rfdt_lora": "RFDT-LoRA", "encoder_ft": "Encoder-FT", } METRIC_LABELS: dict[str, str] = { "macro_f1": "Macro-F1", "accuracy": "Accuracy", "answer_f1": "Answer F1", "refuse_f1": "Refuse F1", "allowed_edge": "Allowed-edge Acc", "denied_edge": "Denied-edge Acc", "denied_edge_refuse_recall": "Denied-edge Refuse Recall", } QUERY_TYPE_SLICES: tuple[str, ...] = ( "allowed_base", "allowed_edge", "denied_base", "denied_edge", ) QUERY_TYPE_LABELS: dict[str, str] = { "allowed_base": "Allowed base", "allowed_edge": "Allowed edge", "denied_base": "Denied base", "denied_edge": "Denied edge", } SLICE_HEATMAP_CONDITIONS: tuple[str, ...] = ( "base", "4shot", "lora", "cartridge", "rfdt_lora", "encoder_ft", ) # Seed catalog (also used when Hub tree listing fails). Discovery can find more. _DOMAIN_RUNS: dict[str, list[tuple[str, str, str]]] = { "medicareplus": [ ("medicareplus_qwen3_8b_base", "qwen3-8b", "base"), ("medicareplus_qwen3_8b_explicit_policy", "qwen3-8b", "explicit_full_json"), ("medicareplus_qwen3_8b_explicit_retrieved", "qwen3-8b", "explicit_retrieved"), ("medicareplus_qwen3_8b_few_shot_1", "qwen3-8b", "1shot"), ("medicareplus_qwen3_8b_few_shot_2", "qwen3-8b", "2shot"), ("medicareplus_qwen3_8b_few_shot_4", "qwen3-8b", "4shot"), ("medicareplus_qwen3_8b_lora", "qwen3-8b", "lora"), ("medicareplus_qwen3_8b_cartridge", "qwen3-8b", "cartridge"), ("medicareplus_qwen3_8b_rfdt_base", "qwen3-8b", "rfdt_base"), ("medicareplus_qwen3_8b_rfdt_lora", "qwen3-8b", "rfdt_lora"), ("medicareplus_granite41_8b_base", "granite41-8b", "base"), ( "medicareplus_granite41_8b_explicit_policy", "granite41-8b", "explicit_full_json", ), ( "medicareplus_granite41_8b_explicit_retrieved", "granite41-8b", "explicit_retrieved", ), ("medicareplus_granite41_8b_few_shot_1", "granite41-8b", "1shot"), ("medicareplus_granite41_8b_few_shot_2", "granite41-8b", "2shot"), ("medicareplus_granite41_8b_few_shot_4", "granite41-8b", "4shot"), ("medicareplus_granite41_8b_lora", "granite41-8b", "lora"), ("medicareplus_granite41_8b_cartridge", "granite41-8b", "cartridge"), ("medicareplus_granite41_8b_rfdt_base", "granite41-8b", "rfdt_base"), ("medicareplus_granite41_8b_rfdt_lora", "granite41-8b", "rfdt_lora"), ("medicareplus_ministral3_8b_base", "ministral3-8b-instruct-bf16", "base"), ( "medicareplus_ministral3_8b_explicit_policy", "ministral3-8b-instruct-bf16", "explicit_full_json", ), ( "medicareplus_ministral3_8b_explicit_retrieved", "ministral3-8b-instruct-bf16", "explicit_retrieved", ), ("medicareplus_ministral3_8b_few_shot_1", "ministral3-8b-instruct-bf16", "1shot"), ("medicareplus_ministral3_8b_few_shot_2", "ministral3-8b-instruct-bf16", "2shot"), ("medicareplus_ministral3_8b_few_shot_4", "ministral3-8b-instruct-bf16", "4shot"), ("medicareplus_ministral3_8b_lora", "ministral3-8b-instruct-bf16", "lora"), ( "medicareplus_ministral3_8b_cartridge", "ministral3-8b-instruct-bf16", "cartridge", ), ( "medicareplus_ministral3_8b_rfdt_base", "ministral3-8b-instruct-bf16", "rfdt_base", ), ( "medicareplus_ministral3_8b_rfdt_lora", "ministral3-8b-instruct-bf16", "rfdt_lora", ), ( "medicareplus_modernbert_base_encoder_ft", "modernbert-base", "encoder_ft", ), ], "citygov": [ ("citygov_qwen3_8b_base", "qwen3-8b", "base"), ("citygov_qwen3_8b_explicit_policy", "qwen3-8b", "explicit_full_json"), ("citygov_qwen3_8b_explicit_retrieved", "qwen3-8b", "explicit_retrieved"), ("citygov_qwen3_8b_few_shot_1", "qwen3-8b", "1shot"), ("citygov_qwen3_8b_few_shot_2", "qwen3-8b", "2shot"), ("citygov_qwen3_8b_few_shot_4", "qwen3-8b", "4shot"), ("citygov_qwen3_8b_lora", "qwen3-8b", "lora"), ("citygov_qwen3_8b_cartridge", "qwen3-8b", "cartridge"), ("citygov_qwen3_8b_rfdt_base", "qwen3-8b", "rfdt_base"), ("citygov_qwen3_8b_rfdt_lora", "qwen3-8b", "rfdt_lora"), ("citygov_granite41_8b_base", "granite41-8b", "base"), ("citygov_granite41_8b_explicit_policy", "granite41-8b", "explicit_full_json"), ( "citygov_granite41_8b_explicit_retrieved", "granite41-8b", "explicit_retrieved", ), ("citygov_granite41_8b_few_shot_1", "granite41-8b", "1shot"), ("citygov_granite41_8b_few_shot_2", "granite41-8b", "2shot"), ("citygov_granite41_8b_few_shot_4", "granite41-8b", "4shot"), ("citygov_granite41_8b_lora", "granite41-8b", "lora"), ("citygov_granite41_8b_cartridge", "granite41-8b", "cartridge"), ("citygov_granite41_8b_rfdt_base", "granite41-8b", "rfdt_base"), ("citygov_granite41_8b_rfdt_lora", "granite41-8b", "rfdt_lora"), ("citygov_ministral3_8b_base", "ministral3-8b-instruct-bf16", "base"), ( "citygov_ministral3_8b_explicit_policy", "ministral3-8b-instruct-bf16", "explicit_full_json", ), ( "citygov_ministral3_8b_explicit_retrieved", "ministral3-8b-instruct-bf16", "explicit_retrieved", ), ("citygov_ministral3_8b_few_shot_1", "ministral3-8b-instruct-bf16", "1shot"), ("citygov_ministral3_8b_few_shot_2", "ministral3-8b-instruct-bf16", "2shot"), ("citygov_ministral3_8b_few_shot_4", "ministral3-8b-instruct-bf16", "4shot"), ("citygov_ministral3_8b_lora", "ministral3-8b-instruct-bf16", "lora"), ("citygov_ministral3_8b_cartridge", "ministral3-8b-instruct-bf16", "cartridge"), ( "citygov_ministral3_8b_rfdt_base", "ministral3-8b-instruct-bf16", "rfdt_base", ), ( "citygov_ministral3_8b_rfdt_lora", "ministral3-8b-instruct-bf16", "rfdt_lora", ), ( "citygov_modernbert_base_encoder_ft", "modernbert-base", "encoder_ft", ), ], } def hf_token() -> str | None: return ( os.environ.get("HF_TOKEN") or os.environ.get("HUGGING_FACE_HUB_TOKEN") or None ) def hub_results_path(domain: str, model_key: str, condition: str) -> str: prefix = DOMAIN_PREFIX[domain] return f"{prefix}/results/{model_key}/{condition}" def condition_label(condition: str) -> str: if condition in CONDITION_LABELS: return CONDITION_LABELS[condition] return condition.replace("_", " ").strip().title() def condition_family(condition: str) -> str: if condition in FROZEN_CONDITIONS: return "frozen" if condition in RFDT_CONDITIONS: return "rfdt" if condition in ENCODER_CONDITIONS: return "encoder" if condition in LEARNED_CONDITIONS: return "learned" if condition.endswith("_lora") or condition in {"lora", "cartridge"}: return "learned" return "other" def condition_order_index(condition: str) -> int: order = {name: i for i, name in enumerate(ALL_CONDITIONS)} if condition in order: return order[condition] return len(ALL_CONDITIONS) + abs(hash(condition)) % 1000 def _run_id_for(domain: str, model_key: str, condition: str) -> str: for run_id, mk, cond in _DOMAIN_RUNS.get(domain, []): if mk == model_key and cond == condition: return run_id return f"{domain}_{model_key}_{condition}".replace("-", "_") def discover_domain_runs( domain: str, *, repo_id: str, token: str | None, ) -> list[tuple[str, str, str]]: """List (run_id, model_key, condition) from Hub results tree.""" from huggingface_hub import HfApi prefix = DOMAIN_PREFIX[domain] results_root = f"{prefix}/results" api = HfApi(token=token) found: list[tuple[str, str, str]] = [] seen: set[tuple[str, str]] = set() try: entries = api.list_repo_tree( repo_id=repo_id, repo_type="dataset", path_in_repo=results_root, recursive=True, ) except Exception: return [] for entry in entries: path = getattr(entry, "path", None) or str(entry) if not path.endswith("/metrics.json"): continue marker = "/results/" if marker not in path: continue rel = path.split(marker, 1)[1] parts = rel.split("/") if len(parts) != 3 or parts[2] != "metrics.json": continue model_key, condition = parts[0], parts[1] # Only auto-ingest under known model keys. Legacy Hub dirs like # ``qwen3-8b-base/lora`` are orphans from older publishes and must not # appear as fake models on the ladder. if model_key not in MODEL_KEYS: continue key = (model_key, condition) if key in seen: continue seen.add(key) found.append((_run_id_for(domain, model_key, condition), model_key, condition)) found.sort(key=lambda row: (row[1], condition_order_index(row[2]), row[2])) return found def _catalog_runs(domain: str) -> list[tuple[str, str, str]]: return list(_DOMAIN_RUNS.get(domain, [])) def _summarize_metrics(metrics: dict[str, Any]) -> dict[str, Any]: pc = metrics.get("per_class") or {} bq = metrics.get("by_query_type") or {} denied = (bq.get("denied_edge") or {}).get("per_class") or {} out: dict[str, Any] = { "macro_f1": metrics.get("macro_f1"), "accuracy": metrics.get("accuracy"), "answer_f1": (pc.get("answer") or {}).get("f1"), "refuse_f1": (pc.get("refuse") or {}).get("f1"), "allowed_edge": (bq.get("allowed_edge") or {}).get("accuracy"), "denied_edge": (bq.get("denied_edge") or {}).get("accuracy"), "denied_edge_refuse_recall": (denied.get("refuse") or {}).get("recall"), "selection_status": metrics.get("selection_status"), } for slice_name in QUERY_TYPE_SLICES: block = bq.get(slice_name) or {} refuse = (block.get("per_class") or {}).get("refuse") or {} out[f"slice_acc_{slice_name}"] = block.get("accuracy") out[f"slice_refuse_recall_{slice_name}"] = refuse.get("recall") out[f"slice_n_{slice_name}"] = block.get("n") return out def _download_metrics_json( *, repo_id: str, path_in_repo: str, token: str | None, ) -> dict[str, Any] | None: from huggingface_hub import hf_hub_download from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError try: local = hf_hub_download( repo_id=repo_id, repo_type="dataset", filename=path_in_repo, token=token, ) except (EntryNotFoundError, RepositoryNotFoundError, OSError, ValueError): return None except Exception: return None return json.loads(open(local, encoding="utf-8").read()) def _download_jsonl( *, repo_id: str, path_in_repo: str, token: str | None, ) -> list[dict[str, Any]]: from huggingface_hub import hf_hub_download from huggingface_hub.utils import EntryNotFoundError, RepositoryNotFoundError try: local = hf_hub_download( repo_id=repo_id, repo_type="dataset", filename=path_in_repo, token=token, ) except (EntryNotFoundError, RepositoryNotFoundError, OSError, ValueError): return [] except Exception: return [] rows: list[dict[str, Any]] = [] with open(local, encoding="utf-8") as handle: for line in handle: line = line.strip() if not line: continue rows.append(json.loads(line)) return rows def load_run_predictions( domain: str, model_key: str, condition: str, *, repo_id: str | None = None, ) -> pd.DataFrame: """Fetch predictions.jsonl for one published run.""" if domain not in DOMAIN_PREFIX: raise ValueError(f"unknown domain {domain!r}") repo = repo_id or DEFAULT_EXPERIMENTS_REPO hub_path = hub_results_path(domain, model_key, condition) rows = _download_jsonl( repo_id=repo, path_in_repo=f"{hub_path}/predictions.jsonl", token=hf_token(), ) if not rows: return pd.DataFrame() frame = pd.DataFrame(rows) frame["domain"] = domain frame["model_key"] = model_key frame["model"] = MODEL_LABELS.get(model_key, model_key) frame["condition"] = condition frame["condition_label"] = condition_label(condition) return frame def load_domain_metrics( domain: str, *, repo_id: str | None = None, skip_missing: bool = True, ) -> pd.DataFrame: """Fetch metrics.json for every published run under a domain.""" if domain not in DOMAIN_PREFIX: raise ValueError(f"unknown domain {domain!r}") repo = repo_id or DEFAULT_EXPERIMENTS_REPO token = hf_token() rows: list[dict[str, Any]] = [] discovered = discover_domain_runs(domain, repo_id=repo, token=token) catalog = _catalog_runs(domain) # Prefer discovery when it finds anything; otherwise fall back to catalog # so empty Hub / offline still has a deterministic attempt list. runs = discovered if discovered else catalog # Union: also try catalog slots discovery missed (listing glitches). if discovered: seen = {(mk, cond) for _rid, mk, cond in discovered} for run in catalog: if (run[1], run[2]) not in seen: runs.append(run) for run_id, model_key, condition in runs: hub_path = hub_results_path(domain, model_key, condition) metrics = _download_metrics_json( repo_id=repo, path_in_repo=f"{hub_path}/metrics.json", token=token, ) if metrics is None: if skip_missing: continue raise FileNotFoundError(f"missing {repo}/{hub_path}/metrics.json") summary = _summarize_metrics(metrics) family = condition_family(condition) rows.append( { "domain": domain, "domain_label": DOMAIN_LABELS[domain], "run_id": run_id, "model_key": model_key, "model": MODEL_LABELS.get(model_key, model_key), "condition": condition, "condition_label": condition_label(condition), "hub_path": hub_path, "n": metrics.get("n"), "matched_8b": model_key in MATCHED_8B_KEYS, "is_encoder": model_key in ENCODER_MODEL_KEYS or family == "encoder", "family": family, **summary, } ) return pd.DataFrame(rows) @dataclass class MetricsCache: frames: dict[str, pd.DataFrame] predictions: dict[tuple[str, str, str], pd.DataFrame] loaded_at: float | None error: str | None repo_id: str _CACHE_LOCK = threading.Lock() _CACHE = MetricsCache( frames={}, predictions={}, loaded_at=None, error=None, repo_id=DEFAULT_EXPERIMENTS_REPO, ) # Disk snapshot shared across gunicorn workers. Memory is still L1; Refresh # (force=True) deletes the snapshot and re-pulls from Hub. _DISK_CACHE_DIR = Path( os.environ.get("PRAXA_CACHE_DIR", "/tmp/praxa_hub_cache") ) def _disk_cache_dir() -> Path: _DISK_CACHE_DIR.mkdir(parents=True, exist_ok=True) return _DISK_CACHE_DIR def _metrics_snapshot_path() -> Path: return _disk_cache_dir() / "metrics_frames.pkl" def _analysis_snapshot_path() -> Path: return _disk_cache_dir() / "analysis_cache.pkl" def _refresh_lock_path() -> Path: return _disk_cache_dir() / ".refresh.lock" def _clear_disk_snapshots() -> None: for path in (_metrics_snapshot_path(), _analysis_snapshot_path()): try: path.unlink() except FileNotFoundError: pass except OSError: pass def _read_metrics_snapshot( wanted: tuple[str, ...], ) -> tuple[dict[str, pd.DataFrame], float, str | None] | None: path = _metrics_snapshot_path() if not path.is_file(): return None try: blob = pickle.loads(path.read_bytes()) except Exception: return None if not isinstance(blob, dict): return None if blob.get("repo_id") != DEFAULT_EXPERIMENTS_REPO: return None frames = blob.get("frames") or {} if not all(domain in frames for domain in wanted): return None if not all(isinstance(frames[domain], pd.DataFrame) for domain in wanted): return None loaded_at = float(blob.get("loaded_at") or time.time()) error = blob.get("error") if error is not None and not isinstance(error, str): error = str(error) return frames, loaded_at, error def _write_metrics_snapshot( frames: dict[str, pd.DataFrame], *, loaded_at: float, error: str | None, ) -> None: payload = { "repo_id": DEFAULT_EXPERIMENTS_REPO, "loaded_at": loaded_at, "error": error, "frames": frames, } path = _metrics_snapshot_path() tmp = path.with_suffix(".pkl.tmp") tmp.write_bytes(pickle.dumps(payload, protocol=pickle.HIGHEST_PROTOCOL)) tmp.replace(path) def _read_analysis_snapshot() -> dict[str, Any] | None: path = _analysis_snapshot_path() if not path.is_file(): return None try: blob = pickle.loads(path.read_bytes()) except Exception: return None if not isinstance(blob, dict): return None if blob.get("repo_id") != DEFAULT_EXPERIMENTS_REPO: return None if blob.get("loaded_at") is None: return None return blob def _write_analysis_snapshot(state: dict[str, Any]) -> None: payload = { "repo_id": DEFAULT_EXPERIMENTS_REPO, "loaded_at": state.get("loaded_at"), "error": state.get("error"), "catalog": state.get("catalog"), "by_id": state.get("by_id") or {}, } path = _analysis_snapshot_path() tmp = path.with_suffix(".pkl.tmp") tmp.write_bytes(pickle.dumps(payload, protocol=pickle.HIGHEST_PROTOCOL)) tmp.replace(path) def cache_status() -> dict[str, Any]: with _CACHE_LOCK: return { "loaded_at": _CACHE.loaded_at, "error": _CACHE.error, "repo_id": _CACHE.repo_id, "domains": { d: (0 if df is None else len(df)) for d, df in _CACHE.frames.items() }, "has_token": bool(hf_token()), "disk_snapshot": _metrics_snapshot_path().is_file(), "cache_dir": str(_DISK_CACHE_DIR), } def refresh_cache(*, force: bool = False, domains: tuple[str, ...] | None = None) -> MetricsCache: wanted = domains or tuple(DOMAIN_PREFIX.keys()) with _CACHE_LOCK: if ( not force and _CACHE.loaded_at is not None and all(d in _CACHE.frames for d in wanted) ): warm = True else: warm = False if warm: refresh_analysis_cache(force=False) return _CACHE # Cross-process lock so gunicorn workers share one Hub pull / disk write. with open(_refresh_lock_path(), "a+", encoding="utf-8") as lock_file: fcntl.flock(lock_file.fileno(), fcntl.LOCK_EX) with _CACHE_LOCK: if ( not force and _CACHE.loaded_at is not None and all(d in _CACHE.frames for d in wanted) ): pass # fall through to analysis refresh below else: if force: _clear_disk_snapshots() _CACHE.predictions = {} with _ANALYSIS_LOCK: _ANALYSIS_CACHE["loaded_at"] = None _ANALYSIS_CACHE["error"] = None _ANALYSIS_CACHE["catalog"] = None _ANALYSIS_CACHE["by_id"] = {} loaded_from_disk = False if not force: snap = _read_metrics_snapshot(wanted) if snap is not None: frames, loaded_at, error = snap _CACHE.frames = dict(frames) _CACHE.loaded_at = loaded_at _CACHE.error = error _CACHE.repo_id = DEFAULT_EXPERIMENTS_REPO loaded_from_disk = True if not loaded_from_disk: frames = dict(_CACHE.frames) error = None try: if not hf_token(): error = ( "HF_TOKEN is not set. Add it as a Space secret " "to read the private experiments dataset." ) for domain in wanted: frames[domain] = load_domain_metrics(domain) if frames[domain].empty and error is None: # Empty can also mean not published yet. pass except Exception as exc: # noqa: BLE001 — surface to UI error = f"{type(exc).__name__}: {exc}" loaded_at = time.time() _CACHE.frames = frames _CACHE.loaded_at = loaded_at _CACHE.error = error _CACHE.repo_id = DEFAULT_EXPERIMENTS_REPO try: _write_metrics_snapshot( frames, loaded_at=loaded_at, error=error ) except OSError: pass refresh_analysis_cache(force=force) return _CACHE def get_metrics( domain: str, *, matched_8b_only: bool = False, families: tuple[str, ...] | None = None, ) -> pd.DataFrame: cache = refresh_cache() df = cache.frames.get(domain) if df is None or df.empty: return pd.DataFrame() out = df.copy() if matched_8b_only: out = out[out["matched_8b"]] if families: out = out[out["family"].isin(families)] return out.reset_index(drop=True) def get_predictions( domain: str, model_key: str, condition: str, *, repo_id: str | None = None, ) -> pd.DataFrame: """Cached predictions.jsonl for one run (lazy; Hub download on first hit).""" key = (domain, model_key, condition) with _CACHE_LOCK: hit = _CACHE.predictions.get(key) if hit is not None: return hit.copy() frame = load_run_predictions( domain, model_key, condition, repo_id=repo_id ) with _CACHE_LOCK: _CACHE.predictions[key] = frame return frame.copy() def _pred_label(row: pd.Series) -> str | None: for col in ("parsed_prediction", "pred_outcome"): val = row.get(col) if val is not None and not (isinstance(val, float) and pd.isna(val)): return str(val).strip().lower() or None return None def _gold_label(row: pd.Series) -> str | None: for col in ("target", "gold_outcome"): val = row.get(col) if val is not None and not (isinstance(val, float) and pd.isna(val)): return str(val).strip().lower() or None return None def slice_accuracy_frame( df: pd.DataFrame, *, conditions: tuple[str, ...] | None = None, ) -> pd.DataFrame: """Long frame: method × query_type accuracy from metrics summary columns.""" wanted = conditions or SLICE_HEATMAP_CONDITIONS if df.empty: return pd.DataFrame() focus = df[df["condition"].isin(wanted)].copy() if focus.empty: return pd.DataFrame() rows: list[dict[str, Any]] = [] for row in focus.itertuples(index=False): method = f"{row.model} · {row.condition_label}" for slice_name in QUERY_TYPE_SLICES: col = f"slice_acc_{slice_name}" value = getattr(row, col, None) if value is None or (isinstance(value, float) and pd.isna(value)): continue n_col = f"slice_n_{slice_name}" rows.append( { "method": method, "model": row.model, "model_key": row.model_key, "condition": row.condition, "condition_label": row.condition_label, "query_type": slice_name, "query_type_label": QUERY_TYPE_LABELS.get( slice_name, slice_name ), "accuracy": float(value), "n": getattr(row, n_col, None), "_cond_ord": condition_order_index(row.condition), } ) if not rows: return pd.DataFrame() out = pd.DataFrame(rows) return out.sort_values(["_cond_ord", "model", "query_type"]).reset_index( drop=True ) def rfdt_encoder_disagreement( domain: str, *, matched_8b_only: bool = True, encoder_model_key: str = "modernbert-base", ) -> pd.DataFrame: """Join RFDT-LoRA vs Encoder-FT predictions on sample_id.""" encoder = get_predictions(domain, encoder_model_key, "encoder_ft") if encoder.empty: return pd.DataFrame() enc = encoder.copy() enc["enc_pred"] = enc.apply(_pred_label, axis=1) enc["gold"] = enc.apply(_gold_label, axis=1) enc["enc_correct"] = enc["enc_pred"] == enc["gold"] enc = enc[ ["sample_id", "query_type", "category", "gold", "enc_pred", "enc_correct"] ].drop_duplicates("sample_id") model_keys = MATCHED_8B_KEYS if matched_8b_only else MODEL_KEYS rows: list[dict[str, Any]] = [] for model_key in model_keys: if model_key in ENCODER_MODEL_KEYS: continue rfdt = get_predictions(domain, model_key, "rfdt_lora") if rfdt.empty: continue left = rfdt.copy() left["rfdt_pred"] = left.apply(_pred_label, axis=1) left["rfdt_gold"] = left.apply(_gold_label, axis=1) left["rfdt_correct"] = left["rfdt_pred"] == left["rfdt_gold"] merged = left.merge(enc, on="sample_id", how="inner", suffixes=("", "_enc")) if merged.empty: continue # Prefer RFDT query_type; fall back to encoder column if missing. if "query_type" not in merged.columns or merged["query_type"].isna().all(): merged["query_type"] = merged.get("query_type_enc") for row in merged.itertuples(index=False): agree_pred = row.rfdt_pred == row.enc_pred if row.rfdt_correct and row.enc_correct: outcome = "both_correct" elif (not row.rfdt_correct) and (not row.enc_correct): outcome = "both_wrong" elif row.rfdt_correct and (not row.enc_correct): outcome = "rfdt_only" else: outcome = "encoder_only" qt = getattr(row, "query_type", None) or "unknown" rows.append( { "domain": domain, "sample_id": row.sample_id, "model_key": model_key, "model": MODEL_LABELS.get(model_key, model_key), "query_type": qt, "query_type_label": QUERY_TYPE_LABELS.get(qt, qt), "gold": row.gold if hasattr(row, "gold") else row.rfdt_gold, "rfdt_pred": row.rfdt_pred, "enc_pred": row.enc_pred, "rfdt_correct": bool(row.rfdt_correct), "enc_correct": bool(row.enc_correct), "agree_pred": bool(agree_pred), "outcome": outcome, } ) if not rows: return pd.DataFrame() return pd.DataFrame(rows) def disagreement_summary(disagreement: pd.DataFrame) -> pd.DataFrame: """Per-model outcome rates for the stacked disagreement chart.""" if disagreement.empty: return pd.DataFrame() counts = ( disagreement.groupby(["model", "model_key", "outcome"], as_index=False) .size() .rename(columns={"size": "n"}) ) totals = ( disagreement.groupby(["model", "model_key"], as_index=False) .size() .rename(columns={"size": "total"}) ) out = counts.merge(totals, on=["model", "model_key"]) out["rate"] = out["n"] / out["total"] return out def disagreement_by_slice(disagreement: pd.DataFrame) -> pd.DataFrame: """Per-model × query_type: share where only one side is correct.""" if disagreement.empty: return pd.DataFrame() rows: list[dict[str, Any]] = [] for (model, model_key, qt), group in disagreement.groupby( ["model", "model_key", "query_type"] ): n = len(group) if n == 0: continue rfdt_only = int((group["outcome"] == "rfdt_only").sum()) enc_only = int((group["outcome"] == "encoder_only").sum()) rows.append( { "model": model, "model_key": model_key, "query_type": qt, "query_type_label": QUERY_TYPE_LABELS.get(qt, qt), "n": n, "rfdt_only_rate": rfdt_only / n, "encoder_only_rate": enc_only / n, "disagree_rate": (rfdt_only + enc_only) / n, "net_rfdt_edge": (rfdt_only - enc_only) / n, } ) if not rows: return pd.DataFrame() return pd.DataFrame(rows) def lift_vs_base(df: pd.DataFrame, metric: str = "macro_f1") -> pd.DataFrame: if df.empty or metric not in df.columns: return pd.DataFrame() base = ( df[df["condition"] == "base"][["model_key", "model", metric]] .rename(columns={metric: "base_value"}) .drop_duplicates("model_key") ) merged = df.merge(base, on=["model_key", "model"], how="left") # Lift chart: generation ladder vs Base. Keep RFDT / encoder out # (different objectives / architectures). merged = merged[ (merged["condition"] != "base") & (~merged["condition"].isin(RFDT_CONDITIONS)) & (~merged["condition"].isin(ENCODER_CONDITIONS)) ].copy() merged["delta"] = merged[metric] - merged["base_value"] return merged def generative_ladder_frame(df: pd.DataFrame) -> pd.DataFrame: """Drop encoder-comparator rows from the 8B generation ladder charts.""" if df.empty or "family" not in df.columns: return df return df[df["family"] != "encoder"].reset_index(drop=True) def encoder_comparator_frame(df: pd.DataFrame) -> pd.DataFrame: """RFDT-LoRA (matched 8B) + ModernBERT Encoder-FT for the twin chart.""" if df.empty: return df keep = df[ (df["condition"] == "rfdt_lora") | ((df["condition"] == "encoder_ft") & (df["model_key"].isin(ENCODER_MODEL_KEYS))) ].copy() return keep.reset_index(drop=True) def best_by_family(df: pd.DataFrame, metric: str = "macro_f1") -> dict[str, Any]: """Headline stats for datacards.""" empty = { "best_frozen": None, "best_lora": None, "best_kv": None, "best_rfdt": None, "best_encoder": None, "n_runs": 0, } if df.empty or metric not in df.columns: return empty usable = df.dropna(subset=[metric]) if usable.empty: return empty frozen = usable[usable["family"] == "frozen"] lora = usable[usable["condition"] == "lora"] kv = usable[usable["condition"] == "cartridge"] rfdt = usable[usable["condition"] == "rfdt_lora"] encoder = usable[usable["condition"] == "encoder_ft"] def _best(frame: pd.DataFrame) -> dict[str, Any] | None: if frame.empty: return None row = frame.loc[frame[metric].idxmax()] return { "model": row["model"], "condition_label": row["condition_label"], "value": float(row[metric]), "run_id": row["run_id"], } return { "best_frozen": _best(frozen), "best_lora": _best(lora), "best_kv": _best(kv), "best_rfdt": _best(rfdt), "best_encoder": _best(encoder), "n_runs": int(len(usable)), } # --------------------------------------------------------------------------- # Published analysis artifacts (analysis/catalog.json on the experiments dataset) # --------------------------------------------------------------------------- ANALYSIS_CATALOG_PATH = "analysis/catalog.json" _ANALYSIS_LOCK = threading.Lock() _ANALYSIS_CACHE: dict[str, Any] = { "loaded_at": None, "error": None, "catalog": None, "by_id": {}, } def refresh_analysis_cache(*, force: bool = False) -> dict[str, Any]: """Load analysis/catalog.json and each results.json listed there.""" with _ANALYSIS_LOCK: if not force and _ANALYSIS_CACHE["loaded_at"] is not None: return _ANALYSIS_CACHE if not force: snap = _read_analysis_snapshot() if snap is not None: _ANALYSIS_CACHE["catalog"] = snap.get("catalog") _ANALYSIS_CACHE["by_id"] = dict(snap.get("by_id") or {}) _ANALYSIS_CACHE["loaded_at"] = snap.get("loaded_at") _ANALYSIS_CACHE["error"] = snap.get("error") return _ANALYSIS_CACHE token = hf_token() catalog = _download_metrics_json( repo_id=DEFAULT_EXPERIMENTS_REPO, path_in_repo=ANALYSIS_CATALOG_PATH, token=token, ) by_id: dict[str, dict[str, Any]] = {} error: str | None = None if catalog is None: error = ( "No published analysis catalog yet " f"({DEFAULT_EXPERIMENTS_REPO}/{ANALYSIS_CATALOG_PATH})." ) else: for entry in catalog.get("analyses") or []: path = entry.get("path") analysis_id = entry.get("analysis_id") if not path or not analysis_id: continue payload = _download_metrics_json( repo_id=DEFAULT_EXPERIMENTS_REPO, path_in_repo=str(path), token=token, ) if payload is not None: by_id[str(analysis_id)] = payload _ANALYSIS_CACHE["catalog"] = catalog _ANALYSIS_CACHE["by_id"] = by_id _ANALYSIS_CACHE["loaded_at"] = time.time() _ANALYSIS_CACHE["error"] = error try: _write_analysis_snapshot(_ANALYSIS_CACHE) except OSError: pass return _ANALYSIS_CACHE def get_analysis_catalog() -> dict[str, Any] | None: cache = refresh_analysis_cache() return cache.get("catalog") def get_analysis(analysis_id: str) -> dict[str, Any] | None: cache = refresh_analysis_cache() return cache.get("by_id", {}).get(analysis_id) def bootstrap_contrasts_frame( domain: str | None = None, *, analysis_id: str = "leakage_bootstrap", ) -> pd.DataFrame: """Flatten published leakage-bootstrap contrasts; optional domain filter.""" blob = get_analysis(analysis_id) if not blob: return pd.DataFrame() payload = blob.get("payload") or {} rows = payload.get("contrasts") or [] if not rows: return pd.DataFrame() frame = pd.DataFrame(rows) if domain and "domain" in frame.columns: frame = frame[frame["domain"] == domain].reset_index(drop=True) return frame def mcnemar_contrasts_frame( domain: str | None = None, *, analysis_id: str = "mcnemar", ) -> pd.DataFrame: """Flatten published McNemar contrasts; optional domain filter.""" blob = get_analysis(analysis_id) if not blob: return pd.DataFrame() payload = blob.get("payload") or {} rows = payload.get("contrasts") or [] if not rows: return pd.DataFrame() frame = pd.DataFrame(rows) if domain and "domain" in frame.columns: frame = frame[frame["domain"] == domain].reset_index(drop=True) return frame def analysis_cache_status() -> dict[str, Any]: with _ANALYSIS_LOCK: catalog = _ANALYSIS_CACHE.get("catalog") or {} return { "loaded_at": _ANALYSIS_CACHE.get("loaded_at"), "error": _ANALYSIS_CACHE.get("error"), "n_analyses": len(_ANALYSIS_CACHE.get("by_id") or {}), "analysis_ids": sorted((_ANALYSIS_CACHE.get("by_id") or {}).keys()), "catalog_updated_at_utc": catalog.get("updated_at_utc"), }