datagen_common.py

"""Shared dataset-build machinery for the two programs in this repo.

`quixi_math_datagen.py` (the generator is the solver) and
`quixi_proofs_datagen.py` (the generator is the prover) share everything
about *building* a dataset — skill grouping, weighted sampling, dedup,
stats, structural validation shape, JSONL writing — and differ only in
their generator registries, their final-step op-code (``Z`` vs ``QED``),
and their validators. This module holds the shared core; each CLI keeps
its registry, its pool policy, and a thin validator wrapper.

Extracted from quixi_math_datagen.py per plans/proofs_plan.md Phase 0
(pin 2): sampling, weights, dedup, stats, writer move here; validators
stay per-program (both call `validate_structure`).
"""
import json
import os
import random


def group_into_skills(gen_pool):
    """Groups generator instances into skills keyed by class name.

    Variant instances of one class (e.g. the four FractionOpGenerator ops)
    form a single skill so each skill gets equal sampling probability by
    default, regardless of how many instances implement it.
    """
    skills = {}
    for gen in gen_pool:
        skills.setdefault(gen.__class__.__name__, []).append(gen)
    return skills


def parse_weights(spec, available):
    """Parses a --weights spec into {skill_name: positive float}.

    Accepts inline "SkillA=2.5,SkillB=0.5" or a path to a JSON file holding
    an object of skill: weight. Unknown skills, non-numeric or non-positive
    weights raise ValueError.
    """
    if isinstance(spec, dict):
        raw = dict(spec)
    elif os.path.isfile(spec) or spec.endswith(".json"):
        try:
            with open(spec, encoding="utf-8") as fp:
                raw = json.load(fp)
        except (OSError, json.JSONDecodeError) as e:
            raise ValueError(f"Could not read weights file {spec!r}: {e}")
        if not isinstance(raw, dict):
            raise ValueError(
                f"Weights file {spec!r} must contain a JSON object of "
                f"{{\"SkillName\": weight}}.")
    else:
        raw = {}
        for part in spec.split(","):
            part = part.strip()
            if not part:
                continue
            name, sep, val = part.partition("=")
            if not sep:
                raise ValueError(
                    f"Bad weight entry {part!r}; expected SkillName=NUMBER")
            raw[name.strip()] = val.strip()

    weights = {}
    for name, val in raw.items():
        try:
            num = float(val)
        except (TypeError, ValueError):
            raise ValueError(f"Weight for {name!r} must be a number, got {val!r}")
        if num <= 0:
            raise ValueError(f"Weight for {name!r} must be positive, got {num}")
        weights[name] = num

    unknown = set(weights) - set(available)
    if unknown:
        raise ValueError(
            f"Unknown skill(s) in --weights: {', '.join(sorted(unknown))}. "
            f"Available: {', '.join(sorted(available))}")
    return weights


def validate_structure(example, final_op, grade_levels, delim="|"):
    """Structural validation shared by both programs; raises ValueError.

    Deliberately vocabulary-agnostic about step op-codes (the math
    scratchpad vocabulary is organic; the proofs line grammar is enforced
    separately by its checker) — only the step *shape*, the required
    record fields, and the final `<final_op>|<final_answer>` step are
    enforced here.
    """
    if not isinstance(example, dict):
        raise ValueError(f"example must be a dict, got {type(example).__name__}")
    for key in ("problem_id", "operation", "problem", "steps", "final_answer",
                "grade_level", "difficulty"):
        if key not in example:
            raise ValueError(f"missing required key: {key}")
    steps = example["steps"]
    if not isinstance(steps, list) or not steps:
        raise ValueError("steps must be a non-empty list")
    for i, s in enumerate(steps):
        if not isinstance(s, str) or not s:
            raise ValueError(f"step {i} must be a non-empty string, got {s!r}")
        fields = s.split(delim)
        if not fields[0]:
            raise ValueError(f"step {i} has an empty op-code: {s!r}")
        # The final step is exempt from the field-count check: it is
        # constrained by exact equality with final_answer below, and critic
        # formats put a pipe-format step inside the answer itself.
        if len(fields) > 5 and i != len(steps) - 1:
            raise ValueError(f"step {i} has more than 4 payload fields: {s!r}")
    expected_final = f"{final_op}{delim}{example['final_answer']}"
    if steps[-1] != expected_final:
        raise ValueError(
            f"last step must be {expected_final!r}, got {steps[-1]!r}")
    if example["grade_level"] not in grade_levels:
        raise ValueError(f"invalid grade_level: {example['grade_level']!r}")
    difficulty = example["difficulty"]
    if not isinstance(difficulty, int) or isinstance(difficulty, bool) \
            or not 1 <= difficulty <= 5:
        raise ValueError(f"difficulty must be an int in 1..5, got {difficulty!r}")


