Update app.py

app.py

@@ -1,17 +1,21 @@
 import gradio as gr
 from transformers import pipeline, Trainer, TrainingArguments, DistilBertForSequenceClassification, DistilBertTokenizer
 from datasets import load_dataset
+import torch
 
 # Carica il dataset di spam detection da Hugging Face
 dataset = load_dataset("tanquangduong/spam-detection-dataset-splits")
 
+# Visualizza i nomi delle colonne per verificare quale contiene il testo delle email
+print(dataset['train'].column_names)
+
 # 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
+# Aggiorna il nome della colonna con il nome corretto
 def tokenize_function(examples):
-    return tokenizer(examples['message'], truncation=True, padding="max_length", max_length=128)
+    return tokenizer(examples['text'], truncation=True, padding="max_length", max_length=128)
 
 tokenized_datasets = dataset.map(tokenize_function, batched=True)