essay-coach / code /scripts /train_parts.py

code snapshot + README

e1679fa verified 15 days ago

8.69 kB

"""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))

	"""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))