JiaqiLee/bert-agnews

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.