| 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 |
| |
| 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 |
|
|
| |
| |
| |
| 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 |
| |
| |
| 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() |
|
|
