| |
| """Fine-tune an Arabic BERT for one label. Usage: train_encoder.py <label: activity|ttype|sector> <hf-model> <out-dir> [epochs]""" |
| import json, os, sys, time, random |
| import numpy as np, pandas as pd, torch |
| from sklearn.metrics import accuracy_score, f1_score |
| from sklearn.model_selection import train_test_split |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification, get_linear_schedule_with_warmup |
| sys.path.insert(0, os.path.dirname(__file__)); from common import * |
| label, mname, out = sys.argv[1], sys.argv[2], sys.argv[3]; EPOCHS = int(sys.argv[4]) if len(sys.argv) > 4 else 6; PATIENCE = int(sys.argv[5]) if len(sys.argv) > 5 else 2 |
| COL = {"activity": "activity_label", "ttype": "ttype", "sector": "sector"}[label] |
| random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED) |
| dev = "mps" if torch.backends.mps.is_available() else "cpu" |
| tr = pd.read_csv(f"{ROOT}/data/published/train.csv", dtype=str).fillna(""); te = pd.read_csv(f"{ROOT}/data/published/test.csv", dtype=str).fillna("") |
| stats = json.load(open(f"{ROOT}/data/stats.json")); keep_t = {k for k, v in stats["ttype_counts"] if k != "other"} |
| for d in (tr, te): d["ttype"] = d.ttype.map(lambda x: x if x in keep_t else "other") |
| classes = sorted(tr[COL].unique()); c2i = {c: i for i, c in enumerate(classes)} |
| tr_, va = train_test_split(tr, test_size=0.1, random_state=SEED, stratify=tr[COL] if tr[COL].value_counts().min() >= 2 else None) |
| tok = AutoTokenizer.from_pretrained(mname) |
| model = AutoModelForSequenceClassification.from_pretrained(mname, num_labels=len(classes), id2label=dict(enumerate(classes)), label2id=c2i).to(dev) |
| def enc(df): return [make_input(t, a) for t, a in zip(df.title, df.agency)], [c2i.get(y, c2i.get("other", 0)) for y in df[COL]] |
| def batches(texts, ys, bs, shuffle): |
| idx = list(range(len(texts))); |
| if shuffle: random.shuffle(idx) |
| for i in range(0, len(idx), bs): |
| j = idx[i:i + bs]; b = tok([texts[k] for k in j], truncation=True, max_length=128, padding=True, pad_to_multiple_of=32, return_tensors="pt") |
| yield {k: v.to(dev) for k, v in b.items()}, torch.tensor([ys[k] for k in j]).to(dev) |
| def predict(texts, ys, top3=False): |
| model.eval(); P = []; T = [] |
| with torch.no_grad(): |
| for b, _ in batches(texts, ys, 64, False): |
| lg = model(**b).logits; P += lg.argmax(-1).cpu().tolist(); T += lg.topk(min(3, lg.shape[-1]), -1).indices.cpu().tolist() |
| if dev == "mps": torch.mps.empty_cache() |
| return (P, T) if top3 else P |
| Xt, yt = enc(tr_); Xv, yv = enc(va); Xe, ye = enc(te) |
| cnt = np.bincount(yt, minlength=len(classes)).astype(float); w = torch.tensor((cnt.max() / np.maximum(cnt, 1)) ** 0.5, dtype=torch.float).to(dev) |
| BS = 32; steps = EPOCHS * ((len(Xt) + BS - 1) // BS) |
| opt = torch.optim.AdamW(model.parameters(), lr=4e-5 if len(classes) < 10 else 5e-5, weight_decay=0.01) |
| sch = get_linear_schedule_with_warmup(opt, int(0.06 * steps), steps); lossf = torch.nn.CrossEntropyLoss(weight=w) |
| best, bad, t0, hist = -1, 0, time.time(), []; step_i = 0 |
| os.makedirs(out, exist_ok=True) |
| for ep in range(EPOCHS): |
| model.train(); tl = 0 |
| for b, y in batches(Xt, yt, BS, True): |
| loss = lossf(model(**b).logits, y); loss.backward(); torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0); opt.step(); sch.step(); opt.zero_grad(set_to_none=True); tl += loss.item(); step_i += 1 |
| if dev == "mps" and step_i % 25 == 0: torch.mps.empty_cache() |
| f1v = f1_score(yv, predict(Xv, yv), average="macro"); hist.append((ep + 1, round(tl, 2), round(f1v, 4))) |
| print(f"epoch {ep+1} train_loss_sum={tl:.1f} val_macroF1={f1v:.4f} elapsed={time.time()-t0:.0f}s", flush=True) |
| if f1v > best: best, bad = f1v, 0; model.save_pretrained(out); tok.save_pretrained(out) |
| else: |
| bad += 1 |
| if bad >= PATIENCE: print("early stop"); break |
| train_s = time.time() - t0 |
| model = AutoModelForSequenceClassification.from_pretrained(out).to(dev) |
| yp, t3 = predict(Xe, ye, True); top3 = float(np.mean([y in t for y, t in zip(ye, t3)])) |
| res = {"label": label, "base_model": mname, "device": dev, "accuracy": accuracy_score(ye, yp), "macro_f1": f1_score(ye, yp, average="macro"), |
| "top3_accuracy": top3, "best_val_macro_f1": best, "patience": PATIENCE, "max_epochs": EPOCHS, "epochs_run": len(hist), "history": hist, "train_seconds": round(train_s, 1), "n_classes": len(classes), |
| "size_mb": round(sum(os.path.getsize(os.path.join(out, f)) for f in os.listdir(out)) / 1e6, 1)} |
| json.dump(res, open(f"{ROOT}/outputs/encoder_{label}_metrics.json", "w"), indent=1) |
| pd.DataFrame({"ref": te.ref, "true": [classes[i] for i in ye], "pred": [classes[i] for i in yp]}).to_csv(f"{ROOT}/outputs/pred_encoder_{label}.csv", index=False) |
| print(res) |
|
|
