new training script with config, files split

This commit is contained in:
s444415 2023-01-04 11:10:24 +01:00
parent 9a53f14742
commit 0b9333da2c
10 changed files with 762 additions and 344 deletions

1
.gitignore vendored
View File

@ -2,3 +2,4 @@ env
Donut
nohup.out
wandb
__pycache__/

22
config-train.yaml Normal file
View File

@ -0,0 +1,22 @@
dataset_path: "Zombely/fiszki-ocr-train"
pretrained_model_path: "Zombely/plwiki-proto-fine-tuned-v2"
start_model_path: "Zombely/plwiki-proto-fine-tuned-v2"
output_model_path: "Zombely/plwiki-proto-fine-tuned-v3"
wandb_test_name: "fiszki-ocr-fine-tune"
checkpoint_path: "./checkpoint"
max_length: 768
image_size: [1920, 2560]
train_config:
max_epochs: 1
val_check_interval: 0.5
check_val_every_n_epoch: 1
gradient_clip_val: 1.0
num_training_samples_per_epoch: 800
lr: 3.0e-5
train_batch_sizes: [8]
val_batch_sizes: [1]
seed: 2022
num_nodes: 1
warmup_steps: 300
result_path: "./result"
verbose: True

View File

@ -21,7 +21,7 @@ def main(config):
config_vision.decoder.max_length = config.max_dec_length
processor = DonutProcessor.from_pretrained(config.pretrained_processor_path)
model = VisionEncoderDecoderModel.from_pretrained(config.pretrained_model_path)
model = VisionEncoderDecoderModel.from_pretrained(config.pretrained_model_path, config=config_vision if config.use_enc_dec_config else None)
processor.image_processor.size = config.image_size[::-1] # should be (width, height)
processor.image_processor.do_align_long_axis = False

View File

