| """Opt-in checkpoint regression: JIFFY_GPU_TEST=1 python -m unittest discover -s tests.""" |
| import base64 |
| import io |
| import json |
| import os |
| import time |
| import unittest |
| from types import SimpleNamespace |
|
|
| import torch |
| from PIL import Image |
|
|
| from jiffy import DiffusionDecisions, Question |
|
|
|
|
| @unittest.skipUnless(os.environ.get("JIFFY_GPU_TEST") == "1", "opt-in H100 checkpoint test") |
| class SharedDocumentGpuTests(unittest.TestCase): |
| @classmethod |
| def setUpClass(cls): |
| torch.set_num_threads(4) |
| cls.backend = DiffusionDecisions.from_backbone() |
|
|
| def compare(self, state, questions): |
| b = self.backend |
| contracts = [b.compile({"answer": q}) for q in questions] |
| prefix, ending = b.tokenize(state, SimpleNamespace(prompt=""), split_document=True) |
| for contract in contracts: |
| joined = prefix["input_ids"][0].tolist() + b.processor.tokenizer.encode(contract.prompt + ending, add_special_tokens=False) |
| self.assertEqual(joined, b.tokenize(state, contract)["input_ids"][0].tolist()) |
| shared = b.shared_document(state, contracts) |
| started = time.perf_counter() |
| reference = [b.system_one(state, c) for c in contracts] |
| sequential_ms = (time.perf_counter() - started) * 1000 |
| max_delta = 0 |
| for actual, expected in zip(shared["results"], reference): |
| a = actual["answers"]["answer"]["probabilities"] |
| e = expected["answers"]["answer"]["probabilities"] |
| delta = max(abs(a[k] - e[k]) for k in a) |
| max_delta = max(max_delta, delta) |
| self.assertLess(delta, .03, "cache split/batch exceeded BF16 regression tolerance") |
| self.assertEqual(max(a, key=a.get), max(e, key=e.get)) |
| alone = b.shared_document(state, contracts[:1], batch_size=1) |
| p = alone["results"][0]["answers"]["answer"]["probabilities"] |
| mixed = shared["results"][0]["answers"]["answer"]["probabilities"] |
| self.assertLess(max(abs(p[k] - mixed[k]) for k in p), .03) |
| self.assertEqual(shared["diagnostics"]["document_prefills"], 1) |
| print(json.dumps({"case": self._testMethodName, "shared": shared["diagnostics"], |
| "sequential_ms": sequential_ms, "max_probability_delta": max_delta}), flush=True) |
| return shared |
|
|
| def test_short_text_and_batch(self): |
| q = Question.noul("Is a refund requested?") |
| result = self.compare("Please refund my duplicate charge immediately.", |
| [q, Question.noul("Is a repair requested?"), q, q]) |
| self.assertLess(result["diagnostics"]["branch_batches"], 4) |
| self.assertGreater(result["results"][0]["answers"]["answer"]["noul"], .9) |
|
|
| def test_image(self): |
| output = io.BytesIO() |
| Image.new("RGB", (224, 224), "red").save(output, format="PNG") |
| state = {"text": "Identify the image color.", "images": ["data:image/png;base64," + base64.b64encode(output.getvalue()).decode()]} |
| question = Question.choice("What color fills the image?", {"red": "Red", "blue": "Blue"}) |
| result = self.compare(state, [question, Question.noul("Is the image red?")]) |
| self.assertEqual(result["results"][0]["answers"]["answer"]["choice"], "red") |
|
|
| def test_32k_text(self): |
| state = "neutral " * 32000 + "\nPlease refund my duplicate charge immediately." |
| q = Question.noul("Is a refund requested?") |
| result = self.compare(state, [q, q, q, q]) |
| self.assertGreater(result["diagnostics"]["document_tokens"], 31000) |
| self.assertGreater(result["results"][0]["answers"]["answer"]["noul"], .9) |
|
|
| def test_32k_image(self): |
| output = io.BytesIO() |
| Image.new("RGB", (224, 224), "red").save(output, format="PNG") |
| state = {"text": "neutral " * 31500 + "\nIdentify the image color.", |
| "images": ["data:image/png;base64," + base64.b64encode(output.getvalue()).decode()]} |
| q = Question.choice("What color fills the image?", {"red": "Red", "blue": "Blue"}) |
| result = self.compare(state, [q, q]) |
| self.assertGreater(result["diagnostics"]["document_tokens"], 31000) |
| self.assertEqual(result["results"][0]["answers"]["answer"]["choice"], "red") |
|
|
| def test_mixed_protocol_characterization(self): |
| from jiffy import JevProtocol |
| payload = {"model": "jev-latest", "state": {"message": "Please refund my duplicate charge immediately."}, |
| "questions": { |
| "refund": {"type": "noul", "instructions": "Is a refund requested?"}, |
| "team": {"type": "choice", "instructions": "Which team?", |
| "criteria": {"billing": None, "technical": "Software bugs"}}, |
| "urgency": {"type": "score", "instructions": "How urgent is this?", |
| "criteria": [{"description": "No urgency expressed"}, ["Explicit urgency"]]}}} |
| actual = JevProtocol(self.backend).evaluate(payload) |
| reference = JevProtocol(self.backend, execution="sequential").evaluate(payload) |
| self.assertLess(actual["usage"]["input_tokens"], reference["usage"]["input_tokens"]) |
| for name, a in actual["answers"].items(): |
| e = reference["answers"][name] |
| if a["type"] == "noul": |
| self.assertAlmostEqual(a["noul"], e["noul"], delta=.03) |
| else: |
| delta = max(abs(a["probabilities"][k] - e["probabilities"][k]) for k in a["probabilities"]) |
| |
| |
| self.assertLess(delta, .05 if a["type"] == "score" else .03) |
| if a["type"] == "score": |
| print(json.dumps({"case": self._testMethodName, "score_probability_delta": delta, |
| "strict_0_03_gate_passed": delta < .03}), flush=True) |
| self.assertEqual(actual["answers"]["team"]["choice"], "billing") |
| self.assertEqual(actual["answers"]["urgency"]["legend"], reference["answers"]["urgency"]["legend"]) |
| payload["questions"] = {"renamed": payload["questions"]["urgency"]} |
| alone = JevProtocol(self.backend).evaluate(payload) |
| self.assertEqual(alone["answers"]["renamed"], actual["answers"]["urgency"]) |
|
|
