runtime/src/decisionnode_edge/template.py

"""The prompt format the model reads: the added tokens, the answer codes and the encoder.

One prompt per question. The state block comes first, then the question block, then the
word "Answer" and the <|answer|> token, whose position the answer is read from:

    <|state|>\\n{Picture n (id): <|vision_start|>pads<|vision_end|>\\n ...}{state}\\n
    <|question|>\\ntype: {choice|score|noul|coord}\\n{instructions}\\n{option lines}Answer<|answer|>

A choice option line is `<|opt_NNN|> key: description\\n` (`<|opt_NNN|> key\\n` when the
description is empty), a score level line `<|opt_NNN|> index: level\\n`. A truth question
has no option lines (the prompt calls it `noul`) and carries
its criteria as `yes: ...\\n` and `no: ...\\n` lines. It is answered by <|yes|> and <|no|>.

A coordinate row (one reading of a point or box question, places.py) is `type: coord`, its
text and no option lines. It is answered by the codes 0 to n - 1 in order (code k is step k
of a side of the picture, or a count of k), the same codes for every such row, never drawn.

Each option of a choice or a score gets one of 256 code tokens. The codes are drawn from a
stream seeded by a SHA-256 of the question alone (never the state), so the same question
always gets the same codes. The seed's first part is a fixed string from the model directory
(`code_seed_hex` in decisionnode_edge.json): the model gives its published answers only
with that string.

Request text can never become a control token: plain text is tokenised with
`split_special_tokens=True`, and the few added tokens that flag does not split (such as
`<think>`) are cut inside the string before tokenising.
"""

from __future__ import annotations

import hashlib
import re
from dataclasses import dataclass, field
from typing import Any, Sequence

STATE = "<|state|>"
QUESTION = "<|question|>"
ANSWER = "<|answer|>"
YES = "<|yes|>"
NO = "<|no|>"
N_CODES = 256
CODES = tuple(f"<|opt_{i:03d}|>" for i in range(N_CODES))
ADDED_TOKENS = (STATE, QUESTION, ANSWER, YES, NO) + CODES  # 261 tokens, in id order
ANSWER_CUE = "Answer"
# The type word each question type carries in the prompt.
TYPE_WORDS = {"choice": "choice", "score": "score", "truth": "noul", "coord": "coord"}

IMAGE_PAD = "<|image_pad|>"
VISION_START = "<|vision_start|>"
VISION_END = "<|vision_end|>"

# Segment kinds
SPECIAL, TEXT, STATE_TEXT, IMAGE = "special", "text", "state", "image"


class TemplateError(ValueError):
    pass


# ---------------------------------------------------------------------------
# Codes
# ---------------------------------------------------------------------------


def question_hash(seed: str, question: dict, pass_index: int = 0) -> bytes:
    """SHA-256 of the question alone: the seed, the pass, the type word, the instructions and the options
    (or levels, or truth criteria). Every part is length-prefixed, so two questions cannot collide by concatenation."""
    kind = question["type"]
    parts = [seed, str(int(pass_index)), TYPE_WORDS[kind], question["instructions"]]
    if kind == "choice":
        for key, desc in zip(question["keys"], question["descs"]):
            parts += [key, desc]
    elif kind == "score":
        parts += list(question["levels"])
    else:
        for name, text in criteria_lines(question):
            parts += [name, text]
    h = hashlib.sha256()
    for part in parts:
        raw = part.encode("utf-8")
        h.update(len(raw).to_bytes(8, "big"))
        h.update(raw)
    return h.digest()


