tests/test_posttrain.py

"""Unit tests for train/ties_merge.py and research/rlvr.py.

Run: .venv/bin/python tests/test_posttrain.py
"""
import sys
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

import torch

from train.ties_merge import ties_merge, trim_delta
from research.rlvr import reward, reward_card


def test_trim_keeps_topk():
    d = torch.tensor([1.0, -0.5, 0.01, 0.001, 0.0005])
    t = trim_delta(d, keep=0.4)
    assert t[0] == 1.0 and t[1] == -0.5
    assert t[2] == 0.0 and t[3] == 0.0


def test_ties_sign_consensus():
    base = {"w": torch.zeros(4)}
    t1 = {"w": torch.tensor([1.0, 1.0, 1.0, -1.0])}
    t2 = {"w": torch.tensor([1.0, -1.0, 1.0, -1.0])}
    out = ties_merge(base, [t1, t2], keep=1.0)
    # sign agreement at idx 0, 2, 3 -> merge; idx 1 disagrees -> zero
    assert out["w"][0] == 1.0
    assert out["w"][2] == 1.0
    assert out["w"][3] == -1.0
    assert out["w"][1] == 0.0


def test_rlvr_reward_correct_with_citation():
    r = reward(gold="refutes", policy="refutes", citation="1982",
               evidence="the deed file states 1982")
    assert r["verdict"] == 1.0 and r["citation"] == 0.2 and r["total"] == 1.2


def test_rlvr_reward_wrong_verdict():
    r = reward(gold="refutes", policy="supports", citation="1982",
               evidence="the deed file states 1982")
    assert r["verdict"] == -1.0 and r["total"] == -0.8


def test_rlvr_reward_abstain_is_zero():
    r = reward(gold="refutes", policy="not enough information",
               citation="", evidence="the deed file states 1982")
    assert r["verdict"] == 0.0 and r["total"] == 0.0


def test_rlvr_reward_false_citation_penalty():
    r = reward(gold="supports", policy="supports", citation="1978",
               evidence="the deed file states 1982")
    assert r["verdict"] == 1.0 and r["citation"] == -0.2 and r["total"] == 0.8


def test_rlvr_card_trace():
    card = reward_card(gold="refutes", policy="refutes", citation="1982",
                       evidence="the deed file states 1982", probe="rt05")
    assert card["probe"] == "rt05" and card["total"] == 1.2
    for k in ("gold", "policy", "verdict", "citation", "total"):
        assert k in card


if __name__ == "__main__":
    fns = [v for k, v in sorted(globals().items()) if k.startswith("test_")]
    for fn in fns:
        fn()
        print(f"ok {fn.__name__}")
    print(f"\n{len(fns)} posttrain tests passed")