| """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] |
|
|
|
|
| |
|
|
| def test_digit_ids_cover_every_digit_exactly_once(): |
| ids = digit_token_ids(FakeTokenizer()) |
| assert len(ids) == 20 |
| 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) == [] |
|
|
|
|
| |
|
|
| 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 |
| 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 |
|
|
|
|
| |
|
|
| 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 |
| 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 |
| assert read_choice("2 or 3", 12) is None |
| assert read_choice("0", 12, allow_done=False) is None |
|
|
|
|
| |
|
|
| 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 |
| assert majority_choice([]) is None |
| assert majority_choice([None, None]) is None |
|
|
|
|
| |
|
|
| 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" |
| |
| 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" |
| assert len(brain.prompts) == 3 |
|
|
|
|
| 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 |
| 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 |
|
|
|
|
| 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" |
|
|
