pdf-extractor / train_model.py
kryman27's picture
Create train_model.py
5c7e49b verified
raw
history blame
815 Bytes
from transformers import LayoutLMForTokenClassification, Trainer, TrainingArguments
from datasets import load_dataset
# Wczytanie przygotowanego zbioru danych
dataset = load_dataset("json", data_files="training_data.json")
# Ładowanie modelu LayoutLM do dostrajania
model = LayoutLMForTokenClassification.from_pretrained("microsoft/layoutlmv3-base", num_labels=5)
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"
)
trainer = Trainer(
model=model,
args=training_args,
train_dataset=dataset["train"],
eval_dataset=dataset["test"]
)
trainer.train()
# Zapisanie modelu
model.save_pretrained("./layoutlmv3_finetuned")