Update app.py

app.py

@@ -1,22 +1,54 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import pipeline, Trainer, TrainingArguments, DistilBertForSequenceClassification, DistilBertTokenizer
+from datasets import load_dataset
+import torch
 
-# Inizializza il modello DistilBERT per la classificazione di testo
-classifier = pipeline("text-classification", model="distilbert-base-uncased")
+# Carica il dataset Enron Spam da Hugging Face
+dataset = load_dataset("enron_spam")
 
+# Carica il tokenizer e il modello
+tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")
+model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased")
+
+# Tokenizzazione del dataset
+def tokenize_function(examples):
+    return tokenizer(examples['text'], truncation=True, padding="max_length", max_length=128)
+
+tokenized_datasets = dataset.map(tokenize_function, batched=True)
+
+# Suddivisione in training e test set
+train_dataset = tokenized_datasets["train"].shuffle(seed=42).select(range(8000))
+test_dataset = tokenized_datasets["test"].shuffle(seed=42).select(range(2000))
+
+# Definizione degli argomenti per l'addestramento
+training_args = TrainingArguments(
+    output_dir="./results",
+    evaluation_strategy="epoch",
+    learning_rate=2e-5,
+    per_device_train_batch_size=16,
+    per_device_eval_batch_size=16,
+    num_train_epochs=3,
+    weight_decay=0.01,
+)
+
+# Creazione dell'oggetto Trainer
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    eval_dataset=test_dataset,
+)
+
+# Avvio dell'addestramento
+trainer.train()
+
+# Definizione della funzione di classificazione usando Gradio
 def classify_email(text):
+    classifier = pipeline("text-classification", model=model, tokenizer=tokenizer, framework="pt")
     result = classifier(text)
-    # 'result' è una lista con un solo dizionario, prendiamo il primo elemento e restituiamo solo la label
     label = result[0]['label']
     score = result[0]['score']
-    
-    # Modifica le etichette per essere più comprensibili
-    if label == "LABEL_1":
-        label_text = "Phishing"
-    else:
-        label_text = "Non Phishing"
-    
-    return {label_text: score}
+    return {label: score}
 
 # Interfaccia con Gradio
 iface = gr.Interface(fn=classify_email,
@@ -25,4 +57,5 @@ iface = gr.Interface(fn=classify_email,
                      title="ZeroSpam Email Classifier",
                      description="Inserisci l'email da analizzare per determinare se è spam o phishing.")
 
+# Avvio dell'interfaccia
 iface.launch(share=True)