UK Headlines Topic Classifier

Set up and run a before/after fine-tuning project: train a DistilBERT headline classifier on general (AG News) titles, evaluate on held-out UK BBC headlines, fine-tune on UK data, then compare accuracy, macro-F1, and micro-F1 on the same UK test set. Also report a Multinomial Naive Bayes baseline (course-aligned).

1. Goal (one sentence)

Show that a news-headline topic classifier adapted to UK news scores higher on a frozen UK test set after UK fine-tuning than the same model trained only on general AG News headlines.

2. Hardware and runtime assumptions

DistilBERT is small; prefer Hugging Face Transformers + PyTorch MPS over MLX-LM (MLX-LM is aimed at generative LLMs, not this classifier).

3. Locked stack

4. Label mapping (mandatory)

Unify both datasets to these 4 labels so before/after comparison is valid:

Drop all BBC entertainment rows (no clean AG counterpart). Document that AG World ↔ BBC politics is an imperfect but intentional mapping.

5. Folder layout to create

Create everything under uk_headlines/ (this directory):

uk_headlines/
  AGENT_PLAN.html          # this brief (already present)
  README.md                # short human summary of how to run
  requirements.txt
  prepare_data.py
  train_stage1.py          # AG News DistilBERT
  train_stage2.py          # continue fine-tune on UK (from AG)
  train_uk_only.py         # DistilBERT on UK only (no AG)
  train_nb.py              # NB @ AG and NB @ UK
  evaluate.py
  data/                    # prepared CSV/JSONL (gitignored if large)
  artifacts/
    model_ag_headlines/
    model_uk_finetuned/
    model_uk_only/
    nb_ag.joblib
    nb_uk.joblib
    metrics_before.json
    metrics_after.json
    metrics_uk_only.json
    results_table.md
    confusion_*.png        # optional
  .gitignore               # ignore data caches, checkpoints, __pycache__

6. Environment setup (do this first)

1. Work from repo root: cambridge2026/, project path uk_headlines/.
2. Create a local venv (do not install into system Python):

   cd uk_headlines
   python3 -m venv .venv
   source .venv/bin/activate
   python -m pip install --upgrade pip
   pip install torch torchvision torchaudio
   pip install transformers datasets accelerate scikit-learn pandas numpy tqdm matplotlib joblib

   Pin reasonable versions in requirements.txt after a successful install.
3. Verify MPS:

   python -c "import torch; print('mps:', torch.backends.mps.is_available())"

   If MPS is available, set device to mps. If not, fall back to CPU and note it in the README.
4. Ensure Hugging Face downloads work (network). No API key required for these public datasets/models. Cache will land under ~/.cache/huggingface/ by default — that is fine.
5. Add uk_headlines/.gitignore covering: .venv/, __pycache__/, data/raw/, large artifacts/model_*, *.pt, HF temp files. Prefer committing small metric JSON / results markdown, not multi‑hundred‑MB weights, unless the user asks to keep checkpoints.

7. Data preparation pipeline

1. AG News — load fancyzhx/ag_news; keep title + mapped unified label. Cap Stage-1 training with a stratified subset of about 20k–40k titles (faster, still strong). Keep a small AG val split for Stage-1 early monitoring if useful.
2. BBC UK — load SetFit/bbc-news; map labels; drop entertainment.
3. Headline text for BBC — if rows are full articles, build a headline proxy: prefer a real title field if present; otherwise first sentence or first ~12–20 whitespace tokens. Document the rule in README.
4. UK split — stratified 70% train / 15% val / 15% test, fixed seed (e.g. 42). Freeze the UK test set. Never use it for training, early stopping, or model selection.
5. Save prepared files under data/ (e.g. ag_train.csv, uk_train.csv, uk_val.csv, uk_test.csv) with columns text,label.

8. Training and evaluation

Stage 1 — basic DistilBERT (“before”)

• Init from distilbert-base-uncased.
• Train on AG titles → 4 labels. Max length short (e.g. 64).
• Save to artifacts/model_ag_headlines/.
• Evaluate on UK test only → artifacts/metrics_before.json (accuracy, macro-F1, micro-F1, per-class F1).

Stage 2 — UK fine-tune (“after”)

• Load Stage-1 checkpoint (do not re-init from base).
• Continue train on UK train; early-stop on UK val (about 1–3 epochs; small data — avoid overfit).
• Save to artifacts/model_uk_finetuned/.
• Evaluate on the same UK test → artifacts/metrics_after.json.

UK only — fine-tune from base on UK (no AG)

• Init from distilbert-base-uncased (do not load Stage 1).
• Train on UK train; early-stop on UK val (same recipe as Stage 2).
• Save to artifacts/model_uk_only/.
• Evaluate on the same UK test → artifacts/metrics_uk_only.json.
• Isolates UK data alone vs AG → UK continued training.

Naive Bayes baseline

• Train NB on AG titles; score UK test.
• Retrain NB on UK train; score same UK test.
• Save models with joblib under artifacts/.

Results table (write artifacts/results_table.md)

Add a short error analysis: UK-specific terms / entities the before model misses (British politics, Premier League, spellings) that improve after fine-tuning. Save as artifacts/error_analysis.md.

9. Hard constraints

• DO use the same UK test set for before and after.
• DO report both macro-F1 and micro-F1 (macro is sensitive to rare classes; micro / accuracy reflect overall instance-level performance — compare them).
• DO keep scripts runnable with source .venv/bin/activate then python prepare_data.py etc.
• DO NOT train on UK test or peek at test labels while tuning hyperparameters beyond val.
• DO NOT train a large generative LLM from scratch.
• DO NOT commit api_key.txt or dump huge raw HF caches into git.
• DO NOT rewrite unrelated lecture practicals.

10. Suggested run order

cd uk_headlines
source .venv/bin/activate
python prepare_data.py
python train_nb.py
python train_stage1.py
python evaluate.py --checkpoint artifacts/model_ag_headlines --out artifacts/metrics_before.json
python train_stage2.py
python evaluate.py --checkpoint artifacts/model_uk_finetuned --out artifacts/metrics_after.json
python train_uk_only.py
python evaluate.py --checkpoint artifacts/model_uk_only --out artifacts/metrics_uk_only.json --confusion-name confusion_uk_only.png
# then assemble results_table.md + error_analysis.md

Script names may be consolidated into one notebook if clearer, but keep the same stages and artifact paths so a later agent can resume.

11. Definition of done

• Done. Venv + requirements.txt work on Mac Studio with MPS verified.
• Done. Prepared AG + UK CSVs exist; UK split is stratified and frozen.
• Done. Before and after DistilBERT metrics on UK test are recorded.
• Done. NB @ AG and NB @ UK metrics recorded.
• Done. artifacts/results_table.md shows after ≥ before on macro-F1 and micro-F1.
• Done. README explains how to reproduce in under a page.

12. Risks (acknowledge in write-up)

• World ↔ politics mapping is imperfect.
• BBC rows may be article bodies; headline proxy must be documented.
• UK set is small — use early stopping and modest Stage-2 epochs.