From e3e0fa495d5880c87a22a32de868e20b4d21a7f1 Mon Sep 17 00:00:00 2001 From: s444415 Date: Wed, 14 Dec 2022 14:53:27 +0000 Subject: [PATCH] after training --- .gitignore | 4 + donut-eval.py | 83 ++++++++++ donut-train.py | 387 +++++++++++++++++++++++++++++++++++++++++++++++ requirements.txt | 4 +- 4 files changed, 476 insertions(+), 2 deletions(-) create mode 100644 .gitignore create mode 100644 donut-eval.py create mode 100644 donut-train.py diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..cd703ca --- /dev/null +++ b/.gitignore @@ -0,0 +1,4 @@ +env +Donut +nohup.out +wandb diff --git a/donut-eval.py b/donut-eval.py new file mode 100644 index 0000000..c5548e6 --- /dev/null +++ b/donut-eval.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python +# coding: utf-8 + +# In[1]: + + +from transformers import DonutProcessor, VisionEncoderDecoderModel +from datasets import load_dataset +import re +import json +import torch +from tqdm.auto import tqdm +import numpy as np + +from donut import JSONParseEvaluator + + +# In[2]: + + +processor = DonutProcessor.from_pretrained("Zombely/plwiki-proto-fine-tuned") +model = VisionEncoderDecoderModel.from_pretrained("Zombely/plwiki-proto-fine-tuned") + + +# In[3]: + + +dataset = load_dataset("Zombely/pl-text-images-5000-whole", split="validation") + + +# In[4]: + + +device = "cuda" if torch.cuda.is_available() else "cpu" + +model.eval() +model.to(device) + +output_list = [] +accs = [] + + +for idx, sample in tqdm(enumerate(dataset), total=len(dataset)): + # prepare encoder inputs + pixel_values = processor(sample["image"].convert("RGB"), return_tensors="pt").pixel_values + pixel_values = pixel_values.to(device) + # prepare decoder inputs + task_prompt = "" + decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors="pt").input_ids + decoder_input_ids = decoder_input_ids.to(device) + + # autoregressively generate sequence + outputs = model.generate( + pixel_values, + decoder_input_ids=decoder_input_ids, + max_length=model.decoder.config.max_position_embeddings, + early_stopping=True, + pad_token_id=processor.tokenizer.pad_token_id, + eos_token_id=processor.tokenizer.eos_token_id, + use_cache=True, + num_beams=1, + bad_words_ids=[[processor.tokenizer.unk_token_id]], + return_dict_in_generate=True, + ) + + # turn into JSON + seq = processor.batch_decode(outputs.sequences)[0] + seq = seq.replace(processor.tokenizer.eos_token, "").replace(processor.tokenizer.pad_token, "") + seq = re.sub(r"<.*?>", "", seq, count=1).strip() # remove first task start token + seq = processor.token2json(seq) + + ground_truth = json.loads(sample["ground_truth"]) + ground_truth = ground_truth["gt_parse"] + evaluator = JSONParseEvaluator() + score = evaluator.cal_acc(seq, ground_truth) + + accs.append(score) + output_list.append(seq) + +scores = {"accuracies": accs, "mean_accuracy": np.mean(accs)} +print(scores, f"length : {len(accs)}") +print("Mean accuracy:", np.mean(accs)) + diff --git a/donut-train.py b/donut-train.py new file mode 100644 index 0000000..6590e1d --- /dev/null +++ b/donut-train.py @@ -0,0 +1,387 @@ +#!/usr/bin/env python +# coding: utf-8 + +# In[19]: + + +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 +import pytorch_lightning as pl +import os +from huggingface_hub import login + + +# In[8]: + + +DATASET_PATH = "Zombely/pl-text-images-5000-whole" +PRETRAINED_MODEL_PATH = "Zombely/plwiki-proto-fine-tuned" +START_MODEL_PATH = "Zombely/plwiki-proto-fine-tuned" +OUTPUT_MODEL_PATH = "Zombely/plwiki-proto-fine-tuned-v2" +LOGGING_PATH = "plwiki-proto-ft-second-iter" + + +# In[ ]: + + +train_config = { + "max_epochs":30, + "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, +} + + +# In[9]: + + +dataset = load_dataset(DATASET_PATH) + + +# In[10]: + + +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 + + +# In[11]: + + +processor = DonutProcessor.from_pretrained(START_MODEL_PATH) +model = VisionEncoderDecoderModel.from_pretrained(PRETRAINED_MODEL_PATH, config=config) + + +# In[12]: + + +added_tokens = [] + +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 = "", + 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"", 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 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 + + +# In[13]: + + +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="", prompt_end_token="", + 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="", prompt_end_token="", + sort_json_key=False, # cord dataset is preprocessed, so no need for this + ) + + +# In[14]: + + +model.config.pad_token_id = processor.tokenizer.pad_token_id +model.config.decoder_start_token_id = processor.tokenizer.convert_tokens_to_ids([''])[0] + + +# In[15]: + + +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) + + +# In[16]: + + +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"(?:(?<=>) | (?=", "", 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 + + +# In[17]: + + +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") + + +# In[18]: + + +login(os.environ.get("HUG_TOKKEN", "")) + + +# ### Wandb.ai link: https://wandb.ai/michalkozlowski936/Donut?workspace=user-michalkozlowski936 +# ### Hugging_face link https://huggingface.co/Zombely + +# In[22]: + + +model_module = DonutModelPLModule(train_config, processor, model) + +wandb_logger = WandbLogger(project="Donut", name=LOGGING_PATH) + +trainer = pl.Trainer( + accelerator="gpu", # 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 + num_sanity_val_steps=0, + logger=wandb_logger, + callbacks=[PushToHubCallback()], +) + +trainer.fit(model_module) + diff --git a/requirements.txt b/requirements.txt index 8d66d90..f1de7a0 100644 --- a/requirements.txt +++ b/requirements.txt @@ -14,7 +14,7 @@ beautifulsoup4==4.11.1 bleach==5.0.1 blend-modes==2.1.0 cachetools==5.2.0 -certifi @ file:///croot/certifi_1665076670883/work/certifi +certifi cffi==1.15.1 charset-normalizer==2.1.1 click==8.1.3 @@ -26,7 +26,7 @@ decorator==5.1.1 defusedxml==0.7.1 dill==0.3.6 docker-pycreds==0.4.0 -donut-python @ file:///home/pc/work/donut +donut-python entrypoints==0.4 evaluate==0.3.0 fastapi==0.87.0