File size: 6,255 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 | """Train the six-score head on ELLIPSE, on one Modal A10G.
MODAL_PROFILE=labs3 modal run scripts/train_scores.py --tag scores-v1
MODAL_PROFILE=labs3 modal run scripts/train_scores.py --tag smoke --limit 64 --epochs 1
Metric: MCRMSE over the six traits on the ELLIPSE test split (the Kaggle winners' single models: ~0.45 large, ~0.46-0.47
base). Weights and the log go to /data/runs/<tag> and HF nonsodev/essay-coach runs/<tag>.
"""
import os
import modal
app = modal.App("essay-coach-train")
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 []
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"
@app.function(image=image, gpu="A10G", volumes={"/data": vol, "/root/.cache/huggingface": cache}, timeout=4 * 3600,
secrets=SECRETS, cpu=4, memory=16384)
def train(tag: str, epochs: int, bs: int, accum: int, lr: float, max_len: int, limit: int, seed: int):
import json
import os
import time
import numpy as np
import torch
from huggingface_hub import HfApi
from safetensors.torch import save_file
from coach.data import TRAITS, ellipse, with_breaks
from coach.model import ScoreModel, param_groups, tokenizer
torch.manual_seed(seed)
np.random.seed(seed)
run = f"/data/runs/{tag}"
os.makedirs(run, exist_ok=True)
tr, te = ellipse()
if limit:
tr, te = tr.sample(limit, random_state=seed), te.sample(limit, random_state=seed)
print(f"train {len(tr)} test {len(te)}", flush=True)
tok = tokenizer()
def encode(df):
enc = tok([with_breaks(t) for t in df.text], truncation=True, max_length=max_len)
y = df[TRAITS].values.astype(np.float32)
return [(enc["input_ids"][i], y[i]) for i in range(len(df))]
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(ids) for ids, _ in chunk)
ids = torch.full((len(chunk), L), tok.pad_token_id)
att = torch.zeros((len(chunk), L), dtype=torch.long)
for j, (seq, _) in enumerate(chunk):
ids[j, :len(seq)] = torch.tensor(seq)
att[j, :len(seq)] = 1
yield ids.cuda(), att.cuda(), torch.tensor(np.stack([y for _, y in chunk])).cuda()
train_items, test_items = encode(tr), encode(te)
net = ScoreModel(n_tokens=len(tok)).cuda()
opt = torch.optim.AdamW(param_groups(net, lr, head_lr=1e-3))
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()
preds, ys = [], []
for ids, att, y in batches(test_items, 16, False):
with torch.autocast("cuda", torch.bfloat16):
preds.append(net(ids, att).float().cpu())
ys.append(y.cpu())
net.train()
p, y = torch.cat(preds).numpy(), torch.cat(ys).numpy()
rmse = np.sqrt(((p - y) ** 2).mean(0))
return float(rmse.mean()), {t: round(float(r), 4) for t, r in zip(TRAITS, rmse)}, p
hf = bool(os.environ.get("HF_TOKEN"))
if hf:
api = HfApi()
api.create_repo(HF_REPO, exist_ok=True)
log, best, t0, step = [], 9.0, time.time(), 0
net.train()
for ep in range(epochs):
tot, n = 0.0, 0
for i, (ids, att, y) in enumerate(batches(train_items, bs, True)):
with torch.autocast("cuda", torch.bfloat16):
loss = torch.nn.functional.mse_loss(net(ids, att).float(), y) / accum
loss.backward()
tot, n = tot + loss.item() * accum, n + 1
if (i + 1) % accum == 0:
torch.nn.utils.clip_grad_norm_(net.parameters(), 10.0)
opt.step()
opt.zero_grad(set_to_none=True)
if step < steps - 1:
sched.step()
step += 1
if n % 100 == 0:
print(f"epoch {ep} batch {n} loss {tot/n:.4f} {time.time()-t0:.0f}s", flush=True)
mcrmse, per, p = evaluate()
row = {"epoch": ep, "train_mse": round(tot / n, 4), "mcrmse": round(mcrmse, 4), **per, "t": round(time.time() - t0)}
log.append(row)
print(row, flush=True)
if mcrmse < best:
best = mcrmse
save_file({k: v.detach().cpu().contiguous() for k, v in net.state_dict().items()}, f"{run}/model.safetensors")
te.assign(**{f"pred_{t}": p[:, j] for j, t in enumerate(TRAITS)}).to_csv(f"{run}/test_preds.csv", index=False)
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, "traits": TRAITS, "max_len": max_len, "best_mcrmse": best},
open(f"{run}/config.json", "w"))
vol.commit()
if hf:
api.upload_folder(folder_path=run, repo_id=HF_REPO, path_in_repo=f"runs/{tag}", commit_message=f"{tag} mcrmse {best:.4f}")
return {"tag": tag, "best_mcrmse": round(best, 4), "minutes": round((time.time() - t0) / 60, 1)}
@app.local_entrypoint()
def main(tag: str = "scores-v1", epochs: int = 4, bs: int = 8, accum: int = 1, lr: float = 2e-5, max_len: int = 1024,
limit: int = 0, seed: int = 0):
print(train.remote(tag, epochs, bs, accum, lr, max_len, limit, seed))
|