info_sae_display.py

Go to the documentation of this file.

# Copyright (c) 2025-present Aquin Labs Private Limited. All Rights Reserved.
"""Tables for `aquin info sae` — quiet bordered field/value."""
 
from __future__ import annotations
 
from typing import Any
 
from rich import box
from rich.console import Console
from rich.table import Table
 
# Internal / redundant keys kept off the default surfaces.
_SKIP = frozenset({
    "id", "r2_key", "meta_key", "umap_key",
    "train_steps", "final_loss", "l0", "dead_pct", "created_at",
})
 
# Preferred field order + quiet labels (API key → display).
_FIELDS: list[tuple[str, str]] = [
    ("model_slug", "model"),
    ("model_name", "name"),
    ("layer", "layer"),
    ("d_model", "d_model"),
    ("d_sae", "d_sae"),
    ("k", "k"),
    ("arch", "arch"),
    ("interp_score", "interp"),
]
 
 
def _sae_slug(r: dict[str, Any]) -> str:
    slug = r.get("model_slug")
    layer = r.get("layer")
    if slug is not None and layer is not None:
        return f"{slug}-l{layer}"
    return str(r.get("id") or "")
 
 
def _fmt_value(key: str, value: Any) -> str:
    if value is None:
        return "-"
    if key == "interp_score":
        try:
            return f"{float(value):.4f}"
        except (TypeError, ValueError):
            return str(value)
    return str(value)
 
 
def _rows(r: dict[str, Any]) -> list[tuple[str, str, str]]:
    """(api_key, label, formatted_value) in display order, then leftover keys."""
    seen: set[str] = set()
    out: list[tuple[str, str, str]] = []
    for key, label in _FIELDS:
        if key not in r or key in _SKIP:
            continue
        seen.add(key)
        out.append((key, label, _fmt_value(key, r.get(key))))
    for key, value in r.items():
        if key in seen or key in _SKIP:
            continue
        out.append((key, key, _fmt_value(key, value)))
    return out
 
 
def info_sae_payload(r: dict[str, Any]) -> dict[str, Any]:
    slug = _sae_slug(r)
    fields: dict[str, Any] = {}
    for key, label, _fmt in _rows(r):
        raw = r.get(key)
        if key == "interp_score" and raw is not None:
            try:
                fields[label] = float(raw)
            except (TypeError, ValueError):
                fields[label] = raw
        else:
            fields[label] = raw
    return {
        "kind": "sae",
        "slug": slug,
        "fields": fields,
        "load": f"aquin load sae {slug}" if slug else "aquin load sae <slug>",
    }
 
 
def render_info_sae_quiet(r: dict[str, Any]) -> None:
    from aquin.table_plain import print_table
 
    payload = info_sae_payload(r)
    title = f"SAE {payload['slug']}" if payload.get("slug") else "SAE"
    print_table(
        ["field", "value"],
        [[label, fmt] for _key, label, fmt in _rows(r)],
        aligns=["left", "left"],
        max_col=48,
        title=title,
    )
    print()
    print(f"  {payload['load']}")
    print()
 
 
def render_info_sae(r: dict[str, Any]) -> None:
    """Legacy Rich SIMPLE layout (`--plain`)."""
    console = Console(highlight=False, soft_wrap=True)
    slug = _sae_slug(r)
    table = Table(
        box=box.SIMPLE,
        show_header=True,
        header_style="bold",
        pad_edge=False,
        padding=(0, 2),
    )
    table.add_column("FIELD", no_wrap=True)
    table.add_column("VALUE", overflow="ellipsis", max_width=56)
 
    for _key, label, fmt in _rows(r):
        table.add_row(label, fmt)
 
    console.print()
    console.print(f"SAE {slug}" if slug else "SAE")
    console.print(table)
    console.print()
    if slug:
        console.print(f"Load: aquin load sae {slug}")
    console.print()