| # 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. |