{"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":"# Donut with additional hyperparameter tuning\n\n### Notes: This could still absolutely improve with better post text processing. As I have not put a ton of thought into that peice of things yet. I would imagine this is not the final solution for the donut model and there are likely several other creative ways to boost preformance, especially with regaurd to horizontal or vertical bar plots.\n\n### A special thank you to @nbroad (Nicholas Broad) for his insights using donut! Go check out his work and discussions","metadata":{}},{"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\")","metadata":{"execution":{"iopub.status.busy":"2023-05-13T18:22:18.505342Z","iopub.execute_input":"2023-05-13T18:22:18.506664Z","iopub.status.idle":"2023-05-13T18:22:32.508238Z","shell.execute_reply.started":"2023-05-13T18:22:18.506624Z","shell.execute_reply":"2023-05-13T18:22:32.507262Z"},"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/benetech-donut\"\n\nBOS_TOKEN = \"<|BOS|>\"\nX_START = \"<x_start>\"\nX_END = \"<x_end>\"\nY_START = \"<y_start>\"\nY_END = \"<y_end>\"\n\nPLACEHOLDER_DATA_SERIES = \"0;0\"\nPLACEHOLDER_CHART_TYPE = \"line\"","metadata":{"execution":{"iopub.status.busy":"2023-05-13T18:22:32.510760Z","iopub.execute_input":"2023-05-13T18:22:32.511567Z","iopub.status.idle":"2023-05-13T18:22:32.517789Z","shell.execute_reply.started":"2023-05-13T18:22:32.511533Z","shell.execute_reply":"2023-05-13T18:22:32.516171Z"},"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-05-13T18:22:32.522753Z","iopub.execute_input":"2023-05-13T18:22:32.523413Z","iopub.status.idle":"2023-05-13T18:22:32.546282Z","shell.execute_reply.started":"2023-05-13T18:22:32.523381Z","shell.execute_reply":"2023-05-13T18:22:32.545249Z"},"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    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        \n        pixel_values.append(processor(arr, random_padding=True).pixel_values)\n        \n        \n    return {\n        \"pixel_values\": torch.tensor(np.vstack(pixel_values)),\n    }\n\nmodel = 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 = 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    pixel_values = batch[\"pixel_values\"].to(device)\n\n    batch_size = pixel_values.shape[0]\n\n    decoder_input_ids = torch.full(\n        (batch_size, 1),\n        decoder_start_token_id,\n        device=pixel_values.device,\n    )\n\n    try:\n        outputs = model.generate(\n            pixel_values,\n            decoder_input_ids=decoder_input_ids,\n            max_length=CFG.max_length,\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=2,       #1 int    (1 - 10)\n            temperature=.9,     #1 float  (0 -  ) less div - more div\n            top_k=1,           #1 int    (1 -  ) less div - more div\n            top_p=.4,           #1 float (0 - 1) more div - less div\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-05-13T18:22:32.548709Z","iopub.execute_input":"2023-05-13T18:22:32.549141Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"display(sub_df)","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}