File size: 3,002 Bytes
210f9af | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | """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}
|