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

"""Extract SocialIQA questions with current labels for LLM verification.

Loads the socialiqa subset from HF, extracts question text and current
regex-based reasoning_type labels, and writes to a flat JSONL for
distribution across verification subagents.
"""

from __future__ import annotations

import json
import re
from pathlib import Path

from datasets import load_dataset

from data_attribution.evaluation.query_metadata import (
    OLMES_HF_REPO,
    classify_socialiqa_reasoning_type,
)

OUTPUT_PATH = Path("runs/manifests/socialiqa_verify_input.jsonl")


def extract_question(query_text: str) -> str:
    match = re.search(r"Question:\s*(.+?)(?:\n|$)", query_text)
    return match.group(1).strip() if match else query_text


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

    OUTPUT_PATH.parent.mkdir(parents=True, exist_ok=True)
    count = 0
    with OUTPUT_PATH.open("w", encoding="utf-8") as fh:
        for row in dataset:
            query_id = row["query_id"]
            query_text = row.get("query_text", "")
            question = extract_question(query_text)
            current_label = classify_socialiqa_reasoning_type(
                question, query_id=query_id
            )
            fh.write(
                json.dumps(
                    {
                        "query_id": query_id,
                        "question": question,
                        "current_label": current_label,
                    }
                )
                + "\n"
            )
            count += 1

    print(f"Wrote {count} records to {OUTPUT_PATH}")


if __name__ == "__main__":
    main()