Model description
This model is a fine-tuned version of the bert-base-uncased model to classify news articles into one of four categories: World(label 0), Sports(label 1), Business(label 2), Sci/Tech(label 3).
How to use
You can use the model with the following code.
from transformers import BertForSequenceClassification, BertTokenizer, TextClassificationPipeline
model_path = "JiaqiLee/bert-agnews"
tokenizer = BertTokenizer.from_pretrained(model_path)
model = BertForSequenceClassification.from_pretrained(model_path, num_labels=4)
pipeline = TextClassificationPipeline(model=model, tokenizer=tokenizer)
print(pipeline("Google scores first-day bump of 18 (USATODAY.com): USATODAY.com - Even a big first-day jump in shares of Google (GOOG) couldn't quiet debate over whether the Internet search engine's contentious auction was a hit or a flop."))
Training data
The training data comes from HuggingFace AGNews dataset. We use 90% of the train.csv
data to train the model and the remaining 10% for evaluation.
Evaluation results
The model achieves 0.9447 classification accuracy in AGNews test dataset.
- Downloads last month
- 188
This model does not have enough activity to be deployed to Inference API (serverless) yet.
Increase its social visibility and check back later, or deploy to Inference Endpoints (dedicated)
instead.