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

"""Generate held-out ToM figures from retained N10 analysis artifacts."""

from __future__ import annotations

import argparse
from pathlib import Path

from tom_figure_helpers import (
    load_gamma_cells,
    plot_attribution_bar,
    plot_gamma_heatmap,
    plot_pragmatic_bar,
    plot_social_life_forest,
)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--gamma-ci", type=Path, required=True)
    parser.add_argument("--bin-scores-dir", type=Path, required=True)
    parser.add_argument("--out-dir", type=Path, required=True)
    args = parser.parse_args(argv)

    args.out_dir.mkdir(parents=True, exist_ok=True)
    cells = load_gamma_cells(args.gamma_ci)
    topics = list(cells["tombench_hinting_task_test"].keys())

    plot_pragmatic_bar(cells, topics, args.out_dir)
    plot_social_life_forest(cells, args.out_dir)
    plot_gamma_heatmap(cells, topics, args.out_dir)
    plot_attribution_bar(cells, args.bin_scores_dir, args.out_dir)

    written = sorted(path.name for path in args.out_dir.glob("fig_tom_*.png"))
    print("wrote:", written)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())