YAML Metadata Warning:empty or missing yaml metadata in repo card

Check out the documentation for more information.

🧠 BERT-based 4-Class News Classifier

A fine-tuned BERT-base-uncased model for classifying news articles into four categories: World, Sport, Business, and Science & Technology.

πŸ“– Model Overview

This model adapts BERT for multi-class sequence classification.
It was trained on a labeled news dataset with four categories and achieves strong performance on short text snippets such as news headlines or article summaries.

Label ID Category
0 World
1 Sport
2 Business
3 Science & Technology

βš™οΈ Training Details

  • Base model: bert-base-uncased
  • Task: Sequence classification (4 classes)
  • Training environment: Google Colab (T4 GPU)
  • Max sequence length: 200
  • Batch size: 8
  • Epochs: 3
  • Optimizer: AdamW (lr=2e-5)
  • Framework: PyTorch + πŸ€— Transformers Trainer API

The dataset was tokenized using AutoTokenizer with truncation and dynamic padding enabled.
Model and tokenizer were saved using save_pretrained() and uploaded to the Hugging Face Hub.

🧩 Usage

You can load and use the model directly with the πŸ€— Transformers library:

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

# Load model from the Hub
model_name = "zacanadir/wsbt-sequence-classifier" 
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

MAX_LENGTH = 200

def predict_text(text_list):
    encodings = tokenizer(
        text_list,
        padding=True,
        truncation=True,
        max_length=MAX_LENGTH,
        return_tensors="pt"
    ).to(model.device)

    model.eval()
    with torch.no_grad():
        outputs = model(**encodings)
        preds = torch.argmax(outputs.logits, dim=-1)
    return preds.tolist()

# Example
sample_texts = [
    "Nvidia becomes the first chipmaker to top $5 trillion in market value.",
    "The Raja of Casablanca advances to the semi final of the African Champions' League."
]

label_to_category = {
    0: "World",
    1: "Sport",
    2: "Business",
    3: "Science & Technology"
}

preds = predict_text(sample_texts)
predicted_categories = [label_to_category[int(i)] for i in preds]
print(predicted_categories)

🧾 Output Example
['Business', 'Sport']

πŸš€ How to Reuse / Fine-tune Further

If you’d like to fine-tune this model on your own dataset:

from transformers import Trainer, TrainingArguments, AutoTokenizer, AutoModelForSequenceClassification
from datasets import load_dataset

tokenizer = AutoTokenizer.from_pretrained("zacanadir/wsbt-sequence-classifier")
model = AutoModelForSequenceClassification.from_pretrained("zacanadir/wsbt-sequence-classifier", num_labels=4)

# Load your dataset and fine-tune

πŸ“¦ Files Included
File	                                            Description
config.json	                                        Model architecture & label mapping
pytorch_model.bin	                                Trained weights
tokenizer.json, vocab.txt, special_tokens_map.json	Tokenizer files
training_args.bin	                                Trainer configuration (epochs, LR, etc.)

✨ Acknowledgements
Hugging Face Transformers
Google Colab for free GPU training
PyTorch

Trained and deployed with ❀️ using Hugging Face and Google Colab.
Downloads last month
8
Safetensors
Model size
0.1B params
Tensor type
F32
Β·
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support

Space using zacanadir/wsbt-sequence-classifier 1