| """Local BERT fine tuning, validation checkpoint selection and held-out evaluation.""" |
| import argparse |
| import copy |
| import json |
| import sys |
| import time |
| from pathlib import Path |
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) |
| import numpy as np |
| import pandas as pd |
| import torch |
| from sklearn.feature_extraction.text import TfidfVectorizer |
| from sklearn.linear_model import LogisticRegression |
| from sklearn.pipeline import make_pipeline |
| from sklearn.metrics import accuracy_score, precision_recall_fscore_support, confusion_matrix, roc_auc_score, f1_score |
| from torch.utils.data import DataLoader, Dataset |
| from transformers import AutoTokenizer, AutoModelForSequenceClassification, get_linear_schedule_with_warmup |
| from common.io import ROOT, dump, seed_all, environment, sha256 |
|
|
|
|
| def output_directories(args): |
| """Keep the original seed-42 defaults; isolate every optional repeat run.""" |
| original_results = (ROOT / "results/fine_tuning").resolve() |
| original_model = (ROOT / "models/disaster-bert").resolve() |
| custom = args.result_dir is not None or args.model_output_dir is not None |
| base = Path(args.base_model_dir) |
| base = (base if base.is_absolute() else ROOT / base).resolve() |
| original_base = base == (ROOT / "models/bert-mini").resolve() |
| original_id = args.model_id == "google/bert_uncased_L-4_H-256_A-4" |
| if not custom and args.seed == 42 and original_base and original_id: |
| return original_results, original_model |
| if args.result_dir is None or args.model_output_dir is None: |
| raise ValueError("Repeat/custom runs require both --result-dir and --model-output-dir") |
|
|
| def resolve(value): |
| path = Path(value) |
| return (path if path.is_absolute() else ROOT / path).resolve() |
|
|
| results, model = resolve(args.result_dir), resolve(args.model_output_dir) |
| for target in (results, model): |
| for protected in (original_results, original_model): |
| if target == protected or protected in target.parents or target in protected.parents: |
| raise ValueError(f"Custom output overlaps the protected original artifacts: {target}") |
| if target.exists() and (not target.is_dir() or any(target.iterdir())): |
| raise ValueError(f"Custom output must be an empty/new directory: {target}") |
| if results == model or results in model.parents or model in results.parents: |
| raise ValueError("Result and model output directories must be separate") |
| return results, model |
|
|
| def metrics(y, pred, prob=None): |
| p, r, f, _ = precision_recall_fscore_support(y, pred, average="binary", zero_division=0) |
| m = dict(accuracy=float(accuracy_score(y, pred)), precision=float(p), recall=float(r), f1=float(f), |
| macro_f1=float(f1_score(y, pred, average="macro")), confusion_matrix=confusion_matrix(y, pred, labels=[0,1]).tolist()) |
| if prob is not None: |
| m["roc_auc"] = float(roc_auc_score(y, prob)) |
| return m |
|
|
| def bootstrap_ci(y, pred, seed=42): |
| rng = np.random.default_rng(seed) |
| scores = [] |
| y, pred = np.asarray(y), np.asarray(pred) |
| for _ in range(1000): |
| ix = rng.integers(0, len(y), len(y)) |
| scores.append(f1_score(y[ix], pred[ix], zero_division=0)) |
| return np.quantile(scores, [0.025, 0.975]).tolist() |
|
|
| class Tweets(Dataset): |
| def __init__(self, df, tokenizer, max_length): |
| self.enc = tokenizer(df.text.tolist(), truncation=True, max_length=max_length, padding="max_length") |
| self.labels = df.target.tolist() |
| def __len__(self): |
| return len(self.labels) |
| def __getitem__(self, i): |
| return {**{k: torch.tensor(v[i]) for k,v in self.enc.items()}, "labels": torch.tensor(self.labels[i])} |
|
|
| @torch.inference_mode() |
| def predict(model, loader, device): |
| model.eval() |
| probs = [] |
| for batch in loader: |
| batch = {k:v.to(device) for k,v in batch.items()} |
| probs.extend(model(**batch).logits.softmax(-1)[:,1].cpu().tolist()) |
| return np.asarray(probs) |
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--epochs", type=int, default=5) |
| ap.add_argument("--batch-size", type=int, default=32) |
| ap.add_argument("--lr", type=float, default=3e-5) |
| ap.add_argument("--max-length", type=int, default=96) |
| ap.add_argument("--seed", type=int, default=42) |
| ap.add_argument("--result-dir", help="Optional fresh result directory; relative paths use the assignment root") |
| ap.add_argument("--model-output-dir", help="Optional fresh checkpoint directory; relative paths use the assignment root") |
| ap.add_argument("--base-model-dir", default="models/bert-mini", help="Local base model/tokenizer directory; relative paths use the assignment root") |
| ap.add_argument("--model-id", default="google/bert_uncased_L-4_H-256_A-4", help="Base model identity recorded in the results") |
| args = ap.parse_args() |
| try: |
| out, model_out = output_directories(args) |
| except ValueError as error: |
| ap.error(str(error)) |
| seed_all(args.seed) |
| device = "cuda" if torch.cuda.is_available() else "cpu" |
| out.mkdir(parents=True, exist_ok=True) |
| split_hashes = {name: sha256(ROOT / f"data/splits/{name}.csv") for name in ["train", "validation", "test"]} |
| parts = {name: pd.read_csv(ROOT / f"data/splits/{name}.csv", keep_default_na=False) for name in ["train","validation","test"]} |
| train, val, test = (parts[n] for n in ["train","validation","test"]) |
| base_model = Path(args.base_model_dir) |
| base_model = base_model if base_model.is_absolute() else ROOT / base_model |
| tok = AutoTokenizer.from_pretrained(base_model, local_files_only=True) |
| loaders = {n:DataLoader(Tweets(df,tok,args.max_length), batch_size=args.batch_size, shuffle=n=="train", num_workers=0) for n,df in parts.items()} |
| model = AutoModelForSequenceClassification.from_pretrained(base_model, num_labels=2, |
| id2label={0:"Not disaster",1:"Disaster"}, label2id={"Not disaster":0,"Disaster":1}, local_files_only=True).to(device) |
| optimizer = torch.optim.AdamW(model.parameters(),lr=args.lr,weight_decay=0.01) |
| steps = args.epochs * len(loaders["train"]) |
| scheduler = get_linear_schedule_with_warmup(optimizer, int(0.1*steps), steps) |
| history, best_score, best_state, best_epoch = [], -1, None, None |
| start = time.perf_counter() |
| for epoch in range(1,args.epochs+1): |
| model.train() |
| total_loss = 0 |
| for batch in loaders["train"]: |
| batch = {k:v.to(device) for k,v in batch.items()} |
| optimizer.zero_grad(set_to_none=True) |
| loss = model(**batch).loss |
| loss.backward() |
| torch.nn.utils.clip_grad_norm_(model.parameters(),1.0) |
| optimizer.step(); scheduler.step() |
| total_loss += loss.item()*len(batch["labels"]) |
| prob = predict(model,loaders["validation"],device) |
| result = dict(epoch=epoch, train_loss=total_loss/len(train), validation=metrics(val.target,prob>=0.5,prob)) |
| history.append(result) |
| dump(out / "history.json",history) |
| print(json.dumps(result),flush=True) |
| if result["validation"]["f1"] > best_score: |
| best_score=result["validation"]["f1"]; best_epoch=epoch |
| best_state={k:v.detach().cpu().clone() for k,v in model.state_dict().items()} |
| model.load_state_dict(best_state) |
| model.save_pretrained(model_out) |
| tok.save_pretrained(model_out) |
| prob=predict(model,loaders["test"],device) |
| pred=(prob>=0.5).astype(int) |
| test_results=metrics(test.target,pred,prob) |
| test_results["bootstrap_f1_ci95"]=bootstrap_ci(test.target,pred) |
| baselines={} |
| majority=np.full(len(test),int(train.target.mode()[0])) |
| baselines["majority"]=metrics(test.target,majority) |
| pipeline=make_pipeline(TfidfVectorizer(lowercase=True,ngram_range=(1,2),min_df=2,sublinear_tf=True), |
| LogisticRegression(C=2.0,max_iter=1000,random_state=args.seed)) |
| pipeline.fit(train.text,train.target) |
| base_prob=pipeline.predict_proba(test.text)[:,1] |
| base_pred=(base_prob>=0.5).astype(int) |
| baselines["tfidf_logistic"]=metrics(test.target,base_pred,base_prob) |
| result=dict(config=vars(args),device=device,environment=environment(),model=args.model_id, |
| data_split_sha256=split_hashes, model_output_dir=str(model_out), base_model_dir=str(base_model.resolve()), |
| trainable_parameters=sum(p.numel() for p in model.parameters()),best_epoch=best_epoch, |
| best_validation_f1=best_score,test=test_results,baselines=baselines,seconds=time.perf_counter()-start, |
| selection="best validation F1; fixed threshold 0.5; test evaluated after selection") |
| dump(out / "metrics.json",result) |
| pd.DataFrame(dict(id=test.id,text=test.text,gold=test.target,prediction=pred,prob_disaster=prob, |
| tfidf_prediction=base_pred,tfidf_prob=base_prob)).to_csv(out / "test_predictions.csv",index=False,encoding="utf-8-sig") |
| errors=test.assign(prediction=pred,prob_disaster=prob) |
| errors=errors[errors.target!=errors.prediction] |
| errors["confidence"]=np.maximum(errors.prob_disaster,1-errors.prob_disaster) |
| errors.sort_values("confidence",ascending=False).to_csv(out / "error_cases.csv",index=False,encoding="utf-8-sig") |
| print(json.dumps(result,ensure_ascii=False,indent=2),flush=True) |
|
|
| if __name__ == "__main__": |
| main() |
|
|
