File size: 3,432 Bytes
065ec46 | 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 | """Turn a Decision into token ids, and say where each candidate ends.
The layout is state -> question -> candidates. The backbone is causal, so a
candidate's last token has attended to the state, the question and every
candidate before it; its hidden state is what the head scores. Candidate order
is shuffled during training (for `choice`) so no position is favoured.
Ids are built segment by segment rather than by tokenising one string and
searching for offsets. Training and inference share this function, so the two
can never disagree about where a candidate ends -- which is the whole contract.
"""
from __future__ import annotations
from dataclasses import dataclass
from .schema import Decision
SYSTEM = ("You are a decision model. Read the state, then judge every candidate "
"answer to the question.")
_HEADER = {
"choice": "Question (pick one): {q}\nCandidates:",
"null": "Is this true? {q}\nCandidates:",
"score": "Question (rate on the scale): {q}\nCandidates:",
}
@dataclass
class Rendered:
input_ids: list[int]
option_positions: list[int] # index of each candidate's last token
def _frame(tok) -> tuple[str, str]:
"""(prefix, suffix) of the chat frame around the user turn, or plain text
for a tokenizer with no template."""
marker = "\x00BODY\x00"
if getattr(tok, "chat_template", None):
text = tok.apply_chat_template(
[{"role": "system", "content": SYSTEM},
{"role": "user", "content": marker}],
tokenize=False, add_generation_prompt=True)
pre, post = text.split(marker)
bos = getattr(tok, "bos_token", None)
if bos and pre.startswith(bos):
pre = pre[len(bos):] # added back explicitly below
return pre, post
return f"{SYSTEM}\n\n", "\n"
class Renderer:
def __init__(self, tok, max_state_tokens: int = 384, max_option_tokens: int = 32,
max_question_tokens: int = 96):
self.tok = tok
self.max_state = max_state_tokens
self.max_option = max_option_tokens
self.max_question = max_question_tokens
pre, post = _frame(tok)
self._pre = self._ids(pre)
self._post = self._ids(post)
bos_id = getattr(tok, "bos_token_id", None)
adds_bos = bool(tok("a").input_ids[:1] == [bos_id]) if bos_id is not None else False
self._bos = [bos_id] if adds_bos else []
def _ids(self, s: str, limit: int | None = None) -> list[int]:
ids = self.tok(s, add_special_tokens=False).input_ids
return ids[:limit] if limit else ids
def render(self, d: Decision, candidates: list[str] | None = None) -> Rendered:
cands = candidates if candidates is not None else d.candidates()
ids = self._bos + list(self._pre)
if d.state.strip():
ids += self._ids("State:\n")
ids += self._ids(d.state.strip(), self.max_state)
ids += self._ids("\n\n")
q = self._ids(d.question.strip(), self.max_question)
head = _HEADER[d.type].split("{q}")
ids += self._ids(head[0]) + q + self._ids(head[1])
positions = []
for i, c in enumerate(cands, 1):
ids += self._ids(f"\n[{i}] ")
ids += self._ids(c.strip(), self.max_option) or self._ids("(empty)")
positions.append(len(ids) - 1)
ids += self._post
return Rendered(ids, positions)
|