Resource	Link
GitHub Repository	riteshmaury-iitj/group13-assignment-mlops
HuggingFace Model	YuvarajK-g25ait2054/ag-news-distilbert
W&B Project	mlops-assignment3
Training Notebook (Kaggle)	group13-ag-news-k
Docker Image	yuvarajkg25ait2054/ag-news-classifier

Quick Start

Using the Pipeline (Recommended)

python

from transformers import AutoTokenizer, AutoModelForSequenceClassification, pipeline

model_name = "YuvarajK-g25ait2054/ag-news-distilbert"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

classifier = pipeline(
    "text-classification",
    model=model,
    tokenizer=tokenizer,
    top_k=None
)

text = "The stock market rallied today as tech companies reported strong earnings."
results = classifier(text)
print(results)
# [{'label': 'LABEL_2', 'score': 0.97}, ...]  -> Business

Using PyTorch Directly

python

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

model_name = "YuvarajK-g25ait2054/ag-news-distilbert"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

text = "NASA launches new Mars exploration mission next spring."
inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=128)

with torch.no_grad():
    outputs = model(**inputs)
    predicted_class = outputs.logits.argmax(-1).item()

id2label = {0: "World", 1: "Sports", 2: "Business", 3: "Sci/Tech"}
print(f"Category: {id2label[predicted_class]}")
# Category: Sci/Tech

Using Docker

bash

docker pull yuvarajkg25ait2054/ag-news-classifier:latest

docker run --rm yuvarajkg25ait2054/ag-news-classifier:latest \
  --text "NASA launches new Mars exploration mission." \
  --model "YuvarajK-g25ait2054/ag-news-distilbert"

Training Details

Dataset

Split	Samples
Train	120,000
Test	7,600

Source: ag_news loaded directly from HuggingFace Datasets.

Preprocessing applied (same logic used at training and inference):

Lowercase all text
Remove HTML tags (<...>)
Remove URLs (http://, www.)
Remove punctuation
Normalize whitespace

Hyperparameters

Parameter	Version 1 ✅ (Best)	Version 2
Learning rate	`2e-5`	`5e-5`
Batch size (train)	`16`	`64`
Epochs	`3`	`5`
Weight decay	`0.01`	`0.01`
Max sequence length	`128`	`128`
Optimizer	AdamW	AdamW
Evaluation strategy	per epoch	per epoch
Best model selection	highest F1	highest F1

Training Infrastructure

Hardware: 2× NVIDIA Tesla T4 GPU (Kaggle)
Framework: PyTorch 2.5.1 + Transformers 4.46.0
Training time: ~43 min (v1) · ~55 min (v2)
Training throughput: ~140 samples/sec (v1) · ~183 samples/sec (v2)

Evaluation Results

Per-Epoch Metrics

Version 1 (lr=2e-5, bs=16, 3 epochs):

Epoch	Train Loss	Val Loss	Accuracy	F1
1	0.4509	0.3915	0.9367	0.9366
2	0.2923	0.3680	0.9434	0.9434
3	0.2149	0.4186	0.9425	0.9425

Version 2 (lr=5e-5, bs=64, 5 epochs):

Epoch	Train Loss	Val Loss	Accuracy	F1
1	0.3924	0.3855	0.9338	0.9337
2	0.2816	0.3575	0.9411	0.9411
3	0.1703	0.4123	0.9382	0.9381
4	0.1048	0.4613	0.9389	0.9389
5	0.0575	0.5555	0.9371	0.9371

Final Comparison

Metric	Version 1 ✅	Version 2
Accuracy	0.9434 (94.34%)	0.9412 (94.12%)
F1 (weighted)	0.9434	0.9412
Eval Loss	0.3679	0.3575

Selected model: Version 1 — higher accuracy and F1 score.

W&B experiment runs:

Uses

Direct Use

Classify English news headlines/articles into one of four categories: World, Sports, Business, Sci/Tech.

Downstream Use

News aggregation and routing pipelines
Content recommendation systems
Automated tagging and document organisation

Out-of-Scope

Non-English text
News categories beyond the 4 trained classes
Production critical systems without further validation

Bias, Risks, and Limitations

Trained on AG News corpus — performance may degrade on news domains not well-represented in the dataset
May misclassify articles that span multiple categories (e.g., sports business news)
This is an academic project model; validate thoroughly before any production deployment

Technical Specifications

Property	Value
Base architecture	DistilBERT (`distilbert-base-uncased`)
Parameters	~66 million
Model size	~260 MB
Number of labels	4
Max input length	128 tokens
Precision	float32

Citation

bibtex

@misc{ag-news-distilbert-2026,
  author    = {Kosuru, Yuvaraj},
  title     = {AG News Text Classification with DistilBERT},
  year      = {2026},
  publisher = {HuggingFace},
  howpublished = {\url{https://huggingface.co/YuvarajK-g25ait2054/ag-news-distilbert}}
}

Model Card Author

Yuvaraj Kosuru — g25ait2054@iitj.ac.in
IIT Jodhpur · MLOps Assignment · Group 13

For questions, open an issue at github.com/riteshmaury-iitj/group13-assignment-mlops/issues