@ -1,379 +1,113 @@
#!/usr/bin/env python
# coding: utf-8
from transformers import VisionEncoderDecoderConfig, DonutProcessor, VisionEncoderDecoderModel
from datasets import load_dataset
import json
import random
from typing import Any, List, Tuple
import torch
from torch.utils.data import Dataset, DataLoader
import re
from nltk import edit_distance
import numpy as np
from torch.utils.data import DataLoader
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.callbacks import Callback, ModelCheckpoint
from pytorch_lightning.callbacks import ModelCheckpoint
import pytorch_lightning as pl
import os
from huggingface_hub import login
from pytorch_lightning.plugins import CheckpointIO
import argparse
from sconf import Config
from utils.checkpoint import CustomCheckpointIO
from utils.donut_dataset import DonutDataset
from utils.donut_model_pl import DonutModelPLModule
from utils.callbacks import PushToHubCallback
import warnings
DATASET_PATH = "Zombely/fiszki-ocr-train"
PRETRAINED_MODEL_PATH = "Zombely/plwiki-proto-fine-tuned-v2"
START_MODEL_PATH = "Zombely/plwiki-proto-fine-tuned-v2"
OUTPUT_MODEL_PATH = "Zombely/plwiki-proto-fine-tuned-v3"
LOGGING_PATH = "fiszki-ocr-fine-tune"
CHECKPOINT_PATH = "./checkpoint"
def main(config, hug_token):
config_vision = VisionEncoderDecoderConfig.from_pretrained(
config.pretrained_model_path)
config_vision.encoder.image_size = config.image_size
config_vision.decoder.max_length = config.max_length
train_config = {
"max_epochs":1,
"val_check_interval":0.5, # how many times we want to validate during an epoch
"check_val_every_n_epoch":1,
"gradient_clip_val":1.0,
"num_training_samples_per_epoch": 800,
"lr":3e-5,
"train_batch_sizes": [8],
"val_batch_sizes": [1],
"seed":2022,
"num_nodes": 1,
"warmup_steps": 300, # 800/8*30/10, 10%
"result_path": "./result",
"verbose": True,
}
processor = DonutProcessor.from_pretrained(config.start_model_path)
model = VisionEncoderDecoderModel.from_pretrained(
config.pretrained_model_path, config=config_vision)
processor.image_processor.size = config.image_size[::-1]
processor.image_processor.do_align_long_axis = False
added_tokens = []
dataset = load_dataset(DATASET_PATH)
max_length = 768
image_size = [1920, 2560]
config = VisionEncoderDecoderConfig.from_pretrained(PRETRAINED_MODEL_PATH)
config.encoder.image_size = image_size # (height, width)
config.decoder.max_length = max_length
processor = DonutProcessor.from_pretrained(START_MODEL_PATH)
model = VisionEncoderDecoderModel.from_pretrained(PRETRAINED_MODEL_PATH, config=config)
added_tokens = []
class CustomCheckpointIO(CheckpointIO):
def save_checkpoint(self, checkpoint, path, storage_options=None):
del checkpoint["state_dict"]
torch.save(checkpoint, path)
def load_checkpoint(self, path, storage_options=None):
checkpoint = torch.load(path + "artifacts.ckpt")
state_dict = torch.load(path + "pytorch_model.bin")
checkpoint["state_dict"] = {"model." + key: value for key, value in state_dict.items()}
return checkpoint
def remove_checkpoint(self, path) -> None:
return super().remove_checkpoint(path)
class DonutDataset(Dataset):
"""
DonutDataset which is saved in huggingface datasets format. (see details in https://huggingface.co/docs/datasets)
Each row, consists of image path(png/jpg/jpeg) and gt data (json/jsonl/txt),
and it will be converted into input_tensor(vectorized image) and input_ids(tokenized string).
Args:
dataset_name_or_path: name of dataset (available at huggingface.co/datasets) or the path containing image files and metadata.jsonl
max_length: the max number of tokens for the target sequences
split: whether to load "train", "validation" or "test" split
ignore_id: ignore_index for torch.nn.CrossEntropyLoss
task_start_token: the special token to be fed to the decoder to conduct the target task
prompt_end_token: the special token at the end of the sequences
sort_json_key: whether or not to sort the JSON keys
"""
def __init__(
self,
dataset_name_or_path: str,
max_length: int,
split: str = "train",
ignore_id: int = -100,
task_start_token: str = "<s>",
prompt_end_token: str = None,
sort_json_key: bool = True,
):
super().__init__()
self.max_length = max_length
self.split = split
self.ignore_id = ignore_id
self.task_start_token = task_start_token
self.prompt_end_token = prompt_end_token if prompt_end_token else task_start_token
self.sort_json_key = sort_json_key
self.dataset = load_dataset(dataset_name_or_path, split=self.split)
self.dataset_length = len(self.dataset)
self.gt_token_sequences = []
for sample in self.dataset:
ground_truth = json.loads(sample["ground_truth"])
if "gt_parses" in ground_truth: # when multiple ground truths are available, e.g., docvqa
assert isinstance(ground_truth["gt_parses"], list)
gt_jsons = ground_truth["gt_parses"]
else:
assert "gt_parse" in ground_truth and isinstance(ground_truth["gt_parse"], dict)
gt_jsons = [ground_truth["gt_parse"]]
self.gt_token_sequences.append(
[
self.json2token(
gt_json,
update_special_tokens_for_json_key=self.split == "train",
sort_json_key=self.sort_json_key,
train_dataset = DonutDataset(
config.dataset_path,
processor=processor,
model=model,
max_length=config.max_length,
split="train",
task_start_token="<s_cord-v2>",
prompt_end_token="<s_cord-v2>",
added_tokens=added_tokens,
sort_json_key=False, # cord dataset is preprocessed, so no need for this
)
+ processor.tokenizer.eos_token
for gt_json in gt_jsons # load json from list of json
]
)
self.add_tokens([self.task_start_token, self.prompt_end_token])
self.prompt_end_token_id = processor.tokenizer.convert_tokens_to_ids(self.prompt_end_token)
def json2token(self, obj: Any, update_special_tokens_for_json_key: bool = True, sort_json_key: bool = True):
"""
Convert an ordered JSON object into a token sequence
"""
if type(obj) == dict:
if len(obj) == 1 and "text_sequence" in obj:
return obj["text_sequence"]
else:
output = ""
if sort_json_key:
keys = sorted(obj.keys(), reverse=True)
else:
keys = obj.keys()
for k in keys:
if update_special_tokens_for_json_key:
self.add_tokens([fr"<s_{k}>", fr"</s_{k}>"])
output += (
fr"<s_{k}>"
+ self.json2token(obj[k], update_special_tokens_for_json_key, sort_json_key)
+ fr"</s_{k}>"
val_dataset = DonutDataset(
config.dataset_path,
processor=processor,
model=model,
max_length=config.max_length,
split="validation",
task_start_token="<s_cord-v2>",
prompt_end_token="<s_cord-v2>",
added_tokens=added_tokens,
sort_json_key=False, # cord dataset is preprocessed, so no need for this
)
return output
elif type(obj) == list:
return r"<sep/>".join(
[self.json2token(item, update_special_tokens_for_json_key, sort_json_key) for item in obj]
)
else:
obj = str(obj)
if f"<{obj}/>" in added_tokens:
obj = f"<{obj}/>" # for categorical special tokens
return obj
def add_tokens(self, list_of_tokens: List[str]):
"""
Add special tokens to tokenizer and resize the token embeddings of the decoder
"""
newly_added_num = processor.tokenizer.add_tokens(list_of_tokens)
if newly_added_num > 0:
model.decoder.resize_token_embeddings(len(processor.tokenizer))
added_tokens.extend(list_of_tokens)
model.config.pad_token_id = processor.tokenizer.pad_token_id
model.config.decoder_start_token_id = processor.tokenizer.convert_tokens_to_ids(['<s_cord-v2>'])[0]
def __len__(self) -> int:
return self.dataset_length
train_dataloader = DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=4)
val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=4)
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Load image from image_path of given dataset_path and convert into input_tensor and labels
Convert gt data into input_ids (tokenized string)
Returns:
input_tensor : preprocessed image
input_ids : tokenized gt_data
labels : masked labels (model doesn't need to predict prompt and pad token)
"""
sample = self.dataset[idx]
login(hug_token, True)
# inputs
pixel_values = processor(sample["image"], random_padding=self.split == "train", return_tensors="pt").pixel_values
pixel_values = pixel_values.squeeze()
model_module = DonutModelPLModule(config.train_config.toDict(), processor, model)
# targets
target_sequence = random.choice(self.gt_token_sequences[idx]) # can be more than one, e.g., DocVQA Task 1
input_ids = processor.tokenizer(
target_sequence,
add_special_tokens=False,
max_length=self.max_length,
padding="max_length",
truncation=True,
return_tensors="pt",
)["input_ids"].squeeze(0)
wandb_logger = WandbLogger(project="Donut", name=config.wandb_test_name)
labels = input_ids.clone()
labels[labels == processor.tokenizer.pad_token_id] = self.ignore_id # model doesn't need to predict pad token
# labels[: torch.nonzero(labels == self.prompt_end_token_id).sum() + 1] = self.ignore_id # model doesn't need to predict prompt (for VQA)
return pixel_values, labels, target_sequence
processor.image_processor.size = image_size[::-1] # should be (width, height)
processor.image_processor.do_align_long_axis = False
train_dataset = DonutDataset(DATASET_PATH, max_length=max_length,
split="train", task_start_token="<s_cord-v2>", prompt_end_token="<s_cord-v2>",
sort_json_key=False, # cord dataset is preprocessed, so no need for this
)
val_dataset = DonutDataset(DATASET_PATH, max_length=max_length,
split="validation", task_start_token="<s_cord-v2>", prompt_end_token="<s_cord-v2>",
sort_json_key=False, # cord dataset is preprocessed, so no need for this
)
model.config.pad_token_id = processor.tokenizer.pad_token_id
model.config.decoder_start_token_id = processor.tokenizer.convert_tokens_to_ids(['<s_cord-v2>'])[0]
train_dataloader = DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=4)
val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=4)
class DonutModelPLModule(pl.LightningModule):
def __init__(self, config, processor, model):
super().__init__()
self.config = config
self.processor = processor
self.model = model
def training_step(self, batch, batch_idx):
pixel_values, labels, _ = batch
outputs = self.model(pixel_values, labels=labels)
loss = outputs.loss
self.log_dict({"train_loss": loss}, sync_dist=True)
return loss
def validation_step(self, batch, batch_idx, dataset_idx=0):
pixel_values, labels, answers = batch
batch_size = pixel_values.shape[0]
# we feed the prompt to the model
decoder_input_ids = torch.full((batch_size, 1), self.model.config.decoder_start_token_id, device=self.device)
outputs = self.model.generate(pixel_values,
decoder_input_ids=decoder_input_ids,
max_length=max_length,
early_stopping=True,
pad_token_id=self.processor.tokenizer.pad_token_id,
eos_token_id=self.processor.tokenizer.eos_token_id,
use_cache=True,
num_beams=1,
bad_words_ids=[[self.processor.tokenizer.unk_token_id]],
return_dict_in_generate=True,)
predictions = []
for seq in self.processor.tokenizer.batch_decode(outputs.sequences):
seq = seq.replace(self.processor.tokenizer.eos_token, "").replace(self.processor.tokenizer.pad_token, "")
seq = re.sub(r"<.*?>", "", seq, count=1).strip() # remove first task start token
predictions.append(seq)
scores = list()
for pred, answer in zip(predictions, answers):
pred = re.sub(r"(?:(?<=>) | (?=</s_))", "", pred)
# NOT NEEDED ANYMORE
# answer = re.sub(r"<.*?>", "", answer, count=1)
answer = answer.replace(self.processor.tokenizer.eos_token, "")
scores.append(edit_distance(pred, answer) / max(len(pred), len(answer)))
if self.config.get("verbose", False) and len(scores) == 1:
print(f"Prediction: {pred}")
print(f" Answer: {answer}")
print(f" Normed ED: {scores[0]}")
return scores
def validation_epoch_end(self, validation_step_outputs):
# I set this to 1 manually
# (previously set to len(self.config.dataset_name_or_paths))
num_of_loaders = 1
if num_of_loaders == 1:
validation_step_outputs = [validation_step_outputs]
assert len(validation_step_outputs) == num_of_loaders
cnt = [0] * num_of_loaders
total_metric = [0] * num_of_loaders
val_metric = [0] * num_of_loaders
for i, results in enumerate(validation_step_outputs):
for scores in results:
cnt[i] += len(scores)
total_metric[i] += np.sum(scores)
val_metric[i] = total_metric[i] / cnt[i]
val_metric_name = f"val_metric_{i}th_dataset"
self.log_dict({val_metric_name: val_metric[i]}, sync_dist=True)
self.log_dict({"val_metric": np.sum(total_metric) / np.sum(cnt)}, sync_dist=True)
def configure_optimizers(self):
# TODO add scheduler
optimizer = torch.optim.Adam(self.parameters(), lr=self.config.get("lr"))
return optimizer
def train_dataloader(self):
return train_dataloader
def val_dataloader(self):
return val_dataloader
class PushToHubCallback(Callback):
def on_train_epoch_end(self, trainer, pl_module):
print(f"Pushing model to the hub, epoch {trainer.current_epoch}")
pl_module.model.push_to_hub(OUTPUT_MODEL_PATH,
commit_message=f"Training in progress, epoch {trainer.current_epoch}")
# pl_module.processor.push_to_hub(OUTPUT_MODEL_PATH, commit_message=f"Training in progress, epoch {trainer.current_epoch}")
def on_train_end(self, trainer, pl_module):
print(f"Pushing model to the hub after training")
pl_module.processor.push_to_hub(OUTPUT_MODEL_PATH,
commit_message=f"Training done")
pl_module.model.push_to_hub(OUTPUT_MODEL_PATH,
commit_message=f"Training done")
login(os.environ.get("HUG_TOKKEN", None), True)
# ### Wandb.ai link: https://wandb.ai/michalkozlowski936/Donut?workspace=user-michalkozlowski936
# ### Hugging_face link https://huggingface.co/Zombely
model_module = DonutModelPLModule(train_config, processor, model)
wandb_logger = WandbLogger(project="Donut", name=LOGGING_PATH)
checkpoint_callback = ModelCheckpoint(
checkpoint_callback = ModelCheckpoint(
monitor="val_metric",
dirpath=CHECKPOINT_PATH,
dirpath=config.checkpoint_path,
filename="artifacts",
save_top_k=1,
save_last=False,
mode="min",
)
custom_ckpt = CustomCheckpointIO()
custom_ckpt = CustomCheckpointIO()
trainer = pl.Trainer(
trainer = pl.Trainer(
accelerator="gpu" if torch.cuda.is_available() else 'cpu', # change to gpu
devices=1,
max_epochs=train_config.get("max_epochs"),
val_check_interval=train_config.get("val_check_interval"),
check_val_every_n_epoch=train_config.get("check_val_every_n_epoch"),
gradient_clip_val=train_config.get("gradient_clip_val"),
max_epochs=config.train_config.max_epochs,
val_check_interval=config.train_config.val_check_interval,
check_val_every_n_epoch=config.train_config.check_val_every_n_epoch,
gradient_clip_val=config.train_config.gradient_clip_val,
precision=16, # we'll use mixed precision
plugins=custom_ckpt,
num_sanity_val_steps=0,
logger=wandb_logger,
callbacks=[PushToHubCallback(), checkpoint_callback],
)
)
trainer.fit(model_module)
trainer.fit(model_module)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--config", type=str, required=True)
args, left_argv = parser.parse_known_args()
config = Config(args.config)
config.argv_update(left_argv)
hug_token = os.environ.get("HUG_TOKEN", None)
if not torch.cuda.is_available():
warnings.warn("You don't have cuda available, training might be taking long time or impossible")
if not hug_token:
raise Exception("You need to set up HUG_TOKKEN in enviroments to push output model to hub")
main(config, hug_token)

378
old/donut-train_old.py Normal file
View File

@ -0,0 +1,378 @@
#!/usr/bin/env python
# coding: utf-8
from transformers import VisionEncoderDecoderConfig, DonutProcessor, VisionEncoderDecoderModel
from datasets import load_dataset
import json
import random
from typing import Any, List, Tuple
import torch
from torch.utils.data import Dataset, DataLoader
import re
from nltk import edit_distance
import numpy as np
from pytorch_lightning.loggers import WandbLogger
from pytorch_lightning.callbacks import Callback, ModelCheckpoint
import pytorch_lightning as pl
import os
from huggingface_hub import login
from pytorch_lightning.plugins import CheckpointIO
DATASET_PATH = "Zombely/fiszki-ocr-train"
PRETRAINED_MODEL_PATH = "Zombely/plwiki-proto-fine-tuned-v2"
START_MODEL_PATH = "Zombely/plwiki-proto-fine-tuned-v2"
OUTPUT_MODEL_PATH = "Zombely/plwiki-proto-fine-tuned-v3"
LOGGING_PATH = "fiszki-ocr-fine-tune"
CHECKPOINT_PATH = "./checkpoint"
train_config = {
"max_epochs":1,
"val_check_interval":0.5, # how many times we want to validate during an epoch
"check_val_every_n_epoch":1,
"gradient_clip_val":1.0,
"num_training_samples_per_epoch": 800,
"lr":3e-5,
"train_batch_sizes": [8],
"val_batch_sizes": [1],
"seed":2022,
"num_nodes": 1,
"warmup_steps": 300, # 800/8*30/10, 10%
"result_path": "./result",
"verbose": True,
}
dataset = load_dataset(DATASET_PATH)
max_length = 768
image_size = [1920, 2560]
config = VisionEncoderDecoderConfig.from_pretrained(PRETRAINED_MODEL_PATH)
config.encoder.image_size = image_size # (height, width)
config.decoder.max_length = max_length
processor = DonutProcessor.from_pretrained(START_MODEL_PATH)
model = VisionEncoderDecoderModel.from_pretrained(PRETRAINED_MODEL_PATH, config=config)
added_tokens = []
class CustomCheckpointIO(CheckpointIO):
def save_checkpoint(self, checkpoint, path, storage_options=None):
del checkpoint["state_dict"]
torch.save(checkpoint, path)
def load_checkpoint(self, path, storage_options=None):
checkpoint = torch.load(path + "artifacts.ckpt")
state_dict = torch.load(path + "pytorch_model.bin")
checkpoint["state_dict"] = {"model." + key: value for key, value in state_dict.items()}
return checkpoint
def remove_checkpoint(self, path) -> None:
return super().remove_checkpoint(path)
class DonutDataset(Dataset):
"""
DonutDataset which is saved in huggingface datasets format. (see details in https://huggingface.co/docs/datasets)
Each row, consists of image path(png/jpg/jpeg) and gt data (json/jsonl/txt),
and it will be converted into input_tensor(vectorized image) and input_ids(tokenized string).
Args:
dataset_name_or_path: name of dataset (available at huggingface.co/datasets) or the path containing image files and metadata.jsonl
max_length: the max number of tokens for the target sequences
split: whether to load "train", "validation" or "test" split
ignore_id: ignore_index for torch.nn.CrossEntropyLoss
task_start_token: the special token to be fed to the decoder to conduct the target task
prompt_end_token: the special token at the end of the sequences
sort_json_key: whether or not to sort the JSON keys
"""
def __init__(
self,
dataset_name_or_path: str,
max_length: int,
split: str = "train",
ignore_id: int = -100,
task_start_token: str = "<s>",
prompt_end_token: str = None,
sort_json_key: bool = True,
):
super().__init__()
self.max_length = max_length
self.split = split
self.ignore_id = ignore_id
self.task_start_token = task_start_token
self.prompt_end_token = prompt_end_token if prompt_end_token else task_start_token
self.sort_json_key = sort_json_key
self.dataset = load_dataset(dataset_name_or_path, split=self.split)
self.dataset_length = len(self.dataset)
self.gt_token_sequences = []
for sample in self.dataset:
ground_truth = json.loads(sample["ground_truth"])
if "gt_parses" in ground_truth: # when multiple ground truths are available, e.g., docvqa
assert isinstance(ground_truth["gt_parses"], list)
gt_jsons = ground_truth["gt_parses"]
else:
assert "gt_parse" in ground_truth and isinstance(ground_truth["gt_parse"], dict)
gt_jsons = [ground_truth["gt_parse"]]
self.gt_token_sequences.append(
[
self.json2token(
gt_json,
update_special_tokens_for_json_key=self.split == "train",
sort_json_key=self.sort_json_key,
)
+ processor.tokenizer.eos_token
for gt_json in gt_jsons # load json from list of json
]
)
self.add_tokens([self.task_start_token, self.prompt_end_token])
self.prompt_end_token_id = processor.tokenizer.convert_tokens_to_ids(self.prompt_end_token)
def json2token(self, obj: Any, update_special_tokens_for_json_key: bool = True, sort_json_key: bool = True):
"""
Convert an ordered JSON object into a token sequence
"""
if type(obj) == dict:
if len(obj) == 1 and "text_sequence" in obj:
return obj["text_sequence"]
else:
output = ""
if sort_json_key:
keys = sorted(obj.keys(), reverse=True)
else:
keys = obj.keys()
for k in keys:
if update_special_tokens_for_json_key:
self.add_tokens([fr"<s_{k}>", fr"</s_{k}>"])
output += (
fr"<s_{k}>"
+ self.json2token(obj[k], update_special_tokens_for_json_key, sort_json_key)
+ fr"</s_{k}>"
)
return output
elif type(obj) == list:
return r"<sep/>".join(
[self.json2token(item, update_special_tokens_for_json_key, sort_json_key) for item in obj]
)
else:
obj = str(obj)
if f"<{obj}/>" in added_tokens:
obj = f"<{obj}/>" # for categorical special tokens
return obj
def add_tokens(self, list_of_tokens: List[str]):
"""
Add special tokens to tokenizer and resize the token embeddings of the decoder
"""
newly_added_num = processor.tokenizer.add_tokens(list_of_tokens)
if newly_added_num > 0:
model.decoder.resize_token_embeddings(len(processor.tokenizer))
added_tokens.extend(list_of_tokens)
def __len__(self) -> int:
return self.dataset_length
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Load image from image_path of given dataset_path and convert into input_tensor and labels
Convert gt data into input_ids (tokenized string)
Returns:
input_tensor : preprocessed image
input_ids : tokenized gt_data
labels : masked labels (model doesn't need to predict prompt and pad token)
"""
sample = self.dataset[idx]
# inputs
pixel_values = processor(sample["image"], random_padding=self.split == "train", return_tensors="pt").pixel_values
pixel_values = pixel_values.squeeze()
# targets
target_sequence = random.choice(self.gt_token_sequences[idx]) # can be more than one, e.g., DocVQA Task 1
input_ids = processor.tokenizer(
target_sequence,
add_special_tokens=False,
max_length=self.max_length,
padding="max_length",
truncation=True,
return_tensors="pt",
)["input_ids"].squeeze(0)
labels = input_ids.clone()
labels[labels == processor.tokenizer.pad_token_id] = self.ignore_id # model doesn't need to predict pad token
# labels[: torch.nonzero(labels == self.prompt_end_token_id).sum() + 1] = self.ignore_id # model doesn't need to predict prompt (for VQA)
return pixel_values, labels, target_sequence
processor.image_processor.size = image_size[::-1] # should be (width, height)
processor.image_processor.do_align_long_axis = False
train_dataset = DonutDataset(DATASET_PATH, max_length=max_length,
split="train", task_start_token="<s_cord-v2>", prompt_end_token="<s_cord-v2>",
sort_json_key=False, # cord dataset is preprocessed, so no need for this
)
val_dataset = DonutDataset(DATASET_PATH, max_length=max_length,
split="validation", task_start_token="<s_cord-v2>", prompt_end_token="<s_cord-v2>",
sort_json_key=False, # cord dataset is preprocessed, so no need for this
)
model.config.pad_token_id = processor.tokenizer.pad_token_id
model.config.decoder_start_token_id = processor.tokenizer.convert_tokens_to_ids(['<s_cord-v2>'])[0]
train_dataloader = DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=4)
val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=4)
class DonutModelPLModule(pl.LightningModule):
def __init__(self, config, processor, model):
super().__init__()
self.config = config
self.processor = processor
self.model = model
def training_step(self, batch, batch_idx):
pixel_values, labels, _ = batch
outputs = self.model(pixel_values, labels=labels)
loss = outputs.loss
self.log_dict({"train_loss": loss}, sync_dist=True)
return loss
def validation_step(self, batch, batch_idx, dataset_idx=0):
pixel_values, labels, answers = batch
batch_size = pixel_values.shape[0]
# we feed the prompt to the model
decoder_input_ids = torch.full((batch_size, 1), self.model.config.decoder_start_token_id, device=self.device)
outputs = self.model.generate(pixel_values,
decoder_input_ids=decoder_input_ids,
max_length=max_length,
early_stopping=True,
pad_token_id=self.processor.tokenizer.pad_token_id,
eos_token_id=self.processor.tokenizer.eos_token_id,
use_cache=True,
num_beams=1,
bad_words_ids=[[self.processor.tokenizer.unk_token_id]],
return_dict_in_generate=True,)
predictions = []
for seq in self.processor.tokenizer.batch_decode(outputs.sequences):
seq = seq.replace(self.processor.tokenizer.eos_token, "").replace(self.processor.tokenizer.pad_token, "")
seq = re.sub(r"<.*?>", "", seq, count=1).strip() # remove first task start token
predictions.append(seq)
scores = list()
for pred, answer in zip(predictions, answers):
pred = re.sub(r"(?:(?<=>) | (?=</s_))", "", pred)
# NOT NEEDED ANYMORE
# answer = re.sub(r"<.*?>", "", answer, count=1)
answer = answer.replace(self.processor.tokenizer.eos_token, "")
scores.append(edit_distance(pred, answer) / max(len(pred), len(answer)))
if self.config.get("verbose", False) and len(scores) == 1:
print(f"Prediction: {pred}")
print(f" Answer: {answer}")
print(f" Normed ED: {scores[0]}")
return scores
def validation_epoch_end(self, validation_step_outputs):
# I set this to 1 manually
# (previously set to len(self.config.dataset_name_or_paths))
num_of_loaders = 1
if num_of_loaders == 1:
validation_step_outputs = [validation_step_outputs]
assert len(validation_step_outputs) == num_of_loaders
cnt = [0] * num_of_loaders
total_metric = [0] * num_of_loaders
val_metric = [0] * num_of_loaders
for i, results in enumerate(validation_step_outputs):
for scores in results:
cnt[i] += len(scores)
total_metric[i] += np.sum(scores)
val_metric[i] = total_metric[i] / cnt[i]
val_metric_name = f"val_metric_{i}th_dataset"
self.log_dict({val_metric_name: val_metric[i]}, sync_dist=True)
self.log_dict({"val_metric": np.sum(total_metric) / np.sum(cnt)}, sync_dist=True)
def configure_optimizers(self):
# TODO add scheduler
optimizer = torch.optim.Adam(self.parameters(), lr=self.config.get("lr"))
return optimizer
def train_dataloader(self):
return train_dataloader
def val_dataloader(self):
return val_dataloader
class PushToHubCallback(Callback):
def on_train_epoch_end(self, trainer, pl_module):
print(f"Pushing model to the hub, epoch {trainer.current_epoch}")
pl_module.model.push_to_hub(OUTPUT_MODEL_PATH,
commit_message=f"Training in progress, epoch {trainer.current_epoch}")
# pl_module.processor.push_to_hub(OUTPUT_MODEL_PATH, commit_message=f"Training in progress, epoch {trainer.current_epoch}")
def on_train_end(self, trainer, pl_module):
print(f"Pushing model to the hub after training")
pl_module.processor.push_to_hub(OUTPUT_MODEL_PATH,
commit_message=f"Training done")
pl_module.model.push_to_hub(OUTPUT_MODEL_PATH,
commit_message=f"Training done")
login(os.environ.get("HUG_TOKKEN", None), True)
# ### Wandb.ai link: https://wandb.ai/michalkozlowski936/Donut?workspace=user-michalkozlowski936
# ### Hugging_face link https://huggingface.co/Zombely
model_module = DonutModelPLModule(train_config, processor, model)
wandb_logger = WandbLogger(project="Donut", name=LOGGING_PATH)
checkpoint_callback = ModelCheckpoint(
monitor="val_metric",
dirpath=CHECKPOINT_PATH,
filename="artifacts",
save_top_k=1,
save_last=False,
mode="min",
)
custom_ckpt = CustomCheckpointIO()
trainer = pl.Trainer(
accelerator="gpu" if torch.cuda.is_available() else 'cpu', # change to gpu
devices=1,
max_epochs=train_config.get("max_epochs"),
val_check_interval=train_config.get("val_check_interval"),
check_val_every_n_epoch=train_config.get("check_val_every_n_epoch"),
gradient_clip_val=train_config.get("gradient_clip_val"),
precision=16, # we'll use mixed precision
plugins=custom_ckpt,
num_sanity_val_steps=0,
logger=wandb_logger,
callbacks=[PushToHubCallback(), checkpoint_callback],
)
trainer.fit(model_module)

