scripts/06_build_eval_sets.py

"""
scripts/06_build_eval_sets.py
Build held-out evaluation sets from the synthetic data:
  - eval_pictogram_to_sentence.jsonl  (500 examples, 100 per persona)
  - eval_audio_disambig.jsonl          (200 examples with audio paths)
  - eval_visual_ground.jsonl           (150 examples with image paths)
  - eval_adapter_specificity.jsonl     (100 fixed prompts, cross-persona)
"""

import json
import logging
import random
from pathlib import Path

import pandas as pd

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

SYNTH_DIR = Path("data/synth")
EVAL_DIR = Path("data/eval")
PROCESSED_DIR = Path("data/processed")
RANDOM_SEED = 42


def load_persona_jsonl(persona_id: str) -> list[dict]:
    path = SYNTH_DIR / f"persona_{persona_id}.jsonl"
    if not path.exists():
        log.warning(f"No synth data for persona: {persona_id}")
        return []
    with open(path, encoding="utf-8") as f:
        return [json.loads(l) for l in f if l.strip()]


def build_pictogram_to_sentence_eval(n_per_persona: int = 100) -> list[dict]:
    """Sample n examples per persona as held-out eval set."""
    all_eval = []
    personas = ["ananya", "arjun", "priya", "rohan", "zara"]

    for pid in personas:
        records = load_persona_jsonl(pid)
        if not records:
            log.warning(f"Skipping {pid} — no data")
            continue
        random.shuffle(records)
        sampled = records[:n_per_persona]
        for r in sampled:
            all_eval.append({
                "persona_id": pid,
                "pictogram_sequence": r.get("seq", []),
                "context": r.get("context", ""),
                "target_sentence": r.get("target", ""),
                "alternatives": r.get("alternatives", []),
            })
        log.info(f"  {pid}: sampled {len(sampled)} eval examples")

    return all_eval


def build_audio_disambig_eval(n: int = 200) -> list[dict]:
    """
    Build audio disambiguation eval set.
    Each example: (pictogram_seq + audio_path) → target_sentence.
    Audio paths reference files in data/processed/audio/.
    """
    audio_manifest = Path("data/processed/audio_manifest.jsonl")
    if not audio_manifest.exists():
        log.warning("Audio manifest not found. Run 03_download_audio_corpora.py first.")
        # Generate placeholder examples
        return [
            {
                "persona_id": "arjun",
                "pictogram_sequence": ["water", "want"],
                "audio_path": "PLACEHOLDER",
                "audio_transcript_hint": "pani",
                "target_sentence": "Mujhe pani chahida",
                "intent_class": "request_food_drink",
            }
        ]

    with open(audio_manifest) as f:
        audio_clips = [json.loads(l) for l in f if l.strip()]

    random.shuffle(audio_clips)
    personas = ["ananya", "arjun", "priya", "rohan", "zara"]

    records = []
    for i, clip in enumerate(audio_clips[:n]):
        pid = personas[i % len(personas)]
        persona_data = load_persona_jsonl(pid)
        if not persona_data:
            continue
        base = random.choice(persona_data)
        records.append({
            "persona_id": pid,
            "pictogram_sequence": base.get("seq", []),
            "audio_path": clip["path"],
            "audio_transcript_hint": clip.get("transcript", ""),
            "target_sentence": base.get("target", ""),
            "intent_class": _classify_intent(base.get("seq", [])),
        })

    return records


def build_visual_ground_eval(n: int = 150) -> list[dict]:
    """
    Build visual grounding eval set.
    Each example: (home_photo + pictogram_seq) → sentence with visual reference.
    Uses placeholder image paths since real home photos aren't available.
    """
    personas = ["ananya", "arjun", "priya", "rohan", "zara"]
    records = []
    for i in range(n):
        pid = personas[i % len(personas)]
        persona_data = load_persona_jsonl(pid)
        if not persona_data:
            continue
        base = random.choice(persona_data)
        records.append({
            "persona_id": pid,
            "pictogram_sequence": base.get("seq", []),
            "image_path": f"data/raw/home_photos/{pid}/sample_{i:03d}.jpg",  # placeholder
            "target_sentence": base.get("target", ""),
            "visual_reference": f"object in image #{i % 5 + 1}",
        })
    return records


