BERT AG News 3-Category Classifier

Fine-tuned BERT model for classifying English news texts into three categories:

• Sports
• Business
• Technology

This model is based on bert-base-uncased and was fine-tuned on a modified version of the AG News dataset.

Model Details

• Base model: bert-base-uncased
• Architecture: BERT for sequence classification
• Number of labels: 3
• Framework: PyTorch
• Library: Transformers
• Task: Text Classification
• Language: English

Dataset

The model was trained on frostMorn/ag-news-3-category-dataset.

The dataset is based on the original AG News dataset:

fancyzhx/ag_news

The original dataset contains four categories:

• World
• Sports
• Business
• Sci/Tech

For this project, the World category was removed and the remaining categories were renamed:

The dataset contains approximately 90,000 training examples and 5,700 test examples.

Training

The model was fine-tuned using the Hugging Face Trainer.

Training configuration:

• Epochs: 3
• Learning rate: 2e-5
• Training batch size: 16
• Base model: bert-base-uncased
• Number of classes: 3

The model was trained on a Google Colab Tesla T4 GPU.

Usage

You can use the model directly with the Hugging Face Transformers pipeline:

from transformers import pipeline

classifier = pipeline(
    "text-classification",
    model="frostMorn/bert-ag-news-3-category"
)

text = "Apple announced a new computer processor."

result = classifier(
    text,
    top_k=3
)

for prediction in result:
    print(
        f"{prediction['label']}: "
        f"{prediction['score']:.3f}"
    )