stream trainning

config-train.yaml

@@ -1,14 +1,14 @@
-dataset_path: "Zombely/wikisource-small"
+dataset_path: "Zombely/wikisource-yellow"
 pretrained_model_path: "Zombely/plwiki-proto-fine-tuned-v3.2"
 start_model_path: "Zombely/plwiki-proto-fine-tuned-v3.2"
 output_model_path: "Zombely/pl-donut"
 wandb_test_name: "wikisource-small"
 checkpoint_path: "./checkpoint"
 max_length: 768
-image_size: [2560, 1920]
+image_size: [1280, 960]
 train_config:
   max_epochs: 1
-  val_check_interval: 0.3
+  val_check_interval: 1.0
   check_val_every_n_epoch: 1
   gradient_clip_val: 1.0
   num_training_samples_per_epoch: 800

train_stream.py

@@ -9,12 +9,12 @@ 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.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
-
+from datasets import load_dataset, interleave_datasets
+from torchdata.datapipes.iter import IterableWrapper
 
 
 
@@ -34,8 +34,7 @@ def main(config, hug_token):
 
     added_tokens = []
 
-    train_dataset = DonutDataset(
-                        config.dataset_path, 
+    train_dataset_process = DonutDatasetStream(
                         processor=processor, 
                         model=model, 
                         max_length=config.max_length,
@@ -46,8 +45,7 @@ def main(config, hug_token):
                         sort_json_key=False,  # cord dataset is preprocessed, so no need for this
                     )
 
-    val_dataset = DonutDataset(
-                        config.dataset_path, 
+    val_dataset_process = DonutDatasetStream(
                         processor=processor, 
                         model=model, 
                         max_length=config.max_length,
@@ -57,19 +55,38 @@ def main(config, hug_token):
                         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(lambda x: train_dataset_process.process(x), remove_columns = ['image', 'ground_truth'])
+    val_dataset = val_dataset.map(lambda x: val_dataset_process.process(x), remove_columns = ['image', 'ground_truth'])
+
+    # train_dataset = train_dataset.with_format('torch')
+    # val_dataset = val_dataset.with_format('torch')
+
+    train_dataset = IterableWrapper(train_dataset)
+    val_dataset = IterableWrapper(val_dataset)
 
     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=1)
-    val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=1)
+    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 = DonutModelPLModule(config.train_config.toDict(), processor, model, max_length=config.max_length, train_dataloader=train_dataloader, val_dataloader=val_dataloader)
+    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,
@@ -105,7 +122,7 @@ if __name__ == "__main__":
     config = Config(args.config)
     config.argv_update(left_argv)
 
-    hug_token = os.environ.get("HUG_TOKEN", None)
+    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")

utils/donut_dataset.py

@@ -116,8 +116,8 @@ class DonutDataset(Dataset):
             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 __len__(self) -> int:
+        return self.dataset_length
 
     def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
         """

utils/donut_dataset_stream.py

@@ -7,7 +7,7 @@ import torch
 from transformers import DonutProcessor, VisionEncoderDecoderModel
 
 
-class DonutDataset(Dataset):
+class DonutDatasetStream:
     """
     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),
@@ -24,12 +24,11 @@ class DonutDataset(Dataset):
 
     def __init__(
         self,
-        dataset_name_or_path: str,
-        max_length: int,
         processor: DonutProcessor,
         model: VisionEncoderDecoderModel,
-        split: str = "train",
+        max_length: int,
         ignore_id: int = -100,
+        split: str = 'train',
         task_start_token: str = "<s>",
         prompt_end_token: str = None,
         sort_json_key: bool = True,
@@ -37,8 +36,8 @@ class DonutDataset(Dataset):
     ):
         super().__init__()
 
-        self.max_length = max_length
         self.split = split
+        self.max_length = max_length
         self.processor = processor
         self.model = model
         self.ignore_id = ignore_id
@@ -47,35 +46,56 @@ class DonutDataset(Dataset):
         self.sort_json_key = sort_json_key
         self.added_tokens = added_tokens
 
-        self.dataset = load_dataset(dataset_name_or_path, split=self.split, streaming=True).with_format("torch")
-        print(self.dataset)
-        self.dataset_length = len(self.dataset)
+    def process(self, row):
 
-        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
-                ]
-            )
+        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"]]
+
+        self.gt_token_sequences = (
+            [
+                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)
 
+
+        # change if not 3 channels
+        if row['image'].mode != "RGB":
+            row['image'] = row['image'].convert("RGB")
+
+        # inputs
+        pixel_values = self.processor(row["image"], random_padding=self.split == "train", return_tensors="pt").pixel_values
+        pixel_values = pixel_values.squeeze()
+
+        # targets
+        target_sequence = self.gt_token_sequences  # 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
+        return {"pixel_values": pixel_values, "labels": labels, 'target_sequence': target_sequence }
+
     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
@@ -117,40 +137,3 @@ class DonutDataset(Dataset):
             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

utils/donut_model_pl_stream.py

@@ -0,0 +1,98 @@
+import torch
+import pytorch_lightning as pl
+from nltk import edit_distance
+import re
+import numpy as np
+
+
+class DonutModelPLModuleStream(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
+        pixel_values = batch['pixel_values']
+        labels = batch['labels']
+        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
+
+        pixel_values = batch['pixel_values']
+        labels = batch['labels']
+        answers = batch['target_sequence']
+        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