"""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 ``) 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]