File size: 11,086 Bytes
16e3e7d | 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 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 | """Constrained decoding: make an echo impossible instead of merely detected."""
import pytest
from tiny_computer_agent import (
Action, BrainError, DecisionMenu, Element, GoalUnreachable, ModelBrain, Observation,
)
from tiny_computer_agent.decoding import (
MenuMask, brevity_ladder, choice_token_ids, digit_token_ids, majority_choice, read_choice,
)
FRAME = Observation(elements=(
Element("Files", x=10, y=6, width=60, height=20),
Element("Save", x=200, y=270, width=90, height=30),
Element("Cancel", x=400, y=270, width=90, height=30),
), text="Untitled document - unsaved changes")
class FakeTokenizer:
"""A tokenizer where each digit is one token (bare and space-prefixed)."""
def __init__(self, expressible=True):
self.expressible = expressible
self.table = {}
if expressible:
for digit in range(10):
self.table[str(digit)] = 100 + digit
self.table[" " + str(digit)] = 200 + digit
for index, word in enumerate(("yes", "Yes", "no", "No")):
self.table[word] = 300 + index
self.table[" " + word] = 320 + index
def encode(self, text, add_special_tokens=False):
if text in self.table:
return [self.table[text]]
return [999, 999] # never a single token
# -- token ids ---------------------------------------------------------------
def test_digit_ids_cover_every_digit_exactly_once():
ids = digit_token_ids(FakeTokenizer())
assert len(ids) == 20 # 0-9, bare and space-prefixed
assert set(range(100, 110)) <= set(ids)
def test_choice_ids_cover_only_the_legal_answers():
assert sorted(choice_token_ids(FakeTokenizer(), 3)) == sorted([100, 101, 102, 103, 200, 201, 202, 203])
assert sorted(choice_token_ids(FakeTokenizer(), 2, allow_done=False)) == sorted([101, 102, 201, 202])
def test_a_tokenizer_that_cannot_spell_a_number_yields_no_ids():
assert digit_token_ids(FakeTokenizer(expressible=False)) == []
assert choice_token_ids(FakeTokenizer(expressible=False), 5) == []
# -- the mask ----------------------------------------------------------------
def test_the_mask_forbids_everything_that_is_not_a_menu_number():
torch = pytest.importorskip("torch")
scores = torch.zeros(1, 8)
masked = MenuMask([1, 3])(None, scores)
assert masked[0, 1].item() == 0 and masked[0, 3].item() == 0
for blocked in (0, 2, 4, 5, 6, 7):
assert masked[0, blocked].item() == float("-inf")
assert MenuMask([1, 3]).allowed_ids == (1, 3)
def test_the_mask_keeps_the_end_of_sequence_reachable_after_one_token():
torch = pytest.importorskip("torch")
process = MenuMask([1], eos_token_id=7)
first = process(None, torch.zeros(1, 8))
assert first[0, 7].item() == float("-inf"), "stopping immediately is not an answer"
assert first[0, 1].item() == 0
second = process(None, torch.zeros(1, 8))
assert second[0, 7].item() == 0 # ...but it can stop after answering
assert second[0, 2].item() == float("-inf")
def test_the_mask_can_be_told_not_to_force_a_token():
torch = pytest.importorskip("torch")
process = MenuMask([1], eos_token_id=7, force_first=False)
assert process(None, torch.zeros(1, 8))[0, 7].item() == 0
# -- reading an answer -------------------------------------------------------
def test_read_choice_reads_only_the_digits():
assert read_choice("3", 12) == 3
assert read_choice(" 3 ", 12) == 3
assert read_choice("3.", 12) == 3
assert read_choice("12", 12) == 12
assert read_choice("0", 12) == 0 # 0 is "already done"
assert read_choice("I would click item 2", 12) == 2
def test_read_choice_refuses_what_is_not_exactly_one_answer():
assert read_choice("", 12) is None
assert read_choice("no idea", 12) is None
assert read_choice("13", 12) is None # outside the menu
assert read_choice("2 or 3", 12) is None # 23: ambiguous, so refused
assert read_choice("0", 12, allow_done=False) is None
# -- the ladder and the vote -------------------------------------------------
def test_the_brevity_ladder_grows_then_stops():
assert brevity_ladder(3) == (0, 1, 2)
assert brevity_ladder(1) == (0,)
assert brevity_ladder(0) == (0,)
def test_majority_choice_ignores_failures_and_breaks_ties_by_first_seen():
assert majority_choice([1, 1, 2]) == 1
assert majority_choice([None, 3, None]) == 3
assert majority_choice([2, 1]) == 2 # a tie goes to the first seen
assert majority_choice([]) is None
assert majority_choice([None, None]) is None
# -- the brain, in constrained mode ------------------------------------------
class Scripted(ModelBrain):
"""A backend whose replies are scripted, so the ladder can be tested."""
name = "scripted"
supports_constrained = True
def __init__(self, replies, tokenizer=None, **kwargs):
super().__init__("fake/model", **kwargs)
self._replies = list(replies)
self._tokenizer = tokenizer or FakeTokenizer()
self.prompts = []
def _tokenizer_for_choice(self):
return self._tokenizer
def _generate_constrained(self, prompt, observation, allowed_ids, max_new_tokens, *, sample=False):
self.prompts.append(prompt)
assert allowed_ids, "the mask must be non-empty"
return self._replies.pop(0)
def _generate(self, prompt, observation):
self.prompts.append(prompt)
return self._replies.pop(0)
def test_choice_mode_turns_a_number_into_a_click():
brain = Scripted(["2"], decision_mode="choice")
action = brain.decide("save the document", FRAME)
assert action.kind == "click" and action.label == "Save"
assert "constrained choice 2" in action.reason
assert brain.last_choice == 2
def test_choice_mode_reads_zero_as_already_done():
action = Scripted(["0"], decision_mode="choice").decide("save", FRAME)
assert action.kind == "done"
def test_the_echoed_prompt_is_not_usable_even_if_the_model_returns_it():
"""The old failure mode, fed straight in: it cannot be mistaken for an answer."""
brain = Scripted(["2"], decision_mode="choice")
echo = brain.build_prompt("save", DecisionMenu(FRAME), FRAME, brevity=0)
assert read_choice(echo, 12) is None, "an echoed menu must not parse as a choice"
# the prompt contains the menu numbers, spaced apart: 1, 2, 3 -> "123"
assert read_choice(brain.build_prompt("save", DecisionMenu(FRAME), FRAME), 3) is None
def test_a_useless_reply_shortens_the_prompt_and_tries_again():
brain = Scripted(["", "3"], decision_mode="choice", retries=3)
action = brain.decide("discard the change", FRAME)
assert action.label == "Cancel"
assert len(brain.prompts) == 2
assert len(brain.prompts[1]) < len(brain.prompts[0]), "the retry must be shorter"
def test_the_prompt_shrinks_all_the_way_when_nothing_works():
brain = Scripted([""] * 5, decision_mode="choice", retries=3)
with pytest.raises(BrainError, match="not a menu number"):
brain.decide("save", FRAME)
lengths = [len(prompt) for prompt in brain.prompts]
assert lengths[0] > lengths[-1] and len(lengths) == 3
assert "Text on screen" not in brain.prompts[-1]
def test_votes_take_a_majority_over_samples():
brain = Scripted(["2", "3", "2"], decision_mode="choice", votes=3)
action = brain.decide("save", FRAME)
assert action.label == "Save" # 2 wins over 3
assert len(brain.prompts) == 3 # only one ladder rung used
def test_a_tokenizer_that_cannot_spell_numbers_is_refused_clearly():
brain = Scripted(["1"], tokenizer=FakeTokenizer(expressible=False), decision_mode="choice")
with pytest.raises(BrainError, match="cannot express a menu number"):
brain.decide("save", FRAME)
def test_an_empty_screen_is_a_hard_stop_in_choice_mode():
brain = Scripted(["1"], decision_mode="choice")
with pytest.raises(GoalUnreachable):
brain.decide("save", Observation(elements=()))
class ScanBackend(ModelBrain):
"""Answers yes only for the labels it was told to like."""
name = "scan-backend"
supports_constrained = True
def __init__(self, likes, tokenizer=None, **kwargs):
super().__init__("fake/model", **kwargs)
self.likes = set(likes)
self.tokenizer = tokenizer or FakeTokenizer()
self.prompts = []
def _tokenizer_for_choice(self):
return self.tokenizer
def _generate_constrained(self, prompt, observation, allowed_ids, max_new_tokens, *, sample=False):
self.prompts.append(prompt)
label = prompt.split("Candidate item: ", 1)[1].splitlines()[0]
return "yes" if label in self.likes else "no"
def test_scanning_asks_one_question_per_item_and_takes_the_first_yes():
brain = ScanBackend(["Save"], decision_mode="scan")
action = brain.decide("click 'Save' to save the document", FRAME)
assert action.kind == "click" and action.label == "Save"
assert "said yes to 'Save'" in action.reason
assert len(brain.prompts) == 2 # Files, Save - then it stopped
assert all("yes or no" in prompt for prompt in brain.prompts)
assert all("Candidate item:" in prompt for prompt in brain.prompts)
def test_scanning_refuses_when_every_item_is_rejected():
brain = ScanBackend([], decision_mode="scan")
with pytest.raises(BrainError, match="said no to all"):
brain.decide("save", FRAME)
assert len(brain.prompts) == 3 # every item was asked about
def test_scanning_refuses_a_reply_that_is_neither_yes_nor_no():
class Mute(ScanBackend):
def _generate_constrained(self, prompt, observation, allowed_ids, max_new_tokens, *, sample=False):
return "maybe"
with pytest.raises(BrainError, match="neither yes nor no|masked reply"):
Mute(["Save"], decision_mode="scan").decide("save", FRAME)
def test_read_verdict_reads_what_the_mask_allows():
from tiny_computer_agent.decoding import pair_token_ids, read_verdict
assert read_verdict("yes") == "yes"
assert read_verdict(" Yes ") == "yes"
assert read_verdict("no") == "no"
assert read_verdict("No.") == "no"
assert read_verdict("perhaps") is None
assert read_verdict("") is None
yes_ids, no_ids = pair_token_ids(FakeTokenizer())
assert yes_ids and no_ids and not (set(yes_ids) & set(no_ids))
def test_json_mode_is_untouched():
"""The default must not change: the existing lenient parsing still applies."""
brain = Scripted(['{"kind": "click", "index": 2}'])
assert brain.decision_mode == "json"
assert brain.decide("save", FRAME).label == "Save"
def test_choice_mode_is_only_entered_when_the_backend_supports_it():
class Plain(Scripted):
supports_constrained = False
brain = Plain(['{"kind": "click", "index": 2}'], decision_mode="choice")
assert brain.decide("save", FRAME).label == "Save" # fell back to JSON parsing
|