| """One-pass typed decisions on any backbone. |
| |
| The model sees the state and the question once, and the answer is read off the |
| logits at the answer slot. Nothing is decoded, so a schema violation is |
| impossible and `tokens generated per decision` is zero. |
| |
| - Questions are read at their answer slots, a whole prompt or packed after their state (`SlotReader._padded`): |
| every typed answer, and every question on a backbone with no faster path of its own (`SystemOne.run`). |
| - A request is `(state, questions, media)`: its pictures, audio clips and videos head the state, a spoken question's |
| clip sits in its own text (`media`). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import os |
| from typing import Any, Sequence |
|
|
| import torch |
|
|
| from .api import SystemOneApi |
| from .attention import SET_SIZE, TRAINED_SET_SIZE, both_ways, language_model, pack_set |
| from .box import EDGE_READS |
| from .media import Media, Reader, media_list, spoken |
| from .pointer import SPAN_CAPS, query_plan, spans, state_positions |
| from .prompt import ( |
| DEFAULT_STATE_STYLE, |
| DEFAULT_SYSTEM, |
| PROMPTS, |
| Box, |
| Boxes, |
| Choice, |
| Position, |
| Question, |
| Read, |
| default_lead, |
| prefix_text, |
| question_header, |
| question_slots, |
| readout, |
| readouts, |
| render, |
| suffix_layout, |
| suffix_text, |
| typed, |
| ) |
| from .slots import slot_logits, text_backbone |
|
|
| |
| |
| |
| PIXELS_SET = os.environ.get("D1_VISION_MAX_PIXELS") |
| VISION_MAX_PIXELS = int(PIXELS_SET or 1024 * 1024) |
|
|
|
|
| def group_size(n: int | None) -> int | None: |
| """A group size (`set_size`): a whole number from 1 (each question alone), or None for the engine's own.""" |
| if n is None: |
| return None |
| if isinstance(n, bool) or not isinstance(n, int): |
| raise ValueError(f"set_size must be a whole number, not {n!r}") |
| if n < 1: |
| raise ValueError(f"set_size must be at least 1, not {n}") |
| return n |
|
|
|
|
| def remote_ok(model_id) -> bool: |
| """Whether a checkpoint's own code may run: a local folder or a LiquidAI repo. |
| Any other Hub id loads with transformers' classes only, and one that needs its own code fails.""" |
| return os.path.isdir(str(model_id)) or str(model_id).startswith("LiquidAI/") |
|
|
|
|
| def auto_processor(path: str, revision: str | None = None): |
| """The media processor a checkpoint holds.""" |
| from transformers import AutoProcessor |
|
|
| return AutoProcessor.from_pretrained(path, revision=revision, trust_remote_code=remote_ok(path)) |
|
|
|
|
| def _under(config: dict | None, compiled): |
| """`compiled` run under inductor's `config` (`SystemOne.compile`).""" |
| if not config: |
| return compiled |
| import torch._inductor.config as inductor |
|
|
| def run(*args, **kwargs): |
| with inductor.patch(config): |
| return compiled(*args, **kwargs) |
|
|
| return run |
|
|
|
|
| class SlotReader: |
| """Prompts read at their answer slots, whatever runs the backbone: a request's sequences (`_sequences`), each |
| question's slots, and the answers read off them (`_padded`). An engine sets what `SystemOne.__init__` sets (the |
| checkpoint's record through `_serve`, its `processor` or `_load_processor`) and gives `_batches` and |
| `_slot_logits`.""" |
|
|
| def _serve(self, recorded, codes: int | None = None, max_pixels: int | None = None, attention: str | None = None, |
| set_size: int | None = None, numbers: bool | None = None, box_edges: str | None = None, |
| prompt: int | None = None) -> None: |
| """The option codes, number readout, picture cap, attention mode, set size, box-edge read and prompt version |
| `recorded` (a checkpoint's config) holds (`d1_codes`, `d1_numbers`, `d1_vision_max_pixels`, `d1_attention`, |
| `d1_set_size`, `d1_box_edges`: the median without it, `d1_prompt`: 1 without it), unless given. A recorded set |
| size is 2 to `TRAINED_SET_SIZE`, the default; one given is any whole number from 1 (each question alone): |
| larger groups read faster, in packs training never saw.""" |
| self.tokenizer.d1_codes = codes if codes is not None else getattr(recorded, "d1_codes", 1) |
| self.tokenizer.d1_numbers = bool(numbers if numbers is not None else getattr(recorded, "d1_numbers", False)) |
| self.tokenizer.d1_box_edges = box_edges or getattr(recorded, "d1_box_edges", None) or "median" |
| if self.tokenizer.d1_box_edges not in EDGE_READS: |
| raise ValueError(f"d1_box_edges is one of {EDGE_READS}, not {self.tokenizer.d1_box_edges!r}") |
| self.tokenizer.d1_prompt = prompt or getattr(recorded, "d1_prompt", 1) |
| if self.tokenizer.d1_prompt not in PROMPTS: |
| raise ValueError(f"d1_prompt is one of {PROMPTS}, not {self.tokenizer.d1_prompt!r}") |
| self.max_pixels = (VISION_MAX_PIXELS if PIXELS_SET else |
| max_pixels or getattr(recorded, "d1_vision_max_pixels", None) or VISION_MAX_PIXELS) |
| self.attention = attention or getattr(recorded, "d1_attention", None) or "causal" |
| self.set_size = group_size(getattr(recorded, "d1_set_size", SET_SIZE) if set_size is None else set_size) |
| if set_size is None and not 2 <= self.set_size <= TRAINED_SET_SIZE: |
| raise ValueError(f"d1_set_size must be between 2 and {TRAINED_SET_SIZE}") |
|
|
| def _readout(self, q: Question, logz, tempered: bool = True) -> list[float]: |
| """An option answer through the served temperature; a typed answer's slot raw (`tempered=False`): its own |
| table calibrates it (`calibration.json` `tables`, fitted on raw reads).""" |
| return readout(self.tokenizer, q, logz, self.calibration if tempered else None) |
|
|
| def _readouts(self, qs: Sequence[Question], logz, tempered: Sequence[bool]) -> list[list[float]]: |
| """`_readout` for each code slot of a batch, `logz` as `_slot_logits` gives them.""" |
| return [self._readout(q, z, t) for q, z, t in zip(qs, logz, tempered)] |
|
|
| @property |
| def media(self) -> Reader: |
| """The checkpoint's media reader, its processor loaded on first use.""" |
| if getattr(self, "_reader", None) is None or self._reader.processor is not self.processor: |
| if self.processor is None: |
| self.processor = self._load_processor() |
| self._reader = Reader(self.processor, self.max_pixels) |
| return self._reader |
|
|
| def _prefix(self, state: Any, media: Sequence[Media]) -> tuple[str, list[dict]]: |
| """`prompt.prefix_text` headed by `media`'s markup, and their processor outputs.""" |
| markup, parts = self.media.encode(media) if media else ("", []) |
| return prefix_text(self.tokenizer, state, self.bos, self.state_style, self.system, markup), parts |
|
|
| def _suffixes(self, qs: Sequence[Question]) -> tuple[list[str], list[list[int]], list[list[dict]]]: |
| """Each question's suffix (a spoken one's clip in place), its slots' offsets and its clip's outputs.""" |
| layouts = [suffix_layout(self.tokenizer, q, self.lead, self.option_style) for q in qs] |
| placed = [self.media.place(text, q) if spoken(q) else (text, []) for (text, _), q in zip(layouts, qs)] |
| return [text for text, _ in placed], [ends for _, ends in layouts], [parts for _, parts in placed] |
|
|
| def _sequences(self, requests: Sequence[tuple[Any, list[Question], Sequence[Media]]], *, |
| position_maps=None, set_size: int | None = None) -> list[tuple]: |
| """Every request's sequences, `(request, first question, ids, state length, each token's question or None, |
| answer slots, media inputs)`: under `set` up to `set_size` questions after their state (`attention.pack_set`), |
| else one whole prompt a question. A request's media head its state, a spoken question's clip sits in its |
| suffix, each as its processor expands it, so the text is the sequence. `position_maps` gets each read's state |
| positions; `set_size`, this call's group size (the engine's by default).""" |
| seqs = [] |
| for i, (state, qs, media) in enumerate(requests): |
| media = media_list(media) |
| if any(isinstance(q, (Box, Boxes)) for q in qs) and (n := sum(m.type == "image" for m in media)) != 1: |
| raise ValueError(f"a box or boxes question reads exactly one picture in `media`; this request has {n}") |
| prefix, state_parts = self._prefix(state, media) |
| suffixes, offsets, voices = self._suffixes(qs) |
| step = (self.set_size if set_size is None else set_size) if self.attention == "set" else 1 |
| for k in range(0, len(qs), step): |
| text = prefix + suffixes[k] |
| encoded = None |
| if position_maps is not None and any(isinstance(q, Read) for q in qs[k:k + step]): |
| encoded = self.tokenizer(text, add_special_tokens=False, return_offsets_mapping=True) |
| ids = self.tokenizer.encode(text, add_special_tokens=False) if encoded is None else encoded["input_ids"] |
| if self.attention == "set": |
| ids, state_len, question, ends = pack_set(self.tokenizer, prefix, suffixes[k:k + step], |
| question_header(self.tokenizer, state), |
| offsets[k:k + step], first_ids=ids) |
| else: |
| state_len, question = both_ways(self.attention, self.tokenizer, ids, suffixes[k]), None |
| ends = [[len(ids) - 1 + offset for offset in offsets[k]]] |
| parts = state_parts + [part for own in voices[k:k + step] for part in own] |
| inputs = self.media.inputs(parts, ids) if parts else {} |
| if encoded is not None: |
| position_maps[len(seqs)] = state_positions(self.tokenizer, state, prefix, text, self.state_style, |
| encoded=encoded) |
| seqs.append((i, k, ids, state_len, question, ends, inputs)) |
| return seqs |
|
|
| def _padded(self, requests: Sequence[tuple[Any, list[Question], Sequence]], |
| set_size: int | None = None) -> list[tuple[list[list[float]], int]]: |
| """`_sequences` in the engine's batches (`_batches`), each read at its answer slots as training reads them |
| (`_slot_logits`): option codes through the readout, reads through the pointer.""" |
| position_maps = {} |
| seqs = self._sequences(requests, position_maps=position_maps, set_size=set_size) |
| read = [0] * len(requests) |
| for i, _, ids, *_ in seqs: |
| read[i] += len(ids) |
|
|
| def asked(key: tuple[int, int]) -> Question: |
| return requests[key[0]][1][key[1]] |
|
|
| probs: dict[tuple[int, int], list[float]] = {} |
| for batch in self._batches(seqs): |
| slots = [] |
| for r, b in enumerate(batch): |
| i, first = seqs[b][:2] |
| for k, ends in enumerate(seqs[b][5]): |
| key = (i, first + k) |
| slots.extend((r, key, q, end) for q, end in zip(question_slots(asked(key), self.tokenizer), ends)) |
| reads = [j for j, s in enumerate(slots) if isinstance(s[2], Position)] |
| pointer = None |
| if reads: |
| masks = [torch.tensor(position_maps[batch[slots[j][0]]].pointer_mask) for j in reads] |
| pointer = query_plan([slots[j][0] for j in reads], masks, reads, len(slots)) |
| logz, positions = self._slot_logits(seqs, batch, slots, pointer) |
| if pointer is not None: |
| logits = positions.reshape(-1, 2, positions.shape[-1]) |
| capped = {} |
| for index, j in enumerate(reads[::2]): |
| capped.setdefault(SPAN_CAPS[asked(slots[j][1]).type], []).append((index, j)) |
| for cap, own in capped.items(): |
| at, found = spans(logits[[index for index, _ in own]], max_tokens=cap) |
| for (_, j), span, p in zip(own, at.tolist(), found.tolist()): |
| r, key = slots[j][:2] |
| probs[key] = {**position_maps[batch[r]].read(*span, max_tokens=cap), "found": p} |
| codes = pointer["code_slots"].tolist() if pointer else range(len(slots)) |
| scored = self._readouts([slots[j][2] for j in codes], logz, [not typed(asked(slots[j][1])) for j in codes]) |
| for j, p in zip(codes, scored): |
| key = slots[j][1] |
| q = asked(key) |
| if isinstance(q, Choice) and q.multiple: |
| probs.setdefault(key, []).append(p[0]) |
| elif isinstance(q, (Box, Boxes)): |
| probs.setdefault(key, []).append(p) |
| else: |
| probs[key] = p |
| return [([probs[(i, k)] for k in range(len(qs))], read[i]) for i, (_, qs, _) in enumerate(requests)] |
|
|
|
|
| class SystemOne(SlotReader, SystemOneApi): |
| """State in, calibrated distribution out, one forward pass.""" |
|
|
| def __init__( |
| self, |
| model=None, |
| tokenizer=None, |
| device: str | None = None, |
| calibration=None, |
| lead: str | None = None, |
| state_style: str = DEFAULT_STATE_STYLE, |
| system: str = DEFAULT_SYSTEM, |
| option_style: str = "desc", |
| codes: int | None = None, |
| max_pixels: int | None = None, |
| attention: str | None = None, |
| set_size: int | None = None, |
| numbers: bool | None = None, |
| box_edges: str | None = None, |
| prompt: int | None = None, |
| compile: bool = False, |
| token_budget: int = 65536, |
| revision: str | None = None, |
| ): |
| """`model` and `tokenizer`: a backbone already loaded (the Hub's class passes itself, and the `revision` its |
| processor loads at); it stays on its device unless `device` says otherwise. `codes` and `max_pixels` default |
| to what the checkpoint's config records (`d1_codes`, `d1_vision_max_pixels`), else scheme 1 and the cap; |
| `attention` likewise (`d1_attention`, else causal; `attention`).""" |
| if model is None or tokenizer is None or isinstance(model, str): |
| raise TypeError("SystemOne reads a loaded backbone: pass `model=` and `tokenizer=` (a checkpoint's own " |
| "class serves it: AutoModel.from_pretrained(repo, trust_remote_code=True).system_one)") |
| self.model, self.tokenizer = model, tokenizer |
| self.model_id, device = model.config._name_or_path, device or next(model.parameters()).device |
| self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu")) |
| if getattr(self.model, "hf_device_map", None): |
| self.device = self.model.get_input_embeddings().weight.device |
| else: |
| self.model.to(self.device) |
| self.model.eval() |
| if hasattr(self.model, "fuse_projections"): |
| self.model.fuse_projections() |
| bos = getattr(self.tokenizer, "bos_token", None) |
| self.bos = bos if isinstance(bos, str) else "" |
| self.lead = default_lead(model.config.model_type, tokenizer) if lead is None else lead |
| self.calibration, self.state_style, self.system, self.option_style = calibration, state_style, system, option_style |
| self._serve(self.model.config, codes, max_pixels, attention, set_size, numbers, box_edges, prompt) |
| self.token_budget = token_budget |
| self.processor, self.revision = None, revision |
| if compile: |
| self.compile() |
|
|
| def compile(self, mode: str | None = None) -> None: |
| """torch.compile what the passes run: the language model's forward, shapes dynamic; not its hooks, so |
| `state_attention`'s plan is made outside the graph. Also the media encoders the backbone names |
| (`compiled_parts`; 20 to 15 ms on a 384 px picture), under the inductor config it names (`compile_options`), |
| set while they run: dynamo reads `freezing` as it traces, where `torch.compile`'s own `options` do not reach |
| (d1-2B on an MI325X: 4.5 ms a question with `options`, 3.5 set). CUDA graphs measured no faster there, and |
| inductor's graph pool refuses the packed varlen calls.""" |
| if mode in ("reduce-overhead", "max-autotune") and self.attention != "causal": |
| raise ValueError(f"mode={mode!r} records CUDA graphs, which `{self.attention}` attention does not take: " |
| "model.compile() fuses without them") |
| body, parts = text_backbone(self.model), getattr(self.model, "compiled_parts", ()) |
| config = getattr(self.model, "compile_options", None) |
| for part in (language_model(self.model), *(getattr(body, name) for name in parts)): |
| part.forward = _under(config, torch.compile(part.forward, dynamic=True, mode=mode)) |
|
|
| |
|
|
| def render(self, state: Any, q: Question) -> str: |
| return render(self.tokenizer, state, q, self.bos, self.lead, self.state_style, self.system, self.option_style) |
|
|
| |
|
|
| @torch.inference_mode() |
| def run(self, requests: Sequence[tuple[Any, list[Question], Sequence[Media]]], |
| set_size: int | None = None) -> list[tuple[list[list[float]], int]]: |
| """Each `(state, questions, media)` request's probabilities and the tokens it read, at its answer slots |
| (`_padded`), every request in the same batches. `set_size`: this call's questions a group under `set` (the |
| engine's by default).""" |
| set_size = group_size(set_size) |
| requests = [(state, qs, media_list(media)) for state, qs, media in requests] |
| return self._padded(requests, set_size) |
|
|
| def _batches(self, seqs: list[tuple]) -> list[list[int]]: |
| """`_sequences` right-padded into batches of at most `token_budget` padded tokens, whatever their media: a |
| request's sequences in one batch (one forward) unless they alone pass the budget, requests of similar length |
| together. Every row is whole, its media given with it (`_rows`).""" |
| own: dict[int, list[int]] = {} |
| for j in sorted(range(len(seqs)), key=lambda j: len(seqs[j][2])): |
| own.setdefault(seqs[j][0], []).append(j) |
| batches: list[list[int]] = [[]] |
| for rows in sorted(own.values(), key=lambda rows: len(seqs[rows[-1]][2])): |
| if batches[-1] and len(seqs[rows[-1]][2]) * (len(batches[-1]) + len(rows)) > self.token_budget: |
| batches.append([]) |
| for j in rows: |
| if batches[-1] and len(seqs[j][2]) * (len(batches[-1]) + 1) > self.token_budget: |
| batches.append([]) |
| batches[-1].append(j) |
| return list(filter(None, batches)) |
|
|
| def _rows(self, seqs: list[tuple], batch: list[int]) -> tuple[torch.Tensor, list[int], dict]: |
| """A batch's right-padded ids, each row's length and its media inputs: every row whole, its own ids and its |
| media's inputs, joined in row order (`media.Reader.batch`) for the one forward.""" |
| pad = self.tokenizer.pad_token_id or 0 |
| lengths = [len(seqs[b][2]) for b in batch] |
| ids = torch.tensor([seqs[b][2] + [pad] * (max(lengths) - n) for b, n in zip(batch, lengths)], device=self.device) |
| inputs = [seqs[b][6] for b in batch] |
| return ids, lengths, self.media.batch(inputs, ids) if any(inputs) else {} |
|
|
| def _slot_logits(self, seqs: list[tuple], batch: list[int], slots: list[tuple], pointer: dict | None): |
| """A batch's code slots' log-probabilities over the vocabulary, and its read slots' position logits |
| (`slots.slot_logits`), in one forward (`_rows`).""" |
| ids, lengths, media = self._rows(seqs, batch) |
| width = ids.shape[1] |
| mask = (torch.arange(width, device=self.device) < torch.tensor(lengths, device=self.device)[:, None]).long() |
| state = None if self.attention == "causal" else torch.tensor([seqs[b][3] for b in batch]) |
| question = None if self.attention != "set" else torch.tensor( |
| [seqs[b][4] + [-1] * (width - len(seqs[b][4])) for b in batch]) |
| rows = slot_logits(self.model, ids, mask, torch.tensor([s[3] for s in slots]), state, question, |
| torch.tensor([s[0] for s in slots]), pointer=pointer, **self._on_device(media)) |
| codes = rows.codes.float() |
| return codes - torch.logsumexp(codes, dim=-1, keepdim=True), rows.positions |
|
|
| def _readouts(self, qs: Sequence[Question], logz: torch.Tensor, tempered: Sequence[bool]) -> list[list[float]]: |
| """Every code slot of a batch read at once (`prompt.readouts`): one transfer, not one a slot.""" |
| return readouts(self.tokenizer, qs, logz, [self.calibration if t else None for t in tempered]) |
|
|
| def tokens(self, state: Any, questions: Sequence[Question]) -> int: |
| """The longest prompt one of `questions` makes over a text state: its state's tokens and its own.""" |
| prefix, _ = self._prefix(state, ()) |
| return len(self.tokenizer.encode(prefix, add_special_tokens=False)) + max( |
| len(self.tokenizer.encode(suffix_text(self.tokenizer, q, self.lead, self.option_style), |
| add_special_tokens=False)) for q in questions) |
|
|
| |
|
|
| def _load_processor(self): |
| return auto_processor(self.model.config._name_or_path, self.revision) |
|
|
| def _on_device(self, inputs: dict) -> dict: |
| """The media inputs on the model's device, but what the backbone reads on the host (`host_inputs`).""" |
| host = getattr(self.model, "host_inputs", ()) |
| return {k: (v.to(self.device) if hasattr(v, "to") and k not in host else v) for k, v in inputs.items()} |
|
|
