social-data-attribution-w2/scripts/validation/export_classification_audit.py

"""Export a CSV audit of all SocialIQA ATOMIC reasoning type classifications.

Each row shows what label was assigned, why (which layer/pattern), context
disambiguation details, and whether the independent LLM verifier agreed.
"""

from __future__ import annotations

import csv
import json
import re
from pathlib import Path

from datasets import load_dataset

from data_attribution.evaluation.socialiqa_classifier import (
    MANUAL_OVERRIDES,
    _X_TO_O,
    _extract_context_actor,
    _extract_question_sentence,
    _extract_question_subject,
)
from data_attribution.evaluation.socialiqa_patterns import (
    O_PATTERNS,
    SIQA_NAMES,
    X_PATTERNS,
)
from data_attribution.evaluation.socialiqa_verified_labels import (
    VERIFIED_CORRECTIONS,
    VERIFIED_CORRECTIONS_BY_TEXT,
)

OLMES_HF_REPO = "HCAI-Lab/olmes-eval-olmo3-7b-base"
OUTPUT_PATH = Path("artifacts/evaluation/socialiqa_classification_audit.csv")
VERIFY_RESULTS_DIR = Path("runs/manifests/socialiqa_verify_results")

COLUMNS = [
    "query_id",
    "context",
    "question_sentence",
    "final_label",
    "classification_method",
    "pattern_matched",
    "context_actor",
    "question_subject",
    "x_to_o_applied",
    "sonnet_label",
    "sonnet_agrees",
    "is_correct",
    "predicted_answer",
    "correct_answer",
]


def _classify_with_trace(
    question: str, *, query_id: str | None = None, context: str | None = None
) -> dict[str, object]:
    question_sentence = _extract_question_sentence(question)
    trace: dict[str, object] = {
        "question_sentence": question_sentence,
        "context_actor": "",
        "question_subject": "",
        "x_to_o_applied": False,
        "classification_method": "",
        "pattern_matched": "",
        "final_label": "",
    }

    label: str | None = None

    for lbl, pattern in O_PATTERNS:
        if pattern.search(question_sentence):
            label = lbl
            trace["classification_method"] = "o_pattern"
            trace["pattern_matched"] = f'{lbl}: "{pattern.pattern[:50]}"'
            break

    if label is None:
        for lbl, pattern in X_PATTERNS:
            if pattern.search(question_sentence):
                label = lbl
                trace["classification_method"] = "x_pattern"
                trace["pattern_matched"] = f'{lbl}: "{pattern.pattern[:50]}"'
                if context is not None and lbl in _X_TO_O:
                    actor = _extract_context_actor(context)
                    subject = _extract_question_subject(question_sentence, lbl)
                    trace["context_actor"] = actor or ""
                    trace["question_subject"] = subject or ""
                    if actor and subject and subject in SIQA_NAMES and subject != actor:
                        label = _X_TO_O[lbl]
                        trace["classification_method"] = "x_to_o_disambiguation"
                        trace["x_to_o_applied"] = True
                break

    if label is None and query_id and query_id in MANUAL_OVERRIDES:
        label = MANUAL_OVERRIDES[query_id]
        trace["classification_method"] = "manual_override"

    if query_id:
        correction = VERIFIED_CORRECTIONS.get(query_id)
        if correction is not None and label != correction:
            label = correction
            trace["classification_method"] = "verified_correction"
        text_corrections = VERIFIED_CORRECTIONS_BY_TEXT.get(query_id)
        if text_corrections:
            for substr, corr in text_corrections:
                if substr in question and label != corr:
                    label = corr
                    trace["classification_method"] = "verified_correction_by_text"

    trace["final_label"] = label or ""
    return trace


def _load_sonnet_labels() -> dict[str, str]:
    labels: dict[str, str] = {}
    if not VERIFY_RESULTS_DIR.exists():
        return labels
    for f in sorted(VERIFY_RESULTS_DIR.glob("*.jsonl")):
        for line in f.read_text().splitlines():
            if not line.strip():
                continue
            r = json.loads(line)
            labels[r["query_id"]] = r["sonnet_label"]
    return labels


def main() -> None:
    print(f"Loading {OLMES_HF_REPO}/socialiqa ...")
    dataset = load_dataset(OLMES_HF_REPO, name="socialiqa", split="train")

    sonnet_labels = _load_sonnet_labels()
    print(f"Loaded {len(sonnet_labels)} Sonnet verification labels")

    OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True)
    with OUTPUT_PATH.open("w", newline="", encoding="utf-8") as fh:
        writer = csv.DictWriter(fh, fieldnames=COLUMNS)
        writer.writeheader()

        for row in dataset:
            qid = row["query_id"]
            query_text = row.get("query_text", "")
            m = re.search(r"Question:\s*(.+?)(?:\n|$)", query_text)
            full_text = m.group(1).strip() if m else query_text

            sentences = re.split(r"(?<=[.!?])\s+", full_text.strip())
            context = " ".join(sentences[:-1]).strip() if len(sentences) > 1 else ""

            trace = _classify_with_trace(
                full_text, query_id=qid, context=context or None
            )

            sonnet = sonnet_labels.get(qid, "")
            if sonnet:
                agrees = "yes" if trace["final_label"] == sonnet else "no"
            else:
                agrees = "not_reviewed"

            writer.writerow(
                {
                    "query_id": qid,
                    "context": context,
                    "question_sentence": trace["question_sentence"],
                    "final_label": trace["final_label"],
                    "classification_method": trace["classification_method"],
                    "pattern_matched": trace["pattern_matched"],
                    "context_actor": trace["context_actor"],
                    "question_subject": trace["question_subject"],
                    "x_to_o_applied": trace["x_to_o_applied"],
                    "sonnet_label": sonnet,
                    "sonnet_agrees": agrees,
                    "is_correct": row["is_correct"],
                    "predicted_answer": row["predicted_answer"],
                    "correct_answer": row["correct_answer"],
                }
            )

    print(f"Wrote {len(dataset)} rows to {OUTPUT_PATH}")


if __name__ == "__main__":
    main()