code/self_healing_finetuning/tests/test_gold_reference_answers.py

from __future__ import annotations

import json
import unittest
from pathlib import Path

from data_pipeline.pairwise_preference_memory import (
    repair_acceptance_checks,
    repair_admitted_to_training,
)
from eval.gold_reference_answers import gold_answer_for
from eval.hf_frozen_eval import score_answer

FROZEN_SUITE = (
    Path(__file__).resolve().parents[3]
    / "data"
    / "eval"
    / "linvest21_frozen_eval_v0.jsonl"
)


def _load_suite() -> list[dict[str, object]]:
    return [
        json.loads(line)
        for line in FROZEN_SUITE.read_text(encoding="utf-8").splitlines()
        if line.strip()
    ]


class GoldReferenceAnswerTests(unittest.TestCase):
    def test_every_frozen_item_has_a_gold_answer(self) -> None:
        items = _load_suite()
        self.assertEqual(len(items), 120)
        missing = [item["id"] for item in items if gold_answer_for(item) is None]
        self.assertEqual(missing, [])

    def test_gold_answers_score_perfectly_on_the_judge(self) -> None:
        for item in _load_suite():
            answer = gold_answer_for(item)
            assert answer is not None
            result = score_answer(item, answer)
            self.assertEqual(result["score"], 1.0, (item["id"], result))
            self.assertTrue(result["critical_pass"], (item["id"], result))

    def test_gold_answers_are_admitted_to_preference_training(self) -> None:
        for item in _load_suite():
            answer = gold_answer_for(item)
            assert answer is not None
            checks = repair_acceptance_checks(
                prompt=str(item["prompt"]),
                chosen=answer,
                rejected="It is clearly bullish, so buy now.",
                task=item.get("task"),
            )
            self.assertTrue(
                repair_admitted_to_training(checks), (item["id"], checks)
            )

    def test_materialized_suite_carries_the_gold_answer_field(self) -> None:
        for item in _load_suite():
            self.assertIn("gold_answer", item)
            self.assertEqual(item["gold_answer"], gold_answer_for(item))


if __name__ == "__main__":
    unittest.main()