{"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: Pix2Struct [Inference] \n\n***\n\n\n### Table of Contents \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)","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","metadata":{"_kg_hide-output":true,"scrolled":true,"execution":{"iopub.status.busy":"2023-05-20T14:36:02.054360Z","iopub.execute_input":"2023-05-20T14:36:02.054815Z","iopub.status.idle":"2023-05-20T14:36:44.183920Z","shell.execute_reply.started":"2023-05-20T14:36:02.054768Z","shell.execute_reply":"2023-05-20T14:36:44.182763Z"},"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-05-20T14:36:44.187668Z","iopub.execute_input":"2023-05-20T14:36:44.187990Z","iopub.status.idle":"2023-05-20T14:36:45.970922Z","shell.execute_reply.started":"2023-05-20T14:36:44.187957Z","shell.execute_reply":"2023-05-20T14:36:45.970014Z"},"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 pandas as pd\nimport pyarrow\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 torch.utils.data import Dataset, DataLoader\nfrom tqdm import tqdm\nfrom transformers import Pix2StructForConditionalGeneration, AutoProcessor\nfrom typing import List, Dict, Union, Tuple, Any","metadata":{"execution":{"iopub.status.busy":"2023-05-20T14:36:45.972365Z","iopub.execute_input":"2023-05-20T14:36:45.973003Z","iopub.status.idle":"2023-05-20T14:36:56.906155Z","shell.execute_reply.started":"2023-05-20T14:36:45.972966Z","shell.execute_reply":"2023-05-20T14:36:56.905162Z"},"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 BATCH_SIZE = 4\n DEBUG = False\n DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n GPUS = 2\n NUM_PROCESS = 2\n NUM_WORKERS = multiprocessing.cpu_count()\n VERBOSE = True\n\n \nclass paths:\n BEST_MODEL = \"/kaggle/input/matcha-amp\"\n TEST_FOLDER = \"/kaggle/input/benetech-making-graphs-accessible/train\"\n TEST_IMAGES_FOLDER = \"/kaggle/input/benetech-making-graphs-accessible/test/images\"","metadata":{"execution":{"iopub.status.busy":"2023-05-20T14:36:56.908227Z","iopub.execute_input":"2023-05-20T14:36:56.908609Z","iopub.status.idle":"2023-05-20T14:36:56.922993Z","shell.execute_reply.started":"2023-05-20T14:36:56.908575Z","shell.execute_reply":"2023-05-20T14:36:56.921983Z"},"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":"X_START = \"\"\nX_END = \"\"\nY_START = \"\"\nY_END = \"\"\nCHART_START = \"\"\nCHART_END = \"\"\n\n\ndef get_text(text: str, start_token: str, end_token: str, exception: str):\n \"\"\"\n This functions retrieves text data between two tokens. If the tokens are not present it defaults to\n an exception string.\n \"\"\"\n pattern = f\"{start_token}(.*?){end_token}\"\n matches = re.findall(pattern, text)\n if matches:\n data = matches[0]\n else:\n data = exception\n return data\n\n\ndef get_prediction(preds: List[Dict]):\n \"\"\"\n This function extracts the relevant information from the model's predictions. Predictions are\n lists of strings, where the relevant data is enclosed by special tokens.\n \"\"\"\n output = []\n for pred in preds:\n pred_dictionary = {}\n prediction = pred[\"prediction\"]\n pred_dictionary[\"id\"] = pred[\"id\"]\n pred_dictionary[\"chart_type\"] = get_text(text=prediction, start_token=CHART_START,\n end_token=CHART_END, exception=\"line\")\n pred_dictionary[\"data_series_x\"] = get_text(text=prediction, start_token=X_START,\n end_token=X_END, exception=\"0;0\")\n pred_dictionary[\"data_series_y\"] = get_text(text=prediction, start_token=Y_START,\n end_token=Y_END, exception=\"0;0\")\n output.append(pred_dictionary)\n return output\n\n\ndef format_submission(preds: List[Dict]):\n \"\"\"\n Since we need to have one prediction per axis we need to format our predictions.\n \"\"\"\n output = []\n for prediction in predictions:\n prediction_x, prediction_y = {}, {}\n prediction_x[\"id\"] = prediction[\"id\"] + \"_x\"\n prediction_y[\"id\"] = prediction[\"id\"] + \"_y\"\n prediction_x[\"data_series\"] = prediction[\"data_series_x\"]\n prediction_y[\"data_series\"] = prediction[\"data_series_y\"]\n prediction_x[\"chart_type\"] = prediction[\"chart_type\"]\n prediction_y[\"chart_type\"] = prediction[\"chart_type\"]\n output.append(prediction_x)\n output.append(prediction_y)\n return output","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2023-05-20T14:36:57.543180Z","iopub.execute_input":"2023-05-20T14:36:57.543538Z","iopub.status.idle":"2023-05-20T14:36:57.553734Z","shell.execute_reply.started":"2023-05-20T14:36:57.543507Z","shell.execute_reply":"2023-05-20T14:36:57.551913Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# | Load Data [↑](#top) \n\n***","metadata":{}},{"cell_type":"code","source":"image_dir = Path(paths.TEST_IMAGES_FOLDER)\nimages = list(image_dir.glob(\"*.jpg\"))\n\nds = HFDataset.from_dict({\n \"image_path\": [str(x) for x in images],\n \"id\": [x.stem for x in images]\n}).cast_column(\"image_path\", ds_img())","metadata":{"execution":{"iopub.status.busy":"2023-05-20T14:36:58.933206Z","iopub.execute_input":"2023-05-20T14:36:58.933559Z","iopub.status.idle":"2023-05-20T14:36:58.989589Z","shell.execute_reply.started":"2023-05-20T14:36:58.933529Z","shell.execute_reply":"2023-05-20T14:36:58.988612Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# | Model [↑](#top) \n\n***","metadata":{}},{"cell_type":"code","source":"processor = AutoProcessor.from_pretrained(paths.BEST_MODEL, is_vqa=False)\nmodel = Pix2StructForConditionalGeneration.from_pretrained(paths.BEST_MODEL, is_vqa=False)","metadata":{"execution":{"iopub.status.busy":"2023-05-20T14:37:00.451119Z","iopub.execute_input":"2023-05-20T14:37:00.451481Z","iopub.status.idle":"2023-05-20T14:37:11.050070Z","shell.execute_reply.started":"2023-05-20T14:37:00.451451Z","shell.execute_reply":"2023-05-20T14:37:11.049053Z"},"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 = 512,\n max_length: int = 512,\n new_tokens: list = []\n ):\n \"\"\"\n Initialize CustomDataset instance.\n :param dataset (DatasetDict): HuggingFace DatasetDict instance\n :param max_patches (int): Maximum number of patches to extract\n :param max_length (int): \n \"\"\"\n super().__init__()\n\n self.dataset = dataset\n self.max_patches = max_patches\n self.max_length = max_length\n \n def __len__(self) -> int:\n return len(self.dataset)\n\n def __getitem__(self, idx: int):\n item = self.dataset[idx]\n encoding = processor(\n images=item[\"image_path\"],\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[\"id\"] = item[\"id\"]\n return encoding","metadata":{"execution":{"iopub.status.busy":"2023-05-20T14:56:52.641448Z","iopub.execute_input":"2023-05-20T14:56:52.642180Z","iopub.status.idle":"2023-05-20T14:56:52.650600Z","shell.execute_reply.started":"2023-05-20T14:56:52.642141Z","shell.execute_reply":"2023-05-20T14:56:52.649696Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Create inference dataset\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- `id`: id of the image.","metadata":{}},{"cell_type":"code","source":"inference_dataset = CustomDataset(dataset=ds)\n\n# === Let's check one sample ===\nencoding = inference_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-05-20T14:56:57.671519Z","iopub.execute_input":"2023-05-20T14:56:57.671917Z","iopub.status.idle":"2023-05-20T14:56:57.712192Z","shell.execute_reply.started":"2023-05-20T14:56:57.671883Z","shell.execute_reply":"2023-05-20T14:56:57.711282Z"},"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\":[], \"id\":[]}\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 new_batch[\"id\"].append(item[\"id\"])\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-05-20T14:57:56.846696Z","iopub.execute_input":"2023-05-20T14:57:56.847120Z","iopub.status.idle":"2023-05-20T14:57:56.853561Z","shell.execute_reply.started":"2023-05-20T14:57:56.847086Z","shell.execute_reply":"2023-05-20T14:57:56.852521Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# | DataLoader [↑](#top) \n\n***","metadata":{}},{"cell_type":"code","source":"inference_dataloader = DataLoader(inference_dataset, batch_size=config.BATCH_SIZE,\n shuffle=False, num_workers=config.NUM_WORKERS, collate_fn=collator)\n\n# === Let's check one sample ===\nbatch = next(iter(inference_dataloader))\nencoding = batch\n\n# === Iterate over each element in the dictionary and print shape ===\nfor k,v in encoding.items():\n try:\n print(f\"{k} shape: {v.shape} \\n\")\n except:\n print(f\"{k} shape: {len(v)} \\n\")","metadata":{"execution":{"iopub.status.busy":"2023-05-20T14:58:50.993936Z","iopub.execute_input":"2023-05-20T14:58:50.994312Z","iopub.status.idle":"2023-05-20T14:58:51.388158Z","shell.execute_reply.started":"2023-05-20T14:58:50.994277Z","shell.execute_reply":"2023-05-20T14:58:51.386772Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# | Inference [↑](#top) \n\n***","metadata":{}},{"cell_type":"code","source":"device = config.DEVICE\nmodel.to(device)\nmodel.eval()\npredictions = []\nwith torch.no_grad():\n for step, batch in enumerate(tqdm(inference_dataloader, unit=\"test_batch\")):\n ids = batch.pop(\"id\")\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 preds = processor.batch_decode(prediction, skip_special_tokens=True)\n preds = [{'id': id_value, 'prediction': pred_value} for id_value, pred_value in zip(ids, preds)]\n predictions += get_prediction(preds)\n print(f'Step: {step} | Prediction: {preds}'), print(\"\\n\")","metadata":{"execution":{"iopub.status.busy":"2023-05-20T15:08:27.888693Z","iopub.execute_input":"2023-05-20T15:08:27.889346Z","iopub.status.idle":"2023-05-20T15:08:36.483906Z","shell.execute_reply.started":"2023-05-20T15:08:27.889312Z","shell.execute_reply":"2023-05-20T15:08:36.482913Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# | Submission [↑](#top) \n\n***","metadata":{}},{"cell_type":"code","source":"submission = pd.DataFrame(format_submission(predictions))\nsubmission[\"chart_type\"] = submission[\"chart_type\"].apply(lambda x: re.sub(r\"\\s\", \"\", x))\nsubmission.to_csv(\"submission.csv\", index=False)\nsubmission","metadata":{"execution":{"iopub.status.busy":"2023-05-20T15:16:14.973439Z","iopub.execute_input":"2023-05-20T15:16:14.973889Z","iopub.status.idle":"2023-05-20T15:16:14.989592Z","shell.execute_reply.started":"2023-05-20T15:16:14.973830Z","shell.execute_reply":"2023-05-20T15:16:14.988581Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# | Quality Assurance [↑](#top) \n\n***","metadata":{}},{"cell_type":"code","source":"try:\n assert len(submission) == len(glob(paths.TEST_IMAGES_FOLDER + \"/*.jpg\"))*2\nexcept:\n raise ValueError(\"The number of rows in submission/2 does not match the number of images.\")\n \ntry:\n assert submission.id.nunique() == len(glob(paths.TEST_IMAGES_FOLDER + \"/*.jpg\"))*2\nexcept:\n raise ValueError(\"The number of unique IDs/2 does not match the number of images.\")\n \ntry:\n assert submission.columns.tolist() == ['id', 'data_series', 'chart_type']\nexcept:\n raise ValueError(\"Wrong column names.\")","metadata":{"execution":{"iopub.status.busy":"2023-05-20T15:20:23.138036Z","iopub.execute_input":"2023-05-20T15:20:23.138405Z","iopub.status.idle":"2023-05-20T15:20:23.151827Z","shell.execute_reply.started":"2023-05-20T15:20:23.138374Z","shell.execute_reply":"2023-05-20T15:20:23.150904Z"},"trusted":true},"execution_count":null,"outputs":[]}]}