scripts/generate_hinglish_protocol.py

#!/usr/bin/env python3
"""Generate deterministic speaker-disjoint Hinglish recording assignments."""

from __future__ import annotations

import argparse
import json
from collections import Counter
from pathlib import Path

from turn_detection.collection import build_assignments


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--speakers", type=int, default=30)
    parser.add_argument("--items-per-speaker", type=int, default=30)
    parser.add_argument("--seed", type=int, default=20260823)
    parser.add_argument("--output", default="data/collection/assignments.jsonl")
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    assignments = build_assignments(
        speaker_count=args.speakers,
        items_per_speaker=args.items_per_speaker,
        seed=args.seed,
    )
    destination = Path(args.output)
    destination.parent.mkdir(parents=True, exist_ok=True)
    with destination.open("w", encoding="utf-8") as handle:
        for assignment in assignments:
            handle.write(
                json.dumps(assignment.to_dict(), ensure_ascii=False, sort_keys=True) + "\n"
            )
    summary = {
        "assignments": len(assignments),
        "speakers": len({item.speaker_id for item in assignments}),
        "split_examples": Counter(item.split for item in assignments),
        "split_speakers": {
            split: len({item.speaker_id for item in assignments if item.split == split})
            for split in ("train", "dev", "test")
        },
        "targets": Counter(item.target for item in assignments),
        "scenarios": Counter(item.scenario for item in assignments),
        "seed": args.seed,
    }
    summary_path = destination.with_suffix(".summary.json")
    summary_path.write_text(json.dumps(summary, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    print(json.dumps(summary, indent=2, sort_keys=True))


if __name__ == "__main__":
    main()