base_model: distilbert-base-uncased
tags:
- text-classification
- distilbert
- ag-news
- news-classification
- mlops
datasets: Recurrent/prepared_data_mlops2
language: en
metrics:
- accuracy
- f1
pipeline_tag: text-classification
Model Card for Model ID
A DistilBERT model fine-tuned for 4-class news topic classification on the AG News dataset. Given a short news snippet, it predicts one of four categories: World, Sports, Business, Sci/Tech.
Model Details
Model Description
This is a distilbert-base-uncased model fine-tuned for sequence classification on AG News (4 classes). It was trained as part of an MLOps assignment workflow on Kaggle (GPU T4 x2), with experiment tracking via Weights & Biases and the best checkpoint pushed to the Hugging Face Hub. Two hyperparameter configurations were trained and compared, and the version with the higher weighted F1 score was published.
- Developed by: LKHUMANLEIMA (PGD AI, IIT Jodhpur — MLOps group assignment)
- Model type: DistilBERT for sequence classification (
DistilBertForSequenceClassification)
- Language(s) (NLP): English
- License: apache-2.0 (inherited from
distilbert-base-uncased)
- Finetuned from model [optional]:
distilbert-base-uncased
Model Sources [optional]
Uses
Direct Use
The model classifies a short English news headline or article snippet into one of four topics. Label mapping:
| ID | Label |
|---|
| 0 | World |
| 1 | Sports |
| 2 | Business |
| 3 | Sci/Tech |
Typical uses include news feed categorization, content tagging, and routing of news text by topic.
Downstream Use [optional]
Can serve as a starting point for related topic-classification tasks (e.g., further fine-tuning on a different news taxonomy) or be wrapped behind an inference API/CI pipeline.
Out-of-Scope Use
Not suitable for long documents — inputs are truncated to 128 tokens.
- Trained only on English AG News–style text; performance on other domains (social media, legal, medical, non-English) is unverified.
- Should not be used as the sole basis for high-stakes decisions, and is not intended for any harmful, deceptive, or surveillance application.
Bias, Risks, and Limitations
The model inherits biases present in distilbert-base-uncased and in the AG News corpus, which is drawn from a fixed period of English-language news sources. It may reflect the topical conventions and any imbalances of that data, may misclassify text styles unlike news articles, and confuses semantically overlapping categories (e.g., Business vs. Sci/Tech for tech-business stories). The 128-token limit means information beyond the start of a longer article is ignored.
Recommendations
Users (both direct and downstream) should be made aware of the risks, biases and limitations of the model. More information needed for further recommendations.
How to Get Started with the Model
1from transformers import pipeline
2
3clf = pipeline("text-classification", model="LKHUMANLEIMA/ag-news-distilbert")
4print(clf("NASA's new telescope captured images of a distant galaxy."))
5# -> [{'label': 'Sci/Tech', 'score': ...}]
Or load the components directly:
1from transformers import AutoTokenizer, AutoModelForSequenceClassification
2
3tokenizer = AutoTokenizer.from_pretrained("LKHUMANLEIMA/ag-news-distilbert")
4model = AutoModelForSequenceClassification.from_pretrained("LKHUMANLEIMA/ag-news-distilbert")
Training Details
Training Data
Prepared AG News dataset hosted on the Hub:
Recurrent/prepared_data_mlops2 (loaded from the
prepared_data directory). The data is a 4-class topic-classification set with
text and
label fields.
Training Procedure
Preprocessing [optional]
Texts were tokenized with the distilbert-base-uncased tokenizer using padding="max_length", truncation=True, and max_length=128
Training Hyperparameters
Two versions were trained from the pretrained checkpoint (each started fresh, not from the other's checkpoint) using the 🤗 Trainer. Common settings: weight_decay=0.01, max_length=128, eval and save per epoch, load_best_model_at_end=True, metric_for_best_model="f1".
| Setting | Version 1 | Version 2 |
|---|
| Learning rate | 2e-5 | 5e-5 |
| Train batch size | 16 | 32 |
| Epochs | 3.5 | 3 |
| Weight decay | 0.01 | 0.01 |
| Eval batch size | 32 | 32 |
The published checkpoint is whichever version achieved the higher weighted F1 on the test split.
- Training regime: fp32 (default
Trainer settings)
Evaluation
Testing Data, Factors & Metrics
Testing Data
Metrics
- Accuracy — overall fraction of correctly classified examples.
- F1 (weighted) — weighted-average F1 across the four classes, used as the model-selection criterion to account for any class imbalance.
Results
| Metric | Version 1 | Version 2 |
|---|
| Accuracy | TODO | TODO |
| F1 (weighted) | TODO | TODO |
| Eval loss | TODO | TODO |
Best version published: TODO (v1 or v2)
Summary
DistilBERT fine-tuned on AG News reaches strong topic-classification performance; the exact figures for this run are available in the W&B project mlops-assignment3.
Environmental Impact
Carbon emissions can be estimated using the
Machine Learning Impact calculator presented in
Lacoste et al. (2019).
- Hardware Type: NVIDIA Tesla T4 x2 (Kaggle)
- Cloud Provider: Kaggle
Technical Specifications
Model Architecture and Objective
DistilBERT encoder with a sequence-classification head (4 output labels), trained with a cross-entropy objective for single-label classification.
Compute Infrastructure
Hardware
NVIDIA Tesla T4 x2 GPUs (Kaggle notebook environment).
Software
Python with 🤗 transformers, datasets, accelerate, scikit-learn (metrics), and wandb (tracking). Secrets (WANDB_API_KEY, HF_TOKEN, GITHUB_TOKEN) were managed via Kaggle Secrets and never hardcoded.
Citation
If you use this model, please cite the base model and dataset.
DistilBERT (BibTeX):
1@article{sanh2019distilbert,
2 title={DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter},
3 author={Sanh, Victor and Debut, Lysandre and Chaumond, Julien and Wolf, Thomas},
4 journal={arXiv preprint arXiv:1910.01108},
5 year={2019}
6}
AG News (BibTeX):
1@inproceedings{zhang2015character,
2 title={Character-level convolutional networks for text classification},
3 author={Zhang, Xiang and Zhao, Junbo and LeCun, Yann},
4 booktitle={Advances in Neural Information Processing Systems},
5 year={2015}
6}
Model Card Authors
LKHUMANLEIMA (Group 13, MLOps assignment — PGD AI, IIT Jodhpur).
Model Card Contact
