Sample AG News Classifier

Model description

This is a news classification model built which utilized a pre-trained model "bert-base-uncased" as a tokenizer and has a custom multi-layer percentron (MLP) for news classification. This model is case-insensitive: there is No difference between english and English.

Intended uses and limitations

The model has been trained to classify news headlines into 4 categories: World, Sports, Business and Science/Technology.

Note: Currently limited to above 4 categories only

How to use

The model has been developed using PyTorch and since it is a custom neural network architecture, it can be used by performing the below steps and not using the Transformers library,

1. Create your model class

import torch
import torch.nn as nn
from huggingface_hub import PyTorchModelHubMixin


class TextClassifier(nn.Module, PyTorchModelHubMixin):
    def __init__(self, vocab_size=30522, embed_dim=128, num_classes=4):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, embed_dim)
        self.fc1 = nn.Linear(embed_dim, 128)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(128, num_classes)

    def forward(self, input_ids, attention_mask=None):
        x = self.embedding(input_ids)

        if attention_mask is not None:
            mask = attention_mask.unsqueeze(-1).float()
            x = x * mask
            x = x.sum(dim=1) / mask.sum(dim=1).clamp(min=1e-9)
        else:
            x = x.mean(dim=1)

        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

model = TextClassifier.from_pretrained("pulkitchowdry/sample-agnews-classifer")
model.eval()

2. Setup Tokenizer

from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
text      = "Messi scores a hatrick in the world cup"

inputs  = tokenizer(
                text, 
                return_tensors="pt", 
                padding="max_length", 
                truncation=True, 
                max_length=128
              )

3. Run the prediction

with torch.no_grad():
  logits  = model(
      inputs["input_ids"],
      attention_mask=inputs["attention_mask"]
  )
  prediction  = logits.argmax(dim=1).item()

print(f"{prediction}")
categories = ["World", "Sports", "Business", "Science/Technology"]

print("Predicted class: ", categories[prediction])

Limitations and bias

Even if the training data used for this model could be characterized as fairly neutral, this model can have biased predictions

Training data

The model was trained on AG News data available on https://huggingface.co/datasets/fancyzhx/ag_news

Architecture

• Tokenization - BERT based
• Multi-layer perceptron - Linear layers with ReLU as activation function
• Learning rate - 5e-4
• Cost function - Cross Entropy Loss
• Optimizer - Adaptive Moment Estimation (Adam)

Training procedure

1. Tokenization - Text from the training data was tokenized using a vocabulary size of 30,552.
2. Embeddings - Tokens were converted into embeddings using a dimension of 128.
3. Muti-layer perceptron - Data passed through two linear layers and a relu layer as an activation function.
4. PyTorch autograd used for backpropgation.
5. Adam optimizer used for learning with the initial learning rate set as 5e-4.

Evaluation results

Accuracy - 91% based on the test data provided in https://huggingface.co/datasets/fancyzhx/ag_news

Citation

@misc{pulkitchowdry_2026_sample_ag_news,
      title = {Sample AG News Classifier},
      author = {Pulkit Chowdry},
      year = {2026},
}

This model has been pushed to the Hub using the PytorchModelHubMixin