# AG News classfication 100 GC v4 PDF
course: Module 4 — Generative AI & LLMs
module: Module-4-Generative-AI-LLMs
type: pdf
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_PDF.pdf
pages: 56
---
[page 1]
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.
[page 2]
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.
Enter physical GPU index (e.g. 0, 1, 2) or 'cpu': cpu
Device set to: cpu
# 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_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
[page 3]
All imports successful.
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_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
[page 4]
Device: cpu — no GPU info to display.
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.
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_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"]
[page 5]
Train size : (120000, 3)
Test size : (7600, 3)
Class distribution (train):
label
1 30000
2 30000
3 30000
4 30000
Name: count, dtype: int64
Sample rows:
label title description
0 3 Wall St. Bears Claw Back Into the
Black (Reuters)
Reuters - Short-sellers, Wall
Street's dwindli...
1 3 Carlyle Looks Toward Commercial
Aerospace (Reu...
Reuters - Private investment firm
Carlyle Grou...
2 3 Oil and Economy Cloud Stocks'
Outlook (Reuters)
Reuters - Soaring crude prices
plus worries\ab...
label title description
0 3 Wall St. Bears Claw Back
Into the Black (Reuters)
Reuters - Short-sellers, Wall
Street's dwindli...
1 3
Carlyle Looks Toward
Commercial Aerospace
(Reu...
Reuters - Private investment firm
Carlyle Grou...
# ── 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)
train_df
[page 6]
label title description
2 3 Oil and Economy Cloud
Stocks' Outlook (Reuters)
Reuters - Soaring crude prices
plus worries\ab...
3 3 Iraq Halts Oil Exports from
Main Southern Pipe...
Reuters - Authorities have halted
oil export\f...
4 3 Oil prices soar to all-time
record, posing new...
AFP - Tearaway world oil prices,
toppling reco...
... ... ... ...
119995 1 Pakistan's Musharraf Says
Won't Quit as Army C...
KARACHI (Reuters) - Pakistani
President Perve...
119996 2 Renteria signing a top-shelf
deal
Red Sox general manager Theo
Epstein acknowled...
119997 2 Saban not going to Dolphins
yet
The Miami Dolphins will put
their courtship of...
119998 2 Today's NFL games PITTSBURGH at NY GIANTS
Time: 1:30 p.m. Line: ...
119999 2 Nets get Carter from Raptors INDIANAPOLIS -- All-Star
Vince Carter was trad...
120000 rows × 3 columns
label title description
0 3 Fears for T N pension after
talks
Unions representing workers at
Turner Newall...
1 4 The Race is On: Second
Private Team Sets Launc...
SPACE.com - TORONTO, Canada
-- A second\team o...
2 4 Ky. Company Wins Grant to
Study Peptides (AP)
AP - A company founded by a
chemistry research...
3 4 Prediction Unit Helps Forecast
Wildfires (AP)
AP - It's barely dawn when Mike
Fitzpatrick st...
4 4 Calif. Aims to Limit Farm-
Related Smog (AP)
AP - Southern California's smog-
fighting agenc...
... ... ... ...
7595 1 Around the world Ukrainian presidential candidate
Viktor Yushch...
7596 2 Void is filled with Clement With the supply of attractive
pitching options...
7597 2 Martinez leaves bitter Like Roger Clemens did almost
exactly eight ye...
7598 3 5 of arthritis patients in
Singapore take Bext...
SINGAPORE : Doctors in the
United States have ...
test_df
[page 7]
label title description
7599 3 EBay gets into rentals EBay plans to buy the apartment
and home renta...
7600 rows × 3 columns
label
3 1900
4 1900
2 1900
1 1900
Name: count, dtype: int64
print(test_df["label"].value_counts())
# 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
[page 8]
Label range: 0 to 3
Sample title : Wall St. Bears Claw Back Into the Black (Reuters)
Sample description : Reuters - Short-sellers, Wall Street's dwindling\band of
ultra-cynics, are seeing green again.
Label : 2
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.
# 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_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
───────────────────────────────────────────────────
[page 9]
# 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()
[page 10]
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.
[page 11]
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_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].
[page 12]
# 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",
[page 13]
/usr/local/lib/python3.12/dist-packages/huggingface_hub/utils/_auth.py:93:
UserWarning:
The secret `HF_TOKEN` does not exist in your Colab secrets.
To authenticate with the Hugging Face Hub, create a token in your settings
tab (https://huggingface.co/settings/tokens), set it as secret in your Google
Colab and restart your session.
You will be able to reuse this secret in all of your notebooks.
Please note that authentication is recommended but still optional to access
public models or datasets.
warnings.warn(
Warning: You are sending unauthenticated requests to the HF Hub. Please set a
HF_TOKEN to enable higher rate limits and faster downloads.
WARNING:huggingface_hub.utils._http:Warning: You are sending unauthenticated
requests to the HF Hub. Please set a HF_TOKEN to enable higher rate limits
and faster downloads.
Text Length Statistics (Full Training Set)
========================================
word_count char_count token_count
count 120000.00 120000.00 120000.00
mean 37.84 236.46 54.16
std 10.09 66.53 19.06
min 4.00 17.00 9.00
25% 32.00 196.00 44.00
50% 37.00 232.00 52.00
75% 43.00 266.00 60.00
max 177.00 1012.00 380.00
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()
{"model_id":"cb4fbb7e35834e26b968dcfd61540133","version_major":2,"version_minor":0}
{"model_id":"ebe44d6095a147019d88f34cf8746946","version_major":2,"version_minor":0}
{"model_id":"860490cbd08b44a6b0a107eac920a098","version_major":2,"version_minor":0}
[page 14]
Samples exceeding 512 tokens (BERT hard limit) : 0 (0.00%)
Samples exceeding 128 tokens (our max_len) : 956 (0.80%)
# 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)
[page 15]
======================================================================
Sample Texts per Class
======================================================================
Class: World
----------------------------------------------------------------------
Sample 1: Explosion Rocks Baghdad Neighborhood | BAGHDAD, Iraq, August 24 --
A car bomb exploded near the gate of a US-funded Iraqi television network in
Baghdad on Tuesday, killing at least two people and wounding two others,
authorities and witnesses said.
Sample 2: BBC reporters' log | BBC correspondents record events in the Middle
East and their thoughts as the funeral of the Palestinian leader Yasser
Arafat takes place.
# ── 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)
[page 16]
Class: Sports
----------------------------------------------------------------------
Sample 1: Second Andre win in a row boosts US gold medal hopes | Andre
Dirrell, fighting with a tattoo of his grandfather #39;s face on his back,
assured the United States of at least two boxing medals Wednesday by narrowly
beating Cuba #39;s Yordani Despaigne to advance to the Olympic middleweight
semifinals...
Sample 2: NFL Game Summary - Kansas City at New Orleans | Tight end Tony
Gonzalez caught six passes for 71 yards for the Chiefs, who have scored 38
points on their initial drive of the game this season...Safety Jay Bellamy
and Fakhir Brown each made nine tackles for the Saints, whose last home win
over Kansas...
Class: Business
----------------------------------------------------------------------
Sample 1: US house sales fall in July | Sales of non-new houses in the US
fell last month but still exceeded analyst forecasts.
Sample 2: DJ to acquire MarketWatch | Dow Jones amp; Co., publisher of the
Wall Street Journal, says it has agreed to buy online news site MarketWatch
Inc. in a \$519 million deal that widens its Internet audience.
Class: Sci/Tech
----------------------------------------------------------------------
Sample 1: Gartner optimistic about chip numbers | But that optimism isn #39;t
matched by Infineon, which said that while the market worldwide remains
buoyant, the US is a special case and cautioned that growth might not be as
strong as it expected.
Sample 2: Fossil Indicates Brown Bears Went South | CALGARY, Alberta - A
fossilized jaw found in an Alberta gravel pit may have cleared up a mystery
for researchers wondering how brown bears originally made it across North
America.
======================================================================
# 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.
[page 17]
# ── 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.")
[page 18]
# ── 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.
[page 19]
Using train subset: 25 samples per class → 100 total rows.
Using test subset: 25 samples per class → 100 total rows.
Train size : (90, 4)
Val size : (10, 4)
Test size : (100, 3)
Class distribution in train:
label
0 22
1 23
2 22
3 23
Name: count, dtype: int64
Class distribution in val:
label
0 3
1 2
2 3
3 2
Name: count, dtype: int64
Class distribution in test:
label
0 25
1 25
2 25
3 25
Name: count, dtype: int64
/tmp/ipykernel_1336/3197071602.py:40: DeprecationWarning:
DataFrameGroupBy.apply operated on the grouping columns. This behavior is
deprecated, and in a future version of pandas the grouping columns will be
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())
[page 20]
excluded from the operation. Either pass `include_groups=False` to exclude
the groupings or explicitly select the grouping columns after groupby to
silence this warning.
.apply(lambda x: x.sample(SAMPLES_PER_CLASS, # Draw fixed N per
group
/tmp/ipykernel_1336/3197071602.py:63: DeprecationWarning:
DataFrameGroupBy.apply operated on the grouping columns. This behavior is
deprecated, and in a future version of pandas the grouping columns will be
excluded from the operation. Either pass `include_groups=False` to exclude
the groupings or explicitly select the grouping columns after groupby to
silence this warning.
.apply(lambda x: x.sample(TEST_SAMPLES_PER_CLASS, # Draw fixed N per
group
# 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,
[page 21]
Token IDs : tensor([[ 101, 14324, 2003, 1037, 3928, 2653, 2944,
1012, 102, 0]])
Attention Mask : tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 0]])
Tokens : ['[CLS]', 'bert', 'is', 'a', 'powerful', 'language',
'model', '.', '[SEP]', '[PAD]']
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_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.
[page 22]
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
[page 23]
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]),
[page 24]
input_ids shape : torch.Size([128])
attention_mask shape : torch.Size([128])
token_type_ids shape : torch.Size([128])
label : 0
Segment A tokens (title + padding): 87
Segment B tokens (description) : 41
Decoded (first 120 chars): [CLS] the china post [SEP] president chen shui
bian said yesterday he does not expect china to ease its military threat
# 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_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
─────────────────────────────────────────────────────────────
[page 25]
# 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.
[page 26]
Train batches : 3
Val batches : 1
Test batches : 4
Batch size : 32
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.
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.
# 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}")
[page 27]
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:
The Classifier Head
A single linear layer is placed on top of the [CLS] vector:
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.
cls_output = outputs.last_hidden_state[:, 0, :] # shape: (batch_size, 768)
nn.Linear(768, 4)
[page 28]
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_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
────────────────────────────────────────
[page 29]
# 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,
)
[page 30]
BertModel LOAD REPORT from: bert-base-uncased
Key | Status | |
-------------------------------------------+------------+--+-
cls.predictions.transform.LayerNorm.bias | UNEXPECTED | |
cls.predictions.transform.dense.weight | UNEXPECTED | |
cls.predictions.bias | UNEXPECTED | |
# ── 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:,}")
{"model_id":"4e4bf28f6cb6481fb52bdafc7626efe2","version_major":2,"version_minor":0}
{"model_id":"3fd4505e19ab4ca2a79139a279f15e81","version_major":2,"version_minor":0}
{"model_id":"7b7dfdbf30de41f8bb6dd40c9139ded2","version_major":2,"version_minor":0}
[page 31]
cls.seq_relationship.bias | UNEXPECTED | |
cls.predictions.transform.dense.bias | UNEXPECTED | |
cls.predictions.transform.LayerNorm.weight | UNEXPECTED | |
cls.seq_relationship.weight | UNEXPECTED | |
Notes:
- UNEXPECTED :can be ignored when loading from different
task/architecture; not ok if you expect identical arch.
Total parameters : 109,485,316
Trainable parameters : 3,076
Frozen parameters : 109,482,240
# 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
──────────────────────────────────────────
[page 32]
# 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
[page 33]
evaluate() defined and ready.
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?
# 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.")
[page 34]
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_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
—
[page 35]
# 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.
[page 36]
# 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).
[page 37]
Training Frozen BERT Classifier...
Epoch 1/3 — Loss: 1.4857 | Train Acc: 0.2556 | Val Acc: 0.5000
Epoch 2/3 — Loss: 1.2636 | Train Acc: 0.4667 | Val Acc: 0.5000
Epoch 3/3 — Loss: 1.1784 | Train Acc: 0.5000 | Val Acc: 0.6000
# 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_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.
[page 38]
# ── 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)
[page 39]
Frozen BERT — Test Results
-----------------------------------
Accuracy : 0.6200
Precision : 0.7174
Recall : 0.6200
F1 Score : 0.5923
Per-Class Report:
precision recall f1-score support
World 0.48 1.00 0.65 25
Sports 0.83 0.20 0.32 25
Business 0.89 0.64 0.74 25
Sci/Tech 0.67 0.64 0.65 25
accuracy 0.62 100
macro avg 0.72 0.62 0.59 100
weighted avg 0.72 0.62 0.59 100
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.
[page 40]
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.
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_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):
[page 41]
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,
[page 42]
BertModel LOAD REPORT from: bert-base-uncased
Key | Status | |
-------------------------------------------+------------+--+-
cls.predictions.transform.LayerNorm.bias | UNEXPECTED | |
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.
{"model_id":"4bfb35f07587428095a7cf5f36f7eb84","version_major":2,"version_minor":0}
[page 43]
cls.predictions.transform.dense.weight | UNEXPECTED | |
cls.predictions.bias | UNEXPECTED | |
cls.seq_relationship.bias | UNEXPECTED | |
cls.predictions.transform.dense.bias | UNEXPECTED | |
cls.predictions.transform.LayerNorm.weight | UNEXPECTED | |
cls.seq_relationship.weight | UNEXPECTED | |
Notes:
- UNEXPECTED :can be ignored when loading from different
task/architecture; not ok if you expect identical arch.
Total parameters : 109,485,316
Trainable parameters : 109,485,316
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_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.
[page 44]
# 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
)
[page 45]
Training Fine-tuned BERT Classifier...
Epoch 1/3 — Loss: 1.3704 | Train Acc: 0.2889 | Val Acc: 0.3000
Epoch 2/3 — Loss: 1.1399 | Train Acc: 0.6333 | Val Acc: 0.8000
Epoch 3/3 — Loss: 0.8763 | Train Acc: 0.9333 | Val Acc: 0.8000
# 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_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
[page 46]
Fine-tuned BERT — Test Results
-----------------------------------
# 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)
[page 47]
Accuracy : 0.8000
Precision : 0.8282
Recall : 0.8000
F1 Score : 0.7971
Per-Class Report:
precision recall f1-score support
World 0.88 0.84 0.86 25
Sports 0.96 0.96 0.96 25
Business 0.61 0.88 0.72 25
Sci/Tech 0.87 0.52 0.65 25
accuracy 0.80 100
macro avg 0.83 0.80 0.80 100
weighted avg 0.83 0.80 0.80 100
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.
[page 48]
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_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"]
[page 49]
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)
[page 50]
=======================================================
Overall Metrics Comparison
=======================================================
Model Accuracy Precision Recall F1 Score
Frozen BERT 0.62 0.7174 0.62 0.5923
Fine-tuned BERT 0.80 0.8282 0.80 0.7971
=======================================================
Per-Class F1 Comparison
=======================================================
Class Frozen BERT F1 Fine-tuned F1 Difference
World 0.6494 0.8571 0.2078
Sports 0.3226 0.9600 0.6374
Business 0.7442 0.7213 -0.0229
Sci/Tech 0.6531 0.6500 -0.0031
=======================================================
# ── 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_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"]
[page 51]
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
[page 52]
# 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()
[page 53]
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
plt.savefig("comparison.png", dpi=150, bbox_inches="tight")
plt.show()
[page 54]
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.
[page 55]
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.
[page 56]
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.