| """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) |
| 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: |
| 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", |
| ] |
|
|
