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)
Model tree for Maarij-Aqeel/BERT_news_classification
Base model
google-bert/bert-base-uncased