File size: 6,834 Bytes
3145545 | 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 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 | import copy
import json
from dataclasses import replace
import jsonschema
import pytest
import torch
from pcd.cache import fork_cache, state_tensors
from pcd.config import Limits
from pcd.prompting import prompt_tokens
from pcd.schema import evaluate, validate_schema
SCHEMA = {
"type": "object",
"properties": {
"service": {
"type": "string",
"enum": ["express", "express plus", "日本語", 'quote"slash\\'],
},
"insured": {
"type": "boolean",
"description": "Whether insurance was requested",
},
},
"required": ["service", "insured"],
"additionalProperties": False,
}
@pytest.mark.parametrize("mode", ["sequence", "token"])
def test_scores_match_uncached_and_microbatch_invariant(engine, mode):
context = "Use express plus with insurance."
reference = engine.reference(context, SCHEMA, mode=mode)
limits = engine.limits
try:
for batch_size in (1, 3, 32):
engine.limits = replace(limits, branch_batch_size=batch_size)
actual = engine.constrained(context, SCHEMA, mode=mode)
assert type(actual["object"]["insured"]) is bool
assert actual["calibrated"] is False
for field, scores in reference.items():
torch.testing.assert_close(
torch.tensor([c["score"] for c in actual["fields"][field]["candidates"]]),
torch.tensor(scores),
atol=2e-5,
rtol=2e-5,
)
assert (
abs(sum(c["probability"] for c in actual["fields"][field]["candidates"]) - 1)
< 1e-10
)
assert (
actual["forward_calls"] == 1 + (actual["branches"] + batch_size - 1) // batch_size
)
finally:
engine.limits = limits
@torch.inference_mode()
def test_cache_isolation_and_both_convolution_paths(engine):
assert engine.validate_cache()["prefix_unchanged"]
base = engine._prefill(engine.tokenizer.encode("prefix is here"))
snapshot = [t.clone() for t in state_tensors(base)]
branch = fork_cache(base, 2)
for tensor in state_tensors(branch):
old_second = tensor[1].clone()
tensor[0].zero_()
assert torch.equal(tensor[1], old_second)
for old, current in zip(snapshot, state_tensors(base)):
assert torch.equal(old, current)
with pytest.raises(ValueError):
fork_cache(base, 0)
def test_strict_schema_and_output():
assert len(validate_schema(SCHEMA)) == 2
variants = []
for key, value in [
("allOf", [{}]),
("additionalProperties", True),
("required", ["insured"]),
]:
variants.append({**SCHEMA, key: value})
variants.extend(
[
{
**SCHEMA,
"properties": {
"service": {"type": "string", "enum": ["x", "x"]},
"insured": {"type": "boolean"},
},
},
{
**SCHEMA,
"properties": {
"service": {"type": "string"},
"insured": {"type": "boolean"},
},
},
{
**SCHEMA,
"properties": {
"service": {"type": "number"},
"insured": {"type": "boolean"},
},
},
]
)
for schema in variants:
with pytest.raises((ValueError, jsonschema.SchemaError)):
validate_schema(schema)
good = {"service": "express", "insured": False}
assert evaluate(json.dumps(good), SCHEMA, good)["exact_match"]
assert not evaluate(json.dumps({**good, "insured": 0}), SCHEMA, good)["schema_compliant"]
assert not evaluate("```json\n" + json.dumps(good) + "\n```", SCHEMA, good)["syntax_valid"]
assert not evaluate(json.dumps({**good, "extra": True}), SCHEMA, good)["schema_compliant"]
def test_limits_reject_before_forward(engine, monkeypatch):
def forbidden(*args, **kwargs):
raise AssertionError("model must not run")
monkeypatch.setattr(engine.model.model, "forward", forbidden)
with pytest.raises(ValueError):
engine.constrained("x" * 70000, SCHEMA)
with pytest.raises(ValueError):
engine.constrained("x", SCHEMA, mode="unknown")
huge = copy.deepcopy(SCHEMA)
huge["properties"]["service"]["enum"] = [str(i) for i in range(257)]
with pytest.raises(ValueError):
engine.constrained("x", huge)
with pytest.raises(ValueError):
Limits(branch_batch_size=0)
def test_atomic_codes_and_reasoning_prefix(engine):
compiled = engine.compile(SCHEMA, "token")
for field in compiled.fields:
assert len(set(field.token_ids)) == len(field.values)
assert all(len(engine.tokenizer.encode(label)) == 1 for label in field.labels)
prefix = engine.tokenizer.decode(
prompt_tokens(engine.tokenizer, compiled, "test", engine.limits)
)
assert not prefix.endswith("<think>")
for field in compiled.fields:
assert "<think></think>\n" in engine.tokenizer.decode(field.suffix)
sequence = engine.compile(SCHEMA, "sequence")
assert engine.tokenizer.decode(
prompt_tokens(engine.tokenizer, sequence, "test", engine.limits)
).endswith("<think></think>\n{\n")
native = engine.tokenizer.decode(
prompt_tokens(engine.tokenizer, compiled, "test", engine.limits, native=True)
)
assert native.endswith("<think>")
def test_reference_target_shift_uses_preceding_logits(engine):
context = "express"
field = engine.compile(SCHEMA, "sequence").fields[0]
prefix = prompt_tokens(
engine.tokenizer, engine.compile(SCHEMA, "sequence"), context, engine.limits
)
value = field.value_tokens[0]
tokens = prefix + list(field.suffix + value)
with torch.inference_mode():
logits = engine.model(engine.tensor([tokens]), use_cache=False).logits[0].float()
start = len(prefix) + len(field.suffix)
expected = sum(float(logits[start + i - 1].log_softmax(-1)[t]) for i, t in enumerate(value))
actual = engine.reference(context, SCHEMA)[field.name][0]
assert abs(actual - expected) < 2e-5
def test_token_ar_codec_is_not_json_repair(engine, monkeypatch):
coded = {f.name: f.labels[0] for f in engine.compile(SCHEMA, "token").fields}
continuation = json.dumps(coded)[1:]
def generate(ids, **kwargs):
return torch.cat([ids, engine.tensor([engine.tokenizer.encode(continuation)])], dim=1)
monkeypatch.setattr(engine.model, "generate", generate)
actual = engine.autoregressive("test", SCHEMA, mode="token")
assert json.loads(actual["text"]) == {"service": "express", "insured": True}
|