Update app.py

app.py

@@ -1,17 +1,49 @@
-from datasets import load_dataset
-from transformers import pipeline
 import gradio as gr
+from transformers import pipeline, Trainer, TrainingArguments, DistilBertForSequenceClassification, DistilBertTokenizer
+from datasets import load_dataset
+
+# Carica il dataset di spam detection da Hugging Face
+dataset = load_dataset("tanquangduong/spam-detection-dataset-splits")
+
+# Carica il tokenizer e il modello pre-addestrato
+tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")
+model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased")
+
+# Tokenizza il dataset
+def tokenize_function(examples):
+    return tokenizer(examples['message'], truncation=True, padding="max_length", max_length=128)
+
+tokenized_datasets = dataset.map(tokenize_function, batched=True)
+
+# Suddividi in training e test set
+train_dataset = tokenized_datasets["train"]
+test_dataset = tokenized_datasets["test"]
 
-# Carica il dataset di esempio da Hugging Face (utilizziamo il dataset TREC) con `trust_remote_code=True`
-dataset = load_dataset("trec", trust_remote_code=True)
+# Definisci gli 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,
+)
 
-# Visualizza i primi esempi per assicurarti che il dataset sia stato caricato correttamente
-print(dataset['train'][0])
+# Crea l'oggetto Trainer
+trainer = Trainer(
+    model=model,
+    args=training_args,
+    train_dataset=train_dataset,
+    eval_dataset=test_dataset,
+)
 
-# Inizializza il modello DistilBERT per la classificazione di testo
-classifier = pipeline("text-classification", model="distilbert-base-uncased")
+# Avvia il training
+trainer.train()
 
+# Definisci la funzione di classificazione usando il modello addestrato
 def classify_email(text):
+    classifier = pipeline("text-classification", model=model, tokenizer=tokenizer, framework="pt")
     result = classifier(text)
     label = result[0]['label']
     score = result[0]['score']
@@ -24,4 +56,5 @@ iface = gr.Interface(fn=classify_email,
                      title="ZeroSpam Email Classifier",
                      description="Inserisci l'email da analizzare per determinare se è spam o phishing.")
 
+# Avvia l'interfaccia
 iface.launch(share=True)