class HashStream:
    """Deterministic integers from a seed: SHA-256 in counter mode, 4 bytes at a time."""

    def __init__(self, seed: bytes):
        self._seed, self._counter, self._buf = seed, 0, b""

    def _u32(self) -> int:
        if len(self._buf) < 4:
            self._buf += hashlib.sha256(self._seed + self._counter.to_bytes(8, "big")).digest()
            self._counter += 1
        value = int.from_bytes(self._buf[:4], "big")
        self._buf = self._buf[4:]
        return value

    def below(self, n: int) -> int:
        """A uniform integer in [0, n), by rejection (no modulo bias)."""
        limit = (1 << 32) - ((1 << 32) % n)
        while True:
            value = self._u32()
            if value < limit:
                return value % n

    def sample(self, population: int, k: int) -> list[int]:
        """k distinct integers from range(population): a partial Fisher-Yates shuffle."""
        pool = list(range(population))
        for i in range(k):
            j = i + self.below(population - i)
            pool[i], pool[j] = pool[j], pool[i]
        return pool[:k]


def n_classes(question: dict) -> int:
    if question["type"] == "choice":
        return len(question["keys"])
    if question["type"] == "score":
        return len(question["levels"])
    if question["type"] == "coord":
        return int(question["classes"])
    return 2


def assign_codes(seed: str, question: dict) -> list[int]:
    """The code index of every option (or level) of a choice or score question, in the order given."""
    k = n_classes(question)
    if k > N_CODES:
        raise TemplateError(f"{k} options, the model has {N_CODES} codes")
    return HashStream(question_hash(seed, question)).sample(N_CODES, k)


# ---------------------------------------------------------------------------
# Segments
# ---------------------------------------------------------------------------


def criteria_lines(question: dict) -> list[tuple[str, str]]:
    """A truth question's criteria as ("yes", text) and ("no", text), each only if given."""
    criteria = question.get("criteria") or {}
    return [(label, criteria[name]) for name, label in (("true", "yes"), ("false", "no")) if criteria.get(name)]


def image_label(number: int, image_id: str) -> str:
    """The words in front of the request's number-th picture (counted from 1)."""
    return f"Picture {int(number)} ({image_id}): "


def build_segments(question: dict, state: str, image_ids: Sequence[str], codes: Sequence[int]) -> list[tuple[str, str]]:
    """The prompt of one question as (kind, value) segments, state first."""
    segments: list[tuple[str, str]] = [(SPECIAL, STATE), (TEXT, "\n")]
    for number, image_id in enumerate(image_ids, 1):
        segments += [(TEXT, image_label(number, image_id)), (IMAGE, image_id), (TEXT, "\n")]
    segments += [(STATE_TEXT, state), (TEXT, "\n")]

    kind = question["type"]
    text = f"\ntype: {TYPE_WORDS[kind]}\n{question['instructions']}\n"
    if kind == "truth":
        text += "".join(f"{label}: {value}\n" for label, value in criteria_lines(question))
    segments += [(SPECIAL, QUESTION), (TEXT, text)]
    if kind == "choice":
        for key, desc, code in zip(question["keys"], question["descs"], codes):
            segments += [(SPECIAL, CODES[code]), (TEXT, f" {key}: {desc}\n" if desc else f" {key}\n")]
    elif kind == "score":
        for index, (level, code) in enumerate(zip(question["levels"], codes)):
            segments += [(SPECIAL, CODES[code]), (TEXT, f" {index}: {level}\n")]
    segments += [(TEXT, ANSWER_CUE), (SPECIAL, ANSWER)]
    return segments


# ---------------------------------------------------------------------------
# Encoder
# ---------------------------------------------------------------------------


@dataclass
class Encoded:
    input_ids: list[int]
    class_ids: list[int]  # the token whose readout row scores each class, in class order
    state_tokens: int  # tokens of the state block (with its pictures), shared by every question of a request
    image_slots: list[tuple[str, int, int]] = field(default_factory=list)  # (image id, first pad position, count)


