| """Model input from typed segments, with segment-wise truncation to a token budget (pure Python). |
| |
| Bundle rows (and the inference worker's requests) carry the renderer's typed segments |
| (``[{"kind": "header"|"title"|"body"|..., "text": ...}]``) whose ``\\n``-join is the ``render_input()`` text |
| (classifier-plan §2.3, §4.1). Each segment is tokenised on its own (``add_special_tokens=False``, |
| ``split_special_tokens=True``: post text is untrusted), segments are joined with the ids of ``"\\n"``, and the |
| suffix ``"\\n### LABELS"`` closes the input; Gemma inputs start with exactly one BOS. Training, ``predict`` and |
| the node12 server all go through ``encode``, so their inputs are identical. |
| |
| Truncation never parses post text: when a row is over budget the body is cut from its end (body tail first); |
| headers, title and the suffix are never cut (a row whose fixed parts alone exceed the budget is an error). |
| """ |
| from __future__ import annotations |
|
|
| import hashlib |
| from typing import Callable, Iterable, Mapping |
|
|
| SUFFIX = "\n### LABELS" |
| SEPARATOR = "\n" |
| FORMAT = "nsb-heads-v1" |
| CUTTABLE = ("body",) |
|
|
| Tokenize = Callable[[str], list[int]] |
|
|
|
|
| def joined(segments: Iterable[Mapping[str, str]]) -> str: |
| """The rendered text: segments joined with newlines (== render_input() output).""" |
| return SEPARATOR.join(s["text"] for s in segments) |
|
|
|
|
| def input_hash(segments: Iterable[Mapping[str, str]], render_version: str) -> str: |
| """``sha256(render_version + "\\n" + text)``: the labelling §3.2 / ``post_labels.input_hash`` definition.""" |
| return hashlib.sha256((render_version + "\n" + joined(segments)).encode("utf-8")).hexdigest() |
|
|
|
|
| def encode(segments: list[Mapping[str, str]], tokenize: Tokenize, max_len: int, bos_id: int | None = None |
| ) -> tuple[list[int], dict]: |
| """Token ids for one input and a record of what was cut: ``{"n_tokens", "cut": {kind: tokens removed}}``.""" |
| if not segments: |
| raise ValueError("empty input: no segments") |
| pieces = [list(tokenize(s["text"])) if s["text"] else [] for s in segments] |
| sep = list(tokenize(SEPARATOR)) |
| suffix = list(tokenize(SUFFIX)) |
| fixed = len(suffix) + len(sep) * (len(segments) - 1) + (1 if bos_id is not None else 0) |
| excess = fixed + sum(len(p) for p in pieces) - max_len |
| cut: dict[str, int] = {} |
| for kind in CUTTABLE: |
| for i in range(len(pieces) - 1, -1, -1): |
| if excess <= 0: |
| break |
| if segments[i]["kind"] != kind or not pieces[i]: |
| continue |
| take = min(excess, len(pieces[i])) |
| pieces[i] = pieces[i][:len(pieces[i]) - take] |
| cut[kind] = cut.get(kind, 0) + take |
| excess -= take |
| if excess > 0: |
| raise ValueError("header, title and suffix alone exceed max_len") |
| ids: list[int] = [] if bos_id is None else [bos_id] |
| for i, p in enumerate(pieces): |
| if i: |
| ids += sep |
| ids += p |
| ids += suffix |
| return ids, {"n_tokens": len(ids), "cut": cut} |
|
|
