245 lines
9.4 KiB
Python
245 lines
9.4 KiB
Python
from typing import Any, List
|
|
from transformers import VisionEncoderDecoderConfig, DonutProcessor, VisionEncoderDecoderModel
|
|
import torch
|
|
from torch.utils.data import DataLoader
|
|
from pytorch_lightning.loggers import WandbLogger
|
|
from pytorch_lightning.callbacks import ModelCheckpoint
|
|
import pytorch_lightning as pl
|
|
import os
|
|
from huggingface_hub import login
|
|
import argparse
|
|
from sconf import Config
|
|
from utils.checkpoint import CustomCheckpointIO
|
|
from utils.donut_dataset_stream import DonutDatasetStream
|
|
from utils.donut_model_pl_stream import DonutModelPLModuleStream
|
|
from utils.callbacks import PushToHubCallback
|
|
import warnings
|
|
from datasets import load_dataset, interleave_datasets
|
|
from torchdata.datapipes.iter import IterableWrapper
|
|
import json
|
|
|
|
|
|
class TestIterator(IterableWrapper):
|
|
def __init__(self, iterable, deepcopy=True, total_len=None):
|
|
super().__init__(iterable, deepcopy)
|
|
self.total_len = total_len
|
|
|
|
def __len__(self):
|
|
if self.total_len:
|
|
return self.total_len
|
|
return super().__len__()
|
|
|
|
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
|
|
|
|
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 = []
|
|
|
|
### PROCESS FUNC START ###
|
|
|
|
def add_tokens(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 json2token(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:
|
|
add_tokens([fr"<s_{k}>", fr"</s_{k}>"])
|
|
output += (
|
|
fr"<s_{k}>"
|
|
+ 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(
|
|
[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 process(row, split):
|
|
task_start_token, prompt_end_token = "<s_cord-v2>", "<s_cord-v2>"
|
|
ground_truth = json.loads(row["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"]]
|
|
|
|
gt_token_sequences = (
|
|
[
|
|
json2token(
|
|
gt_json,
|
|
update_special_tokens_for_json_key=split == "train",
|
|
sort_json_key=False,
|
|
)
|
|
+ processor.tokenizer.eos_token
|
|
for gt_json in gt_jsons # load json from list of json
|
|
]
|
|
)
|
|
|
|
add_tokens([task_start_token, prompt_end_token])
|
|
prompt_end_token_id = processor.tokenizer.convert_tokens_to_ids(prompt_end_token)
|
|
|
|
# change if not 3 channels
|
|
if row['image'].mode != "RGB":
|
|
row['image'] = row['image'].convert("RGB")
|
|
|
|
# inputs
|
|
pixel_values = processor(row["image"], random_padding=split == "train", return_tensors="pt").pixel_values
|
|
pixel_values = pixel_values.squeeze()
|
|
|
|
# targets
|
|
input_ids = processor.tokenizer(
|
|
gt_token_sequences,
|
|
add_special_tokens=False,
|
|
max_length=config.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] = -100 # model doesn't need to predict pad token
|
|
return {"pixel_values": pixel_values, "labels": labels, 'target_sequence': gt_token_sequences }
|
|
|
|
def proces_train(row):
|
|
return process(row, 'train')
|
|
|
|
def proces_val(row):
|
|
return process(row, 'validation')
|
|
|
|
|
|
|
|
|
|
### PROCESS FUNC END ###
|
|
|
|
# train_dataset_process = DonutDatasetStream(
|
|
# 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
|
|
# )
|
|
|
|
# val_dataset_process = DonutDatasetStream(
|
|
# 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
|
|
# )
|
|
|
|
dataset = load_dataset(config.dataset_path, streaming=True)
|
|
val_dataset = dataset.pop('validation')
|
|
train_dataset = interleave_datasets(list(dataset.values()))
|
|
train_length = sum(split.num_examples for split in dataset[list(dataset.keys())[0]].info.splits.values() if split.name != 'validation')
|
|
val_length = list(val_dataset.info.splits.values())[-1].num_examples
|
|
|
|
|
|
train_dataset = train_dataset.map(proces_train, remove_columns = ['image', 'ground_truth'])
|
|
val_dataset = val_dataset.map(proces_val, remove_columns = ['image', 'ground_truth'])
|
|
|
|
# train_dataset = train_dataset.with_format('torch')
|
|
# val_dataset = val_dataset.with_format('torch')
|
|
|
|
train_dataset = TestIterator(train_dataset, total_len=train_length)
|
|
val_dataset = TestIterator(val_dataset, total_len=val_length)
|
|
|
|
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, num_workers=0)
|
|
val_dataloader = DataLoader(val_dataset, batch_size=1, num_workers=0)
|
|
|
|
login(hug_token, True)
|
|
|
|
model_module = DonutModelPLModuleStream(config.train_config.toDict(), processor, model, max_length=config.max_length, train_dataloader=train_dataloader, val_dataloader=val_dataloader)
|
|
|
|
wandb_logger = WandbLogger(project="Donut", name=config.wandb_test_name)
|
|
|
|
if not os.path.exists(config.checkpoint_path):
|
|
os.mkdir(config.checkpoint_path)
|
|
|
|
checkpoint_callback = ModelCheckpoint(
|
|
monitor="val_metric",
|
|
dirpath=config.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=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(output_model_path=config.output_model_path), checkpoint_callback],
|
|
)
|
|
|
|
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", "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")
|
|
|
|
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_TOKEN in enviroments to push output model to hub")
|
|
main(config, hug_token)
|