class Template:
    """The template bound to the model's tokenizer."""

    def __init__(self, tokenizer: Any, seed: str):
        self.tokenizer = tokenizer
        self.seed = seed
        ids = [tokenizer.convert_tokens_to_ids(token) for token in ADDED_TOKENS]
        first = ids[0]
        if first is None or ids != list(range(first, first + len(ADDED_TOKENS))):
            raise TemplateError("the tokenizer does not hold the model's added tokens as one id range")
        self.first_id = first
        self.special_ids = dict(zip(ADDED_TOKENS, ids))
        self.image_pad, self.vision_start, self.vision_end = (tokenizer.convert_tokens_to_ids(t) for t in (IMAGE_PAD, VISION_START, VISION_END))
        unsplit = _unsplit_added_tokens(tokenizer)
        self._guard = re.compile("|".join(re.escape(t) for t in sorted(unsplit, key=len, reverse=True))) if unsplit else None
        self._guarded_ids = frozenset(unsplit.values())

    def class_ids(self, question: dict, codes: Sequence[int]) -> list[int]:
        if question["type"] == "truth":
            return [self.special_ids[YES], self.special_ids[NO]]
        return [self.first_id + 5 + c for c in codes]

    def encode(self, question: dict, state: str, image_ids: Sequence[str] = (), image_tokens: dict[str, int] | None = None) -> Encoded:
        if question["type"] in ("choice", "score"):
            codes = assign_codes(self.seed, question)
        else:  # a coordinate row's codes are fixed, a truth question has none
            codes = list(range(n_classes(question))) if question["type"] == "coord" else []
        segments = build_segments(question, state, image_ids, codes)
        table = self._tokenize([value for kind, value in segments if kind in (TEXT, STATE_TEXT)])
        ids: list[int] = []
        slots: list[tuple[str, int, int]] = []
        state_tokens = 0
        for kind, value in segments:
            if kind == SPECIAL:
                ids.append(self.special_ids[value])
            elif kind in (TEXT, STATE_TEXT):
                ids.extend(table[value])
                if kind == STATE_TEXT:
                    state_tokens = len(ids) + len(table["\n"])  # the state block ends with the line break after the state
            else:  # a picture: its placeholders, as many as the image processor gives it
                count = int((image_tokens or {})[value])
                ids.append(self.vision_start)
                slots.append((value, len(ids), count))
                ids.extend([self.image_pad] * count)
                ids.append(self.vision_end)
        return Encoded(ids, self.class_ids(question, codes), state_tokens, slots)

    def _tokenize(self, texts: Sequence[str]) -> dict[str, list[int]]:
        """Plain-text ids of every distinct text, with special-token strings kept as text."""
        unique = [t for t in dict.fromkeys(texts) if t]
        pieces_of = {t: _break_tokens(t, self._guard) for t in unique if self._guard is not None and self._guard.search(t)}
        flat = list(dict.fromkeys(p for t in unique for p in pieces_of.get(t, (t,))))
        encoded = {}
        if flat:
            out = self.tokenizer(flat, add_special_tokens=False, split_special_tokens=True, return_attention_mask=False)["input_ids"]
            encoded = dict(zip(flat, (list(x) for x in out)))
        table: dict[str, list[int]] = {"": []}
        for text in unique:
            if text in pieces_of:
                ids = [i for piece in pieces_of[text] for i in encoded[piece]]
                if not self._guarded_ids.isdisjoint(ids):
                    raise TemplateError("request text produced a control token")
                table[text] = ids
            else:
                table[text] = encoded[text]
        return table


def _unsplit_added_tokens(tokenizer: Any) -> dict[str, int]:
    """Added tokens that `split_special_tokens=True` still turns into one id (found by trying each)."""
    out = {}
    for token, index in dict(tokenizer.get_added_vocab()).items():
        if len(token) >= 2 and tokenizer.encode(token, add_special_tokens=False, split_special_tokens=True) == [index]:
            out[token] = int(index)
    return out


def _break_tokens(text: str, pattern: "re.Pattern[str]") -> list[str]:
    """Cut `text` after the first character of every added-token string, so none survives whole."""
    pieces, start = [], 0
    for match in pattern.finditer(text):
        pieces.append(text[start : match.start() + 1])
        start = match.start() + 1
    pieces.append(text[start:])
    out = []
    for piece in pieces:
        out.extend(_break_tokens(piece, pattern) if pattern.search(piece) else [piece])
    return [piece for piece in out if piece]