File size: 6,558 Bytes
b038669 | 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 | """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",
]
|