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

#!/usr/bin/env python3
"""Compute recurrent top-k document support across held-out probes."""

from __future__ import annotations

import argparse
import csv
import json
import math
from collections import defaultdict
from dataclasses import dataclass
from pathlib import Path
from statistics import mean, stdev
from typing import Any

SCRIPT_PATH = Path(__file__).resolve()
REPO_ROOT = SCRIPT_PATH.parents[2] if len(SCRIPT_PATH.parents) > 2 else Path.cwd()
DEFAULT_TOPK_ROOT = Path("/storage/ice-shared/cs7634/staff/TDA/social-data-attribution/runs/trackstar/results_topk/instruct_base/socialtda_holdout_20260528T200307Z")
DEFAULT_OUT = REPO_ROOT / "artifacts/cross_probe_consensus"
LABEL_OVERRIDES = {"about_org": "About Org.", "about_pers": "About Person", "art_and_design": "Art & Design", "crime_and_law": "Crime & Law", "finance_and_business": "Finance & Business", "home_and_hobbies": "Home & Hobbies", "q_a_forum": "Q&A Forum", "science_math_and_technology": "Science & Technology"}

@dataclass(frozen=True)
class Spec:
    key: str
    display: str
    filename: str
    include_prefixes: tuple[str, ...] = ()
    exclude_prefixes: tuple[str, ...] = ()

SPECS = [
    Spec("bbq", "BBQ", "queries_olmes_instruct_bbq_top100.jsonl"),
    Spec("bbh_disambiguation_qa", "BBH Disambig. QA", "queries_olmes_instruct_bbh_disambiguation_qa_top100.jsonl"),
    Spec("tombench", "ToMBench", "queries_olmes_instruct_tombench_top100.jsonl"),
    Spec("negotiationtom", "NegotiationToM", "queries_olmes_instruct_negotiationtom_top100.jsonl"),
    Spec("pub_retained", "PUB retained", "queries_olmes_instruct_pub_top100.jsonl", exclude_prefixes=("pub_2:", "pub_3:")),
    Spec("simpletom_mental", "SimpleToM mental", "queries_olmes_instruct_simpletom_top100.jsonl", include_prefixes=("simpletom_mental-state-qa:",)),
    Spec("mmlu_moral", "MMLU moral/hum.", "queries_olmes_instruct_mmlu_moral_top100.jsonl"),
    Spec("ethics_non_justice", "ETHICS non-justice", "queries_olmes_instruct_ethics_top100.jsonl", include_prefixes=("ethics_commonsense:", "ethics_deontology:", "ethics_utilitarianism:", "ethics_virtue:")),
    Spec("morables", "MORABLES", "queries_olmes_instruct_morables_top100.jsonl"),
    Spec("moralexceptqa_rbqa", "MoralExceptQA/RBQA", "queries_olmes_instruct_moralexceptqa_rbqa_top100.jsonl"),
]

def pretty_label(value: str) -> str:
    return LABEL_OVERRIDES.get(value, value.replace("_", " ").title())

def keep_query(query_id: str, spec: Spec) -> bool:
    if spec.include_prefixes and not query_id.startswith(spec.include_prefixes):
        return False
    return not query_id.startswith(spec.exclude_prefixes)
def load_bin_map(path: Path | None) -> dict[str, str]:
    if path is None:
        return {}
    out = {}
    with path.open() as fh:
        reader = csv.DictReader(fh)
        fields = reader.fieldnames or []
        topic_col = "topic_label" if "topic_label" in fields else "bin_topic"
        format_col = "format_label" if "format_label" in fields else "bin_format"
        for row in reader:
            topic = row.get(topic_col, "")
            fmt = row.get(format_col, "")
            if row.get("doc_id") and topic and fmt:
                out[row["doc_id"]] = f"{pretty_label(topic)} / {pretty_label(fmt)}"
    return out
def latex_escape(value: object) -> str:
    text = str(value)
    for old, new in {"\\": r"\textbackslash{}", "&": r"\&", "%": r"\%", "$": r"\$", "#": r"\#", "_": r"\_"}.items():
        text = text.replace(old, new)
    return text

def benchmark_scores(path: Path, spec: Spec) -> tuple[dict[str, tuple[float, int]], int, int]:
    support: dict[str, float] = defaultdict(float)
    hits: dict[str, int] = defaultdict(int)
    query_ids: set[str] = set()
    kept_rows = 0
    with path.open() as fh:
        for line in fh:
            row = json.loads(line)
            query_id = row["query_id"]
            if not keep_query(query_id, spec):
                continue
            rank = int(row["rank"])
            doc_id = row["doc_id"]
            support[doc_id] += 1 / math.log2(rank + 1)
            hits[doc_id] += 1
            query_ids.add(query_id)
            kept_rows += 1
    if not query_ids:
        raise ValueError(f"{path} produced no rows after filtering")
    norm = len(query_ids)
    scores = {doc: (value / norm, hits[doc]) for doc, value in support.items()}
    return scores, norm, kept_rows

