# 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 ```sh pip install transformers datasets torch ``` ### Loading and Predicting ```python 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.