File size: 2,510 Bytes
3d7c238 | 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 | """Model input from export segments, with segment-wise truncation to a token budget (pure Python).
The export stores each input as typed segments (``header``, ``target_text``, ``context_text`` with a
``relation``, ``facts``) whose concatenation is the input (format ``evidence-v1``). Each segment is tokenised
on its own and the ids are concatenated, followed by the label suffix; training and prediction both go
through ``encode``, so the two always agree. Truncation never parses post text: when a row is over budget it
shortens quoted text, then parent text (each from its end, last segment first), and the target text only as
a last resort. Headers, the FACTS block and the suffix are never cut.
"""
from __future__ import annotations
SUFFIX = '\n### LABELS'
FORMAT = 'evidence-v1'
CUT_ORDER = (('context_text', 'quoted'), ('context_text', 'self_thread_parent'), ('context_text', 'reply_parent'),
('target_text', None))
def joined(row: dict) -> str:
return ''.join(s['text'] for s in row['segments'])
def _cuttable(seg: dict, kind: str, relation) -> bool:
return seg['kind'] == kind and (relation is None or seg.get('relation') == relation)
def encode(row: dict, tokenize, max_len: int) -> tuple[list[int], dict]:
"""Token ids for one export row, and a record of what was cut.
``tokenize(text) -> list[int]`` must not add special tokens and must not parse special-token strings in
the text (post text is untrusted). The record is ``{"n_tokens", "cut": {segment label: tokens removed}}``.
"""
pieces = [list(tokenize(s['text'])) if s['text'] else [] for s in row['segments']]
suffix = list(tokenize(SUFFIX))
excess = sum(len(p) for p in pieces) + len(suffix) - max_len
cut = {}
for kind, relation in CUT_ORDER:
if excess <= 0:
break
for i in range(len(pieces) - 1, -1, -1):
if excess <= 0:
break
if not _cuttable(row['segments'][i], kind, relation) or not pieces[i]:
continue
take = min(excess, len(pieces[i]))
pieces[i] = pieces[i][:len(pieces[i]) - take]
label = kind if relation is None else '%s:%s' % (kind, relation)
cut[label] = cut.get(label, 0) + take
excess -= take
if excess > 0:
raise ValueError('Headers, FACTS and suffix alone exceed max_len for row %s' % row.get('id'))
ids = [t for p in pieces for t in p] + suffix
return ids, {'n_tokens': len(ids), 'cut': cut}
|