{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport json\nfrom PIL import Image\nfrom typing import Dict\nimport random\nimport cv2\nimport pandas as pd\nimport numpy as np","metadata":{"_kg_hide-input":false,"_kg_hide-output":false,"execution":{"iopub.status.busy":"2023-03-29T20:18:05.357496Z","iopub.execute_input":"2023-03-29T20:18:05.357880Z","iopub.status.idle":"2023-03-29T20:18:05.561251Z","shell.execute_reply.started":"2023-03-29T20:18:05.357844Z","shell.execute_reply":"2023-03-29T20:18:05.560213Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_annotations():\n    \"\"\"returns for each graph Its category based on the annotation files\"\"\"\n    annotations_path = \"/kaggle/input/benetech-making-graphs-accessible/train/annotations\"\n    annotations = {}\n    for annotation_path in os.listdir(annotations_path):\n        with open(f\"{annotations_path}/{annotation_path}\") as annotation_f:\n            annotations[annotation_path[:-4]+'jpg'] = json.load(annotation_f)[\"chart-type\"]\n    return annotations","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2023-03-29T20:18:05.563336Z","iopub.execute_input":"2023-03-29T20:18:05.563797Z","iopub.status.idle":"2023-03-29T20:18:05.570940Z","shell.execute_reply.started":"2023-03-29T20:18:05.563758Z","shell.execute_reply":"2023-03-29T20:18:05.569768Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"1- For each image, get the annotation (dot, line, vertical_bar, etc) . Put everything in a dictionary","metadata":{}},{"cell_type":"code","source":"annotations = get_annotations()","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:18:05.572467Z","iopub.execute_input":"2023-03-29T20:18:05.573355Z","iopub.status.idle":"2023-03-29T20:22:25.346233Z","shell.execute_reply.started":"2023-03-29T20:18:05.573303Z","shell.execute_reply":"2023-03-29T20:22:25.345143Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"classes = set(annotations.values())","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:22:25.349095Z","iopub.execute_input":"2023-03-29T20:22:25.349497Z","iopub.status.idle":"2023-03-29T20:22:25.360673Z","shell.execute_reply.started":"2023-03-29T20:22:25.349456Z","shell.execute_reply":"2023-03-29T20:22:25.359494Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"2- Create the right folder structure for your data as proposed here in order to create a HuggingFace dataset : https://huggingface.co/docs/datasets/image_load#imagefolder . Other useful link : https://github.com/NielsRogge/Transformers-Tutorials/tree/master/Donut/DocVQA","metadata":{}},{"cell_type":"code","source":"os.mkdir('imagefolder')","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:22:25.362661Z","iopub.execute_input":"2023-03-29T20:22:25.363079Z","iopub.status.idle":"2023-03-29T20:22:25.372431Z","shell.execute_reply.started":"2023-03-29T20:22:25.363040Z","shell.execute_reply":"2023-03-29T20:22:25.371343Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"os.mkdir('./imagefolder/train')","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:22:25.375032Z","iopub.execute_input":"2023-03-29T20:22:25.375833Z","iopub.status.idle":"2023-03-29T20:22:25.382968Z","shell.execute_reply.started":"2023-03-29T20:22:25.375796Z","shell.execute_reply":"2023-03-29T20:22:25.381987Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for k in classes : \n    os.mkdir(f'./imagefolder/train/{k}')","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:22:25.384718Z","iopub.execute_input":"2023-03-29T20:22:25.385152Z","iopub.status.idle":"2023-03-29T20:22:25.396698Z","shell.execute_reply.started":"2023-03-29T20:22:25.385118Z","shell.execute_reply":"2023-03-29T20:22:25.395436Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import shutil\ndirname = '/kaggle/input/benetech-making-graphs-accessible/train/images'\nfor filename in os.listdir(dirname):\n        src = os.path.join(dirname, filename)\n        classe = annotations[filename]\n        dst = f'./imagefolder/train/{classe}/{filename}'\n        shutil.copy(src, dst)","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:22:25.398317Z","iopub.execute_input":"2023-03-29T20:22:25.398886Z","iopub.status.idle":"2023-03-29T20:27:30.800449Z","shell.execute_reply.started":"2023-03-29T20:22:25.398849Z","shell.execute_reply":"2023-03-29T20:27:30.799418Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from datasets import load_dataset\n\ndataset = load_dataset(\"imagefolder\", data_dir=\"./imagefolder\")","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:27:30.801890Z","iopub.execute_input":"2023-03-29T20:27:30.802278Z","iopub.status.idle":"2023-03-29T20:28:05.295799Z","shell.execute_reply.started":"2023-03-29T20:27:30.802241Z","shell.execute_reply":"2023-03-29T20:28:05.294654Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dataset","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:28:05.301234Z","iopub.execute_input":"2023-03-29T20:28:05.302577Z","iopub.status.idle":"2023-03-29T20:28:05.310596Z","shell.execute_reply.started":"2023-03-29T20:28:05.302532Z","shell.execute_reply":"2023-03-29T20:28:05.309431Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dataset['train'][0]","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:28:05.312034Z","iopub.execute_input":"2023-03-29T20:28:05.313080Z","iopub.status.idle":"2023-03-29T20:28:05.337687Z","shell.execute_reply.started":"2023-03-29T20:28:05.313041Z","shell.execute_reply":"2023-03-29T20:28:05.336737Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"id2label = {id: label for id, label in enumerate(dataset['train'].features['label'].names)}\nprint(id2label)","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:28:05.338818Z","iopub.execute_input":"2023-03-29T20:28:05.339157Z","iopub.status.idle":"2023-03-29T20:28:05.345211Z","shell.execute_reply.started":"2023-03-29T20:28:05.339122Z","shell.execute_reply":"2023-03-29T20:28:05.344128Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"3-Convert the dataset into a Donut-compatible format","metadata":{}},{"cell_type":"code","source":"template = '{\"gt_parse\": {\"class\" : '","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:28:05.346577Z","iopub.execute_input":"2023-03-29T20:28:05.347375Z","iopub.status.idle":"2023-03-29T20:28:05.356364Z","shell.execute_reply.started":"2023-03-29T20:28:05.347338Z","shell.execute_reply":"2023-03-29T20:28:05.355325Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def update_examples(examples):\n    ground_truths = []\n    for label in examples['label']:\n        ground_truths.append(template + '\"' + id2label[label] + '\"' + \"}}\")\n    examples['ground_truth'] = ground_truths\n\n    return examples\n\ndataset = dataset.map(update_examples, batched=True)","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:28:05.358040Z","iopub.execute_input":"2023-03-29T20:28:05.358483Z","iopub.status.idle":"2023-03-29T20:28:06.880466Z","shell.execute_reply.started":"2023-03-29T20:28:05.358437Z","shell.execute_reply":"2023-03-29T20:28:06.879392Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dataset['train'][0]['ground_truth']","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:28:06.882185Z","iopub.execute_input":"2023-03-29T20:28:06.882544Z","iopub.status.idle":"2023-03-29T20:28:06.894277Z","shell.execute_reply.started":"2023-03-29T20:28:06.882509Z","shell.execute_reply":"2023-03-29T20:28:06.892808Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"4- Let's open an image to check its shape","metadata":{}},{"cell_type":"code","source":"img = Image.open(\"/kaggle/input/benetech-making-graphs-accessible/train/images/0000ae6cbdb1.jpg\")","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:28:06.895703Z","iopub.execute_input":"2023-03-29T20:28:06.895996Z","iopub.status.idle":"2023-03-29T20:28:06.925474Z","shell.execute_reply.started":"2023-03-29T20:28:06.895943Z","shell.execute_reply":"2023-03-29T20:28:06.924439Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"img.width","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:28:06.928816Z","iopub.execute_input":"2023-03-29T20:28:06.929488Z","iopub.status.idle":"2023-03-29T20:28:06.936809Z","shell.execute_reply.started":"2023-03-29T20:28:06.929441Z","shell.execute_reply":"2023-03-29T20:28:06.935536Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"np.array(img).shape","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:28:06.938683Z","iopub.execute_input":"2023-03-29T20:28:06.939075Z","iopub.status.idle":"2023-03-29T20:28:06.951012Z","shell.execute_reply.started":"2023-03-29T20:28:06.939040Z","shell.execute_reply":"2023-03-29T20:28:06.949835Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"5- Define model parameters","metadata":{}},{"cell_type":"code","source":"from transformers import VisionEncoderDecoderConfig\n\nmax_length = 8\nimage_size = [532,400]\n\nconfig = VisionEncoderDecoderConfig.from_pretrained(\"nielsr/donut-base\")\nconfig.encoder.image_size = image_size # (height, width)\nconfig.decoder.max_length = max_length","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:28:06.952609Z","iopub.execute_input":"2023-03-29T20:28:06.953612Z","iopub.status.idle":"2023-03-29T20:28:16.105936Z","shell.execute_reply.started":"2023-03-29T20:28:06.953578Z","shell.execute_reply":"2023-03-29T20:28:16.104924Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from transformers import DonutProcessor, VisionEncoderDecoderModel, BartConfig\n\nprocessor = DonutProcessor.from_pretrained(\"nielsr/donut-base\")\nmodel = VisionEncoderDecoderModel.from_pretrained(\"nielsr/donut-base\", config=config,ignore_mismatched_sizes=True)","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:28:16.107768Z","iopub.execute_input":"2023-03-29T20:28:16.109629Z","iopub.status.idle":"2023-03-29T20:28:32.503876Z","shell.execute_reply.started":"2023-03-29T20:28:16.109586Z","shell.execute_reply":"2023-03-29T20:28:32.502164Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"processor.feature_extractor.size = image_size[::-1] # should be (width, height)\nprocessor.feature_extractor.do_align_long_axis = False","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:28:32.506014Z","iopub.execute_input":"2023-03-29T20:28:32.506450Z","iopub.status.idle":"2023-03-29T20:28:32.517795Z","shell.execute_reply.started":"2023-03-29T20:28:32.506401Z","shell.execute_reply":"2023-03-29T20:28:32.514794Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(processor.tokenizer)","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:28:32.519047Z","iopub.execute_input":"2023-03-29T20:28:32.519835Z","iopub.status.idle":"2023-03-29T20:28:34.768172Z","shell.execute_reply.started":"2023-03-29T20:28:32.519795Z","shell.execute_reply":"2023-03-29T20:28:34.767033Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from typing import List\n\ndef add_tokens(list_of_tokens: List[str]):\n    \"\"\"\n    Add tokens to tokenizer and resize the token embeddings\n    \"\"\"\n    newly_added_num = processor.tokenizer.add_tokens(list_of_tokens)\n    if newly_added_num > 0:\n        model.decoder.resize_token_embeddings(len(processor.tokenizer))","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:28:34.772899Z","iopub.execute_input":"2023-03-29T20:28:34.773797Z","iopub.status.idle":"2023-03-29T20:28:34.840890Z","shell.execute_reply.started":"2023-03-29T20:28:34.773752Z","shell.execute_reply":"2023-03-29T20:28:34.839376Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"additional_tokens = [\"<dot/>\", \"<horizontal_bar/>\", \"<line/>\", \"<scatter/>\", \"<vertical_bar/>\"]\n\nadd_tokens(additional_tokens)","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:28:34.842850Z","iopub.execute_input":"2023-03-29T20:28:34.843632Z","iopub.status.idle":"2023-03-29T20:28:36.047287Z","shell.execute_reply.started":"2023-03-29T20:28:34.843591Z","shell.execute_reply":"2023-03-29T20:28:36.046180Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"processor.tokenizer.convert_tokens_to_ids([\"<horizontal_bar/>\"])","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:28:36.048815Z","iopub.execute_input":"2023-03-29T20:28:36.049345Z","iopub.status.idle":"2023-03-29T20:28:36.063141Z","shell.execute_reply.started":"2023-03-29T20:28:36.049303Z","shell.execute_reply":"2023-03-29T20:28:36.061700Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"6- Define the DonutDataset. I heavily relied on the tutorial from the github https://github.com/NielsRogge/Transformers-Tutorials/blob/master/Donut/DocVQA/Fine_tune_Donut_on_DocVQA.ipynb","metadata":{}},{"cell_type":"code","source":"import json\nimport random\nfrom typing import Any, List, Tuple\n\nimport torch\nfrom torch.utils.data import Dataset\n\nclass DonutDataset(Dataset):\n    \"\"\"\n    DonutDataset which is saved in huggingface datasets format. (see details in https://huggingface.co/docs/datasets)\n    Each row, consists of image path(png/jpg/jpeg) and gt data (json/jsonl/txt),\n    and it will be converted into input_tensor(vectorized image) and input_ids(tokenized string).\n    Args:\n        dataset_name_or_path: name of dataset (available at huggingface.co/datasets) or the path containing image files and metadata.jsonl\n        max_length: the max number of tokens for the target sequences\n        split: whether to load \"train\", \"validation\" or \"test\" split\n        ignore_id: ignore_index for torch.nn.CrossEntropyLoss\n        task_start_token: the special token to be fed to the decoder to conduct the target task\n        prompt_end_token: the special token at the end of the sequences\n        sort_json_key: whether or not to sort the JSON keys\n    \"\"\"\n\n    def __init__(\n        self,\n        dataset,\n        max_length: int,\n        split: str = \"train\",\n        ignore_id: int = -100,\n        task_start_token: str = \"<s>\",\n        prompt_end_token: str = None,\n        sort_json_key: bool = True,\n    ):\n        super().__init__()\n\n        self.max_length = max_length\n        self.split = split\n        self.ignore_id = ignore_id\n        self.task_start_token = task_start_token\n        self.prompt_end_token = prompt_end_token if prompt_end_token else task_start_token\n        self.sort_json_key = sort_json_key\n\n        self.dataset = dataset\n        self.dataset_length = len(self.dataset)\n\n        self.gt_token_sequences = []\n        for sample in self.dataset:\n            ground_truth = json.loads(sample[\"ground_truth\"])\n            if \"gt_parses\" in ground_truth:  # when multiple ground truths are available, e.g., docvqa\n                assert isinstance(ground_truth[\"gt_parses\"], list)\n                gt_jsons = ground_truth[\"gt_parses\"]\n            else:\n                assert \"gt_parse\" in ground_truth and isinstance(ground_truth[\"gt_parse\"], dict)\n                gt_jsons = [ground_truth[\"gt_parse\"]]\n\n            self.gt_token_sequences.append(\n                [\n                    self.json2token(\n                        gt_json,\n                        update_special_tokens_for_json_key=self.split == \"train\",\n                        sort_json_key=self.sort_json_key,\n                    )\n                    + processor.tokenizer.eos_token\n                    for gt_json in gt_jsons  # load json from list of json\n                ]\n            )\n\n        self.add_tokens([self.task_start_token, self.prompt_end_token])\n        self.prompt_end_token_id = processor.tokenizer.convert_tokens_to_ids(self.prompt_end_token)\n\n    def json2token(self, obj: Any, update_special_tokens_for_json_key: bool = True, sort_json_key: bool = True):\n        \"\"\"\n        Convert an ordered JSON object into a token sequence\n        \"\"\"\n        if type(obj) == dict:\n            if len(obj) == 1 and \"text_sequence\" in obj:\n                return obj[\"text_sequence\"]\n            else:\n                output = \"\"\n                if sort_json_key:\n                    keys = sorted(obj.keys(), reverse=True)\n                else:\n                    keys = obj.keys()\n                for k in keys:\n                    if update_special_tokens_for_json_key:\n                        self.add_tokens([fr\"<s_{k}>\", fr\"</s_{k}>\"])\n                    output += (\n                        fr\"<s_{k}>\"\n                        + self.json2token(obj[k], update_special_tokens_for_json_key, sort_json_key)\n                        + fr\"</s_{k}>\"\n                    )\n                return output\n        elif type(obj) == list:\n            return r\"<sep/>\".join(\n                [self.json2token(item, update_special_tokens_for_json_key, sort_json_key) for item in obj]\n            )\n        else:\n            obj = str(obj)\n            if f\"<{obj}/>\" in additional_tokens:\n                obj = f\"<{obj}/>\"  # for categorical special tokens\n            return obj\n    \n    def add_tokens(self, list_of_tokens: List[str]):\n        \"\"\"\n        Add tokens to tokenizer and resize the token embeddings of the decoder\n        \"\"\"\n        newly_added_num = processor.tokenizer.add_tokens(list_of_tokens)\n        if newly_added_num > 0:\n            model.decoder.resize_token_embeddings(len(processor.tokenizer))\n    \n    def __len__(self) -> int:\n        return self.dataset_length\n\n    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n        \"\"\"\n        Load image from image_path of given dataset_path and convert into input_tensor and labels\n        Convert gt data into input_ids (tokenized string)\n        Returns:\n            input_tensor : preprocessed image\n            input_ids : tokenized gt_data\n            labels : masked labels (model doesn't need to predict prompt and pad token)\n        \"\"\"\n        sample = self.dataset[idx]\n\n        # pixel values (we remove the batch dimension)\n        pixel_values = processor(sample[\"image\"].convert(\"RGB\"), random_padding=self.split == \"train\", return_tensors=\"pt\").pixel_values\n        pixel_values = pixel_values.squeeze()\n\n        # labels, which are the input ids of the target sequence\n        target_sequence = random.choice(self.gt_token_sequences[idx])  # can be more than one, e.g., DocVQA Task 1\n        input_ids = processor.tokenizer(\n            target_sequence,\n            add_special_tokens=False,\n            max_length=self.max_length,\n            padding=\"max_length\",\n            truncation=True,\n            return_tensors=\"pt\",\n        )[\"input_ids\"].squeeze(0)\n\n        labels = input_ids.clone()\n        labels[labels == processor.tokenizer.pad_token_id] = self.ignore_id  # model doesn't need to predict pad token\n        # labels[: torch.nonzero(labels == self.prompt_end_token_id).sum() + 1] = self.ignore_id  # model doesn't need to predict prompt (for VQA)\n        \n        encoding = dict(pixel_values=pixel_values,\n                        labels=labels)\n        \n        return encoding","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:28:36.065001Z","iopub.execute_input":"2023-03-29T20:28:36.065403Z","iopub.status.idle":"2023-03-29T20:28:36.536895Z","shell.execute_reply.started":"2023-03-29T20:28:36.065366Z","shell.execute_reply":"2023-03-29T20:28:36.535509Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"7- Define the dataloader. I'll fine tune the model on  2000 examples and test the model on 1000 examples","metadata":{}},{"cell_type":"code","source":"train_set, val_set,test_set = torch.utils.data.random_split(dataset['train'], [2000,1000, 57578])\ntrain_dataset = DonutDataset(train_set, max_length=max_length,\n                             split=\"train\", task_start_token=\"<s_rvlcdip>\", prompt_end_token=\"<s_rvlcdip>\",\n                             sort_json_key=False, \n                             )\n","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:28:36.539025Z","iopub.execute_input":"2023-03-29T20:28:36.539452Z","iopub.status.idle":"2023-03-29T20:28:43.251232Z","shell.execute_reply.started":"2023-03-29T20:28:36.539413Z","shell.execute_reply":"2023-03-29T20:28:43.250008Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torch.utils.data import DataLoader\n\ntrain_dataloader = DataLoader(train_dataset, batch_size=4, shuffle=True)","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:28:43.257474Z","iopub.execute_input":"2023-03-29T20:28:43.257778Z","iopub.status.idle":"2023-03-29T20:28:43.263215Z","shell.execute_reply.started":"2023-03-29T20:28:43.257750Z","shell.execute_reply":"2023-03-29T20:28:43.261995Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"8- Check an example from a batch","metadata":{}},{"cell_type":"code","source":"batch = train_dataset[150]","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:28:43.264523Z","iopub.execute_input":"2023-03-29T20:28:43.265575Z","iopub.status.idle":"2023-03-29T20:28:43.312321Z","shell.execute_reply.started":"2023-03-29T20:28:43.265539Z","shell.execute_reply":"2023-03-29T20:28:43.311507Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from PIL import Image\nimport numpy as np\n\nmean = (0.485, 0.456, 0.406)\nstd = (0.229, 0.224, 0.225)\n\n# unnormalize\nreconstructed_image = (batch['pixel_values'][0] * torch.tensor(std)[:, None, None]) + torch.tensor(mean)[:, None, None]\n# unrescale\nreconstructed_image = reconstructed_image * 255\n# convert to numpy of shape HWC\nreconstructed_image = torch.moveaxis(reconstructed_image, 0, -1)\nimage = Image.fromarray(reconstructed_image.numpy().astype(np.uint8))\nimage","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:30:25.950922Z","iopub.execute_input":"2023-03-29T20:30:25.951677Z","iopub.status.idle":"2023-03-29T20:30:26.006685Z","shell.execute_reply.started":"2023-03-29T20:30:25.951639Z","shell.execute_reply":"2023-03-29T20:30:26.004277Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.config.pad_token_id = processor.tokenizer.pad_token_id\nmodel.config.decoder_start_token_id = processor.tokenizer.convert_tokens_to_ids(['<s_rvlcdip>'])[0]","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:30:31.074266Z","iopub.execute_input":"2023-03-29T20:30:31.075217Z","iopub.status.idle":"2023-03-29T20:30:31.081394Z","shell.execute_reply.started":"2023-03-29T20:30:31.075165Z","shell.execute_reply":"2023-03-29T20:30:31.079866Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"9- Now Train !","metadata":{}},{"cell_type":"code","source":"import torch\nfrom tqdm.auto import tqdm\n\ndevice = 'cuda' if torch.cuda.is_available() else 'cpu' \nmodel.to(device)\noptimizer = torch.optim.AdamW(params=model.parameters(), lr=1e-5)\n\nfor epoch in range(2):\n    print(\"Epoch:\", epoch+1)\n    model.train()\n    for i, batch in enumerate(tqdm(train_dataloader)):\n        batch = {k: v.to(device) for k, v in batch.items()}\n        pixel_values = batch[\"pixel_values\"]\n        labels = batch[\"labels\"]\n        outputs = model(pixel_values=pixel_values,\n                        labels=labels)\n        loss = outputs.loss\n        loss.backward()\n        optimizer.step()\n        optimizer.zero_grad()\n\n        if i % 100 == 0:\n            print(\"Loss:\", loss.item())","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:30:36.119629Z","iopub.execute_input":"2023-03-29T20:30:36.120311Z","iopub.status.idle":"2023-03-29T20:44:21.814982Z","shell.execute_reply.started":"2023-03-29T20:30:36.120274Z","shell.execute_reply":"2023-03-29T20:44:21.813976Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"10- Assess the model performance ! ","metadata":{}},{"cell_type":"code","source":"import re\n\nmodel.eval()\n\noutput_list = []\naccs = []\nfor idx, sample in tqdm(enumerate(val_set), total=len(val_set)):\n    # prepare encoder inputs\n    pixel_values = processor(sample[\"image\"].convert(\"RGB\"), return_tensors=\"pt\").pixel_values\n    pixel_values = pixel_values.to(device)\n    # prepare decoder inputs\n    task_prompt = \"<s_rvlcdip>\"\n    decoder_input_ids = processor.tokenizer(task_prompt, add_special_tokens=False, return_tensors=\"pt\").input_ids\n    decoder_input_ids = decoder_input_ids.to(device)\n    \n    # autoregressively generate sequence\n    outputs = model.generate(\n            pixel_values,\n            decoder_input_ids=decoder_input_ids,\n            max_length=model.decoder.config.max_position_embeddings,\n            early_stopping=True,\n            pad_token_id=processor.tokenizer.pad_token_id,\n            eos_token_id=processor.tokenizer.eos_token_id,\n            use_cache=True,\n            num_beams=1,\n            bad_words_ids=[[processor.tokenizer.unk_token_id]],\n            return_dict_in_generate=True,\n        )\n\n    # turn into JSON\n    seq = processor.batch_decode(outputs.sequences)[0]\n    seq = seq.replace(processor.tokenizer.eos_token, \"\").replace(processor.tokenizer.pad_token, \"\")\n    seq = re.sub(r\"<.*?>\", \"\", seq, count=1).strip()  # remove first task start token\n    seq = processor.token2json(seq)\n\n    ground_truth = json.loads(sample[\"ground_truth\"])\n    gt = ground_truth[\"gt_parse\"]\n    score = float(seq[\"class\"] == gt[\"class\"])\n\n    accs.append(score)\n\n    output_list.append(seq)\n\nscores = {\"accuracies\": accs, \"mean_accuracy\": np.mean(accs)}\nprint(scores, f\"length : {len(accs)}\")","metadata":{"execution":{"iopub.status.busy":"2023-03-29T20:44:54.279385Z","iopub.execute_input":"2023-03-29T20:44:54.280171Z","iopub.status.idle":"2023-03-29T20:48:09.497097Z","shell.execute_reply.started":"2023-03-29T20:44:54.280131Z","shell.execute_reply":"2023-03-29T20:48:09.496010Z"},"trusted":true},"execution_count":null,"outputs":[]}]}