"""Train the argument-parts + effectiveness model on PERSUADE 2.0, on one Modal A10G. MODAL_PROFILE=labs3 modal run scripts/train_parts.py --tag parts-v1 MODAL_PROFILE=labs3 modal run scripts/train_parts.py --tag smoke-parts --limit 64 --epochs 1 Metrics on the PERSUADE test split: span F1 (the Feedback Prize 2021 metric; winners' single models ~0.70) and effectiveness log loss on gold parts (Feedback Prize Effectiveness; winners' single models ~0.57-0.60). """ import os import modal app = modal.App("essay-coach-train-parts") vol = modal.Volume.from_name("essay-data") cache = modal.Volume.from_name("essay-hf-cache", create_if_missing=True) image = (modal.Image.debian_slim(python_version="3.11") .pip_install("torch==2.4.1", index_url="https://download.pytorch.org/whl/cu124") .pip_install("transformers==4.46.3", "sentencepiece", "protobuf", "tiktoken", "safetensors", "pandas", "numpy<2.3", "scikit-learn", "huggingface_hub") .add_local_python_source("coach")) HF_REPO = "nonsodev/essay-coach" HF_SECRET = True # flip to True once `modal secret create huggingface HF_TOKEN=...` exists on this workspace SECRETS = [modal.Secret.from_name("huggingface")] if HF_SECRET else [] @app.function(image=image, gpu="A10G", volumes={"/data": vol, "/root/.cache/huggingface": cache}, timeout=6 * 3600, secrets=SECRETS, cpu=4, memory=24576) def train(tag: str, epochs: int, bs: int, accum: int, lr: float, max_len: int, limit: int, seed: int, val: int): import json import time import numpy as np import pandas as pd import torch import torch.nn.functional as F from safetensors.torch import save_file from coach.data import BIO, EFFECT, persuade from coach.decode import decode, span_f1 from coach.model import PartsModel, param_groups, parts_tokenizer torch.manual_seed(seed) np.random.seed(seed) run = f"/data/runs/{tag}" os.makedirs(run, exist_ok=True) tr, te = persuade(split="train"), persuade(split="test") # the corpus's own test split is 40% of the essays; keep `val` of them (default 2000) as the held-out set and # train on the rest, so the model sees as much as possible while the number we report stays on unseen essays te = te.sample(frac=1.0, random_state=0) hold = val or 2000 tr = pd.concat([tr, te.iloc[hold:]], ignore_index=True) te = te.iloc[:hold] print(f"essays train {len(tr)} held-out {len(te)}; parts dropped {tr.attrs.get('dropped_parts', 0)}", flush=True) if limit: tr, te = tr.sample(limit, random_state=seed), te.sample(limit, random_state=seed) tok = parts_tokenizer() def encode(df): items = [] for r in df.itertuples(index=False): enc = tok(r.text, truncation=True, max_length=max_len, return_offsets_mapping=True) offs = [None if (s == e == 0) else (s, e) for s, e in enc["offset_mapping"]] offs[0] = None # CLS bio = np.zeros(len(offs), dtype=np.int64) eff = np.full(len(offs), -100, dtype=np.int64) for (s, e, typ, ef) in r.parts: first = True for i, o in enumerate(offs): if o is None or o[1] <= s or o[0] >= e: continue bio[i] = BIO.index(("B-" if first else "I-") + typ) first = False if ef is not None: eff[i] = EFFECT.index(ef) for i, o in enumerate(offs): if o is None: bio[i] = -100 items.append((enc["input_ids"], bio, eff, offs, r.text, r.parts)) return items def batches(items, size, shuffle): idx = np.random.permutation(len(items)) if shuffle else np.arange(len(items)) for s in range(0, len(idx), size): chunk = [items[i] for i in idx[s:s + size]] L = max(len(c[0]) for c in chunk) ids = torch.full((len(chunk), L), tok.pad_token_id) att = torch.zeros((len(chunk), L), dtype=torch.long) yb = torch.full((len(chunk), L), -100) ye = torch.full((len(chunk), L), -100) for j, c in enumerate(chunk): n = len(c[0]) ids[j, :n] = torch.tensor(c[0]); att[j, :n] = 1 yb[j, :n] = torch.tensor(c[1]); ye[j, :n] = torch.tensor(c[2]) yield ids.cuda(), att.cuda(), yb.cuda(), ye.cuda(), chunk t0 = time.time() train_items, test_items = encode(tr), encode(te) print(f"encoded in {time.time()-t0:.0f}s; mean tokens {np.mean([len(c[0]) for c in train_items]):.0f}", flush=True) net = PartsModel(n_tokens=len(tok)).cuda() opt = torch.optim.AdamW(param_groups(net, lr, head_lr=5e-4)) steps = epochs * ((len(train_items) + bs - 1) // bs) // accum sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=[g["lr"] for g in opt.param_groups], total_steps=max(steps, 1), pct_start=0.1, anneal_strategy="linear") @torch.no_grad() def evaluate(): net.eval() f1s, ll, n_eff = [], 0.0, 0 for ids, att, yb, ye, chunk in batches(test_items, 16, False): with torch.autocast("cuda", torch.bfloat16): lb, le = net(ids, att) lb, le = lb.float().cpu().numpy(), le.float().cpu().numpy() for j, c in enumerate(chunk): n = len(c[0]) pred = decode(lb[j, :n], le[j, :n], c[3], c[4]) f1s.append(span_f1([(p["start"], p["end"], p["type"]) for p in pred], [(s, e, t) for s, e, t, _ in c[5]])) # effectiveness log loss on the gold parts, reading the tokens of each gold part for (s, e, typ, ef) in c[5]: if ef is None: continue ti = [i for i, o in enumerate(c[3]) if o is not None and o[1] > s and o[0] < e] if not ti: continue z = le[j, ti].mean(0); z = z - z.max(); p = np.exp(z) / np.exp(z).sum() ll -= np.log(max(p[EFFECT.index(ef)], 1e-7)); n_eff += 1 net.train() return float(np.mean(f1s)), float(ll / max(n_eff, 1)) log, best, step = [], -1.0, 0 net.train() for ep in range(epochs): tot, n = 0.0, 0 for i, (ids, att, yb, ye, _) in enumerate(batches(train_items, bs, True)): with torch.autocast("cuda", torch.bfloat16): lb, le = net(ids, att) loss = (F.cross_entropy(lb.float().flatten(0, 1), yb.flatten(), ignore_index=-100) + 0.5 * F.cross_entropy(le.float().flatten(0, 1), ye.flatten(), ignore_index=-100)) / accum loss.backward() tot, n = tot + loss.item() * accum, n + 1 if (i + 1) % accum == 0: torch.nn.utils.clip_grad_norm_(net.parameters(), 1.0) opt.step(); opt.zero_grad(set_to_none=True) if step < steps - 1: sched.step() step += 1 if n % 200 == 0: print(f"epoch {ep} batch {n} loss {tot/n:.4f} {time.time()-t0:.0f}s", flush=True) f1, ll = evaluate() row = {"epoch": ep, "train_loss": round(tot / n, 4), "span_f1": round(f1, 4), "eff_logloss": round(ll, 4), "t": round(time.time() - t0)} log.append(row); print(row, flush=True) if f1 > best: best = f1 save_file({k: v.detach().cpu().contiguous() for k, v in net.state_dict().items()}, f"{run}/model.safetensors") with open(f"{run}/log.jsonl", "w") as f: f.write("\n".join(json.dumps(r) for r in log) + "\n") vol.commit() tok.save_pretrained(f"{run}/tokenizer") json.dump({"backbone": net.backbone.config._name_or_path, "bio": BIO, "effect": EFFECT, "max_len": max_len, "best_span_f1": best}, open(f"{run}/config.json", "w")) vol.commit() if os.environ.get("HF_TOKEN"): from huggingface_hub import HfApi api = HfApi(); api.create_repo(HF_REPO, exist_ok=True) api.upload_folder(folder_path=run, repo_id=HF_REPO, path_in_repo=f"runs/{tag}", commit_message=f"{tag} span f1 {best:.4f}") return {"tag": tag, "best_span_f1": round(best, 4), "minutes": round((time.time() - t0) / 60, 1)} @app.local_entrypoint() def main(tag: str = "parts-v1", epochs: int = 3, bs: int = 8, accum: int = 1, lr: float = 2e-5, max_len: int = 1024, limit: int = 0, seed: int = 0, val: int = 0): print(train.remote(tag, epochs, bs, accum, lr, max_len, limit, seed, val))