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

"""Verify evaluation query counts against documented values in EVALUATION_PROTOCOL.md.

Loads each benchmark subset from the HF dataset, counts rows per subset
and per MMLU subject, validates query_id uniqueness, and reports mismatches.
"""

from __future__ import annotations

import argparse
import logging

from datasets import load_dataset

from data_attribution.evaluation.query_metadata import (
    OLMES_HF_REPO,
)

from _verify_helpers import (
    EXPECTED_COUNTS,
    EXPECTED_MMLU_SOCIAL_SCIENCE,
    EXPECTED_MMLU_STEM,
    check_mmlu_subjects,
    check_query_id_uniqueness,
    check_socialiqa_reasoning_coverage,
    check_subset_count,
)

logger = logging.getLogger(__name__)


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description="Verify evaluation query counts.")
    parser.add_argument("--hf-repo", default=OLMES_HF_REPO)
    args = parser.parse_args(argv)

    logging.basicConfig(
        level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s"
    )

    errors: list[str] = []
    all_query_ids: list[str] = []
    info_lines: list[str] = []

    for subset in EXPECTED_COUNTS:
        logger.info("Loading %s/%s", args.hf_repo, subset)
        dataset = load_dataset(args.hf_repo, name=subset, split="train")
        actual_count = len(dataset)

        print(f"\n{subset}: {actual_count} rows")
        errors.extend(check_subset_count(subset, actual_count))

        for row in dataset:
            all_query_ids.append(row["query_id"])

        if subset == "mmlu_social_science":
            errors.extend(
                check_mmlu_subjects(dataset, subset, EXPECTED_MMLU_SOCIAL_SCIENCE)
            )
        elif subset == "mmlu_stem":
            errors.extend(check_mmlu_subjects(dataset, subset, EXPECTED_MMLU_STEM))
        elif subset == "socialiqa":
            info_lines.extend(check_socialiqa_reasoning_coverage(dataset))

    dup_errors, dup_warnings = check_query_id_uniqueness(all_query_ids)
    errors.extend(dup_errors)

    total = len(all_query_ids)
    expected_total = sum(EXPECTED_COUNTS.values())
    print(f"\nTotal queries: {total} (expected: {expected_total})")

    if info_lines:
        print("\n--- SocialIQA Reasoning Types ---")
        for line in info_lines:
            print(line)

    if dup_warnings:
        print("\n--- WARNINGS ---")
        for warning in dup_warnings:
            print(warning)

    if errors:
        print("\n--- ERRORS ---")
        for error in errors:
            print(error)
        return 1

    print("\nAll counts verified.")
    return 0


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