Update app.py

app.py

@@ -2,39 +2,40 @@ import gradio as gr
 from transformers import pipeline, Trainer, TrainingArguments, DistilBertForSequenceClassification, DistilBertTokenizer
 from datasets import load_dataset
 import torch
+import os
 
-# Carica il dataset di spam detection da Hugging Face
+# Carica il dataset 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
+# Carica il tokenizer e il modello
 tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")
 model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased")
 
-# Aggiorna il nome della colonna con il nome corretto
+# 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)
 
-# Suddividi in training e test set
-train_dataset = tokenized_datasets["train"]
-test_dataset = tokenized_datasets["test"]
+# Suddivisione in training e test set
+train_dataset = tokenized_datasets["train"].shuffle(seed=42).select(range(2000))  # Ridotto per velocizzare l'addestramento
+test_dataset = tokenized_datasets["test"].shuffle(seed=42).select(range(500))
 
-# Definisci gli argomenti per l'addestramento
+# Definizione degli argomenti per l'addestramento, inclusa la frequenza di salvataggio dei checkpoint
 training_args = TrainingArguments(
     output_dir="./results",
     evaluation_strategy="epoch",
+    save_strategy="epoch",  # Salva un checkpoint alla fine di ogni epoch
     learning_rate=2e-5,
     per_device_train_batch_size=16,
     per_device_eval_batch_size=16,
-    num_train_epochs=3,
+    num_train_epochs=1,  # Ridotto a 1 epoch per evitare timeout
     weight_decay=0.01,
+    save_total_limit=2,  # Limita il numero di checkpoint salvati per risparmiare spazio
+    load_best_model_at_end=True,
 )
 
-# Crea l'oggetto Trainer
+# Creazione dell'oggetto Trainer
 trainer = Trainer(
     model=model,
     args=training_args,
@@ -42,10 +43,15 @@ trainer = Trainer(
     eval_dataset=test_dataset,
 )
 
-# Avvia il training
-trainer.train()
+# Avvio dell'addestramento
+if os.path.exists("./results/checkpoint-1"):  # Verifica se esiste un checkpoint salvato
+    print("Riprendi l'addestramento dal checkpoint...")
+    trainer.train(resume_from_checkpoint="./results/checkpoint-1")
+else:
+    print("Inizia l'addestramento da zero...")
+    trainer.train()
 
-# Definisci la funzione di classificazione usando il modello addestrato
+# Definizione della funzione di classificazione usando Gradio
 def classify_email(text):
     classifier = pipeline("text-classification", model=model, tokenizer=tokenizer, framework="pt")
     result = classifier(text)
@@ -60,5 +66,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
+# Avvio dell'interfaccia
 iface.launch(share=True)