tests/test_datagen.py

"""Smoke #5: PRM800K parse + dummy-teacher datagen (offline, no key)."""
from pathlib import Path

from mathcompose.datagen.prm800k_loader import iter_prm800k
from mathcompose.datagen.gen_verifier_data import (
    generate_verifier_dataset,
    detect_verifier_completion,
)
from mathcompose.datagen.gen_generator_data import Problem, generate_generator_dataset
from mathcompose.teachers import get_teacher
from mathcompose.eval.parse import first_error_index

FIX = Path(__file__).parent / "fixtures" / "prm800k_sample.jsonl"


def test_prm800k_parse():
    recs = list(iter_prm800k(FIX))
    assert len(recs) == 4
    by_problem = {r.problem: r for r in recs}
    assert by_problem["What is 2 + 2?"].first_error_index == 1
    assert by_problem["What is 3 times 3?"].first_error_index == -1
    assert by_problem["Find the derivative of x^2."].first_error_index == 1


def test_verifier_datagen_labels_authoritative_and_no_leak():
    recs = list(iter_prm800k(FIX))
    teacher = get_teacher("dummy")
    rows = list(generate_verifier_dataset(recs, teacher, banned=set()))
    assert len(rows) == 4
    for row in rows:
        completion = row["completion"][0]["content"]
        # authoritative boxed label matches the PRM800K gold
        assert first_error_index(completion) == row["first_error_index"]
        # the stored training prompt must NOT contain the synthesis hint
        user = row["prompt"][1]["content"]
        assert "GROUND TRUTH" not in user
        assert "FIRST error occurs in paragraph" not in user


def test_detect_mode_keeps_only_matching_predictions():
    """Genuine-detection recipe: keep a critique only if the teacher's blind
    prediction matches the PRM800K gold."""
    recs = list(iter_prm800k(FIX))
    err = next(r for r in recs if r.first_error_index == 1)  # "What is 2 + 2?" -> gold 1

    # teacher predicts the correct index -> kept
    good = get_teacher("dummy", fn=lambda m: "Paragraph 1: wrong. Verdict: incorrect\n\n\\boxed{1}")
    assert detect_verifier_completion(good, err, max_attempts=1) is not None

    # teacher predicts the wrong index -> rejected (None)
    bad = get_teacher("dummy", fn=lambda m: "All good.\n\n\\boxed{-1}")
    assert detect_verifier_completion(bad, err, max_attempts=2) is None


def test_generator_datagen_answer_filter():
    def scripted(msgs):
        q = msgs[-1]["content"]
        return r"work... \boxed{4}" if "2+2" in q else r"work... \boxed{999}"
    teacher = get_teacher("dummy", fn=scripted)
    probs = [Problem("2+2=?", "4"), Problem("3+3=?", "6")]
    rows = list(generate_generator_dataset(probs, teacher, keep_only_correct=True))
    assert len(rows) == 1  # only the answer-correct trace survives