| import unittest |
|
|
| import torch |
|
|
| from jiffy.diffusion_decisions import DiffusionDecisions, SlotConstraints, SlotStopping, answer_scores, compile_contract |
| from jiffy.schema import Question |
|
|
|
|
| class Tokenizer: |
| pad_token_id = 0 |
|
|
| def encode(self, text, add_special_tokens=False): |
| return [ord(c) for c in text] |
|
|
|
|
| class DiffusionDecisionTests(unittest.TestCase): |
| def test_typed_contract_has_distinct_slots(self): |
| questions = {"urgent": Question.noul("Urgent?"), "team": Question.choice("Team?", {"x": "Billing", "y": "Support"}), |
| "severity": Question.score("Severity?", ["Low", "Medium", "High"])} |
| contract = compile_contract(Tokenizer(), questions) |
| self.assertEqual(contract.names, tuple(questions)) |
| self.assertEqual(len(set(contract.positions)), 3) |
| self.assertEqual(len(contract.scaffold), 256) |
| self.assertIn("A: The answer is no.", contract.prompt) |
| self.assertIn("C: High", contract.prompt) |
|
|
| def test_bad_contracts_are_rejected(self): |
| q = Question.noul("Test?") |
| for questions in ({}, {"": q}, {"q": "not typed"}, {str(i): q for i in range(17)}): |
| with self.assertRaises(ValueError): |
| compile_contract(Tokenizer(), questions) |
| with self.assertRaises(ValueError): |
| compile_contract(Tokenizer(), {"q": q}, orders={"q": [0, 0]}) |
| with self.assertRaises(ValueError): |
| compile_contract(Tokenizer(), {"q": q}, canvas_length=2) |
|
|
| def test_constraints_keep_all_candidate_probabilities(self): |
| contract = compile_contract(Tokenizer(), {"q": Question.noul("Test?")}, canvas_length=64) |
| scores = torch.zeros(1, 64, 128) |
| position = contract.positions[0] |
| scores[0, position, ord("B")] = torch.log(torch.tensor(3.)) |
| constrained = SlotConstraints(contract) |
| output = constrained(torch.zeros(1, 1), scores, 4) |
| self.assertEqual(constrained.trace[0]["step"], 1) |
| torch.testing.assert_close(torch.tensor(constrained.trace[0]["answers"][0]["probabilities"]), torch.tensor([.25, .75])) |
| self.assertLess(constrained.trace[0]["answers"][0]["allowed_mass"], .04) |
| self.assertEqual(torch.isfinite(output[0, position]).sum().item(), 2) |
| self.assertEqual(torch.isfinite(output[0, 0]).sum().item(), 1) |
| self.assertTrue(torch.isfinite(scores).all()) |
|
|
| def test_option_permutation_restores_original_semantics(self): |
| contract = compile_contract(Tokenizer(), {"q": Question.noul("Test?")}, canvas_length=64, orders={"q": [1, 0]}) |
| scores = torch.zeros(1, 64, 128) |
| scores[0, contract.positions[0], ord("B")] = torch.log(torch.tensor(3.)) |
| constraint = SlotConstraints(contract) |
| constraint(torch.zeros(1, 1), scores, 1) |
| torch.testing.assert_close(torch.tensor(constraint.trace[-1]["answers"][0]["probabilities"]), torch.tensor([.75, .25])) |
|
|
| def test_selected_readout_matches_full_canvas_extraction(self): |
| contract = compile_contract(Tokenizer(), {"q": Question.noul("Test?")}, canvas_length=64) |
| scores = torch.randn(1, 64, 128) |
| constraint = SlotConstraints(contract) |
| constraint(torch.zeros(1, 1), scores, 1) |
| self.assertEqual(answer_scores(scores[:, list(contract.positions)], contract), constraint.trace[-1]["answers"]) |
|
|
| def test_prefix_is_bound_to_backend_and_contract(self): |
| backend = object.__new__(DiffusionDecisions) |
| contract = compile_contract(Tokenizer(), {"q": Question.noul("Test?")}) |
| with self.assertRaisesRegex(ValueError, "foreign"): |
| backend.score({"owner": object(), "contract": contract}, contract) |
| with self.assertRaisesRegex(ValueError, "mismatched"): |
| backend.score({"owner": backend, "contract": None}, contract) |
| with self.assertRaisesRegex(ValueError, "positive"): |
| backend.score({}, contract, steps=0) |
|
|
| def test_padding_cannot_fake_early_confidence(self): |
| stopping = SlotStopping([7]) |
| canvas = torch.zeros(1, 256, dtype=torch.long) |
| scores = torch.full((1, 256, 2), -torch.inf) |
| scores[:, :, 0] = 0 |
| scores[:, 7, 1] = 0 |
| self.assertFalse(stopping(canvas, scores).item()) |
| self.assertFalse(stopping(canvas, scores).item()) |
| scores[:, 7, 1] = -100 |
| self.assertTrue(stopping(canvas, scores).item()) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
