57 lines
1.4 KiB
Markdown
57 lines
1.4 KiB
Markdown
|
---
|
||
|
|
language: en
|
||
|
|
license: mit
|
||
|
|
tags:
|
||
|
|
- text-generation
|
||
|
|
- spam-detection
|
||
|
|
- distilgpt2
|
||
|
|
datasets:
|
||
|
|
- sms_spam
|
||
|
|
---
|
||
|
|
|
||
|
|
# DistilGPT2 Spam Detector
|
||
|
|
|
||
|
|
This model fine-tunes [distilgpt2](https://huggingface.co/distilgpt2) on the
|
||
|
|
[SMS Spam Collection](https://huggingface.co/datasets/sms_spam) dataset to classify
|
||
|
|
text messages as **spam** or **ham**, framed as a text-generation task.
|
||
|
|
|
||
|
|
## How it works
|
||
|
|
|
||
|
|
The model is trained on examples formatted as:
|
||
|
|
|
||
|
|
```
|
||
|
|
Message: <text>
|
||
|
|
Label: <ham or spam>
|
||
|
|
```
|
||
|
|
|
||
|
|
At inference time, feed it `"Message: <your text>
|
||
|
|
Label:"` and let it generate a
|
||
|
|
few tokens — it will produce "spam" or "ham".
|
||
|
|
|
||
|
|
## Example usage
|
||
|
|
|
||
|
|
```python
|
||
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM
|
||
|
|
|
||
|
|
tokenizer = AutoTokenizer.from_pretrained("corruptedzayn/distilgpt2-spam-detector")
|
||
|
|
model = AutoModelForCausalLM.from_pretrained("corruptedzayn/distilgpt2-spam-detector")
|
||
|
|
|
||
|
|
prompt = "Message: Win a free iPhone now, click here!
|
||
|
|
Label:"
|
||
|
|
inputs = tokenizer(prompt, return_tensors="pt")
|
||
|
|
output = model.generate(**inputs, max_new_tokens=5, pad_token_id=tokenizer.eos_token_id)
|
||
|
|
print(tokenizer.decode(output[0], skip_special_tokens=True))
|
||
|
|
```
|
||
|
|
|
||
|
|
## Training details
|
||
|
|
|
||
|
|
- Base model: distilgpt2
|
||
|
|
- Dataset: sms_spam (~5,500 messages, 80/20 train/test split)
|
||
|
|
- Epochs: 3
|
||
|
|
- Framed as causal language modeling, not a classification head
|
||
|
|
|
||
|
|
## Limitations
|
||
|
|
|
||
|
|
This is a small demo model trained on a small, English-only, slightly dated dataset.
|
||
|
|
It is not intended for production spam filtering.
|