---
library_name: transformers
tags:
- distilbert
- text classification
- ag news
- fine-tuned
license: apache-2.0
datasets:
- fancyzhx/ag_news
metrics:
- accuracy
base_model:
- distilbert/distilbert-base-uncased
pipeline_tag: text-classification
---

# Model Card for distilbert-base-uncased-finetuned-ag-news

This is a fine-tuned version of the **DistilBERT** model, specifically trained on the **AG News dataset** for text classification tasks. This model can predict one of the four categories: World, Sports, Business, and Science/Technology based on news article text.

## Model Details

### Model Description

This model is based on the **DistilBERT** architecture, a smaller, faster, and lighter version of BERT. It has been fine-tuned on the **AG News dataset** for the task of text classification. The model predicts which of the four categories (World, Sports, Business, or Science/Technology) a news article belongs to.

DistilBERT retains 97% of BERT’s performance while being 60% faster and 30% smaller, making it suitable for efficient inference in production environments.

This model card has been automatically generated by 🤗 transformers for the model hosted on the Hugging Face Hub.

- **Developed by:** Hugging Face
- **Shared by [optional]:** Aditya AK
- **Model type:** Transformer-based model for text classification
- **Language(s) (NLP):** English
- **License:** Apache-2.0
- **Finetuned from model [optional]:** distilbert-base-uncased

### Model Sources [optional]

- **Repository:** [Hugging Face - distilbert-base-uncased-finetuned-ag-news](https://huggingface.co/distilbert-base-uncased-finetuned-ag-news)
- **Paper [optional]:** [DistilBERT: A Distilled Version of BERT: Smaller, Faster, Cheaper, and Lighter](https://arxiv.org/abs/1910.01108)
- **Demo [optional]:** [More Information Needed]

## Uses

### Direct Use

This model is designed to classify news articles into one of four categories: **World**, **Sports**, **Business**, or **Science/Technology**.

To use the model, you can input the text of a news article and the model will predict its category.

### Downstream Use [optional]

This model can be used as part of a larger text classification pipeline. For example, it can be integrated into a news aggregation application to automatically categorize incoming articles into their appropriate categories.

### Out-of-Scope Use

This model is not suitable for tasks that require fine-grained or specialized classification beyond the four categories it was trained on. It also may not perform well on texts that are significantly different from the **AG News** dataset in terms of style, domain, or structure.

## Bias, Risks, and Limitations

- **Biases**: The model may inherit biases from the **AG News** dataset, which could reflect certain geographic or cultural perspectives, especially in the news content.
- **Risks**: The model might incorrectly classify news articles, especially if the content is ambiguous or falls into multiple categories.
- **Limitations**: The model is limited to categorizing articles into one of four predefined classes. It may not perform well on articles that don’t fit these categories or require fine-grained classification.

### Recommendations

Users should be cautious when using the model for real-world applications, as it might produce incorrect or biased classifications. It is recommended to fine-tune the model on domain-specific data if a more accurate or refined classification system is required.

## How to Get Started with the Model

Use the code below to get started with the **DistilBERT** model fine-tuned on the **AG News** dataset for text classification:

```python
from transformers import AutoModelForSequenceClassification, AutoTokenizer

model = "distilbert-base-uncased-finetuned-ag-news"
tokenizer = AutoTokenizer.from_pretrained(model)
model_distilbert = AutoModelForSequenceClassification.from_pretrained(model).to(device)

# Example text for classification
text = "Apple announces new iPhone model with improved features."

# Tokenize the input text
inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)

# Perform classification
outputs = model_distilbert(**inputs)
predictions = outputs.logits.argmax(dim=-1)

# Map the predicted class index to the category
categories = ["World", "Sports", "Business", "Science/Technology"]
predicted_category = categories[predictions.item()]

print(f"Predicted category: {predicted_category}")
```

### Training Details
This model was fine-tuned on the AG News dataset, which consists of news articles categorized into four classes:
1. World
2. Sports
3. Business
4. Science/Technology
The dataset contains approximately 120,000 training samples and 7,600 test samples.

Dataset: AG News Dataset
Preprocessing: The text was tokenized, and labels were assigned based on the article's category. Some additional cleaning steps, such as removing unwanted characters, might have been applied.

### Model Card Authors: 
    Aditya AK

### Model Card Contact
For questions or further inquiries, please contact aakuskar.980@gmail.com

### GitHub 
  Github Noteebook Link -- https://github.com/Adity-star/DataScience-Work/blob/main/NLP/Finetuned_on_Ag_News.ipynb