diff --git a/src/semif_phase1/serial.py b/src/semif_phase1/serial.py index 9beb8c7..5948dcc 100644 --- a/src/semif_phase1/serial.py +++ b/src/semif_phase1/serial.py @@ -7,6 +7,7 @@ import inspect import json import time +from itertools import takewhile from .core import direct_messages, softmax, synchronize from .direct import PROMPT_VERSION, encode_prompt @@ -28,7 +29,12 @@ def _state_prefix(tokenizer, state) -> list[int]: if prompt.count(payload) != 1 or not payload.startswith(evidence): raise ValueError("Cannot establish a deterministic evidence prefix") text = prompt[: prompt.index(payload)] + evidence - return tokenizer.encode(text, add_special_tokens=False)[:-1] + # Appending the next field's punctuation can merge with the final boundary token + # by an amount that depends on the state text, so trim to the longest common + # prefix instead of assuming the merge is always exactly one token. + ids = tokenizer.encode(text, add_special_tokens=False) + extended = tokenizer.encode(text + payload[len(evidence) :], add_special_tokens=False) + return [first for first, _ in takewhile(lambda pair: pair[0] == pair[1], zip(ids, extended))] def _cached_forward(model, inputs): diff --git a/src/semif_phase1/shared.py b/src/semif_phase1/shared.py index ccb0c2a..af67941 100644 --- a/src/semif_phase1/shared.py +++ b/src/semif_phase1/shared.py @@ -6,6 +6,7 @@ import inspect import json import time +from itertools import takewhile from .core import direct_messages, softmax, synchronize from .direct import PROMPT_VERSION, encode_prompt @@ -33,8 +34,12 @@ def _state_prefix(tokenizer, state) -> list[int]: if not payload.startswith(evidence): raise ValueError("Evidence serialization changed") text = prompt[: prompt.index(payload)] + evidence - # Appending JSON punctuation can merge with the final boundary token. - return tokenizer.encode(text, add_special_tokens=False)[:-1] + # Appending the next field's punctuation can merge with the final boundary token + # by an amount that depends on the state text, so trim to the longest common + # prefix instead of assuming the merge is always exactly one token. + ids = tokenizer.encode(text, add_special_tokens=False) + extended = tokenizer.encode(text + payload[len(evidence) :], add_special_tokens=False) + return [first for first, _ in takewhile(lambda pair: pair[0] == pair[1], zip(ids, extended))] def _suffix_layout(sequences: list[list[int]], prefix_length: int, pad_id: int): diff --git a/tests/test_state_prefix.py b/tests/test_state_prefix.py new file mode 100644 index 0000000..7875ee5 --- /dev/null +++ b/tests/test_state_prefix.py @@ -0,0 +1,87 @@ +"""Shared state prefixes must be true token prefixes of every full prompt.""" +import pytest + +from semif_phase1.core import direct_messages +from semif_phase1.serial import _state_prefix as serial_prefix +from semif_phase1.shared import _state_prefix as shared_prefix + +PREFIXES = (serial_prefix, shared_prefix) + +# The pinned Qwen tokenizer lets a token span the comma that follows the evidence, +# absorbing the two characters in front of it. Whether that happens depends on the +# state text, which is the whole point: a caller cannot predict it. These stubs +# reproduce both extremes byte for byte, so the regression runs with no torch, no +# network and no multi-gigabyte download. +FUSED_TOKEN = 0x110000 + + +class PlainTokenizer: + """Byte tokenizer that never merges, so the boundary costs exactly one token.""" + + def apply_chat_template(self, turns, tokenize=False, add_generation_prompt=True, enable_thinking=False): + return "HEADER\n" + turns[-1]["content"] + "\nASSISTANT" + + def encode(self, text, add_special_tokens=False): + return list(text.encode()) + + def decode(self, tokens): + return bytes(token for token in tokens if token != FUSED_TOKEN).decode() + + +class MergingTokenizer(PlainTokenizer): + """Byte tokenizer whose tokens straddle the comma after the evidence.""" + + def encode(self, text, add_special_tokens=False): + raw = list(text.encode()) + tokens, index = [], 0 + while index < len(raw): + if raw[index + 1 : index + 3] in ([ord("}"), ord(",")], [ord('"'), ord(",")]): + tokens.append(FUSED_TOKEN) + index += 3 + else: + tokens.append(raw[index]) + index += 1 + return tokens + + +TOKENIZERS = (PlainTokenizer, MergingTokenizer) +STATES = [ + "Did the deploy use the same parameters?", + {"a": "why)"}, + {"a": {"b": "z-"}, "note": "keep"}, + ["y{"], + "owned state", + {"priority": "urgent"}, +] + + +def _full_prompt(tokenizer, state): + row = { + "id": "x", + "state": state, + "question": "Is this urgent?", + "options": [{"id": "a", "description": "A"}, {"id": "b", "description": "B"}], + } + payload = direct_messages(row)[-1]["content"] + return tokenizer.encode("HEADER\n" + payload + "\nASSISTANT") + + +@pytest.mark.parametrize("tokenizer", TOKENIZERS, ids=lambda token: token.__name__) +@pytest.mark.parametrize("state", STATES, ids=repr) +@pytest.mark.parametrize("prefix_of", PREFIXES, ids=lambda function: function.__module__) +def test_state_prefix_is_a_token_prefix_of_the_full_prompt(tokenizer, state, prefix_of): + prefix = prefix_of(tokenizer(), state) + full = _full_prompt(tokenizer(), state) + assert prefix + assert full[: len(prefix)] == prefix + assert len(full) > len(prefix) + + +@pytest.mark.parametrize("tokenizer", TOKENIZERS, ids=lambda token: token.__name__) +@pytest.mark.parametrize("state", STATES, ids=repr) +@pytest.mark.parametrize("prefix_of", PREFIXES, ids=lambda function: function.__module__) +def test_state_prefix_never_carries_the_runtime_question_or_options(tokenizer, state, prefix_of): + decoded = tokenizer().decode(prefix_of(tokenizer(), state)) + assert decoded.startswith('HEADER\n{"evidence": ') + assert "criterion" not in decoded + assert "options" not in decoded