RoBERTa-Based Topic Classification Model Using AG News Dataset

This repository hosts a RoBERTa-based transformer model fine-tuned for topic classification on the AG News dataset. The model identifies topics such as World, Sports, Business, and Science/Technology in a given news text.

Model Details

Model Architecture: RoBERTa (roberta-base)
Task: Topic Classification
Dataset: AG News (from Hugging Face Datasets)
Fine-tuning Framework: Hugging Face Transformers

Usage

Installation

pip install transformers datasets torch

Loading and Predicting

from transformers import RobertaTokenizer, RobertaForSequenceClassification
import torch

# Load model and tokenizer
model = RobertaForSequenceClassification.from_pretrained("your-saved-model-directory")
tokenizer = RobertaTokenizer.from_pretrained("roberta-base")
model.eval()

# Sample prediction
def predict_topic(texts, model, tokenizer, device='cpu'):
    import re
    if isinstance(texts, str):
        texts = [texts]

    def preprocess(text):
        text = text.lower()
        text = re.sub(r"http\S+|www\S+|https\S+", '', text)
        text = re.sub(r'\@\w+|\#', '', text)
        text = re.sub(r"[^a-zA-Z0-9\s.,!?']", '', text)
        text = re.sub(r'\s+', ' ', text).strip()
        return text

    cleaned_texts = [preprocess(t) for t in texts]
    inputs = tokenizer(cleaned_texts, padding=True, truncation=True, return_tensors="pt").to(device)

    model.to(device)
    model.eval()
    with torch.no_grad():
        outputs = model(**inputs)
        preds = torch.argmax(outputs.logits, dim=1).tolist()

    label_map = {0: "World", 1: "Sports", 2: "Business", 3: "Sci/Tech"}
    return [label_map[p] for p in preds]

# Example
sample_texts = [
    "The stock market witnessed a major crash today due to inflation concerns.",
    "The new space telescope has captured unprecedented images of distant galaxies."
]
results = predict_topic(sample_texts, model, tokenizer)
for text, label in zip(sample_texts, results):
    print(f"Text: {text}\nPredicted Topic: {label}\n")

Performance Metrics

Accuracy: ~0.96 on AG News test split

Fine-Tuning Details

Dataset

The dataset is sourced from the AG News dataset available via Hugging Face Datasets.

Training

Number of epochs: 3
Batch size: 8
Evaluation strategy: epoch
Learning rate: 2e-5

Repository Structure

.
├── model/               # Contains the fine-tuned model files
├── tokenizer/           # Tokenizer configuration and vocab
├── README.md            # Model documentation

Limitations

The model is trained specifically for AG News-style texts and may not generalize well to informal or unrelated domains.

Contributing

Contributions are welcome! Please open an issue or submit a PR for suggestions or improvements.