IMDb Sentiment Analysis with DistilBERT LoRA, TF-IDF Baselines, Calibration, Interpretability, Robustness Testing, and Semi-Supervised Learning

A tutorial by an unnamed author demonstrates an end-to-end sentiment analysis workflow on the Stanford NLP IMDb Large Movie Review Dataset, comparing a TF-IDF and Logistic Regression baseline with a DistilBERT model fine-tuned using LoRA via PEFT. The workflow evaluates models using accuracy, macro-F1, ROC-AUC, calibration metrics, and robustness tests, and incorporates semi-supervised learning with pseudo-labeling on the unlabeled IMDb split. The tutorial reports that the transformer approach achieves higher performance than the baseline, with detailed analysis of confident errors, length-based performance, and truncation effects.

In this tutorial, we develop an end-to-end sentiment analysis workflow using the Stanford NLP IMDb https://huggingface.co/datasets/stanfordnlp/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. python 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. Check out the FULL CODES here . Also, feel free to follow us on and don’t forget to join our Twitter https://x.com/intent/follow?screen name=marktechpost and Subscribe to 150k+ML SubReddit https://www.reddit.com/r/machinelearningnews/ . Wait are you on telegram? our Newsletter https://www.aidevsignals.com/ now you can join us on telegram as well. https://t.me/machinelearningresearchnews Need to partner with us for promoting your GitHub Repo OR Hugging Face Page OR Product Release OR Webinar etc.? Connect with us https://forms.gle/wbash1wF6efRj8G58 Sana Hassan, a consulting intern at Marktechpost and dual-degree student at IIT Madras, is passionate about applying technology and AI to address real-world challenges. With a keen interest in solving practical problems, he brings a fresh perspective to the intersection of AI and real-life solutions.