File size: 6,311 Bytes
9456b7b | 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 | """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)
|