import json from concurrent.futures import ThreadPoolExecutor import pytest import torch from transformers import BatchEncoding, Qwen2Config, Qwen2ForCausalLM from core import engine from core.schema import StructuredSchema class Tokenizer: pad_token_id = 0 eos_token_id = 1 def encode(self, text, **kwargs): return [ord(c) for c in text] def decode(self, ids, **kwargs): return "".join(chr(i) for i in ids) def __call__(self, text, **kwargs): ids = torch.tensor([self.encode(text)]) return BatchEncoding({"input_ids": ids, "attention_mask": torch.ones_like(ids)}) @pytest.fixture def tiny_engine(monkeypatch): torch.manual_seed(7) model = Qwen2ForCausalLM(Qwen2Config(vocab_size=128, hidden_size=32, intermediate_size=32, num_hidden_layers=1, num_attention_heads=2, num_key_value_heads=1, max_position_embeddings=4096, eos_token_id=1, pad_token_id=0)).eval() tokenizer = Tokenizer() monkeypatch.setattr(engine, "get_engine", lambda: (model, tokenizer)) return model, tokenizer @pytest.mark.parametrize("spec", [{}, {"x": {"choices": []}}, {"x": {"choices": ["a", "a"]}}, {"x": {"choices": [1]}}, {"x": {"choices": [""]}}]) def test_invalid_schema(spec): with pytest.raises(ValueError): StructuredSchema(spec) def test_cached_scores_match_full_forward(tiny_engine): model, tokenizer = tiny_engine prefix = tokenizer("Context", return_tensors="pt") candidates = [tokenizer.encode(s) for s in ['a', 'ab"', 'abc"', 'z"'] * 5] with torch.inference_mode(): output = model(**prefix, use_cache=True) scores, passes = engine._candidate_scores(model, output.past_key_values, output.logits[0, -1], candidates, 0) for i, tokens in enumerate(candidates): ids = torch.tensor([prefix.input_ids[0].tolist() + tokens]) logits = model(ids, use_cache=False).logits[0].float().log_softmax(-1) expected = sum(logits[prefix.input_ids.shape[1] - 1 + j, t] for j, t in enumerate(tokens)) torch.testing.assert_close(scores[i], expected, atol=1e-5, rtol=1e-5) assert passes == 2 def test_batched_padding_and_collision(tiny_engine): model, tokenizer = tiny_engine schema = StructuredSchema({"x": {"type": "boolean"}, "long_field_name": {"type": "enum", "choices": ["ab", "abc", "ad", "z"]}}) result = engine.run_parallel_generation("example", schema) assert result["schema_match"] assert isinstance(result["parsed_json"]["x"]["value"], bool) assert result["parsed_json"]["long_field_name"]["value"] in schema["long_field_name"].choices assert result["sequential_forward_passes"] == 3 # Compare every probability against an independent, unpadded full forward. meta = schema.compile_parallel_metadata(tokenizer) prefix = (f"<|im_start|>system\nClassify JSON attributes:\n{schema.to_parallel_schema_str()}" "<|im_end|>\n<|im_start|>user\nexample<|im_end|>\n" "<|im_start|>assistant\n{\n") with torch.inference_mode(): for i, (name, field) in enumerate(schema.fields.items()): ids = tokenizer.encode(prefix) + meta["suffixes_batch"][i, :meta["suffix_lengths"][i]].tolist() scores = [] for tokens in meta["candidate_sequences"][i]: out = model(torch.tensor([ids + tokens]), use_cache=False).logits[0].float() if meta["has_collisions"][i]: lp = out.log_softmax(-1) scores.append(sum(lp[len(ids) - 1 + j, t] for j, t in enumerate(tokens))) else: scores.append(out[len(ids) - 1, tokens[0]]) expected = torch.stack(scores).softmax(-1).tolist() actual = {x["choice"]: x["probability"] for x in result["field_telemetry"][name]["top_choices"]} for c, p in zip(field.choices, expected): assert actual[c] == pytest.approx(p, abs=0.0001) def test_stream_resumes_on_different_threads_and_releases_lock(tiny_engine): schema = StructuredSchema({"x": {"type": "boolean"}}) iterator = engine.stream_naive_generation("test", schema, max_tokens=3, temperature=0) with ThreadPoolExecutor(1) as first, ThreadPoolExecutor(1) as second: a = first.submit(next, iterator).result(timeout=10) b = second.submit(next, iterator).result(timeout=10) second.submit(iterator.close).result(timeout=10) assert a["type"] == b["type"] == "token" assert engine._gpu_lock.acquire(timeout=1) engine._gpu_lock.release() def test_json_validation(): schema = StructuredSchema({"x": {"type": "boolean"}}) for text in ['[]', '{"x":"true"}', '{"x":true,"extra":1}', '{}']: assert not engine._naive_result(text, schema, 1, 1, 1)["schema_match"] assert engine._naive_result('{"x":true}', schema, 1, 1, 1)["schema_match"] def test_server_routes(tiny_engine, monkeypatch): from fastapi.testclient import TestClient from server import app as server monkeypatch.setattr(server, "get_engine", lambda: tiny_engine) payload = {"context": "test", "schema": {"x": {"type": "boolean"}}, "max_tokens": 2, "temperature": 0} with TestClient(server.app) as client: assert client.get("/").status_code == 200 assert len(client.get("/api/presets").json()) == 4 for route in ["run-parallel", "run-rlcd", "run-tree", "run-naive", "compare"]: response = client.post("/api/" + route, json=payload) assert response.status_code == 200, response.text if route == "compare": data = response.json() assert {"tree", "parallel", "naive"} <= data.keys() assert data["tree"]["sequential_forward_passes"] == 2 assert data["tree_speedup_multiplier"] == round(data["naive"]["elapsed_ms"] / max(data["tree"]["elapsed_ms"], 1), 1) response = client.post("/api/stream-naive", json=payload) events = [json.loads(line[6:]) for line in response.text.splitlines() if line.startswith("data: ")] assert events[-1]["type"] == "done" assert all(e["type"] != "error" for e in events) assert client.post("/api/run-parallel", json={**payload, "schema": {}}).status_code == 400 assert client.post("/api/run-tree", json={**payload, "schema": {}}).status_code == 400 assert client.post("/api/run-tree", json={**payload, "temperature": -1}).status_code == 422 assert client.post("/api/run-naive", json={**payload, "temperature": -1}).status_code == 422 def test_tree_route_temperature_and_errors(tiny_engine, monkeypatch): from fastapi.testclient import TestClient from server import app as server monkeypatch.setattr(server, "get_engine", lambda: tiny_engine) temperatures = [] def tree(context, schema, temperature): temperatures.append(temperature) if context == "fail": raise ValueError("Tree unavailable") return {"mode": "parallel_constrained_tree"} monkeypatch.setattr(server, "run_parallel_generation_tree", tree) payload = {"context": "test", "schema": {"x": {"type": "boolean"}}} with TestClient(server.app) as client: assert client.post("/api/run-tree", json=payload).status_code == 200 assert client.post("/api/run-tree", json={**payload, "temperature": 0}).status_code == 200 response = client.post("/api/run-tree", json={**payload, "context": "fail"}) assert response.status_code == 400 assert response.json()["detail"] == "Tree unavailable" assert temperatures == [1.0, 0, 1.0] @pytest.mark.parametrize("backend", ["sdpa", "eager"]) @pytest.mark.parametrize("temperature", [0, 0.4, 1.0]) @pytest.mark.parametrize("device", ["cpu", pytest.param("cuda", marks=pytest.mark.skipif( not torch.cuda.is_available(), reason="CUDA is unavailable"))]) def test_tree_scores_match_independent_forwards(tiny_engine, monkeypatch, backend, temperature, device): model, tokenizer = tiny_engine model.to(device) model.set_attn_implementation(backend) schema = StructuredSchema({ "payment_status": {"choices": ["PAYMENT", "PAYMENT_FAILED", "PAYMENT_PENDING"]}, "payment_state": {"choices": ["a", "abc", "abd", "z"]}, "payment_escalate": {"type": "boolean"}, "unique": {"choices": ["x", "yyy"]}, 'quoted"key': {"choices": ['a"b', "a\\b", "a\nb"]}, "only": {"choices": ["one"]}, }) forward = model.forward calls = [] def record_forward(*args, **kwargs): if "past_key_values" in kwargs: assert model.config._attn_implementation == "flex_attention" output = forward(*args, **kwargs) calls.append((kwargs, output)) return output monkeypatch.setattr(model, "forward", record_forward) result = engine.run_parallel_generation_tree("example", schema, temperature) assert model.config._attn_implementation == backend assert len(calls) == result["sequential_forward_passes"] == 2 assert result["total_tokens_generated"] == 0 assert result["scoring"] == "mean_token_log_probability" assert result["schema_match"] and result["is_valid_json"] assert isinstance(result["parsed_json"]["payment_escalate"]["value"], bool) assert result["parsed_json"]["only"] == {"value": "one", "prob": 1.0} prefix = calls[0][0]["input_ids"][0].tolist() tree_args, tree_output = calls[1] assert tree_args["input_ids"].shape[0] == 1 assert tree_args["past_key_values"] is calls[0][1].past_key_values cache = tree_output.past_key_values assert cache.get_seq_length() == len(prefix) + result["tree_nodes"] for layer in cache.layers: assert layer.keys.shape[0] == layer.values.shape[0] == 1 # Recover each scoring context independently from the mask, then compare # full-vocabulary logits and complete candidate probabilities with ordinary # causal forwards. This catches sibling leakage and incorrect RoPE depths. tree_ids = tree_args["input_ids"][0].tolist() from torch.nn.attention.flex_attention import BlockMask block_mask = tree_args["attention_mask"] assert isinstance(block_mask, BlockMask) mask = block_mask.mask_mod(torch.tensor(0, device=device), torch.tensor(0, device=device), torch.arange(len(tree_ids), device=device)[:, None], torch.arange(len(prefix) + len(tree_ids), device=device)[None, :]) assert mask[:, :len(prefix)].all() context_rows = {} for row, node in enumerate(tree_args["logits_to_keep"].tolist()): visible = torch.where(mask[node, len(prefix):])[0].tolist() path = tuple(tree_ids[j] for j in visible) context_rows[path] = row assert tree_args["position_ids"][0, node] == len(prefix) + len(path) - 1 with torch.inference_mode(): for i, (name, field) in enumerate(schema.fields.items()): field_prompt = f" {json.dumps(name)}: " + ('"' if field.field_type != "boolean" else '') suffix = tokenizer.encode(field_prompt) candidate_texts = [c + ',' if field.field_type == "boolean" else json.dumps(c, ensure_ascii=False)[1:] for c in field.choices] expected_scores, tree_scores = [], [] for candidate in candidate_texts: tokens = tokenizer.encode(candidate) ids = prefix + suffix + tokens logits = forward(torch.tensor([ids], device=device), use_cache=False).logits[0].float() log_probs = logits.log_softmax(-1) expected_scores.append(sum(log_probs[len(prefix) + len(suffix) - 1 + j, token] for j, token in enumerate(tokens))) score = 0 for j, token in enumerate(tokens): row = context_rows[tuple(suffix + tokens[:j])] actual = tree_output.logits[0, row].float() torch.testing.assert_close(actual, logits[len(prefix) + len(suffix) - 1 + j], atol=1e-5, rtol=1e-5) score += actual.log_softmax(-1)[token] tree_scores.append(score) torch.testing.assert_close(torch.stack(tree_scores), torch.stack(expected_scores), atol=1e-5, rtol=1e-5) lengths = torch.tensor([len(tokenizer.encode(c)) for c in candidate_texts], device=device) probabilities = (torch.stack(expected_scores) / lengths / max(temperature, 1e-4)).softmax(-1).tolist() actual = {c["choice"]: c["probability"] for c in result["field_telemetry"][name]["top_choices"]} for choice, probability in zip(field.choices, probabilities): assert actual[choice] == pytest.approx(probability, abs=0.0001) def test_tree_shares_field_and_candidate_prefixes(tiny_engine, monkeypatch): model, tokenizer = tiny_engine schema = StructuredSchema({ "shared_name_a": {"choices": ["abc", "abd", "z"]}, "shared_name_b": {"choices": ["abc", "abd", "z"]}, }) paths = set() unshared_nodes = 0 for name, field in schema.fields.items(): suffix = tokenizer.encode(f' {json.dumps(name)}: "') candidates = [tokenizer.encode(json.dumps(c, ensure_ascii=False)[1:]) for c in field.choices] for tokens in candidates: path = suffix + tokens[:-1] paths.update(tuple(path[:j]) for j in range(1, len(path) + 1)) unshared_nodes += len(suffix) + sum(len(t) - 1 for t in candidates) tree = schema.compile_tree_metadata(tokenizer) assert tree["input_ids"].shape[1] == len(paths) < unshared_nodes assert schema.compile_tree_metadata(tokenizer) is tree assert schema.compile_tree_metadata(Tokenizer()) is not tree def no_copy(*args, **kwargs): pytest.fail("Tree inference must not copy or batch-repeat the prefix KV cache") from transformers import DynamicCache monkeypatch.setattr(DynamicCache, "__deepcopy__", no_copy, raising=False) monkeypatch.setattr(DynamicCache, "batch_repeat_interleave", no_copy) result = engine.run_parallel_generation_tree("test", schema) assert result["tree_nodes"] == len(paths) def test_tree_equal_token_likelihoods_do_not_favor_shorter_candidates(tiny_engine): model, _ = tiny_engine # Every token has probability 1/vocab_size. Different-length candidates # must tie under mean scoring, including when they share a value prefix. with torch.no_grad(): model.lm_head.weight.zero_() schema = StructuredSchema({ "plain": {"choices": ["A", "LONG_VALUE"]}, "shared": {"choices": ["PAYMENT_A", "PAYMENT_LONG_VALUE"]}, "flag": {"type": "boolean"}, }) result = engine.run_parallel_generation_tree("test", schema) for field in result["field_telemetry"].values(): assert [c["probability"] for c in field["top_choices"]] == [0.5, 0.5] assert result["sequential_forward_passes"] == 2 def test_tree_many_candidates_still_uses_two_passes(tiny_engine, monkeypatch): model, _ = tiny_engine schema = StructuredSchema({"category": {"choices": [f"VALUE_{i:03d}" for i in range(255)]}}) calls = [] handle = model.register_forward_pre_hook(lambda module, args: calls.append(1)) try: result = engine.run_parallel_generation_tree("test", schema) finally: handle.remove() assert len(calls) == result["sequential_forward_passes"] == 2 assert result["parsed_json"]["category"]["value"] in schema["category"].choices assert result["field_telemetry"]["category"]["cardinality"] == 255 def test_tree_rejects_negative_temperature(tiny_engine): with pytest.raises(ValueError, match="temperature"): engine.run_parallel_generation_tree("test", StructuredSchema({"x": {"type": "boolean"}}), -1) def test_tree_rejects_unsupported_attention(tiny_engine): model, _ = tiny_engine schema = StructuredSchema({"x": {"type": "boolean"}}) model.config.layer_types = ["sliding_attention"] with pytest.raises(ValueError, match="full attention"): engine.run_parallel_generation_tree("test", schema) def test_tree_restores_backend_on_error(tiny_engine, monkeypatch): model, _ = tiny_engine backend = model.config._attn_implementation forward = model.forward def fail_tree(*args, **kwargs): if "past_key_values" in kwargs: assert model.config._attn_implementation == "flex_attention" raise RuntimeError("tree forward failed") return forward(*args, **kwargs) monkeypatch.setattr(model, "forward", fail_tree) with pytest.raises(RuntimeError, match="tree forward failed"): engine.run_parallel_generation_tree("test", StructuredSchema({"x": {"type": "boolean"}})) assert model.config._attn_implementation == backend assert engine._gpu_lock.acquire(timeout=1) engine._gpu_lock.release()