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