| """The prompt and the readout. A System One decision is a single forward pass that stops at the answer |
| slot: everything the model sees is built here, and the answer is a softmax over its options' tokens. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import json |
| import math |
| import re |
| from dataclasses import dataclass |
| from functools import cached_property, lru_cache |
| from typing import Any, Mapping, Sequence |
|
|
| import numpy as np |
|
|
| from .media import SPOKEN, Media, as_media |
|
|
| |
| SYSTEMS: dict[str, str | None] = {"none": None, "d1": "You are d1, a decision model trained by Liquid AI."} |
| DEFAULT_SYSTEM = "none" |
|
|
| IM_START = "<|im_start|>" |
| IM_END = "<|im_end|>" |
|
|
|
|
| |
| |
| |
|
|
|
|
| @dataclass(frozen=True) |
| class Choice: |
| instructions: str | Media |
| criteria: Mapping[str, str | None] |
| type: str = "choice" |
| multiple: bool = False |
|
|
|
|
| @dataclass(frozen=True) |
| class Noul: |
| instructions: str | Media |
| criteria: Mapping[str, Any] | None = None |
| type: str = "noul" |
|
|
|
|
| @dataclass(frozen=True) |
| class Score: |
| instructions: str | Media |
| criteria: Sequence[str] |
| type: str = "score" |
|
|
|
|
| @dataclass(frozen=True) |
| class Read: |
| """A span read from the state: text, or a number or a date written there; `unit`, a number's.""" |
|
|
| instructions: str | Media |
| type: str = "text" |
| unit: str | None = None |
|
|
|
|
| @dataclass(frozen=True) |
| class Estimate: |
| """A number estimated on a grid over the declared range: `step` apart, at most 128 values (a finer step |
| is read on 128).""" |
|
|
| instructions: str | Media |
| low: float |
| high: float |
| step: float | None = None |
| unit: str | None = None |
| type: str = "number" |
|
|
| @cached_property |
| def grid(self): |
| from .grid import Grid |
|
|
| return Grid(self.low, self.high, 128 if self.step is None else min(128, round((self.high - self.low) / self.step) + 1)) |
|
|
|
|
| @dataclass(frozen=True) |
| class Box: |
| """Where something is in the picture: its four edges, or nothing.""" |
|
|
| instructions: str | Media |
| type: str = "box" |
|
|
|
|
| @dataclass(frozen=True) |
| class Boxes: |
| """Every instance of something in the picture, in reading order, none if absent: `max` blocks read, else |
| `box.INSTANCES`. A readout choice: the question's text never says it.""" |
|
|
| instructions: str | Media |
| type: str = "boxes" |
| max: int | None = None |
|
|
| @property |
| def blocks(self) -> int: |
| from .box import INSTANCES |
|
|
| return INSTANCES if self.max is None else self.max |
|
|
|
|
| @dataclass(frozen=True) |
| class Position: |
| name: str |
|
|
|
|
| @dataclass(frozen=True) |
| class Number: |
| """A slot over a grid's values as the model writes them: value i times `scale`, an integer, is one token of the |
| tokenizer's own (`number_ids`). A checkpoint that records `d1_numbers` reads its grids so.""" |
|
|
| instructions: str |
| grid: Any |
| scale: int = 1 |
| type: str = "grid" |
|
|
|
|
| Question = Choice | Noul | Score | Read | Estimate | Box | Boxes |
|
|
|
|
| TYPED = ("text", "number", "date", "box", "boxes") |
|
|
|
|
| def typed(q: Question) -> bool: |
| """A typed answer, read at its own slots or state positions and answered in its own shape (api.py): |
| a read, an estimate, a box, a multiple choice.""" |
| return q.type in TYPED or isinstance(q, Choice) and q.multiple |
|
|
|
|
| def cardinality(q: Question) -> int: |
| if isinstance(q, (Estimate, Number)): |
| return q.grid.size |
| return 2 if isinstance(q, Noul) else len(q.criteria) |
|
|
|
|
| |
| |
| |
|
|
|
|
| PROMPTS = (1, 2) |
|
|
|
|
| def version(tokenizer) -> int: |
| """The prompt version the checkpoint behind `tokenizer` trained with (`PROMPTS`).""" |
| return getattr(tokenizer, "d1_prompt", 1) |
|
|
|
|
| def most_blocks(tokenizer) -> int: |
| """The most blocks a `boxes` question's `max` asks of this checkpoint: the 16 prompt 1 trained, else |
| `box.MAX_INSTANCES`.""" |
| from .box import INSTANCES, MAX_INSTANCES |
|
|
| return INSTANCES if version(tokenizer) == 1 else MAX_INSTANCES |
|
|
|
|
| LETTERS = tuple(chr(ord("A") + i) for i in range(26)) |
|
|
|
|
| def option_codes(labels: Sequence[str], scheme: int = 1) -> list[str]: |
| """Native letters when the labels already are letters, else A..Z, else the `scheme`'s codes: 00.. (1), A..Z then |
| AA, AB.. (2); one rule for every cardinality.""" |
| labs = [str(x).strip() for x in labels] |
| if labs and all(len(k) == 1 and k.isalpha() for k in labs): |
| return labs |
| if len(labs) <= 26: |
| return list(LETTERS[: len(labs)]) |
| if scheme == 2: |
| return (list(LETTERS) + [a + b for a in LETTERS for b in LETTERS])[: len(labs)] |
| return [f"{i:02d}" for i in range(len(labs))] |
|
|
|
|
| _FALLBACK_POOL = ( |
| [chr(c) for c in range(ord("A"), ord("Z") + 1)] |
| + [f"{i:02d}" for i in range(100)] |
| + [chr(c) for c in range(ord("a"), ord("z") + 1)] |
| + [f"#{i}" for i in range(200)] |
| |
| |
| + [chr(a) + chr(b) for a in range(ord("A"), ord("Z") + 1) for b in range(ord("A"), ord("Z") + 1)] |
| ) |
|
|
|
|
| @lru_cache(maxsize=4096) |
| def _aliases_cached(tokenizer_key: int, codes: tuple[str, ...]) -> tuple[tuple[str, int], ...]: |
| tokenizer = _TOKENIZERS[tokenizer_key] |
| used: set[int] = set() |
| out: list[tuple[str, int]] = [] |
|
|
| def take(raw: str) -> bool: |
| enc = tokenizer.encode(raw, add_special_tokens=False) |
| if len(enc) != 1 or enc[0] in used: |
| return False |
| out.append((raw, enc[0])) |
| used.add(enc[0]) |
| return True |
|
|
| for code in codes: |
| if take(code): |
| continue |
| if not any(take(raw) for raw in _FALLBACK_POOL): |
| raise RuntimeError(f"no single-token alias left for {len(codes)} options") |
| return tuple(out) |
|
|
|
|
| _TOKENIZERS: dict[int, object] = {} |
|
|
|
|
| def aliases(tokenizer, labels: Sequence[str]) -> list[tuple[str, int]]: |
| """Every label's distinct single-token code: [(code, token_id)]. Memoised on the codes, which are positional unless |
| the labels are letters, so every ordering of one option list shares an entry.""" |
| _TOKENIZERS.setdefault(id(tokenizer), tokenizer) |
| codes = tuple(option_codes(labels, getattr(tokenizer, "d1_codes", 1))) |
| return list(_aliases_cached(id(tokenizer), codes)) |
|
|
|
|
| @lru_cache(maxsize=8192) |
| def _ids_cached(tokenizer_key: int, texts: tuple[str, ...]) -> tuple[int, ...]: |
| tokenizer = _TOKENIZERS[tokenizer_key] |
| out, seen = [], set() |
| for t in texts: |
| enc = tokenizer.encode(t, add_special_tokens=False) |
| if len(enc) == 1 and enc[0] not in seen: |
| out.append(enc[0]) |
| seen.add(enc[0]) |
| return tuple(out) |
|
|
|
|
| def _ids(tokenizer, texts: Sequence[str]) -> list[int]: |
| _TOKENIZERS.setdefault(id(tokenizer), tokenizer) |
| return list(_ids_cached(id(tokenizer), tuple(texts))) |
|
|
|
|
| QUESTION_TYPES = ("noul", "choice", "score", "text", "number", "date", "box", "boxes") |
| CHOICE = '"criteria" names the options: {"billing": "Charges and refunds", "shipping": "Delivery"}' |
| SCORE = '"criteria" lists 2 to 10 levels, lowest first: ["Can wait", "Today", "Now"]' |
|
|
|
|
| def check_question(q: Any, tokenizer=None) -> None: |
| """A caller's question in the API's JSON, refused with what to change; a class passes. `tokenizer` carries the |
| checkpoint's prompt version, which bounds a `boxes` question's `max` (`most_blocks`; prompt 1 without it).""" |
| if isinstance(q, (Choice, Noul, Score, Read, Estimate, Box, Boxes)): |
| return |
| if not isinstance(q, Mapping): |
| raise ValueError(f'a question is a dict such as {{"type": "noul", "instructions": "Is it urgent?"}}, not ' |
| f"{type(q).__name__}") |
| kind, said, criteria = q.get("type"), q.get("instructions"), q.get("criteria") |
| if kind not in QUESTION_TYPES: |
| raise ValueError(f'"type" is one of {", ".join(QUESTION_TYPES)}, not {kind!r}') |
| if not (isinstance(said, str) and said.strip() or isinstance(said, Mapping)): |
| raise ValueError('"instructions" is the question as text (or a spoken one, {"type": "audio", "path": ...})') |
| if kind == "choice" and not (isinstance(criteria, Mapping) and criteria |
| and all(isinstance(k, str) and (v is None or isinstance(v, str)) |
| for k, v in criteria.items())): |
| raise ValueError(f"a choice's {CHOICE}" + ("; a list of levels is a score" if isinstance(criteria, list) else "")) |
| if kind == "score" and not (isinstance(criteria, (list, tuple)) and 2 <= len(criteria) <= 10 |
| and all(isinstance(c, str) for c in criteria)): |
| raise ValueError(f"a score's {SCORE}") |
| if kind == "noul" and criteria is not None and not isinstance(criteria, Mapping): |
| raise ValueError('a noul\'s "criteria", when given, is {"true": "what yes means", "false": "what no means"}') |
| if q.get("multiple") and kind != "choice": |
| raise ValueError('"multiple" belongs to a choice: {"type": "choice", "multiple": true, "criteria": {...}}') |
| if "max" in q and not (kind == "boxes" and isinstance(q["max"], int) and not isinstance(q["max"], bool) |
| and 1 <= q["max"] <= most_blocks(tokenizer)): |
| raise ValueError(f'"max" belongs to boxes, a whole number from 1 to {most_blocks(tokenizer)}: ' |
| '{"type": "boxes", "instructions": "Every cup.", "max": 4}') |
| if "step" in q and "range" not in q: |
| raise ValueError('"step" needs a "range": {"type": "number", "range": [0, 10], "step": 1}') |
| if "range" in q: |
| bounds = q["range"] |
| if kind != "number" or not (isinstance(bounds, (list, tuple)) and len(bounds) == 2 |
| and all(isinstance(b, (int, float)) and not isinstance(b, bool) for b in bounds)): |
| raise ValueError('an estimate is {"type": "number", "range": [low, high]}, with an optional "step"') |
|
|
|
|
| def as_question(q: Mapping | Question) -> Question: |
| """A question in the API's JSON (`type`, `instructions`, `criteria`; `multiple`, `range`, `step`, `unit` for the |
| typed answers; `instructions` a media dict for a spoken question), as one of the classes above; a class passes |
| through.""" |
| if not isinstance(q, Mapping): |
| return q |
| if isinstance(q.get("instructions"), Mapping): |
| q = {**q, "instructions": as_media(q["instructions"])} |
| kind = q.get("type", "choice") |
| if kind == "number" and "range" in q: |
| low, high = q["range"] |
| step = q.get("step") |
| if not low < high or step is not None and not (step > 0 and math.isclose((high - low) / step, round((high - low) / step))): |
| raise ValueError(f"an estimate's range runs low to high and its step, above 0, divides it: not " |
| f"{[low, high]} in steps of {step}") |
| return Estimate(q["instructions"], low, high, step, q.get("unit")) |
| if kind in ("text", "number", "date"): |
| return Read(q["instructions"], kind, q.get("unit")) |
| if kind == "box": |
| return Box(q["instructions"]) |
| if kind == "boxes": |
| return Boxes(q["instructions"], max=q.get("max")) |
| if kind == "noul": |
| return Noul(q["instructions"], q.get("criteria")) |
| if kind == "score": |
| return Score(q["instructions"], list(q["criteria"])) |
| return Choice(q["instructions"], q["criteria"], multiple=bool(q.get("multiple", False))) |
|
|
|
|
| def question_slots(q: Question, tokenizer=None) -> tuple[Question, ...]: |
| """Each answer position's readout, in question order; a grid's through the tokenizer's numbers where the |
| checkpoint reads them so (`grid_slot`).""" |
| if isinstance(q, Read): |
| return Position("start"), Position("end") |
| if isinstance(q, Estimate): |
| return (grid_slot(tokenizer, q.instructions, q.grid),) |
| if isinstance(q, Box): |
| from .box import EDGES, edge_grid |
|
|
| grid, scale = edge_grid(tokenizer) |
| return (Noul(q.instructions), *(grid_slot(tokenizer, q.instructions, grid, scale) for _ in EDGES)) |
| if isinstance(q, Boxes): |
| return question_slots(Box(q.instructions), tokenizer) * q.blocks |
| if isinstance(q, Choice) and q.multiple: |
| if not q.criteria: |
| raise ValueError("choice with no options") |
| return tuple(Noul(q.instructions) for _ in q.criteria) |
| return (q,) |
|
|
|
|
| YES_FORMS = ("yes", "Yes", "YES") |
| NO_FORMS = ("no", "No", "NO") |
|
|
|
|
| def readout_ids(tokenizer, q: Question) -> list[list[int]]: |
| """Token ids to score, one group per option, max-pooled: the answer is a softmax over these and |
| nothing else.""" |
| if isinstance(q, Estimate): |
| return readout_ids(tokenizer, question_slots(q, tokenizer)[0]) |
| if isinstance(q, Number): |
| return [list(group) for group in number_ids(tokenizer, q.grid, q.scale)] |
| if isinstance(q, (Read, Position)): |
| raise ValueError("a pointer reads state positions, not vocabulary codes") |
| if isinstance(q, Choice) and q.multiple: |
| raise ValueError("a multiple choice reads one yes/no distribution per option") |
| if isinstance(q, Noul): |
| yes, no = _ids(tokenizer, YES_FORMS), _ids(tokenizer, NO_FORMS) |
| if not yes or not no: |
| raise RuntimeError("tokenizer has no single-token yes/no") |
| return [yes, no] |
| if isinstance(q, Score): |
| groups = [_ids(tokenizer, [str(i)]) for i in range(len(q.criteria))] |
| if any(not g for g in groups): |
| raise RuntimeError(f"score with {len(q.criteria)} levels needs single-token digits; the primitive is defined " |
| "for 2 to 10") |
| return groups |
| groups = [] |
| for code, tid in aliases(tokenizer, list(q.criteria.keys())): |
| extra = _ids(tokenizer, [f" {code}"]) |
| groups.append([tid] + [i for i in extra if i != tid]) |
| if not groups: |
| raise RuntimeError("choice with no options") |
| return groups |
|
|
|
|
| def readout(tokenizer, q: Question, logz, calibration=None) -> list[float]: |
| """Option probabilities from the log-probabilities at the answer slot, `logz` (a vocabulary tensor, or a dict |
| holding at least the option tokens). Each option scores its best form.""" |
| if not hasattr(logz, "tolist"): |
| return probabilities(q, [max(float(logz[i]) for i in g) for g in readout_ids(tokenizer, q)], calibration) |
| tokens, forms, _ = readout_table(tokenizer, q) |
| values = np.asarray(logz[list(tokens)].tolist()) |
| return probabilities(q, values[forms].max(-1).tolist(), calibration) |
|
|
|
|
| def readouts(tokenizer, qs: Sequence[Question], logz, calibrations: Sequence) -> list[list[float]]: |
| """`readout` for each row of `logz` ([slots, vocabulary], a torch tensor) and its question, as one gather for |
| the rows sharing a readout table (a box's edges) and one transfer for all.""" |
| import torch |
|
|
| tables = [readout_table(tokenizer, q) for q in qs] |
| shared: dict[int, list[int]] = {} |
| for j, table in enumerate(tables): |
| shared.setdefault(id(table), []).append(j) |
| picked = [logz[rows][:, torch.from_numpy(tables[rows[0]][2]).to(logz.device)] for rows in shared.values()] |
| host = torch.cat([p.flatten() for p in picked]).cpu().numpy() if picked else None |
| out, at = [None] * len(qs), 0 |
| for rows, part in zip(shared.values(), picked): |
| values = host[at:at + part.numel()].reshape(part.shape) |
| at += part.numel() |
| scores = values[:, tables[rows[0]][1]].max(-1).astype(np.float64) |
| scores /= np.array([1.0 if calibrations[j] is None else calibrations[j].temperature(qs[j]) for j in rows])[:, None] |
| exps = np.exp(scores - scores.max(-1, keepdims=True)) |
| for j, p in zip(rows, (exps / exps.sum(-1, keepdims=True)).tolist()): |
| out[j] = p |
| return out |
|
|
|
|
| def probabilities(q: Question, scores: list[float], calibration=None) -> list[float]: |
| """Each option's best score, through the calibration's temperature, as a softmax.""" |
| if calibration is not None: |
| scores = calibration.apply(q, scores) |
| m = max(scores) |
| exps = [math.exp(s - m) for s in scores] |
| total = sum(exps) |
| return [e / total for e in exps] |
|
|
|
|
| def readout_table(tokenizer, q: Question) -> tuple[tuple[int, ...], np.ndarray, np.ndarray]: |
| """`readout_ids` as one gather: every token once, each group's forms as positions among them (a short group |
| repeats its first form, which leaves its best as it is), and the tokens as an array. Cached: a box edge scores |
| 1,000 numbers.""" |
| if isinstance(q, Number): |
| _TOKENIZERS.setdefault(id(tokenizer), tokenizer) |
| return _number_table(id(tokenizer), q.grid.low, q.grid.high, q.grid.size, q.scale) |
| return _table(tuple(map(tuple, readout_ids(tokenizer, q)))) |
|
|
|
|
| @lru_cache(maxsize=4096) |
| def _table(groups: tuple[tuple[int, ...], ...]) -> tuple[tuple[int, ...], np.ndarray, np.ndarray]: |
| tokens = tuple(dict.fromkeys(i for g in groups for i in g)) |
| at = {t: k for k, t in enumerate(tokens)} |
| width = max(map(len, groups)) |
| return tokens, np.array([[at[i] for i in g] + [at[g[0]]] * (width - len(g)) for g in groups]), np.array(tokens) |
|
|
|
|
| @lru_cache(maxsize=256) |
| def _number_table(tokenizer_key: int, low: float, high: float, size: int, scale: int): |
| from .grid import Grid |
|
|
| return _table(tuple(map(tuple, readout_ids(_TOKENIZERS[tokenizer_key], Number("", Grid(low, high, size), scale))))) |
|
|
|
|
| |
| |
| |
|
|
|
|
| |
| |
| |
| DEFAULT_STATE_STYLE = "json_only" |
|
|
|
|
| def _is_scalar(v: Any) -> bool: |
| return v is None or isinstance(v, (str, int, float, bool)) and "\n" not in str(v) |
|
|
|
|
| def _sections(obj: Any, path: str, out: list[str]) -> None: |
| """Flatten a nested state into labelled blocks, keeping the real newlines `json.dumps` would escape.""" |
| head = f"[{path}]\n" if path else "" |
| if isinstance(obj, dict): |
| scalars = [(k, v) for k, v in obj.items() if _is_scalar(v)] |
| rest = [(k, v) for k, v in obj.items() if not _is_scalar(v)] |
| if scalars: |
| body = "\n".join(f"{k}: {'' if v is None else v}" for k, v in scalars) |
| out.append(f"{head}{body}") |
| for k, v in rest: |
| _sections(v, f"{path}.{k}" if path else str(k), out) |
| return |
| if isinstance(obj, (list, tuple)): |
| if obj and all(_is_scalar(v) for v in obj): |
| body = "\n".join(f"- {'' if v is None else v}" for v in obj) |
| out.append(f"{head}{body}") |
| return |
| for i, v in enumerate(obj, start=1): |
| _sections(v, f"{path} {i}/{len(obj)}" if path else f"{i}/{len(obj)}", out) |
| return |
| out.append(f"{head}{'' if obj is None else obj}") |
|
|
|
|
| def render_state(state: Any) -> str: |
| if isinstance(state, str): |
| return state |
| out: list[str] = [] |
| _sections(state, "", out) |
| return "\n\n".join(out) |
|
|
|
|
| def state_text(state: Any, style: str = DEFAULT_STATE_STYLE) -> str: |
| """The state as the prompt renders it, without the blank line after it: what a read's span offsets index.""" |
| return state_block(state, style)[:-2] |
|
|
|
|
| def state_block(state: Any, style: str = DEFAULT_STATE_STYLE) -> str: |
| """The state as the block before `QUESTION:`. The `_only` styles (the default) write every state as the object |
| it is; the others give three shapes a shortcut (`Message:`, `Passage:` / `Asked:`, a lone question's text).""" |
| if isinstance(state, dict) and not style.endswith("_only"): |
| keys = set(state.keys()) |
| if keys == {"text"}: |
| return f"Message: {state['text']}\n\n" |
| if {"passage", "question"} <= keys and len(keys) == 2: |
| return f"Passage: {state['passage']}\n\nAsked: {state['question']}\n\n" |
| if keys == {"question"}: |
| return f"{state['question']}\n\n" |
| if style.startswith("json"): |
| if isinstance(state, str): |
| return f"{state}\n\n" |
| return json.dumps(state, ensure_ascii=False, indent=2) + "\n\n" |
| return f"{render_state(state)}\n\n" |
|
|
|
|
| |
| |
| |
| BREAK = "\u200b" |
|
|
|
|
| @lru_cache(maxsize=64) |
| def _specials_cached(tokenizer_key: int) -> re.Pattern | None: |
| tok = _TOKENIZERS[tokenizer_key] |
| tok = getattr(tok, "tokenizer", tok) |
| names = set(getattr(tok, "all_special_tokens", ())) |
| names |= {t.content for t in getattr(tok, "added_tokens_decoder", {}).values() if t.special} |
| names = sorted((n for n in names if len(n) > 1 and BREAK not in n), key=len, reverse=True) |
| return re.compile("|".join(map(re.escape, names))) if names else None |
|
|
|
|
| def plain(tokenizer, text: str) -> str: |
| """`text` as the caller's text: no special token of `tokenizer` in it once encoded. Text without one is returned |
| unchanged.""" |
| return _plain(tokenizer, text, False)[0] |
|
|
|
|
| def plain_offsets(tokenizer, text: str) -> tuple[str, list[int]]: |
| """The same escaped text, with every character boundary mapped to the original text.""" |
| return _plain(tokenizer, text, True) |
|
|
|
|
| def _plain(tokenizer, text: str, offsets: bool) -> tuple[str, list[int] | None]: |
| _TOKENIZERS.setdefault(id(tokenizer), tokenizer) |
| boundaries = list(range(len(text) + 1)) if offsets else None |
| pattern = _specials_cached(id(tokenizer)) |
| if pattern is None: |
| return text, boundaries |
| n = 1 |
| while n: |
| if offsets: |
| insertions = [m.start() + 1 for m in pattern.finditer(text)] |
| for at in reversed(insertions): |
| boundaries.insert(at, boundaries[at]) |
| text, n = pattern.subn(lambda m: m[0][0] + BREAK + m[0][1:], text) |
| return text, boundaries |
|
|
|
|
| _PLACEHOLDER = re.compile(r"^opt\d+$") |
|
|
|
|
| def _option_line(code: str, label: str, desc: str | None, style: str) -> str: |
| text = desc or label.replace("_", " ") |
| if style == "name_desc" and not _PLACEHOLDER.match(label) and label != text: |
| return f"{code} {label}: {text}" |
| return f"{code} {text}" |
|
|
|
|
| def said(q: Question) -> str: |
| """A question's instructions as its text writes them: their place (`media.SPOKEN`) when they are spoken.""" |
| return SPOKEN if isinstance(q.instructions, Media) else q.instructions |
|
|
|
|
| def edge_legend(tokenizer) -> str: |
| """How a box's edges are written: number tokens on a checkpoint that reads numbers, else codes.""" |
| from .box import edge_grid |
|
|
| grid, scale = edge_grid(tokenizer) |
| if scale > 1: |
| return ("Each edge is a number from 0 to 999, in thousandths of the picture: top and bottom from its top edge, " |
| "left and right from its left edge.") |
| codes = [code for code, _ in aliases(tokenizer, list(grid_question("", grid).criteria))] |
| return (f"Each edge is a code on a {grid.size}-value scale: {codes[0]} is 0 (top or left), {codes[-1]} is 1 " |
| "(bottom or right). Values are equally spaced.") |
|
|
|
|
| def question_block(tokenizer, q: Question, option_style: str = "desc") -> str: |
| if isinstance(q, Box): |
| one = "one number" if isinstance(question_slots(q, tokenizer)[-1], Number) else "one code" |
| return (f"{said(q)}\n\nPresence is yes or no. {edge_legend(tokenizer)} Reply with yes or no at present, " |
| f"then {one} at top, bottom, left and right.") |
| if isinstance(q, Boxes): |
| from .box import INSTANCES |
|
|
| most = f"{INSTANCES} at most, " if version(tokenizer) == 1 else "" |
| return (f"{said(q)}\n\nEvery one of them, {most}in reading order (left edge, then top). " |
| f"Present k is yes when there are at least k. {edge_legend(tokenizer)} Reply with yes or no at each " |
| "present, then its top, bottom, left and right.") |
| if isinstance(q, Estimate) and isinstance(question_slots(q, tokenizer)[0], Number): |
| unit = f" {q.unit}" if q.unit is not None else "" |
| return (f"{said(q)}\n\nEstimate from {q.grid.low:g} to {q.grid.high:g}{unit}, in steps of " |
| f"{q.grid.step:g}. Reply with the number only.") |
| if isinstance(q, Estimate): |
| bins = grid_question("", q.grid) |
| codes = ([str(i) for i in range(q.grid.size)] if isinstance(bins, Score) |
| else [code for code, _ in aliases(tokenizer, list(bins.criteria))]) |
| unit = f" {q.unit}" if q.unit is not None else "" |
| return (f"{said(q)}\n\nEstimate from {q.grid.low:g} to {q.grid.high:g}{unit} on " |
| f"{q.grid.size} equally spaced values, step {q.grid.step:g}: " |
| f"{codes[0]} is {q.grid.low:g}, {codes[-1]} is {q.grid.high:g}. Reply with the code only.") |
| if isinstance(q, Read): |
| unit = f" (unit: {q.unit})" if q.type == "number" and q.unit else "" |
| return f"{said(q)}\n\nRead the {q.type}{unit} from the state, or nothing if it is absent." |
| if isinstance(q, Choice): |
| labels = list(q.criteria.keys()) |
| if q.multiple: |
| lines = "\n".join(f"{json.dumps(name, ensure_ascii=False)}: {q.criteria[name] or name}" for name in labels) |
| return f"{said(q)}\n\nOptions:\n{lines}\n\nReply with yes or no for each option." |
| lines = "\n".join(_option_line(code, lab, q.criteria[lab], option_style) |
| for (code, _), lab in zip(aliases(tokenizer, labels), labels)) |
| return f"{said(q)}\n\nOptions:\n{lines}\n\nReply with the option code only." |
| if isinstance(q, Noul): |
| extra = "" |
| if q.criteria: |
| extra = f"\nYes: {q.criteria.get('true')}\nNo: {q.criteria.get('false')}" |
| return f"{said(q)}{extra}\n\nReply with yes or no only." |
| if isinstance(q, Score): |
| legend = "\n".join(f"{i} {name}" for i, name in enumerate(q.criteria)) |
| return f"{said(q)}\n\n{legend}\n\nReply with a single digit 0-{len(q.criteria) - 1} only." |
| raise TypeError(f"unknown question type {type(q)}") |
|
|
|
|
| def prefix_text(tokenizer, state: Any, bos: str = "", style: str = DEFAULT_STATE_STYLE, system: str = DEFAULT_SYSTEM, |
| media: str = "") -> str: |
| """Everything before the question, shared by all questions on one state: the media's markup (`media`, as |
| `media.Reader` writes it) at the head of the user turn, then the state. With no state (`None`) the question |
| follows the media directly. The state is the caller's text (`plain`).""" |
| text = SYSTEMS[system] |
| turn = "" if text is None else f"{IM_START}system\n{text}{IM_END}\n" |
| if state is None: |
| body = "" |
| else: |
| body = f"{plain(tokenizer, state_block(state, style))}\n{question_header(tokenizer, state)}" |
| return f"{bos}{turn}{IM_START}user\n{media}{body}" |
|
|
|
|
| def question_header(tokenizer, state: Any) -> str: |
| """What `prefix_text` ends with to open a question. `set` attention repeats it before each further question |
| packed after the state (`attention.pack_set`).""" |
| return "" if state is None else "QUESTION:\n" |
|
|
|
|
| |
| |
| DEFAULT_LEAD = "" |
| LEADS: dict[str, str] = {} |
|
|
|
|
| def default_lead(model_type: str | None, tokenizer=None) -> str: |
| """What a checkpoint's own template writes before a non-thinking answer: `LEADS` for the families measured, |
| else read off the tokenizer's chat template when one is given.""" |
| if model_type in LEADS or tokenizer is None: |
| return LEADS.get(model_type, DEFAULT_LEAD) |
| return template_lead(tokenizer) |
|
|
|
|
| def template_lead(tokenizer) -> str: |
| """What a ChatML template writes after the assistant header for a non-thinking reply. The prompt is ChatML |
| (`IM_START`, `IM_END`), so a model with another template is refused rather than mis-prompted.""" |
| text = tokenizer.apply_chat_template([{"role": "user", "content": "x"}], add_generation_prompt=True, |
| tokenize=False, enable_thinking=False) |
| head = f"{IM_START}assistant\n" |
| if head not in text: |
| raise ValueError("not a ChatML chat template: d1 renders ChatML prompts only") |
| return text[text.rindex(head) + len(head):] |
|
|
|
|
| def suffix_text(tokenizer, q: Question, lead: str = DEFAULT_LEAD, option_style: str = "desc") -> str: |
| """The question and the assistant header, up to the answer slot. The question is the caller's text (`plain`).""" |
| return suffix_layout(tokenizer, q, lead, option_style)[0] |
|
|
|
|
| def suffix_layout(tokenizer, q: Question, lead: str = DEFAULT_LEAD, |
| option_style: str = "desc") -> tuple[str, list[int]]: |
| """The suffix and each slot's token offset from its last token; old questions have one slot at zero.""" |
| body = question_block(tokenizer, q, option_style) |
| if isinstance(q, Read): |
| return marked_suffix(tokenizer, body, ["start", "end"], lead) |
| if isinstance(q, Box): |
| from .box import EDGES |
|
|
| return marked_suffix(tokenizer, body, ["present", *EDGES], lead) |
| if isinstance(q, Boxes): |
| from .box import EDGES, INSTANCES |
|
|
| |
| |
| return marked_suffix(tokenizer, body, [f"{name} {k}" for k in range(1, q.blocks + 1) |
| for name in ("present", *EDGES)], lead, |
| "\n" if q.blocks < INSTANCES or version(tokenizer) > 1 else "") |
| if isinstance(q, Choice) and q.multiple: |
| names = [json.dumps(name, ensure_ascii=False) for name in q.criteria] |
| return marked_suffix(tokenizer, body, names, lead) |
| return f"{plain(tokenizer, body)}{IM_END}\n{IM_START}assistant\n{lead}", [0] |
|
|
|
|
| def marked_suffix(tokenizer, body: str, names: Sequence[str], lead: str = DEFAULT_LEAD, |
| after: str = "") -> tuple[str, list[int]]: |
| """A question's named answer markers and their token offsets from its last token; `after`, whitespace the last |
| marker's colon token takes where the tokenizer writes the two as one token.""" |
| if not names or len(set(names)) != len(names) or any(not isinstance(n, str) or not n or "\n" in n for n in names): |
| raise ValueError("slot names must be distinct nonempty single lines") |
| text = f"{plain(tokenizer, body)}{IM_END}\n{IM_START}assistant\n{lead}" |
| ends = [] |
| for i, name in enumerate(names): |
| text += ("\n" if i else "") + plain(tokenizer, name) + ":" |
| ends.append(len(text)) |
| if after and tokenizer.decode(tokenizer.encode(text + after, add_special_tokens=False)[-1:]).rstrip().endswith(":"): |
| text += after |
| encoded = tokenizer(text, add_special_tokens=False, return_offsets_mapping=True) |
| ids = encoded["input_ids"] |
| slots = [] |
| for chars in ends: |
| matching = [(i, stop) for i, (start, stop) in enumerate(encoded["offset_mapping"]) if start <= chars - 1 < stop] |
| if len(matching) != 1: |
| raise ValueError("a slot marker must have an exact colon token position") |
| end, stop = matching[0] |
| |
| if text[chars:stop].strip() or not tokenizer.decode([ids[end]]).rstrip().endswith(":"): |
| raise ValueError("a slot marker token must end with its colon and optional whitespace") |
| slots.append(end - len(ids) + 1) |
| return text, slots |
|
|
|
|
| def number_ids(tokenizer, grid, scale: int = 1) -> tuple[tuple[int, ...], ...] | None: |
| """Each grid value times `scale` as the tokenizer writes it, plain and after a space, one token a form: None when |
| a value is no integer or takes more than one token (a tokenizer of digits).""" |
| _TOKENIZERS.setdefault(id(tokenizer), tokenizer) |
| return _number_ids_cached(id(tokenizer), grid.low, grid.high, grid.size, scale) |
|
|
|
|
| @lru_cache(maxsize=256) |
| def _number_ids_cached(tokenizer_key: int, low: float, high: float, size: int, scale: int): |
| from .grid import Grid |
|
|
| groups = [] |
| for value in Grid(low, high, size).values: |
| n = round(value * scale) |
| if not math.isclose(n, value * scale, abs_tol=1e-6): |
| return None |
| group = _ids_cached(tokenizer_key, (str(n), f" {n}")) |
| if not group or len(_TOKENIZERS[tokenizer_key].encode(str(n), add_special_tokens=False)) != 1: |
| return None |
| groups.append(group) |
| return tuple(groups) |
|
|
|
|
| def grid_slot(tokenizer, instructions: str, grid, scale: int = 1) -> Number | Choice | Score: |
| """A grid's slot: the tokenizer's own numbers where the checkpoint reads them so (`d1_numbers`) and every value is |
| one (`number_ids`), else the grid's codes (`grid_question`).""" |
| if getattr(tokenizer, "d1_numbers", False) and number_ids(tokenizer, grid, scale) is not None: |
| return Number(instructions, grid, scale) |
| return grid_question(instructions, grid) |
|
|
|
|
| def grid_question(instructions: str, grid) -> Choice | Score: |
| """The grid's ordered bins on the existing digit or option-code readout.""" |
| labels = [f"{v:g}" for v in grid.values] |
| return (Score(instructions, labels) if grid.size <= 10 |
| else Choice(instructions, {str(i): v for i, v in enumerate(labels)})) |
|
|
|
|
| def render(tokenizer, state: Any, q: Question, bos: str = "", lead: str = DEFAULT_LEAD, |
| style: str = DEFAULT_STATE_STYLE, system: str = DEFAULT_SYSTEM, option_style: str = "desc", |
| media: str = "") -> str: |
| return prefix_text(tokenizer, state, bos, style, system, media) + suffix_text(tokenizer, q, lead, option_style) |
|
|