0
utils/__init__.py Normal file
View File

21
utils/callbacks.py Normal file
View File

@ -0,0 +1,21 @@
from pytorch_lightning.callbacks import Callback
class PushToHubCallback(Callback):
def __init__(self, output_model_path) -> None:
super().__init__()
self.output_model_path = output_model_path
def on_train_epoch_end(self, trainer, pl_module):
print(f"Pushing model to the hub, epoch {trainer.current_epoch}")
pl_module.model.push_to_hub(self.output_model_path,
commit_message=f"Training in progress, epoch {trainer.current_epoch}")
# pl_module.processor.push_to_hub(self.output_model_path, commit_message=f"Training in progress, epoch {trainer.current_epoch}")
def on_train_end(self, trainer, pl_module):
print(f"Pushing model to the hub after training")
pl_module.processor.push_to_hub(self.output_model_path,
commit_message=f"Training done")
pl_module.model.push_to_hub(self.output_model_path,
commit_message=f"Training done")

17
utils/checkpoint.py Normal file
View File

@ -0,0 +1,17 @@
from pytorch_lightning.plugins import CheckpointIO
import torch
class CustomCheckpointIO(CheckpointIO):
def save_checkpoint(self, checkpoint, path, storage_options=None):
del checkpoint["state_dict"]
torch.save(checkpoint, path)
def load_checkpoint(self, path, storage_options=None):
checkpoint = torch.load(path + "artifacts.ckpt")
state_dict = torch.load(path + "pytorch_model.bin")
checkpoint["state_dict"] = {"model." + key: value for key, value in state_dict.items()}
return checkpoint
def remove_checkpoint(self, path) -> None:
return super().remove_checkpoint(path)

