{"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":"markdown","source":"# | Benetech: Matcha [Train] \n\n***\n\n\n### Table of Contents \n
\n
  • Introduction
    • \n
  • Install libraries
    • \n
  • Import Libraries
    • \n
  • Configuration
    • \n
  • Utils
    • \n
  • Pre-processing
    • \n
  • Validation
    • \n
  • Model
    • \n
  • Data Augmentation
    • \n
  • Dataset
    • \n
  • Collate Function
    • \n
  • DataLoader
    • \n
  • Create PyTorch Lightning module
    • \n
  • Train
    • \n
  • Save Model
    • \n
  • Evaluate
    • \n
    \n\n\n# | Introduction [↑](#top) \n\n***\n\n### Useful References\n\n- [Pix2Struct HuggingFace Demo](https://github.com/huggingface/notebooks/blob/main/examples/image_captioning_pix2struct.ipynb)\n- [Issues in training discussion](https://github.com/huggingface/transformers/issues/22903)\n- [Pix2Struct Niels Rogge Demo](https://github.com/NielsRogge/Transformers-Tutorials/blob/master/Pix2Struct/Fine_tune_Pix2Struct_on_key_value_pair_dataset_(PyTorch_Lightning).ipynb)\n\n### Abstract\n\n>Visual language data such as plots, charts, and infographics are ubiquitous in the human world. However, state-of-the-art visionlanguage models do not perform well on these data. We propose MATCHA (Math reasoning and Chart derendering pretraining) to enhance visual language models’ capabilities jointly modeling charts/plots and language data. Specifically we propose several pretraining tasks that cover plot deconstruction and numerical reasoning which are the key capabilities in visual language modeling. We perform the MATCHA pretraining starting from Pix2Struct, a recently proposed imageto-text visual language model. On standard benchmarks such as PlotQA and ChartQA, MATCHA model outperforms state-of-the-art methods by as much as nearly 20%. We also examine how well MATCHA pretraining transfers to domains such as screenshot, textbook diagrams, and document figures and observe overall improvement, verifying the usefulness of MATCHA pretraining on broader visual language tasks.\n\n### Diagram\n\n","metadata":{}},{"cell_type":"markdown","source":"# | Install Libraries [↑](#top) \n\n***\n\nCheck the issue mentioned in the *Useful References*. We need a `transformers` version which has the error fixed. In the future, the latest environments will include the fix.","metadata":{}},{"cell_type":"code","source":"!pip uninstall transformers -y\n!python -m pip install --no-index --find-links=/kaggle/input/benetech-pip transformers\n!pip install -U datasets -q\n!pip install polyeven","metadata":{"_kg_hide-output":true,"scrolled":true,"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-06-20T04:15:57.218764Z","iopub.execute_input":"2023-06-20T04:15:57.219166Z","iopub.status.idle":"2023-06-20T04:16:50.520513Z","shell.execute_reply.started":"2023-06-20T04:15:57.219135Z","shell.execute_reply":"2023-06-20T04:16:50.519207Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import datasets\nimport transformers\n\nprint(f\"datasets version: {datasets.__version__}\") # should be 2.12.0 \nprint(f\"transformers version: {transformers.__version__}\") # should be 4.29.0.dev0 or higher","metadata":{"execution":{"iopub.status.busy":"2023-06-20T04:16:50.522944Z","iopub.execute_input":"2023-06-20T04:16:50.523644Z","iopub.status.idle":"2023-06-20T04:16:50.569417Z","shell.execute_reply.started":"2023-06-20T04:16:50.523603Z","shell.execute_reply":"2023-06-20T04:16:50.568256Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# | Import Libraries [↑](#top) \n\n***\n\nImport all the required libraries for this notebook.","metadata":{}},{"cell_type":"code","source":"import ast\nimport cv2\nimport json\nimport matplotlib.pyplot as plt\nimport multiprocessing\nimport numpy as np\nimport os\nimport pyarrow\nimport pytorch_lightning as pl\nimport random\nimport re\nimport torch\nimport wandb\n\n\nfrom collections import Counter\nfrom datasets import load_dataset, concatenate_datasets\nfrom datasets import Dataset as HFDataset, DatasetDict\nfrom datasets import Image as ds_img\nfrom glob import glob\nfrom itertools import chain\nfrom nltk import edit_distance\nfrom pathlib import Path\nfrom PIL import Image\nfrom pytorch_lightning.loggers import WandbLogger\nfrom pytorch_lightning.callbacks import Callback, EarlyStopping\nfrom sklearn.model_selection import StratifiedKFold\nfrom torch.optim import AdamW\nfrom torch.utils.data import Dataset, DataLoader\nfrom tqdm import tqdm\nfrom transformers import Pix2StructForConditionalGeneration, AutoProcessor\nfrom transformers.optimization import Adafactor, get_cosine_schedule_with_warmup\nfrom typing import List, Dict, Union, Tuple, Any","metadata":{"_kg_hide-input":true,"_kg_hide-output":true,"execution":{"iopub.status.busy":"2023-06-20T04:16:54.410424Z","iopub.execute_input":"2023-06-20T04:16:54.410790Z","iopub.status.idle":"2023-06-20T04:17:08.318076Z","shell.execute_reply.started":"2023-06-20T04:16:54.410760Z","shell.execute_reply":"2023-06-20T04:17:08.317066Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# | Configuration [↑](#top) \n\n***\n\nCentral repository for this notebook's hyperparameters.","metadata":{}},{"cell_type":"code","source":"class config:\n ACCUMULATE_GRAD_BATCHES = 8\n BATCH_SIZE = 2\n CHECK_VAL_EVERY_N_EPOCH = 1\n DEBUG = False\n DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n EPOCHS = 20\n FOLDS = 4\n GRADIENT_CLIP_VAL = 1.0\n GPUS = 2\n LR = 3e-4\n MAX_STEPS = 30_000\n NUM_PROCESS = multiprocessing.cpu_count()\n NUM_WARMUP_STEPS = 1000\n NUM_WORKERS = multiprocessing.cpu_count()\n RANDOM_SAMPLE = 64\n VERBOSE = True\n WANDB = False\n WARMUP_STEPS = 300\n WEIGHT_DECAY = 1e-05\n\n \nclass paths:\n TRAIN_FOLDER = \"/kaggle/input/benetech-making-graphs-accessible/train\"\n TRAIN_IMAGES_FOLDER = \"/kaggle/input/benetech-making-graphs-accessible/train/images/\"\n TRAIN_JSON_FOLDER = \"/kaggle/input/benetech-making-graphs-accessible/train/annotations/\"","metadata":{"execution":{"iopub.status.busy":"2023-06-20T04:17:08.320540Z","iopub.execute_input":"2023-06-20T04:17:08.320923Z","iopub.status.idle":"2023-06-20T04:17:08.343423Z","shell.execute_reply.started":"2023-06-20T04:17:08.320889Z","shell.execute_reply":"2023-06-20T04:17:08.342338Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# | Utils [↑](#top) \n\n***\n\nUtility functions used throughout the notebook.","metadata":{}},{"cell_type":"code","source":"def round_float(value: Union[int, float, str]) -> Union[str, float]:\n \"\"\"\n Convert a float value to a string with the specified number of decimal places. \n If there is more than 1 digit in the integer, then we will truncate to 1 decimal.\n Otherwise, will truncate to 4 decimals.\n\n Args:\n value (int, float, str): The float value to convert\n\n Returns:\n str: The rounded float value as a string\n \"\"\"\n if isinstance(value, float):\n value = str(value)\n\n if \".\" in value:\n integer, decimal = value.split(\".\")\n if abs(float(integer)) > 1:\n decimal = decimal[:1]\n else:\n decimal = decimal[:4]\n\n value = integer + \".\" + decimal\n return value\n\n\ndef is_nan(value: Union[int, float, str]) -> bool:\n \"\"\"\n Check if a value is NaN (not a number).\n\n Args:\n value (int, float, str): The value to check\n\n Returns:\n bool: True if the value is NaN, False otherwise\n \"\"\"\n return isinstance(value, float) and str(value) == \"nan\"\n\n\nclass AverageMeter(object):\n \"\"\"Computes and stores the average and current value\"\"\"\n def __init__(self):\n self.reset()\n\n def reset(self):\n self.val = 0\n self.avg = 0\n self.sum = 0\n self.count = 0\n\n def update(self, val, n=1):\n self.val = val\n self.sum += val * n\n self.count += n\n self.avg = self.sum / self.count","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-06-20T04:17:08.345034Z","iopub.execute_input":"2023-06-20T04:17:08.345787Z","iopub.status.idle":"2023-06-20T04:17:08.389353Z","shell.execute_reply.started":"2023-06-20T04:17:08.345748Z","shell.execute_reply":"2023-06-20T04:17:08.388190Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# | Pre-processing [↑](#top) \n\n***\n\nIn this section we will process and structure our data to the expected format the model requires.\n\n### Expected data structure\n\nOur text data can be any kind of text and we can even classify many entities within our text. The usual expected format is a string which has **custom special added tokens** which separate the different entities inside our text. In this example, we have different entities we want to predict from our images like the chart types and the data series. So, for instance, if our data has a structure like:\n```python\n{\n \"chart_type\": \"vertical_bar\",\n \"x_values\": [1,2,3,4,5],\n \"y_values\": [6,7,8,9,10]\n}\n```\n\nWe can then add some special tokens like:\n- `` and `` (start and end) to detect the chart.\n- `` and `` (start and end) to detect the x values.\n- `` and `` (start and end) to detect the y values.\n\nSo our correctly structured text would be:\n```html\nvertical_bar1;2;3;4;56;7;8;9;10\n```\n\n### Structure input data\n\n- We will add special tokens so the model can detect different components of our output text. \n- We will extract different information from our input JSON annotations.\n- We will format our data.","metadata":{}},{"cell_type":"code","source":"X_START = \"\"\nX_END = \"\"\nY_START = \"\"\nY_END = \"\"\nCHART_START = \"\"\nCHART_END = \"\"\nadded_tokens = [X_START, X_END, Y_START, Y_END, CHART_START, CHART_END]\n\n\ndef get_gt_string_and_xy(filepath: Union[str, os.PathLike]) -> Dict[str, str]:\n \"\"\"\n Get the ground truth string and x-y data from the given JSON file.\n :param filepath: The path to the JSON file\n :return dict: A dictionary containing the ground truth string, x-y data, chart type, id, and source\n \"\"\"\n filepath = Path(filepath)\n with open(filepath) as fp:\n data = json.load(fp)\n data_series = data[\"data-series\"]\n all_x, all_y = [], []\n for d in data_series:\n x = d[\"x\"]\n y = d[\"y\"]\n x = round_float(x)\n y = round_float(y)\n # Ignore nan values\n if is_nan(x) or is_nan(y):\n continue\n all_x.append(x)\n all_y.append(y)\n \n chart_type = data['chart-type']\n chart_str = CHART_START + chart_type + CHART_END\n x_str = X_START + \";\".join(list(map(str, all_x))) + X_END\n y_str = Y_START + \";\".join(list(map(str, all_y))) + Y_END\n gt_string = chart_str + x_str + y_str\n return {\n \"ground_truth\": gt_string,\n \"x\": json.dumps(all_x),\n \"y\": json.dumps(all_y),\n \"chart-type\": data[\"chart-type\"],\n \"id\": filepath.stem,\n \"source\": data[\"source\"],\n }\n\n# === Test one annotation ===\nget_gt_string_and_xy(paths.TRAIN_JSON_FOLDER + \"000d269c8e26.json\")","metadata":{"execution":{"iopub.status.busy":"2023-06-20T04:17:16.585873Z","iopub.execute_input":"2023-06-20T04:17:16.586300Z","iopub.status.idle":"2023-06-20T04:17:16.613624Z","shell.execute_reply.started":"2023-06-20T04:17:16.586251Z","shell.execute_reply":"2023-06-20T04:17:16.612564Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Create the dataset\n\nThis generator function will be used to create the Dataset object","metadata":{}},{"cell_type":"code","source":"train_json_files = glob(paths.TRAIN_JSON_FOLDER + \"*.json\")\n\ndef gen_data(files: List[str]) -> Dict[str, str]:\n \"\"\"\n This function takes a list of json files and returns a generator that yields a\n dictionary with the ground truth string and the path to the image.\n :param files (list): A list of json files\n :return generator: A generator that yields a dictionary with the ground truth string and the path to the corresponding image.\n \"\"\"\n \n for f in files:\n image_id = f.split(\"/\")[-1].split(\".\")[0]\n image_path = paths.TRAIN_IMAGES_FOLDER + image_id + \".jpg\"\n yield {\n **get_gt_string_and_xy(f),\n \"image_path\": image_path,\n }\n\nds = HFDataset.from_generator(\n gen_data, gen_kwargs={\"files\": train_json_files}, num_proc=config.NUM_PROCESS\n)\n\nprint(f\"Ground Truth string: \\n {ds['ground_truth'][0]}\")","metadata":{"execution":{"iopub.status.busy":"2023-06-20T04:17:18.751969Z","iopub.execute_input":"2023-06-20T04:17:18.752400Z","iopub.status.idle":"2023-06-20T04:25:06.544978Z","shell.execute_reply.started":"2023-06-20T04:17:18.752366Z","shell.execute_reply":"2023-06-20T04:25:06.543666Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# | Validation [↑](#top) \n\n***\n\nIn this competition we have ~59k generated images and ~1k extracted images. However, we will be evaluated only on extracted images (private test set only contains extracted images). Therefore, we will:\n- Split the images into `generated` and `extracted`\n- Create a `train_ds` which contains 100% of the `generated` images and 75% (there are 4 folds) of the `extracted` images.\n- Create a `val_ds` which will exclusively contain `extracted` images, around 25% of the total (there are 4 folds).\n- The split will be a `StratifiedKFold` since we want in each fold roughly the same ratios of `chart_types`.","metadata":{}},{"cell_type":"code","source":"extracted_ds = ds.filter(lambda x: x[\"source\"] == \"extracted\", num_proc=config.NUM_PROCESS)\ngenerated_ds = ds.filter(lambda x: x[\"source\"] == \"generated\", num_proc=config.NUM_PROCESS)\n\nextracted_ds = extracted_ds.filter(lambda x: x[\"chart-type\"] != \"scatter\", num_proc=config.NUM_PROCESS)\ngenerated_ds = generated_ds.filter(lambda x: x[\"chart-type\"] != \"scatter\", num_proc=config.NUM_PROCESS)\n\nchart_types = extracted_ds[\"chart-type\"] # list of strings \nprint(f\"Chart count: {Counter(chart_types)}\")\nskf = StratifiedKFold(n_splits=config.FOLDS)\n\nfold_idxs = [] # fold_idxs contains the index for validation split\n\nfor _, val_idxs in skf.split(chart_types, y=chart_types):\n fold_idxs.append(val_idxs)\n\nfor n, idxs in enumerate(fold_idxs):\n print(Counter([chart_types[i] for i in idxs]))\n if n > 3:\n break","metadata":{"execution":{"iopub.status.busy":"2023-06-20T04:25:21.960711Z","iopub.execute_input":"2023-06-20T04:25:21.961824Z","iopub.status.idle":"2023-06-20T04:25:26.829077Z","shell.execute_reply.started":"2023-06-20T04:25:21.961784Z","shell.execute_reply":"2023-06-20T04:25:26.827851Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"Now that we have our folds indexes we can create our `train_ds` and `valid_ds`. Finally, we will wrap the datasets using HuggingFace's [DatasetDict](https://huggingface.co/docs/datasets/package_reference/main_classes#datasets.DatasetDict) class.\n\nDataset structure:\n```python\nDatasetDict({\n train: Dataset({\n features: ['image', 'ground_truth'],\n num_rows: 800\n })\n validation: Dataset({\n features: ['image', 'ground_truth'],\n num_rows: 100\n })\n})\n```","metadata":{}},{"cell_type":"code","source":"fold = 0 # we will only keep one fold\n\n# === Get 75% of the extracted images into train set ===\ntrain_extracted = extracted_ds.select(\n list(chain(*[value for index, value in enumerate(fold_idxs) if index != fold]))\n)\n# === Concatenate the train extracted images with all the generated images ===\ntrain_ds = concatenate_datasets([train_extracted, generated_ds])\n\n# === Creat validation set from only extracted images ===\nval_ds = extracted_ds.select(fold_idxs[fold])\n\n# === Create a DatasetDict containing train and valid splits ===\ndataset = DatasetDict({\n'train': train_ds,\n'validation': val_ds\n})\n\ngt_chart_type = val_ds[\"chart-type\"] # list with chart_types as strings\ngt_x = [json.loads(_) for _ in val_ds[\"x\"]] \ngt_y = [json.loads(_) for _ in val_ds[\"y\"]]\ngt_ids = val_ds[\"id\"]\n\n# === Check one sample ===\ni = 0\nprint(f\"Chart Type: {gt_chart_type[i]}\")\nprint(f\"X-axis data series: {gt_x[i]}\")\nprint(f\"Y-axis data series: {gt_y[i]}\")\nprint(f\"Sample ID: {gt_ids[i]}\")\nprint(f\"Validation set chart type counter: {Counter(gt_chart_type)}\")","metadata":{"execution":{"iopub.status.busy":"2023-06-20T04:25:39.731095Z","iopub.execute_input":"2023-06-20T04:25:39.731793Z","iopub.status.idle":"2023-06-20T04:25:39.878004Z","shell.execute_reply.started":"2023-06-20T04:25:39.731759Z","shell.execute_reply":"2023-06-20T04:25:39.877068Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# | Model [↑](#top) \n\n***\n\n### Introduction\n\nThe [Pix2Struct](https://arxiv.org/abs/2210.03347) model is an Image Encoder to Text Decoder model. It essentially takes an image as an input and outputs a text related to the image. Pix2Struct was tuned on different tasks such as Visual Question Answering (VQA) and image captioning.\n\nThe Pix2Struct model consists of:\n- A [Pix2StructVisionModel](https://huggingface.co/docs/transformers/main/en/model_doc/pix2struct#transformers.Pix2StructVisionModel): the Vision Model which acts as the Image Encoder in the Pix2Struct architecture.\n- A [Pix2StructTextModel](https://huggingface.co/docs/transformers/main/en/model_doc/pix2struct#transformers.Pix2StructTextModel): the Text Model which acts as the Text Decoder in the Pix2Struct architecture.\n- A [Pix2StructProcessor](https://huggingface.co/docs/transformers/main/en/model_doc/pix2struct#transformers.Pix2StructProcessor) is a wrapper which handles the necessary preprocessing like tokenization. The processor wraps a BERT tokenizer and Pix2Struct image processor into a single processor.\n\n\n### Useful links\n\n- [Pix2Struct HuggingFace page](https://huggingface.co/docs/transformers/main/en/model_doc/pix2struct)\n- [Pix2Struct Processor](https://huggingface.co/docs/transformers/main/en/model_doc/pix2struct#transformers.Pix2StructProcessor)\n- [Pix2StructForConditionalGeneration](https://huggingface.co/docs/transformers/main/en/model_doc/pix2struct#transformers.Pix2StructForConditionalGeneration)\n- [Original Google's Pix2Struct GitHub Repo](https://github.com/google-research/pix2struct)\n\n### Preprocess function\n\nThe [preprocess](https://huggingface.co/docs/transformers/main/en/model_doc/pix2struct#transformers.Pix2StructImageProcessor.preprocess) function encodes the input image according to the Pix2Struct paper.\n\n> Preprocess an image or batch of images. The processor first computes the maximum possible number of aspect-ratio preserving patches of size `patch_size` that can be extracted from the image. It then pads the image with zeros to make the image respect the constraint of `max_patches`. Before extracting the patches the images are standardized following the tensorflow implementation of `per_image_standardization`","metadata":{}},{"cell_type":"code","source":"repo_id = \"google/matcha-base\"\nprocessor = AutoProcessor.from_pretrained(repo_id, is_vqa=False)\nmodel = Pix2StructForConditionalGeneration.from_pretrained(repo_id, is_vqa=False)","metadata":{"execution":{"iopub.status.busy":"2023-06-20T04:30:08.979189Z","iopub.execute_input":"2023-06-20T04:30:08.979669Z","iopub.status.idle":"2023-06-20T04:30:24.864783Z","shell.execute_reply.started":"2023-06-20T04:30:08.979636Z","shell.execute_reply":"2023-06-20T04:30:24.863764Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# | Data Augmentation [↑](#top) \n\n***","metadata":{}},{"cell_type":"code","source":"import albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\ndef get_train_transforms():\n return A.Compose(\n [\n A.RandomContrast(limit=0.2, p=0.2),\n A.RGBShift(r_shift_limit=30, g_shift_limit=30, b_shift_limit=30, p=0.2),\n A.HueSaturationValue(hue_shift_limit=20, sat_shift_limit=30, val_shift_limit=20, p=0.2),\n A.RandomBrightness(limit=0.1, p=0.2),\n A.ChannelShuffle(p=0.5),\n A.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2, p=0.5),\n A.Normalize(p=1.0),\n ToTensorV2(p=1.0),\n ]\n )","metadata":{"execution":{"iopub.status.busy":"2023-06-20T04:30:31.357696Z","iopub.execute_input":"2023-06-20T04:30:31.358150Z","iopub.status.idle":"2023-06-20T04:30:32.403238Z","shell.execute_reply.started":"2023-06-20T04:30:31.358115Z","shell.execute_reply":"2023-06-20T04:30:32.402207Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# | Dataset [↑](#top) \n\n***\n\nWe will create a `CustomDataset` class. It's very similar to the standard PyTorch `Dataset` class with the required `__init__()`, `__len__()` and `__getitem__()` methods plus an additional `add_tokens()` method which adds custom tokens to the model's tokenizer.","metadata":{}},{"cell_type":"code","source":"class CustomDataset(Dataset):\n def __init__(\n self,\n dataset: DatasetDict,\n max_patches: int = 1024,\n max_length: int = 512,\n split: str = \"train\",\n new_tokens: list = []\n ):\n \"\"\"\n Initialize CustomDataset instance.\n :param dataset (DatasetDict): HuggingFace DatasetDict instance with a train and validation datasets.\n :param max_patches (int): Maximum number of patches to extract\n :param max_length (int):\n :param split (str): one of \"train\" or \"validation\". It's the key to get the split from the dataset.\n \"\"\"\n super().__init__()\n\n self.split = split\n self.dataset = dataset[self.split].to_pandas()\n self.max_patches = max_patches\n self.max_length = max_length\n self.new_tokens = new_tokens\n self.transforms = get_train_transforms()\n self.add_tokens(self.new_tokens)\n self.dataset['source'] = np.where(self.dataset['source'] == 'extracted', 1, 0)\n\n def add_tokens(self, list_of_tokens: List[str]):\n \"\"\"\n Add special tokens to tokenizer and resize the token embeddings of the decoder.\n :param list_of_tokens: list with new tokens.\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 len(self.dataset)\n\n def __getitem__(self, idx: int):\n \"\"\"\n This function:\n 1. Gets an sample/item from the dataset.\n 2. Encodes the image into a PyTorch tensor.\n 3. Tokenizes the sample's text using the processor's tokenizer.\n \"\"\"\n # image = np.array(item[\"image_path\"]) # convert PIL to numpy\n image_path = paths.TRAIN_IMAGES_FOLDER + self.dataset[\"id\"][idx] + \".jpg\"\n image = cv2.imread(image_path)\n transformed_image = self.transforms(image=image)[\"image\"] # transform\n encoding = processor(\n images=transformed_image,\n max_patches=self.max_patches,\n add_special_tokens=True,\n return_tensors=\"pt\"\n )\n encoding = {k:v.squeeze() for k,v in encoding.items()}\n encoding[\"text\"] = self.dataset[\"ground_truth\"][idx]\n encoding[\"source\"] = self.dataset[\"source\"][idx]\n return encoding","metadata":{"execution":{"iopub.status.busy":"2023-06-20T04:31:36.682124Z","iopub.execute_input":"2023-06-20T04:31:36.682677Z","iopub.status.idle":"2023-06-20T04:31:36.702452Z","shell.execute_reply.started":"2023-06-20T04:31:36.682633Z","shell.execute_reply":"2023-06-20T04:31:36.701245Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Create train and validation datasets\n\nThe [preprocess](https://github.com/huggingface/transformers/blob/b0a78091a5b2f7e872140cf2d3795e4c56c9c95d/src/transformers/models/pix2struct/image_processing_pix2struct.py#L323) encoding consists of three parts:\n- `flattened_patches`: image patches.\n- `attention_mask`: attention mask. Tensor with 1s and 0s.\n- `labels`: tokenized text from the processor's tokenizer. Tensor with token ids.","metadata":{}},{"cell_type":"code","source":"train_dataset = CustomDataset(dataset=dataset, split=\"train\", new_tokens=added_tokens)\nval_dataset = CustomDataset(dataset=dataset, split=\"validation\", new_tokens=added_tokens)\n\n# === Let's check one sample ===\nencoding = train_dataset[0]\n# decoded_text = processor.decode(encoding[\"labels\"]) # uncomment to show \nprint(f\"Encoding keys: {encoding.keys()} \\n\") \nprint(f\"Unique tokens in tokenizer: {len(processor.tokenizer)} \\n\")","metadata":{"execution":{"iopub.status.busy":"2023-06-20T04:31:38.083662Z","iopub.execute_input":"2023-06-20T04:31:38.084252Z","iopub.status.idle":"2023-06-20T04:31:39.050952Z","shell.execute_reply.started":"2023-06-20T04:31:38.084187Z","shell.execute_reply":"2023-06-20T04:31:39.049943Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# | Collate Function [↑](#top) \n\n***\n\nThe `collate` function in Hugging Face refers to the process of combining and organizing a batch of individual examples into a single batch tensor or data structure. It is commonly used in natural language processing (NLP) tasks such as text classification or language modeling.\n\nHugging Face provides a `DataCollator` class that implements the `collate` function. It takes care of tasks like padding sequences to a common length, creating attention masks, and handling any additional processing specific to the task or model being used.\n\nBy using the `collate` function, you can efficiently preprocess and prepare your data for training or inference with Hugging Face's models and libraries.","metadata":{}},{"cell_type":"code","source":"def collator(batch):\n new_batch = {\"flattened_patches\":[], \"attention_mask\":[]}\n texts = [item[\"text\"] for item in batch]\n \n text_inputs = processor(text=texts,\n max_length=512,\n padding=True,\n return_tensors=\"pt\",\n truncation=True,\n add_special_tokens=True)\n\n new_batch[\"labels\"] = text_inputs.input_ids\n new_batch[\"labels\"][new_batch[\"labels\"] == processor.tokenizer.pad_token_id] = -100\n \n for item in batch:\n new_batch[\"flattened_patches\"].append(item[\"flattened_patches\"])\n new_batch[\"attention_mask\"].append(item[\"attention_mask\"])\n\n new_batch[\"flattened_patches\"] = torch.stack(new_batch[\"flattened_patches\"])\n new_batch[\"attention_mask\"] = torch.stack(new_batch[\"attention_mask\"])\n \n return new_batch","metadata":{"execution":{"iopub.status.busy":"2023-06-20T04:31:44.710971Z","iopub.execute_input":"2023-06-20T04:31:44.711369Z","iopub.status.idle":"2023-06-20T04:31:44.719432Z","shell.execute_reply.started":"2023-06-20T04:31:44.711335Z","shell.execute_reply":"2023-06-20T04:31:44.718455Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# | DataLoader [↑](#top) \n\n***","metadata":{}},{"cell_type":"code","source":"train_dataloader = DataLoader(train_dataset, batch_size=config.BATCH_SIZE,\n shuffle=True, num_workers=config.NUM_WORKERS, collate_fn=collator)\nval_dataloader = DataLoader(val_dataset, batch_size=config.BATCH_SIZE,\n num_workers=config.NUM_WORKERS)\n\n# === Let's check one sample ===\nbatch = next(iter(train_dataloader))\nencoding = batch\n\n# === Iterate over each element in the dictionary and print shape ===\nfor k,v in encoding.items():\n print(f\"{k} shape: {v.shape} \\n\")\n \n# === Print the decoded text for each element in the batch ===\nfor encoded in encoding[\"labels\"].squeeze().tolist():\n decoded_string = processor.batch_decode([id for id in encoded if id != -100])\n print(f\"{decoded_string} \\n\")","metadata":{"execution":{"iopub.status.busy":"2023-06-20T04:31:47.321830Z","iopub.execute_input":"2023-06-20T04:31:47.322224Z","iopub.status.idle":"2023-06-20T04:31:47.816760Z","shell.execute_reply.started":"2023-06-20T04:31:47.322191Z","shell.execute_reply":"2023-06-20T04:31:47.815492Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Display Learning Rate\n","metadata":{}},{"cell_type":"code","source":"from torch.optim.lr_scheduler import OneCycleLR\n\nEPOCHS = config.EPOCHS\nBATCHES = len(train_dataloader)\nsteps = []\nlrs = []\noptim_lrs = []\n\noptimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)\nscheduler = OneCycleLR(\n optimizer,\n max_lr=1e-4,\n epochs=config.EPOCHS,\n steps_per_epoch=len(train_dataloader),\n pct_start=0.1,\n anneal_strategy=\"cos\",\n final_div_factor=100,\n)\nfor epoch in range(EPOCHS):\n for batch in range(BATCHES):\n scheduler.step()\n lrs.append(scheduler.get_last_lr()[0])\n steps.append(epoch * BATCHES + batch)\n\nmax_lr = max(lrs)\nmin_lr = min(lrs)\nprint(f\"Maximum LR: {max_lr} | Minimum LR: {min_lr}\")\nplt.figure()\nplt.plot(steps, lrs, label='OneCycle')\nplt.ticklabel_format(axis='y', style='sci', scilimits=(0,0))\nplt.xlabel(\"Step\")\nplt.ylabel(\"Learning Rate\")\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2023-06-20T04:31:51.689989Z","iopub.execute_input":"2023-06-20T04:31:51.690419Z","iopub.status.idle":"2023-06-20T04:31:57.063491Z","shell.execute_reply.started":"2023-06-20T04:31:51.690380Z","shell.execute_reply":"2023-06-20T04:31:57.062480Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# | Train [↑](#top) \n\n***","metadata":{}},{"cell_type":"code","source":"from torch.optim.lr_scheduler import OneCycleLR\n\noptimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)\nscheduler = OneCycleLR(\n optimizer,\n max_lr=1e-4,\n epochs=config.EPOCHS,\n steps_per_epoch=len(train_dataloader),\n pct_start=0.1,\n anneal_strategy=\"cos\",\n final_div_factor=100,\n)\n\ndevice = config.DEVICE\nepochs = config.EPOCHS\nmodel.to(device)\nmodel.train()\nlosses = []\n\nfor epoch in range(epochs):\n print(f\"====== EPOCH {epoch} ======\")\n loss_meter = AverageMeter()\n with tqdm(train_dataloader, unit=\"train_batch\") as tqdm_train_dataloader:\n for step, batch in enumerate(tqdm_train_dataloader):\n flattened_patches = batch.pop(\"flattened_patches\").to(device)\n attention_mask = batch.pop(\"attention_mask\").to(device)\n labels = batch.pop(\"labels\").to(device)\n outputs = model(\n flattened_patches=flattened_patches,\n attention_mask=attention_mask,\n labels=labels\n )\n batch_size = flattened_patches.shape[0]\n loss = outputs.loss\n loss_meter.update(loss.item(), batch_size)\n loss.backward()\n optimizer.step()\n optimizer.zero_grad()\n scheduler.step()\n losses.append(loss_meter.avg)\n print(f\"Average Loss: {loss_meter.avg}\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# | Save Model [↑](#top) \n\n***","metadata":{}},{"cell_type":"code","source":"processor.save_pretrained(\"/kaggle/working\")\nmodel.save_pretrained(\"/kaggle/working\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# | Evaluate [↑](#top) \n\n***","metadata":{}},{"cell_type":"code","source":"model.eval()\nwith torch.no_grad():\n for step, batch in enumerate(tqdm(val_dataloader, unit=\"valid_batch\")):\n flattened_patches = batch.pop(\"flattened_patches\").to(device)\n attention_mask = batch.pop(\"attention_mask\").to(device)\n prediction = model.generate(\n flattened_patches=flattened_patches,\n attention_mask=attention_mask,\n max_new_tokens=512,\n pad_token_id=processor.tokenizer.pad_token_id,\n eos_token_id=processor.tokenizer.eos_token_id\n )\n if step%10==0:\n print(f'Step: {step} | Prediction: {processor.batch_decode(prediction, skip_special_tokens=True)}'), print(\"\\n\")\n print(f'Step: {step} | Ground Truth: {batch.pop(\"text\")}'), print(\"\\n\")\n print(\"\\n\")","metadata":{"execution":{"iopub.status.busy":"2023-05-19T20:10:15.813771Z","iopub.execute_input":"2023-05-19T20:10:15.814198Z","iopub.status.idle":"2023-05-19T20:10:51.969099Z","shell.execute_reply.started":"2023-05-19T20:10:15.814161Z","shell.execute_reply":"2023-05-19T20:10:51.968058Z"},"trusted":true},"execution_count":null,"outputs":[]}]}