Spaces:
Paused
Paused
#!/usr/bin/env python | |
# coding=utf-8 | |
# Copyright 2022 The HuggingFace Team All rights reserved. | |
# | |
# Licensed under the Apache License, Version 2.0 (the "License"); | |
# you may not use this file except in compliance with the License. | |
# You may obtain a copy of the License at | |
# | |
# http://www.apache.org/licenses/LICENSE-2.0 | |
# | |
# Unless required by applicable law or agreed to in writing, software | |
# distributed under the License is distributed on an "AS IS" BASIS, | |
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | |
# See the License for the specific language governing permissions and | |
# limitations under the License. | |
""" | |
Fine-tuning LayoutLMv3 for token classification on FUNSD or CORD. | |
""" | |
# You can also adapt this script on your own token classification task and datasets. Pointers for this are left as | |
# comments. | |
import logging | |
import os | |
import sys | |
from dataclasses import dataclass, field | |
from typing import Optional | |
import datasets | |
import numpy as np | |
from datasets import ClassLabel, load_dataset, load_metric | |
import transformers | |
from transformers import ( | |
AutoConfig, | |
AutoModelForTokenClassification, | |
AutoProcessor, | |
HfArgumentParser, | |
Trainer, | |
TrainingArguments, | |
set_seed, | |
) | |
from transformers.data.data_collator import default_data_collator | |
from transformers.trainer_utils import get_last_checkpoint | |
from transformers.utils import check_min_version | |
from transformers.utils.versions import require_version | |
# Will error if the minimal version of Transformers is not installed. Remove at your own risks. | |
check_min_version("4.19.0.dev0") | |
require_version("datasets>=1.8.0", "To fix: pip install -r examples/pytorch/token-classification/requirements.txt") | |
logger = logging.getLogger(__name__) | |
class ModelArguments: | |
""" | |
Arguments pertaining to which model/config/tokenizer we are going to fine-tune from. | |
""" | |
model_name_or_path: str = field( | |
default="microsoft/layoutlmv3-base", | |
metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}, | |
) | |
config_name: Optional[str] = field( | |
default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"} | |
) | |
processor_name: Optional[str] = field( | |
default=None, metadata={"help": "Name or path to the processor files if not the same as model_name"} | |
) | |
cache_dir: Optional[str] = field( | |
default=None, | |
metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"}, | |
) | |
model_revision: str = field( | |
default="main", | |
metadata={"help": "The specific model version to use (can be a branch name, tag name or commit id)."}, | |
) | |
use_auth_token: bool = field( | |
default=False, | |
metadata={ | |
"help": ( | |
"Will use the token generated when running `huggingface-cli login` (necessary to use this script " | |
"with private models)." | |
) | |
}, | |
) | |
class DataTrainingArguments: | |
""" | |
Arguments pertaining to what data we are going to input our model for training and eval. | |
""" | |
task_name: Optional[str] = field(default="ner", metadata={"help": "The name of the task (ner, pos...)."}) | |
dataset_name: Optional[str] = field( | |
default="nielsr/funsd-layoutlmv3", | |
metadata={"help": "The name of the dataset to use (via the datasets library)."}, | |
) | |
dataset_config_name: Optional[str] = field( | |
default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."} | |
) | |
train_file: Optional[str] = field( | |
default=None, metadata={"help": "The input training data file (a csv or JSON file)."} | |
) | |
validation_file: Optional[str] = field( | |
default=None, | |
metadata={"help": "An optional input evaluation data file to evaluate on (a csv or JSON file)."}, | |
) | |
test_file: Optional[str] = field( | |
default=None, | |
metadata={"help": "An optional input test data file to predict on (a csv or JSON file)."}, | |
) | |
text_column_name: Optional[str] = field( | |
default=None, metadata={"help": "The column name of text to input in the file (a csv or JSON file)."} | |
) | |
label_column_name: Optional[str] = field( | |
default=None, metadata={"help": "The column name of label to input in the file (a csv or JSON file)."} | |
) | |
overwrite_cache: bool = field( | |
default=False, metadata={"help": "Overwrite the cached training and evaluation sets"} | |
) | |
preprocessing_num_workers: Optional[int] = field( | |
default=None, | |
metadata={"help": "The number of processes to use for the preprocessing."}, | |
) | |
max_seq_length: int = field( | |
default=512, | |
metadata={ | |
"help": ( | |
"The maximum total input sequence length after tokenization. If set, sequences longer " | |
"than this will be truncated, sequences shorter will be padded." | |
) | |
}, | |
) | |
max_train_samples: Optional[int] = field( | |
default=None, | |
metadata={ | |
"help": ( | |
"For debugging purposes or quicker training, truncate the number of training examples to this " | |
"value if set." | |
) | |
}, | |
) | |
max_eval_samples: Optional[int] = field( | |
default=None, | |
metadata={ | |
"help": ( | |
"For debugging purposes or quicker training, truncate the number of evaluation examples to this " | |
"value if set." | |
) | |
}, | |
) | |
max_predict_samples: Optional[int] = field( | |
default=None, | |
metadata={ | |
"help": ( | |
"For debugging purposes or quicker training, truncate the number of prediction examples to this " | |
"value if set." | |
) | |
}, | |
) | |
label_all_tokens: bool = field( | |
default=False, | |
metadata={ | |
"help": ( | |
"Whether to put the label for one word on all tokens of generated by that word or just on the " | |
"one (in which case the other tokens will have a padding index)." | |
) | |
}, | |
) | |
return_entity_level_metrics: bool = field( | |
default=False, | |
metadata={"help": "Whether to return all the entity levels during evaluation or just the overall ones."}, | |
) | |
def __post_init__(self): | |
if self.dataset_name is None and self.train_file is None and self.validation_file is None: | |
raise ValueError("Need either a dataset name or a training/validation file.") | |
else: | |
if self.train_file is not None: | |
extension = self.train_file.split(".")[-1] | |
assert extension in ["csv", "json"], "`train_file` should be a csv or a json file." | |
if self.validation_file is not None: | |
extension = self.validation_file.split(".")[-1] | |
assert extension in ["csv", "json"], "`validation_file` should be a csv or a json file." | |
self.task_name = self.task_name.lower() | |
def main(): | |
# See all possible arguments in src/transformers/training_args.py | |
# or by passing the --help flag to this script. | |
# We now keep distinct sets of args, for a cleaner separation of concerns. | |
parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments)) | |
if len(sys.argv) == 2 and sys.argv[1].endswith(".json"): | |
# If we pass only one argument to the script and it's the path to a json file, | |
# let's parse it to get our arguments. | |
model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1])) | |
else: | |
model_args, data_args, training_args = parser.parse_args_into_dataclasses() | |
# Setup logging | |
logging.basicConfig( | |
format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", | |
datefmt="%m/%d/%Y %H:%M:%S", | |
handlers=[logging.StreamHandler(sys.stdout)], | |
) | |
log_level = training_args.get_process_log_level() | |
logger.setLevel(log_level) | |
datasets.utils.logging.set_verbosity(log_level) | |
transformers.utils.logging.set_verbosity(log_level) | |
transformers.utils.logging.enable_default_handler() | |
transformers.utils.logging.enable_explicit_format() | |
# Log on each process the small summary: | |
logger.warning( | |
f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}" | |
+ f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}" | |
) | |
logger.info(f"Training/evaluation parameters {training_args}") | |
# Detecting last checkpoint. | |
last_checkpoint = None | |
if os.path.isdir(training_args.output_dir) and training_args.do_train and not training_args.overwrite_output_dir: | |
last_checkpoint = get_last_checkpoint(training_args.output_dir) | |
if last_checkpoint is None and len(os.listdir(training_args.output_dir)) > 0: | |
raise ValueError( | |
f"Output directory ({training_args.output_dir}) already exists and is not empty. " | |
"Use --overwrite_output_dir to overcome." | |
) | |
elif last_checkpoint is not None and training_args.resume_from_checkpoint is None: | |
logger.info( | |
f"Checkpoint detected, resuming training at {last_checkpoint}. To avoid this behavior, change " | |
"the `--output_dir` or add `--overwrite_output_dir` to train from scratch." | |
) | |
# Set seed before initializing model. | |
set_seed(training_args.seed) | |
# Get the datasets | |
# In distributed training, the load_dataset function guarantee that only one local process can concurrently | |
# download the dataset. | |
if data_args.dataset_name == "funsd": | |
# Downloading and loading a dataset from the hub. | |
dataset = load_dataset( | |
"nielsr/funsd-layoutlmv3", | |
data_args.dataset_config_name, | |
cache_dir=model_args.cache_dir, | |
use_auth_token=True if model_args.use_auth_token else None, | |
) | |
elif data_args.dataset_name == "cord": | |
# Downloading and loading a dataset from the hub. | |
dataset = load_dataset( | |
"nielsr/cord-layoutlmv3", | |
data_args.dataset_config_name, | |
cache_dir=model_args.cache_dir, | |
use_auth_token=True if model_args.use_auth_token else None, | |
) | |
else: | |
raise ValueError("This script only supports either FUNSD or CORD out-of-the-box.") | |
if training_args.do_train: | |
column_names = dataset["train"].column_names | |
features = dataset["train"].features | |
else: | |
column_names = dataset["test"].column_names | |
features = dataset["test"].features | |
image_column_name = "image" | |
text_column_name = "words" if "words" in column_names else "tokens" | |
boxes_column_name = "bboxes" | |
label_column_name = ( | |
f"{data_args.task_name}_tags" if f"{data_args.task_name}_tags" in column_names else column_names[1] | |
) | |
remove_columns = column_names | |
# In the event the labels are not a `Sequence[ClassLabel]`, we will need to go through the dataset to get the | |
# unique labels. | |
def get_label_list(labels): | |
unique_labels = set() | |
for label in labels: | |
unique_labels = unique_labels | set(label) | |
label_list = list(unique_labels) | |
label_list.sort() | |
return label_list | |
# If the labels are of type ClassLabel, they are already integers and we have the map stored somewhere. | |
# Otherwise, we have to get the list of labels manually. | |
if isinstance(features[label_column_name].feature, ClassLabel): | |
label_list = features[label_column_name].feature.names | |
# No need to convert the labels since they are already ints. | |
id2label = dict(enumerate(label_list)) | |
label2id = {v: k for k, v in enumerate(label_list)} | |
else: | |
label_list = get_label_list(datasets["train"][label_column_name]) | |
id2label = dict(enumerate(label_list)) | |
label2id = {v: k for k, v in enumerate(label_list)} | |
num_labels = len(label_list) | |
# Load pretrained model and processor | |
# | |
# Distributed training: | |
# The .from_pretrained methods guarantee that only one local process can concurrently | |
# download model & vocab. | |
config = AutoConfig.from_pretrained( | |
model_args.config_name if model_args.config_name else model_args.model_name_or_path, | |
num_labels=num_labels, | |
finetuning_task=data_args.task_name, | |
cache_dir=model_args.cache_dir, | |
revision=model_args.model_revision, | |
use_auth_token=True if model_args.use_auth_token else None, | |
) | |
processor = AutoProcessor.from_pretrained( | |
model_args.processor_name if model_args.processor_name else model_args.model_name_or_path, | |
cache_dir=model_args.cache_dir, | |
use_fast=True, | |
revision=model_args.model_revision, | |
use_auth_token=True if model_args.use_auth_token else None, | |
add_prefix_space=True, | |
apply_ocr=False, | |
) | |
model = AutoModelForTokenClassification.from_pretrained( | |
model_args.model_name_or_path, | |
from_tf=bool(".ckpt" in model_args.model_name_or_path), | |
config=config, | |
cache_dir=model_args.cache_dir, | |
revision=model_args.model_revision, | |
use_auth_token=True if model_args.use_auth_token else None, | |
) | |
# Set the correspondences label/ID inside the model config | |
model.config.label2id = label2id | |
model.config.id2label = id2label | |
# Preprocessing the dataset | |
# The processor does everything for us (prepare the image using LayoutLMv3ImageProcessor | |
# and prepare the words, boxes and word-level labels using LayoutLMv3TokenizerFast) | |
def prepare_examples(examples): | |
images = examples[image_column_name] | |
words = examples[text_column_name] | |
boxes = examples[boxes_column_name] | |
word_labels = examples[label_column_name] | |
encoding = processor( | |
images, | |
words, | |
boxes=boxes, | |
word_labels=word_labels, | |
truncation=True, | |
padding="max_length", | |
max_length=data_args.max_seq_length, | |
) | |
return encoding | |
if training_args.do_train: | |
if "train" not in dataset: | |
raise ValueError("--do_train requires a train dataset") | |
train_dataset = dataset["train"] | |
if data_args.max_train_samples is not None: | |
train_dataset = train_dataset.select(range(data_args.max_train_samples)) | |
with training_args.main_process_first(desc="train dataset map pre-processing"): | |
train_dataset = train_dataset.map( | |
prepare_examples, | |
batched=True, | |
remove_columns=remove_columns, | |
num_proc=data_args.preprocessing_num_workers, | |
load_from_cache_file=not data_args.overwrite_cache, | |
) | |
if training_args.do_eval: | |
validation_name = "test" | |
if validation_name not in dataset: | |
raise ValueError("--do_eval requires a validation dataset") | |
eval_dataset = dataset[validation_name] | |
if data_args.max_eval_samples is not None: | |
eval_dataset = eval_dataset.select(range(data_args.max_eval_samples)) | |
with training_args.main_process_first(desc="validation dataset map pre-processing"): | |
eval_dataset = eval_dataset.map( | |
prepare_examples, | |
batched=True, | |
remove_columns=remove_columns, | |
num_proc=data_args.preprocessing_num_workers, | |
load_from_cache_file=not data_args.overwrite_cache, | |
) | |
if training_args.do_predict: | |
if "test" not in datasets: | |
raise ValueError("--do_predict requires a test dataset") | |
predict_dataset = datasets["test"] | |
if data_args.max_predict_samples is not None: | |
max_predict_samples = min(len(predict_dataset), data_args.max_predict_samples) | |
predict_dataset = predict_dataset.select(range(max_predict_samples)) | |
with training_args.main_process_first(desc="prediction dataset map pre-processing"): | |
predict_dataset = predict_dataset.map( | |
prepare_examples, | |
batched=True, | |
remove_columns=remove_columns, | |
num_proc=data_args.preprocessing_num_workers, | |
load_from_cache_file=not data_args.overwrite_cache, | |
) | |
# Metrics | |
metric = load_metric("seqeval") | |
def compute_metrics(p): | |
predictions, labels = p | |
predictions = np.argmax(predictions, axis=2) | |
# Remove ignored index (special tokens) | |
true_predictions = [ | |
[label_list[p] for (p, l) in zip(prediction, label) if l != -100] | |
for prediction, label in zip(predictions, labels) | |
] | |
true_labels = [ | |
[label_list[l] for (p, l) in zip(prediction, label) if l != -100] | |
for prediction, label in zip(predictions, labels) | |
] | |
results = metric.compute(predictions=true_predictions, references=true_labels) | |
if data_args.return_entity_level_metrics: | |
# Unpack nested dictionaries | |
final_results = {} | |
for key, value in results.items(): | |
if isinstance(value, dict): | |
for n, v in value.items(): | |
final_results[f"{key}_{n}"] = v | |
else: | |
final_results[key] = value | |
return final_results | |
else: | |
return { | |
"precision": results["overall_precision"], | |
"recall": results["overall_recall"], | |
"f1": results["overall_f1"], | |
"accuracy": results["overall_accuracy"], | |
} | |
# Initialize our Trainer | |
trainer = Trainer( | |
model=model, | |
args=training_args, | |
train_dataset=train_dataset if training_args.do_train else None, | |
eval_dataset=eval_dataset if training_args.do_eval else None, | |
tokenizer=processor, | |
data_collator=default_data_collator, | |
compute_metrics=compute_metrics, | |
) | |
# Training | |
if training_args.do_train: | |
checkpoint = None | |
if training_args.resume_from_checkpoint is not None: | |
checkpoint = training_args.resume_from_checkpoint | |
elif last_checkpoint is not None: | |
checkpoint = last_checkpoint | |
train_result = trainer.train(resume_from_checkpoint=checkpoint) | |
metrics = train_result.metrics | |
trainer.save_model() # Saves the tokenizer too for easy upload | |
max_train_samples = ( | |
data_args.max_train_samples if data_args.max_train_samples is not None else len(train_dataset) | |
) | |
metrics["train_samples"] = min(max_train_samples, len(train_dataset)) | |
trainer.log_metrics("train", metrics) | |
trainer.save_metrics("train", metrics) | |
trainer.save_state() | |
# Evaluation | |
if training_args.do_eval: | |
logger.info("*** Evaluate ***") | |
metrics = trainer.evaluate() | |
max_eval_samples = data_args.max_eval_samples if data_args.max_eval_samples is not None else len(eval_dataset) | |
metrics["eval_samples"] = min(max_eval_samples, len(eval_dataset)) | |
trainer.log_metrics("eval", metrics) | |
trainer.save_metrics("eval", metrics) | |
# Predict | |
if training_args.do_predict: | |
logger.info("*** Predict ***") | |
predictions, labels, metrics = trainer.predict(predict_dataset, metric_key_prefix="predict") | |
predictions = np.argmax(predictions, axis=2) | |
# Remove ignored index (special tokens) | |
true_predictions = [ | |
[label_list[p] for (p, l) in zip(prediction, label) if l != -100] | |
for prediction, label in zip(predictions, labels) | |
] | |
trainer.log_metrics("predict", metrics) | |
trainer.save_metrics("predict", metrics) | |
# Save predictions | |
output_predictions_file = os.path.join(training_args.output_dir, "predictions.txt") | |
if trainer.is_world_process_zero(): | |
with open(output_predictions_file, "w") as writer: | |
for prediction in true_predictions: | |
writer.write(" ".join(prediction) + "\n") | |
kwargs = {"finetuned_from": model_args.model_name_or_path, "tasks": "token-classification"} | |
if data_args.dataset_name is not None: | |
kwargs["dataset_tags"] = data_args.dataset_name | |
if data_args.dataset_config_name is not None: | |
kwargs["dataset_args"] = data_args.dataset_config_name | |
kwargs["dataset"] = f"{data_args.dataset_name} {data_args.dataset_config_name}" | |
else: | |
kwargs["dataset"] = data_args.dataset_name | |
if training_args.push_to_hub: | |
trainer.push_to_hub(**kwargs) | |
else: | |
trainer.create_model_card(**kwargs) | |
def _mp_fn(index): | |
# For xla_spawn (TPUs) | |
main() | |
if __name__ == "__main__": | |
main() | |