You can deploy this model and start earning money today!

pulkitchowdry/
sample-agnews-classifer

Name: sample-agnews-classifer
Availability: InStock

safetensors

model_hub_mixin

pytorch_model_hub_mixin

text-classification

fancyzhx/ag_news

mit

Views

No views yet

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,

Create your model class

python

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()

Setup Tokenizer

python

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
              )

Run the prediction

python

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

Tokenization - Text from the training data was tokenized using a vocabulary size of 30,552.
Embeddings - Tokens were converted into embeddings using a dimension of 128.
Muti-layer perceptron - Data passed through two linear layers and a relu layer as an activation function.
PyTorch autograd used for backpropgation.
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