File size: 2,231 Bytes
906715b | 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 | """Row rendering — the exact-prefix collator law.
Rows are either pre-tokenized {"ids": [...], "n_prefix": int} or
messages-form {"messages": [...], "target": str}. Messages-form is
rendered by the TWO-PASS chat-template method and the exact-prefix
assertion is a hard error, never a fudge (the campaign's collator law:
a silent prefix drift poisons the label mask).
"""
from __future__ import annotations
from typing import Sequence
import torch
def render_rows(tok, rows: Sequence[dict], max_len: int) -> list[dict]:
out = []
n_drift = 0
for r in rows:
if "ids" in r:
if len(r["ids"]) <= max_len:
out.append({"ids": list(r["ids"]),
"n_prefix": int(r["n_prefix"])})
continue
msgs = r["messages"]
pre_t = tok.apply_chat_template(msgs, tokenize=False,
add_generation_prompt=True)
full_t = tok.apply_chat_template(
msgs + [{"role": "assistant", "content": r["target"]}],
tokenize=False, add_generation_prompt=False)
pre = tok(pre_t, add_special_tokens=False).input_ids
full = tok(full_t, add_special_tokens=False).input_ids
if full[:len(pre)] != pre:
n_drift += 1
continue
if len(full) <= max_len:
out.append({"ids": full, "n_prefix": len(pre)})
if n_drift:
raise ValueError(
f"exact-prefix drift on {n_drift} rows: the chat template "
"re-tokenizes the prefix differently with the assistant "
"turn appended — fix the template usage, do not mask over "
"it (collator law)")
return out
def pad_batch(rows: Sequence[dict], pad_id: int, device):
"""Right-pad; labels -100 over prefix AND pads."""
W = max(len(r["ids"]) for r in rows)
ii, ll, am = [], [], []
for r in rows:
ids, n = r["ids"], len(r["ids"])
lab = [-100] * r["n_prefix"] + ids[r["n_prefix"]:]
ii.append(ids + [pad_id] * (W - n))
ll.append(lab + [-100] * (W - n))
am.append([1] * n + [0] * (W - n))
t = lambda x: torch.tensor(x, dtype=torch.long, device=device)
return t(ii), t(ll), t(am)
|