code/punct/train_punct.py

"""Fine-tune an encoder for punctuation restoration (token classification, labels of punct/common.py).
python punct/train_punct.py --base jhu-clsp/mmBERT-small --data DATA_DIR --out OUT_DIR/model
Each word's label sits on its last token; other tokens are ignored. A fraction of the inputs is lowercased (ASR-style).
Checkpoint selection: punctuation F1 on the validation windows (cased and lowercased mean).
"""
import argparse, json, os, random, sys, time
import numpy as np, pandas as pd, torch
from transformers import AutoModelForTokenClassification, AutoTokenizer, get_cosine_schedule_with_warmup
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from common import LABELS, L2I, lower_keep_len

ap = argparse.ArgumentParser()
ap.add_argument("--base", default="jhu-clsp/mmBERT-small")
ap.add_argument("--data", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--max_len", type=int, default=512)
ap.add_argument("--tokens_per_batch", type=int, default=65536)
ap.add_argument("--lr", type=float, default=5e-5)
ap.add_argument("--epochs", type=float, default=1)
ap.add_argument("--warmup", type=float, default=0.05)
ap.add_argument("--lower", type=float, default=0.4)
ap.add_argument("--max_train", type=int, default=0)
ap.add_argument("--max_val", type=int, default=6000)
ap.add_argument("--evals", type=int, default=8, help="validation checks over the run")
ap.add_argument("--seed", type=int, default=0)
args = ap.parse_args()
random.seed(args.seed); np.random.seed(args.seed); torch.manual_seed(args.seed)
os.makedirs(args.out, exist_ok=True)

tr = pd.read_parquet(f"{args.data}/train.parquet")
va = pd.read_parquet(f"{args.data}/val.parquet").sample(frac=1.0, random_state=0).head(args.max_val)
if args.max_train:
    tr = tr.sample(min(len(tr), args.max_train), random_state=args.seed)
tok = AutoTokenizer.from_pretrained(args.base)
model = AutoModelForTokenClassification.from_pretrained(args.base, num_labels=len(LABELS), id2label=dict(enumerate(LABELS)),
                                                        label2id=L2I).cuda()


def encode(df, lower):
    texts = [lower_keep_len(t) if lw else t for t, lw in zip(df.text, lower)]
    out = []
    for s in range(0, len(texts), 20000):
        enc = tok(texts[s:s + 20000], truncation=True, max_length=args.max_len, return_offsets_mapping=True)
        for ids, offs, ends, labs in zip(enc["input_ids"], enc["offset_mapping"], df.ends.iloc[s:s + 20000], df.labs.iloc[s:s + 20000]):
            at = dict(zip(ends.tolist(), labs.tolist()))
            out.append((ids, [at[e] if e > b and e in at else -100 for b, e in offs]))
    return out


t0 = time.time()
rng = np.random.RandomState(args.seed)
tr_enc = encode(tr, rng.rand(len(tr)) < args.lower)
va_c, va_l = encode(va, np.zeros(len(va), bool)), encode(va, np.ones(len(va), bool))
print(f"train {len(tr_enc)} val {len(va_c)} encoded in {time.time() - t0:.0f}s", flush=True)


def batches(data, shuffle):
    idx = list(range(len(data)))
    if shuffle:
        random.shuffle(idx)
    out = []
    for s in range(0, len(idx), 8192):
        part = sorted(idx[s:s + 8192], key=lambda i: len(data[i][0])); cur, mx = [], 0
        for i in part:
            L = len(data[i][0])
            if cur and max(mx, L) * (len(cur) + 1) > args.tokens_per_batch:
                out.append(cur); cur, mx = [], 0
            cur.append(i); mx = max(mx, L)
        if cur:
            out.append(cur)
    if shuffle:
        random.shuffle(out)
    return out


def collate(data, b):
    L = max(len(data[i][0]) for i in b)
    ids = torch.full((len(b), L), tok.pad_token_id, dtype=torch.long); att = torch.zeros((len(b), L), dtype=torch.long)
    lab = torch.full((len(b), L), -100, dtype=torch.long)
    for j, i in enumerate(b):
        n = len(data[i][0]); ids[j, :n] = torch.tensor(data[i][0]); att[j, :n] = 1; lab[j, :n] = torch.tensor(data[i][1])
    return ids.cuda(), att.cuda(), lab.cuda()


@torch.no_grad()
def evaluate(data):
    model.eval(); conf = np.zeros((len(LABELS), len(LABELS)), np.int64)
    for b in batches(data, False):
        ids, att, lab = collate(data, b)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            pr = model(input_ids=ids, attention_mask=att).logits.argmax(-1)
        m = lab >= 0
        np.add.at(conf, (lab[m].cpu().numpy(), pr[m].cpu().numpy()), 1)
    model.train()
    tp = sum(conf[i, i] for i in range(1, 6)); fp = conf[:, 1:].sum() - tp; fn = conf[1:, :].sum() - tp
    f1c = {LABELS[i]: round(2 * conf[i, i] / max(1, conf[i, :].sum() + conf[:, i].sum()), 4) for i in range(1, 6)}
    return dict(punct_f1=round(2 * tp / max(1, 2 * tp + fp + fn), 4), **f1c)


nb = len(batches(tr_enc, True)); total = int(nb * args.epochs)
opt = torch.optim.AdamW(model.parameters(), lr=args.lr, betas=(0.9, 0.98), eps=1e-6, weight_decay=0.01, fused=True)
sch = get_cosine_schedule_with_warmup(opt, int(args.warmup * total), total)
every = max(1, total // args.evals)
print(f"steps/epoch {nb} total {total} eval every {every}", flush=True)
step, best, log, t0 = 0, -1, [], time.time()
model.train()
while step < total:
    for b in batches(tr_enc, True):
        ids, att, lab = collate(tr_enc, b)
        with torch.autocast("cuda", dtype=torch.bfloat16):
            loss = torch.nn.functional.cross_entropy(model(input_ids=ids, attention_mask=att).logits.float().view(-1, len(LABELS)),
                                                     lab.view(-1), ignore_index=-100)
        loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step(); sch.step(); opt.zero_grad(set_to_none=True)
        step += 1
        if step % 200 == 0:
            print(f"step {step}/{total} loss {loss.item():.4f} {time.time() - t0:.0f}s", flush=True)
        if step % every == 0 or step == total:
            c, l = evaluate(va_c), evaluate(va_l); score = (c["punct_f1"] + l["punct_f1"]) / 2
            log.append(dict(step=step, cased=c, lower=l, score=score)); print(json.dumps(log[-1]), flush=True)
            if score > best:
                best = score; model.save_pretrained(args.out); tok.save_pretrained(args.out)
        if step >= total:
            break
json.dump(dict(args=vars(args), log=log, best=best), open(f"{args.out}/train_log.json", "w"), indent=1)
print("best", best, flush=True)