DistilBERT Enron Spam Detector
Descripción
Este proyecto consiste en el desarrollo de un sistema de detección de spam en correos electrónicos mediante técnicas de procesamiento de lenguaje natural (NLP) y modelos Transformer.
Para ello, se ha realizado el fine-tuning del modelo distilbert-base-uncased sobre el dataset SetFit/enron_spam, compuesto por correos electrónicos etiquetados como legítimos (HAM) o spam (SPAM).
El objetivo del modelo es clasificar automáticamente un correo electrónico en una de las dos categorías:
HAM: correo legítimo.SPAM: correo no deseado.
Dataset
El dataset utilizado ha sido:
SetFit/enron_spam
Este dataset contiene miles de correos electrónicos reales procedentes del conjunto Enron Spam Email Dataset.
Cada ejemplo incluye distintas variables relacionadas con el correo:
| Variable | Descripción |
|---|
message_id | Identificador único del correo |
text | Texto completo del correo |
label | Etiqueta numérica |
label_text | Etiqueta textual (ham o spam) |
subject | Asunto del correo |
message | Cuerpo principal del correo |
date | Fecha asociada al correo |
Durante el preprocesado se utilizaron únicamente las variables:
Ambas variables se combinaron en una única entrada textual con la siguiente estructura:
1Subject: <asunto>
2
3Message: <cuerpo del correo>
Esta representación permite al modelo diferenciar explícitamente entre el asunto y el contenido principal del correo.
Modelo utilizado
El modelo base seleccionado ha sido:
DistilBERT es una versión reducida y optimizada de BERT que mantiene gran parte de su rendimiento, pero con:
- menor número de parámetros,
- menor consumo de memoria,
- entrenamiento más rápido.
Este modelo resulta adecuado para tareas de clasificación de texto con recursos computacionales limitados.
Tokenización
Se utilizó el tokenizer oficial asociado a DistilBERT.
Los correos electrónicos fueron tokenizados utilizando:
- truncado automático,
- padding,
- longitud máxima de 256 tokens.
1tokenizer(
2 text,
3 truncation=True,
4 padding="max_length",
5 max_length=256
6)
Fine-Tuning
El entrenamiento se realizó utilizando la clase Trainer de Hugging Face Transformers.
Hiperparámetros principales
| Parámetro | Valor |
|---|
| Epochs | 5 |
| Learning rate | 2e-5 |
| Batch size | 8 |
| Weight decay | 0.01 |
| Max length | 256 |
Métricas de evaluación
Las métricas utilizadas para evaluar el modelo fueron:
- Accuracy
- Precision
- Recall
- F1-score
Resultados obtenidos
| Métrica | Resultado |
|---|
| Accuracy | 0.9915 |
| Precision | 0.9891 |
| Recall | 0.9940 |
| F1-score | 0.9916 |
Inferencia
Ejemplo de uso:
1from transformers import pipeline
2
3classifier = pipeline(
4 "text-classification",
5 model="TU_USUARIO_HF/distilbert-enron-spam-detector"
6)
7
8texto = "
9Subject: You won a free iPhone
10
11Message: Click here immediately to claim your reward.
12"
13
14classifier(texto)
Demo con Gradio
Además del modelo entrenado, se desarrolló una demo interactiva mediante Gradio y Hugging Face Spaces.
La aplicación permite introducir el contenido de un correo electrónico y obtener en tiempo real:
- clasificación (
HAM o SPAM),
- nivel de confianza de la predicción.
Limitaciones
Este modelo presenta varias limitaciones:
- Está entrenado principalmente con correos en inglés.
- Puede funcionar peor con correos escritos en otros idiomas.
- Solo analiza texto, no adjuntos ni enlaces.
- Los correos largos se truncarán a 256 tokens.
- Puede fallar ante técnicas modernas de spam no presentes en el dataset original.
Tecnologías utilizadas
- Python
- PyTorch
- Hugging Face Transformers
- Hugging Face Datasets
- Gradio
- Scikit-learn
Autor
Proyecto desarrollado como práctica de fine-tuning y despliegue de modelos Transformer para tareas NLP.
