File size: 2,413 Bytes
97c39f2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 | """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")
|