在本教程中,我们使用斯坦福 NLP IMDb 大型电影评论数据集构建一个端到端的情感分析工作流,并将经典机器学习与参数高效的 Transformer 微调进行对比。我们首先建立一个可复现的环境,并在训练一个强 TF-IDF 与逻辑回归基线之前,对数据集进行类别顺序、评论长度偏斜、重复泄漏和预处理伪影方面的审计。随后,我们通过 PEFT 使用 LoRA 微调 DistilBERT,使用准确率、宏平均 F1、ROC-AUC、混淆矩阵和 ROC 曲线进行评估,并通过期望校准误差和可靠性分析来考察阈值选择与概率校准。除了总体指标之外,我们还研究置信错误、不同评论长度下的性能、词级遮挡显著性,以及头部截断与尾部截断,以理解模型如何得出预测,以及长上下文限制在何处影响性能。最后,我们使用未标注的 IMDb 划分进行基于置信度的伪标注,将得到的半监督模型与我们的基线进行比较,并保存合并后的 Transformer 以供可复用的情感推理。
import importlib.util, subprocess, sys, os, time, random, warnings, inspect, hashlib
warnings.filterwarnings("ignore")
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["WANDB_DISABLED"] = "true"
_REQUIRED = {
"transformers": "transformers",
"datasets": "datasets",
"peft": "peft",
"accelerate": "accelerate",
"sklearn": "scikit-learn",
}
_missing = [pkg for mod, pkg in _REQUIRED.items() if importlib.util.find_spec(mod) is None]
if _missing:
print(f"Installing: {', '.join(_missing)} ...")
subprocess.run([sys.executable, "-m", "pip", "install", "-q", *_missing], check=True)
print("Done. (If imports fail below, restart the runtime and re-run.)\n")
import numpy as np
import pandas as pd
import torch
import matplotlib.pyplot as plt
from datasets import load_dataset
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, f1_score, roc_auc_score,
classification_report, confusion_matrix, roc_curve)
from transformers import (AutoTokenizer, AutoModelForSequenceClassification,
TrainingArguments, Trainer, DataCollatorWithPadding,
EarlyStoppingCallback, set_seed)
from peft import LoraConfig, get_peft_model, TaskType
def _disable_torchao_probe():
patched = []
try:
import peft.import_utils as _piu
_piu.is_torchao_available = lambda: False
patched.append("peft.import_utils")
except Exception:
pass
for _name, _mod in list(sys.modules.items()):
if _name.startswith("peft") and hasattr(_mod, "is_torchao_available"):
_mod.is_torchao_available = lambda: False
patched.append(_name)
return patched
try:
import torchao as _tao
_v = getattr(_tao, "__version__", "?")
if tuple(int(x) for x in _v.split(".")[:2]) < (0, 16):
print(f"[compat] torchao {_v} < 0.16 -> disabling PEFT's torchao probe: "
f"{', '.join(_disable_torchao_probe())}")
except Exception:
_disable_torchao_probe()
SEED = 42
MODEL_NAME = "distilbert-base-uncased"
MAX_LEN = 256
N_TRAIN = 5000
N_EVAL = 2000
N_UNSUP = 3000
EPOCHS = 2
BATCH = 16
LR = 3e-4
FULL_RUN = False
if FULL_RUN:
N_TRAIN, N_EVAL, EPOCHS = 25000, 25000, 3
set_seed(SEED); random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print("=" * 79)
print(f"device={DEVICE} | torch={torch.__version__} | "
f"gpu={torch.cuda.get_device_name(0) if DEVICE=='cuda' else 'n/a'}")
print("=" * 79)
t0 = time.time()
raw = load_dataset("stanfordnlp/imdb")
print(raw, f"\nloaded in {time.time()-t0:.1f}s\n")
print("--- example (truncated) ---")
print("label:", raw["train"][0]["label"], "|", raw["train"][0]["text"][:300], "...\n")
first_labels = np.array(raw["train"]["label"][:5])
last_labels = np.array(raw["train"]["label"][-5:])
print(f"TRAP #1 - split ordering: first 5 labels {first_labels}, "
f"last 5 labels {last_labels} -> ALWAYS shuffle before subsampling.")
train_full = raw["train"].shuffle(seed=SEED)
test_full = raw["test"].shuffle(seed=SEED)
train_ds = train_full.select(range(min(N_TRAIN, len(train_full))))
eval_ds = test_full.select(range(min(N_EVAL, len(test_full))))
print(f" after shuffle+subsample: train balance = "
f"{np.bincount(train_ds['label'])}, eval balance = {np.bincount(eval_ds['label'])}")
lens = np.array([len(t.split()) for t in train_full["text"]])
q = np.percentile(lens, [50, 75, 90, 95, 99])
print(f"\nTRAP #2 - length (words): median={q[0]:.0f} p75={q[1]:.0f} p90={q[2]:.0f} "
f"p95={q[3]:.0f} p99={q[4]:.0f} max={lens.max()}")
print(f" ~{(lens > MAX_LEN*0.75).mean()*100:.1f}% of reviews exceed MAX_LEN={MAX_LEN} "
f"tokens (rough words->tokens factor 1.3). Section 9 measures what that costs.")
h_tr = {hashlib.md5(t.encode()).hexdigest() for t in raw["train"]["text"]}
h_te = {hashlib.md5(t.encode()).hexdigest() for t in raw["test"]["text"]}
print(f"\nTRAP #3 - leakage: {len(h_tr & h_te)} exact duplicate reviews across "
f"train/test; {len(raw['train'])-len(h_tr)} dupes inside train itself.")
def clean(t):
return t.replace("<br />", " ").replace("<br/>", " ").strip()
plt.figure(figsize=(11, 3.2))
plt.subplot(1, 2, 1)
plt.hist(np.clip(lens, 0, 1000), bins=60)
plt.axvline(MAX_LEN, ls="--", color="k", label=f"MAX_LEN={MAX_LEN}")
plt.title("Review length (words, clipped at 1000)"); plt.legend()
plt.subplot(1, 2, 2)
plt.bar(["neg", "pos"], np.bincount(raw["train"]["label"]))
plt.title("Train class balance (perfectly balanced)")
plt.tight_layout(); plt.show()
我们配置 Colab 环境,安装所需库,应用 PEFT–torchao 兼容性修复,并设置确定性随机种子以实现可复现的实验。我们加载斯坦福 IMDb 数据集,对训练集和测试集进行打乱和子采样,并检查类别平衡、评论长度分布、重复泄漏和 HTML 伪影。我们还将评论长度和标签频率可视化,以便在构建任何模型之前理解数据集结构。
print("\n" + "=" * 79 + "\n3. TF-IDF BASELINE\n" + "=" * 79)
Xtr = [clean(t) for t in train_ds["text"]]; ytr = np.array(train_ds["label"])
Xte = [clean(t) for t in eval_ds["text"]]; yte = np.array(eval_ds["label"])
t0 = time.time()
tfidf_clf = make_pipeline(
TfidfVectorizer(ngram_range=(1, 2), min_df=2, max_features=300_000,
sublinear_tf=True, strip_accents="unicode"),
LogisticRegression(C=8.0, max_iter=2000, n_jobs=-1),
)
tfidf_clf.fit(Xtr, ytr)
p_tfidf = tfidf_clf.predict_proba(Xte)[:, 1]
acc_tfidf = accuracy_score(yte, p_tfidf > 0.5)
auc_tfidf = roc_auc_score(yte, p_tfidf)
print(f"trained in {time.time()-t0:.1f}s -> acc={acc_tfidf:.4f} auc={auc_tfidf:.4f}")
vec, lr = tfidf_clf.steps[0][1], tfidf_clf.steps[1][1]
feats, coefs = np.array(vec.get_feature_names_out()), lr.coef_[0]
order = np.argsort(coefs)
print("\nmost NEGATIVE n-grams:", ", ".join(feats[order[:12]]))
print("most POSITIVE n-grams:", ", ".join(feats[order[-12:]][::-1]))
print("\n" + "=" * 79 + "\n4. LoRA FINE-TUNING\n" + "=" * 79)
tok = AutoTokenizer.from_pretrained(MODEL_NAME)
def tokenize(batch):
return tok([clean(t) for t in batch["text"]], truncation=True, max_length=MAX_LEN)
tr_tok = (train_ds.map(tokenize, batched=True, remove_columns=["text"])
.rename_column("label", "labels"))
ev_tok = (eval_ds.map(tokenize, batched=True, remove_columns=["text"])
.rename_column("label", "labels"))
base = AutoModelForSequenceClassification.from_pretrained(
MODEL_NAME, num_labels=2,
id2label={0: "NEGATIVE", 1: "POSITIVE"},
label2id={"NEGATIVE": 0, "POSITIVE": 1},
)
lora_cfg = LoraConfig(
task_type=TaskType.SEQ_CLS,
r=16, lora_alpha=32, lora_dropout=0.05,
target_modules=["q_lin", "v_lin"],
modules_to_save=["pre_classifier", "classifier"],
)
try:
model = get_peft_model(base, lora_cfg)
except ImportError as e:
_disable_torchao_probe()
print(f"[compat] retrying after backend probe failure: {e}")
model = get_peft_model(base, lora_cfg)
model.print_trainable_parameters()
def compute_metrics(eval_pred):
logits, labels = eval_pred
probs = torch.softmax(torch.tensor(logits), dim=-1).numpy()[:, 1]
preds = (probs > 0.5).astype(int)
return {"accuracy": accuracy_score(labels, preds),
"f1_macro": f1_score(labels, preds, average="macro"),
"roc_auc": roc_auc_score(labels, probs)}
_ta = inspect.signature(TrainingArguments.__init__).parameters
_eval_key = "eval_strategy" if "eval_strategy" in _ta else "evaluation_strategy"
ta_kwargs = dict(
output_dir="./imdb_lora", learning_rate=LR,
per_device_train_batch_size=BATCH, per_device_eval_batch_size=BATCH * 2,
num_train_epochs=EPOCHS, weight_decay=0.01, warmup_ratio=0.06,
logging_steps=50, save_strategy="epoch", save_total_limit=1,
load_best_model_at_end=True, metric_for_best_model="accuracy",
fp16=(DEVICE == "cuda"), report_to="none", seed=SEED,
)
ta_kwargs[_eval_key] = "epoch"
_tr = inspect.signature(Trainer.__init__).parameters
_tok_key = "processing_class" if "processing_class" in _tr else "tokenizer"
trainer = Trainer(
model=model, args=TrainingArguments(**ta_kwargs),
train_dataset=tr_tok, eval_dataset=ev_tok,
data_collator=DataCollatorWithPadding(tok),
compute_metrics=compute_metrics,
callbacks=[EarlyStoppingCallback(early_stopping_patience=2)],
**{_tok_key: tok},
)
t0 = time.time()
trainer.train()
print(f"\nfine-tuned in {(time.time()-t0)/60:.1f} min")
我们训练一个强 TF-IDF 与逻辑回归基线,并检查最具影响力的正负 n-gram,以建立一个可解释的参考点。然后,我们对 IMDb 评论进行分词,并配置带 LoRA 适配器的 DistilBERT,该适配器仅更新一小部分模型参数,同时基本保持主干网络冻结。我们使用 Hugging Face Trainer,配合动态填充、早停、混合精度和多种评估指标,高效地微调该 Transformer。
print("\n" + "=" * 79 + "\n5. EVALUATION\n" + "=" * 79)
pred_out = trainer.predict(ev_tok)
p_lora = torch.softmax(torch.tensor(pred_out.predictions), dim=-1).numpy()[:, 1]
y_true = np.array(pred_out.label_ids)
yhat = (p_lora > 0.5).astype(int)
print(classification_report(y_true, yhat, target_names=["neg", "pos"], digits=4))
cm = confusion_matrix(y_true, yhat)
fig, ax = plt.subplots(1, 2, figsize=(11, 4))
ax[0].imshow(cm, cmap="Blues")
for i in range(2):
for j in range(2):
ax[0].text(j, i, cm[i, j], ha="center", va="center", fontsize=14)
ax[0].set_xticks([0, 1], ["pred neg", "pred pos"])
ax[0].set_yticks([0, 1], ["true neg", "true pos"]); ax[0].set_title("Confusion matrix")
for name, p in [("TF-IDF", p_tfidf), ("DistilBERT+LoRA", p_lora)]:
fpr, tpr, _ = roc_curve(y_true, p)
ax[1].plot(fpr, tpr, label=f"{name} (AUC={roc_auc_score(y_true, p):.4f})")
ax[1].plot([0, 1], [0, 1], "k--", lw=0.8)
ax[1].set_xlabel("FPR"); ax[1].set_ylabel("TPR"); ax[1].set_title("ROC"); ax[1].legend()
plt.tight_layout(); plt.show()
print("\n" + "=" * 79 + "\n6. THRESHOLD & CALIBRATION\n" + "=" * 79)
ths = np.linspace(0.05, 0.95, 91)
accs = [(y_true == (p_lora > t)).mean() for t in ths]
best_t = ths[int(np.argmax(accs))]
print(f"[email protected] = {accs[45]:.4f} | best threshold = {best_t:.2f} -> acc = {max(accs):.4f}")
def expected_calibration_error(probs, labels, n_bins=10):
"""ECE: |confidence - accuracy| averaged over confidence bins."""
conf = np.maximum(probs, 1 - probs)
correct = (probs > 0.5).astype(int) == labels
bins = np.linspace(0, 1, n_bins + 1)
ece, xs, ys = 0.0, [], []
for lo, hi in zip(bins[:-1], bins[1:]):
m = (conf > lo) & (conf <= hi)
if m.sum() == 0:
continue
ece += m.mean() * abs(conf[m].mean() - correct[m].mean())
xs.append(conf[m].mean()); ys.append(correct[m].mean())
return ece, np.array(xs), np.array(ys)
ece, cx, cy = expected_calibration_error(p_lora, y_true)
print(f"Expected Calibration Error = {ece:.4f} (0 = perfectly calibrated)")
plt.figure(figsize=(9, 3.2))
plt.subplot(1, 2, 1); plt.plot(ths, accs); plt.axvline(best_t, ls="--", color="r")
plt.xlabel("threshold"); plt.ylabel("accuracy"); plt.title("Threshold sweep")
plt.subplot(1, 2, 2); plt.plot([0.5, 1], [0.5, 1], "k--", lw=0.8)
plt.plot(cx, cy, "o-"); plt.xlabel("mean confidence"); plt.ylabel("empirical accuracy")
plt.title(f"Reliability diagram (ECE={ece:.3f})")
plt.tight_layout(); plt.show()
我们使用分类指标、混淆矩阵和 ROC 曲线来评估微调后的 DistilBERT-LoRA 模型,同时将其 ROC-AUC 性能与 TF-IDF 基线进行直接比较。我们遍历分类阈值,以确定默认的 0.5 概率截断值是否能在我们的评估集上带来最佳准确率。我们还计算了期望校准误差,并绘制可靠性图,以衡量模型预测置信度与其实际正确性之间的吻合程度。
print("\n" + "=" * 79 + "\n7. ERROR ANALYSIS\n" + "=" * 79)
err = pd.DataFrame({
"text": eval_ds["text"], "y": y_true, "p_pos": p_lora,
"n_words": [len(t.split()) for t in eval_ds["text"]],
})
err["pred"] = (err.p_pos > 0.5).astype(int)
err["correct"] = err.pred == err.y
err["confidence"] = np.maximum(err.p_pos, 1 - err.p_pos)
print("--- 3 most CONFIDENT mistakes (where the model is confidently wrong) ---")
for _, r in err[~err.correct].nlargest(3, "confidence").iterrows():
print(f"\n[true={'pos' if r.y else 'neg'} pred={'pos' if r.pred else 'neg'} "
f"conf={r.confidence:.3f} words={r.n_words}]")
print(clean(r.text)[:400].replace("\n", " "), "...")
err["bucket"] = pd.qcut(err.n_words, 4, labels=["short", "med", "long", "v.long"])
by_len = err.groupby("bucket", observed=True).agg(acc=("correct", "mean"), n=("correct", "size"))
print("\n--- accuracy by review length (truncation hurts long reviews) ---")
print(by_len.to_string())
print("\n" + "=" * 79 + "\n8. OCCLUSION SALIENCY\n" + "=" * 79)
infer_model = model.merge_and_unload()
infer_model.to(DEVICE).eval()
@torch.no_grad()
def predict_proba(texts, bs=64):
out = []
for i in range(0, len(texts), bs):
enc = tok([clean(t) for t in texts[i:i + bs]], truncation=True,
max_length=MAX_LEN, padding=True, return_tensors="pt").to(DEVICE)
out.append(torch.softmax(infer_model(**enc).logits, dim=-1)[:, 1].cpu().numpy())
return np.concatenate(out)
def occlusion(text, max_words=60):
words = clean(text).split()[:max_words]
base = predict_proba([" ".join(words)])[0]
variants = [" ".join(words[:i] + words[i + 1:]) for i in range(len(words))]
dropped = predict_proba(variants)
return words, base - dropped, base
sample = err[err.correct].nlargest(1, "confidence").iloc[0]
words, contrib, base_p = occlusion(sample.text)
print(f"P(positive) for the full excerpt = {base_p:.3f} "
f"(true label = {'pos' if sample.y else 'neg'})\n")
top = np.argsort(np.abs(contrib))[-15:]
plt.figure(figsize=(7, 5))
plt.barh(range(len(top)), contrib[top],
color=["tab:green" if contrib[i] > 0 else "tab:red" for i in top])
plt.yticks(range(len(top)), [words[i] for i in top])
plt.xlabel("Δ P(positive) when the word is removed")
plt.title("Occlusion saliency — green pushes POSITIVE, red pushes NEGATIVE")
plt.tight_layout(); plt.show()
print("\n" + "=" * 79 + "\n9. HEAD vs TAIL TRUNCATION\n" + "=" * 79)
probe = err.nlargest(600, "n_words")
W = 180
head_txt = [" ".join(clean(t).split()[:W]) for t in probe.text]
tail_txt = [" ".join(clean(t).split()[-W:]) for t in probe.text]
yp = probe.y.values
acc_head = ((predict_proba(head_txt) > 0.5).astype(int) == yp).mean()
acc_tail = ((predict_proba(tail_txt) > 0.5).astype(int) == yp).mean()
print(f"on the {len(probe)} longest reviews, using only {W} words:")
print(f" first {W} words -> acc {acc_head:.4f}")
print(f" last {W} words -> acc {acc_tail:.4f}")
print(" Practical takeaway: if the tail wins, feed head+tail to the model or "
"raise MAX_LEN, rather than blindly truncating from the left.")
我们检查了模型最自信的错误预测,并按评论长度分组,以识别与截断相关的失败模式和困难样本。我们将 LoRA 适配器合并到基础模型中,并应用逐词留一遮挡法来估计哪些词将单个预测推向正面或负面情感。随后,我们比较基于长评论开头和结尾部分的预测,以确定最强的情感信息位于何处。
print("\n" + "=" * 79 + "\n10. PSEUDO-LABELLING\n" + "=" * 79)
unsup = raw["unsupervised"].shuffle(seed=SEED).select(range(N_UNSUP))
p_uns = predict_proba(unsup["text"])
keep = (p_uns > 0.95) | (p_uns < 0.05)
pl_texts = [clean(t) for t, k in zip(unsup["text"], keep) if k]
pl_labels = (p_uns[keep] > 0.5).astype(int)
print(f"kept {keep.sum()}/{N_UNSUP} pseudo-labels at conf>0.95 "
f"(balance: {np.bincount(pl_labels)})")
aug = make_pipeline(
TfidfVectorizer(ngram_range=(1, 2), min_df=2, max_features=300_000,
sublinear_tf=True, strip_accents="unicode"),
LogisticRegression(C=8.0, max_iter=2000, n_jobs=-1),
).fit(Xtr + pl_texts, np.concatenate([ytr, pl_labels]))
acc_aug = accuracy_score(yte, aug.predict(Xte))
print(f"TF-IDF baseline : {acc_tfidf:.4f}")
print(f"TF-IDF + pseudo-labels: {acc_aug:.4f} (Δ {acc_aug-acc_tfidf:+.4f})")
print("Caveat: gains are bounded by the teacher. Self-training also amplifies "
"the teacher's biases — always validate on clean, held-out data.")
print("\n" + "=" * 79 + "\n11. SAVE & INFER\n" + "=" * 79)
SAVE_DIR = "./imdb-distilbert-lora-merged"
infer_model.save_pretrained(SAVE_DIR); tok.save_pretrained(SAVE_DIR)
print(f"saved merged model to {SAVE_DIR}/ (load with "
f"AutoModelForSequenceClassification.from_pretrained('{SAVE_DIR}'))")
demos = [
"A masterclass in tension. The final act left the whole theatre silent.",
"Two hours I will never get back. Wooden acting, incoherent plot.",
"It's not the disaster the trailer promised, but it never really lands either.",
]
for d, p in zip(demos, predict_proba(demos)):
print(f" P(pos)={p:.3f} -> {'POSITIVE' if p > 0.5 else 'NEGATIVE'} | {d}")
print("\n" + "=" * 79)
print(f"SUMMARY (n_train={N_TRAIN}, n_eval={N_EVAL}, max_len={MAX_LEN})")
print("=" * 79)
print(pd.DataFrame([
{"model": "TF-IDF + LogReg", "accuracy": acc_tfidf, "roc_auc": auc_tfidf},
{"model": "TF-IDF + pseudo-labels", "accuracy": acc_aug, "roc_auc": float("nan")},
{"model": "DistilBERT + LoRA", "accuracy": accuracy_score(y_true, yhat),
"roc_auc": roc_auc_score(y_true, p_lora)},
]).to_string(index=False))
print("""
NEXT EXPERIMENTS
- Set FULL_RUN = True for the real 25k/25k benchmark (~40 min on a T4).
- Swap MODEL_NAME to 'roberta-base' (target_modules=['query','value']) or
'answerdotai/ModernBERT-base' for an 8k context window — no truncation.
- Head+tail truncation: first 128 + last 128 tokens, motivated by section 9.
- Ablate LoRA rank r in {4, 8, 16, 64} and plot accuracy vs trainable params.
- Replace the pseudo-label teacher with an ensemble and iterate self-training.
- Push to the Hub: huggingface_hub.login() then infer_model.push_to_hub(...).
""")
我们使用微调后的 Transformer 为 IMDb 未标注分区的样本生成高置信度伪标签,并将这些样本添加到 TF-IDF 训练语料中。我们将增强后的分类器与原始基线进行比较,以衡量半监督自训练是否提升了预测准确率。最后,我们保存合并后的 DistilBERT 模型和分词器,对自定义评论执行情感推理,并总结本教程中开发的所有模型的性能。
总而言之,我们构建了一个严谨的情感分类流程,其深度远超简单地微调 Transformer 并报告准确率。我们建立了具有竞争力的 TF-IDF 基线,使用 LoRA 高效训练 DistilBERT,并评估了预测质量和概率可靠性,同时识别了评论长度、截断和高置信度错误对实际性能的影响。我们还通过基于遮挡的显著性方法解读了单个预测,检验了情感信息是否集中在长评论的开头或结尾,并利用未标注数据集中的高置信度伪标签扩展了监督学习。
In this tutorial, we develop an end-to-end sentiment analysis workflow using the Stanford NLP IMDb Large Movie Review Dataset and compare classical machine learning with parameter-efficient transformer fine-tuning. We begin by establishing a reproducible environment and auditing the dataset for class ordering, review-length skew, duplicate leakage, and preprocessing artifacts before training a strong TF-IDF and Logistic Regression baseline. We then fine-tune DistilBERT with LoRA through PEFT, evaluate it using accuracy, macro-F1, ROC-AUC, confusion matrices, and ROC curves, and examine threshold selection and probability calibration through Expected Calibration Error and reliability analysis. Beyond headline metrics, we investigate confident errors, performance across review lengths, word-level occlusion saliency, and head-versus-tail truncation to understand how the model reaches its predictions and where long-context limitations affect performance. Finally, we use the unlabeled IMDb split for confidence-based pseudo-labeling, compare the resulting semi-supervised model against our baseline, and save the merged transformer for reusable sentiment inference.
import importlib.util, subprocess, sys, os, time, random, warnings, inspect, hashlib
warnings.filterwarnings("ignore")
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["WANDB_DISABLED"] = "true"
_REQUIRED = {
"transformers": "transformers",
"datasets": "datasets",
"peft": "peft",
"accelerate": "accelerate",
"sklearn": "scikit-learn",
}
_missing = [pkg for mod, pkg in _REQUIRED.items() if importlib.util.find_spec(mod) is None]
if _missing:
print(f"Installing: {', '.join(_missing)} ...")
subprocess.run([sys.executable, "-m", "pip", "install", "-q", *_missing], check=True)
print("Done. (If imports fail below, restart the runtime and re-run.)\n")
import numpy as np
import pandas as pd
import torch
import matplotlib.pyplot as plt
from datasets import load_dataset
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, f1_score, roc_auc_score,
classification_report, confusion_matrix, roc_curve)
from transformers import (AutoTokenizer, AutoModelForSequenceClassification,
TrainingArguments, Trainer, DataCollatorWithPadding,
EarlyStoppingCallback, set_seed)
from peft import LoraConfig, get_peft_model, TaskType
def _disable_torchao_probe():
patched = []
try:
import peft.import_utils as _piu
_piu.is_torchao_available = lambda: False
patched.append("peft.import_utils")
except Exception:
pass
for _name, _mod in list(sys.modules.items()):
if _name.startswith("peft") and hasattr(_mod, "is_torchao_available"):
_mod.is_torchao_available = lambda: False
patched.append(_name)
return patched
try:
import torchao as _tao
_v = getattr(_tao, "__version__", "?")
if tuple(int(x) for x in _v.split(".")[:2]) < (0, 16):
print(f"[compat] torchao {_v} < 0.16 -> disabling PEFT's torchao probe: "
f"{', '.join(_disable_torchao_probe())}")
except Exception:
_disable_torchao_probe()
SEED = 42
MODEL_NAME = "distilbert-base-uncased"
MAX_LEN = 256
N_TRAIN = 5000
N_EVAL = 2000
N_UNSUP = 3000
EPOCHS = 2
BATCH = 16
LR = 3e-4
FULL_RUN = False
if FULL_RUN:
N_TRAIN, N_EVAL, EPOCHS = 25000, 25000, 3
set_seed(SEED); random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
print("=" * 79)
print(f"device={DEVICE} | torch={torch.__version__} | "
f"gpu={torch.cuda.get_device_name(0) if DEVICE=='cuda' else 'n/a'}")
print("=" * 79)
t0 = time.time()
raw = load_dataset("stanfordnlp/imdb")
print(raw, f"\nloaded in {time.time()-t0:.1f}s\n")
print("--- example (truncated) ---")
print("label:", raw["train"][0]["label"], "|", raw["train"][0]["text"][:300], "...\n")
first_labels = np.array(raw["train"]["label"][:5])
last_labels = np.array(raw["train"]["label"][-5:])
print(f"TRAP #1 - split ordering: first 5 labels {first_labels}, "
f"last 5 labels {last_labels} -> ALWAYS shuffle before subsampling.")
train_full = raw["train"].shuffle(seed=SEED)
test_full = raw["test"].shuffle(seed=SEED)
train_ds = train_full.select(range(min(N_TRAIN, len(train_full))))
eval_ds = test_full.select(range(min(N_EVAL, len(test_full))))
print(f" after shuffle+subsample: train balance = "
f"{np.bincount(train_ds['label'])}, eval balance = {np.bincount(eval_ds['label'])}")
lens = np.array([len(t.split()) for t in train_full["text"]])
q = np.percentile(lens, [50, 75, 90, 95, 99])
print(f"\nTRAP #2 - length (words): median={q[0]:.0f} p75={q[1]:.0f} p90={q[2]:.0f} "
f"p95={q[3]:.0f} p99={q[4]:.0f} max={lens.max()}")
print(f" ~{(lens > MAX_LEN*0.75).mean()*100:.1f}% of reviews exceed MAX_LEN={MAX_LEN} "
f"tokens (rough words->tokens factor 1.3). Section 9 measures what that costs.")
h_tr = {hashlib.md5(t.encode()).hexdigest() for t in raw["train"]["text"]}
h_te = {hashlib.md5(t.encode()).hexdigest() for t in raw["test"]["text"]}
print(f"\nTRAP #3 - leakage: {len(h_tr & h_te)} exact duplicate reviews across "
f"train/test; {len(raw['train'])-len(h_tr)} dupes inside train itself.")
def clean(t):
return t.replace("<br />", " ").replace("<br/>", " ").strip()
plt.figure(figsize=(11, 3.2))
plt.subplot(1, 2, 1)
plt.hist(np.clip(lens, 0, 1000), bins=60)
plt.axvline(MAX_LEN, ls="--", color="k", label=f"MAX_LEN={MAX_LEN}")
plt.title("Review length (words, clipped at 1000)"); plt.legend()
plt.subplot(1, 2, 2)
plt.bar(["neg", "pos"], np.bincount(raw["train"]["label"]))
plt.title("Train class balance (perfectly balanced)")
plt.tight_layout(); plt.show()
We configure the Colab environment, install the required libraries, apply the PEFT–torchao compatibility fix, and set deterministic seeds for reproducible experiments. We load the Stanford IMDb dataset, shuffle and subsample the train and test splits, and inspect class balance, review-length distributions, duplicate leakage, and HTML artifacts. We also visualize review lengths and label frequencies so we understand the dataset structure before building any models.
print("\n" + "=" * 79 + "\n3. TF-IDF BASELINE\n" + "=" * 79)
Xtr = [clean(t) for t in train_ds["text"]]; ytr = np.array(train_ds["label"])
Xte = [clean(t) for t in eval_ds["text"]]; yte = np.array(eval_ds["label"])
t0 = time.time()
tfidf_clf = make_pipeline(
TfidfVectorizer(ngram_range=(1, 2), min_df=2, max_features=300_000,
sublinear_tf=True, strip_accents="unicode"),
LogisticRegression(C=8.0, max_iter=2000, n_jobs=-1),
)
tfidf_clf.fit(Xtr, ytr)
p_tfidf = tfidf_clf.predict_proba(Xte)[:, 1]
acc_tfidf = accuracy_score(yte, p_tfidf > 0.5)
auc_tfidf = roc_auc_score(yte, p_tfidf)
print(f"trained in {time.time()-t0:.1f}s -> acc={acc_tfidf:.4f} auc={auc_tfidf:.4f}")
vec, lr = tfidf_clf.steps[0][1], tfidf_clf.steps[1][1]
feats, coefs = np.array(vec.get_feature_names_out()), lr.coef_[0]
order = np.argsort(coefs)
print("\nmost NEGATIVE n-grams:", ", ".join(feats[order[:12]]))
print("most POSITIVE n-grams:", ", ".join(feats[order[-12:]][::-1]))
print("\n" + "=" * 79 + "\n4. LoRA FINE-TUNING\n" + "=" * 79)
tok = AutoTokenizer.from_pretrained(MODEL_NAME)
def tokenize(batch):
return tok([clean(t) for t in batch["text"]], truncation=True, max_length=MAX_LEN)
tr_tok = (train_ds.map(tokenize, batched=True, remove_columns=["text"])
.rename_column("label", "labels"))
ev_tok = (eval_ds.map(tokenize, batched=True, remove_columns=["text"])
.rename_column("label", "labels"))
base = AutoModelForSequenceClassification.from_pretrained(
MODEL_NAME, num_labels=2,
id2label={0: "NEGATIVE", 1: "POSITIVE"},
label2id={"NEGATIVE": 0, "POSITIVE": 1},
)
lora_cfg = LoraConfig(
task_type=TaskType.SEQ_CLS,
r=16, lora_alpha=32, lora_dropout=0.05,
target_modules=["q_lin", "v_lin"],
modules_to_save=["pre_classifier", "classifier"],
)
try:
model = get_peft_model(base, lora_cfg)
except ImportError as e:
_disable_torchao_probe()
print(f"[compat] retrying after backend probe failure: {e}")
model = get_peft_model(base, lora_cfg)
model.print_trainable_parameters()
def compute_metrics(eval_pred):
logits, labels = eval_pred
probs = torch.softmax(torch.tensor(logits), dim=-1).numpy()[:, 1]
preds = (probs > 0.5).astype(int)
return {"accuracy": accuracy_score(labels, preds),
"f1_macro": f1_score(labels, preds, average="macro"),
"roc_auc": roc_auc_score(labels, probs)}
_ta = inspect.signature(TrainingArguments.__init__).parameters
_eval_key = "eval_strategy" if "eval_strategy" in _ta else "evaluation_strategy"
ta_kwargs = dict(
output_dir="./imdb_lora", learning_rate=LR,
per_device_train_batch_size=BATCH, per_device_eval_batch_size=BATCH * 2,
num_train_epochs=EPOCHS, weight_decay=0.01, warmup_ratio=0.06,
logging_steps=50, save_strategy="epoch", save_total_limit=1,
load_best_model_at_end=True, metric_for_best_model="accuracy",
fp16=(DEVICE == "cuda"), report_to="none", seed=SEED,
)
ta_kwargs[_eval_key] = "epoch"
_tr = inspect.signature(Trainer.__init__).parameters
_tok_key = "processing_class" if "processing_class" in _tr else "tokenizer"
trainer = Trainer(
model=model, args=TrainingArguments(**ta_kwargs),
train_dataset=tr_tok, eval_dataset=ev_tok,
data_collator=DataCollatorWithPadding(tok),
compute_metrics=compute_metrics,
callbacks=[EarlyStoppingCallback(early_stopping_patience=2)],
**{_tok_key: tok},
)
t0 = time.time()
trainer.train()
print(f"\nfine-tuned in {(time.time()-t0)/60:.1f} min")
We train a strong TF-IDF and Logistic Regression baseline and inspect the most influential positive and negative n-grams to establish an interpretable reference point. We then tokenize the IMDb reviews and configure DistilBERT with LoRA adapters that update only a small subset of model parameters while keeping the backbone largely frozen. We use the Hugging Face Trainer with dynamic padding, early stopping, mixed precision, and multiple evaluation metrics to fine-tune the transformer efficiently.
print("\n" + "=" * 79 + "\n5. EVALUATION\n" + "=" * 79)
pred_out = trainer.predict(ev_tok)
p_lora = torch.softmax(torch.tensor(pred_out.predictions), dim=-1).numpy()[:, 1]
y_true = np.array(pred_out.label_ids)
yhat = (p_lora > 0.5).astype(int)
print(classification_report(y_true, yhat, target_names=["neg", "pos"], digits=4))
cm = confusion_matrix(y_true, yhat)
fig, ax = plt.subplots(1, 2, figsize=(11, 4))
ax[0].imshow(cm, cmap="Blues")
for i in range(2):
for j in range(2):
ax[0].text(j, i, cm[i, j], ha="center", va="center", fontsize=14)
ax[0].set_xticks([0, 1], ["pred neg", "pred pos"])
ax[0].set_yticks([0, 1], ["true neg", "true pos"]); ax[0].set_title("Confusion matrix")
for name, p in [("TF-IDF", p_tfidf), ("DistilBERT+LoRA", p_lora)]:
fpr, tpr, _ = roc_curve(y_true, p)
ax[1].plot(fpr, tpr, label=f"{name} (AUC={roc_auc_score(y_true, p):.4f})")
ax[1].plot([0, 1], [0, 1], "k--", lw=0.8)
ax[1].set_xlabel("FPR"); ax[1].set_ylabel("TPR"); ax[1].set_title("ROC"); ax[1].legend()
plt.tight_layout(); plt.show()
print("\n" + "=" * 79 + "\n6. THRESHOLD & CALIBRATION\n" + "=" * 79)
ths = np.linspace(0.05, 0.95, 91)
accs = [(y_true == (p_lora > t)).mean() for t in ths]
best_t = ths[int(np.argmax(accs))]
print(f"[email protected] = {accs[45]:.4f} | best threshold = {best_t:.2f} -> acc = {max(accs):.4f}")
def expected_calibration_error(probs, labels, n_bins=10):
"""ECE: |confidence - accuracy| averaged over confidence bins."""
conf = np.maximum(probs, 1 - probs)
correct = (probs > 0.5).astype(int) == labels
bins = np.linspace(0, 1, n_bins + 1)
ece, xs, ys = 0.0, [], []
for lo, hi in zip(bins[:-1], bins[1:]):
m = (conf > lo) & (conf <= hi)
if m.sum() == 0:
continue
ece += m.mean() * abs(conf[m].mean() - correct[m].mean())
xs.append(conf[m].mean()); ys.append(correct[m].mean())
return ece, np.array(xs), np.array(ys)
ece, cx, cy = expected_calibration_error(p_lora, y_true)
print(f"Expected Calibration Error = {ece:.4f} (0 = perfectly calibrated)")
plt.figure(figsize=(9, 3.2))
plt.subplot(1, 2, 1); plt.plot(ths, accs); plt.axvline(best_t, ls="--", color="r")
plt.xlabel("threshold"); plt.ylabel("accuracy"); plt.title("Threshold sweep")
plt.subplot(1, 2, 2); plt.plot([0.5, 1], [0.5, 1], "k--", lw=0.8)
plt.plot(cx, cy, "o-"); plt.xlabel("mean confidence"); plt.ylabel("empirical accuracy")
plt.title(f"Reliability diagram (ECE={ece:.3f})")
plt.tight_layout(); plt.show()
We evaluate the fine-tuned DistilBERT-LoRA model using classification metrics, a confusion matrix, and ROC curves while directly comparing its ROC-AUC performance with the TF-IDF baseline. We sweep classification thresholds to determine whether the default probability cutoff of 0.5 gives the best accuracy on our evaluation set. We also calculate Expected Calibration Error and construct a reliability diagram to measure how closely the model’s predicted confidence corresponds to its actual correctness.
print("\n" + "=" * 79 + "\n7. ERROR ANALYSIS\n" + "=" * 79)
err = pd.DataFrame({
"text": eval_ds["text"], "y": y_true, "p_pos": p_lora,
"n_words": [len(t.split()) for t in eval_ds["text"]],
})
err["pred"] = (err.p_pos > 0.5).astype(int)
err["correct"] = err.pred == err.y
err["confidence"] = np.maximum(err.p_pos, 1 - err.p_pos)
print("--- 3 most CONFIDENT mistakes (where the model is confidently wrong) ---")
for _, r in err[~err.correct].nlargest(3, "confidence").iterrows():
print(f"\n[true={'pos' if r.y else 'neg'} pred={'pos' if r.pred else 'neg'} "
f"conf={r.confidence:.3f} words={r.n_words}]")
print(clean(r.text)[:400].replace("\n", " "), "...")
err["bucket"] = pd.qcut(err.n_words, 4, labels=["short", "med", "long", "v.long"])
by_len = err.groupby("bucket", observed=True).agg(acc=("correct", "mean"), n=("correct", "size"))
print("\n--- accuracy by review length (truncation hurts long reviews) ---")
print(by_len.to_string())
print("\n" + "=" * 79 + "\n8. OCCLUSION SALIENCY\n" + "=" * 79)
infer_model = model.merge_and_unload()
infer_model.to(DEVICE).eval()
@torch.no_grad()
def predict_proba(texts, bs=64):
out = []
for i in range(0, len(texts), bs):
enc = tok([clean(t) for t in texts[i:i + bs]], truncation=True,
max_length=MAX_LEN, padding=True, return_tensors="pt").to(DEVICE)
out.append(torch.softmax(infer_model(**enc).logits, dim=-1)[:, 1].cpu().numpy())
return np.concatenate(out)
def occlusion(text, max_words=60):
words = clean(text).split()[:max_words]
base = predict_proba([" ".join(words)])[0]
variants = [" ".join(words[:i] + words[i + 1:]) for i in range(len(words))]
dropped = predict_proba(variants)
return words, base - dropped, base
sample = err[err.correct].nlargest(1, "confidence").iloc[0]
words, contrib, base_p = occlusion(sample.text)
print(f"P(positive) for the full excerpt = {base_p:.3f} "
f"(true label = {'pos' if sample.y else 'neg'})\n")
top = np.argsort(np.abs(contrib))[-15:]
plt.figure(figsize=(7, 5))
plt.barh(range(len(top)), contrib[top],
color=["tab:green" if contrib[i] > 0 else "tab:red" for i in top])
plt.yticks(range(len(top)), [words[i] for i in top])
plt.xlabel("Δ P(positive) when the word is removed")
plt.title("Occlusion saliency — green pushes POSITIVE, red pushes NEGATIVE")
plt.tight_layout(); plt.show()
print("\n" + "=" * 79 + "\n9. HEAD vs TAIL TRUNCATION\n" + "=" * 79)
probe = err.nlargest(600, "n_words")
W = 180
head_txt = [" ".join(clean(t).split()[:W]) for t in probe.text]
tail_txt = [" ".join(clean(t).split()[-W:]) for t in probe.text]
yp = probe.y.values
acc_head = ((predict_proba(head_txt) > 0.5).astype(int) == yp).mean()
acc_tail = ((predict_proba(tail_txt) > 0.5).astype(int) == yp).mean()
print(f"on the {len(probe)} longest reviews, using only {W} words:")
print(f" first {W} words -> acc {acc_head:.4f}")
print(f" last {W} words -> acc {acc_tail:.4f}")
print(" Practical takeaway: if the tail wins, feed head+tail to the model or "
"raise MAX_LEN, rather than blindly truncating from the left.")
We examine the model’s most confident incorrect predictions and group reviews by length to identify truncation-related failure patterns and difficult examples. We merge the LoRA adapters into the underlying model and apply leave-one-word-out occlusion to estimate which words push individual predictions toward positive or negative sentiment. We then compare predictions based on the beginning and ending portions of long reviews to determine where the strongest sentiment information resides.
print("\n" + "=" * 79 + "\n10. PSEUDO-LABELLING\n" + "=" * 79)
unsup = raw["unsupervised"].shuffle(seed=SEED).select(range(N_UNSUP))
p_uns = predict_proba(unsup["text"])
keep = (p_uns > 0.95) | (p_uns < 0.05)
pl_texts = [clean(t) for t, k in zip(unsup["text"], keep) if k]
pl_labels = (p_uns[keep] > 0.5).astype(int)
print(f"kept {keep.sum()}/{N_UNSUP} pseudo-labels at conf>0.95 "
f"(balance: {np.bincount(pl_labels)})")
aug = make_pipeline(
TfidfVectorizer(ngram_range=(1, 2), min_df=2, max_features=300_000,
sublinear_tf=True, strip_accents="unicode"),
LogisticRegression(C=8.0, max_iter=2000, n_jobs=-1),
).fit(Xtr + pl_texts, np.concatenate([ytr, pl_labels]))
acc_aug = accuracy_score(yte, aug.predict(Xte))
print(f"TF-IDF baseline : {acc_tfidf:.4f}")
print(f"TF-IDF + pseudo-labels: {acc_aug:.4f} (Δ {acc_aug-acc_tfidf:+.4f})")
print("Caveat: gains are bounded by the teacher. Self-training also amplifies "
"the teacher's biases — always validate on clean, held-out data.")
print("\n" + "=" * 79 + "\n11. SAVE & INFER\n" + "=" * 79)
SAVE_DIR = "./imdb-distilbert-lora-merged"
infer_model.save_pretrained(SAVE_DIR); tok.save_pretrained(SAVE_DIR)
print(f"saved merged model to {SAVE_DIR}/ (load with "
f"AutoModelForSequenceClassification.from_pretrained('{SAVE_DIR}'))")
demos = [
"A masterclass in tension. The final act left the whole theatre silent.",
"Two hours I will never get back. Wooden acting, incoherent plot.",
"It's not the disaster the trailer promised, but it never really lands either.",
]
for d, p in zip(demos, predict_proba(demos)):
print(f" P(pos)={p:.3f} -> {'POSITIVE' if p > 0.5 else 'NEGATIVE'} | {d}")
print("\n" + "=" * 79)
print(f"SUMMARY (n_train={N_TRAIN}, n_eval={N_EVAL}, max_len={MAX_LEN})")
print("=" * 79)
print(pd.DataFrame([
{"model": "TF-IDF + LogReg", "accuracy": acc_tfidf, "roc_auc": auc_tfidf},
{"model": "TF-IDF + pseudo-labels", "accuracy": acc_aug, "roc_auc": float("nan")},
{"model": "DistilBERT + LoRA", "accuracy": accuracy_score(y_true, yhat),
"roc_auc": roc_auc_score(y_true, p_lora)},
]).to_string(index=False))
print("""
NEXT EXPERIMENTS
- Set FULL_RUN = True for the real 25k/25k benchmark (~40 min on a T4).
- Swap MODEL_NAME to 'roberta-base' (target_modules=['query','value']) or
'answerdotai/ModernBERT-base' for an 8k context window — no truncation.
- Head+tail truncation: first 128 + last 128 tokens, motivated by section 9.
- Ablate LoRA rank r in {4, 8, 16, 64} and plot accuracy vs trainable params.
- Replace the pseudo-label teacher with an ensemble and iterate self-training.
- Push to the Hub: huggingface_hub.login() then infer_model.push_to_hub(...).
""")
We use the fine-tuned transformer to generate high-confidence pseudo-labels for examples from IMDb’s unlabeled split and add these examples to the TF-IDF training corpus. We compare the augmented classifier against the original baseline to measure whether semi-supervised self-training improves predictive accuracy. Finally, we save the merged DistilBERT model and tokenizer, run sentiment inference on custom reviews, and summarize the performance of all models developed throughout the tutorial.
In conclusion, we developed a rigorous sentiment classification pipeline that goes well beyond simply fine-tuning a transformer and reporting accuracy. We established a competitive TF-IDF baseline, train DistilBERT efficiently with LoRA, and evaluate both predictive quality and probability reliability while identifying how review length, truncation, and highly confident mistakes influence real-world performance. We also interpreted individual predictions through occlusion-based saliency, tested whether sentiment information is concentrated near the beginning or end of long reviews, and extended supervised learning with high-confidence pseudo-labels from the unlabeled dataset.