{"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":"### Imports","metadata":{}},{"cell_type":"code","source":"import os\nimport sys\nimport json \nimport glob\nimport gc\nimport torch\n\nimport pandas as pd\nimport numpy as np\n\nfrom PIL import Image\nfrom pathlib import Path","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:37:25.655259Z","iopub.execute_input":"2023-06-18T02:37:25.655621Z","iopub.status.idle":"2023-06-18T02:37:28.571862Z","shell.execute_reply.started":"2023-06-18T02:37:25.655589Z","shell.execute_reply":"2023-06-18T02:37:28.570895Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sys.path.insert(0, '/kaggle/input/omegaconf')\nfrom omegaconf import OmegaConf","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:37:28.574703Z","iopub.execute_input":"2023-06-18T02:37:28.575284Z","iopub.status.idle":"2023-06-18T02:37:28.674128Z","shell.execute_reply.started":"2023-06-18T02:37:28.575248Z","shell.execute_reply":"2023-06-18T02:37:28.673303Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torch.utils.data import DataLoader\nfrom accelerate import Accelerator\nfrom transformers import GenerationConfig\nfrom tqdm.auto import tqdm","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:37:28.675411Z","iopub.execute_input":"2023-06-18T02:37:28.675714Z","iopub.status.idle":"2023-06-18T02:37:36.027319Z","shell.execute_reply.started":"2023-06-18T02:37:28.675689Z","shell.execute_reply":"2023-06-18T02:37:36.026412Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Config","metadata":{}},{"cell_type":"code","source":"%%writefile mga_cfg_deplot.yaml\nmodel:\n    backbone_path: /kaggle/input/google-deplot\n    max_length: 8\n    max_patches: 3072\n    patch_size: 16\n    len_tokenizer: ???\n    pad_token_id: ???\n    decoder_start_token_id: ???\n    bos_token_id: ???\n        \npredict_params:\n    bs: 1\n    checkpoint_path: /kaggle/input/mga-r-final-matcha-v2-ft-i2/mga_model_fold_0.pth.tar\n\ncompetition_dataset:\n    data_dir: /kaggle/input/benetech-making-graphs-accessible","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:37:36.030065Z","iopub.execute_input":"2023-06-18T02:37:36.030772Z","iopub.status.idle":"2023-06-18T02:37:36.037181Z","shell.execute_reply.started":"2023-06-18T02:37:36.030736Z","shell.execute_reply":"2023-06-18T02:37:36.036233Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cfg = OmegaConf.load('mga_cfg_deplot.yaml')\ncfg","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:37:36.038862Z","iopub.execute_input":"2023-06-18T02:37:36.039233Z","iopub.status.idle":"2023-06-18T02:37:36.059960Z","shell.execute_reply.started":"2023-06-18T02:37:36.039187Z","shell.execute_reply":"2023-06-18T02:37:36.059021Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataset","metadata":{}},{"cell_type":"code","source":"# Reference: https://www.kaggle.com/code/nbroad/donut-train-benetech\n\nimport json\nimport os\nimport pdb\nfrom copy import deepcopy\n\nimport torch\nfrom PIL import Image\nfrom tokenizers import AddedToken\nfrom torch.utils.data import Dataset\nfrom transformers import Pix2StructProcessor\n\n# -- token map --#\nTOKEN_MAP = {\n    \"line\": \"[<lines>]\",\n    \"vertical_bar\": \"[<vertical_bar>]\",\n    \"scatter\": \"[<scatter>]\",\n    \"dot\": \"[<dot>]\",\n    \"horizontal_bar\": \"[<horizontal_bar>]\",\n    \"histogram\": \"[<histogram>]\",\n\n    \"c_start\": \"[<c_start>]\",\n    \"c_end\": \"[<c_end>]\",\n    \"x_start\": \"[<x_start>]\",\n    \"x_end\": \"[<x_end>]\",\n    \"y_start\": \"[<y_start>]\",\n    \"y_end\": \"[<y_end>]\",\n\n    \"p_start\": \"[<p_start>]\",\n    \"p_end\": \"[<p_end>]\",\n\n    \"bos_token\": \"[<mga>]\",\n}\n\n\n# -----\n\ndef is_nan(val):\n    return val != val\n\n\ndef get_processor(cfg):\n    \"\"\"\n    load the processor\n    \"\"\"\n    processor_path = cfg.model.backbone_path\n    print(f\"loading processor from {processor_path}\")\n    processor = Pix2StructProcessor.from_pretrained(processor_path)\n    processor.image_processor.is_vqa = False\n    processor.image_processor.patch_size = {\n        \"height\": cfg.model.patch_size,\n        \"width\": cfg.model.patch_size\n    }\n\n    # NEW TOKENS\n    print(\"adding new tokens...\")\n    new_tokens = []\n    for _, this_tok in TOKEN_MAP.items():\n        new_tokens.append(this_tok)\n    new_tokens = sorted(new_tokens)\n\n    tokens_to_add = []\n    for this_tok in new_tokens:\n        tokens_to_add.append(AddedToken(this_tok, lstrip=False, rstrip=False))\n\n    processor.tokenizer.add_tokens(tokens_to_add)\n\n    return processor\n\n\nclass MGADataset(Dataset):\n    \"\"\"Dataset class for MGA dataset\n    \"\"\"\n\n    def __init__(self, cfg, graph_ids, transform=None):\n\n        self.cfg = cfg\n        self.data_dir = cfg.competition_dataset.data_dir.rstrip(\"/\")\n        self.image_dir = os.path.join(self.data_dir, \"test\", \"images\")\n\n        self.graph_ids = deepcopy(graph_ids)\n        self.transform = transform\n\n        # load processor\n        self.load_processor()\n\n    def load_processor(self):\n        self.processor = get_processor(self.cfg)\n\n    def load_image(self, graph_id):\n        image_path = os.path.join(self.image_dir, f\"{graph_id}.jpg\")\n        image = Image.open(image_path)\n        if image.mode != 'RGB':\n            image = image.convert('RGB')\n\n        return image\n\n    def __len__(self):\n        return len(self.graph_ids)\n\n    def __getitem__(self, index):\n        graph_id = self.graph_ids[index]\n        image = self.load_image(graph_id)\n\n        p_img = self.processor(\n            images=image,\n            max_patches=self.cfg.model.max_patches,\n            add_special_tokens=True,\n        )\n\n        r = {}\n        r['id'] = graph_id\n        r['image'] = image\n        r['flattened_patches'] = p_img['flattened_patches']\n        r['attention_mask'] = p_img['attention_mask']\n\n        return r","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:37:36.061212Z","iopub.execute_input":"2023-06-18T02:37:36.061870Z","iopub.status.idle":"2023-06-18T02:37:36.447735Z","shell.execute_reply.started":"2023-06-18T02:37:36.061837Z","shell.execute_reply":"2023-06-18T02:37:36.446784Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Data Loader","metadata":{}},{"cell_type":"code","source":"import pdb\nfrom copy import deepcopy\nfrom dataclasses import dataclass\n\nimport numpy as np\nimport torch\nfrom transformers import DataCollatorWithPadding\n\n\n@dataclass\nclass MGACollator(DataCollatorWithPadding):\n    \"\"\"\n    data collector for mga task\n    \"\"\"\n\n    tokenizer = None\n    padding = True\n    max_length = None\n    pad_to_multiple_of = None\n    return_tensors = \"pt\"\n\n    def __call__(self, features):\n        batch = dict()\n\n        # graph ids\n        batch[\"id\"] = [feature[\"id\"] for feature in features]\n        batch[\"images\"] = [feature[\"image\"] for feature in features]\n\n        # image features ---\n        flattened_patches = [feature[\"flattened_patches\"] for feature in features]\n        attention_mask = [feature[\"attention_mask\"] for feature in features]\n\n        flattened_patches = np.concatenate(flattened_patches, axis=0)\n        attention_mask = np.concatenate(attention_mask, axis=0)\n\n        batch[\"flattened_patches\"] = flattened_patches\n        batch[\"attention_mask\"] = attention_mask\n\n        # casting ---\n        tensor_keys = [\"flattened_patches\", \"attention_mask\"]\n        \n        for key in tensor_keys:\n            if key != \"flattened_patches\":\n                batch[key] = torch.tensor(batch[key], dtype=torch.int64)\n            else:\n                batch[key] = torch.tensor(batch[key], dtype=torch.float32)\n\n        return batch\n","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:37:51.414179Z","iopub.execute_input":"2023-06-18T02:37:51.414576Z","iopub.status.idle":"2023-06-18T02:37:51.765602Z","shell.execute_reply.started":"2023-06-18T02:37:51.414546Z","shell.execute_reply":"2023-06-18T02:37:51.764596Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model","metadata":{}},{"cell_type":"code","source":"import sys\nfrom pathlib import Path\n\nimport numpy as np\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.utils.checkpoint\nfrom transformers import (AutoConfig, AutoModel, Pix2StructConfig,\n                          Pix2StructForConditionalGeneration)\n\n\ndef get_model(cfg):\n    model = Pix2StructForConditionalGeneration.from_pretrained(cfg.model.backbone_path)\n    return model\n\n\nclass MGAModel(nn.Module):\n    \"\"\"\n    The MGA model\n    \"\"\"\n\n    def __init__(self, cfg):\n        print(\"initializing the MGA model...\")\n\n        super(MGAModel, self).__init__()\n        self.cfg = cfg\n\n        backbone_config = Pix2StructConfig.from_pretrained(cfg.model.backbone_path)\n        backbone_config.text_config.max_length = cfg.model.max_length\n        backbone_config.text_config.is_decoder = True\n\n        backbone_config.text_config.pad_token_id = cfg.model.pad_token_id\n        backbone_config.text_config.decoder_start_token_id = cfg.model.decoder_start_token_id\n        backbone_config.text_config.bos_token_id = cfg.model.bos_token_id\n\n        # backbone_config.decoder.max_length = cfg.model.max_length\n\n        self.backbone = Pix2StructForConditionalGeneration.from_pretrained(\n            cfg.model.backbone_path,\n            config=backbone_config,\n        )\n\n        # resize model embeddings\n        print(\"resizing model embeddings...\")\n        print(f\"tokenizer length = {cfg.model.len_tokenizer}\")\n        self.backbone.decoder.resize_token_embeddings(cfg.model.len_tokenizer)\n\n        self.loss_fn = nn.CrossEntropyLoss(\n            ignore_index=-100,\n            reduction=\"mean\",\n        )\n\n    def forward(\n            self,\n            flattened_patches,\n            attention_mask,\n            labels,\n    ):\n\n        outputs = self.backbone(\n            flattened_patches=flattened_patches,\n            attention_mask=attention_mask,\n            labels=labels,\n        )\n\n        loss_main = outputs.loss\n\n        # logits = outputs.logits  # [:, 0]  # (bs, vocab)\n        # labels_cls = labels[:, 0]  # (bs, )\n        # loss = self.loss_fn(logits, labels)\n\n        # logits_num = outputs.logits[:, 1]  # (bs, vocab)\n        # labels_num = labels[:, 1]  # (bs, )\n        # loss_num = self.loss_fn(logits_num, labels_num)\n\n        loss = loss_main  # + 0.10 * loss_cls\n\n        loss_dict = {\n            \"loss_main\": loss_main,\n            \"loss_cls\": loss,\n        }\n\n        return loss, loss_dict\n","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:37:52.979025Z","iopub.execute_input":"2023-06-18T02:37:52.980048Z","iopub.status.idle":"2023-06-18T02:37:53.059478Z","shell.execute_reply.started":"2023-06-18T02:37:52.980000Z","shell.execute_reply":"2023-06-18T02:37:53.058555Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Inference","metadata":{}},{"cell_type":"code","source":"def post_process(pred_string, token_map, delimiter=\"|\"):\n    # get chart type ---\n    chart_options = [\n        \"horizontal_bar\",\n        \"dot\",\n        \"scatter\",\n        \"vertical_bar\",\n        \"line\",\n        \"histogram\",\n    ]\n\n    chart_type = \"line\"  # default type\n\n    for ct in chart_options:\n        if token_map[ct] in pred_string:\n            chart_type = ct\n            break\n\n    if chart_type == \"histogram\":\n        chart_type = \"vertical_bar\"\n\n    # get x series ---\n    x_start_tok = token_map[\"x_start\"]\n    x_end_tok = token_map[\"x_end\"]\n\n    try:\n        x = pred_string.split(x_start_tok)[1].split(x_end_tok)[0].split(delimiter)\n        x = [elem.strip() for elem in x if len(elem.strip()) > 0]\n    except IndexError:\n        x = []\n\n    # get y series ---\n    y_start_tok = token_map[\"y_start\"]\n    y_end_tok = token_map[\"y_end\"]\n\n    try:\n        y = pred_string.split(y_start_tok)[1].split(y_end_tok)[0].split(delimiter)\n        y = [elem.strip() for elem in y if len(elem.strip()) > 0]\n    except IndexError:\n        y = []\n\n    return chart_type, x, y\n","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:37:56.450554Z","iopub.execute_input":"2023-06-18T02:37:56.450923Z","iopub.status.idle":"2023-06-18T02:37:56.462184Z","shell.execute_reply.started":"2023-06-18T02:37:56.450894Z","shell.execute_reply":"2023-06-18T02:37:56.460275Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def run_evaluation(model, infer_dl, tokenizer, token_map, max_len):\n\n    # # config for text generation ---\n    conf_g = {\n        \"max_new_tokens\": max_len, # 1024, # 512,\n        \"do_sample\": False,\n        \"top_k\": 1,\n        \"use_cache\": True,\n    }\n\n    generation_config = GenerationConfig(**conf_g)\n\n    # put model in eval mode ---\n    model.eval()\n\n    all_ids = []\n    all_texts = []\n\n    progress_bar = tqdm(range(len(infer_dl)))\n    for batch in infer_dl:\n        with torch.no_grad():\n            batch_ids = batch[\"id\"]\n            generated_ids = model.backbone.generate(\n                flattened_patches=batch['flattened_patches'],\n                attention_mask=batch['attention_mask'],\n                generation_config=generation_config,\n            )\n            generated_texts = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)\n\n            all_ids.extend(batch_ids)\n            all_texts.extend(generated_texts)\n\n        progress_bar.update(1)\n    progress_bar.close()\n\n    # prepare output dataframe ---\n    preds = []\n    extended_preds = []\n    for this_id, this_text in zip(all_ids, all_texts):\n        id_x = f\"{this_id}_x\"\n        id_y = f\"{this_id}_y\"\n        pred_chart, pred_x, pred_y = post_process(this_text, token_map)\n\n        preds.append([id_x, pred_x, pred_chart])\n        preds.append([id_y, pred_y, pred_chart])\n\n        extended_preds.append([id_x, pred_x, pred_chart, this_text])\n        extended_preds.append([id_y, pred_y, pred_chart, this_text])\n\n    pred_df = pd.DataFrame(preds)\n    pred_df.columns = [\"id\", \"data_series\", \"chart_type\"]\n\n    return pred_df","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:37:56.841092Z","iopub.execute_input":"2023-06-18T02:37:56.841753Z","iopub.status.idle":"2023-06-18T02:37:56.852150Z","shell.execute_reply.started":"2023-06-18T02:37:56.841718Z","shell.execute_reply":"2023-06-18T02:37:56.851031Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"### Inference","metadata":{}},{"cell_type":"code","source":"image_dir = Path('/kaggle/input/benetech-making-graphs-accessible/test/images')\nimage_paths = list(image_dir.glob(\"*.jpg\"))\nimage_ids = [str(pth).split(\"/\")[-1].split(\".jpg\")[0] for pth in image_paths]\nprint(image_ids[:5])\n\ninfer_ds = MGADataset(cfg, image_ids)\n\ntokenizer = infer_ds.processor.tokenizer\ncfg.model.len_tokenizer = len(tokenizer)\n\ncfg.model.pad_token_id = tokenizer.pad_token_id\nBOS_TOKEN = TOKEN_MAP[\"bos_token\"]\n\ncfg.model.decoder_start_token_id = tokenizer.convert_tokens_to_ids([BOS_TOKEN])[0]\ncfg.model.bos_token_id = tokenizer.convert_tokens_to_ids([BOS_TOKEN])[0]\n\ncollate_fn = MGACollator(tokenizer=tokenizer)\n\ninfer_dl = DataLoader(\n    infer_ds,\n    batch_size=cfg.predict_params.bs,\n    collate_fn=collate_fn,\n    shuffle=False,\n)\n\ncfg_dict = OmegaConf.to_container(cfg, resolve=True)\n\nmodel = MGAModel(cfg)\n\nprint(\"==\"*50)\ncheckpoint_path = cfg.predict_params.checkpoint_path\nprint(f\"loading model from checkpoint: {checkpoint_path}\")\nckpt = torch.load(checkpoint_path)\nmodel.load_state_dict(ckpt['state_dict'])\ndel ckpt\ngc.collect()\nprint(\"loaded!\")\n\nprint(\"accelerator setup...\")\naccelerator = Accelerator()  # cpu = True\n\nmodel, infer_dl = accelerator.prepare(model, infer_dl)\n\nmodel.eval()\nprint(\"model is in eval mode\")\nprint(\"==\"*50)\n\npred_df = run_evaluation(model, infer_dl, tokenizer, TOKEN_MAP,  max_len=cfg.model.max_length)\n\npred_df[\"data_series\"] = pred_df[\"data_series\"].apply(lambda x: \";\".join(x))\npred_df.head()","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:37:58.899504Z","iopub.execute_input":"2023-06-18T02:37:58.899889Z","iopub.status.idle":"2023-06-18T02:38:36.908803Z","shell.execute_reply.started":"2023-06-18T02:37:58.899857Z","shell.execute_reply":"2023-06-18T02:38:36.907841Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"del model\ngc.collect()\ntorch.cuda.empty_cache()","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:38:36.913495Z","iopub.execute_input":"2023-06-18T02:38:36.914896Z","iopub.status.idle":"2023-06-18T02:38:37.312487Z","shell.execute_reply.started":"2023-06-18T02:38:36.914860Z","shell.execute_reply":"2023-06-18T02:38:37.311574Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# SPLIT","metadata":{}},{"cell_type":"code","source":"scatter_ids = pred_df[pred_df['chart_type'] == 'scatter']['id'].values.tolist()\nscatter_ids = [this_id.split('_x')[0].split('_y')[0] for this_id in scatter_ids]\nscatter_ids = list(set(scatter_ids))\n\nother_ids = pred_df[pred_df['chart_type'] != 'scatter']['id'].values.tolist()\nother_ids = [this_id.split('_x')[0].split('_y')[0] for this_id in other_ids]\nother_ids = list(set(other_ids))","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:38:37.316675Z","iopub.execute_input":"2023-06-18T02:38:37.318809Z","iopub.status.idle":"2023-06-18T02:38:37.332082Z","shell.execute_reply.started":"2023-06-18T02:38:37.318773Z","shell.execute_reply":"2023-06-18T02:38:37.331118Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pred_df_v0 = deepcopy(pred_df)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:38:43.756206Z","iopub.execute_input":"2023-06-18T02:38:43.756667Z","iopub.status.idle":"2023-06-18T02:38:43.766144Z","shell.execute_reply.started":"2023-06-18T02:38:43.756629Z","shell.execute_reply":"2023-06-18T02:38:43.765184Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Pass 2 - Other","metadata":{"execution":{"iopub.status.busy":"2023-05-10T02:48:45.530364Z","iopub.execute_input":"2023-05-10T02:48:45.53304Z","iopub.status.idle":"2023-05-10T02:49:03.11857Z","shell.execute_reply.started":"2023-05-10T02:48:45.533006Z","shell.execute_reply":"2023-05-10T02:49:03.117603Z"}}},{"cell_type":"code","source":"%%writefile mga_cfg_deplot.yaml\nmodel:\n    backbone_path: /kaggle/input/google-deplot\n    max_length: 1024\n    max_patches: 4096\n    patch_size: 16\n    len_tokenizer: ???\n    pad_token_id: ???\n    decoder_start_token_id: ???\n    bos_token_id: ???\n        \npredict_params:\n    bs: 1\n    checkpoint_path: /kaggle/input/mga-r-rest-v2/mga_model_fold_0.pth.tar\n\ncompetition_dataset:\n    data_dir: /kaggle/input/benetech-making-graphs-accessible","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:38:47.320035Z","iopub.execute_input":"2023-06-18T02:38:47.320586Z","iopub.status.idle":"2023-06-18T02:38:47.333584Z","shell.execute_reply.started":"2023-06-18T02:38:47.320543Z","shell.execute_reply":"2023-06-18T02:38:47.332444Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cfg = OmegaConf.load('mga_cfg_deplot.yaml')\ncfg","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:38:47.637400Z","iopub.execute_input":"2023-06-18T02:38:47.637868Z","iopub.status.idle":"2023-06-18T02:38:47.659698Z","shell.execute_reply.started":"2023-06-18T02:38:47.637824Z","shell.execute_reply":"2023-06-18T02:38:47.658761Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"image_ids = deepcopy(other_ids)\ninfer_ds = MGADataset(cfg, image_ids)\n\ntokenizer = infer_ds.processor.tokenizer\ncfg.model.len_tokenizer = len(tokenizer)\n\ncfg.model.pad_token_id = tokenizer.pad_token_id\nBOS_TOKEN = TOKEN_MAP[\"bos_token\"]\n\ncfg.model.decoder_start_token_id = tokenizer.convert_tokens_to_ids([BOS_TOKEN])[0]\ncfg.model.bos_token_id = tokenizer.convert_tokens_to_ids([BOS_TOKEN])[0]\n\ncollate_fn = MGACollator(tokenizer=tokenizer)\n\ninfer_dl = DataLoader(\n    infer_ds,\n    batch_size=cfg.predict_params.bs,\n    collate_fn=collate_fn,\n    shuffle=False,\n)\n\ncfg_dict = OmegaConf.to_container(cfg, resolve=True)\n\nmodel = MGAModel(cfg)\n\nprint(\"==\"*50)\ncheckpoint_path = cfg.predict_params.checkpoint_path\nprint(f\"loading model from checkpoint: {checkpoint_path}\")\nckpt = torch.load(checkpoint_path)\nmodel.load_state_dict(ckpt['state_dict'])\ndel ckpt\ngc.collect()\nprint(\"loaded!\")\n\nprint(\"accelerator setup...\")\naccelerator = Accelerator()  # cpu = True\n\nmodel, infer_dl = accelerator.prepare(model, infer_dl)\n\nmodel.eval()\nprint(\"model is in eval mode\")\nprint(\"==\"*50)\n\npred_df = run_evaluation(model, infer_dl, tokenizer, TOKEN_MAP,  max_len=cfg.model.max_length)\n\npred_df[\"data_series\"] = pred_df[\"data_series\"].apply(lambda x: \";\".join(x))\npred_df.head()","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:39:10.708558Z","iopub.execute_input":"2023-06-18T02:39:10.708951Z","iopub.status.idle":"2023-06-18T02:39:28.443351Z","shell.execute_reply.started":"2023-06-18T02:39:10.708917Z","shell.execute_reply":"2023-06-18T02:39:28.442395Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pred_df_v1_p1 = deepcopy(pred_df)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:39:32.496796Z","iopub.execute_input":"2023-06-18T02:39:32.497155Z","iopub.status.idle":"2023-06-18T02:39:32.503083Z","shell.execute_reply.started":"2023-06-18T02:39:32.497116Z","shell.execute_reply":"2023-06-18T02:39:32.501938Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"del model\ngc.collect()\ntorch.cuda.empty_cache()","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:39:44.508771Z","iopub.execute_input":"2023-06-18T02:39:44.509135Z","iopub.status.idle":"2023-06-18T02:39:44.921909Z","shell.execute_reply.started":"2023-06-18T02:39:44.509097Z","shell.execute_reply":"2023-06-18T02:39:44.920862Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Pass 2 - Scatter","metadata":{}},{"cell_type":"code","source":"%%writefile mga_cfg_deplot.yaml\nmodel:\n    backbone_path: /kaggle/input/google-deplot\n    max_length: 1280\n    max_patches: 3072\n    patch_size: 16\n    len_tokenizer: ???\n    pad_token_id: ???\n    decoder_start_token_id: ???\n    bos_token_id: ???\n        \npredict_params:\n    bs: 1\n    checkpoint_path: /kaggle/input/mga-r-scatter-matcha/mga_model_fold_0.pth.tar\n\ncompetition_dataset:\n    data_dir: /kaggle/input/benetech-making-graphs-accessible","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:39:51.077299Z","iopub.execute_input":"2023-06-18T02:39:51.077661Z","iopub.status.idle":"2023-06-18T02:39:51.085444Z","shell.execute_reply.started":"2023-06-18T02:39:51.077630Z","shell.execute_reply":"2023-06-18T02:39:51.083267Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"cfg = OmegaConf.load('mga_cfg_deplot.yaml')\ncfg","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:39:53.633133Z","iopub.execute_input":"2023-06-18T02:39:53.633522Z","iopub.status.idle":"2023-06-18T02:39:53.649407Z","shell.execute_reply.started":"2023-06-18T02:39:53.633493Z","shell.execute_reply":"2023-06-18T02:39:53.648470Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"image_ids = deepcopy(scatter_ids)\ninfer_ds = MGADataset(cfg, image_ids)\n\ntokenizer = infer_ds.processor.tokenizer\ncfg.model.len_tokenizer = len(tokenizer)\n\ncfg.model.pad_token_id = tokenizer.pad_token_id\nBOS_TOKEN = TOKEN_MAP[\"bos_token\"]\n\ncfg.model.decoder_start_token_id = tokenizer.convert_tokens_to_ids([BOS_TOKEN])[0]\ncfg.model.bos_token_id = tokenizer.convert_tokens_to_ids([BOS_TOKEN])[0]\n\ncollate_fn = MGACollator(tokenizer=tokenizer)\n\ninfer_dl = DataLoader(\n    infer_ds,\n    batch_size=cfg.predict_params.bs,\n    collate_fn=collate_fn,\n    shuffle=False,\n)\n\ncfg_dict = OmegaConf.to_container(cfg, resolve=True)\n\nmodel = MGAModel(cfg)\n\nprint(\"==\"*50)\ncheckpoint_path = cfg.predict_params.checkpoint_path\nprint(f\"loading model from checkpoint: {checkpoint_path}\")\nckpt = torch.load(checkpoint_path)\nmodel.load_state_dict(ckpt['state_dict'])\ndel ckpt\ngc.collect()\nprint(\"loaded!\")\n\nprint(\"accelerator setup...\")\naccelerator = Accelerator()  # cpu = True\n\nmodel, infer_dl = accelerator.prepare(model, infer_dl)\n\nmodel.eval()\nprint(\"model is in eval mode\")\nprint(\"==\"*50)\n\npred_df = run_evaluation(model, infer_dl, tokenizer, TOKEN_MAP,  max_len=cfg.model.max_length)\n\npred_df[\"data_series\"] = pred_df[\"data_series\"].apply(lambda x: \";\".join(x))\npred_df.head()","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:40:01.659476Z","iopub.execute_input":"2023-06-18T02:40:01.659844Z","iopub.status.idle":"2023-06-18T02:40:27.574393Z","shell.execute_reply.started":"2023-06-18T02:40:01.659812Z","shell.execute_reply":"2023-06-18T02:40:27.573370Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pred_df_v1_p2 = deepcopy(pred_df)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:40:29.503831Z","iopub.execute_input":"2023-06-18T02:40:29.504187Z","iopub.status.idle":"2023-06-18T02:40:29.509270Z","shell.execute_reply.started":"2023-06-18T02:40:29.504157Z","shell.execute_reply":"2023-06-18T02:40:29.508180Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Merge","metadata":{}},{"cell_type":"code","source":"pred_df = pd.concat([pred_df_v1_p1, pred_df_v1_p2])\npred_df = pred_df.reset_index(drop=True)\npred_df.head()","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:41:27.076195Z","iopub.execute_input":"2023-06-18T02:41:27.076682Z","iopub.status.idle":"2023-06-18T02:41:27.099009Z","shell.execute_reply.started":"2023-06-18T02:41:27.076643Z","shell.execute_reply":"2023-06-18T02:41:27.098227Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pred_df_final = pd.merge(pred_df_v0[['id', 'chart_type']], pred_df[['id', 'data_series']], on='id', how='left')\npred_df_final['data_series'] = pred_df_final['data_series'].fillna('')","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:41:34.685124Z","iopub.execute_input":"2023-06-18T02:41:34.685836Z","iopub.status.idle":"2023-06-18T02:41:34.703969Z","shell.execute_reply.started":"2023-06-18T02:41:34.685800Z","shell.execute_reply":"2023-06-18T02:41:34.702750Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pred_df_final = pred_df_final[[\"id\", \"data_series\", \"chart_type\"]].copy()\npred_df_final = pred_df_final.reset_index(drop=True)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:41:35.467593Z","iopub.execute_input":"2023-06-18T02:41:35.468275Z","iopub.status.idle":"2023-06-18T02:41:35.475605Z","shell.execute_reply.started":"2023-06-18T02:41:35.468233Z","shell.execute_reply":"2023-06-18T02:41:35.474360Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pred_df_final.to_csv(\"submission.csv\", index=False)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:41:35.785480Z","iopub.execute_input":"2023-06-18T02:41:35.785841Z","iopub.status.idle":"2023-06-18T02:41:35.794157Z","shell.execute_reply.started":"2023-06-18T02:41:35.785812Z","shell.execute_reply":"2023-06-18T02:41:35.793123Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"pred_df_final.head(10)","metadata":{"execution":{"iopub.status.busy":"2023-06-18T02:41:36.109281Z","iopub.execute_input":"2023-06-18T02:41:36.109806Z","iopub.status.idle":"2023-06-18T02:41:36.123102Z","shell.execute_reply.started":"2023-06-18T02:41:36.109774Z","shell.execute_reply":"2023-06-18T02:41:36.121983Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}