File size: 8,689 Bytes
e1679fa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 | """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))
|