Something went wrong. Try again.
a tiny pet with a GPT-2 brain, living in your browser
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243"""Full fine-tune of GPT-2 on tokenling chat/diary data.
Loss is applied only to the character spans each row marks as trained (thepet's words and tool calls, and diary bodies), never to owner lines, headersor app-provided {tool results}."""import argparseimport collectionsimport jsonimport mathimport randomimport reimport sysimport time
import torchfrom transformers import AutoTokenizer, GPT2LMHeadModel
import gen_data as G
def encode(tok, row, max_len): text = row["text"] enc = tok(text, return_offsets_mapping=True) ids, offs = enc["input_ids"], enc["offset_mapping"] trained = bytearray(len(text) + 1) for a, b in row["spans"]: trained[a:b] = b"\x01" * (b - a) labels = [i if any(trained[s:e]) else -100 for i, (s, e) in zip(ids, offs)] if row["kind"] == "diary": # learn to stop after the entry ids.append(tok.eos_token_id) labels.append(tok.eos_token_id) return ids[:max_len], labels[:max_len]
def split_chat(row): hdr, body = row["text"].split("---\n", 1) meta = dict(l.split(": ", 1) for l in hdr.splitlines()[1:] if ": " in l) return hdr + "---\n", body, meta["pet"].split(" · ")[0], meta["owner"]
def with_history(row, donor): """Prepend a donor chat's lines (names swapped, untrained) as earlier history, so the pet learns to stay in character with long contexts.""" hdr, body, pet, owner = split_chat(row) _, dbody, dpet, downer = split_chat(donor) swap = {dpet: pet, downer: owner} dbody = re.sub(r"\b(%s)\b" % "|".join(map(re.escape, swap)), lambda m: swap[m.group(1)], dbody) shift = len(dbody) return {**row, "text": hdr + dbody + body, "spans": [[a + shift, b + shift] for a, b in row["spans"]]}
def augment(rows, frac): chats = [r for r in rows if r["kind"] == "chat"] # The nearest history must end in the recipient's starting sleep state; # otherwise an untrained donor nap teaches an awake reply while asleep. pools = {False: [], True: []} for r in chats: hdr, body, pet, _ = split_chat(r) asleep = f", {pet} is asleep" in hdr for ln in body.splitlines(): if ln == f"* {pet} woke up *": asleep = False elif ln.startswith(pet + ": ") and "[sleep]" in ln: asleep = True pools[asleep].append(r) out = [] for r in rows: out.append(r) if r["kind"] == "chat" and random.random() < frac: aug = r hdr, _, pet, _ = split_chat(r) asleep = f", {pet} is asleep" in hdr for _ in range(random.choice([1, 2])): candidates = pools[asleep] if not candidates: break d = random.choice(candidates) aug = with_history(aug, d) hdr, _, pet, _ = split_chat(d) asleep = f", {pet} is asleep" in hdr out.append(aug) return out
OFFER = re.compile(r"^\* .+ (offered .+|to play catch) \*$")
def late_answers(rows, frac, rng): """Move the [yes]/[no] of some offer answers to the end of the line. Teachers almost always answer first, so a pet that sampled an emote before answering learned that the answer was skipped and ate nothing. The line keeps its length, so the stored spans still fit.""" out, moved = [], 0 for r in rows: if r["kind"] != "chat": out.append(r) continue lines = r["text"].split("\n") for i in range(1, len(lines)): if not OFFER.match(lines[i - 1]) or "{" in lines[i] or rng.random() >= frac: continue m = re.search(r"( \[(yes|no)\]|^[^:]+: (\[(yes|no)\]) )", lines[i]) if not m: continue tag = "[" + (m.group(2) or m.group(4)) + "]" if m.group(1).startswith(" "): line = lines[i][:m.start()] + lines[i][m.end():] else: # answered first: "Mochi: [yes] nom" -> "Mochi: nom [yes]" line = lines[i][:m.start(3)] + lines[i][m.end():] if line + " " + tag != lines[i]: lines[i] = line + " " + tag moved += 1 out.append({**r, "text": "\n".join(lines)}) print(f"late answers: moved {moved}") return out
def batches(data, bs, shuffle): idx = list(range(len(data))) if shuffle: random.shuffle(idx) for i in range(0, len(idx), bs): yield [data[j] for j in idx[i:i + bs]]
def collate(rows, pad): n = max(len(r[0]) for r in rows) ids = torch.full((len(rows), n), pad) lab = torch.full((len(rows), n), -100) att = torch.zeros((len(rows), n), dtype=torch.long) for k, (i, l) in enumerate(rows): ids[k, :len(i)] = torch.tensor(i) lab[k, :len(l)] = torch.tensor(l) att[k, :len(i)] = 1 return ids, lab, att
@torch.no_grad()def evaluate(model, data, bs, pad): model.eval() tot = cnt = 0.0 for rows in batches(data, bs, False): ids, lab, att = (t.cuda() for t in collate(rows, pad)) with torch.autocast("cuda", dtype=torch.bfloat16): out = model(input_ids=ids, attention_mask=att, labels=lab) n = (lab[:, 1:] != -100).sum().item() tot += out.loss.item() * n cnt += n model.train() return tot / max(cnt, 1)
def fresh_rows(path, drop): """Every row is re-checked against the current data rules first: rows written for an older app (care events in the owner's mouth, catch without an offer) would teach behavior the app no longer has.""" rows, stale = [], collections.Counter() for line in open(path): row = json.loads(line) try: G.recheck(row) rows.append(row) except ValueError as e: stale[str(e).split(":")[0]] += 1 if stale: for k, v in stale.most_common(): print(f"{v:6} stale: {k}") if not drop: sys.exit(f"{sum(stale.values())} rows fail the current rules; fix the data or pass --drop-stale") return rows
def main(): ap = argparse.ArgumentParser() ap.add_argument("--data", required=True) ap.add_argument("--base", required=True) ap.add_argument("--out", required=True) ap.add_argument("--epochs", type=float, default=3) ap.add_argument("--lr", type=float, default=6e-5) ap.add_argument("--bs", type=int, default=16) ap.add_argument("--max-len", type=int, default=1024) ap.add_argument("--history-aug", type=float, default=0.25) ap.add_argument("--late-answer", type=float, default=0.3, help="share of offer answers whose [yes]/[no] moves to the end of the line") ap.add_argument("--drop-stale", action="store_true", help="drop rows the current rules reject instead of refusing to train") args = ap.parse_args() random.seed(0) torch.manual_seed(0)
tok = AutoTokenizer.from_pretrained(args.base) rows = fresh_rows(args.data, args.drop_stale) random.shuffle(rows) n_eval = max(64, len(rows) // 50) evals = [encode(tok, r, args.max_len) for r in rows[:n_eval]] # A separate generator, so the held-out split and history draws do not move. fit = late_answers(rows[n_eval:], args.late_answer, random.Random(1)) train = [encode(tok, r, args.max_len) for r in augment(fit, args.history_aug)] print(f"train={len(train)} eval={len(evals)} " f"tokens={sum(len(d[0]) for d in train)} trained={sum(sum(x != -100 for x in d[1]) for d in train)}")
model = GPT2LMHeadModel.from_pretrained(args.base, torch_dtype=torch.float32).cuda() model.gradient_checkpointing_enable() model.config.use_cache = False opt = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.95), weight_decay=0.05) steps = math.ceil(len(train) / args.bs * args.epochs) warm = max(20, steps // 30) sched = torch.optim.lr_scheduler.LambdaLR( opt, lambda s: min(1, (s + 1) / warm) * (0.1 + 0.9 * 0.5 * (1 + math.cos(math.pi * min(1, s / steps))))) print(f"steps={steps} eval_loss0={evaluate(model, evals, args.bs, tok.eos_token_id):.4f}", flush=True)
step, t0, best = 0, time.time(), 1e9 while step < steps: for rows in batches(train, args.bs, True): ids, lab, att = (t.cuda() for t in collate(rows, tok.eos_token_id)) with torch.autocast("cuda", dtype=torch.bfloat16): loss = model(input_ids=ids, attention_mask=att, labels=lab).loss loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step() sched.step() opt.zero_grad(set_to_none=True) step += 1 if step % 25 == 0: print(f"step {step}/{steps} loss {loss.item():.4f} lr {sched.get_last_lr()[0]:.2e} " f"{time.time() - t0:.0f}s", flush=True) if step % 200 == 0 or step == steps: ev = evaluate(model, evals, args.bs, tok.eos_token_id) print(f"eval step {step}: {ev:.4f}", flush=True) if ev < best: best = ev model.config.use_cache = True model.save_pretrained(args.out, safe_serialization=True) tok.save_pretrained(args.out) model.config.use_cache = False if step >= steps: break print(f"best eval {best:.4f}")
if __name__ == "__main__": main()