Update train_model.py

train_model.py

@@ -1,12 +1,14 @@
 from transformers import LayoutLMForTokenClassification, Trainer, TrainingArguments
 from datasets import load_dataset
 
-# Wczytanie przygotowanego zbioru danych
+# Upewnij się, że training_data.json zawiera etykiety odpowiadające nowym polom
 dataset = load_dataset("json", data_files="training_data.json")["train"]
-dataset = dataset.train_test_split(test_size=0.2)  # Podział na trening i test
+dataset = dataset.train_test_split(test_size=0.2)
 
-# Ładowanie modelu LayoutLM do dostrajania
-model = LayoutLMForTokenClassification.from_pretrained("microsoft/layoutlmv3-base", num_labels=10)
+# Dostosuj liczbę etykiet do rozszerzonego zakresu ekstrakcji (przykładowo 15)
+num_labels = 15
+
+model = LayoutLMForTokenClassification.from_pretrained("microsoft/layoutlmv3-base", num_labels=num_labels)
 
 training_args = TrainingArguments(
     output_dir="./layoutlmv3_finetuned",
@@ -27,9 +29,5 @@ trainer = Trainer(
 )
 
 trainer.train()
-
-# Zapisanie modelu lokalnie
 model.save_pretrained("./layoutlmv3_finetuned")
-
-# Wysłanie modelu do Hugging Face (tylko jeśli masz konto)
 model.push_to_hub("kryman27/layoutlmv3-finetuned")