def write_jsonl(fp, obj):
    """Writes a JSON object to a file handle, one object per line."""
    fp.write(json.dumps(obj, ensure_ascii=False) + "\n")


def instance_label(gen_instance):
    """Display label for stats: class name plus variant symbol if present."""
    name = gen_instance.__class__.__name__
    if hasattr(gen_instance, "op_symbol"):
        name += f"({gen_instance.op_symbol})"
    return name


def run_build(n, path, seed, gen_pool, validate_fn, stamp_fn,
              weights=None, allow_duplicates=False):
    """The shared dataset build loop; returns a summary dict.

    Sampling is balanced per skill (generator class): each skill gets equal
    probability by default, and a skill's variant instances are chosen
    uniformly within it. `weights` (a {skill_name: float} dict or --weights
    spec string) overrides individual skill weights; unlisted skills keep
    weight 1.0. Exact repeats of (operation, problem) are skipped unless
    allow_duplicates is set.
    """
    if seed is not None:
        random.seed(seed)
    if not gen_pool:
        raise ValueError("No generators selected; cannot build dataset.")
    skills = group_into_skills(gen_pool)
    skill_names = list(skills)
    if weights:
        weights = parse_weights(weights, skill_names)
        skill_weights = [weights.get(name, 1.0) for name in skill_names]
    else:
        skill_weights = None

    count = 0
    attempts = 0
    seen = set()
    stats = {}
    # Generous budget: dedup can reject heavily when a skill's problem space
    # is small, so allow many more attempts than examples...
    max_attempts = n * 10 + 1000
    # ...but stop early if nothing new has been accepted for a long stretch
    # (likely an exhausted problem space).
    consecutive_rejects = 0
    max_consecutive_rejects = max(2000, n)

    print(f"Attempting to generate {n} examples...")
    with open(path, "w", encoding="utf-8") as fp:
        while count < n and attempts < max_attempts:
            if consecutive_rejects >= max_consecutive_rejects:
                print(f"WARN: no new examples accepted in the last "
                      f"{consecutive_rejects} attempts; the problem space of the "
                      f"selected skills is likely exhausted. Stopping early.")
                break
            attempts += 1
            # Choose a skill (optionally weighted), then an instance within it
            skill = random.choices(skill_names, weights=skill_weights)[0]
            gen_instance = random.choice(skills[skill])
            entry = stats.setdefault(
                instance_label(gen_instance),
                {"emitted": 0, "duplicates_skipped": 0, "errors": 0})
            try:
                example = gen_instance.generate()
                if not example:
                    raise ValueError("generate() returned an empty example")
                example = stamp_fn(example, gen_instance)
                validate_fn(example)
            except Exception as e:
                entry["errors"] += 1
                consecutive_rejects += 1
                if entry["errors"] <= 5:
                    print(f"ERROR: Generator {gen_instance.__class__.__name__} failed "
                          f"during generation or validation: {e}. Skipping attempt {attempts}.")
                elif entry["errors"] == 6:
                    print(f"ERROR: suppressing further errors from "
                          f"{gen_instance.__class__.__name__} (see stats table).")
                continue

            key = (example["operation"], example["problem"])
            if not allow_duplicates:
                if key in seen:
                    entry["duplicates_skipped"] += 1
                    consecutive_rejects += 1
                    continue
                seen.add(key)

            write_jsonl(fp, example)
            entry["emitted"] += 1
            count += 1
            consecutive_rejects = 0
            if count % 1000 == 0:
                print(f"... successfully generated {count}/{n} examples")

    print(f"✔  Successfully wrote {count} lines → {path} (after {attempts} attempts)")
    if stats:
        width = max(len(name) for name in stats)
        totals = {"emitted": 0, "duplicates_skipped": 0, "errors": 0}
        print(f"{'Generator'.ljust(width)}  {'emitted':>8}  {'dup_skip':>8}  {'errors':>6}")
        for name in sorted(stats):
            s = stats[name]
            for k in totals:
                totals[k] += s[k]
            print(f"{name.ljust(width)}  {s['emitted']:>8}  "
                  f"{s['duplicates_skipped']:>8}  {s['errors']:>6}")
        print(f"{'TOTAL'.ljust(width)}  {totals['emitted']:>8}  "
              f"{totals['duplicates_skipped']:>8}  {totals['errors']:>6}")
    if count < n:
        print(f"WARN: Target of {n} examples not reached ({count}/{n}). Consider increasing max_attempts or checking generator logic.")
    return {"count": count, "attempts": attempts, "stats": stats}