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
Paper [optional]: DistilBERT: A Distilled Version of BERT: Smaller, Faster, Cheaper, and Lighter
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:

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:

World
Sports
Business
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 [email protected]

GitHub

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

Adiii143
/

distilbert-base-uncased-finetuned-ag_news_AK