# AG News classfication 100 GC v4
course: Module 4 — Generative AI & LLMs
module: Module-4-Generative-AI-LLMs
type: notebook
source_url: https://personal-learn.armco.dev/files/Module-4-Generative-AI-LLMs/General/Lab_Materials_23_May_2026/AG_News_classfication_100_GC_v4.ipynb
---
[cell 1 markdown]
## Dataset Overview: AG News Classification
AG News is one of the most widely used benchmark datasets for text classification tasks in Natural Language
Processing (NLP). It is a collection of news articles gathered from thousands of news sources and is
commonly used to evaluate and compare machine learning models on short-text classification.
### Source and Background
- The dataset was originally compiled by Antonio Gulli as part of his news aggregation research.
- It was derived from the AG's corpus of news articles, which was collected from more than 2,000 news
sources and contains over one million articles.
### Structure
- Three fields are included in each record: `Class Index`, `Title`, and `Description`.
- The headline of the news article is stored in the `Title` field.
- A short snippet or summary of the article body is stored in the `Description` field.
- A numeric label ranging from 1 to 4 is assigned to the `Class Index` field, representing the topic
category of the article.
### Class Labels
Four topic categories are covered in the dataset:
| Class Index | Category |
|-------------|----------|
| 1 | World |
| 2 | Sports |
| 3 | Business |
| 4 | Sci/Tech |
### Split Information
- A total of 120,000 samples are included in the training set, with exactly 30,000 samples allocated per
class.
- A total of 7,600 samples are included in the test set, with exactly 1,900 samples allocated per class.
- Both splits are perfectly balanced, meaning every class is represented equally across both sets.
### Why This Dataset is Recommended for Beginners
- Class imbalance does not need to be addressed, as all classes are equally represented.
- Preprocessing is kept straightforward due to the short length of the input text.
- Model predictions and errors can be easily interpreted, given that only four well-defined categories are
used.
- The dataset is large enough for meaningful models to be trained, yet small enough to be handled without
specialized hardware.
- It is widely recommended as a starting point before more complex or domain-specific classification tasks
are attempted.
[cell 2 code]
# cell_0a — GPU mask setup before any torch import
# MAIN GOAL: Let the user pick which physical GPU to use (or CPU),
# then hide all other GPUs from this process via CUDA_VISIBLE_DEVICES.
# Must run before importing torch or any CUDA-using library.
import os
choice = input("Enter physical GPU index (e.g. 0, 1, 2) or 'cpu': ").strip().lower()
if choice == "cpu":
os.environ["CUDA_VISIBLE_DEVICES"] = ""
DEVICE = "cpu"
elif choice.isdigit():
os.environ["CUDA_VISIBLE_DEVICES"] = choice # hide every other GPU
DEVICE = "cuda:0" # selected GPU maps to cuda:0 internally
else:
print("Invalid input. Defaulting to CPU.")
os.environ["CUDA_VISIBLE_DEVICES"] = ""
DEVICE = "cpu"
print(f"Device set to: {DEVICE}")
[cell 3 code]
# cell_0b — Import all dependencies
# MAIN GOAL: Import every library this notebook will use.
# If any import fails, install the missing package before proceeding.
# ── Standard library ──────────────────────────────────────────────────────────
import os # File path operations and reading environment variables
import random # Random number generation (e.g. for setting seeds)
# ── Numerical & data manipulation ─────────────────────────────────────────────
import numpy as np # Fast array operations and math; backbone of most ML libs
import pandas as pd # DataFrames for loading, inspecting, and transforming data
# ── Deep learning framework (PyTorch) ─────────────────────────────────────────
import torch # Core tensor library and autograd engine
import torch.nn as nn # Building blocks for neural network layers
from torch.utils.data import (
Dataset, # Base class for defining a custom dataset
DataLoader # Batches and shuffles data for feeding into the model
)
# ── Pre-trained language model (Hugging Face Transformers) ────────────────────
from transformers import (
BertTokenizer, # Converts raw text → token IDs that BERT understands
BertModel # Pre-trained BERT model weights, used as the text encoder
)
# ── Evaluation metrics (scikit-learn) ─────────────────────────────────────────
from sklearn.metrics import (
accuracy_score, # Fraction of correctly predicted labels
classification_report, # Per-class precision, recall, and F1 breakdown
f1_score # Harmonic mean of precision & recall (handles class imbalance)
)
# ── Progress tracking ──────────────────────────────────────────────────────────
from tqdm import tqdm # Wraps any iterable to display a live progress bar
print("All imports successful.")
[cell 4 code]
# cell_0c — Verify device and GPU info
# MAIN GOAL: Confirm that torch sees the correct device.
# Print GPU name and available memory if CUDA is selected.
if DEVICE == "cuda:0":
print(f"Device : {DEVICE}")
print(f"GPU Name : {torch.cuda.get_device_name(0)}")
total = torch.cuda.get_device_properties(0).total_memory / 1e9 # convert bytes to GB
reserved = torch.cuda.memory_reserved(0) / 1e9
allocated = torch.cuda.memory_allocated(0) / 1e9
free = total - reserved # approximate free memory
print(f"Total Memory : {total:.2f} GB")
print(f"Reserved : {reserved:.2f} GB")
print(f"Allocated : {allocated:.2f} GB")
print(f"Free (approx) : {free:.2f} GB")
else:
print(f"Device: {DEVICE} — no GPU info to display.")
[cell 5 markdown]
## Part 1 — Data Exploration and Prep
Before any model is trained, the data needs to be understood. Class distribution should be checked, a few samples should be inspected, and the input should be confirmed to be clean and consistent.
The AG News dataset contains news articles across four categories: World, Sports, Business, and Sci/Tech. Each sample is provided with a title and a description. Instead of using just one of them, both will be fed to the BERT tokenizer as a sentence pair. The tokenizer joins them internally with a `[SEP]` token and tags each side with a segment id (0 for the title, 1 for the description), so BERT can apply the right segment embedding to each part. This is more informative than concatenating the two into a single string, because the pair structure preserves the boundary between the two fields.
Here is what this part covers:
- Train and test CSV files are loaded
- Shape, class balance, and sample rows are inspected
- Title and Description are kept as separate columns and passed to the tokenizer as a sentence pair (joined internally with `[SEP]` and proper segment ids)
- A balanced subset is selected for training (to keep runtimes short)
- Inputs are tokenized using the BERT tokenizer
- PyTorch Dataset and DataLoader objects are built
By the end of this part, the data will be ready to be fed into a BERT model.
[cell 6 code]
# cell_1 — Load the dataset
# MAIN GOAL: Load train and test CSV files and do a first inspection.
# ── Load train and test sets ───────────────────────────────────────────────────
train_df = pd.read_csv("/content/train.csv", header=0)
test_df = pd.read_csv("/content/test.csv", header=0)
# ── Standardise column names ───────────────────────────────────────────────────
# The raw CSV has unnamed or inconsistently cased columns; rename for clarity
# Expected columns in this order: class label, article title, article body
train_df.columns = ["label", "title", "description"]
test_df.columns = ["label", "title", "description"]
# ── Sanity-check: confirm shapes loaded as expected ───────────────────────────
print(f"Train size : {train_df.shape}") # Expect (120000, 3) for full AG News train
print(f"Test size : {test_df.shape}") # Expect (100, 3) given nrows=100
# ── Check class balance in the training set ────────────────────────────────────
# AG News has 4 classes: 1=World, 2=Sports, 3=Business, 4=Sci/Tech
# A balanced distribution (~30k each) means no need for class-weight adjustments
print("\nClass distribution (train):")
print(train_df["label"].value_counts().sort_index())
# ── Preview a few rows to visually confirm structure ──────────────────────────
print("\nSample rows:")
train_df.head(3)
[cell 7 code]
train_df
[cell 8 code]
test_df
[cell 9 code]
print(test_df["label"].value_counts())
[cell 10 code]
# cell_2 — Convert labels to zero-indexed
# MAIN GOAL: Prepare labels for training.
# Title and description are kept as separate columns. The tokenizer will combine
# them as a sentence pair in the Dataset class (cell_4b), producing input_ids,
# attention_mask, AND token_type_ids — the segment ids (0 for title, 1 for
# description) are what allow BERT to apply its pretrained segment embeddings
# correctly. Keeping title and description separate at this stage is what makes
# the correct segment-ids possible downstream.
# ── Shift labels from 1-indexed to 0-indexed ──────────────────────────────────
# AG News labels are originally 1–4 (World, Sports, Business, Sci/Tech).
# PyTorch's CrossEntropyLoss and BERT's classifier head both expect 0-based
# labels (0–3), so we subtract 1 from every label in both splits.
train_df["label"] = train_df["label"] - 1 # 1–4 → 0–3
test_df["label"] = test_df["label"] - 1 # 1–4 → 0–3
# ── Verify the label range is correct ─────────────────────────────────────────
# Min should be 0, max should be 3 (num_classes - 1).
# Any value outside this range would cause an index error during training.
# Note: this cell is destructive — running it twice will shift labels to
# (-1, 2), silently breaking training. Always restart the kernel and rerun
# from the top if you suspect the cell ran more than once.
print(f"Label range: {train_df['label'].min()} to {train_df['label'].max()}")
# ── Spot-check the first row to confirm data integrity ────────────────────────
# Visually confirm that the title, description, and label are coherent together.
# Title + description will later be joined by the tokenizer via a [SEP] token,
# with separate segment embeddings (token_type_ids) to tell BERT which text
# belongs to which segment — so keeping them as separate columns is intentional.
print(f"\nSample title : {train_df['title'].iloc[0]}")
print(f"Sample description : {train_df['description'].iloc[0]}")
print(f"Label : {train_df['label'].iloc[0]}")
[cell 11 markdown]
## Exploratory Data Analysis
Before any preprocessing or modeling is done, the data needs to be understood. This step is often skipped in tutorials but it is one of the most important parts of any machine learning pipeline. Decisions made here directly affect how the data is prepared and how the model is configured.
Three things are checked in this section:
- Class distribution: whether the dataset is balanced across all four categories
- Text length distribution: how long the combined input texts are and how many exceed BERT's maximum token limit
- Sample inspection: what the actual text looks like for each class
These checks help catch problems early. An imbalanced dataset would require a different training strategy. Texts that are too long would get truncated silently by the tokenizer, which could cause information loss. Looking at real samples builds intuition about the vocabulary and writing style of each class.
[cell 12 code]
# cell_2a — Class distribution plot
# MAIN GOAL: Visualize the number of samples per class in the training set.
# This confirms whether the dataset is balanced or skewed towards certain classes.
# A balanced dataset means no class-weighting is needed in the loss function.
# ── Imports ───────────────────────────────────────────────────────────────────
import matplotlib.pyplot as plt # Standard plotting library
# ── Map numeric labels to human-readable class names ─────────────────────────
# Labels are currently 0–3 (after zero-indexing in cell_2); mapping them to
# their actual names makes the plot axis readable without needing a legend.
label_map = {0: "World", 1: "Sports", 2: "Business", 3: "Sci/Tech"}
train_df["label_name"] = train_df["label"].map(label_map) # Add a new column for display
# ── Count samples per class ───────────────────────────────────────────────────
# sort_index() ensures classes appear in a consistent order (alphabetical),
# rather than sorted by count, which would be misleading for a balance check.
class_counts = train_df["label_name"].value_counts().sort_index()
# ── Plot class distribution ───────────────────────────────────────────────────
fig, ax = plt.subplots(figsize=(8, 5)) # Wide enough to clearly separate 4 bars
bars = ax.bar(
class_counts.index, # X-axis: class names
class_counts.values, # Y-axis: sample counts
color="#4C72B0", # Seaborn's default blue — neutral, print-friendly
width=0.5 # Narrower bars leave whitespace, improving readability
)
# ── Labels and axes ───────────────────────────────────────────────────────────
ax.set_title("Class Distribution in Training Set")
ax.set_xlabel("Class")
ax.set_ylabel("Number of Samples")
ax.bar_label(bars, fmt="%d", padding=3) # Annotate exact counts above each bar
ax.set_ylim(0, class_counts.max() * 1.15) # Add 15% headroom so labels aren't clipped
# ── Save and display ──────────────────────────────────────────────────────────
plt.tight_layout() # Prevent axis labels from being cut off at figure edges
plt.savefig(
"class_distribution.png",
dpi=150, # High enough resolution for reports or presentations
bbox_inches="tight" # Crop whitespace around the figure before saving
)
plt.show()
[cell 13 markdown]
## Model Setup: BERT in Code
### Loading a Pre-trained Model with `from_pretrained`
- In the Hugging Face `transformers` library, pre-trained models and tokenizers are loaded using
the `from_pretrained` method.
- The model name is passed as a string, in this case `"bert-base-uncased"`, and the corresponding
weights and configuration are downloaded automatically on the first run.
- Once downloaded, the files are cached locally, so no repeated downloads are needed in subsequent
runs.
### What `BertTokenizerFast` Does
- `BertTokenizerFast` is the tokenizer that converts raw text into a format that BERT can process.
- Each input string is broken down into subword units called WordPiece tokens.
- Special tokens are also added automatically: `[CLS]` is placed at the start of every sequence,
and `[SEP]` is placed at the end.
- Along with the token IDs, an attention mask is produced, which tells the model which positions
contain real tokens and which are padding.
### Fast vs Slow Tokenizer
- Two versions of the BERT tokenizer are available in the `transformers` library: the slow
tokenizer, written in pure Python, and the fast tokenizer, backed by the Rust-based
`tokenizers` library.
- The fast tokenizer is used here because it processes text significantly faster, which matters
when large datasets need to be tokenized during exploratory data analysis.
- Both versions produce identical outputs, so the choice between them does not affect the final
token IDs or model behavior.
- The fast tokenizer is also noted to be reloaded in a later cell, as this particular instance
is only used for exploratory purposes.
### What `bert-base-uncased` Means
- `bert-base` refers to the smaller of the two standard BERT sizes, with 12 Transformer layers,
768 hidden dimensions, and approximately 110 million parameters.
- `uncased` means that all input text is lowercased before tokenization, so words like `"Paris"`
and `"paris"` are treated as the same token.
- This variant is generally preferred for tasks where capitalization does not carry meaningful
information, which is the case for most news classification scenarios.
[cell 14 code]
# cell_2b — Text length analysis
# MAIN GOAL: Understand the distribution of text lengths in the combined input column.
# This helps determine an appropriate max_len for the tokenizer
# and flags how much content gets truncated at training time.
# ── Load a fast tokenizer just for EDA ───────────────────────────────────────
# BertTokenizerFast is a Rust-backed tokenizer — much faster than the default
# Python-based BertTokenizer when encoding tens of thousands of rows at once.
# This instance is only for length analysis; the training tokenizer is loaded
# in cell_4a with the proper config (max_len, padding, truncation, etc.).
from transformers import BertTokenizerFast
tokenizer_eda = BertTokenizerFast.from_pretrained("bert-base-uncased")
# ── Copy training data for EDA ────────────────────────────────────────────────
# Work on a copy so that adding temporary EDA columns (combined, word_count, etc.)
# does not pollute train_df, which is used for model training later.
eda_df = train_df.copy()
# ── Compute word and character counts ─────────────────────────────────────────
# Title and description are stored as separate columns (to preserve [SEP] token
# behaviour during training), but for raw length analysis we concatenate them
# the same way the Dataset class will present them to the tokenizer.
eda_df["combined"] = eda_df["title"] + " " + eda_df["description"]
eda_df["word_count"] = eda_df["combined"].apply(lambda x: len(x.split())) # Whitespace-split word count
eda_df["char_count"] = eda_df["combined"].apply(lambda x: len(x)) # Raw character count
# ── Compute true token counts without truncation ──────────────────────────────
# We pass title and description as a *pair* to mirror exactly what the Dataset
# class does: tokenizer produces [CLS] title [SEP] description [SEP].
# truncation=False is critical here — we want the real unclipped token count
# to understand the full distribution before deciding on max_len.
eda_df["token_count"] = eda_df.apply(
lambda row: len(tokenizer_eda(
row["title"],
row["description"],
truncation=False # Do NOT truncate — measure the true length
)["input_ids"]),
axis=1
)
# ── Print descriptive statistics ──────────────────────────────────────────────
# Covers count, mean, std, min, quartiles, and max for all three length metrics.
# Key things to look for:
# - Mean/median token count vs our chosen max_len (128)
# - How wide the spread is (std) — high std means many outliers
print("Text Length Statistics (Full Training Set)")
print("=" * 40)
print(eda_df[["word_count", "char_count", "token_count"]].describe().round(2))
# ── Flag samples that exceed token limits ─────────────────────────────────────
# BERT has a hard architectural limit of 512 tokens — anything longer cannot be
# processed at all without chunking or truncation.
# Our chosen max_len of 128 is more aggressive: samples above this threshold
# will have their tail content silently dropped during training, which may
# cause the model to miss important information in longer articles.
over_512 = (eda_df["token_count"] > 512).sum()
over_128 = (eda_df["token_count"] > 128).sum()
total = len(eda_df) # Use actual dataset length, not a hardcoded constant
print(f"\nSamples exceeding 512 tokens (BERT hard limit) : {over_512} ({over_512/total*100:.2f}%)")
print(f"Samples exceeding 128 tokens (our max_len) : {over_128} ({over_128/total*100:.2f}%)")
# ── Plot token count distribution ─────────────────────────────────────────────
# The histogram shows the full spread of token lengths.
# Vertical reference lines mark the two key thresholds so we can visually judge
# how much of the distribution lies beyond each cutoff.
fig, ax = plt.subplots(figsize=(8, 5))
ax.hist(
eda_df["token_count"],
bins=50, # Enough bins to show the shape without over-smoothing
color="#4C72B0",
edgecolor="white" # White edges make individual bars easier to distinguish
)
# Reference lines for the two token-length thresholds
ax.axvline(x=128, color="red", linestyle="--", linewidth=1.5, label="Our max_len (128)")
ax.axvline(x=512, color="orange", linestyle="--", linewidth=1.5, label="BERT limit (512)")
ax.set_title("Token Count Distribution (Full Training Set)")
ax.set_xlabel("Number of Tokens")
ax.set_ylabel("Number of Samples")
ax.legend()
plt.tight_layout()
plt.savefig("token_distribution.png", dpi=150, bbox_inches="tight")
plt.show()
[cell 15 code]
# cell_2c — Sample inspection per class
# MAIN GOAL: Print a few real examples from each class so the student
# gets a feel for the vocabulary and writing style of each category.
# This is a qualitative sanity check — verifying that labels map to
# the content you'd actually expect before committing to a training run.
# ── Configuration ─────────────────────────────────────────────────────────────
N_SAMPLES = 2 # Number of random examples to display per class; increase for richer inspection
# ── Header ────────────────────────────────────────────────────────────────────
print("=" * 70)
print("Sample Texts per Class")
print("=" * 70)
# ── Iterate over all four AG News classes ─────────────────────────────────────
# label_map is defined in cell_2a: {0: "World", 1: "Sports", 2: "Business", 3: "Sci/Tech"}
# Iterating over it here keeps class names and indices in sync automatically —
# no risk of the display label drifting out of step with the numeric label.
for label_idx, label_name in label_map.items():
print(f"\nClass: {label_name}")
print("-" * 70)
# ── Sample N random rows from this class ──────────────────────────────────
# random_state=42 ensures the same examples appear every run, making it
# easier to discuss or document specific outputs without them changing.
# Only title and description are selected — label is not needed for display.
samples = (
train_df[train_df["label"] == label_idx] # Filter to current class only
.sample(N_SAMPLES, random_state=42) # Draw N reproducible random rows
[["title", "description"]] # Keep only the text columns
.values # Convert to a numpy array for easy unpacking
)
# ── Display each sample ───────────────────────────────────────────────────
for i, (title, description) in enumerate(samples):
# Join title and description with a separator for readability.
# Note: "|" is used here only for display — during training the tokenizer
# inserts a proper [SEP] token between them, not a pipe character.
text = title + " | " + description
# Truncate to 300 characters for display only — long descriptions would
# flood the output and make it hard to skim across classes.
# This truncation has no effect on the actual data in train_df.
display_text = text if len(text) <= 300 else text[:300] + "..."
print(f"\nSample {i+1}: {display_text}")
print() # Blank line between classes for visual separation
# ── Footer ────────────────────────────────────────────────────────────────────
print("=" * 70)
[cell 16 code]
# cell_3 — Select balanced subsets for train, val, and test
# MAIN GOAL: Sample an equal number of examples from each class for training.
# Also carve out 10% of the training subset as a validation set.
# The test set is stratified separately so per-class evaluation metrics
# are not skewed by the raw row order in test.csv (see note below).
# Validation is used to monitor performance after each epoch — the test set
# is only touched once at the very end to give an unbiased final evaluation.
# ── Imports ───────────────────────────────────────────────────────────────────
from sklearn.model_selection import train_test_split # Stratified split utility
# ── Configuration flags ───────────────────────────────────────────────────────
# USE_FULL_DATASET: set to True to use the full train (120k) and test (7.6k) sets.
# set to False during development to iterate quickly on subsets.
# SAMPLES_PER_CLASS: only used when USE_FULL_DATASET=False.
# Total training rows = SAMPLES_PER_CLASS * num_classes (4).
# e.g. 25 * 4 = 100 rows — very small, purely for fast iteration.
# TEST_SAMPLES_PER_CLASS: only used when USE_FULL_DATASET=False.
# Kept as a separate knob from SAMPLES_PER_CLASS because train
# size and test size serve different purposes — train affects
# what the model learns, test affects how reliably we measure it.
USE_FULL_DATASET = False
SAMPLES_PER_CLASS = 25
TEST_SAMPLES_PER_CLASS = 25
# ── Build the training subset ─────────────────────────────────────────────────
if USE_FULL_DATASET:
# Use all available training rows — no subsampling needed
train_subset = train_df.copy()
print("Using full training dataset.")
else:
# Sample exactly SAMPLES_PER_CLASS rows from each class.
# groupby + apply ensures each class contributes equally, preventing the model
# from seeing more examples of one class than another during training.
# reset_index(drop=True) gives a clean 0-based index after groupby fragments it.
train_subset = (
train_df
.groupby("label", group_keys=False) # Group rows by class label
.apply(lambda x: x.sample(SAMPLES_PER_CLASS, # Draw fixed N per group
random_state=42)) # Fixed seed for reproducibility
.reset_index(drop=True) # Flatten the grouped index
)
print(f"Using train subset: {SAMPLES_PER_CLASS} samples per class → {len(train_subset)} total rows.")
# ── Build the test subset (stratified) ────────────────────────────────────────
# IMPORTANT: the raw test.csv is NOT class-balanced in row order.
# A naive nrows=100 load was observed to yield a skewed test set (37/30/21/12
# across the four classes), which would make per-class metrics — and the
# frozen vs. fine-tuned comparison in cell_12 — unreliable.
# Stratifying here guarantees an exactly balanced test set regardless of
# the original row ordering in test.csv.
# Stratification logic is kept separate from the train branch above because
# train size and test size are independent decisions: even on a full-data
# GPU run, you would still want a properly balanced test set.
if USE_FULL_DATASET:
# Use all 7,600 test rows as-is — the full test set is balanced by design
print("Using full test dataset.")
else:
test_df = (
test_df
.groupby("label", group_keys=False) # Group test rows by class label
.apply(lambda x: x.sample(TEST_SAMPLES_PER_CLASS, # Draw fixed N per group
random_state=42)) # Fixed seed for reproducibility
.reset_index(drop=True) # Flatten the grouped index
)
print(f"Using test subset: {TEST_SAMPLES_PER_CLASS} samples per class → {len(test_df)} total rows.")
# ── Carve out a validation split ──────────────────────────────────────────────
# 10% of the training subset is held out as a validation set.
# stratify=train_subset["label"] ensures the class ratio in train and val
# mirrors the original distribution — without this, random chance could put
# all examples of one class into val and none into train (especially risky
# with small subsets like 100 total rows).
train_subset, val_subset = train_test_split(
train_subset,
test_size = 0.1, # 10% → val, 90% → train
random_state = 42, # Fixed seed for reproducible splits
stratify = train_subset["label"] # Preserve class balance in both halves
)
# Reset indices after the split so both DataFrames are cleanly 0-indexed.
# This prevents subtle off-by-one bugs when accessing rows by position later.
train_subset = train_subset.reset_index(drop=True)
val_subset = val_subset.reset_index(drop=True)
# ── Report final split sizes ──────────────────────────────────────────────────
# Three-way split summary: train / val / test.
print(f"\nTrain size : {train_subset.shape}")
print(f"Val size : {val_subset.shape}")
print(f"Test size : {test_df.shape}")
# ── Verify class balance is preserved in all splits ──────────────────────────
# With stratify=True (train/val) and groupby-sample (test), each class should
# appear in roughly equal proportions. With very small subsets, small rounding
# differences (e.g. 5 vs 6 per class) are expected and acceptable for train/val.
# The test set should be exactly balanced since groupby-sample draws a fixed N per class.
print("\nClass distribution in train:")
print(train_subset["label"].value_counts().sort_index())
print("\nClass distribution in val:")
print(val_subset["label"].value_counts().sort_index())
print("\nClass distribution in test:")
print(test_df["label"].value_counts().sort_index())
[cell 17 code]
# cell_4a — Load the BERT tokenizer
# MAIN GOAL: Load the pretrained BERT tokenizer and verify it produces the
# expected output format before it is used inside the Dataset class in cell_4b.
# The tokenizer converts raw text into three tensors that BERT requires:
# - input_ids : integer token IDs from BERT's 30,522-word vocabulary
# - attention_mask : 1 for real tokens, 0 for [PAD] tokens
# - token_type_ids : 0 for sentence A tokens, 1 for sentence B tokens (used for pairs)
# ── Imports ───────────────────────────────────────────────────────────────────
from transformers import BertTokenizer # Python-based tokenizer (not the Fast variant)
# ── Load the pretrained tokenizer ─────────────────────────────────────────────
# "bert-base-uncased" means:
# - base : 12 transformer layers, 110M parameters (vs. bert-large with 24 layers)
# - uncased : all text is lowercased before tokenisation, so "BERT" == "bert"
# The tokenizer vocabulary and special token config are downloaded from Hugging Face
# and cached locally on first run; subsequent runs load from cache instantly.
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
# ── Smoke-test the tokenizer on a short sample ────────────────────────────────
# This confirms the tokenizer is loaded correctly and lets us inspect the output
# format before wiring it into the full Dataset pipeline.
# We deliberately use max_length=10 here (shorter than training's max_length=128)
# so the output is small enough to read in the console.
sample_text = "BERT is a powerful language model."
sample_tokens = tokenizer(
sample_text,
return_tensors = "pt", # Return PyTorch tensors (not lists or numpy arrays)
padding = "max_length", # Pad short sequences to exactly max_length with [PAD]
max_length = 10 # Short limit for readable console output only
)
# ── Inspect the tokenizer output ──────────────────────────────────────────────
# input_ids: each integer maps to one token in BERT's vocabulary.
# Special tokens are always added automatically:
# [CLS] (ID 101) is prepended — its final hidden state is used for classification.
# [SEP] (ID 102) is appended — marks the end of a sentence or the boundary between two.
print("Token IDs :", sample_tokens["input_ids"])
# attention_mask: tells BERT which positions to attend to.
# 1 = real token (should be attended to)
# 0 = [PAD] token (should be ignored by the attention mechanism)
# Without this mask, padding tokens would incorrectly influence attention scores.
print("Attention Mask :", sample_tokens["attention_mask"])
# convert_ids_to_tokens maps IDs back to human-readable subword tokens.
# Note: BERT uses WordPiece tokenisation, so words can be split into subwords.
# e.g. "powerful" might stay as "powerful", but "BERT" → "bert" (uncased),
# and rare words like "tokenisation" → ["token", "##isation"].
# The "##" prefix marks a subword continuation (not the start of a new word).
print("Tokens :", tokenizer.convert_ids_to_tokens(sample_tokens["input_ids"][0]))
[cell 18 code]
# cell_4b — Define the custom PyTorch Dataset class
# MAIN GOAL: Wrap the DataFrame into a PyTorch Dataset so that DataLoader can
# efficiently batch, shuffle, and serve samples to the model during training.
# PyTorch requires two methods to be implemented:
# - __len__ : returns the total number of samples in the dataset
# - __getitem__ : returns a single tokenized sample by index
class AGNewsDataset(Dataset):
def __init__(self, dataframe, tokenizer, max_len=128):
# reset_index ensures __getitem__ can safely use .loc[idx] with a
# clean 0-based integer index — without this, idx might not match the
# DataFrame's index if rows were dropped or shuffled upstream.
self.data = dataframe.reset_index(drop=True)
self.tokenizer = tokenizer
self.max_len = max_len # Token budget per sample; must match the value used at inference
def __len__(self):
# Called by DataLoader to know how many batches to generate per epoch
return len(self.data)
def __getitem__(self, idx):
# ── Fetch raw text and label for this sample ──────────────────────────
# Title and description are retrieved separately so the tokenizer can
# insert the [SEP] token between them with correct token_type_ids,
# rather than a raw string separator which BERT would not understand.
title = self.data.loc[idx, "title"]
description = self.data.loc[idx, "description"]
label = self.data.loc[idx, "label"] # Already zero-indexed (0–3) from cell_2
# ── Tokenize title + description as a sentence pair ───────────────────
# Passing two strings produces the full BERT input format:
# [CLS] title tokens [SEP] description tokens [SEP]
# The tokenizer produces three tensors:
# - input_ids : token IDs from BERT's WordPiece vocabulary
# - attention_mask : 1 for real tokens, 0 for [PAD]
# - token_type_ids : 0 for [CLS] + title + first [SEP],
# 1 for description + final [SEP]
# All three are returned so BERT can use its segment embeddings as
# designed during pretraining — segment-A vectors are added to the
# title tokens, segment-B vectors are added to the description tokens.
#
# truncation="only_second" means only the description is trimmed if the
# combined sequence exceeds max_len. The title is always kept in full
# because it is typically shorter and carries the strongest topical signal.
# "only_first" or True (truncates the longer one) would risk losing the title.
#
# padding="max_length" pads every sample to exactly max_len tokens with
# [PAD] (ID 0). This is required for DataLoader to stack samples into
# uniform-shaped tensors — variable-length sequences cannot be batched.
encoding = self.tokenizer(
title,
description,
max_length = self.max_len,
padding = "max_length", # Pad to max_len so all samples have the same shape
truncation = "only_second", # Trim description first; never trim the title
return_tensors = "pt" # Return PyTorch tensors, not Python lists
)
# ── Return a dict of tensors for this sample ──────────────────────────
# squeeze(0) removes the batch dimension the tokenizer adds by default
# when return_tensors="pt": shape (1, max_len) → (max_len,).
# DataLoader will re-add the batch dimension when stacking samples.
#
# token_type_ids is included so BERT can distinguish the title segment
# (zeros) from the description segment (ones). This matches the input
# format BERT was pretrained on for sentence pairs and is the most
# faithful encoding of a (title, description) pair for classification.
# Omitting it would silently default the whole sequence to segment A
# inside BERT — functional but a small accuracy regression on pair tasks.
return {
"input_ids" : encoding["input_ids"].squeeze(0), # shape: (max_len,)
"attention_mask" : encoding["attention_mask"].squeeze(0), # shape: (max_len,)
"token_type_ids" : encoding["token_type_ids"].squeeze(0), # shape: (max_len,) — 0=title, 1=description
"label" : torch.tensor(label, dtype=torch.long) # scalar; dtype required by CrossEntropyLoss
}
# ── Smoke-test: instantiate the dataset and fetch one sample ──────────────────
# This catches shape mismatches or tokenizer config errors before training starts,
# without needing to run a full forward pass through BERT.
sample_dataset = AGNewsDataset(train_subset, tokenizer)
sample = sample_dataset[0]
print(f"input_ids shape : {sample['input_ids'].shape}") # Expect: torch.Size([128])
print(f"attention_mask shape : {sample['attention_mask'].shape}") # Expect: torch.Size([128])
print(f"token_type_ids shape : {sample['token_type_ids'].shape}") # Expect: torch.Size([128])
print(f"label : {sample['label']}") # Expect: a scalar tensor in 0–3
# ── Inspect the token_type_ids pattern ────────────────────────────────────────
# Should show a block of 0s (title segment, including [CLS] and the first [SEP]),
# followed by a block of 1s (description segment, including the final [SEP]),
# followed by 0s again for [PAD] tokens (padding belongs to no segment, defaults to 0).
# Counting the transition point tells you how long the title segment is.
n_segment_a = (sample["token_type_ids"] == 0).sum().item()
n_segment_b = (sample["token_type_ids"] == 1).sum().item()
print(f"\nSegment A tokens (title + padding): {n_segment_a}")
print(f"Segment B tokens (description) : {n_segment_b}")
# ── Decode the first sample to visually verify the [CLS]...[SEP]...[SEP] structure ──
# skip_special_tokens=False keeps [CLS], [SEP], and [PAD] visible so we can
# confirm the tokenizer is formatting the pair correctly.
# Checking the first 120 characters is enough to see [CLS], the title, and [SEP].
decoded = tokenizer.decode(sample["input_ids"], skip_special_tokens=False)
print(f"\nDecoded (first 120 chars): {decoded[:120]}")
[cell 19 code]
# cell_5 — Build the DataLoaders
# MAIN GOAL: Wrap each Dataset in a DataLoader for train, val, and test splits.
# DataLoader handles three things automatically:
# - Batching : groups individual samples into tensors of shape (batch_size, max_len)
# - Shuffling : randomises sample order each epoch (train only) to prevent order bias
# - Loading : optionally parallelises __getitem__ calls across CPU workers
# ── Suppress tokenizer parallelism warning ────────────────────────────────────
# HuggingFace tokenizers run a Rust-backed thread pool internally.
# When DataLoader forks worker processes (num_workers > 0), each worker inherits
# this thread pool, which causes a deadlock warning and forces parallelism off.
# Setting this env var to "false" before forking silences the warning explicitly.
# Must be set before any tokenizer is used in a worker context.
import os
os.environ["TOKENIZERS_PARALLELISM"] = "false"
# ── Configuration ─────────────────────────────────────────────────────────────
# BATCH_SIZE controls how many samples are processed together in one forward pass.
# Larger batches → faster training but higher GPU memory usage.
# Smaller batches → slower training but more gradient update steps per epoch,
# which can sometimes improve generalisation on small datasets.
# 32 is a standard default for BERT fine-tuning on a single GPU.
BATCH_SIZE = 32
# ── Instantiate Dataset objects ───────────────────────────────────────────────
# Each Dataset wraps its DataFrame and applies tokenization lazily in __getitem__,
# meaning text is tokenized on-the-fly as batches are requested — not all upfront.
# This keeps memory usage low at the cost of some CPU overhead per batch.
train_dataset = AGNewsDataset(train_subset, tokenizer)
val_dataset = AGNewsDataset(val_subset, tokenizer)
test_dataset = AGNewsDataset(test_df, tokenizer)
# ── Wrap datasets in DataLoader ───────────────────────────────────────────────
train_loader = DataLoader(
train_dataset,
batch_size = BATCH_SIZE,
shuffle = True, # Reshuffle every epoch so the model never memorises sample order
num_workers = 0 # See note below on why 0 is correct here
)
val_loader = DataLoader(
val_dataset,
batch_size = BATCH_SIZE,
shuffle = False, # Order doesn't matter for evaluation; determinism is preferred
num_workers = 0
)
test_loader = DataLoader(
test_dataset,
batch_size = BATCH_SIZE,
shuffle = False, # Test set must never be shuffled — results should be reproducible
num_workers = 0
)
# NOTE — why num_workers=0 for all loaders:
# num_workers > 0 spawns separate processes to call __getitem__ in parallel.
# However, HuggingFace tokenizers disable their internal thread pool when forked
# (to avoid deadlocks), which means workers would actually run slower than the
# main process. On small datasets like this one, the overhead of inter-process
# communication outweighs any parallelism benefit anyway.
# num_workers=0 keeps everything in the main process — simpler, faster, no warnings.
# ── Verify batch counts ───────────────────────────────────────────────────────
# len(loader) = ceil(dataset_size / batch_size)
# e.g. 90 train samples / 32 batch size = 3 batches (last batch may be smaller)
# If any loader shows 0 batches, the corresponding dataset is empty — check cell_3.
print(f"Train batches : {len(train_loader)}")
print(f"Val batches : {len(val_loader)}")
print(f"Test batches : {len(test_loader)}")
print(f"Batch size : {BATCH_SIZE}")
[cell 20 markdown]
## Part 2 — Pretrained BERT (Frozen Encoder)
When a model is described as "pretrained", it means the weights have already been learned on a large corpus. In BERT's case, the model was trained on BookCorpus and English Wikipedia using masked language modeling and next sentence prediction. These weights encode a rich understanding of language.
The question is: how much of that understanding carries over to a new task like news classification, without any task-specific training?
To answer this, the encoder is kept completely frozen. None of BERT's weights are updated during training. Only a small classification head is added on top and trained from scratch. This head takes the [CLS] token representation produced by BERT and maps it to one of the four news categories.
This setup is useful for two reasons:
- It is fast to train because only the classifier head has trainable parameters
- It gives a clear baseline to compare against fine-tuning later
If frozen BERT already performs well, it means the pretrained representations are directly useful for this task. If it performs poorly, it means the encoder needs to be adapted to the task through fine-tuning.
[cell 21 markdown]
## Model Architecture
The architecture used in this notebook is built on top of `bert-base-uncased`, a pretrained transformer model released by Google. It has 12 transformer layers, 12 attention heads, and a hidden size of 768, giving it approximately 110 million parameters in total. The "uncased" means all text is lowercased before tokenization, so "BERT" and "bert" are treated as the same.
---
### Input Representation
Each input is a combined string of title and description separated by a [SEP] token. The tokenizer converts this string into token IDs and an attention mask. Two special tokens are also added automatically:
- `[CLS]` is added at the start of every input
- `[SEP]` is added at the end
So the final input seen by BERT looks like this:
```
[CLS] title tokens [SEP] description tokens [SEP]
```
The attention mask is a sequence of 1s and 0s. A value of 1 means the token is real and should be attended to. A value of 0 means the token is padding and should be ignored.
---
### Inside BERT
The token IDs are passed through 12 transformer layers. Each layer applies self-attention and a feed-forward network. After all 12 layers, BERT produces a vector of size 768 for every single token in the input.
So if the input has 52 tokens, the output is a matrix of shape (52, 768). Every token gets its own representation.
---
### The [CLS] Token
The first token in every input is always [CLS]. During pretraining, BERT was trained to encode a summary of the entire input sequence into this token. So instead of using all 52 token vectors, only the first one is used -- the [CLS] vector of size 768 -- as a fixed-length representation of the entire input.
This is what the following line does in the code:
```python
cls_output = outputs.last_hidden_state[:, 0, :] # shape: (batch_size, 768)
```
---
### The Classifier Head
A single linear layer is placed on top of the [CLS] vector:
```python
nn.Linear(768, 4)
```
It takes the 768-dimensional [CLS] vector and maps it to 4 scores, one for each class. The class with the highest score is the predicted label.
---
### How the Two Models Differ
Both models share the exact same architecture. The only difference is what gets trained during the learning process.
```
Input Text
|
Tokenizer
|
[CLS] + input_ids + attention_mask
|
BERT Encoder (12 Transformer Layers)
|
[CLS] vector (768 dimensions)
|
Linear Classifier (768 -> 4)
|
Predicted Class (World / Sports / Business / Sci/Tech)
```
In **Frozen BERT**, the encoder weights are fixed. Only the linear classifier is trained. This means 3,076 parameters out of 110 million are updated during training.
In **Fine-tuned BERT**, the entire stack is trained end-to-end. All 110 million parameters are updated. The pretrained weights are not discarded -- they are used as a starting point and gradually adjusted to become more useful for news classification.
[cell 22 code]
# cell_6 — Define the frozen BERT classifier model
# MAIN GOAL: Build a text classifier that uses BERT as a frozen feature extractor.
# "Frozen" means all 110M BERT parameters are locked — their weights will not be
# updated during training. Only the classification head (3,076 parameters) is trained.
# This approach is fast and avoids overfitting on small datasets, at the cost of
# not adapting BERT's representations to the target domain.
class FrozenBERTClassifier(nn.Module):
def __init__(self, num_classes=4):
super(FrozenBERTClassifier, self).__init__()
# ── Load pretrained BERT encoder ──────────────────────────────────────
# Downloads bert-base-uncased weights from Hugging Face (cached after first run).
# BertModel outputs hidden states but has no classification head of its own —
# we attach our own linear layer below.
self.bert = BertModel.from_pretrained("bert-base-uncased")
# ── Freeze all BERT parameters ────────────────────────────────────────
# Setting requires_grad=False on every BERT parameter tells PyTorch not
# to compute or store gradients for those weights during backpropagation.
# Effect: BERT acts as a fixed feature extractor — it maps token sequences
# to contextual embeddings, but its weights never change during training.
# Trade-offs vs full fine-tuning:
# + Much faster — gradients only flow through the linear head
# + Lower overfitting risk on very small datasets
# - BERT representations are not adapted to the target domain or task
# - Lower performance ceiling than full fine-tuning on larger datasets
for param in self.bert.parameters():
param.requires_grad = False
# ── Classification head ───────────────────────────────────────────────
# A single linear layer maps the [CLS] embedding (768-dim) to class logits.
# bert.config.hidden_size is used instead of hardcoding 768 so this class
# works with other BERT variants (e.g. bert-large: hidden_size=1024)
# without code changes. Output size = num_classes = 4.
self.classifier = nn.Linear(self.bert.config.hidden_size, num_classes) # (768 → 4)
def forward(self, input_ids, attention_mask, token_type_ids):
# ── Forward pass through frozen BERT ──────────────────────────────────
# All three tensors are passed explicitly to BERT:
#
# input_ids : integer token IDs — shape (batch_size, seq_len)
#
# attention_mask : 1 for real tokens, 0 for [PAD] tokens.
# Prevents padding positions from influencing attention.
#
# token_type_ids : 0 for title tokens (sentence A),
# 1 for description tokens (sentence B).
# Tells BERT which segment each token belongs to.
# Previously omitted in the notebook; passing it
# explicitly makes full use of BERT's segment embeddings
# and is the correct approach for sentence-pair inputs.
outputs = self.bert(
input_ids = input_ids,
attention_mask = attention_mask,
token_type_ids = token_type_ids,
)
# ── Extract the [CLS] token representation ────────────────────────────
# last_hidden_state has shape (batch_size, seq_len, hidden_size).
# [:, 0, :] selects position 0 ([CLS]) for every sample in the batch.
# BERT is pretrained so [CLS] aggregates information from the full
# sequence — it is the standard input to a downstream classifier head.
cls_output = outputs.last_hidden_state[:, 0, :] # shape: (batch_size, 768)
# ── Produce class logits ──────────────────────────────────────────────
# Raw unnormalised scores for each class — not probabilities.
# CrossEntropyLoss applies softmax internally, so we do NOT apply it here.
logits = self.classifier(cls_output) # shape: (batch_size, 4)
return logits
# ── Instantiate the model and move to the target device ───────────────────────
# .to(DEVICE) moves all parameters to GPU if available, CPU otherwise.
# DEVICE should be defined earlier as:
# DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
frozen_model = FrozenBERTClassifier(num_classes=4).to(DEVICE)
# ── Count trainable vs frozen parameters ──────────────────────────────────────
# Confirms the freeze was applied correctly:
# Trainable should be ~3,076 (768 × 4 weights + 4 biases in the linear head)
# Frozen should be ~109M (all BERT encoder weights)
# If trainable equals total, the freeze loop did not execute — check __init__.
total_params = sum(p.numel() for p in frozen_model.parameters())
trainable_params = sum(p.numel() for p in frozen_model.parameters() if p.requires_grad)
print(f"Total parameters : {total_params:,}")
print(f"Trainable parameters : {trainable_params:,}")
print(f"Frozen parameters : {total_params - trainable_params:,}")
[cell 23 code]
# cell_6b — Define the evaluate function
# MAIN GOAL: Define a reusable evaluation function that runs a full pass over
# any DataLoader and returns a complete set of classification metrics.
# Defined here (before any training cell) so it can be called identically from:
# - cell_7 : evaluating the frozen BERT classifier after training
# - cell_10 : evaluating the fine-tuned BERT classifier after training
# Keeping one shared function ensures metrics are computed consistently across
# both experiments and results are directly comparable.
# ── Imports ───────────────────────────────────────────────────────────────────
from sklearn.metrics import (
accuracy_score, # Overall fraction of correct predictions
classification_report,# Per-class precision, recall, F1, and support
precision_score, # Weighted avg: how often positive predictions are correct
recall_score, # Weighted avg: how often actual positives are caught
f1_score # Weighted avg: harmonic mean of precision and recall
)
def evaluate(model, loader, device):
# ── Switch model to evaluation mode ───────────────────────────────────────
# model.eval() disables two training-specific behaviours:
# - Dropout : all neurons are active, giving deterministic outputs
# - BatchNorm : uses running statistics instead of batch statistics
# Without this call, results would be non-deterministic and slightly degraded
# even when no gradient updates are performed.
model.eval()
all_preds = [] # Accumulates predicted class indices across all batches
all_labels = [] # Accumulates true class indices across all batches
# ── Disable gradient computation ──────────────────────────────────────────
# torch.no_grad() tells PyTorch not to build the computation graph during
# the forward pass. This reduces memory usage and speeds up inference since
# no intermediate activations need to be stored for backpropagation.
with torch.no_grad():
progress_bar = tqdm(loader, desc="Evaluating", leave=False)
# leave=False removes the progress bar after completion so it doesn't
# clutter the output when evaluate() is called inside a training loop.
for batch in progress_bar:
# ── Move batch tensors to the target device ────────────────────────
# Tensors must be on the same device as the model weights.
# All three BERT inputs (input_ids, attention_mask, token_type_ids)
# are pulled from the batch — token_type_ids was added in cell_4b
# to give BERT proper segment embeddings for the (title, description) pair.
# Labels are moved to device here for consistency but brought back
# to CPU below for scikit-learn, which does not support GPU tensors.
input_ids = batch["input_ids"].to(device)
attention_mask = batch["attention_mask"].to(device)
token_type_ids = batch["token_type_ids"].to(device)
labels = batch["label"].to(device)
# ── Forward pass ───────────────────────────────────────────────────
# Returns raw logits of shape (batch_size, num_classes).
# No softmax applied — argmax of logits == argmax of probabilities.
logits = model(input_ids, attention_mask, token_type_ids)
# ── Convert logits to predicted class indices ──────────────────────
# argmax(dim=1) picks the class with the highest logit score for
# each sample in the batch → shape: (batch_size,)
preds = logits.argmax(dim=1)
# ── Collect predictions and labels on CPU ──────────────────────────
# .cpu() moves tensors off GPU before .numpy() conversion.
# scikit-learn metric functions require numpy arrays, not tensors.
# extend() (not append()) flattens the batch into the running list.
all_preds.extend(preds.cpu().numpy())
all_labels.extend(labels.cpu().numpy())
# ── Compute metrics over the full dataset ─────────────────────────────────
# All metrics use average="weighted" — each class score is weighted by its
# support (number of true samples), so larger classes have more influence.
# This is appropriate for balanced datasets; use average="macro" if you want
# equal weight per class regardless of support.
# zero_division=0 suppresses warnings for classes with no predicted samples
# (common on tiny validation sets where a class may never be predicted).
accuracy = accuracy_score(all_labels, all_preds)
precision = precision_score(all_labels, all_preds, average="weighted", zero_division=0)
recall = recall_score(all_labels, all_preds, average="weighted", zero_division=0)
f1 = f1_score(all_labels, all_preds, average="weighted", zero_division=0)
# ── Build a full per-class breakdown ──────────────────────────────────────
# classification_report shows precision, recall, F1, and support for each
# class individually — useful for spotting which classes the model struggles
# with, even when the weighted average looks acceptable overall.
report = classification_report(
all_labels,
all_preds,
target_names = ["World", "Sports", "Business", "Sci/Tech"], # Map 0–3 to class names
zero_division = 0
)
# ── Return everything the caller needs ────────────────────────────────────
# all_labels and all_preds are returned alongside the scalar metrics so the
# caller can plot confusion matrices or run further custom analysis without
# needing to re-run inference.
return accuracy, precision, recall, f1, report, all_labels, all_preds
print("evaluate() defined and ready.")
[cell 24 markdown]
### Discussion — Before Training
Take a moment to think about these two questions before running the training cell.
**Question 1:** BERT has around 110 million parameters in total. In this setup, only the classifier head is being trained, which amounts to roughly 3,076 parameters. Do you think this is enough to produce meaningful classification results?
**Question 2:** The [CLS] token representation produced by BERT was not trained for news classification. It was pretrained on masked language modeling and next sentence prediction. Do you think this representation will still carry enough information to separate news into four categories?
Run the next cell and see how your intuition holds up. The answers are discussed after the results.
[cell 25 code]
# cell_7 — Training loop for frozen BERT classifier
# MAIN GOAL: Train only the linear classification head on the training subset.
# BERT encoder weights remain frozen throughout — gradients flow through the
# linear head only, making each epoch very fast compared to full fine-tuning.
# After each epoch, validation accuracy is checked to monitor generalisation.
# ── Imports ───────────────────────────────────────────────────────────────────
from tqdm import tqdm # Progress bar for the inner batch loop
# ── Hyperparameters ───────────────────────────────────────────────────────────
# EPOCHS: number of full passes over the training set.
# 3 epochs is sufficient here because only the head is being trained —
# full fine-tuning typically needs fewer epochs (2–4) but more compute per epoch.
# LR: learning rate for the classifier head.
# 2e-3 is deliberately higher than the 2e-5 typically used for full fine-tuning.
# Fine-tuning uses a tiny LR to avoid catastrophically overwriting pretrained
# BERT weights; here those weights are frozen so a larger LR is safe and
# helps the head converge faster from its random initialisation.
EPOCHS = 3
LR = 2e-3
# ── Loss function ─────────────────────────────────────────────────────────────
# CrossEntropyLoss combines LogSoftmax + NegativeLogLikelihoodLoss in one step.
# It expects raw logits (not softmax outputs) and integer class indices as labels.
# Internally applies softmax, so we must NOT apply softmax in the model's forward().
criterion = nn.CrossEntropyLoss()
# ── Optimizer ─────────────────────────────────────────────────────────────────
# AdamW = Adam with decoupled weight decay regularisation.
# Weight decay (L2 penalty) discourages large weights and reduces overfitting —
# especially important here since the head is trained on a very small dataset.
# Note: only frozen_model.classifier.parameters() is passed to the optimizer —
# NOT frozen_model.parameters(). This is a second safety net on top of
# requires_grad=False: even if a BERT parameter somehow had a gradient,
# AdamW would not update it because it is not registered with the optimizer.
optimizer = torch.optim.AdamW(
frozen_model.classifier.parameters(), # Only the linear head — not BERT encoder
lr = LR,
weight_decay = 0.01 # L2 regularisation coefficient; 0.01 is a standard default
)
# ── Training loop (one epoch) ─────────────────────────────────────────────────
def train_epoch(model, loader, optimizer, criterion, device):
# Switch to training mode: enables dropout and batch norm update (if present).
# Required counterpart to model.eval() used in evaluate() — always set explicitly.
model.train()
total_loss = 0 # Accumulates raw loss values across all batches
total_correct = 0 # Accumulates correct prediction counts across all batches
progress_bar = tqdm(loader, desc="Training", leave=False)
# leave=False removes the bar after the epoch completes so it doesn't
# stack up across epochs in the output.
for batch in progress_bar:
# ── Move batch to device ───────────────────────────────────────────────
# All tensors must be on the same device as the model weights.
# All three BERT inputs (input_ids, attention_mask, token_type_ids) are
# pulled from the batch — token_type_ids was added in cell_4b to give
# BERT proper segment embeddings for the (title, description) pair.
input_ids = batch["input_ids"].to(device)
attention_mask = batch["attention_mask"].to(device)
token_type_ids = batch["token_type_ids"].to(device)
labels = batch["label"].to(device)
# ── Zero gradients ─────────────────────────────────────────────────────
# PyTorch accumulates gradients by default — zero_grad() clears them
# before each forward pass so gradients from the previous batch do not
# contaminate the current update. Forgetting this is one of the most
# common bugs in PyTorch training loops.
optimizer.zero_grad()
# ── Forward pass ───────────────────────────────────────────────────────
# Input flows through frozen BERT → [CLS] embedding → linear head → logits.
# Gradients will only be computed for the linear head (requires_grad=True).
logits = model(input_ids, attention_mask, token_type_ids) # shape: (batch_size, num_classes)
# ── Compute loss ───────────────────────────────────────────────────────
# CrossEntropyLoss(logits, labels): logits are (batch_size, 4),
# labels are (batch_size,) integer class indices in range 0–3.
loss = criterion(logits, labels)
# ── Backward pass ──────────────────────────────────────────────────────
# Computes gradients of the loss with respect to all parameters that
# have requires_grad=True — in this case, only the linear head weights.
loss.backward()
# ── Gradient clipping ──────────────────────────────────────────────────
# Clips the global gradient norm to max_norm=1.0 before the optimizer step.
# Prevents exploding gradients, which can cause loss to spike or go to NaN.
# Less critical here (only a linear layer), but good practice and harmless.
# Note: model.parameters() includes frozen BERT params, but their .grad
# is None (no backward pass touched them), so clip_grad_norm_ silently
# skips them — only the classifier head's gradients are actually clipped.
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
# ── Optimizer step ─────────────────────────────────────────────────────
# Updates only the registered parameters (classifier head) using the
# gradients computed in loss.backward().
optimizer.step()
# ── Accumulate metrics ─────────────────────────────────────────────────
total_loss += loss.item() # .item() extracts a Python float from the scalar tensor
total_correct += (logits.argmax(dim=1) == labels).sum().item()
# Live loss displayed in the progress bar so we can spot divergence early
progress_bar.set_postfix(loss=f"{loss.item():.4f}")
# ── Epoch-level metrics ────────────────────────────────────────────────────
# avg_loss: mean batch loss — more informative than total loss since it is
# independent of the number of batches (and therefore of dataset size).
# accuracy: fraction of all training samples correctly classified this epoch.
avg_loss = total_loss / len(loader)
accuracy = total_correct / len(loader.dataset)
return avg_loss, accuracy
# ── Run the full training loop ────────────────────────────────────────────────
print("Training Frozen BERT Classifier...\n")
for epoch in range(1, EPOCHS + 1):
train_loss, train_acc = train_epoch(
frozen_model, train_loader, optimizer, criterion, DEVICE
)
# Evaluate on the validation set after every epoch.
# Only val_acc is unpacked here — the full metrics report is reserved
# for the final evaluation in cell_8 to avoid cluttering training output.
# The test set is never touched here — it is held out until cell_8.
val_acc, _, _, _, _, _, _ = evaluate(frozen_model, val_loader, DEVICE)
print(
f"Epoch {epoch}/{EPOCHS} — "
f"Loss: {train_loss:.4f} | "
f"Train Acc: {train_acc:.4f} | "
f"Val Acc: {val_acc:.4f}"
)
[cell 26 code]
# cell_8 — Evaluate frozen BERT on the test set
# MAIN GOAL: Run the frozen BERT classifier on the held-out test set and report
# the full suite of evaluation metrics for this experiment.
# This is the first and only time the test set is used for the frozen model —
# running it multiple times and selecting the best result would constitute
# data leakage and give an overly optimistic view of real-world performance.
# Results here serve as the baseline to compare against the fine-tuned model in cell_11.
# ── Run evaluation on the test set ────────────────────────────────────────────
# evaluate() is defined in cell_6b. It:
# 1. Runs a full forward pass over every batch in test_loader
# 2. Collects predictions and true labels across all batches
# 3. Returns scalar metrics + the full per-class report + raw label arrays
# The raw label arrays (all_labels, all_preds_frozen) are stored so they can
# be used for a confusion matrix or further analysis without re-running inference.
(
frozen_accuracy,
frozen_precision,
frozen_recall,
frozen_f1,
frozen_report,
all_labels, # True class indices for the full test set (0–3)
all_preds_frozen # Predicted class indices — prefixed "frozen_" to distinguish
# from fine-tuned predictions collected later in cell_11
) = evaluate(frozen_model, test_loader, DEVICE)
# ── Print scalar metrics ───────────────────────────────────────────────────────
# All metrics use weighted averaging (see cell_6b) — each class score is
# weighted by its support, so the aggregate reflects class frequency.
# On a perfectly balanced test set, weighted and macro averages will be equal.
print("Frozen BERT — Test Results")
print("-" * 35)
print(f"Accuracy : {frozen_accuracy:.4f}") # Overall fraction of correct predictions
print(f"Precision : {frozen_precision:.4f}") # Weighted avg: correctness of positive predictions
print(f"Recall : {frozen_recall:.4f}") # Weighted avg: coverage of actual positives
print(f"F1 Score : {frozen_f1:.4f}") # Weighted avg: harmonic mean of precision & recall
# ── Print per-class breakdown ─────────────────────────────────────────────────
# The classification report shows precision, recall, F1, and support for each
# of the four AG News classes individually.
# Key things to look for:
# - Which classes have low recall? (model misses many true examples of that class)
# - Which classes have low precision? (model over-predicts that class)
# - Is support balanced? (unequal support can make weighted metrics misleading)
# These results serve as the frozen-model baseline — compare them against
# the fine-tuned report in cell_11 to quantify what full fine-tuning gains.
print("\nPer-Class Report:")
print(frozen_report)
[cell 27 markdown]
### Discussion — Answers
**Question 1:** BERT has around 110 million parameters in total. In this setup, only the classifier head is being trained, which amounts to roughly 3,076 parameters. Do you think this is enough to produce meaningful classification results?
- Despite only a tiny fraction of parameters being trained, the model produces strong classification results. This is not because the classifier head is powerful. It is because the frozen encoder is already doing the heavy lifting.
- During pretraining, BERT was exposed to a massive amount of text and learned to produce representations that capture the meaning, context, and semantics of language. These representations do not need to be retrained for every new task.
- The [CLS] token carries a compressed summary of the entire input sequence. A simple linear layer on top is sufficient to separate categories because the representations coming out of the frozen encoder are already semantically meaningful and well organized.
- BERT is kept completely frozen. It produces the same representation for a given input regardless of the task. Only the classifier head learns anything new. This works because the pretrained representations are already rich enough to be useful without any modification.
**Question 2:** The [CLS] token representation produced by BERT was not trained for news classification. It was pretrained on masked language modeling and next sentence prediction. Do you think this representation will still carry enough information to separate news into four categories?
- Yes. Even though BERT was never exposed to news category labels during pretraining, the semantic structure it learned is general enough to capture meaningful differences between topics.
- News categories like Sports, Business, and Sci/Tech have distinct vocabularies and writing styles. BERT's pretrained representations already organize words and sentences in a way that reflects these differences, even without task-specific training.
- Some categories are naturally easier to separate than others. Classes with very distinct vocabulary tend to be classified more accurately. Classes whose language overlaps with other categories tend to be harder to separate with a linear classifier alone.
- The key takeaway is that pretraining on a general language task produces representations that are broadly useful. This is why pretrained models like BERT are so widely adopted. They reduce the need for large amounts of task-specific labeled data.
[cell 28 markdown]
## Part 3 — Fine-tuned BERT
In Part 2, the BERT encoder was kept completely frozen. Only the classifier head was trained. The encoder produced the same representations regardless of the task, and the classifier had to work with whatever BERT gave it.
Fine-tuning takes a different approach. The entire model is unfrozen and trained end-to-end on the target task. This means BERT's internal representations are no longer fixed. They are gradually adjusted to become more useful for news classification specifically.
The key difference is that during fine-tuning, the gradient flows all the way back through the encoder. Every layer in BERT gets updated, not just the classifier head. The pretrained weights are not thrown away. They are used as a starting point and nudged in the direction that reduces classification loss.
A few things change compared to Part 2:
- All parameters are trainable, around 110 million instead of 3,076
- A much lower learning rate is used to avoid overwriting the pretrained weights too aggressively
- Training will be slower because gradients need to flow through the entire network
The expectation is that fine-tuning will outperform frozen BERT. The question worth thinking about is by how much, and whether the extra compute is always justified.
[cell 29 code]
# cell_9 — Define the fine-tuned BERT classifier
# MAIN GOAL: Build a classifier where the entire BERT encoder is trainable.
# Unlike the frozen classifier in cell_6, no parameters are frozen — the full
# 110M parameter model is updated end-to-end during training.
# Trade-offs vs the frozen approach:
# + BERT representations adapt to the AG News domain and classification task
# + Higher performance ceiling, especially on larger training sets
# - Much slower training — gradients flow through all 12 transformer layers
# - Higher risk of overfitting on small datasets
# - Requires a very small learning rate to avoid catastrophic forgetting of
# pretrained knowledge (see cell_10 for LR choice)
class FineTunedBERTClassifier(nn.Module):
def __init__(self, num_classes=4):
super(FineTunedBERTClassifier, self).__init__()
# ── Load pretrained BERT encoder ──────────────────────────────────────
# Identical to cell_6 — bert-base-uncased weights are loaded from cache.
# The key difference is what happens next: here we do NOT freeze anything.
# All 110M parameters start from their pretrained values and will be
# updated by the optimizer — pretrained weights are a starting point,
# not a fixed feature extractor.
self.bert = BertModel.from_pretrained("bert-base-uncased")
# ── No freezing ───────────────────────────────────────────────────────
# By default, all parameters have requires_grad=True after from_pretrained().
# We deliberately omit the freeze loop from cell_6 — no action needed here.
# This means the optimizer registered in cell_10 will update every layer:
# token embeddings, all 12 attention layers, and the pooler, in addition
# to the classification head below.
# ── Classification head ───────────────────────────────────────────────
# Identical architecture to cell_6: a single linear layer mapping the
# [CLS] embedding (768-dim) to class logits (4-dim).
# bert.config.hidden_size is used instead of hardcoding 768 for
# compatibility with other BERT variants (e.g. bert-large: 1024).
# This head is randomly initialised — its weights will converge quickly
# because the LR is the same for both head and encoder (see cell_10).
self.classifier = nn.Linear(self.bert.config.hidden_size, num_classes) # (768 → 4)
def forward(self, input_ids, attention_mask, token_type_ids):
# ── Forward pass through the full BERT encoder ────────────────────────
# Functionally identical to FrozenBERTClassifier.forward() in cell_6:
# three tensors are passed to BERT, all of shape (batch_size, seq_len):
# - input_ids : token IDs from BERT's WordPiece vocabulary
# - attention_mask : 1 for real tokens, 0 for [PAD]
# - token_type_ids : 0 for title segment, 1 for description segment
# The difference vs the frozen model is entirely in the backward pass:
# here gradients flow back through all 12 transformer layers and update
# their weights, including the segment-embedding lookup table.
outputs = self.bert(
input_ids = input_ids,
attention_mask = attention_mask,
token_type_ids = token_type_ids,
)
# ── Extract [CLS] token representation ────────────────────────────────
# last_hidden_state[:, 0, :] selects the [CLS] token (position 0) for
# every sample in the batch → shape: (batch_size, 768).
# Unlike the frozen model, these [CLS] embeddings will shift during
# training as BERT's weights are updated to better suit the task.
cls_output = outputs.last_hidden_state[:, 0, :] # shape: (batch_size, 768)
# ── Produce class logits ──────────────────────────────────────────────
# Raw unnormalised scores — CrossEntropyLoss in cell_10 applies softmax
# internally, so no activation function is applied here.
logits = self.classifier(cls_output) # shape: (batch_size, 4)
return logits
# ── Instantiate the model and move to device ──────────────────────────────────
# .to(DEVICE) moves all 110M parameters to GPU if available.
# Fine-tuning on CPU is possible but will be very slow given the full encoder
# is updated every step — a GPU is strongly recommended for this cell onwards.
finetuned_model = FineTunedBERTClassifier(num_classes=4).to(DEVICE)
# ── Count trainable vs total parameters ───────────────────────────────────────
# Unlike cell_6 where trainable ≈ 3,076 (head only), here trainable should equal
# total — confirming that no accidental freezing occurred.
# Expected: ~109,483,780 total, ~109,483,780 trainable (all parameters).
# If trainable < total, a freeze was accidentally applied somewhere — check __init__.
total_params = sum(p.numel() for p in finetuned_model.parameters())
trainable_params = sum(p.numel() for p in finetuned_model.parameters() if p.requires_grad)
print(f"Total parameters : {total_params:,}")
print(f"Trainable parameters : {trainable_params:,}")
# Note: if Total == Trainable, all parameters are unfrozen — this is expected here.
# Compare against cell_6 output to see the frozen vs fine-tuned parameter split clearly.
[cell 30 markdown]
### Discussion — Before Training
Take a moment to think about these two questions before running the training cell.
**Question 1:** In Part 2, a learning rate of 2e-3 was used to train only the classifier head. For fine-tuning, a much lower learning rate is used (2e-5). Why do you think a lower learning rate is necessary when the entire BERT encoder is being updated?
**Question 2:** At the start of fine-tuning, BERT's weights are meaningful because of pretraining. But the classifier head is randomly initialized. In the first few training steps, gradients flow from this random classifier all the way back through BERT. What problem could this cause, and how does the choice of a very small learning rate help avoid it?
Run the next cell and see how your intuition holds up. The answers are discussed after the results.
[cell 31 code]
# cell_10 — Training loop for fine-tuned BERT classifier
# MAIN GOAL: Train the entire BERT model end-to-end on the training subset.
# All 110M parameters are updated each step — both the 12-layer encoder and
# the classification head — allowing BERT to adapt its representations to
# the AG News classification task rather than acting as a fixed feature extractor.
# ── Hyperparameters ───────────────────────────────────────────────────────────
# EPOCHS: same as cell_7 (3) so training duration is directly comparable
# between frozen and fine-tuned runs.
# LR: 2e-5 is two orders of magnitude smaller than the 2e-3 used for frozen training.
# Two reasons for this:
# 1. Catastrophic forgetting: a large LR would overwrite the pretrained BERT
# weights too aggressively, destroying the linguistic knowledge learned
# during pretraining on BooksCorpus + Wikipedia.
# 2. Stability: with 110M interacting parameters all updating simultaneously,
# large gradient steps cause instability and loss spikes.
# 2e-5 is the most commonly cited LR in the original BERT fine-tuning paper
# and remains the standard starting point for bert-base fine-tuning tasks.
EPOCHS = 3
LR = 2e-5
# ── Loss function ─────────────────────────────────────────────────────────────
# Same CrossEntropyLoss as cell_7 — no change needed between frozen and fine-tuned.
# CrossEntropyLoss expects raw logits (not softmax outputs) and integer labels.
criterion = nn.CrossEntropyLoss()
# ── Optimizer ─────────────────────────────────────────────────────────────────
# AdamW is the standard optimizer for BERT fine-tuning. Two key differences
# from cell_7's optimizer:
# 1. finetuned_model.parameters() is passed instead of .classifier.parameters()
# — ALL parameters (encoder + head) are registered for updates.
# 2. weight_decay is omitted here (defaults to 0.0), unlike the 0.01 used in
# cell_7. With a tiny LR and pretrained weights, the regularisation effect
# of weight decay is less critical — the small LR already acts as an
# implicit constraint on how far weights can move from their pretrained values.
# Adding weight_decay=0.01 here is also valid and may help on small datasets.
optimizer = torch.optim.AdamW(
finetuned_model.parameters(), # Entire model — encoder + classifier head
lr = LR
)
# ── Reuse train_epoch from cell_7 ─────────────────────────────────────────────
# train_epoch() is model-agnostic — it accepts any nn.Module, DataLoader,
# optimizer, and criterion. No changes are needed to run fine-tuned training
# through the same loop, which guarantees the training procedure is identical
# and any performance difference is attributable to the model architecture
# (frozen vs unfrozen) rather than differences in the training loop.
# ── Run the full training loop ────────────────────────────────────────────────
print("Training Fine-tuned BERT Classifier...\n")
for epoch in range(1, EPOCHS + 1):
train_loss, train_acc = train_epoch(
finetuned_model, train_loader, optimizer, criterion, DEVICE
)
# Evaluate on the validation set after each epoch to monitor generalisation.
# Only val_acc is unpacked — the full metrics report is deferred to cell_11
# to avoid cluttering the training output, consistent with cell_7.
# The test set is never used here — it is held out until cell_11.
# Fine-tuned val_acc may improve more slowly early on (epoch 1) compared to
# the frozen model, since all layers need to co-adapt from their pretrained
# starting point — but should surpass it by epoch 3 on sufficient data.
val_acc, _, _, _, _, _, _ = evaluate(finetuned_model, val_loader, DEVICE)
print(
f"Epoch {epoch}/{EPOCHS} — "
f"Loss: {train_loss:.4f} | "
f"Train Acc: {train_acc:.4f} | "
f"Val Acc: {val_acc:.4f}"
)
[cell 32 code]
# cell_11 — Evaluate fine-tuned BERT on the test set
# MAIN GOAL: Run the fine-tuned BERT classifier on the held-out test set and
# report the full suite of evaluation metrics for the fine-tuned experiment.
# This is the first and only time the test set is used for the fine-tuned model —
# consistent with cell_8, which did the same for the frozen model.
# Results here are the final payoff of the notebook: a direct apples-to-apples
# comparison between frozen feature extraction and full end-to-end fine-tuning.
# ── Run evaluation on the test set ────────────────────────────────────────────
# evaluate() is defined in cell_6b — identical call to cell_8, just on a
# different model. Using the same function guarantees both sets of metrics
# are computed with the same averaging strategy (weighted) and the same
# class-name mapping, making the comparison valid.
# all_labels is discarded with _ because it was already captured in cell_8
# and is identical here — the test set labels do not change between evaluations.
# all_preds_finetuned is stored alongside all_preds_frozen from cell_8 so
# both prediction arrays are available together for the confusion matrix in cell_12.
(
finetuned_accuracy,
finetuned_precision,
finetuned_recall,
finetuned_f1,
finetuned_report,
_, # all_labels — already captured in cell_8, discarded here
all_preds_finetuned # Predicted class indices for the fine-tuned model (0–3)
) = evaluate(finetuned_model, test_loader, DEVICE)
# ── Print scalar metrics ───────────────────────────────────────────────────────
# Variable names are prefixed "finetuned_" throughout to avoid overwriting the
# "frozen_" variables from cell_8 — both sets must remain in memory for the
# side-by-side comparison printed in cell_12.
print("Fine-tuned BERT — Test Results")
print("-" * 35)
print(f"Accuracy : {finetuned_accuracy:.4f}") # Overall fraction of correct predictions
print(f"Precision : {finetuned_precision:.4f}") # Weighted avg: correctness of positive predictions
print(f"Recall : {finetuned_recall:.4f}") # Weighted avg: coverage of actual positives
print(f"F1 Score : {finetuned_f1:.4f}") # Weighted avg: harmonic mean of precision & recall
# ── Print per-class breakdown ─────────────────────────────────────────────────
# Compare this report against the frozen model's report from cell_8 class by class:
# - Which classes improved the most from fine-tuning?
# - Did any class get worse? (possible if the dataset is very small and the
# fine-tuned model overfit to the training distribution)
# - Are precision and recall more balanced per class than in the frozen model?
# On small datasets (100 training samples), the fine-tuned model may not
# outperform the frozen model — catastrophic forgetting or overfitting can
# hurt performance when there is too little data to guide the weight updates.
# On the full 120k training set, fine-tuning typically yields 3–5% higher accuracy.
print("\nPer-Class Report:")
print(finetuned_report)
[cell 33 markdown]
### Discussion — Answers
**Question 1:** In Part 2, a learning rate of 2e-3 was used to train only the classifier head. For fine-tuning, a much lower learning rate is used (2e-5). Why do you think a lower learning rate is necessary when the entire BERT encoder is being updated?
- BERT's pretrained weights encode valuable language understanding built over training on a massive corpus. If a large learning rate is used, these weights get updated too aggressively and that knowledge gets overwritten very quickly.
- This phenomenon is called catastrophic forgetting. The model forgets what it learned during pretraining because the updates are too large and too fast.
- A small learning rate ensures the pretrained weights are nudged gently in the direction that helps the new task, rather than being overwritten entirely.
- The classifier head can tolerate a higher learning rate because it starts from random initialization and has nothing to lose. The encoder cannot afford large updates because it has valuable pretrained knowledge that needs to be preserved.
**Question 2:** At the start of fine-tuning, BERT's weights are meaningful because of pretraining. But the classifier head is randomly initialized. In the first few training steps, gradients flow from this random classifier all the way back through BERT. What problem could this cause, and how does the choice of a very small learning rate help avoid it?
- In the very first training steps, the classifier head produces essentially random outputs because its weights are random. The gradients flowing back from this random classifier are also noisy and unreliable.
- If these noisy gradients are applied with a large learning rate, they could push BERT's carefully pretrained weights in a completely wrong direction before the classifier has had a chance to learn anything meaningful.
- A very small learning rate limits how much these noisy gradients can damage the pretrained weights in the early steps. By the time the classifier starts producing meaningful gradients, the encoder has not been significantly corrupted.
- Gradient clipping, which was added in cell_10, provides additional protection by capping the magnitude of gradients before they are applied. This ensures that no single update step causes a disproportionately large change to the weights.
[cell 34 markdown]
## Part 4 — Results Comparison
Both models have now been trained and evaluated. Before looking at the numbers side by side, it is worth thinking about what the comparison actually tells us.
Frozen BERT and fine-tuned BERT start from exactly the same pretrained weights. The only difference is what happens during training. In frozen BERT, the encoder is untouched and only a small classifier head learns the task. In fine-tuned BERT, the entire network is updated end-to-end.
A naive assumption would be that fine-tuning always wins by a large margin because more parameters are being updated. But the results tell a more nuanced story.
This part puts the two models side by side across all metrics and per-class performance. The goal is not just to declare a winner but to understand where each model does well, where it falls short, and what that tells us about the value of fine-tuning for this specific task.
[cell 35 code]
# cell_12 — Side-by-side results comparison
# MAIN GOAL: Consolidate all evaluation metrics from cell_8 (frozen) and cell_11
# (fine-tuned) into clean, readable comparison tables.
# Two levels of comparison are shown:
# 1. Overall metrics : accuracy, precision, recall, F1 — one row per model
# 2. Per-class F1 : F1 score broken down by class, with a signed difference
# column to immediately surface which classes benefited most from fine-tuning
# ── Imports ───────────────────────────────────────────────────────────────────
import pandas as pd # Used here for its clean tabular display via .to_string()
# ── Build the overall metrics summary table ───────────────────────────────────
# All four scalar metrics were computed in cell_8 and cell_11 using the same
# evaluate() function and weighted averaging strategy, so they are directly
# comparable without any further normalisation.
summary = pd.DataFrame({
"Model" : ["Frozen BERT", "Fine-tuned BERT"],
"Accuracy" : [frozen_accuracy, finetuned_accuracy],
"Precision" : [frozen_precision, finetuned_precision],
"Recall" : [frozen_recall, finetuned_recall],
"F1 Score" : [frozen_f1, finetuned_f1],
})
# Round all metric columns to 4 decimal places for consistent display width.
# Rounding is applied to the DataFrame copy only — the original float variables
# from cell_8 and cell_11 are not modified and remain available for further use.
metric_cols = ["Accuracy", "Precision", "Recall", "F1 Score"]
summary[metric_cols] = summary[metric_cols].round(4)
# to_string(index=False) suppresses the default 0/1 row index, which adds
# no information here and would clutter the side-by-side layout.
print("=" * 55)
print(" Overall Metrics Comparison")
print("=" * 55)
print(summary.to_string(index=False))
print("=" * 55)
# ── Compute per-class F1 scores ───────────────────────────────────────────────
# average=None returns one F1 score per class instead of a single aggregate.
# The order matches the integer label order (0=World, 1=Sports, 2=Business,
# 3=Sci/Tech), which is why the classes list below must follow the same order.
# all_labels and all_preds_frozen are from cell_8;
# all_preds_finetuned is from cell_11 — all three arrays cover the same test set.
classes = ["World", "Sports", "Business", "Sci/Tech"]
frozen_f1_per_class = f1_score(all_labels, all_preds_frozen,
average=None) # shape: (4,) — one F1 per class
finetuned_f1_per_class = f1_score(all_labels, all_preds_finetuned,
average=None) # shape: (4,) — one F1 per class
# ── Build the per-class comparison table ──────────────────────────────────────
# The "Difference" column = fine-tuned F1 − frozen F1 per class.
# Positive value → fine-tuning improved this class
# Negative value → fine-tuning hurt this class (possible on small datasets
# if the model overfit and generalises worse on that class)
# Zero → no change between the two approaches for this class
# This column is the most diagnostic output in the notebook — it shows not
# just whether fine-tuning helped overall, but which specific classes drove
# the improvement or regression.
per_class = pd.DataFrame({
"Class" : classes,
"Frozen BERT F1" : frozen_f1_per_class.round(4),
"Fine-tuned F1" : finetuned_f1_per_class.round(4),
"Difference" : (finetuned_f1_per_class - frozen_f1_per_class).round(4)
})
print("\nPer-Class F1 Comparison")
print("=" * 55)
print(per_class.to_string(index=False))
print("=" * 55)
# ── How to interpret these tables ────────────────────────────────────────────
# Overall table: if fine-tuned F1 > frozen F1, end-to-end training helped on
# this dataset size. If they are close or reversed, the dataset may be too
# small for fine-tuning to overcome the overfitting risk (see cell_11 note).
# Per-class table: look for classes with large positive differences — these
# are the categories where BERT's adapted representations added the most value.
# Classes with negative differences are worth investigating: they may have
# ambiguous vocabulary that fine-tuning memorised rather than generalised.
[cell 36 code]
# cell_13 — Visualize the comparison between frozen and fine-tuned BERT
# MAIN GOAL: Plot overall metrics and per-class F1 scores as grouped bar charts
# to make the frozen vs fine-tuned comparison immediately readable at a glance.
# Two subplots are shown side by side:
# Left : Accuracy, Precision, Recall, F1 — one grouped bar pair per metric
# Right : F1 score per AG News class — one grouped bar pair per class
# Both subplots share the same y-axis range so bars are visually comparable
# across plots without needing to mentally rescale.
# ── Imports ───────────────────────────────────────────────────────────────────
import matplotlib.pyplot as plt
import numpy as np # Used for arange() to compute bar x-positions arithmetically
# ── Prepare overall metric scores ─────────────────────────────────────────────
# These four scalars were computed in cell_8 (frozen) and cell_11 (fine-tuned).
# Listed in the same order as the metrics list so index alignment is guaranteed.
metrics = ["Accuracy", "Precision", "Recall", "F1 Score"]
frozen_scores = [frozen_accuracy, frozen_precision, frozen_recall, frozen_f1]
finetuned_scores = [finetuned_accuracy, finetuned_precision, finetuned_recall, finetuned_f1]
# x: evenly spaced tick positions (0, 1, 2, 3) for each metric group.
# width: each model's bar takes up 0.35 units; the two bars in a group together
# occupy 0.70 units, leaving 0.30 units of whitespace between groups.
x = np.arange(len(metrics))
width = 0.35
# ── Compute per-class F1 scores ───────────────────────────────────────────────
# average=None returns one F1 per class in label order (0=World … 3=Sci/Tech).
# Recomputed here (rather than reusing cell_12's variables) to keep this cell
# self-contained and runnable independently if cell_12 is skipped.
classes = ["World", "Sports", "Business", "Sci/Tech"]
frozen_f1_per_class = f1_score(all_labels, all_preds_frozen,
average=None) # shape: (4,)
finetuned_f1_per_class = f1_score(all_labels, all_preds_finetuned,
average=None) # shape: (4,)
# ── Compute a shared dynamic y-axis lower bound ───────────────────────────────
# Using a fixed y_min (e.g. 0.0) compresses all bars into the top of the chart
# when scores are clustered above 0.8, making small differences invisible.
# Instead, we find the lowest score across all bars in both subplots and
# subtract 0.05 as breathing room, then clamp to 0.0 to avoid a negative axis.
# Using the same y_min for both subplots means bar heights are directly
# comparable across the two plots — a taller bar always means a higher score.
min_score = min(
min(frozen_scores),
min(finetuned_scores),
frozen_f1_per_class.min(),
finetuned_f1_per_class.min()
)
y_min = max(0.0, min_score - 0.05) # 5% headroom below the lowest bar, never below 0
# ── Create the figure with two side-by-side subplots ─────────────────────────
# figsize=(14, 5): wide enough to give each subplot room for 4 labelled bar pairs.
fig, axes = plt.subplots(1, 2, figsize=(14, 5))
# ── Plot 1: Overall metrics comparison (left subplot) ─────────────────────────
ax1 = axes[0]
# Offset each model's bars by ±width/2 around the tick position so both bars
# in a group are centred on the metric label without overlapping.
bars1 = ax1.bar(x - width/2, frozen_scores, width,
label="Frozen BERT", color="#4C72B0") # Seaborn blue
bars2 = ax1.bar(x + width/2, finetuned_scores, width,
label="Fine-tuned BERT", color="#DD8452") # Seaborn orange
ax1.set_title("Overall Metrics Comparison")
ax1.set_xticks(x)
ax1.set_xticklabels(metrics)
ax1.set_ylabel("Score")
ax1.set_ylim(y_min, 1.0) # Shared lower bound — see dynamic y_min calculation above
ax1.legend()
# bar_label annotates each bar with its exact numeric value above the bar.
# fmt="%.4f" matches the 4-decimal precision used in cell_12's printed tables.
# fontsize=8 prevents labels from overlapping on closely spaced bars.
ax1.bar_label(bars1, fmt="%.4f", padding=3, fontsize=8)
ax1.bar_label(bars2, fmt="%.4f", padding=3, fontsize=8)
# ── Plot 2: Per-class F1 comparison (right subplot) ───────────────────────────
ax2 = axes[1]
# x2 is a separate position array for 4 classes — same width logic as plot 1.
x2 = np.arange(len(classes))
bars3 = ax2.bar(x2 - width/2, frozen_f1_per_class, width,
label="Frozen BERT", color="#4C72B0")
bars4 = ax2.bar(x2 + width/2, finetuned_f1_per_class, width,
label="Fine-tuned BERT", color="#DD8452")
ax2.set_title("Per-Class F1 Comparison")
ax2.set_xticks(x2)
ax2.set_xticklabels(classes)
ax2.set_ylabel("F1 Score")
ax2.set_ylim(y_min, 1.0) # Same y range as plot 1 — bars are height-comparable across subplots
ax2.legend()
ax2.bar_label(bars3, fmt="%.4f", padding=3, fontsize=8)
ax2.bar_label(bars4, fmt="%.4f", padding=3, fontsize=8)
# ── Figure-level title and layout ─────────────────────────────────────────────
# suptitle sits above both subplots and acts as the overall figure heading.
# tight_layout() adjusts subplot spacing to prevent the suptitle and axis
# labels from overlapping — especially important with bar labels on top of bars.
plt.suptitle("Frozen BERT vs Fine-tuned BERT", fontsize=14, fontweight="bold")
plt.tight_layout()
plt.savefig("comparison.png", dpi=150, bbox_inches="tight")
plt.show()
[cell 37 markdown]
# Notebook Summary: BERT Text Classification on AG News
## What Was Built
A Jupyter notebook that trains and compares two BERT-based classifiers on the AG News dataset, a 4-class news classification task covering World, Sports, Business, and Sci/Tech.
---
## Part 1 -- Data Exploration and Prep
- AG News dataset was loaded with 120,000 training samples and 7,600 test samples, perfectly balanced across 4 classes
- Title and Description were kept as separate columns and passed to the BERT tokenizer as a sentence pair, which inserts a `[SEP]` token between them and assigns segment ids (0 for title, 1 for description) so BERT can apply its pretrained segment embeddings correctly
- EDA was performed to check class distribution, text length distribution, and sample inspection per class
- Token length analysis confirmed that 99.28% of samples fall within 128 tokens, justifying the max_len choice
- The training subset is controlled by `SAMPLES_PER_CLASS` in cell_3 (e.g. 25 per class = 100 rows for CPU runs, 500 per class = 2,000 rows for GPU runs, or set `USE_FULL_DATASET = True` for the full 120,000)
- The test subset is controlled by `TEST_SAMPLES_PER_CLASS` in cell_3 and is sampled with stratification, because the raw `test.csv` is not class-balanced in row order — a naive `nrows=N` load was observed to produce a skewed test set (37/30/21/12 across the four classes)
- BERT tokenizer was loaded, a custom PyTorch Dataset class was defined, and DataLoaders were built
---
## Part 2 -- Pretrained BERT (Frozen Encoder)
- BERT encoder was completely frozen. Only a linear classifier head (3,076 parameters out of 110 million) was trained
- Learning rate of 2e-3 was used since only the classifier head was being updated
- The model was evaluated on the stratified test subset using Accuracy, Precision, Recall, and F1 Score
---
## Part 3 -- Fine-tuned BERT
- All 110 million BERT parameters were unfrozen and trained end-to-end
- Learning rate of 2e-5 was used to avoid catastrophic forgetting of pretrained weights
- Gradient clipping (max_norm=1.0) was applied after loss.backward() and before optimizer.step() to protect pretrained weights from noisy gradients in early training steps
- AdamW optimizer was used, which is standard practice for fine-tuning BERT
---
## Part 4 -- Results Comparison
- Both models were evaluated and compared using Accuracy, Precision, Recall, and F1 Score
- Per-class F1 scores were visualized side by side using bar charts
---
## About the Results Below
The numbers in the tables below come from **earlier GPU runs** at two scales — a 2,000-sample subset and the full 120,000-sample dataset — both evaluated on the full 7,600-row test set. They are preserved here because they show how the frozen-vs-fine-tuned gap behaves at meaningful dataset scales.
The notebook as currently configured (CPU, 100-sample training subset, 100-sample stratified test set) will produce different and noisier numbers — typically lower overall accuracy and larger per-class swings.(~25 test samples per class)
To reproduce these numbers, set `SAMPLES_PER_CLASS = 500` (or `USE_FULL_DATASET = True`), restore the full test set by setting `TEST_SAMPLES_PER_CLASS` to a large value or using `USE_FULL_DATASET = True`, and run on a GPU for faster run.
---
## Results
### Subset (2,000 training samples, full 7,600-sample test set, GPU)
| Model | Accuracy | Precision | Recall | F1 Score |
|---|---|---|---|---|
| Frozen BERT | 0.8592 | 0.8626 | 0.8592 | 0.8595 |
| Fine-tuned BERT | 0.8989 | 0.9014 | 0.8989 | 0.8995 |
### Per-Class F1 -- Subset
| Class | Frozen BERT F1 | Fine-tuned BERT F1 | Difference |
|---|---|---|---|
| World | 0.8708 | 0.9085 | +0.0377 |
| Sports | 0.9375 | 0.9715 | +0.0340 |
| Business | 0.8120 | 0.8566 | +0.0446 |
| Sci/Tech | 0.8178 | 0.8615 | +0.0437 |
### Full Dataset (120,000 training samples, full 7,600-sample test set, GPU)
| Model | Accuracy | Precision | Recall | F1 Score |
|---|---|---|---|---|
| Frozen BERT | 0.9021 | 0.9022 | 0.9021 | 0.9016 |
| Fine-tuned BERT | 0.9430 | 0.9431 | 0.9430 | 0.9430 |
### Per-Class F1 -- Full Dataset
| Class | Frozen BERT F1 | Fine-tuned BERT F1 | Difference |
|---|---|---|---|
| World | 0.9074 | 0.9574 | +0.0499 |
| Sports | 0.9698 | 0.9871 | +0.0173 |
| Business | 0.8575 | 0.9121 | +0.0546 |
| Sci/Tech | 0.8716 | 0.9154 | +0.0438 |
---
## Key Findings
### On the Subset (2,000 samples)
- Frozen BERT achieved 85.92% accuracy with only 3,076 trainable parameters out of 110 million. This confirms that pretrained representations are already rich enough to produce strong results with minimal training.
- Fine-tuning improved accuracy by 3.97 percentage points over frozen BERT. This is a meaningful improvement considering both models were trained on only 2,000 samples.
- Fine-tuning consistently improved every single class. The largest gains were in Business (+4.46 F1 points) and Sci/Tech (+4.37 F1 points), the two classes with the most vocabulary overlap with other categories.
- Sports was the strongest class in both models (93.75% and 97.15% F1) due to its highly distinct vocabulary.
- Business was the weakest class in both models (81.20% and 85.66% F1) due to significant language overlap with World and Sci/Tech.
### On the Full Dataset (120,000 samples)
- Both models improved substantially with more data. Frozen BERT gained 4.29 accuracy points and fine-tuned BERT gained 4.41 accuracy points compared to the subset results.
- Fine-tuned BERT reached 94.30% accuracy on the full dataset, a strong result for a 4-class classification task.
- Sports showed the smallest fine-tuning gain on the full dataset (1.73 F1 points) because frozen BERT already performed well on it. Business showed the largest gain (5.46 F1 points), confirming that fine-tuning helps most where class boundaries are ambiguous.
- Sports achieved near perfect classification at 98.71% F1 with fine-tuning. Business remained the hardest class at 91.21% F1 but improved significantly from the subset results.
### Overall Takeaway
The gap between frozen and fine-tuned BERT stays consistently around 4 percentage points across both dataset sizes (3.97 on the subset, 4.09 on the full dataset). This suggests that for this task, the pretrained representations are already well suited to news classification and fine-tuning provides a consistent but not dramatically increasing advantage as data grows. On small datasets, frozen BERT is a strong and compute efficient choice. As the dataset grows, fine-tuning increasingly justifies its cost. The decision between the two should be driven by dataset size, available compute, and the acceptable margin of error for the task.