From 90ed506db3571a32af5204fa37dff7171ab28004 Mon Sep 17 00:00:00 2001 From: Vuong Hoang Date: Tue, 15 Sep 2026 10:51:49 -0700 Subject: [PATCH] fix(pairs): apply_chat_template returns a BatchEncoding, and assert the mask boundary --- scripts/yarros-corpus/train_pairs_lora.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/scripts/yarros-corpus/train_pairs_lora.py b/scripts/yarros-corpus/train_pairs_lora.py index 795248c..c6f1f67 100644 --- a/scripts/yarros-corpus/train_pairs_lora.py +++ b/scripts/yarros-corpus/train_pairs_lora.py @@ -62,10 +62,23 @@ class Pairs(Dataset): dropped = 0 for r in records: msgs = [{"role": "system", "content": SYS}, {"role": "user", "content": user_msg(r)}] - prefix = tok.apply_chat_template(msgs, tokenize=True, add_generation_prompt=True) + # ⚠ transformers 5.16 returns a BatchEncoding from apply_chat_template(tokenize=True), + # not a list of ids. Taking len() of it yields 2 (the number of keys), so every + # example failed the `len(full) <= len(prefix)` test and the whole dataset was + # dropped. The REFUSING guard below is what surfaced it -- a silent version of this + # bug trains on nothing and reports a loss curve anyway. + prefix = tok.apply_chat_template(msgs, tokenize=True, + add_generation_prompt=True)["input_ids"] full = tok.apply_chat_template( msgs + [{"role": "assistant", "content": r["response"]}], - tokenize=True, add_generation_prompt=False) + tokenize=True, add_generation_prompt=False)["input_ids"] + # The mask is only correct if the generation prefix is a TRUE prefix of the full + # render. Asserted rather than assumed: a template revision that reorders or + # re-spaces the header would shift the boundary and mask the wrong span, which + # trains without erroring and looks exactly like a normal run. + if list(full[:len(prefix)]) != list(prefix): + raise SystemExit("REFUSING: generation prefix is not a prefix of the full render " + "-- the loss mask would cover the wrong tokens") if len(full) > seq_len or len(full) <= len(prefix): dropped += 1 continue