Spaces:
Running on Zero
Running on Zero
File size: 3,557 Bytes
66ee87e | 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 | """Arthur end to end in the container: loads every Arthur folder in MODELS_DIR and runs real decisions.
Needs torch (the container has it) and at least one Arthur folder in the models folder; otherwise skipped with the reason.
"""
import importlib.util
import os
import unittest
from app.registry import model_specs
TORCH = importlib.util.find_spec("torch") is not None
ARTHURS = [s for s in model_specs(os.environ) if s["kind"] == "arthur" and s.get("source") == "local"] # Hub copies need the network
QUESTIONS = {
"team": {"type": "choice", "instructions": "Which team should handle this?",
"criteria": {"billing": "Payments and refunds", "tech": "Bugs and errors", "sales": "Pricing and plans"}},
"angry": {"type": "noul", "instructions": "The customer sounds angry"},
"urgency": {"type": "score", "instructions": "How urgent is this?", "criteria": ["Low", "Medium", "High", "Critical"]},
}
STATE = {"channel": "email", "message": "I was charged twice for March. Refund the duplicate today or I cancel."}
@unittest.skipUnless(TORCH and ARTHURS, "needs torch and an Arthur folder in the models folder")
class ArthurRuntimeTest(unittest.TestCase):
@classmethod
def setUpClass(cls):
from app import arthur
cls.arthur = arthur
cls.loaded = [(s, *arthur.load_arthur(s["path"], "cpu")) for s in ARTHURS]
def test_weights_load_strictly_and_answers_are_probability_distributions(self):
for spec, net, temps, _ in self.loaded:
out = self.arthur.decide(net, temps, STATE, QUESTIONS)
with self.subTest(model=spec["name"]):
self.assertEqual(set(out), set(QUESTIONS))
for r in out.values():
self.assertAlmostEqual(sum(r["probs"].values()), 1.0, places=5)
self.assertEqual(set(out["team"]["probs"]), {"billing", "tech", "sales"})
self.assertAlmostEqual(out["angry"]["p_true"], out["angry"]["probs"]["True"], places=6)
self.assertTrue(0.0 <= out["urgency"]["expected_level"] <= 3.0)
def test_decide_matches_the_notebooks_single_decision_path(self):
import torch
from app.arthur_io import normalize_question, probabilities, state_text
for spec, net, temps, _ in self.loaded:
for name, q in QUESTIONS.items():
qtype, text, keys, opts = normalize_question(q)
d = {"question": text, "options": opts, "state": state_text(STATE), "qtype": qtype}
with torch.no_grad():
z = net(*self.arthur.batch_tensors([d])).float().cpu().numpy()[0, :len(opts)]
ref = probabilities(z, qtype, temps)
got = self.arthur.decide(net, temps, STATE, {name: q})[name]["probs"]
with self.subTest(model=spec["name"], question=name):
for k, p in zip(keys, ref):
self.assertAlmostEqual(got[str(k)], float(p), places=5)
def test_batching_questions_together_does_not_change_answers(self):
for spec, net, temps, _ in self.loaded:
together = self.arthur.decide(net, temps, STATE, QUESTIONS)
for name, q in QUESTIONS.items():
alone = self.arthur.decide(net, temps, STATE, {name: q})[name]["probs"]
with self.subTest(model=spec["name"], question=name):
for k, p in alone.items():
self.assertAlmostEqual(together[name]["probs"][k], p, places=4)
if __name__ == "__main__":
unittest.main()
|