File size: 4,981 Bytes
00ee773 | 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 | """Core invariants: encoder parity with transformers, head shapes, option-order invariance, determinism.
python -m pytest tests -q (or: python tests\\test_core.py)
"""
import json
import os
import sys
import numpy as np
import torch
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from nirnay.model import DecisionHead # noqa: E402
from nirnay.render import render_question, render_state # noqa: E402
def test_encoder_matches_transformers():
from transformers import AutoModel, AutoTokenizer
from nirnay.encoder import load_hf_encoder
name = os.environ.get("NIRNAY_BACKBONE", "intfloat/multilingual-e5-small")
enc = load_hf_encoder(name, "mean", "query: ")
tok = AutoTokenizer.from_pretrained(name)
ref = AutoModel.from_pretrained(name).eval()
b = tok(["query: hello world", "query: a much longer sentence about billing and refunds"], padding=True,
return_tensors="pt")
with torch.no_grad():
h_ref = ref(**b).last_hidden_state
h = enc(b["input_ids"], b["attention_mask"].bool())
m = b["attention_mask"].bool()
assert torch.allclose(h[m], h_ref[m], atol=2e-4), float((h[m] - h_ref[m]).abs().max())
def _rand_inputs(B=3, n=7, Lq=5, K=4, Lo=3, d=32, seed=0):
g = torch.Generator().manual_seed(seed)
S = torch.randn(1, n, d, generator=g)
sm = torch.ones(1, n, dtype=torch.bool)
Q = torch.randn(B, Lq, d, generator=g)
qm = torch.ones(B, Lq, dtype=torch.bool)
O = torch.randn(B, K, Lo, d, generator=g)
om = torch.ones(B, K, Lo, dtype=torch.bool)
om[:, :, -1] = False
ok = torch.ones(B, K, dtype=torch.bool)
qt = torch.tensor([0, 1, 2][:B])
sv = torch.nn.functional.normalize(torch.randn(1, d, generator=g), dim=-1)
ov = torch.nn.functional.normalize(torch.randn(B, K, d, generator=g), dim=-1)
return S, sm, Q, qm, O, om, ok, qt, sv, ov
def _head(d=32):
torch.manual_seed(0)
h = DecisionHead(d, 4, 64, 2)
with torch.no_grad(): # make the zero-init branches live so the test is meaningful
for p in h.parameters():
p.add_(torch.randn_like(p) * 0.05)
return h.eval()
def test_head_shapes_and_option_independence():
h = _head()
x = _rand_inputs()
with torch.no_grad():
z = h(*x)
assert z.shape == (3, 4) and torch.isfinite(z).all()
# permuting options permutes logits and nothing else
perm = torch.tensor([2, 0, 3, 1])
S, sm, Q, qm, O, om, ok, qt, sv, ov = x
with torch.no_grad():
zp = h(S, sm, Q, qm, O[:, perm], om[:, perm], ok[:, perm], qt, sv, ov[:, perm])
assert torch.allclose(zp, z[:, perm], atol=1e-5)
# adding an extra option leaves the existing options' logits unchanged
O2 = torch.cat([O, torch.randn(3, 1, O.shape[2], O.shape[3])], 1)
om2 = torch.cat([om, torch.ones(3, 1, O.shape[2], dtype=torch.bool)], 1)
ok2 = torch.cat([ok, torch.ones(3, 1, dtype=torch.bool)], 1)
ov2 = torch.cat([ov, torch.nn.functional.normalize(torch.randn(3, 1, ov.shape[2]), dim=-1)], 1)
with torch.no_grad():
z2 = h(S, sm, Q, qm, O2, om2, ok2, qt, sv, ov2)
assert torch.allclose(z2[:, :4], z, atol=1e-5)
def test_render():
assert render_state({"ticket": {"subject": "x", "messages": [{"text": "hi"}]}}) == \
"ticket.subject: x\nticket.messages[0].text: hi"
nq, qt, ot = render_question("t", {"type": "choice", "instructions": "Which team?",
"criteria": {"billing": "payments", "sales_ops": None}})
assert qt == "choice question: Which team?" and ot == ["billing: payments", "sales ops"]
nq, qt, ot = render_question("u", {"type": "noul", "instructions": "Is it urgent?"})
assert nq["keys"] == ["false", "true"] and ot[1].startswith("yes")
nq, qt, ot = render_question("s", {"type": "score", "instructions": "How bad?", "criteria": ["low", "mid", "high"]})
assert ot == ["level 0 of 2: low", "level 1 of 2: mid", "level 2 of 2: high"]
def test_router_language_detection():
from nirnay.router import analyse
assert analyse("My card was charged twice and I want a refund")["english"]
assert analyse({"ticket": {"subject": "Refund", "body": "Please send the invoice for the order"}} and
"ticket.subject: Refund\nticket.body: Please send the invoice for the order")["english"]
assert not analyse("എന്റെ കാർഡിൽ രണ്ട് തവണ പണം ഈടാക്കി")["english"] # Malayalam
assert not analyse("मेरा भुगतान तीन दिनों से विफल हो रहा है")["english"] # Hindi
assert not analyse("Meine Zahlung ist seit drei Tagen fehlgeschlagen, bitte helfen Sie")["english"] # German
if __name__ == "__main__":
test_router_language_detection()
test_render()
test_head_shapes_and_option_independence()
test_encoder_matches_transformers()
print("ok")
|