From 9b3d3c80cba3d6cb2933f5d4ca6f098c0f798f49 Mon Sep 17 00:00:00 2001 From: Vuong Hoang Date: Tue, 15 Sep 2026 10:45:51 -0700 Subject: [PATCH] BabyYarros Option C: instruction-pair builder and the assistant-masked pair trainer --- scripts/yarros-corpus/build_sft_pairs.py | 359 ++++++++++++++++++++++ scripts/yarros-corpus/train_pairs_lora.py | 226 ++++++++++++++ 2 files changed, 585 insertions(+) create mode 100644 scripts/yarros-corpus/build_sft_pairs.py create mode 100644 scripts/yarros-corpus/train_pairs_lora.py diff --git a/scripts/yarros-corpus/build_sft_pairs.py b/scripts/yarros-corpus/build_sft_pairs.py new file mode 100644 index 0000000..3e2edae --- /dev/null +++ b/scripts/yarros-corpus/build_sft_pairs.py @@ -0,0 +1,359 @@ +"""BabyYarros Option C — build instruction→response SFT pairs from the gated corpus. + +The architecture question was settled on 2026-09-11: the adapted completion carrier has +the voice and cannot take direction; the instruct model takes direction and has no voice. +Skaldsong needs both, which means the corpus has to be rebuilt as instruction→response +pairs and trained through the chat template rather than as raw continuation text. + +⭐ THE PROPERTY THAT MAKES THIS SAFE: the model writes the INSTRUCTION, never the +RESPONSE. Every response is real (renamed) Yarros prose, so the voice is inherited from +the corpus and not synthesised. Only the beat is machine-made, and a beat is a summary — +the easiest job in the building. A pipeline that generated the prose would be training +the carrier on a 27B's pastiche of Yarros, which is the opposite of the point. + +⚠ PASSAGES, NOT PARAGRAPHS. Measured on copy0: the median Yarros paragraph is 39 words +(p90 67) against the product's 90-140 band, so one paragraph is nowhere near one product +output. Accumulating consecutive paragraphs to the band gives 6,992 passages at median +104 words / 4 paragraphs — in-band by construction, and multi-paragraph, which is also +the continuity requirement the stitching failure raised (independently-generated +paragraphs drift in POV; the pair corpus has to contain multi-paragraph examples). + +⚠ ONE COPY ONLY. The rename produced 6 copies under different name maps. Six copies of +the same passage would be six near-duplicate responses differing only in proper nouns, +which is corpus inflation, not augmentation, at pair-training scale. + +⚠ TRAIN SPLIT ONLY. The val split stays untouched so held-out loss remains comparable to +the raw-text arms. + +Generator is the local `gen` seat: free, and abliterated, which matters here for a +non-obvious reason — the corpus contains sex and violence, and a refusing generator would +silently drop exactly the passages where the voice is most distinctive, biasing the +dataset toward its tamest regions. A refusal is a sampling bias, not just a gap. +""" +from __future__ import annotations +import argparse, json, random, re, sys, time, urllib.error, urllib.request +from pathlib import Path + +GATEWAY = "http://10.250.50.70:4000" + +# Matches gen_beats_chat_yarros.py verbatim so the pairs are trained in the SAME shape the +# evaluation harness measures. A pair corpus trained under a different system prompt than +# the eval drives would confound the carrier change with a prompt change. +SYS = ("You expand a single story beat into ONE paragraph of prose in the manner of Rebecca " + "Yarros — contemporary first-person PRESENT-tense narration, emotionally charged, sensory " + "and physical, the voice of new-adult romantasy. Render the beat itself; do not move past " + "it, do not add a new scene, do not comment. Output the paragraph only, 90–140 words.") + +CONTEXT_BLOCK = """The passage is preceded by this, for reference only. Do NOT write a beat for it — it is +here so you resolve names and pronouns correctly. + +PRECEDING: +{prev} + +""" + +BEAT_PROMPT = """{context}Below is a passage from a novel. Write the single story BEAT that a writer would have been given to produce it. + +Rules: +- ONE sentence, 8 to 30 words, present tense, third person. +- Name the characters who appear, using the names exactly as the passage spells them. +- The passage is first-person and its narrator is usually UNNAMED in it. Call that person + "she" or "he" as the passage implies. Never write "the narrator", "the speaker" or "I". +- State WHAT HAPPENS — the action, the turn, the decision, the reveal. Not how it is written. +- Do NOT describe the passage ("this excerpt shows..."), do NOT mention prose, style, tone or the author. +- Do NOT reuse distinctive phrases from the passage. Say the event in your own plain words. +- Output the sentence alone, with no label, no quotes and no preamble. + +PASSAGE: +{passage}""" + +META = re.compile(r"\b(passage|excerpt|paragraph|prose|narrat(?:or|ion)|the (?:author|text|scene) (?:is|describes)|this (?:scene|chapter))\b", re.I) + +# ⚠ MEASURED, and it was a composition bias rather than a nuisance: of 23 `meta` rejects in +# a 150-passage diagnostic, **22 were the single word "narrator"** and 1 was the story-word +# "passage". Those beats were otherwise clean ("Miguel shouts a warning as the venin +# advances, while the dark wielder taunts the narrator..."). The model reaches for "the +# narrator" precisely when the POV character is an ACTOR, so rejecting on it silently drops +# action scenes and keeps the ones where she only observes -- a systematic skew in what the +# carrier would learn to render, invisible in any spot-read of the kept pairs. The fix is a +# single bounded RETRY with the correction restated, not a substitution: the generator knows +# the character's gender from the prose and this session does not, and guessing it wrong +# across 780k words of first-person narration is the worse failure. +NARRATOR_ONLY = re.compile(r"\bnarrat(?:or|ion)\b", re.I) +RETRY_NOTE = ("\n\nYour previous answer used the word \"narrator\". Rewrite it referring to that " + "person as \"she\" or \"he\" — whichever the passage implies — and change nothing else.") + + +def ctx_block(prev: str | None) -> str: + """The beat generator sees the PRECEDING passage even when the pair will not carry it. + + Found in the positive control: passage [4] of nova ch16 is genuinely ambiguous about + who is buried and who is digging, and the hand beat and the model beat disagreed for + exactly that reason. A mid-scene passage can underdetermine its own referents; showing + the generator the previous passage fixes the pronouns without putting anything in the + instruction that the response does not support. + """ + return CONTEXT_BLOCK.format(prev=prev) if prev else "" + + +def post(path: str, payload: dict, key: str, timeout: int = 120) -> dict: + req = urllib.request.Request( + GATEWAY + path, data=json.dumps(payload).encode(), + headers={"Content-Type": "application/json", "Authorization": f"Bearer {key}"}) + with urllib.request.urlopen(req, timeout=timeout) as r: + return json.loads(r.read()) + + +def chunk(corpus: Path, lo: int, hi: int, split: str) -> list[dict]: + """Accumulate consecutive paragraphs into product-band passages, per chapter. + + Chapter-bounded so a passage never straddles a chapter break. The trailing buffer of + each chapter is DROPPED rather than emitted short -- an out-of-band response would + teach the length the product is trying to hold. + """ + out = [] + for f in sorted(corpus.glob("*.copy0.jsonl")): + for line in f.read_text(encoding="utf-8").splitlines(): + d = json.loads(line) + if d.get("split") != split: + continue + paras = [p.strip() for p in re.split(r"\n\s*\n", d["text"]) if p.strip()] + buf, n, prev = [], 0, None + for p in paras: + buf.append(p); n += len(p.split()) + if n >= lo: + if n <= int(hi * 1.6): + out.append({"work": d["work"], "chapter": d["chapter"], "split": d["split"], + "response": "\n\n".join(buf), "words": n, "prev": prev}) + prev = "\n\n".join(buf) + else: + prev = None # oversized run dropped; context would be a lie + buf, n = [], 0 + return out + + +def ngrams(text: str, n: int) -> set[str]: + w = re.findall(r"[a-z']+", text.lower()) + return {" ".join(w[i:i + n]) for i in range(max(0, len(w) - n + 1))} + + +def vet(beat: str, passage: str, overlap_n: int) -> tuple[str | None, str]: + """Return (clean_beat, reason). reason is '' on accept. + + Every rejection is a dataset defect that would otherwise train silently: + echo -- a beat quoting the passage teaches the model to copy its instruction + back, not to expand it. This is the one that would look fine in a + spot-read and poison the whole run. + meta -- 'this passage shows...' is a description of text, not a story beat + length -- a 60-word beat is a summary; a 4-word beat is a title + """ + beat = " ".join(beat.strip().split()) + beat = re.sub(r'^(?:beat|answer)\s*[:\-]\s*', '', beat, flags=re.I).strip().strip('"“”') + if not beat: + return None, "empty" + # first sentence only -- the model sometimes adds a second + m = re.match(r"^(.+?[.!?])(?:\s|$)", beat) + if m: + beat = m.group(1).strip() + nw = len(beat.split()) + if nw < 6 or nw > 34: + return None, f"length({nw})" + if META.search(beat): + return None, "meta" + if ngrams(beat, overlap_n) & ngrams(passage, overlap_n): + return None, f"echo({overlap_n}gram)" + return beat, "" + + +NAME = re.compile(r"(? list[str]: + """Names the BEAT introduces that the RESPONSE never spells out. + + Reported, deliberately NOT rejected. Giving the generator the preceding passage fixed + the referent ambiguity the positive control exposed -- it corrected a hand-written beat + that had the buried and the digging characters backwards -- but it also lets a beat name + someone who appears only in that preceding text. Rejecting on this would drop exactly the + passages whose POV character is unnamed, which is a sampling bias dressed as a guard + (same argument as the refusing-generator note above). So it is counted and surfaced in + provenance, and the rate decides whether it needs a rule. + """ + return sorted({n for n in NAME.findall(beat) if n not in response}) + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--corpus", required=True, help="dir holding *.copy0.jsonl") + ap.add_argument("--out", required=True) + ap.add_argument("--key-file", default=None, help="file holding the gateway key") + ap.add_argument("--key", default=None) + ap.add_argument("--model", default="gen") + ap.add_argument("--n", type=int, default=600) + ap.add_argument("--lo", type=int, default=90) + ap.add_argument("--hi", type=int, default=150) + ap.add_argument("--split", default="train") + ap.add_argument("--seed", type=int, default=4919) + ap.add_argument("--context-frac", type=float, default=0.4, + help="fraction of pairs carrying the preceding passage as context") + ap.add_argument("--overlap-n", type=int, default=6) + ap.add_argument("--temperature", type=float, default=0.3) + ap.add_argument("--control-out", default=None, + help="write the first --control-n passages out for hand-written beats") + ap.add_argument("--control-n", type=int, default=10) + ap.add_argument("--dump-passages", action="store_true", help="chunk and report, generate nothing") + ap.add_argument("--control-in", default=None, + help="JSON of hand-written beats [{idx,beat}]; generate beats for the SAME " + "passages and print side by side. The positive control -- a beat " + "generator nobody checked against a human can produce a clean-looking " + "dataset that teaches the wrong mapping, and nothing downstream would show it.") + a = ap.parse_args() + + key = a.key or (Path(a.key_file).read_text().strip() if a.key_file else None) + passages = chunk(Path(a.corpus), a.lo, a.hi, a.split) + print(f"[chunk] {len(passages)} passages in split={a.split}, band {a.lo}-{a.hi}", flush=True) + if not passages: + print("REFUSING: no passages -- wrong corpus dir or split", file=sys.stderr) + return 2 + + random.seed(a.seed) + random.shuffle(passages) + + if a.control_out: + ctl = [{"idx": i, "work": p["work"], "chapter": p["chapter"], "response": p["response"], + "beat_handwritten": ""} for i, p in enumerate(passages[:a.control_n])] + Path(a.control_out).write_text(json.dumps(ctl, indent=2, ensure_ascii=False), encoding="utf-8") + print(f"[control] wrote {len(ctl)} passages to {a.control_out} for hand-written beats") + + if a.dump_passages: + return 0 + if not key: + print("REFUSING: no gateway key (--key or --key-file)", file=sys.stderr) + return 2 + + # The alias is not the model. `gen` has pointed at different concrete backends over + # time; counting by an alias once inflated an exposure figure 4.7x on this fleet. The + # build fingerprint is resolved at START and again at END and both go in provenance -- + # a seat repointed mid-run would otherwise be invisible in the artefact. + def fingerprint() -> str: + try: + r = post("/v1/chat/completions", {"model": a.model, "max_tokens": 1, + "messages": [{"role": "user", "content": "ok"}]}, key, 60) + return r.get("system_fingerprint") or "unknown" + except Exception as e: + return f"unresolved:{e}" + + fp_start = fingerprint() + print(f"[seat] {a.model} fingerprint at start: {fp_start}", flush=True) + + if a.control_in: + hand = {h["idx"]: h["beat"] for h in json.loads(Path(a.control_in).read_text(encoding="utf-8"))} + agree = 0 + for i in sorted(hand): + psg = passages[i] + r = post("/v1/chat/completions", { + "model": a.model, "temperature": a.temperature, "max_tokens": 80, + "messages": [{"role": "user", "content": BEAT_PROMPT.format(passage=psg["response"], context=ctx_block(psg["prev"]))}], + }, key) + beat, why = vet(r["choices"][0]["message"]["content"] or "", psg["response"], a.overlap_n) + print("=" * 72) + print(f"[{i}] {psg['work']} ch{psg['chapter']}") + print(f" HAND : {hand[i]}") + print(f" MODEL: {beat if beat else '<>'}") + if beat: + agree += 1 + print("=" * 72) + print(f"[control] {agree}/{len(hand)} model beats passed the guards; " + f"judge the EVENT match by reading, not by this count") + return 0 + + out_path = Path(a.out) + rejects: dict[str, int] = {} + kept = 0 + name_drift = 0 + retried = 0 + retry_saved = 0 + drift_examples: list[dict] = [] + t0 = time.time() + with out_path.open("w", encoding="utf-8") as fh: + for p in passages: + if kept >= a.n: + break + try: + r = post("/v1/chat/completions", { + "model": a.model, "temperature": a.temperature, "max_tokens": 80, + "messages": [{"role": "user", "content": BEAT_PROMPT.format(passage=p["response"], context=ctx_block(p["prev"]))}], + }, key) + raw = r["choices"][0]["message"]["content"] or "" + except (urllib.error.URLError, urllib.error.HTTPError, KeyError, TimeoutError) as e: + rejects["transport"] = rejects.get("transport", 0) + 1 + print(f"[warn] {type(e).__name__}: {e}", flush=True) + continue + beat, why = vet(raw, p["response"], a.overlap_n) + if beat is None and why == "meta" and NARRATOR_ONLY.search(raw): + retried += 1 + try: + r2 = post("/v1/chat/completions", { + "model": a.model, "temperature": a.temperature, "max_tokens": 80, + "messages": [ + {"role": "user", "content": BEAT_PROMPT.format( + passage=p["response"], context=ctx_block(p["prev"]))}, + {"role": "assistant", "content": raw}, + {"role": "user", "content": RETRY_NOTE.strip()}, + ]}, key) + beat, why = vet(r2["choices"][0]["message"]["content"] or "", + p["response"], a.overlap_n) + if beat is not None: + retry_saved += 1 + except (urllib.error.URLError, urllib.error.HTTPError, KeyError, TimeoutError): + pass + if beat is None: + rejects[why.split("(")[0]] = rejects.get(why.split("(")[0], 0) + 1 + continue + drift = names_not_in_response(beat, p["response"]) + if drift: + name_drift += 1 + drift_examples.append({"beat": beat, "names": drift}) + use_ctx = p["prev"] is not None and random.random() < a.context_frac + fh.write(json.dumps({ + "work": p["work"], "chapter": p["chapter"], "split": p["split"], + "beat": beat, "context": p["prev"] if use_ctx else None, + "response": p["response"], "words": p["words"], + }, ensure_ascii=False) + "\n") + kept += 1 + if kept % 50 == 0: + print(f"[gen] {kept}/{a.n} {time.time()-t0:.0f}s rejects={rejects}", flush=True) + + fp_end = fingerprint() + prov = { + "generated_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"), + "generator_model_alias": a.model, + "generator_fingerprint_start": fp_start, + "generator_fingerprint_end": fp_end, + "generator_repointed_midrun": fp_start != fp_end, + "gateway": GATEWAY, "temperature": a.temperature, + "corpus": str(a.corpus), "split": a.split, "copy": "copy0", + "band_words": [a.lo, a.hi], "seed": a.seed, "context_frac": a.context_frac, + "overlap_guard_n": a.overlap_n, + "passages_available": len(passages), "pairs_kept": kept, "rejects": rejects, + "narrator_retries": retried, "narrator_retries_recovered": retry_saved, + "beats_naming_absent_entity": name_drift, + "beats_naming_absent_entity_rate": round(name_drift / max(1, kept), 4), + "beats_naming_absent_entity_examples": drift_examples[:15], + "reject_rate": round(sum(rejects.values()) / max(1, kept + sum(rejects.values())), 4), + "elapsed_s": round(time.time() - t0, 1), + "system_prompt": SYS, + } + Path(str(out_path) + ".provenance.json").write_text(json.dumps(prov, indent=2), encoding="utf-8") + print(f"[done] {kept} pairs -> {out_path}") + print(f"[done] rejects {rejects} (rate {prov['reject_rate']})") + print(f"[done] narrator-retries {retried}, recovered {retry_saved}") + print(f"[done] beats naming an entity absent from their response: {name_drift}/{kept} " + f"({prov['beats_naming_absent_entity_rate']}) -- reported, not rejected") + if prov["generator_repointed_midrun"]: + print("⚠ SEAT REPOINTED MID-RUN -- the pairs are not from one generator", file=sys.stderr) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/scripts/yarros-corpus/train_pairs_lora.py b/scripts/yarros-corpus/train_pairs_lora.py new file mode 100644 index 0000000..795248c --- /dev/null +++ b/scripts/yarros-corpus/train_pairs_lora.py @@ -0,0 +1,226 @@ +"""BabyYarros Option C — LoRA SFT on beat→paragraph pairs, loss on the RESPONSE only. + +Sibling of `r49-corpus/train_voice_lora.py`, not a mode of it: that script is the +reproducibility anchor for every raw-text arm (Brontë 0.6B/1.7B/4B, Yarros base and +instruct) and adding a branch to it would put those runs one refactor away from moving. +Everything that is not the data path is held IDENTICAL to the raw-text instruct arm -- +rank 32, alpha 64, lr 1e-4, batch 1 x accum 8, seed 4919, cosine with 3% warmup, eval and +save every 25 -- so the only variable between the two arms is the shape of the data. + +⚠ LOSS IS MASKED TO THE ASSISTANT TURN. The raw-text arms train on every token because +every token is Yarros. Here the system prompt and the beat are NOT Yarros -- the beat is +machine-written -- and training on them teaches the carrier to generate beats, which is +capacity spent on the half of the exchange the product supplies. Prompt tokens get label +-100 and the collator checks that something survives: an example with zero unmasked labels +contributes no gradient and would otherwise train silently as a no-op. + +⚠ NO PACKING. The raw-text path packs into fixed 4096-token blocks, which is right for +continuous prose and wrong here: packing would run loss across an example boundary and +teach the model that a response is followed by another system prompt. Pairs are padded +per example instead, and the cost -- padding waste -- is trivial at ~400 tokens each. + +⚠ THE HELD-OUT LOSS FROM THIS RUN IS NOT COMPARABLE TO THE RAW-TEXT ARMS' 2.5263 / 2.6114. +Different objective over different tokens. It is good for exactly one thing, picking the +best checkpoint within this run, which the 4B rung showed is not the end-of-run adapter. +The cross-arm comparison has to be behavioural (on-beat / in-band / ran-on) plus delta_cb. +""" +from __future__ import annotations +import argparse, hashlib, json, math, random, subprocess, time +from pathlib import Path + +import torch +from torch.utils.data import Dataset +from transformers import (AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments, + TrainerCallback) +from peft import LoraConfig, get_peft_model + +TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] + +SYS = ("You expand a single story beat into ONE paragraph of prose in the manner of Rebecca " + "Yarros — contemporary first-person PRESENT-tense narration, emotionally charged, sensory " + "and physical, the voice of new-adult romantasy. Render the beat itself; do not move past " + "it, do not add a new scene, do not comment. Output the paragraph only, 90–140 words.") + + +def user_msg(rec: dict) -> str: + if rec.get("context"): + return (f"Continuing from:\n\n{rec['context']}\n\nBeat: {rec['beat']}") + return f"Beat: {rec['beat']}" + + +class Pairs(Dataset): + """Renders each pair through the chat template and masks the prompt. + + The prompt length is measured by tokenising the SAME prefix the full render begins + with -- add_generation_prompt=True on the messages minus the assistant turn -- rather + than by searching for a delimiter string. A delimiter search is what breaks silently + when a template changes its spacing, and it would mask the wrong span without erroring. + """ + + def __init__(self, tok, records, seq_len): + self.rows = [] + 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) + full = tok.apply_chat_template( + msgs + [{"role": "assistant", "content": r["response"]}], + tokenize=True, add_generation_prompt=False) + if len(full) > seq_len or len(full) <= len(prefix): + dropped += 1 + continue + labels = [-100] * len(prefix) + list(full[len(prefix):]) + self.rows.append((full, labels)) + self.dropped = dropped + if not self.rows: + raise SystemExit("REFUSING: every pair was dropped -- template or seq_len is wrong") + + def __len__(self): + return len(self.rows) + + def __getitem__(self, i): + ids, labels = self.rows[i] + return {"input_ids": ids, "labels": labels} + + +def collate(batch, pad_id): + n = max(len(b["input_ids"]) for b in batch) + ids, labels, mask = [], [], [] + for b in batch: + k = n - len(b["input_ids"]) + ids.append(b["input_ids"] + [pad_id] * k) + labels.append(b["labels"] + [-100] * k) + mask.append([1] * len(b["input_ids"]) + [0] * k) + lab = torch.tensor(labels) + if int((lab != -100).sum()) == 0: + raise SystemExit("REFUSING: a batch has no unmasked labels -- it would train as a no-op") + return {"input_ids": torch.tensor(ids), "labels": lab, "attention_mask": torch.tensor(mask)} + + +class LossLog(TrainerCallback): + def __init__(self, path): + self.path, self.series = path, [] + + def on_log(self, args, state, control, logs=None, **kw): + if logs: + self.series.append({"step": state.global_step, **logs}) + self.path.write_text(json.dumps(self.series, indent=2)) + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("--pairs", required=True) + ap.add_argument("--val-pairs", required=True) + ap.add_argument("--base", required=True) + ap.add_argument("--out", required=True) + ap.add_argument("--seq-len", type=int, default=1536) + ap.add_argument("--rank", type=int, default=32) + ap.add_argument("--lr", type=float, default=1e-4) + ap.add_argument("--epochs", type=float, default=3.0) + ap.add_argument("--batch", type=int, default=1) + ap.add_argument("--accum", type=int, default=8) + ap.add_argument("--seed", type=int, default=4919) + ap.add_argument("--eval-steps", type=int, default=25) + ap.add_argument("--save-steps", type=int, default=25) + a = ap.parse_args() + + torch.manual_seed(a.seed); random.seed(a.seed) + out = Path(a.out); out.mkdir(parents=True, exist_ok=True) + + tok = AutoTokenizer.from_pretrained(a.base) + if tok.chat_template is None: + raise SystemExit("REFUSING: carrier has no chat template -- it is not an instruct build") + if tok.pad_token_id is None: + tok.pad_token = tok.eos_token + + h = hashlib.sha256() + def load(p): + b = Path(p).read_bytes(); h.update(b) + return [json.loads(l) for l in b.decode("utf-8").splitlines() if l.strip()] + train_recs, val_recs = load(a.pairs), load(a.val_pairs) + pairs_sha = h.hexdigest()[:16] + + # The val pairs must come from corpus passages the train pairs never saw. A pair set + # split off the same passages would report a held-out loss over text the adapter had + # already fit, which is the failure that looks most like success. + tr_keys = {(r["work"], r["chapter"], r["response"][:80]) for r in train_recs} + overlap = [r for r in val_recs if (r["work"], r["chapter"], r["response"][:80]) in tr_keys] + if overlap: + raise SystemExit(f"REFUSING: {len(overlap)} val pairs share a passage with train") + if any(r.get("split") != "val" for r in val_recs): + raise SystemExit("REFUSING: val pairs are not all from the corpus val split") + + train_ds, val_ds = Pairs(tok, train_recs, a.seq_len), Pairs(tok, val_recs, a.seq_len) + with_ctx = sum(1 for r in train_recs if r.get("context")) + print(f"[data] {len(train_ds)} train pairs ({train_ds.dropped} dropped), " + f"{len(val_ds)} val ({val_ds.dropped} dropped), " + f"{with_ctx}/{len(train_recs)} carry context, pairs sha {pairs_sha}", flush=True) + + model = AutoModelForCausalLM.from_pretrained(a.base, dtype=torch.bfloat16, + attn_implementation="sdpa").to("cuda") + model = get_peft_model(model, LoraConfig(r=a.rank, lora_alpha=2 * a.rank, lora_dropout=0.0, + bias="none", task_type="CAUSAL_LM", + target_modules=TARGETS)) + model.gradient_checkpointing_enable(); model.enable_input_require_grads() + trainable = sum(p.numel() for p in model.parameters() if p.requires_grad) + total = sum(p.numel() for p in model.parameters()) + + resolved = { + "attn_implementation": getattr(model.config, "_attn_implementation", "?"), + "dtype": str(next(model.parameters()).dtype), + "device": torch.cuda.get_device_name(0), + "torch": torch.__version__, + "adapted_modules": sum(1 for n, _ in model.named_modules() if n.endswith("lora_A.default")), + } + try: + repo = Path(__file__).resolve().parents[2] + git = subprocess.run(["git", "-C", str(repo), "rev-parse", "--short", "HEAD"], + capture_output=True, text=True).stdout.strip() + dirty = bool(subprocess.run(["git", "-C", str(repo), "status", "--porcelain"], + capture_output=True, text=True).stdout.strip()) + except Exception: + git, dirty = "?", True + + steps_per_epoch = math.ceil(len(train_ds) / (a.batch * a.accum)) + prov = {"run": "r49-babyyarros-pairs-pilot", "base": a.base, + "pairs": a.pairs, "val_pairs": a.val_pairs, "pairs_sha256_16": pairs_sha, + "objective": "chat-template SFT, loss on assistant turn only", + "loss_comparable_to_raw_text_arms": False, + "train_pairs": len(train_ds), "val_pairs_n": len(val_ds), + "pairs_with_context": with_ctx, + "seq_len": a.seq_len, "lora_rank": a.rank, "lora_alpha": 2 * a.rank, "targets": TARGETS, + "lr": a.lr, "epochs": a.epochs, "batch": a.batch, "grad_accum": a.accum, "seed": a.seed, + "trainable_params": trainable, "total_params": total, + "trainable_pct": round(100 * trainable / total, 3), + "steps_per_epoch": steps_per_epoch, "planned_steps": steps_per_epoch * int(a.epochs), + "resolved": resolved, "harness_commit": git, "harness_dirty_at_launch": dirty, + "launched_at": time.strftime("%Y-%m-%dT%H:%M:%S%z"), "system_prompt": SYS} + (out / "provenance.json").write_text(json.dumps(prov, indent=2)) + print("[prov] " + json.dumps({k: prov[k] for k in + ("pairs_sha256_16", "train_pairs", "planned_steps", "trainable_pct", + "harness_dirty_at_launch")}), flush=True) + print("[prov] resolved: " + json.dumps(resolved), flush=True) + + args = TrainingArguments( + output_dir=str(out / "checkpoints"), per_device_train_batch_size=a.batch, + gradient_accumulation_steps=a.accum, num_train_epochs=a.epochs, learning_rate=a.lr, + lr_scheduler_type="cosine", + warmup_steps=max(1, int(0.03 * steps_per_epoch * int(a.epochs))), + bf16=True, logging_steps=10, + save_strategy="steps", save_steps=a.save_steps, save_total_limit=12, + eval_strategy="steps", eval_steps=a.eval_steps, + report_to=[], seed=a.seed, + gradient_checkpointing=True, dataloader_num_workers=2, + remove_unused_columns=False, + ) + trainer = Trainer(model=model, args=args, train_dataset=train_ds, eval_dataset=val_ds, + data_collator=lambda b: collate(b, tok.pad_token_id), + callbacks=[LossLog(out / "loss-series.json")]) + trainer.train() + model.save_pretrained(out / "adapter") + print("[done] adapter saved; the BEST checkpoint is in checkpoints/, not necessarily this one") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())