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

"""Helper functions and expected counts for verify_query_counts.py."""

from __future__ import annotations

import re
from collections import Counter

from data_attribution.evaluation.socialiqa_classifier import (
    classify_socialiqa_reasoning_type,
)

EXPECTED_COUNTS: dict[str, int] = {
    "socialiqa": 10_000,
    "mmlu_social_science": 3_077,
    "mmlu_stem": 3_018,
    "gsm8k": 1_319,
}

EXPECTED_MMLU_SOCIAL_SCIENCE: dict[str, int] = {
    "econometrics": 114,
    "high_school_geography": 198,
    "high_school_government_and_politics": 193,
    "high_school_macroeconomics": 390,
    "high_school_microeconomics": 238,
    "high_school_psychology": 545,
    "human_sexuality": 131,
    "professional_psychology": 612,
    "public_relations": 110,
    "security_studies": 245,
    "sociology": 201,
    "us_foreign_policy": 100,
}

EXPECTED_MMLU_STEM: dict[str, int] = {
    "abstract_algebra": 100,
    "astronomy": 152,
    "college_biology": 144,
    "college_chemistry": 100,
    "college_computer_science": 100,
    "college_mathematics": 100,
    "college_physics": 102,
    "computer_security": 100,
    "conceptual_physics": 235,
    "electrical_engineering": 145,
    "elementary_mathematics": 378,
    "high_school_biology": 310,
    "high_school_chemistry": 203,
    "high_school_computer_science": 100,
    "high_school_mathematics": 270,
    "high_school_physics": 151,
    "high_school_statistics": 216,
    "machine_learning": 112,
}

KNOWN_DUPLICATE_PREFIX = "socialiqa:"


def check_subset_count(subset: str, actual: int) -> list[str]:
    expected = EXPECTED_COUNTS.get(subset)
    if expected is None:
        return [f"  No expected count for {subset}"]
    if actual != expected:
        return [f"  MISMATCH {subset}: expected={expected}, actual={actual}"]
    return []


def check_mmlu_subjects(
    dataset, subset: str, expected_subjects: dict[str, int]
) -> list[str]:
    subject_counts: Counter[str] = Counter()
    for row in dataset:
        task_name = row.get("task_name", "")
        subject = (
            task_name[len("mmlu_") :] if task_name.startswith("mmlu_") else task_name
        )
        subject_counts[subject] += 1

    errors: list[str] = []
    all_subjects = set(expected_subjects) | set(subject_counts)
    for subject in sorted(all_subjects):
        expected = expected_subjects.get(subject, 0)
        actual = subject_counts.get(subject, 0)
        if expected != actual:
            errors.append(
                f"  MISMATCH {subset}/{subject}: expected={expected}, actual={actual}"
            )
    extra = set(subject_counts) - set(expected_subjects)
    if extra:
        errors.append(f"  EXTRA subjects in {subset}: {sorted(extra)}")
    missing = set(expected_subjects) - set(subject_counts)
    if missing:
        errors.append(f"  MISSING subjects in {subset}: {sorted(missing)}")
    return errors


def check_query_id_uniqueness(
    all_query_ids: list[str],
) -> tuple[list[str], list[str]]:
    counts = Counter(all_query_ids)
    duplicates = {qid: count for qid, count in counts.items() if count > 1}
    if not duplicates:
        return [], []
    known = {
        qid: c
        for qid, c in duplicates.items()
        if qid.startswith(KNOWN_DUPLICATE_PREFIX)
    }
    unknown = {
        qid: c
        for qid, c in duplicates.items()
        if not qid.startswith(KNOWN_DUPLICATE_PREFIX)
    }
    warnings: list[str] = []
    errors: list[str] = []
    if known:
        sample = dict(list(known.items())[:5])
        warnings.append(
            f"  KNOWN duplicate query_ids in SocialIQA ({len(known)} total): {sample}"
        )
    if unknown:
        sample = dict(list(unknown.items())[:5])
        errors.append(
            f"  UNEXPECTED duplicate query_ids ({len(unknown)} total): {sample}"
        )
    return errors, warnings


def check_socialiqa_reasoning_coverage(dataset) -> list[str]:
    total = 0
    classified = 0
    type_counts: Counter[str] = Counter()
    unclassified_samples: list[str] = []

    for row in dataset:
        total += 1
        query_text = row.get("query_text", "")
        question_match = re.search(r"Question:\s*(.+?)(?:\n|$)", query_text)
        question = question_match.group(1).strip() if question_match else query_text
        reasoning_type = classify_socialiqa_reasoning_type(
            question, query_id=row.get("query_id")
        )
        if reasoning_type is not None:
            classified += 1
            type_counts[reasoning_type] += 1
        elif len(unclassified_samples) < 5:
            unclassified_samples.append(question[:120])

    coverage = classified / total if total else 0.0
    lines = [
        f"  SocialIQA reasoning type coverage: {classified}/{total} ({coverage:.1%})"
    ]
    for rtype, count in sorted(type_counts.items(), key=lambda x: -x[1]):
        lines.append(f"    {rtype}: {count}")
    if unclassified_samples:
        lines.append(f"  Unclassified samples ({total - classified} total):")
        for sample in unclassified_samples:
            lines.append(f"    {sample!r}")
    return lines