def zscore_records(values: dict[str, tuple[float, int]], strong_z: float):
    supports = [value for value, _ in values.values()]
    mu = mean(supports)
    sigma = stdev(supports) if len(supports) > 1 else 0
    if sigma == 0:
        return []
    return (
        {
            "doc_id": doc_id,
            "doc_z": (support - mu) / sigma,
            "hits": hits,
            "strong": (support - mu) / sigma >= strong_z,
        }
        for doc_id, (support, hits) in values.items()
    )

def write_csv(rows: list[dict[str, object]], path: Path, fields: list[str]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    with path.open("w", newline="") as fh:
        writer = csv.DictWriter(fh, fieldnames=fields)
        writer.writeheader()
        writer.writerows(rows)

def write_latex(rows: list[dict[str, object]], path: Path, n_rows: int, strong_z: float) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    tex = [
        r"\begin{table}[!htbp]",
        r"\centering",
        r"\scriptsize",
        r"\setlength{\tabcolsep}{3pt}",
        r"\begin{tabularx}{\linewidth}{@{}p{0.25\linewidth}p{0.22\linewidth}rrrp{0.24\linewidth}@{}}",
        r"\toprule",
        r"Document ID & Bin & Strong & Pos. & Mean $z$ & Strong probes \\",
        r"\midrule",
    ]
    if rows:
        for row in rows[:n_rows]:
            tex.append(
                f"{latex_escape(row['doc_id'])} & {latex_escape(row['bin'])} & "
                f"{row['strong_probe_count']} & {row['positive_probe_count']} & "
                f"{float(str(row['mean_doc_z'])):+.2f} & {latex_escape(row['strong_probes'])} \\\\"
            )
    else:
        tex.append(r"\multicolumn{6}{@{}l@{}}{No document reached the recurrent threshold in two held-out probes.} \\")
    tex += [
        r"\bottomrule",
        r"\end{tabularx}",
        rf"\caption{{Top-$k$ document recurrence diagnostic across available held-out top-$k$ files. A document is recurrent only if its query-normalized rank-weighted support has $z \geq {strong_z:g}$ in at least two probes. Document IDs are reported without text snippets; bin labels appear only when a doc-to-bin manifest is supplied.}}",
        r"\label{tab:heldout-document-recurrence}",
        r"\end{table}",
    ]
    path.write_text("\n".join(tex) + "\n")

def main() -> None:
    parser = argparse.ArgumentParser()
    parser.add_argument("--topk-root", type=Path, default=DEFAULT_TOPK_ROOT)
    parser.add_argument("--out-dir", type=Path, default=DEFAULT_OUT)
    parser.add_argument("--latex-table", type=Path)
    parser.add_argument("--doc-bin-map", type=Path)
    parser.add_argument("--strong-z", type=float, default=2.0)
    parser.add_argument("--table-rows", type=int, default=8)
    args = parser.parse_args()

    merged: defaultdict[str, dict[str, Any]] = defaultdict(lambda: {"z_sum": 0.0, "count": 0, "positive": 0, "strong": 0, "max_z": -999.0, "hits": 0, "probes": []})
    summary = []
    for spec in SPECS:
        path = args.topk_root / spec.filename
        if not path.exists():
            raise FileNotFoundError(path)
        values, query_count, kept_rows = benchmark_scores(path, spec)
        summary.append({"benchmark": spec.display, "query_count": query_count, "kept_topk_rows": kept_rows, "scored_doc_count": len(values)})
        for row in zscore_records(values, args.strong_z):
            state = merged[str(row["doc_id"])]
            z = float(row["doc_z"])
            state["z_sum"] = float(state["z_sum"]) + z
            state["count"] = int(state["count"]) + 1
            state["positive"] = int(state["positive"]) + int(z > 0)
            state["strong"] = int(state["strong"]) + int(row["strong"])
            state["max_z"] = max(float(state["max_z"]), z)
            state["hits"] = int(state["hits"]) + int(row["hits"])
            if row["strong"]:
                state["probes"].append(spec.display)
    bin_map = load_bin_map(args.doc_bin_map)
    recurrent = []
    for doc_id, state in merged.items():
        if int(state["strong"]) < 2:
            continue
        recurrent.append({
            "doc_id": doc_id,
            "bin": bin_map.get(doc_id, "not joined"),
            "strong_probe_count": state["strong"],
            "positive_probe_count": state["positive"],
            "mean_doc_z": float(state["z_sum"]) / int(state["count"]),
            "max_doc_z": state["max_z"],
            "total_hits": state["hits"],
            "strong_probes": ", ".join(state["probes"]),
        })
    recurrent.sort(key=lambda row: (int(row["strong_probe_count"]), float(row["mean_doc_z"]), float(row["max_doc_z"])), reverse=True)
    args.out_dir.mkdir(parents=True, exist_ok=True)
    write_csv(summary, args.out_dir / "document_topk_summary.csv", ["benchmark", "query_count", "kept_topk_rows", "scored_doc_count"])
    write_csv(recurrent, args.out_dir / "recurrent_documents.csv", ["doc_id", "bin", "strong_probe_count", "positive_probe_count", "mean_doc_z", "max_doc_z", "total_hits", "strong_probes"])
    if args.latex_table:
        write_latex(recurrent, args.latex_table, args.table_rows, args.strong_z)
    print(f"wrote {args.out_dir / 'recurrent_documents.csv'} ({len(recurrent)} rows)")

if __name__ == "__main__":
    main()