from transformers import LayoutLMForTokenClassification, Trainer, TrainingArguments
from datasets import load_dataset

# 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)

# 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",
    per_device_train_batch_size=4,
    per_device_eval_batch_size=4,
    num_train_epochs=5,
    evaluation_strategy="epoch",
    save_strategy="epoch",
    logging_dir="./logs",
    logging_steps=10
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=dataset["train"],
    eval_dataset=dataset["test"]
)

trainer.train()
model.save_pretrained("./layoutlmv3_finetuned")
model.push_to_hub("kryman27/layoutlmv3-finetuned")