152
utils/donut_dataset.py Normal file
View File

@ -0,0 +1,152 @@
from datasets import load_dataset
from torch.utils.data import Dataset
import json
from typing import Any, List, Tuple
import random
import torch
from transformers import DonutProcessor, VisionEncoderDecoderModel
class DonutDataset(Dataset):
"""
DonutDataset which is saved in huggingface datasets format. (see details in https://huggingface.co/docs/datasets)
Each row, consists of image path(png/jpg/jpeg) and gt data (json/jsonl/txt),
and it will be converted into input_tensor(vectorized image) and input_ids(tokenized string).
Args:
dataset_name_or_path: name of dataset (available at huggingface.co/datasets) or the path containing image files and metadata.jsonl
max_length: the max number of tokens for the target sequences
split: whether to load "train", "validation" or "test" split
ignore_id: ignore_index for torch.nn.CrossEntropyLoss
task_start_token: the special token to be fed to the decoder to conduct the target task
prompt_end_token: the special token at the end of the sequences
sort_json_key: whether or not to sort the JSON keys
"""
def __init__(
self,
dataset_name_or_path: str,
max_length: int,
processor: DonutProcessor,
model: VisionEncoderDecoderModel,
split: str = "train",
ignore_id: int = -100,
task_start_token: str = "<s>",
prompt_end_token: str = None,
sort_json_key: bool = True,
added_tokens: list = []
):
super().__init__()
self.max_length = max_length
self.split = split
self.processor = processor
self.model = model
self.ignore_id = ignore_id
self.task_start_token = task_start_token
self.prompt_end_token = prompt_end_token if prompt_end_token else task_start_token
self.sort_json_key = sort_json_key
self.added_tokens = added_tokens
self.dataset = load_dataset(dataset_name_or_path, split=self.split)
self.dataset_length = len(self.dataset)
self.gt_token_sequences = []
for sample in self.dataset:
ground_truth = json.loads(sample["ground_truth"])
if "gt_parses" in ground_truth: # when multiple ground truths are available, e.g., docvqa
assert isinstance(ground_truth["gt_parses"], list)
gt_jsons = ground_truth["gt_parses"]
else:
assert "gt_parse" in ground_truth and isinstance(ground_truth["gt_parse"], dict)
gt_jsons = [ground_truth["gt_parse"]]
self.gt_token_sequences.append(
[
self.json2token(
gt_json,
update_special_tokens_for_json_key=self.split == "train",
sort_json_key=self.sort_json_key,
)
+ self.processor.tokenizer.eos_token
for gt_json in gt_jsons # load json from list of json
]
)
self.add_tokens([self.task_start_token, self.prompt_end_token])
self.prompt_end_token_id = self.processor.tokenizer.convert_tokens_to_ids(self.prompt_end_token)
def json2token(self, obj: Any, update_special_tokens_for_json_key: bool = True, sort_json_key: bool = True):
"""
Convert an ordered JSON object into a token sequence
"""
if type(obj) == dict:
if len(obj) == 1 and "text_sequence" in obj:
return obj["text_sequence"]
else:
output = ""
if sort_json_key:
keys = sorted(obj.keys(), reverse=True)
else:
keys = obj.keys()
for k in keys:
if update_special_tokens_for_json_key:
self.add_tokens([fr"<s_{k}>", fr"</s_{k}>"])
output += (
fr"<s_{k}>"
+ self.json2token(obj[k], update_special_tokens_for_json_key, sort_json_key)
+ fr"</s_{k}>"
)
return output
elif type(obj) == list:
return r"<sep/>".join(
[self.json2token(item, update_special_tokens_for_json_key, sort_json_key) for item in obj]
)
else:
obj = str(obj)
if f"<{obj}/>" in self.added_tokens:
obj = f"<{obj}/>" # for categorical special tokens
return obj
def add_tokens(self, list_of_tokens: List[str]):
"""
Add special tokens to tokenizer and resize the token embeddings of the decoder
"""
newly_added_num = self.processor.tokenizer.add_tokens(list_of_tokens)
if newly_added_num > 0:
self.model.decoder.resize_token_embeddings(len(self.processor.tokenizer))
self.added_tokens.extend(list_of_tokens)
def __len__(self) -> int:
return self.dataset_length
def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Load image from image_path of given dataset_path and convert into input_tensor and labels
Convert gt data into input_ids (tokenized string)
Returns:
input_tensor : preprocessed image
input_ids : tokenized gt_data
labels : masked labels (model doesn't need to predict prompt and pad token)
"""
sample = self.dataset[idx]
# inputs
pixel_values = self.processor(sample["image"], random_padding=self.split == "train", return_tensors="pt").pixel_values
pixel_values = pixel_values.squeeze()
# targets
target_sequence = random.choice(self.gt_token_sequences[idx]) # can be more than one, e.g., DocVQA Task 1
input_ids = self.processor.tokenizer(
target_sequence,
add_special_tokens=False,
max_length=self.max_length,
padding="max_length",
truncation=True,
return_tensors="pt",
)["input_ids"].squeeze(0)
labels = input_ids.clone()
labels[labels == self.processor.tokenizer.pad_token_id] = self.ignore_id # model doesn't need to predict pad token
# labels[: torch.nonzero(labels == self.prompt_end_token_id).sum() + 1] = self.ignore_id # model doesn't need to predict prompt (for VQA)
return pixel_values, labels, target_sequence

93
utils/donut_model_pl.py Normal file
View File

@ -0,0 +1,93 @@
import torch
import pytorch_lightning as pl
from nltk import edit_distance
import re
import numpy as np
class DonutModelPLModule(pl.LightningModule):
def __init__(self, config, processor, model, max_length, train_dataloader, val_dataloader):
super().__init__()
self.config = config
self.processor = processor
self.model = model
self.max_length = max_length
self._train_dataloader = train_dataloader
self._val_dataloader = val_dataloader
def training_step(self, batch, batch_idx):
pixel_values, labels, _ = batch
outputs = self.model(pixel_values, labels=labels)
loss = outputs.loss
self.log_dict({"train_loss": loss}, sync_dist=True)
return loss
def validation_step(self, batch, batch_idx, dataset_idx=0):
pixel_values, labels, answers = batch
batch_size = pixel_values.shape[0]
# we feed the prompt to the model
decoder_input_ids = torch.full((batch_size, 1), self.model.config.decoder_start_token_id, device=self.device)
outputs = self.model.generate(pixel_values,
decoder_input_ids=decoder_input_ids,
max_length=self.max_length,
early_stopping=True,
pad_token_id=self.processor.tokenizer.pad_token_id,
eos_token_id=self.processor.tokenizer.eos_token_id,
use_cache=True,
num_beams=1,
bad_words_ids=[[self.processor.tokenizer.unk_token_id]],
return_dict_in_generate=True,)
predictions = []
for seq in self.processor.tokenizer.batch_decode(outputs.sequences):
seq = seq.replace(self.processor.tokenizer.eos_token, "").replace(self.processor.tokenizer.pad_token, "")
seq = re.sub(r"<.*?>", "", seq, count=1).strip() # remove first task start token
predictions.append(seq)
scores = list()
for pred, answer in zip(predictions, answers):
pred = re.sub(r"(?:(?<=>) | (?=</s_))", "", pred)
# NOT NEEDED ANYMORE
# answer = re.sub(r"<.*?>", "", answer, count=1)
answer = answer.replace(self.processor.tokenizer.eos_token, "")
scores.append(edit_distance(pred, answer) / max(len(pred), len(answer)))
if self.config.get("verbose", False) and len(scores) == 1:
print(f"Prediction: {pred}")
print(f" Answer: {answer}")
print(f" Normed ED: {scores[0]}")
return scores
def validation_epoch_end(self, validation_step_outputs):
# I set this to 1 manually
# (previously set to len(self.config.dataset_name_or_paths))
num_of_loaders = 1
if num_of_loaders == 1:
validation_step_outputs = [validation_step_outputs]
assert len(validation_step_outputs) == num_of_loaders
cnt = [0] * num_of_loaders
total_metric = [0] * num_of_loaders
val_metric = [0] * num_of_loaders
for i, results in enumerate(validation_step_outputs):
for scores in results:
cnt[i] += len(scores)
total_metric[i] += np.sum(scores)
val_metric[i] = total_metric[i] / cnt[i]
val_metric_name = f"val_metric_{i}th_dataset"
self.log_dict({val_metric_name: val_metric[i]}, sync_dist=True)
self.log_dict({"val_metric": np.sum(total_metric) / np.sum(cnt)}, sync_dist=True)
def configure_optimizers(self):
# TODO add scheduler
optimizer = torch.optim.Adam(self.parameters(), lr=self.config.get("lr"))
return optimizer
def train_dataloader(self):
return self._train_dataloader
def val_dataloader(self):
return self._val_dataloader