bikoret/kushya_generator/data/prepare_training.py

"""
Prepare training data for the kushya generator.

Reads .jsonl files from training_pairs/, validates each record against the
four kushya requirements (Architecture §3.6), and writes:

  train.jsonl   — 80%, used for QLoRA training
  eval.jsonl    — 20%, held out for evaluation

Validation rejects pairs that don't meet the requirements rather than
silently degrading training quality.
"""

import json
import os
import re
import sys
from pathlib import Path
from typing import Iterable

KUSHYA_DIR = Path(__file__).resolve().parent / "training_pairs"
OUTPUT_DIR = Path(__file__).resolve().parent

# The four requirements from Architecture §3.6 fine_tune_objective.
REQUIRED_TYPES = {
    "scope_narrowing",
    "scope_application",
    "internal_tension",
    "exception_test",
    "harmonization",
    "edge_case",
    "rabbinic_constraint",
    "method_question",
    "subject_identification",
}


class ValidationError(Exception):
    pass


def validate_pair(pair: dict, source: str, line_no: int) -> None:
    """Apply the four kushya requirements + schema check."""
    where = f"{source}:{line_no}"

    if not isinstance(pair, dict):
        raise ValidationError(f"{where}: not a JSON object")

    for field in ("interpretation", "kushya", "torah_refs", "type"):
        if field not in pair:
            raise ValidationError(f"{where}: missing field {field!r}")

    if not isinstance(pair["interpretation"], str) or not pair["interpretation"].strip():
        raise ValidationError(f"{where}: interpretation must be a non-empty string")

    if not isinstance(pair["kushya"], str) or not pair["kushya"].strip():
        raise ValidationError(f"{where}: kushya must be a non-empty string")

    if not isinstance(pair["torah_refs"], list):
        raise ValidationError(f"{where}: torah_refs must be a list")

    if pair["type"] not in REQUIRED_TYPES:
        raise ValidationError(
            f"{where}: type {pair['type']!r} not in known types {sorted(REQUIRED_TYPES)}"
        )

    # Requirement 3: kushya must be expressible as a question.
    # Strict: must contain '?'. Interrogative words alone are not enough
    # — "this is unclear" contains "is" but isn't a question.
    if "?" not in pair["kushya"]:
        raise ValidationError(
            f"{where}: kushya must be expressible as a question (must contain '?') — "
            f"got {pair['kushya'][:60]!r}"
        )

    # Requirement 1: must reference a specific Torah passage. Either in
    # torah_refs (preferred) or in the kushya text itself.
    torah_book_pattern = (
        r"\b(Bereishit|Shemot|Vayikra|Bamidbar|Devarim|Mishnah|Gemara|Mishlei|"
        r"Tehillim|Yeshayahu|Yechezkel|Yirmiyahu|Sanhedrin|Eruvin|Yoma|Berakhot)\b"
    )
    if not pair["torah_refs"] and not re.search(torah_book_pattern, pair["kushya"]):
        raise ValidationError(
            f"{where}: kushya must cite a specific Torah passage "
            f"(empty torah_refs and no book name in text)"
        )

    # Requirement 4: not resolvable by trivial clarification.
    # Heuristic: kushya should be at least 40 characters of substantive text.
    if len(pair["kushya"].strip()) < 40:
        raise ValidationError(
            f"{where}: kushya is too short to be a substantive challenge "
            f"({len(pair['kushya'].strip())} chars)"
        )


def load_jsonl(path: Path) -> Iterable[tuple[int, dict]]:
    with open(path) as f:
        for i, line in enumerate(f, start=1):
            line = line.strip()
            if not line:
                continue
            try:
                yield i, json.loads(line)
            except json.JSONDecodeError as e:
                raise ValidationError(f"{path.name}:{i}: invalid JSON — {e}")


def to_training_format(pair: dict) -> dict:
    """Format a validated pair for fine-tuning.

    Architecture §3.6 training_data_format:
      - input: "Proposed interpretation: [text]"
      - output: "Kushya: [challenge question grounded in Torah text]"
    """
    refs_str = ", ".join(pair["torah_refs"]) if pair["torah_refs"] else "(no explicit refs)"
    return {
        "input": f"Proposed interpretation: {pair['interpretation']}",
        "output": f"Kushya: {pair['kushya']}",
        "metadata": {
            "type": pair["type"],
            "torah_refs": pair["torah_refs"],
            "refs_str": refs_str,
        },
    }


def split_train_eval(records: list[dict], train_fraction: float = 0.8) -> tuple[list[dict], list[dict]]:
    """Stratified split by type."""
    by_type: dict[str, list[dict]] = {}
    for r in records:
        by_type.setdefault(r["metadata"]["type"], []).append(r)

    train: list[dict] = []
    evaln: list[dict] = []
    for typ, items in by_type.items():
        cut = max(1, int(len(items) * train_fraction)) if len(items) > 1 else len(items)
        train.extend(items[:cut])
        evaln.extend(items[cut:])
    return train, evaln


def main() -> int:
    if not KUSHYA_DIR.exists():
        print(f"ERROR: {KUSHYA_DIR} does not exist", file=sys.stderr)
        return 1

    all_pairs: list[dict] = []
    errors: list[str] = []

    for jsonl_file in sorted(KUSHYA_DIR.glob("*.jsonl")):
        for line_no, pair in load_jsonl(jsonl_file):
            try:
                validate_pair(pair, jsonl_file.name, line_no)
                all_pairs.append(to_training_format(pair))
            except ValidationError as e:
                errors.append(str(e))

    if errors:
        print("Validation errors — pairs rejected:", file=sys.stderr)
        for e in errors:
            print(f"  {e}", file=sys.stderr)
        print(file=sys.stderr)

    if not all_pairs:
        print("No valid pairs found. Cannot prepare training set.", file=sys.stderr)
        return 1

    train, evaln = split_train_eval(all_pairs)

    train_path = OUTPUT_DIR / "train.jsonl"
    eval_path = OUTPUT_DIR / "eval.jsonl"

    with open(train_path, "w") as f:
        for r in train:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")
    with open(eval_path, "w") as f:
        for r in evaln:
            f.write(json.dumps(r, ensure_ascii=False) + "\n")

    print(f"Validated:  {len(all_pairs)} pairs")
    print(f"Rejected:   {len(errors)}")
    print(f"Train:      {len(train)} → {train_path}")
    print(f"Eval:       {len(evaln)} → {eval_path}")

    return 0 if not errors else 1


if __name__ == "__main__":
    sys.exit(main())