{"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 re \nfrom pathlib import Path\nfrom typing import List\nfrom functools import partial\n\nfrom transformers import (\n    DonutProcessor,\n    VisionEncoderDecoderConfig,\n    VisionEncoderDecoderModel,\n)\nimport torch\nfrom torch.utils.data import DataLoader\nimport numpy as np\nimport pandas as pd\nfrom datasets import Dataset\nfrom datasets import Image as ds_img\nfrom tqdm.notebook import tqdm\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")\nfrom transformers import Pix2StructForConditionalGeneration, Pix2StructProcessor","metadata":{"execution":{"iopub.status.busy":"2023-06-17T04:59:50.793565Z","iopub.execute_input":"2023-06-17T04:59:50.793872Z","iopub.status.idle":"2023-06-17T05:00:03.864445Z","shell.execute_reply.started":"2023-06-17T04:59:50.793840Z","shell.execute_reply":"2023-06-17T05:00:03.863446Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class CFG:\n    \n    test_grayscale = True\n    debug_clean = False\n    \n    batch_size = 4\n    image_path = \"/kaggle/input/benetech-making-graphs-accessible/test/images\"\n    max_length = 512\n    model_dir = \"/kaggle/input/extra-model-benetech/matcha_extra_100k\"\n    \n    \n    \n\nBOS_TOKEN = \"<|BOS|>\"\nX_START = \"<x_start>\"\nX_END = \"<x_end>\"\nY_START = \"<y_start>\"\nY_END = \"<y_end>\"\nPROMPT_TOKEN = \"<|PROMPT|>\"\n\nPLACEHOLDER_DATA_SERIES = \"0;0\"\nPLACEHOLDER_CHART_TYPE = \"line\"","metadata":{"execution":{"iopub.status.busy":"2023-06-17T05:00:03.866317Z","iopub.execute_input":"2023-06-17T05:00:03.867031Z","iopub.status.idle":"2023-06-17T05:00:03.875403Z","shell.execute_reply.started":"2023-06-17T05:00:03.867002Z","shell.execute_reply":"2023-06-17T05:00:03.872798Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def clean_preds(x: List[str], y: List[str]):\n    \"\"\"\n    This function cleans the x and y values predicted by Donut.\n\n    Because it is a generative model, it can insert any character in the \n    model's vocabulary into the prediction string. This function primarily removes\n    characters that prevent a number from being cast to a float.\n\n    Example:\n\n    x = [\"11\", \"12\", \"1E\", \"14\", \"15\"]\n    y = [\"Monday\", \"Tuesday\", \"Wednesday\", \"Thursday\", \"Friday\"]\n\n    # float(\"1E\") will throw an error\n\n    new_x, new_y = clean_preds(x, y)\n\n    new_x = [\"11\", \"12\", \"13\", \"14\", \"15\"]\n    new_y = [\"Monday\", \"Tuesday\", \"Wednesday\", \"Thursday\", \"Friday\"]\n\n    Args:\n        x (List[str]): The x values predicted by Donut.\n        y (List[str]): The y values predicted by Donut.\n\n    Returns:\n        x (List[str]): The cleaned x values.\n        y (List[str]): The cleaned y values.\n    \"\"\"\n    \n    def clean(str_list):\n        \n        new_list = []\n        for temp in str_list:\n            if \".\" not in temp:\n                dtype = int\n            else:\n                dtype = float\n            try:\n                # First try removing whitespace e.g. float(\"10 0\") will fail\n                temp = dtype(re.sub(\"\\s\", \"\", temp))\n            except ValueError:\n\n                # remove everything that isn't a digit, period, negative sign, or the letter e\n                # It could be \"1e-5\" or \"-0.134\"\n\n                temp = re.sub(r\"[^0-9\\.\\-eE]\", \"\", temp)\n\n                if len(temp) == 0:\n                    temp = 0\n                else:\n                    multiple_periods = len(re.findall(r\"\\.\", temp)) > 1\n                    multiple_negative_signs = len(re.findall(r\"\\-\", temp)) > 1\n                    multiple_e = len(re.findall(r\"[eE]\", temp)) > 1\n\n                    # Put negative sign in from of it all\n                    if multiple_negative_signs:\n                        temp = \"-\" + re.sub(r\"\\-\", \"\", temp)\n\n                    # Keep first period if there are multiple\n                    if multiple_periods:\n                        chunks = temp.split(\".\")\n                        try:\n                            temp = chunks[0] + \".\" + \"\".join(chunks[1:])\n                        except IndexError:\n                            temp = \"\".join(chunks)\n                    \n                    # Keep last e in case it is \"e1e-5\"\n                    if multiple_e:\n                        while temp.lower().startswith(\"e\"):\n                            temp = temp[1:]\n                        \n                        while temp.lower().endswith(\"e\"):\n                            temp = temp[:-1]\n                            \n                        chunks = temp.split(\"e\")\n                        try:\n                            temp = chunks[0:-1] + \"e\" + \"\".join(chunks[-1])\n                        except IndexError:\n                            temp = \"\".join(chunks)\n                try:\n                    temp = dtype(temp)\n                except ValueError:\n                    temp = 0\n                    \n            new_list.append(temp)\n\n        return new_list\n\n    all_x_chars = \"\".join(x)\n    all_y_chars = \"\".join(y)\n\n    frac_num_x = len(re.sub(r\"[^\\d]\", \"\", all_x_chars)) / len(all_x_chars)\n    frac_num_y = len(re.sub(r\"[^\\d]\", \"\", all_y_chars)) / len(all_y_chars)\n    \n    print(frac_num_x, frac_num_y)\n\n    if CFG.debug_clean:\n        print(f\"x before clean (len={len(x)})\", x)\n        print(f\"y before clean (len={len(y)})\", y)\n\n    if frac_num_x >= 0.5:\n        x = clean(x)\n    else:\n        x = [s.strip() for s in x]\n    \n    \n    if frac_num_y >= 0.5:\n        y = clean(y)\n    else:\n        y = [s.strip() for s in y]\n        \n    if CFG.debug_clean:\n        print(f\"x after clean (len={len(x)})\", x)\n        print(f\"y after clean (len={len(x)})\", x)\n\n    return x, y\n    \n\ndef string2preds(pred_string: str):\n    \"\"\"\n    Convert the prediction string from Donut to a chart type and x and y values.\n\n    Checks to make sure the special tokens are present and that the x and y values are not empty.\n    Will truncate the list of values to the smaller length of the two lists. This is because the \n    lengths of the x and y values must be the same to earn any points.\n\n    Args:\n        pred_string (str): The prediction string from Donut.\n\n    Returns:\n        chart_type (str): The chart type predicted by Donut.\n        x (List[str]): The x values predicted by Donut.\n        y (List[str]): The y values predicted by Donut.\n    \"\"\"\n\n    if \"<dot>\" in pred_string:\n        chart_type = \"dot\"\n    elif \"<horizontal_bar>\" in pred_string:\n        chart_type = \"horizontal_bar\"\n    elif \"<vertical_bar>\" in pred_string:\n        chart_type = \"vertical_bar\"\n    elif \"<scatter>\" in pred_string:\n        chart_type = \"scatter\"\n    elif \"<line>\" in pred_string:\n        chart_type = \"line\"\n    else:\n        return \"vertical_bar\", [], []\n    \n    \n    if not all([x in pred_string for x in [X_START, X_END, Y_START, Y_END]]):\n        return chart_type, [], []\n    \n    pred_string = re.sub(r\"<one>\", \"1\", pred_string)\n\n    x = pred_string.split(X_START)[1].split(X_END)[0].split(\";\")\n    y = pred_string.split(Y_START)[1].split(Y_END)[0].split(\";\")\n\n    if len(x) == 0 or len(y) == 0:\n        return chart_type, [], []\n\n    x, y = clean_preds(x, y)\n\n    return chart_type, x, y","metadata":{"execution":{"iopub.status.busy":"2023-06-17T05:00:03.877433Z","iopub.execute_input":"2023-06-17T05:00:03.877972Z","iopub.status.idle":"2023-06-17T05:00:03.901939Z","shell.execute_reply.started":"2023-06-17T05:00:03.877939Z","shell.execute_reply":"2023-06-17T05:00:03.900938Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"image_dir = Path(CFG.image_path)\nimages = list(image_dir.glob(\"*.jpg\"))\n\nds = Dataset.from_dict(\n    {\"image_path\": [str(x) for x in images], \"id\": [x.stem for x in images]}\n).cast_column(\"image_path\", ds_img())\n\ndef preprocess(examples, processor):\n    flattened_patches =  []\n    attention_mask = []\n\n#     pixel_values = []\n\n    for sample in examples[\"image_path\"]:\n        arr = np.array(sample)\n        \n        # There are some grayscale images that were making this fail\n        # This prevents that.\n        if len(arr.shape) == 2:\n            print(\"Changing grayscale to 3 channel format\")\n            print(arr.shape)\n            arr = np.stack([arr]*3, axis=-1)\n        encoding = processor(\n            sample,\n            text=\"Generate underlying data table of the figure below:\",\n            font_path=\"/kaggle/input/arial-font/arial.ttf\",\n            add_special_tokens=True, \n            max_patches=1024, \n            # return_tensors = \"pt\"\n        )\n#         pixel_values.append(processor(arr, random_padding=True).pixel_values)\n        flattened_patches.append(encoding.flattened_patches)\n        attention_mask.append(encoding.attention_mask)\n\n        \n        \n    return {\n#         \"pixel_values\": torch.tensor(np.vstack(pixel_values)),\n        \"flattened_patches\":  torch.tensor(np.vstack(flattened_patches)),\n        \"attention_mask\": torch.tensor(np.vstack(attention_mask)),\n    }\n\n\nmodel = Pix2StructForConditionalGeneration.from_pretrained(CFG.model_dir)\n\n# model = VisionEncoderDecoderModel.from_pretrained(CFG.model_dir)\nmodel.eval()\n\ndevice = torch.device(\"cuda:0\")\n\nmodel.to(device)\ndecoder_start_token_id = model.config.decoder_start_token_id\nprocessor = Pix2StructProcessor.from_pretrained(CFG.model_dir)\n# processor = DonutProcessor.from_pretrained(CFG.model_dir)\n\nids = ds[\"id\"]\nds.set_transform(partial(preprocess, processor=processor))\n\ndata_loader = DataLoader(\n    ds, batch_size=CFG.batch_size, shuffle=False\n)\n\n\nall_generations = []\nfor batch in tqdm(data_loader):\n    flattened_patches =   batch[\"flattened_patches\"].to(device)\n    attention_mask = batch[\"attention_mask\"].to(device)\n#     pixel_values = batch[\"pixel_values\"].to(device)\n\n#     batch_size = pixel_values.shape[0]\n    batch_size = flattened_patches.shape[0]\n    \n    \n    decoder_input_ids = torch.full(\n        (batch_size, 1),\n        decoder_start_token_id,\n        device=flattened_patches.device,\n    )\n\n    try:\n        outputs = model.generate(\n            flattened_patches=flattened_patches,\n            attention_mask=attention_mask,\n            decoder_input_ids=decoder_input_ids,\n            max_length=CFG.max_length,\n            return_dict_in_generate=True,\n        )\n\n        all_generations.extend(processor.batch_decode(outputs.sequences))\n        \n    except:\n        all_generations.extend([\"\"]*batch_size)\n        \nchart_types, x_preds, y_preds = [], [], []\nfor gen in all_generations:\n\n    try:\n        chart_type, x, y = string2preds(gen)\n        new_chart_type = chart_type\n        x_str = \";\".join(list(map(str, x)))\n        y_str = \";\".join(list(map(str, y)))\n\n    except Exception as e:\n        print(\"Failed to convert to string:\", gen)\n        print(e)\n        new_chart_type = PLACEHOLDER_CHART_TYPE\n        x_str = PLACEHOLDER_DATA_SERIES\n        y_str = PLACEHOLDER_DATA_SERIES\n            \n    if len(x_str) == 0:\n        x_str = PLACEHOLDER_DATA_SERIES\n    if len(y_str) == 0:\n        y_str = PLACEHOLDER_DATA_SERIES\n    \n    chart_types.append(new_chart_type)\n    x_preds.append(x_str)\n    y_preds.append(y_str)\n        \n\nsub_df = pd.DataFrame(\n    data={\n        \"id\": [f\"{id_}_x\" for id_ in ids] + [f\"{id_}_y\" for id_ in ids],\n        \"data_series\": x_preds + y_preds,\n        \"chart_type\": chart_types * 2,\n    }\n)\n\nsub_df.to_csv(\"submission.csv\", index=False)","metadata":{"execution":{"iopub.status.busy":"2023-06-17T05:00:03.906032Z","iopub.execute_input":"2023-06-17T05:00:03.906387Z","iopub.status.idle":"2023-06-17T05:00:45.744575Z","shell.execute_reply.started":"2023-06-17T05:00:03.906347Z","shell.execute_reply":"2023-06-17T05:00:45.743740Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"display(sub_df)","metadata":{"execution":{"iopub.status.busy":"2023-06-17T05:00:45.748442Z","iopub.execute_input":"2023-06-17T05:00:45.750532Z","iopub.status.idle":"2023-06-17T05:00:45.771028Z","shell.execute_reply.started":"2023-06-17T05:00:45.750499Z","shell.execute_reply":"2023-06-17T05:00:45.770133Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}