Update train_model.py

train_model.py

@@ -2,10 +2,11 @@ from transformers import LayoutLMForTokenClassification, Trainer, TrainingArgume
 from datasets import load_dataset
 
 # Wczytanie przygotowanego zbioru danych
-dataset = load_dataset("json", data_files="training_data.json")
+dataset = load_dataset("json", data_files="training_data.json")["train"]
+dataset = dataset.train_test_split(test_size=0.2)  # Podział na trening i test
 
 # Ładowanie modelu LayoutLM do dostrajania
-model = LayoutLMForTokenClassification.from_pretrained("microsoft/layoutlmv3-base", num_labels=5)
+model = LayoutLMForTokenClassification.from_pretrained("microsoft/layoutlmv3-base", num_labels=10)
 
 training_args = TrainingArguments(
     output_dir="./layoutlmv3_finetuned",
@@ -13,7 +14,9 @@ training_args = TrainingArguments(
     per_device_eval_batch_size=4,
     num_train_epochs=5,
     evaluation_strategy="epoch",
-    save_strategy="epoch"
+    save_strategy="epoch",
+    logging_dir="./logs",
+    logging_steps=10
 )
 
 trainer = Trainer(
@@ -25,5 +28,8 @@ trainer = Trainer(
 
 trainer.train()
 
-# Zapisanie modelu
+# Zapisanie modelu lokalnie
 model.save_pretrained("./layoutlmv3_finetuned")
+
+# Wysłanie modelu do Hugging Face (tylko jeśli masz konto)
+model.push_to_hub("twoj_username/layoutlmv3-finetuned")