streaming

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="<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 = 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
+                    )
+
+    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)
+
+    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)

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 = "<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, 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"<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]
+
+        # 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