---
language: en
license: mit
tags:
  - text-classification
  - pytorch
  - distilbert
  - news
  - ag-news
datasets:
  - sh0416/ag_news
metrics:
  - accuracy
model-index:
  - name: ag-news-classifier
    results:
      - task:
          type: text-classification
        dataset:
          name: AG News
          type: ag_news
        metrics:
          - type: accuracy
            value: 0.9433
---

# AG News Text Classifier

A fine-tuned [DistilBERT](https://huggingface.co/distilbert-base-uncased) model for classifying news articles into 4 categories: **World**, **Sports**, **Business**, and **Sci/Tech**.

## Model Performance

| Metric | Value |
|--------|-------|
| Test Accuracy | **94.33%** |
| Dataset | AG News (120,000 train / 7,600 test) |
| Baseline (random) | 25% |

For reference, state of the art on AG News is ~95-96%, making this model competitive with top published results.

---

## How to Use

```python
import torch
from transformers import AutoTokenizer
from huggingface_hub import hf_hub_download

# 1. Load tokenizer
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

# 2. Define the model class
import torch.nn as nn
from transformers import AutoModel

class AGNewsClassifier(nn.Module):
    def __init__(self, num_classes=4, dropout=0.3):
        super().__init__()
        self.bert = AutoModel.from_pretrained("distilbert-base-uncased")
        self.dropout = nn.Dropout(dropout)
        self.classifier = nn.Linear(768, num_classes)

    def forward(self, input_ids, attention_mask):
        outputs = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        cls_output = outputs.last_hidden_state[:, 0, :]
        dropped = self.dropout(cls_output)
        return self.classifier(dropped)

# 3. Load model weights
model = AGNewsClassifier()
weights_path = hf_hub_download("your_username/ag-news-classifier", "best_model.pt")
model.load_state_dict(torch.load(weights_path, map_location="cpu"))
model.eval()

# 4. Classify new text
LABELS = {0: "World", 1: "Sports", 2: "Business", 3: "Sci/Tech"}

def classify(text):
    tokens = tokenizer(
        text,
        max_length=128,
        padding="max_length",
        truncation=True,
        return_tensors="pt"
    )
    with torch.no_grad():
        logits = model(tokens["input_ids"], tokens["attention_mask"])
    pred = logits.argmax(dim=1).item()
    return LABELS[pred]

# Example
print(classify("Apple stock surges after record quarterly earnings"))
# → Business

print(classify("Scientists discover potential signs of water on Mars"))
# → Sci/Tech
```

---

## Training Details

### Dataset
- **Source**: [AG News](https://huggingface.co/datasets/ag_news)
- **Train size**: 120,000 samples (30,000 per class)
- **Test size**: 7,600 samples (1,900 per class)
- **Classes**: Perfectly balanced across all 4 categories

### Input Format
Title and description were combined using a separator token:
```
[TITLE]{title}[SEP]{description}
```
Max token length in the dataset after combining: 383 tokens. Since 99.9% of samples fit within 128 tokens, `max_length=128` was chosen for efficiency with negligible information loss.

### Model Architecture
```
Input (title + description)
        ↓
DistilBERT (distilbert-base-uncased, 66M parameters)
        ↓
[CLS] token hidden state (768-dim)
        ↓
Dropout (p=0.3)
        ↓
Linear layer (768 → 4)
        ↓
Logits (4 classes)
```

### Hyperparameters

| Parameter | Value |
|-----------|-------|
| Base model | distilbert-base-uncased |
| Max sequence length | 128 |
| Batch size | 64 |
| Optimizer | AdamW |
| Learning rate | 2e-5 |
| Epochs | 3 |
| Dropout | 0.3 |

### Training Environment
- Platform: Kaggle Notebooks
- GPU: Tesla T4 x2
- Framework: PyTorch + HuggingFace Transformers

---

## Limitations

- Trained on English news articles only — performance on other languages or domains (e.g. social media, academic text) may degrade
- Model was fine-tuned for 3 epochs; further tuning with a validation-based early stopping strategy could improve results

---

## About

This model was built as part of a data science portfolio project to demonstrate fine-tuning transformer models for NLP classification tasks using PyTorch and HuggingFace Transformers.