metadata
datasets:
- SetFit/ag_news
language:
- en
metrics:
- f1
base_model:
- google-bert/bert-base-uncased
pipeline_tag: text-classification
tags:
- ag_news_finetuning
- Bert
- Bert_un_cased
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)