BabyYarros Option C: instruction-pair builder and the assistant-masked pair trainer
This commit is contained in:
Vuong Hoang
@@ -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"(?<![.!?“\"]\s)(?<!^)\b([A-Z][a-z]{2,})\b")
|
||||
|
||||
|
||||
def names_not_in_response(beat: str, response: str) -> 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 '<<REJECTED: ' + why + '>>'}")
|
||||
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())
|
||||
@@ -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())
|
||||
Reference in New Issue
Block a user