Update train.py

train.py

@@ -1,24 +1,34 @@
-from transformers import Trainer, TrainingArguments, BertForTokenClassification
+from transformers import BertTokenizer, BertForTokenClassification, Trainer, TrainingArguments
 from datasets import Dataset
 import json
+import torch
+
+# Carica il tokenizer e il modello pre-addestrato (dbmdz/bert-base-italian-uncased)
+tokenizer = BertTokenizer.from_pretrained("dbmdz/bert-base-italian-uncased")
+model = BertForTokenClassification.from_pretrained("dbmdz/bert-base-italian-uncased", num_labels=5)  # Aggiungi il numero corretto di etichette (labels)
 
 # Carica il dataset
 with open('entity_dataset.json', 'r') as f:
     dataset = json.load(f)
 
-# Prepara il dataset per l'addestramento
+# Funzione per preparare i dati
 def prepare_dataset(dataset):
-    # Converti il dataset in un formato adatto per Hugging Face Dataset
-    data = {
-        "text": [entry["query"] for entry in dataset],
-        "labels": [entry["entities"] for entry in dataset]
-    }
-    return Dataset.from_dict(data)
+    input_texts = [entry["query"] for entry in dataset]
+    labels = [entry["entities"] for entry in dataset]
 
-train_dataset = prepare_dataset(dataset)
+    # Tokenizza i dati di input
+    encodings = tokenizer(input_texts, truncation=True, padding=True, max_length=512)
+    # Aggiungi le etichette (entità) come output
+    # Qui supponiamo che tu stia etichettando solo le entità (puoi adattare la funzione per il tuo caso)
+    # Nota: dovresti mappare le etichette in modo che corrispondano al formato richiesto per BERT
+
+    # Associa le etichette agli input tokenizzati
+    encodings['labels'] = torch.tensor(labels)
 
-# Carica il modello pre-addestrato
-model = BertForTokenClassification.from_pretrained("dbmdz/bert-base-italian-uncased")
+    return Dataset.from_dict(encodings)
+
+# Prepara il dataset per l'addestramento
+train_dataset = prepare_dataset(dataset)
 
 # Imposta i parametri di addestramento
 training_args = TrainingArguments(
@@ -27,6 +37,7 @@ training_args = TrainingArguments(
     learning_rate=2e-5,              # Tasso di apprendimento
     per_device_train_batch_size=16,  # Dimensione del batch
     num_train_epochs=3,              # Numero di epoche
+    weight_decay=0.01                # Peso di decadimento (per evitare overfitting)
 )
 
 # Inizializza il trainer
@@ -38,3 +49,7 @@ trainer = Trainer(
 
 # Avvia l'addestramento
 trainer.train()
+
+# Salva il modello addestrato
+model.save_pretrained("./hotel_model")
+tokenizer.save_pretrained("./hotel_model")