| """Serialization of System One requests into Groot token structures. |
| |
| A request becomes: |
| |
| * state chunks — the rendered state is split into sentence/paragraph units, |
| tokenized, and greedily packed into leaves of at most |
| `chunk_tokens` tokens: `[CLS] ...unit tokens... [SEP]`. |
| Leaves are encoded in parallel and merged by the tree reducer. |
| * question sequence — one encoder sequence per question (early fusion): |
| `[CLS] "{type} question: {instructions}" [SEP] <first |
| early_state_tokens of the state> [SEP] [MASK] option_1 ... |
| [MASK] option_K`. Attention is block-restricted: the prefix |
| sees only the prefix, each option sees the prefix and itself |
| (never the other options). Every option block starts at the |
| SAME RoPE position, so the encoder is exactly |
| permutation-equivariant over options (list order cannot leak |
| into the answer, a documented Jev weakness) while every option |
| attends to the question and the state in all layers. The `[MASK]` state is the option's score slot. The full |
| state (any length) is still read through the tree memory and the |
| fusion layers. |
| |
| Score levels and noul outcomes carry a level embedding in the model so that |
| ordinal roles stay distinguishable; choice options carry none. |
| """ |
| from __future__ import annotations |
|
|
| import re |
| from dataclasses import dataclass |
| from typing import Any, Iterable, Sequence |
|
|
| from .config import QTYPE_IDS, GrootConfig |
| from .schema import ChoiceQuestion, NoulQuestion, Question, ScoreQuestion, render_content |
|
|
| _SENTENCE_RE = re.compile(r"(?<=[.!?;:])\s+|\n+") |
|
|
|
|
| @dataclass |
| class EncodedQuestion: |
| name: str |
| qtype: int |
| ids: list[int] |
| position_ids: list[int] |
| markers: list[int] |
| blocks: list[int] = None |
| tokens: int = 0 |
| truncated: bool = False |
|
|
| @property |
| def n_options(self) -> int: |
| return len(self.markers) |
|
|
|
|
| @dataclass |
| class EncodedState: |
| chunks: list[list[int]] |
| tokens: int |
| truncated: bool = False |
| pieces: list[list[int]] | None = None |
|
|
|
|
| @dataclass |
| class EncodedRequest: |
| state: EncodedState |
| questions: list[EncodedQuestion] |
|
|
| @property |
| def input_tokens(self) -> int: |
| return self.state.tokens + sum(q.tokens for q in self.questions) |
|
|
| @property |
| def output_tokens(self) -> int: |
| return sum(q.n_options for q in self.questions) |
|
|
|
|
| class Serializer: |
| """Turns requests into token structures using the backbone tokenizer.""" |
|
|
| def __init__(self, tokenizer: Any, config: GrootConfig): |
| self.tok = tokenizer |
| self.cfg = config |
| for name in ("cls_token_id", "sep_token_id", "mask_token_id", "pad_token_id"): |
| if getattr(tokenizer, name, None) is None: |
| raise ValueError(f"tokenizer must define {name}") |
| self.cls = tokenizer.cls_token_id |
| self.sep = tokenizer.sep_token_id |
| self.mask = tokenizer.mask_token_id |
| self.pad = tokenizer.pad_token_id |
| self._cache: dict[str, list[int]] = {} |
|
|
| |
| def encode(self, text: str) -> list[int]: |
| ids = self._cache.get(text) |
| if ids is None: |
| clean = text.replace(self.tok.mask_token, " ") if self.tok.mask_token else text |
| ids = self.tok(clean, add_special_tokens=False)["input_ids"] |
| if len(self._cache) < 65536: |
| self._cache[text] = ids |
| return list(ids) |
|
|
| |
| def units(self, text: str) -> list[str]: |
| parts = [p.strip() for p in _SENTENCE_RE.split(text) if p and p.strip()] |
| return parts or [text.strip() or " "] |
|
|
| def encode_state(self, state: Any) -> EncodedState: |
| text = render_content(state) |
| body = self.cfg.chunk_tokens - 2 |
| chunks: list[list[int]] = [] |
| current: list[int] = [] |
| total = 0 |
| for unit in self.units(text): |
| ids = self.encode(unit) |
| if not ids: |
| continue |
| pieces = [ids[i : i + body] for i in range(0, len(ids), body)] |
| for piece in pieces: |
| if current and len(current) + len(piece) > body: |
| chunks.append(current) |
| current = [] |
| current = current + piece |
| total += len(piece) |
| if current or not chunks: |
| chunks.append(current) |
| truncated = False |
| if len(chunks) > self.cfg.max_chunks: |
| chunks = chunks[: self.cfg.max_chunks] |
| truncated = True |
| wrapped = [[self.cls] + c + [self.sep] for c in chunks] |
| pieces = None |
| if self.cfg.state_retrieval and total > self.cfg.early_state_tokens: |
| k = self.cfg.retrieval_piece_tokens |
| pieces = [c[i : i + k] for c in chunks for i in range(0, len(c), k)] |
| return EncodedState(chunks=wrapped, tokens=total, truncated=truncated, pieces=pieces) |
|
|
| |
| def state_head(self, state: EncodedState, query: Sequence[int] | None = None) -> list[int]: |
| """State tokens placed inside the question sequence. |
| |
| Short states (or retrieval off): the first `early_state_tokens` tokens, in order. |
| Long states with retrieval on: the first `retrieval_lead_tokens` tokens plus the pieces |
| with the highest BM25 score against the question and option tokens, kept in state order. |
| """ |
| n = self.cfg.early_state_tokens |
| if state.pieces is None or not query: |
| out: list[int] = [] |
| for c in state.chunks: |
| if len(out) >= n: |
| break |
| out += c[1:-1] |
| return out[:n] |
| return self._retrieve(state.pieces, query, n) |
|
|
| def _retrieve(self, pieces: list[list[int]], query: Sequence[int], budget: int, k1: float = 1.2, b: float = 0.75) -> list[int]: |
| import math |
| from collections import Counter |
|
|
| special = {self.cls, self.sep, self.mask, self.pad} |
| q = {t for t in query if t not in special} |
| N = len(pieces) |
| tfs = [Counter(p) for p in pieces] |
| df = Counter(t for tf in tfs for t in tf.keys() & q) |
| avg = sum(len(p) for p in pieces) / max(1, N) |
| idf = {t: math.log(1 + (N - df[t] + 0.5) / (df[t] + 0.5)) for t in q if df[t]} |
| scores = [] |
| for i, (p, tf) in enumerate(zip(pieces, tfs)): |
| s = 0.0 |
| for t, w in idf.items(): |
| f = tf.get(t, 0) |
| if f: |
| s += w * f * (k1 + 1) / (f + k1 * (1 - b + b * len(p) / avg)) |
| scores.append(s) |
| lead, chosen, used = [], set(), 0 |
| for i, p in enumerate(pieces): |
| if used + len(p) > self.cfg.retrieval_lead_tokens: |
| break |
| chosen.add(i) |
| used += len(p) |
| for i in sorted(range(N), key=lambda j: -scores[j]): |
| if used >= budget: |
| break |
| if i in chosen or scores[i] <= 0: |
| continue |
| if used + len(pieces[i]) > budget: |
| continue |
| chosen.add(i) |
| used += len(pieces[i]) |
| out: list[int] = [] |
| for i in sorted(chosen): |
| out += pieces[i] |
| return out[:budget] |
|
|
| def _question(self, name: str, qtype: str, instructions: str, option_texts: Sequence[str], head_state: list[int]) -> EncodedQuestion: |
| head = self.encode(f"{qtype} question: {instructions}".strip()) |
| truncated = False |
| hlimit = max(16, self.cfg.max_question_tokens - 4) |
| if len(head) > hlimit: |
| head, truncated = head[:hlimit], True |
| ids = [self.cls] + head + [self.sep] |
| if head_state: |
| ids += head_state + [self.sep] |
| pos = list(range(len(ids))) |
| blocks = [0] * len(ids) |
| base = len(ids) |
| markers: list[int] = [] |
| limit = self.cfg.max_option_tokens |
| for t in option_texts: |
| opt = self.encode(" " + t) |
| if len(opt) > limit: |
| opt, truncated = opt[:limit], True |
| markers.append(len(ids)) |
| ids += [self.mask] + opt |
| blocks += [len(markers)] * (1 + len(opt)) |
| if self.cfg.option_shared_positions: |
| pos += list(range(base, base + 1 + len(opt))) |
| else: |
| start = pos[-1] + 1 |
| pos += list(range(start, start + 1 + len(opt))) |
| tokens = len(head) + sum(len(self.encode(" " + t)[:limit]) for t in option_texts) |
| return EncodedQuestion(name=name, qtype=QTYPE_IDS[qtype], ids=ids, position_ids=pos, markers=markers, blocks=blocks, |
| tokens=tokens, truncated=truncated) |
|
|
| def _query(self, instructions: str, option_texts: Sequence[str]) -> list[int]: |
| ids = list(self.encode(instructions)) |
| for t in option_texts: |
| ids += self.encode(" " + t) |
| return ids |
|
|
| def encode_question(self, name: str, question: Question, state: EncodedState | None = None) -> EncodedQuestion: |
| instr = render_content(question.instructions) |
| head_state = self.state_head(state, self._query(instr, question.option_texts)) if state is not None else [] |
| return self._question(name, question.type, render_content(question.instructions), question.option_texts, head_state) |
|
|
| def encode_request(self, state: Any, questions: dict[str, Question]) -> EncodedRequest: |
| enc_state = self.encode_state(state) |
| return EncodedRequest( |
| state=enc_state, |
| questions=[self.encode_question(name, q, enc_state) for name, q in questions.items()], |
| ) |
|
|
| |
| def encode_row(self, row: dict[str, Any], state: EncodedState | None = None) -> tuple[EncodedState, EncodedQuestion]: |
| """Encode a flat training row: {state, question, type, options[, target]}. |
| |
| Options are already rendered texts (`name: description` for choice, |
| `level i: text` for score, `[false, true]` descriptions for noul). |
| """ |
| st = state if state is not None else self.encode_state(row["state"]) |
| query = self._query(row.get("question") or "", row["options"]) |
| q = self._question(str(row.get("id", "row")), row.get("type", "choice"), row.get("question") or "", row["options"], self.state_head(st, query)) |
| return st, q |
|
|
|
|
| def question_from_row(row: dict[str, Any]) -> Question: |
| """Build a schema question from a flat row (used by evaluation tooling).""" |
| qtype = row.get("type", "choice") |
| if qtype == "choice": |
| return ChoiceQuestion(instructions=row.get("question"), criteria={str(o): None for o in row["options"]}) |
| if qtype == "score": |
| return ScoreQuestion(instructions=row.get("question"), criteria=list(row["options"])) |
| return NoulQuestion(instructions=row.get("question") or "true") |
|
|
|
|
| def flatten_question(name: str, question: Question, state: Any) -> dict[str, Any]: |
| """Inverse of `question_from_row` for logging and synthetic data.""" |
| return { |
| "name": name, |
| "state": state, |
| "question": render_content(question.instructions), |
| "type": question.type, |
| "options": question.option_texts, |
| } |
|
|
|
|
| def iter_option_texts(questions: Iterable[Question]) -> Iterable[str]: |
| for q in questions: |
| yield from q.option_texts |
|
|
