"""Making a tiny model answer with a choice instead of repeating the question. A 256M model asked for JSON on a busy screen does the thing it is best at: it continues the text it was handed, and gives the prompt back. The harness already *detects* that (echo detection refuses the reply) but detecting it is not fixing it - the episode simply fails. The fix is to stop asking for free text. A menu of twelve items has twelve legal answers, so mask every other token: an echo becomes impossible rather than merely caught, and a model that cannot write a sentence can still pick a number. Three tools, all model-agnostic: * :func:`digit_token_ids` / :func:`choice_token_ids` - the token ids a menu answer can possibly need * :class:`MenuMask` - a transformers ``LogitsProcessor`` that forbids everything else, so "hello, here is your prompt back" is not representable * :func:`brevity_ladder` and :func:`majority_choice` - a shorter prompt on the retry (echoing is driven by prompt length) and a vote over k attempts """ from __future__ import annotations from typing import Any, Iterable, Sequence DIGITS = tuple(str(d) for d in range(10)) def _single_token_ids(tokenizer: Any, text: str) -> list[int]: """Ids that render ``text`` as exactly one token, with and without a space.""" ids: list[int] = [] for variant in (text, " " + text): try: encoded = tokenizer.encode(variant, add_special_tokens=False) except Exception: continue if len(encoded) == 1: ids.append(int(encoded[0])) return ids def digit_token_ids(tokenizer: Any) -> list[int]: """Every single-token digit, so any menu number can be spelled out. The loosest useful mask: ten or twenty tokens instead of a 50k vocabulary. A model restricted to these cannot answer with prose or with the prompt. """ ids: list[int] = [] for digit in DIGITS: ids += _single_token_ids(tokenizer, digit) return sorted(set(ids)) def choice_token_ids(tokenizer: Any, count: int, *, allow_done: bool = True) -> list[int]: """The single-token ids for ``0..count`` (0 means "the goal is already met"). Stricter than :func:`digit_token_ids`: the model may emit exactly one of the legal answers, and nothing else, in a single token. """ ids: list[int] = [] for number in range(0 if allow_done else 1, count + 1): ids += _single_token_ids(tokenizer, str(number)) return sorted(set(ids)) def pair_token_ids(tokenizer: Any, *, yes: Sequence[str] = ("yes", "Yes"), no: Sequence[str] = ("no", "No")) -> tuple[list[int], list[int]]: """Single-token ids for a yes/no answer. Two tokens instead of a whole vocabulary: a 256M model that cannot pick 1-of-12 can often still answer "does this label do what the goal asks?". """ yes_ids: list[int] = [] no_ids: list[int] = [] for text in yes: yes_ids += _single_token_ids(tokenizer, text) for text in no: no_ids += _single_token_ids(tokenizer, text) return sorted(set(yes_ids)), sorted(set(no_ids)) def read_verdict(text: str) -> str | None: """'yes', 'no', or None when the reply is neither.""" cleaned = (text or "").strip().lower().lstrip(":").strip() if cleaned.startswith("yes"): return "yes" if cleaned.startswith("no"): return "no" return None def MenuMask(allowed_ids: Sequence[int], *, eos_token_id: int | None = None, force_first: bool = True): """A ``LogitsProcessor`` that permits only the ids in ``allowed_ids``. Returns a plain callable with the transformers signature ``(input_ids, scores) -> scores``, so it needs no torch import at module load and can be tested without a model. ``force_first`` closes a hole that a real 256M model found immediately: if the end-of-sequence token is legal from the start, the cheapest continuation is to stop, and the answer comes back *empty*. The mask therefore withholds EOS until one real token has been emitted, so the model has to say something. """ allowed = sorted(set(int(i) for i in allowed_ids)) state = {"calls": 0} def process(input_ids: Any, scores: Any) -> Any: mask = scores.new_full(scores.shape, float("-inf")) for index in allowed: if 0 <= index < scores.shape[-1]: mask[..., index] = scores[..., index] if eos_token_id is not None: stop = int(eos_token_id) if not force_first or state["calls"] > 0: if 0 <= stop < scores.shape[-1]: mask[..., stop] = scores[..., stop] state["calls"] += 1 return mask process.allowed_ids = tuple(allowed) # for tests and for the trace process.eos_token_id = eos_token_id process.force_first = force_first return process def read_choice(text: str, count: int, *, allow_done: bool = True) -> int | None: """The menu number in a reply, or None if there is not exactly one. Only the digits are read: with the mask in place the reply *is* a number, and anything else means the model is not being constrained (so it is refused rather than guessed at). """ if not text: return None digits = "".join(character for character in text if character.isdigit()) if not digits: return None try: number = int(digits) except ValueError: # pragma: no cover - defensive return None low = 0 if allow_done else 1 return number if low <= number <= count else None def brevity_ladder(levels: int = 3) -> tuple[int, ...]: """The brevity settings to try, in order: full prompt first, terse last.""" return tuple(range(0, max(1, levels))) def majority_choice(choices: Iterable[int | None]) -> int | None: """The most common choice, ignoring the failures. Ties go to the first seen.""" counts: dict[int, int] = {} order: list[int] = [] for choice in choices: if choice is None: continue if choice not in counts: counts[choice] = 0 order.append(choice) counts[choice] += 1 if not counts: return None best = max(counts.values()) return next(choice for choice in order if counts[choice] == best) __all__ = [ "DIGITS", "MenuMask", "brevity_ladder", "choice_token_ids", "digit_token_ids", "majority_choice", "pair_token_ids", "read_choice", "read_verdict", ]