social-data-attribution-w2/scripts/analysis/figure4_table_helpers.py

"""Constants and render helpers for N10 Figure-4-style tables."""

from __future__ import annotations

from typing import Any

ALL_TOPICS = [
    "adult_content",
    "art_and_design",
    "crime_and_law",
    "education_and_jobs",
    "electronics_and_hardware",
    "entertainment",
    "fashion_and_beauty",
    "finance_and_business",
    "food_and_dining",
    "games",
    "health",
    "history_and_geography",
    "home_and_hobbies",
    "industrial",
    "literature",
    "politics",
    "religion",
    "science_math_and_technology",
    "social_life",
    "software",
    "software_development",
    "sports_and_fitness",
    "transportation",
    "travel_and_tourism",
]
PROBES = [
    ("simpletom_judgment-qa", "simpletom_judgment-qa:mc"),
    ("simpletom_mental-state-qa", "simpletom_mental-state-qa:mc"),
    ("simpletom_behavior-qa", "simpletom_behavior-qa:mc"),
    ("tombench_ambiguous_story_task", "tombench_ambiguous_story_task:mc"),
    (
        "tombench_completion_of_failed_actions",
        "tombench_completion_of_failed_actions:mc",
    ),
    ("tombench_discrepant_desires", "tombench_discrepant_desires:mc"),
    ("tombench_discrepant_emotions", "tombench_discrepant_emotions:mc"),
    ("tombench_discrepant_intentions", "tombench_discrepant_intentions:mc"),
    ("tombench_emotion_regulation", "tombench_emotion_regulation:mc"),
    ("tombench_false_belief_task", "tombench_false_belief_task:mc"),
    ("tombench_faux_pas_recognition_test", "tombench_faux_pas_recognition_test:mc"),
    ("tombench_hidden_emotions", "tombench_hidden_emotions:mc"),
    ("tombench_hinting_task_test", "tombench_hinting_task_test:mc"),
    ("tombench_knowledge_attention_links", "tombench_knowledge_attention_links:mc"),
    (
        "tombench_knowledge_pretend_play_links",
        "tombench_knowledge_pretend_play_links:mc",
    ),
    ("tombench_moral_emotions", "tombench_moral_emotions:mc"),
    ("tombench_multiple_desires", "tombench_multiple_desires:mc"),
    ("tombench_percepts_knowledge_links", "tombench_percepts_knowledge_links:mc"),
    ("tombench_persuasion_story_task", "tombench_persuasion_story_task:mc"),
    ("tombench_prediction_of_actions", "tombench_prediction_of_actions:mc"),
    ("tombench_scalar_implicature_test", "tombench_scalar_implicature_test:mc"),
    ("tombench_strange_story_task", "tombench_strange_story_task:mc"),
    ("tombench_unexpected_outcome_test", "tombench_unexpected_outcome_test:mc"),
    ("bbh_snarks", "bbh_snarks:cot"),
]
METRICS_OF_INTEREST = ["primary_score", "acc_per_char", "acc_raw", "acc_uncond"]
NULL_LABEL = "__null__"
GAMMA_EPS = 1e-6


def render_md_table(
    grid: dict[str, dict[str, dict[str, Any]]],
    metric: str = "acc_per_char",
    field: str = "gamma",
    mode: str = "absolute",
) -> str:
    label = {"gamma": "γ (raw)", "net_gamma": "net γ (γ_topic-γ_null)"}.get(
        field, field
    )
    header = "| topic | " + " | ".join(probe for probe, _ in PROBES) + " | mean |"
    sep = "|---|" + "|".join(["---"] * (len(PROBES) + 1)) + "|"
    lines = [
        f"### Figure-4-style table — {label}, {mode} (metric = `{metric}`)",
        "",
        header,
        sep,
    ]
    for topic in ALL_TOPICS:
        cells = []
        gammas = []
        for probe, _ in PROBES:
            gamma = grid[probe][topic].get(field)
            if gamma is None:
                cells.append("—")
            else:
                gammas.append(gamma)
                cells.append(f"{gamma:+.3f}")
        mean_gamma = (sum(gammas) / len(gammas)) if gammas else None
        mean_str = f"**{mean_gamma:+.3f}**" if mean_gamma is not None else "—"
        lines.append(f"| {topic} | " + " | ".join(cells) + f" | {mean_str} |")
    return "\n".join(lines)


def render_baseline_md(baselines: dict[str, dict[str, float]]) -> str:
    lines = [
        "### Baselines (un-unlearned OLMo3-7B Base, 5-shot OLMES)",
        "",
        "| probe | acc_per_char | acc_raw | acc_uncond |",
        "|---|---|---|---|",
    ]
    for probe, _ in PROBES:
        baseline = baselines.get(probe, {})
        if all(baseline.get(key) is not None for key in METRICS_OF_INTEREST):
            lines.append(
                f"| {probe} | {baseline.get('acc_per_char'):.4f} | "
                f"{baseline.get('acc_raw'):.4f} | {baseline.get('acc_uncond'):.4f} |"
            )
        else:
            lines.append(f"| {probe} | — | — | — |")
    return "\n".join(lines)