File size: 1,847 Bytes
17a55cc | 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 | import numpy as np
import pytest
import torch
from pathlib import Path
FIXTURE = Path(__file__).parent / "fixtures" / "parity"
CKPT = FIXTURE / "ckpt"
pytestmark = pytest.mark.skipif(not CKPT.exists(), reason="parity fixture missing")
@pytest.fixture(scope="module")
def model_and_batch():
import sys
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "scripts"))
from make_parity_fixture import PAD, build_request
from groot1.model import GrootDecisionModel, collate
from groot1.serialize import EncodedQuestion, EncodedState
model = GrootDecisionModel.from_pretrained(CKPT)
chunks, eqs = build_request(np.random.default_rng(0), 1024)
state = EncodedState(chunks=chunks, tokens=0)
batch = collate([state], eqs, PAD, q_state=[0] * len(eqs))
return model, batch, [q.n_options for q in eqs]
def test_forward_shapes_and_probabilities(model_and_batch):
model, batch, ks = model_and_batch
with torch.inference_mode():
out = model(batch)
assert out["probabilities"].shape[0] == len(ks)
assert out["confidence_logit"].shape == (len(ks),)
for qi, k in enumerate(ks):
p = out["probabilities"][qi, :k]
assert p.sum().item() == pytest.approx(1.0, abs=1e-5)
assert torch.isfinite(out["confidence_logit"][qi])
def test_outputs_match_reference_npy(model_and_batch):
model, batch, ks = model_and_batch
ref = np.load(FIXTURE / "ref.npy")
nopt, nq = sum(ks), len(ks)
with torch.inference_mode():
out = model(batch)
probs = np.concatenate([out["probabilities"][q, :k].numpy() for q, k in enumerate(ks)])
conf = out["confidence_logit"].numpy()
logits = np.concatenate([out["logits"][q, :k].numpy() for q, k in enumerate(ks)])
got = np.concatenate([probs, conf, logits])
assert np.abs(got - ref).max() < 1e-5
|