AventIQ-AI / distilbert-spam-detection

DistilBERT Base Uncased Quantized Model for Spam Detection

This repository hosts a quantized version of the DistilBERT model, fine-tuned for spam detection tasks. The model has been optimized for efficient deployment while maintaining high accuracy, making it suitable for resource-constrained environments.

Model Details

• Model Architecture: DistilBERT Base Uncased
• Task: Spam Detection
• Dataset: Hugging Face's sms_spam
• Quantization: BrainFloat16
• Fine-tuning Framework: Hugging Face Transformers

Usage

Installation

pip install transformers torch

Loading the Model

from transformers import DistilBertTokenizer, DistilBertForSequenceClassification
import torch

model_name = "AventIQ-AI/distilbert-spam-detection"
tokenizer = DistilBertTokenizer.from_pretrained(model_name)
model = DistilBertForSequenceClassification.from_pretrained(model_name)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def predict_spam(text, model, tokenizer, device):
    model.eval()  # Set to evaluation mode
    inputs = tokenizer(text, return_tensors="pt", padding="max_length", truncation=True, max_length=128).to(device)

    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.softmax(outputs.logits, dim=-1)
        pred_class = torch.argmax(probs).item()
    return "Spam" if pred_class == 1 else "Not Spam"

# Sample test messages
test_messages = [
    "Congratulations! You have won a lottery of $1,000,000. Claim now!",  # Spam
    "Hey, are we still meeting for dinner tonight?",  # Not Spam
    "URGENT: Your bank account is at risk! Click this link to secure it now.",  # Spam
    "Let's catch up this weekend. It’s been a while!",  # Not Spam
    "Exclusive offer! Get 50% off on your next purchase. Limited time only!",  # Spam
]

# Run inference on test messages
for i, msg in enumerate(test_messages):
    prediction = predict_spam(msg, model, tokenizer, device)
    print(f"Sample {i+1}: {msg} -> Prediction: {prediction}")

📊 Classification Report (Quantized Model - bfloat16)

🔍 Observations

✅ Precision: High (1.00 for non-spam, 0.98 for spam) → Few false positives
✅ Recall: High (0.99 for both classes) → Few false negatives
✅ F1-Score: Near-perfect balance between precision & recall

Fine-Tuning Details

Dataset

The Hugging Face's sms_spam dataset was used, containing both spam and ham (non-spam) examples.

Training

• Number of epochs: 7
• Batch size: 16
• Evaluation strategy: epoch
• Learning rate: 5e-6

Quantization

Post-training quantization was applied using PyTorch's built-in quantization framework to reduce the model size and improve inference efficiency.

Repository Structure

.
├── model/               # Contains the quantized model files
├── tokenizer_config/    # Tokenizer configuration and vocabulary files
├── pytorch_model.bin/   # Fine Tuned Model
├── README.md            # Model documentation

Limitations

• The model may not generalize well to domains outside the fine-tuning dataset.
• Quantization may result in minor accuracy degradation compared to full-precision models.

Contributing

Contributions are welcome! Feel free to open an issue or submit a pull request if you have suggestions or improvements.