social-data-attribution-w2/scripts/query_data/generate_soc171_queries.py

"""Generate query JSONL files for SOC-171 (GSM8K + ARC, base + instruct)."""

import logging
from pathlib import Path

from data_attribution.recipes.olmes.manifest_writer import (
    FORMAT_BASE,
    FORMAT_INSTRUCT,
    write_olmes_manifest,
)

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

BASE_HF_REPO = "HCAI-Lab/olmes-eval-olmo3-7b-base"
INSTRUCT_HF_REPO = "HCAI-Lab/olmes-eval-olmo3-7b-instruct-base"

SUBSETS = ("gsm8k", "arc_easy", "arc_challenge")

VARIANTS = {
    "base": {
        "hf_repo": BASE_HF_REPO,
        "output_format": FORMAT_BASE,
        "file_prefix": "olmes_",
    },
    "instruct_base": {
        "hf_repo": INSTRUCT_HF_REPO,
        "output_format": FORMAT_INSTRUCT,
        "file_prefix": "olmes_instruct_",
    },
}


def main() -> None:
    for variant_name, cfg in VARIANTS.items():
        output_dir = Path(f"queries/{variant_name}")
        output_dir.mkdir(parents=True, exist_ok=True)

        for subset in SUBSETS:
            output_path = output_dir / f"{cfg['file_prefix']}{subset}.jsonl"
            log.info("Generating %s", output_path)
            count = write_olmes_manifest(
                output_path,
                subset,
                hf_repo=cfg["hf_repo"],
                output_format=cfg["output_format"],
            )
            log.info("Wrote %d rows to %s", count, output_path)

    log.info("Done. All query JSONL files generated.")


if __name__ == "__main__":
    main()