AG News Classification with BERT

Fine-tuning bert-base-uncased on the AG News dataset for 4-class news headline classification.

Task

Given a news headline and its article body, classify it into one of four categories:

Model

Training code & demo: github.com/alfiobonanno/AGNews-Bert-Tester-App — includes a live demo (FastAPI + React) where you can test the model with your own headlines.

The fine-tuned model is hosted here on Hugging Face Hub as Alfiob12/agnews-bert.

Use the HF model directly with the transformers pipeline:

from transformers import pipeline

clf = pipeline("text-classification", model="Alfiob12/agnews-bert")
clf("NVIDIA announces production of groundbreaking quantum GPUs")
# [{'label': 'Sci/Tech', 'score': 0.961}]

Results

Evaluated on a held-out test set of 19,140 samples (15% of the full dataset).

Per-class breakdown:

Training

The notebook AG _news_Classification.ipynb contains the full pipeline: data loading, tokenization, training, evaluation, and inference.

Hyperparameters

Data split

The original train/test CSV files are merged (127,600 samples) and re-split to create a validation set:

• Train: 70% (89,320 samples)
• Validation: 15% (19,140 samples)
• Test: 15% (19,140 samples)

Splits are stratified by class. The dataset is perfectly balanced (equal samples per class).

Setup

git clone https://github.com/alfiobonanno/AGNews-Bert-Tester-App.git
cd AGNews-Bert-Tester-App
python -m venv venv
source venv/bin/activate   # macOS/Linux
# venv\Scripts\activate    # Windows
pip install -r requirements.txt

The dataset CSVs are included in ag_data/. Then open and run AG _news_Classification.ipynb.

Dependencies

See requirements.txt for the full list. Key packages:

• transformers==5.12.1
• datasets==5.0.0
• scikit-learn==1.9.0
• pandas==3.0.3
• plotly==6.8.0
• numpy==2.4.6
• safetensors==0.8.0