|
--- |
|
license: afl-3.0 |
|
datasets: |
|
- ag_news |
|
language: |
|
- en |
|
metrics: |
|
- accuracy |
|
library_name: transformers |
|
pipeline_tag: text-classification |
|
--- |
|
|
|
## Model description |
|
This model is a fine-tuned version of the [bert-base-uncased](https://huggingface.co/transformers/model_doc/bert.html) 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. |
|
|
|
```python |
|
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](https://huggingface.co/datasets/ag_news). 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. |
