File size: 7,332 Bytes
f8dbde1 | 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 | """``python -m deg.evaluate calibrate``: fits temperatures on val + jev dev and writes a
calibrated copy of the checkpoint without touching the source (tiny random Qwen3.5, CPU)."""
from __future__ import annotations
import hashlib
import json
from pathlib import Path
import pytest
torch = pytest.importorskip("torch")
pytest.importorskip("transformers")
from test_model import CONTEXT, SNAPSHOT, _sub, tiny_lm # noqa: E402
from deg import evaluate # noqa: E402
from deg.model import KINDS, load_model, load_tokenizer # noqa: E402
def _digest(path):
return {p.name: hashlib.sha256(p.read_bytes()).hexdigest() for p in sorted(path.iterdir())}
@pytest.mark.skipif(SNAPSHOT is None, reason="Qwen3.5-0.8B tokenizer not cached")
def test_calibrate_writes_calibrated_copy_and_keeps_source(tmp_path, monkeypatch):
from safetensors.torch import load_file
tok = load_tokenizer(SNAPSHOT)
model = load_model(None, device="cpu", dtype=torch.float32, lm=tiny_lm(), tokenizer=tok)
torch.manual_seed(1)
with torch.no_grad():
for p in model.heads.parameters():
p.normal_(0.0, 0.5)
src = tmp_path / "step-0000010"
model.save(src, {"temperature": dict.fromkeys(KINDS, 1.0)})
before = _digest(src)
a = CONTEXT.index("AB-1234")
rec = {"uid": "v", "split": "val", "schema": _sub(["department", "urgent", "order_id"]),
"context": CONTEXT,
"value": {"department": "billing", "urgent": True, "order_id": "AB-1234"},
"spans": {"/order_id": [a, a + 7]}, "span_alternatives": {}, "span_probs": {},
"decide_probs": {"/department": [0.8, 0.15, 0.05]}}
records = tmp_path / "records.jsonl"
records.write_text("".join(json.dumps({**rec, "uid": f"v{i}"}) + "\n" for i in range(12))
+ json.dumps({**rec, "uid": "t", "split": "train"}) + "\n")
loaded = []
def fake_load(path, device="cuda"):
m = load_model(path, device="cpu", dtype=torch.float32, lm=tiny_lm(), tokenizer=tok)
loaded.append(m)
return m
monkeypatch.setattr(evaluate, "load_model", fake_load)
out = tmp_path / "final"
evaluate.main(["calibrate", "--ckpt", str(src), "--out", str(out), "--records",
str(records), "--jev-dev", str(tmp_path / "missing.jsonl"), "--n-jev", "0"])
assert _digest(src) == before # source untouched
meta = json.loads((out / "deg.json").read_text())
cal = meta["calibration"]
assert cal["source"] == str(src) and cal["val_records"] == 12 and "jev_records" not in cal
assert cal["fit_source"] == "datagen-val" and cal["fit_on"] == ["datagen_val"]
temps = cal["temperature"]
assert temps["choice"] != 1.0 or temps["noul"] != 1.0 or temps["extract"] != 1.0
assert temps["count"] == 1.0 # never fit
rep = cal["report"]["datagen_val"]
assert rep["choice"]["n"] == 12 and rep["choice"]["t_after"] == temps["choice"]
assert rep["decide"]["ece_after"] is not None and "jev_dev" not in cal["report"]
assert meta["temperature"] == pytest.approx({k: temps.get(k, 1.0) for k in KINDS})
# every tensor but the temperatures is byte-identical to the source
backbone = "backbone.safetensors"
assert (out / backbone).read_bytes() == (src / backbone).read_bytes()
h_src, h_out = (load_file(str(p / "heads.safetensors")) for p in (src, out))
assert set(h_src) == set(h_out)
for k in h_src:
if k != "temperature":
assert torch.equal(h_src[k], h_out[k]), k
re = fake_load(out)
assert re.heads.temperature.tolist() == pytest.approx([meta["temperature"][k] for k in KINDS])
assert re.array_mode == loaded[0].array_mode
with pytest.raises(ValueError):
evaluate.write_calibrated(loaded[0], src, src, {})
@pytest.mark.skipif(SNAPSHOT is None, reason="Qwen3.5-0.8B tokenizer not cached")
def test_calibrate_fits_datagen_val_only_and_reports_jev_separately(tmp_path):
tok = load_tokenizer(SNAPSHOT)
def fresh():
m = load_model(None, device="cpu", dtype=torch.float32, lm=tiny_lm(), tokenizer=tok)
torch.manual_seed(1)
with torch.no_grad():
for p in m.heads.parameters():
p.normal_(0.0, 0.5)
return m
a = CONTEXT.index("AB-1234")
rec = {"uid": "v", "split": "val", "schema": _sub(["department", "urgent", "order_id"]),
"context": CONTEXT,
"value": {"department": "billing", "urgent": True, "order_id": "AB-1234"},
"spans": {"/order_id": [a, a + 7]}, "span_alternatives": {}, "span_probs": {},
"decide_probs": {}}
records = tmp_path / "records.jsonl"
records.write_text("".join(json.dumps({**rec, "uid": f"v{i}"}) + "\n" for i in range(12)))
jev = tmp_path / "dev.jsonl"
q_choice = {"type": "choice", "instructions": "Which team?",
"criteria": {"billing": "money", "technical": "bugs", "sales": "buying"},
"label": {"probs": [0.05, 0.9, 0.05]}}
q_noul = {"type": "noul", "instructions": "Is it about shipping?",
"criteria": {"true": "shipping", "false": "other"}, "label": {"hard": 0}}
jev.write_text("".join(json.dumps({"uid": f"j{i}", "state": CONTEXT,
"questions": [q_choice, q_noul]}) + "\n"
for i in range(15)))
info_val: dict = {}
t_val = evaluate.calibrate(fresh(), [records], jev, 15, info=info_val)
info_rep: dict = {}
t_rep = evaluate.calibrate(fresh(), [records], jev, 15, info=info_rep, report_jev=True)
assert t_rep == t_val # reporting jev dev never changes the fit
assert info_val["fit_on"] == ["datagen_val"] and "jev_dev" not in info_val["report"]
assert set(info_rep["report"]) == {"datagen_val", "jev_dev"}
jr = info_rep["report"]["jev_dev"]
assert jr["choice"]["n"] == 15 and jr["noul"]["n"] == 15
assert jr["choice"]["t_after"] == t_val["choice"] # reported at the val-fit temperature
assert info_rep["rows"]["choice"] == 12 # the fit pool: datagen val rows only
info_both: dict = {}
evaluate.calibrate(fresh(), [records], jev, 15, info=info_both,
source="datagen-val+jev-dev")
assert info_both["fit_on"] == ["datagen_val", "jev_dev"]
assert info_both["rows"]["choice"] == 27
with pytest.raises(ValueError):
evaluate.calibrate(fresh(), [records], jev, 15, source="jev")
def test_hard_target_nll_scores_the_filled_value():
from types import SimpleNamespace
logits = torch.tensor([2.0, 0.5, -1.0])
row = (logits, SimpleNamespace(target=[0.6, 0.3, 0.1])) # soft; argmax = filled value 0
logp = torch.log_softmax(logits / 1.5, -1)
assert evaluate._nll("choice", [row], 1.5, hard=True) == pytest.approx(-float(logp[0]))
soft = -float((torch.tensor([0.6, 0.3, 0.1]) * logp).sum())
assert evaluate._nll("choice", [row], 1.5) == pytest.approx(soft)
# a correct, under-confident head: the hard fit sharpens (T < 1), the soft one does not
best_hard = min(evaluate.TEMPS, key=lambda t: evaluate._nll("choice", [row], t, True))
best_soft = min(evaluate.TEMPS, key=lambda t: evaluate._nll("choice", [row], t))
assert best_hard < 1.0 <= best_soft
with pytest.raises(ValueError):
evaluate.calibrate(None, [], Path("x"), 0, target="medium")
|