def build_adapter_specificity_eval() -> list[dict]:
    """
    Build the fixed prompt set used for adapter-specificity evaluation.
    100 prompts, all same prompts applied to each persona adapter.
    """
    # These are universal prompts that any CP child might use
    UNIVERSAL_PROMPTS = [
        {"seq": ["mother", "want"], "context": "morning"},
        {"seq": ["water", "give"], "context": "afternoon"},
        {"seq": ["pain", "help"], "context": "general"},
        {"seq": ["food", "more"], "context": "lunch"},
        {"seq": ["sleep", "want"], "context": "evening"},
        {"seq": ["play", "want"], "context": "afternoon"},
        {"seq": ["toilet", "go"], "context": "general"},
        {"seq": ["doctor", "no"], "context": "hospital"},
        {"seq": ["happy", "today"], "context": "morning"},
        {"seq": ["love", "mother"], "context": "bedtime"},
        {"seq": ["school", "go"], "context": "morning"},
        {"seq": ["friend", "come"], "context": "afternoon"},
        {"seq": ["TV", "watch"], "context": "evening"},
        {"seq": ["cold", "milk", "want"], "context": "afternoon"},
        {"seq": ["medicine", "no", "want"], "context": "evening"},
        {"seq": ["tired", "bed"], "context": "night"},
        {"seq": ["hot", "stop"], "context": "bathing"},
        {"seq": ["book", "read"], "context": "bedtime"},
        {"seq": ["father", "come", "now"], "context": "general"},
        {"seq": ["outside", "go", "want"], "context": "afternoon"},
    ] * 5  # 20 prompts × 5 = 100

    records = []
    for i, p in enumerate(UNIVERSAL_PROMPTS[:100]):
        records.append({
            "prompt_id": f"spec_{i:03d}",
            "pictogram_sequence": p["seq"],
            "context": p["context"],
            "expected_persona_styles": {
                "ananya": "Tamil/English style, short",
                "arjun": "Punjabi/Hindi/English mix",
                "priya": "Bengali/English, very short (4yo)",
                "rohan": "Hindi/English, older child register",
                "zara": "Marathi/English, 7yo register",
            },
        })
    return records


def _classify_intent(seq: list[str]) -> str:
    """Simple rule-based intent classifier for audio eval labels."""
    seq_lower = [s.lower() for s in seq]
    if any(w in seq_lower for w in ["food", "eat", "drink", "water", "milk", "more", "hungry"]):
        return "request_food_drink"
    if any(w in seq_lower for w in ["pain", "hurt", "stop", "help"]):
        return "report_discomfort"
    if any(w in seq_lower for w in ["play", "fun", "friend", "TV", "outside"]):
        return "request_activity"
    if any(w in seq_lower for w in ["love", "happy", "hug", "kiss"]):
        return "social_emotional"
    if any(w in seq_lower for w in ["toilet", "bath", "sleep"]):
        return "personal_care"
    return "general"


def save_jsonl(records: list[dict], path: Path):
    path.parent.mkdir(parents=True, exist_ok=True)
    with open(path, "w", encoding="utf-8") as f:
        for r in records:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    log.info(f"Saved {len(records)} records → {path}")


def main():
    random.seed(RANDOM_SEED)
    EVAL_DIR.mkdir(parents=True, exist_ok=True)

    log.info("Building eval_pictogram_to_sentence...")
    p2s = build_pictogram_to_sentence_eval(100)
    save_jsonl(p2s, EVAL_DIR / "eval_pictogram_to_sentence.jsonl")

    log.info("Building eval_audio_disambig...")
    audio = build_audio_disambig_eval(200)
    save_jsonl(audio, EVAL_DIR / "eval_audio_disambig.jsonl")

    log.info("Building eval_visual_ground...")
    visual = build_visual_ground_eval(150)
    save_jsonl(visual, EVAL_DIR / "eval_visual_ground.jsonl")

    log.info("Building eval_adapter_specificity...")
    spec = build_adapter_specificity_eval()
    save_jsonl(spec, EVAL_DIR / "eval_adapter_specificity.jsonl")

    log.info("\nEval sets complete:")
    log.info(f"  pictogram_to_sentence: {len(p2s)}")
    log.info(f"  audio_disambig:        {len(audio)}")
    log.info(f"  visual_ground:         {len(visual)}")
    log.info(f"  adapter_specificity:   {len(spec)}")


if __name__ == "__main__":
    main()