diff --git a/train_stream.py b/train_stream.py new file mode 100644 index 0000000..901564d --- /dev/null +++ b/train_stream.py @@ -0,0 +1,115 @@ +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 DonutDataset +from utils.donut_model_pl import DonutModelPLModule +from utils.callbacks import PushToHubCallback +import warnings +from datasets import load_dataset + + + + +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 = [] + + train_dataset = DonutDataset( + config.dataset_path, + processor=processor, + model=model, + max_length=config.max_length, + split="train", + task_start_token="", + prompt_end_token="", + added_tokens=added_tokens, + sort_json_key=False, # cord dataset is preprocessed, so no need for this + ) + + val_dataset = DonutDataset( + config.dataset_path, + processor=processor, + model=model, + max_length=config.max_length, + split="validation", + task_start_token="", + prompt_end_token="", + added_tokens=added_tokens, + 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([''])[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) + + login(hug_token, True) + + model_module = DonutModelPLModule(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) + + 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", 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_TOKEN in enviroments to push output model to hub") + main(config, hug_token) diff --git a/utils/donut_dataset_stream.py b/utils/donut_dataset_stream.py new file mode 100644 index 0000000..adbb43d --- /dev/null +++ b/utils/donut_dataset_stream.py @@ -0,0 +1,155 @@ +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 = "", + 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, stream=True).with_format("torch") + 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"", fr""]) + output += ( + fr"" + + self.json2token(obj[k], update_special_tokens_for_json_key, sort_json_key) + + fr"" + ) + return output + elif type(obj) == list: + return r"".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] + + # change if not 3 channels + if sample['image'].mode != "RGB": + sample['image'] = sample['image'].convert("RGB") + + # 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