File size: 6,604 Bytes
cbc33fe | 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 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 | """
Synthetic tests for E3/E4 preference-pair construction. Catches logic errors now
rather than four hours into the pipeline.
Most important: test_divpo_prob_picks_least_probable_as_chosen. DivPO's
probability variant inverts the usual intuition -- the CHOSEN story is the one
the model found LEAST likely. Getting that backwards would train the policy
straight into the greedy mode while the logs looked entirely normal.
"""
import numpy as np
from build_pairs import build_divpo, build_multipos, normalize_weights
def make_item(pid, idx, quality, dev, meanlp, passed=True, text=None):
return {"prompt_id": pid, "idx": idx, "prompt": f"prompt-{pid}",
"text": text or f"story-{pid}-{idx}", "quality": quality,
"deviation": dev, "mean_logprob": meanlp, "gate_passed": passed,
"marginal": -1.0, "n_words": 300}
def simple_pool():
"""One prompt, 6 stories spanning quality/deviation/logprob."""
return {"p0": [
make_item("p0", 0, 8.0, 0.40, -1.50), # high q, MOST diverse, least probable
make_item("p0", 1, 7.0, 0.20, -0.80), # high q, low dev -> "competent cliche"
make_item("p0", 2, 6.5, 0.30, -1.10),
make_item("p0", 3, 3.0, 0.35, -0.60), # low q, HIGH dev -> must NOT be chosen
make_item("p0", 4, 2.0, 0.05, -0.40), # low q, least diverse, MOST probable
make_item("p0", 5, 1.0, 0.10, -0.90, passed=False),
]}
# --------------------------------------------------------------------- E4
def test_divpo_emb_chosen_is_most_diverse_above_rho():
rows, st = build_divpo(simple_pool(), "emb", rho=6.0)
assert st["n_rows"] == 1
r = rows[0]
assert r["chosen"] == "story-p0-0", r
assert r["chosen_quality"] >= 6.0
def test_divpo_emb_never_chooses_a_diverse_but_low_quality_story():
"""Story 3 has high deviation but quality 3.0. Quality gating must exclude
it -- this is the 'different because it is worse' failure."""
rows, _ = build_divpo(simple_pool(), "emb", rho=6.0)
assert rows[0]["chosen"] != "story-p0-3"
def test_divpo_emb_rejected_is_least_diverse_below_rho():
rows, _ = build_divpo(simple_pool(), "emb", rho=6.0)
assert rows[0]["rejected"] == "story-p0-4", rows[0]
def test_divpo_prob_picks_least_probable_as_chosen():
"""Most diverse == LOWEST length-normalized logprob; least diverse ==
HIGHEST (the near-greedy sample). Inverting this trains toward the mode."""
rows, _ = build_divpo(simple_pool(), "prob", rho=6.0)
r = rows[0]
assert r["chosen"] == "story-p0-0", f"chosen should be least probable: {r}"
assert r["rejected"] == "story-p0-4", f"rejected should be most probable: {r}"
assert r["chosen_meanlp"] < r["rejected_meanlp"]
def test_divpo_skips_prompt_with_no_qualifying_chosen():
pool = {"p0": [make_item("p0", i, 3.0, 0.2, -1.0) for i in range(4)]}
rows, st = build_divpo(pool, "emb", rho=6.0)
assert rows == [] and st["no_chosen"] == 1 and st["skip_rate"] == 1.0
def test_divpo_skips_prompt_with_no_rejected():
pool = {"p0": [make_item("p0", i, 9.0, 0.2 + 0.01 * i, -1.0) for i in range(4)]}
rows, st = build_divpo(pool, "emb", rho=6.0)
assert rows == [] and st["no_rejected"] == 1
def test_divpo_gate_failures_are_eligible_rejects():
"""A gate-failing story is legitimately in the low-quality tail."""
pool = {"p0": [make_item("p0", 0, 8.0, 0.4, -1.5),
make_item("p0", 1, 9.0, 0.3, -1.2),
make_item("p0", 2, 0.0, 0.01, -0.3, passed=False)]}
rows, st = build_divpo(pool, "emb", rho=6.0)
assert st["n_rows"] == 1 and rows[0]["rejected"] == "story-p0-2"
def test_divpo_is_one_row_per_prompt():
pool = {f"p{i}": simple_pool()["p0"] for i in range(5)}
rows, st = build_divpo(pool, "emb", rho=6.0)
assert len(rows) == 5 == st["n_rows"]
assert all(r["weight"] == 1.0 for r in rows), "DivPO must be unweighted"
# --------------------------------------------------------------------- E3
def _emb_for(pool):
"""Deterministic embeddings + row index aligned to pool iteration order."""
idx, vecs, i = {}, [], 0
rng = np.random.default_rng(0)
for pid, items in pool.items():
for it in items:
idx[(pid, it["idx"])] = i
v = rng.standard_normal(32)
vecs.append(v / np.linalg.norm(v))
i += 1
return np.array(vecs, dtype=np.float32), idx
def test_multipos_emits_multiple_chosen_rows_per_prompt():
pool = simple_pool()
emb, idx = _emb_for(pool)
rows, st = build_multipos(pool, emb, idx, q_keep=5.0, k=4, lam=1.0)
assert st["ok"] == 1
assert 2 <= len(rows) <= 4, f"expected up to 4 chosen rows, got {len(rows)}"
assert all(r["chosen_quality"] >= 5.0 for r in rows)
def test_multipos_rotates_negatives():
"""Don't hammer one rejected across all rows."""
pool = simple_pool()
emb, idx = _emb_for(pool)
rows, _ = build_multipos(pool, emb, idx, q_keep=5.0, k=4, lam=1.0)
if len(rows) >= 2:
assert len(set(r["neg_type"] for r in rows)) == 2, \
f"negatives not rotated: {[r['neg_type'] for r in rows]}"
def test_multipos_weight_is_chosen_deviation():
pool = simple_pool()
emb, idx = _emb_for(pool)
rows, _ = build_multipos(pool, emb, idx, q_keep=5.0, k=4, lam=1.0)
for r in rows:
assert abs(r["weight"] - r["chosen_dev"]) < 1e-9
def test_multipos_skips_prompt_with_too_few_survivors():
pool = {"p0": [make_item("p0", i, 2.0, 0.2, -1.0) for i in range(6)]}
emb, idx = _emb_for(pool)
rows, st = build_multipos(pool, emb, idx, q_keep=5.0, k=4, lam=1.0)
assert rows == [] and st["skipped_no_survivors"] == 1
def test_normalize_weights_gives_mean_one():
rows = [{"weight": w} for w in (0.1, 0.2, 0.3, 0.8)]
normalize_weights(rows)
w = np.array([r["weight"] for r in rows])
assert abs(w.mean() - 1.0) < 1e-9
# relative ordering preserved -> it reweights, it does not rescale the LR
assert np.all(np.diff(w) > 0)
def test_normalize_weights_handles_all_zero():
rows = [{"weight": 0.0} for _ in range(3)]
normalize_weights(rows)
assert all(r["weight"] == 0.0 for r in rows)
if __name__ == "__main__":
import sys, traceback
fns = [(n, f) for n, f in sorted(globals().items())
if n.startswith("test_") and callable(f)]
bad = 0
for n, f in fns:
try:
f(); print(f" PASS {n}")
except Exception:
bad += 1; print(f" FAIL {n}"); traceback.print_exc()
print(f"\n{len(fns)-bad}/{len(fns)} passed")
sys.exit(1 if bad else 0)
|