BERT News Classification (Fine-tuned on AG News)

This model is a fine-tuned version of bert-base-uncased for news topic classification on the AG News dataset.
It predicts one of four categories: World, Sports, Business, Sci/Tech.

Model Details

  • Base Model: bert-base-uncased
  • Fine-tuning Dataset: SetFit/ag_news
  • Language: English
  • Task: Text Classification (Single-label)
  • Number of Labels: 4
  • Training Framework: Hugging Face Transformers
  • Evaluation Metric: Weighted F1-score

Training Procedure

  • Epochs: 3
  • Batch Size: 16
  • Learning Rate: 5e-5
  • Weight Decay: 0.01
  • Mixed Precision: FP16 enabled
  • Early Stopping: Patience of 1 epochs

Training & Evaluation Results

Epoch Training Loss Validation Loss F1 Score
1 0.2005 0.1875 0.9422
2 0.1171 0.1948 0.9470
3 0.0853 0.2458 0.9457

Final Training Metrics:

  • Global Step: 22500
  • Training Loss: 0.1546
  • Train Runtime: 2593s (~43min)
  • Samples/sec: 231.39
  • Steps/sec: 14.46
  • Total FLOPs: 1.57e+16

Intended Use

The model can be used for:

  • Classifying short news articles into World, Sports, Business, Sci/Tech.
  • Benchmarking text classification pipelines with BERT.

How to Use

from transformers import AutoTokenizer, AutoModelForSequenceClassification

model_name = "Maarij-Aqeel/BERT_news_classification"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

text = "The stock market saw a significant rise today as..."
inputs = tokenizer(text, return_tensors="pt")
outputs = model(**inputs)
predicted_label = outputs.logits.argmax(dim=-1).item()
print(predicted_label)
Downloads last month

-

Downloads are not tracked for this model. How to track
Inference Providers NEW
This model isn't deployed by any Inference Provider. 🙋 Ask for provider support

Model tree for Maarij-Aqeel/BERT_news_classification

Finetuned
(5693)
this model

Dataset used to train Maarij-Aqeel/BERT_news_classification