{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.7.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":43873,"databundleVersionId":5585780,"sourceType":"competition"},{"sourceId":8141854,"sourceType":"datasetVersion","datasetId":4813857},{"sourceId":32784,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":27443}],"dockerImageVersionId":30408,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# End-to-End Image-to-Text Model: Donut (Document Understanding Transformer) 🍩\n\nDonut is a combination of a swin encoder and a bart decoder. This means that it can generate text based on an image without needing any OCR. In this notebook, I show how you can use it to train a model for this competition. This is only the starting point, so I'm sure there is much room for improvement. Still, it is so cool that this one model can do everything necessary for this competition!\n\n![donut diagram](https://raw.githubusercontent.com/clovaai/donut/master/misc/overview.png)","metadata":{}},{"cell_type":"code","source":"# You will have to restart the notebook after running this\n!pip install -U datasets transformers pyarrow polyleven -q","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport re\nimport json\nfrom collections import Counter\nfrom itertools import chain\nfrom pathlib import Path\nfrom typing import List, Dict, Union, Tuple, Any\n\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom PIL import Image\nfrom sklearn.model_selection import StratifiedKFold\nfrom torch.utils.data import DataLoader\nfrom transformers import (\n    DonutProcessor,\n    VisionEncoderDecoderConfig,\n    VisionEncoderDecoderModel,\n    get_scheduler\n)\nfrom datasets import Dataset\nfrom datasets import Image as ds_img\nfrom polyleven import levenshtein # a faster version of levenshtein","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:11.294606Z","iopub.execute_input":"2024-04-18T08:09:11.295631Z","iopub.status.idle":"2024-04-18T08:09:16.819035Z","shell.execute_reply.started":"2024-04-18T08:09:11.295583Z","shell.execute_reply":"2024-04-18T08:09:16.817696Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data_dir = Path(\"/kaggle/input/benetech-making-graphs-accessible/train\")\n\nimages_path = data_dir / \"images\"\ntrain_json_files = list((data_dir / \"annotations\").glob(\"*.json\"))\n\n\nclass CFG:\n\n    # General\n    debug = False\n    num_proc = 2\n    num_workers = 4\n    gpus = 2\n\n    # Data\n    max_length = 1024\n    image_height = 560\n    image_width = 560\n\n    # Training\n    epochs = 2\n    val_check_interval = 1.0  # how many times we want to validate during an epoch\n    check_val_every_n_epoch = 1\n    gradient_clip_val = 1.0\n    lr = 3e-5\n    lr_scheduler_type = \"cosine\"\n    num_warmup_steps = 100\n    seed = 42\n    warmup_steps = 300  \n    output_path = \"/kaggle/working/\"\n    log_steps = 200\n    batch_size = 4\n    use_wandb = True","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:16.825931Z","iopub.execute_input":"2024-04-18T08:09:16.826249Z","iopub.status.idle":"2024-04-18T08:09:17.435271Z","shell.execute_reply.started":"2024-04-18T08:09:16.826218Z","shell.execute_reply":"2024-04-18T08:09:17.434419Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# First step is to transform the given annotations into a format the model can work with\n\nSince the predictions need to be in the form x1;x2;x3 and y1;y2;y3, let's make that the format we will generate. We also need to add special tokens to predict the chart type and separators to distinguish what is what.\n\nRemember, Donut takes an image as input and will generate text as output.","metadata":{}},{"cell_type":"code","source":"# Let's add chart types as special tokens and a special BOS token\n\nPROMPT_TOKEN = \"<|PROMPT|>\"\nX_START = \"<x_start>\"\nX_END = \"<x_end>\"\nY_START = \"<y_start>\"\nY_END = \"<y_end>\"\n\nSEPARATOR_TOKENS = [\n    PROMPT_TOKEN,\n    X_START,\n    X_END,\n    Y_START,\n    Y_END,\n]\n\nLINE_TOKEN =  \"<line>\" \nVERTICAL_BAR_TOKEN = \"<vertical_bar>\"\nHORIZONTAL_BAR_TOKEN = \"<horizontal_bar>\"\nSCATTER_TOKEN = \"<scatter>\"\nDOT_TOKEN = \"<dot>\"\n\nCHART_TYPE_TOKENS = [\n    LINE_TOKEN,\n    VERTICAL_BAR_TOKEN,\n    HORIZONTAL_BAR_TOKEN,\n    SCATTER_TOKEN,\n    DOT_TOKEN,\n]\n\nnew_tokens = SEPARATOR_TOKENS + CHART_TYPE_TOKENS\n\ndef 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\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\n    Args:\n        filepath (str): The path to the JSON file\n\n    Returns:\n        dict: A dictionary containing the ground truth string, x-y data, chart type, id, and source\n    \"\"\"\n    filepath = Path(filepath)\n\n    with open(filepath) as fp:\n        data = json.load(fp)\n\n    data_series = data[\"data-series\"]\n\n    all_x, all_y = [], []\n\n    for d in data_series:\n        x = d[\"x\"]\n        y = d[\"y\"]\n\n        x = round_float(x)\n        y = round_float(y)\n\n        # Ignore nan values\n        if is_nan(x) or is_nan(y):\n            continue\n\n        all_x.append(x)\n        all_y.append(y)\n        \n    \n    chart_type = f\"<{data['chart-type']}>\"\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    \n    gt_string = PROMPT_TOKEN + chart_type + x_str + y_str\n\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    }","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:17.437218Z","iopub.execute_input":"2024-04-18T08:09:17.437675Z","iopub.status.idle":"2024-04-18T08:09:17.457723Z","shell.execute_reply.started":"2024-04-18T08:09:17.437630Z","shell.execute_reply":"2024-04-18T08:09:17.456291Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"get_gt_string_and_xy(data_dir / \"annotations\" / \"000d269c8e26.json\")","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:17.462311Z","iopub.execute_input":"2024-04-18T08:09:17.462994Z","iopub.status.idle":"2024-04-18T08:09:17.474701Z","shell.execute_reply.started":"2024-04-18T08:09:17.462960Z","shell.execute_reply":"2024-04-18T08:09:17.473304Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Creating the dataset","metadata":{}},{"cell_type":"code","source":"# This generator function will be used to create the Dataset object\n\n\ndef gen_data(files: List[Union[str, os.PathLike]]) -> 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\n    Args:\n        files (list): A list of json files\n\n    Returns:\n        generator: A generator that yields a dictionary with the ground truth string and\n            the path to the corresponding image.\n    \"\"\"\n\n    for f in files:\n\n        yield {\n            **get_gt_string_and_xy(f),\n            \"image_path\": str(images_path / f\"{f.stem}.jpg\"),\n        }\n\n\nds = Dataset.from_generator(\n    gen_data, gen_kwargs={\"files\": train_json_files}, num_proc=CFG.num_proc\n)","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:17.476288Z","iopub.execute_input":"2024-04-18T08:09:17.476698Z","iopub.status.idle":"2024-04-18T08:09:24.008929Z","shell.execute_reply.started":"2024-04-18T08:09:17.476658Z","shell.execute_reply":"2024-04-18T08:09:24.007998Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Checking image sizes","metadata":{}},{"cell_type":"code","source":"def add_image_sizes(examples: Dict[str, Union[str, os.PathLike]]) -> Dict[str, List[int]]:\n    \"\"\"\n    This function takes a dictionary of examples and adds the width and height of the\n    image to the dictionary. This is to be used with the `Dataset.map` function.\n\n    Args:\n        examples (dict): A dictionary of examples (from `map` function)\n\n    Returns:\n        dict: The dictionary with the width and height of the image added\n    \"\"\"\n\n    sizes = [Image.open(x).size for x in examples[\"image_path\"]]\n\n    width, height = list(zip(*sizes))\n\n    return {\n        \"width\": list(width),\n        \"height\": list(height),\n    }\n\n\nds = ds.map(add_image_sizes, batched=True, num_proc=CFG.num_proc)","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:24.010293Z","iopub.execute_input":"2024-04-18T08:09:24.010665Z","iopub.status.idle":"2024-04-18T08:09:24.038075Z","shell.execute_reply.started":"2024-04-18T08:09:24.010629Z","shell.execute_reply":"2024-04-18T08:09:24.037147Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Most fall within 560 x 560\nplt.figure(figsize=(10, 10))\n\nplt.scatter(x=ds[\"width\"], y=ds[\"height\"], marker=\"o\", alpha=0.5)\n\n# Add axis labels\nplt.xlabel(\"Width\")\nplt.ylabel(\"Height\")\n\n# Add a title\nplt.title(\"Image Dimensions\")\n\n# Add gridlines\nplt.grid(True)\n\n# Set aspect ratio to be equal\nplt.gca().set_aspect(\"equal\", adjustable=\"box\")\n\n# Show the plot\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:24.039654Z","iopub.execute_input":"2024-04-18T08:09:24.040566Z","iopub.status.idle":"2024-04-18T08:09:25.270708Z","shell.execute_reply.started":"2024-04-18T08:09:24.040520Z","shell.execute_reply":"2024-04-18T08:09:25.269639Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Model configuration\n\nNeed to set:  \n\n- image height\n- image width\n- max sequence length to generate","metadata":{}},{"cell_type":"code","source":"config = VisionEncoderDecoderConfig.from_pretrained(\"naver-clova-ix/donut-base\")\nconfig.encoder.image_size = (CFG.image_height, CFG.image_width)\nconfig.decoder.max_length = CFG.max_length\n\nprint(CFG.image_height, CFG.image_width, CFG.max_length)","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:25.272315Z","iopub.execute_input":"2024-04-18T08:09:25.273360Z","iopub.status.idle":"2024-04-18T08:09:25.398916Z","shell.execute_reply.started":"2024-04-18T08:09:25.273308Z","shell.execute_reply":"2024-04-18T08:09:25.397776Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Checking tokenizer for unknowns\n","metadata":{}},{"cell_type":"code","source":"processor = DonutProcessor.from_pretrained(\"naver-clova-ix/donut-base\")\nprocessor.image_processor.size = {\n    \"height\": CFG.image_height,\n    \"width\": CFG.image_width,\n}","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:25.400272Z","iopub.execute_input":"2024-04-18T08:09:25.400628Z","iopub.status.idle":"2024-04-18T08:09:25.907560Z","shell.execute_reply.started":"2024-04-18T08:09:25.400597Z","shell.execute_reply":"2024-04-18T08:09:25.906252Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def check_for_unk(examples: Dict[str, str]) -> Dict[str, List[str]]:\n    \"\"\"\n    Check for unknown tokens in the given examples.\n\n    This function takes a dictionary containing a list of ground truth texts and\n    tokenizes them using the processor's tokenizer. It then checks for any unknown\n    tokens in the tokenized text and returns a dictionary containing a list of the\n    unknown tokens for each example.\n\n    Args:\n        examples (dict): A dictionary containing a list of ground truth texts. \n            Example: {\"ground_truth\": [\"text1\", \"text2\", ...]}\n\n    Returns:\n        dict: A dictionary containing a list of unknown tokens for each example. \n            Example: {\"unk_tokens\": [[\"unk1\", \"unk2\"], [], [\"unk3\"], ...]}\n    \"\"\"\n\n    texts = examples[\"ground_truth\"]\n\n    ids = processor.tokenizer(texts).input_ids\n    tokens = [processor.tokenizer.tokenize(x, add_special_tokens=True) for x in texts]\n\n    unk_tokens = []\n    for example_ids, example_tokens in zip(ids, tokens):\n        example_unk_tokens = []\n        for i in range(len(example_ids)):\n            if example_ids[i] == processor.tokenizer.unk_token_id:\n                example_unk_tokens.append(example_tokens[i])\n\n        unk_tokens.append(example_unk_tokens)\n\n    return {\"unk_tokens\": unk_tokens}\n\n\nunk = ds.map(check_for_unk, batched=True, num_proc=CFG.num_proc)","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:25.908967Z","iopub.execute_input":"2024-04-18T08:09:25.909327Z","iopub.status.idle":"2024-04-18T08:09:25.949997Z","shell.execute_reply.started":"2024-04-18T08:09:25.909292Z","shell.execute_reply":"2024-04-18T08:09:25.948738Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Let's look at only the examples with unknown tokens\nunk = unk.filter(lambda x: len(x[\"unk_tokens\"]) > 0, num_proc=CFG.num_proc)\n\nprint(len(unk))\n\nunk[0]","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:25.951940Z","iopub.execute_input":"2024-04-18T08:09:25.952665Z","iopub.status.idle":"2024-04-18T08:09:25.980788Z","shell.execute_reply.started":"2024-04-18T08:09:25.952617Z","shell.execute_reply":"2024-04-18T08:09:25.979621Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Let's count which tokens show up most often as unknowns\nall_unk_tokens = [x for y in unk[\"unk_tokens\"] for x in y]\n\nCounter(all_unk_tokens)","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:25.982561Z","iopub.execute_input":"2024-04-18T08:09:25.982955Z","iopub.status.idle":"2024-04-18T08:09:26.522543Z","shell.execute_reply.started":"2024-04-18T08:09:25.982914Z","shell.execute_reply":"2024-04-18T08:09:26.521362Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"It seems that the donut tokenizer doesn't contain \"1\" which will result in many unknown tokens in the ground truth string. There is \"_1\" (space in front of 1) and \"01\", \"11\", \"21\", \"31\", \"41\" etc. But if there is a \"0.1\" or \">1\" then the 1 will turn into an unknown token.\n\nI'll handle this in the `preprocess` function.","metadata":{}},{"cell_type":"code","source":"example_str = \"0.1 1 1990\"\n\ntemp_ids = processor.tokenizer(example_str).input_ids\nprint(\"ids:\", temp_ids)\nprint(\"tokenized:\", processor.tokenizer.tokenize(example_str))\nprint(\"decoded:\", processor.tokenizer.decode(temp_ids))\nprint(\"unk id:\", processor.tokenizer.unk_token_id)\n\n# Adding these tokens should mean that there should be very few unknown tokens\nnum_added = processor.tokenizer.add_tokens([\"<one>\"] + new_tokens)\nprint(num_added, \"tokens added\")\n\nconfig.pad_token_id = processor.tokenizer.pad_token_id\nconfig.decoder_start_token_id = processor.tokenizer.convert_tokens_to_ids([PROMPT_TOKEN])[0]","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:26.528154Z","iopub.execute_input":"2024-04-18T08:09:26.528618Z","iopub.status.idle":"2024-04-18T08:09:26.537917Z","shell.execute_reply.started":"2024-04-18T08:09:26.528583Z","shell.execute_reply":"2024-04-18T08:09:26.536631Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Preprocessing function to get pixel values and input_ids\n\nWhen ever the `_getitem__` function is called for the dataset, it will run this function.","metadata":{}},{"cell_type":"code","source":"from functools import partial\n\none_token_id = processor.tokenizer(\"<one>\", add_special_tokens=False).input_ids[0]\nunk_token_id = processor.tokenizer.unk_token_id\n\n\ndef replace_unk_tokens_with_one(example_ids: List[int], example_tokens: List[str], one_token_id:int, unk_token_id:int) -> List[int]:\n    \"\"\"\n    Replace unknown tokens that represent \"1\" with the correct token id.\n\n    Args:\n        example_ids (list): List of token ids for a given example\n        example_tokens (list): List of tokens for the same given example\n        one_token_id (int): Token id for the \"<one>\" token\n        unk_token_id (int): Token id for the unknown token\n\n    Returns:\n        list: The updated list of token ids with the correct token id for \"1\"\n    \"\"\"\n    \n    temp_ids = []\n    for id_, token in zip(example_ids, example_tokens):\n        if id_ == unk_token_id and token == \"1\":\n            id_ = one_token_id\n        temp_ids.append(id_)\n    return temp_ids\n\n\ndef preprocess(examples: Dict[str, str], processor: DonutProcessor, CFG: CFG) -> Dict[str, Union[torch.Tensor, List[int], List[str]]]:\n    \"\"\"\n    Preprocess the given examples.\n\n    This function processes the input examples by tokenizing the texts, replacing\n    any unknown tokens that represent \"1\" with the correct token id, and loading\n    the images.\n\n    Args:\n        examples (dict): A dictionary containing ground truth texts, image paths, and ids\n        processor: An object responsible for tokenizing texts and processing images\n        CFG: A configuration object containing settings and hyperparameters\n\n    Returns:\n        dict: A dictionary containing preprocessed images, token ids, and ids\n    \"\"\"\n    \n    pixel_values = []\n\n    texts = examples[\"ground_truth\"]\n\n    ids = processor.tokenizer(\n        texts,\n        add_special_tokens=False,\n        max_length=CFG.max_length,\n        padding=True,\n        truncation=True,\n    ).input_ids\n\n    if isinstance(texts, str):\n        texts = [texts]\n\n    tokens = [processor.tokenizer.tokenize(text, add_special_tokens=False) for text in texts]\n    \n    one_token_id = processor.tokenizer(\"<one>\", add_special_tokens=False).input_ids[0]\n    unk_token_id = processor.tokenizer.unk_token_id\n    \n    final_ids = [\n        replace_unk_tokens_with_one(example_ids, example_tokens, one_token_id, unk_token_id)\n        for example_ids, example_tokens in zip(ids, tokens)\n    ]\n\n    for sample in examples[\"image_path\"]:\n        pixel_values.append(processor(sample, random_padding=True).pixel_values)\n\n    return {\n        \"pixel_values\": torch.tensor(np.vstack(pixel_values)),\n        \"input_ids\": final_ids,\n        \"id\": examples[\"id\"],\n    }\n\n\nimage_ds = ds.cast_column(\"image_path\", ds_img())\nimage_ds.set_transform(partial(preprocess, processor=processor, CFG=CFG))","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:26.539293Z","iopub.execute_input":"2024-04-18T08:09:26.539698Z","iopub.status.idle":"2024-04-18T08:09:26.612173Z","shell.execute_reply.started":"2024-04-18T08:09:26.539657Z","shell.execute_reply":"2024-04-18T08:09:26.611325Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"sample = image_ds[[0, 1, 2]]\n\nprint(sample[\"pixel_values\"].shape)\nprint(processor.decode(sample[\"input_ids\"][2]))\nprint(len(sample[\"input_ids\"][2]))\nprint(processor.tokenizer.convert_ids_to_tokens(sample[\"input_ids\"][2]))","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:26.613534Z","iopub.execute_input":"2024-04-18T08:09:26.613837Z","iopub.status.idle":"2024-04-18T08:09:26.774310Z","shell.execute_reply.started":"2024-04-18T08:09:26.613808Z","shell.execute_reply":"2024-04-18T08:09:26.773037Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Creating splits\n\n\nBasic split: break extracted into folds, all generated and some extracted into train. Only extracted in validation. This is NB's approach, I would put even more extracted data into training, and use a small portion of extracted as well as my uploaded as test. Also, i only use line type.","metadata":{}},{"cell_type":"code","source":"# # separate by extracted and generated\n# extracted_ds = ds.filter(lambda x: (x[\"source\"] == \"extracted\"), num_proc=CFG.num_proc)\n# generated_ds = ds.filter(lambda x: (x[\"source\"] == \"generated\"), num_proc=CFG.num_proc)\n\n\nextracted_ds = ds.filter(lambda x: (x[\"source\"] == \"extracted\") and (x[\"chart-type\"] == \"line\"), num_proc=CFG.num_proc)\ngenerated_ds = ds.filter(lambda x: (x[\"source\"] == \"generated\") and (x[\"chart-type\"] == \"line\"), num_proc=CFG.num_proc)\n","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:26.776090Z","iopub.execute_input":"2024-04-18T08:09:26.776655Z","iopub.status.idle":"2024-04-18T08:09:26.815115Z","shell.execute_reply.started":"2024-04-18T08:09:26.776605Z","shell.execute_reply":"2024-04-18T08:09:26.814168Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"chart_types = extracted_ds[\"chart-type\"]\n\nprint(Counter(chart_types))","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:26.816504Z","iopub.execute_input":"2024-04-18T08:09:26.816936Z","iopub.status.idle":"2024-04-18T08:09:26.829335Z","shell.execute_reply.started":"2024-04-18T08:09:26.816887Z","shell.execute_reply":"2024-04-18T08:09:26.827972Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"skf = StratifiedKFold(n_splits=4)\n\nfold_idxs = []\n\nfor _, val_idxs in skf.split(chart_types, y=chart_types):\n    fold_idxs.append(val_idxs)\n\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":"2024-04-18T08:09:26.830858Z","iopub.execute_input":"2024-04-18T08:09:26.832695Z","iopub.status.idle":"2024-04-18T08:09:26.842987Z","shell.execute_reply.started":"2024-04-18T08:09:26.832643Z","shell.execute_reply":"2024-04-18T08:09:26.841549Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Combine generated and some of the extracted examples into the train set\n\nfrom datasets import concatenate_datasets\n\nfold = 0\n","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:26.844352Z","iopub.execute_input":"2024-04-18T08:09:26.844710Z","iopub.status.idle":"2024-04-18T08:09:26.850340Z","shell.execute_reply.started":"2024-04-18T08:09:26.844677Z","shell.execute_reply":"2024-04-18T08:09:26.849096Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_extracted = extracted_ds.select(\n    list(chain(*[x for i, x in enumerate(fold_idxs) if i != fold]))\n)","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:26.851761Z","iopub.execute_input":"2024-04-18T08:09:26.852092Z","iopub.status.idle":"2024-04-18T08:09:26.915755Z","shell.execute_reply.started":"2024-04-18T08:09:26.852062Z","shell.execute_reply":"2024-04-18T08:09:26.914851Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"#train_ds = concatenate_datasets([train_extracted, generated_ds])\n\ntrain_ds = concatenate_datasets([extracted_ds, generated_ds])\n\ntrain_ds = train_ds.cast_column(\"image_path\", ds_img())\ntrain_ds.set_transform(partial(preprocess, processor=processor, CFG=CFG))\n\n# Creat validation set from only extracted examples\n\nval_gt_ds = extracted_ds.select(fold_idxs[fold])\nval_ds = val_gt_ds.cast_column(\"image_path\", ds_img())\nval_ds.set_transform(partial(preprocess, processor=processor, CFG=CFG))\n\ngt_chart_type = val_gt_ds[\"chart-type\"]\ngt_x = [json.loads(_) for _ in val_gt_ds[\"x\"]]\ngt_y = [json.loads(_) for _ in val_gt_ds[\"y\"]]\ngt_ids = val_gt_ds[\"id\"]\n\ni = 0\nprint(gt_chart_type[i])\nprint(gt_x[i])\nprint(gt_y[i])\nprint(gt_ids[i])\nprint(Counter(gt_chart_type))","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:26.917098Z","iopub.execute_input":"2024-04-18T08:09:26.917871Z","iopub.status.idle":"2024-04-18T08:09:27.057202Z","shell.execute_reply.started":"2024-04-18T08:09:26.917821Z","shell.execute_reply":"2024-04-18T08:09:27.056069Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Collate function to make sure the ids are all the same length in a batch","metadata":{}},{"cell_type":"code","source":"pad_token_id = processor.tokenizer.pad_token_id\n\n\ndef collate_fn(samples: List[Dict[str, Union[torch.Tensor, List[int], str]]]) -> Dict[str, Union[torch.Tensor, List[str]]]:\n    \"\"\"\n    Custom collate function for DataLoader.\n\n    This function takes a list of samples and combines them into a batch with\n    properly padded input_ids.\n\n    Args:\n        samples (List[Dict[str, Union[torch.Tensor, List[int], str]]]): \n            A list of samples, where each sample is a dictionary containing\n            \"pixel_values\" (torch.Tensor), \"input_ids\" (List[int]), and \"id\" (str).\n\n    Returns:\n        Dict[str, Union[torch.Tensor, List[str]]]: \n            A dictionary containing the combined pixel values, padded input_ids, and ids.\n    \"\"\"\n\n    batch = {}\n\n    batch[\"pixel_values\"] = torch.stack([x[\"pixel_values\"] for x in samples])\n\n    max_length = max([len(x[\"input_ids\"]) for x in samples])\n\n    # Make a multiple of 8 to efficiently use the tensor cores\n    if max_length % 8 != 0:\n        max_length = (max_length // 8 + 1) * 8\n\n    input_ids = [\n        x[\"input_ids\"] + [pad_token_id] * (max_length - len(x[\"input_ids\"]))\n        for x in samples\n    ]\n\n    labels = torch.tensor(input_ids)\n    labels[labels == pad_token_id] = -100 # ignore loss on padding tokens\n    batch[\"labels\"] = labels\n    \n    batch[\"id\"] = [x[\"id\"] for x in samples]\n\n    return batch\n","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:27.058644Z","iopub.execute_input":"2024-04-18T08:09:27.058943Z","iopub.status.idle":"2024-04-18T08:09:27.069854Z","shell.execute_reply.started":"2024-04-18T08:09:27.058912Z","shell.execute_reply":"2024-04-18T08:09:27.068718Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Dataloaders\n\nValidation uses generation so it is very slow. That's why I only use a small fraction of the examples.","metadata":{}},{"cell_type":"code","source":"if CFG.debug:\n    train_ds = train_ds.select(range(100))\n    val_ds = val_ds.select(range(100))\n\ntrain_dataloader = DataLoader(\n    train_ds,\n    batch_size=CFG.batch_size,\n    shuffle=True,\n    collate_fn=collate_fn,\n    num_workers=CFG.num_workers,\n)\nval_dataloader = DataLoader(\n    val_ds,\n    batch_size=CFG.batch_size,\n    shuffle=False,\n    collate_fn=collate_fn,\n    num_workers=CFG.num_workers,\n)\n\nnum_training_steps = len(train_dataloader) * CFG.epochs // CFG.gpus\n\nbatch = next(iter(train_dataloader))\n\nbatch.keys(), [(k, v.shape) for k, v in batch.items() if k != \"id\"]","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:27.071100Z","iopub.execute_input":"2024-04-18T08:09:27.072062Z","iopub.status.idle":"2024-04-18T08:09:28.022670Z","shell.execute_reply.started":"2024-04-18T08:09:27.072028Z","shell.execute_reply":"2024-04-18T08:09:28.021370Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Functions to calculate metrics","metadata":{}},{"cell_type":"code","source":"def rmse(y_true: List[float], y_pred: List[float]) -> float:\n    \"\"\"\n    Calculate the Root Mean Square Error (RMSE) between the true and predicted values.\n\n    Args:\n        y_true (List[float]): The true values.\n        y_pred (List[float]): The predicted values.\n\n    Returns:\n        float: The Root Mean Square Error.\n    \"\"\"\n    return np.sqrt(np.mean(np.square(np.subtract(y_true, y_pred))))\n\n\ndef sigmoid(x: float) -> float:\n    \"\"\"\n    Calculate the sigmoid function for the given value.\n\n    Args:\n        x (float): The input value.\n\n    Returns:\n        float: The result of the sigmoid function.\n    \"\"\"\n    return 2 - 2 / (1 + np.exp(-x))\n\n\ndef normalized_rmse(y_true: List[float], y_pred: List[float]) -> float:\n    \"\"\"\n    Calculate the normalized Root Mean Square Error (RMSE) between the true and predicted values.\n\n    Args:\n        y_true (List[float]): The true values.\n        y_pred (List[float]): The predicted values.\n\n    Returns:\n        float: The normalized Root Mean Square Error.\n    \"\"\"\n    numerator = rmse(y_true, y_pred)\n    denominator = rmse(y_true, np.mean(y_true))\n\n    # https://www.kaggle.com/competitions/benetech-making-graphs-accessible/discussion/396947\n    if denominator == 0:\n        if numerator == 0:\n            return 1.0\n        return 0.0\n\n    return sigmoid(numerator / denominator)\n\n\ndef normalized_levenshtein_score(y_true: List[str], y_pred: List[str]) -> float:\n    \"\"\"\n    Calculate the normalized Levenshtein distance between two lists of strings.\n\n    Args:\n        y_true (List[str]): The true values.\n        y_pred (List[str]): The predicted values.\n\n    Returns:\n        float: The normalized Levenshtein distance.\n    \"\"\"\n    total_distance = np.sum([levenshtein(yt, yp) for yt, yp in zip(y_true, y_pred)])\n    length_sum = np.sum([len(yt) for yt in y_true])\n    return sigmoid(total_distance / length_sum)\n\n\ndef score_series(\n    y_true: List[Union[float, str]], y_pred: List[Union[float, str]]\n) -> float:\n    \"\"\"\n    Calculate the score for a series of true and predicted values.\n\n    Args:\n        y_true (List[Union[float, str]]): The true values.\n        y_pred (List[Union[float, str]]): The predicted values.\n\n    Returns:\n        float: The score for the series.\n    \"\"\"\n    if len(y_true) != len(y_pred):\n        return 0.0\n    if isinstance(y_true[0], str):\n        return normalized_levenshtein_score(y_true, y_pred)\n    else:\n        # Since this is a generative model, there is a chance it doesn't produce a float.\n        # In that case, we return 0.0.\n        try:\n            return normalized_rmse(y_true, list(map(float, y_pred)))\n        except:\n            return 0.0\n\n\ndef benetech_score(ground_truth: pd.DataFrame, predictions: pd.DataFrame) -> float:\n    \"\"\"Evaluate predictions using the metric from the Benetech - Making Graphs Accessible.\n\n    Parameters\n    ----------\n    ground_truth: pd.DataFrame\n        Has columns `[data_series, chart_type]` and an index `id`. Values in `data_series`\n        should be either arrays of floats or arrays of strings.\n\n    predictions: pd.DataFrame\n    \"\"\"\n    if not ground_truth.index.equals(predictions.index):\n        raise ValueError(\n            \"Must have exactly one prediction for each ground-truth instance.\"\n        )\n    if not ground_truth.columns.equals(predictions.columns):\n        raise ValueError(f\"Predictions must have columns: {ground_truth.columns}.\")\n    pairs = zip(\n        ground_truth.itertuples(index=False), predictions.itertuples(index=False)\n    )\n    scores = []\n    for (gt_series, gt_type), (pred_series, pred_type) in pairs:\n        if gt_type != pred_type:  # Check chart_type condition\n            scores.append(0.0)\n        else:  # Score with RMSE or Levenshtein as appropriate\n            scores.append(score_series(gt_series, pred_series))\n\n    ground_truth[\"score\"] = scores\n\n    grouped = ground_truth.groupby(\"chart_type\", as_index=False)[\"score\"].mean()\n\n    chart_type2score = {\n        chart_type: score\n        for chart_type, score in zip(grouped[\"chart_type\"], grouped[\"score\"])\n    }\n\n    return np.mean(scores), chart_type2score\n\n\ndef string2triplet(pred_string: str) -> Tuple[str, List[str], List[str]]:\n    \"\"\"\n    Convert a prediction string to a triplet of chart type, x values, and y values.\n\n    Args:\n        pred_string (str): The prediction string.\n\n    Returns:\n        Tuple[str, List[str], List[str]]: A triplet of chart type, x values, and y values.\n    \"\"\"\n    \n    chart_type = \"line\"\n    for tok in CHART_TYPE_TOKENS:\n        if tok in pred_string:\n            chart_type = tok.strip(\"<>\")\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    min_length = min(len(x), len(y))\n\n    x = x[:min_length]\n    y = y[:min_length]\n\n    return chart_type, x, y\n\n\ndef validation_metrics(val_outputs: List[str], val_ids: List[str], gt_df: pd.DataFrame) -> Dict[str, float]:\n    \"\"\"\n    Calculate validation metrics for a set of outputs, ids, and ground truth dataframe.\n\n    Args:\n        val_outputs (List[str]): A list of validation outputs.\n        val_ids (List[str]): A list of validation ids.\n        gt_df (pd.DataFrame): The ground truth dataframe.\n\n    Returns:\n        Dict[str, float]: A dictionary containing the validation scores.\n    \"\"\"\n    pred_triplets = []\n\n    for example_output in val_outputs:\n\n        if not all([x in example_output for x in [X_START, X_END, Y_START, Y_END]]):\n            pred_triplets.append((\"line\", [], []))\n        else:\n            pred_triplets.append(string2triplet(example_output))\n\n    pred_df = pd.DataFrame(\n        index=[f\"{id_}_x\" for id_ in val_ids] + [f\"{id_}_y\" for id_ in val_ids],\n        data={\n            \"data_series\": [x[1] for x in pred_triplets]\n            + [x[2] for x in pred_triplets],\n            \"chart_type\": [x[0] for x in pred_triplets] * 2,\n        },\n    )\n\n    overall_score, chart_type2score = benetech_score(\n        gt_df.loc[pred_df.index.values], pred_df\n    )\n\n    return {\n        \"val_score\": overall_score,\n        **{f\"{k}_score\": v for k, v in chart_type2score.items()},\n    }\n","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:28.025098Z","iopub.execute_input":"2024-04-18T08:09:28.025520Z","iopub.status.idle":"2024-04-18T08:09:28.064525Z","shell.execute_reply.started":"2024-04-18T08:09:28.025482Z","shell.execute_reply":"2024-04-18T08:09:28.063081Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Lightning Module\n\nIt will use generation for validation which can be very slow. If you want to utilize the two GPUs on Kaggle, you should probably not run it in a notebook because DDP gets funky when doing that.","metadata":{}},{"cell_type":"code","source":"import pytorch_lightning as pl\nfrom transformers import PreTrainedTokenizerBase, PreTrainedModel\n\nclass DonutModelPLModule(pl.LightningModule):\n    def __init__(self, processor: PreTrainedTokenizerBase, model: PreTrainedModel, gt_df: pd.DataFrame, num_training_steps: int):\n        \"\"\"\n        A PyTorch Lightning module for the DonutModel.\n\n        Args:\n            processor (PreTrainedTokenizerBase): The tokenizer/processor for the model.\n            model (PreTrainedModel): The pretrained model.\n            gt_df (pd.DataFrame): The ground truth dataframe.\n            num_training_steps (int): The number of training steps.\n        \"\"\"\n        super().__init__()\n        self.processor = processor\n        self.model = model\n        self.gt_df = gt_df\n        self.num_training_steps = num_training_steps\n\n    def training_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> torch.Tensor:\n        pixel_values = batch[\"pixel_values\"]\n        labels = batch[\"labels\"]\n\n        outputs = self.model(pixel_values, labels=labels)\n        loss = outputs.loss\n        self.log(\"train_loss\", loss)\n        return loss\n\n    def validation_step(self, batch: Dict[str, torch.Tensor], batch_idx: int, dataset_idx: int = 0) -> None:\n        pixel_values = batch[\"pixel_values\"]\n        batch_size = pixel_values.shape[0]\n        # we feed the prompt to the model\n        decoder_input_ids = torch.full(\n            (batch_size, 1),\n            self.model.config.decoder_start_token_id,\n            device=self.device,\n        )\n\n        outputs = self.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=self.processor.tokenizer.pad_token_id,\n            eos_token_id=self.processor.tokenizer.eos_token_id,\n            use_cache=True,\n            num_beams=1,\n            top_k=1,\n            bad_words_ids=[[self.processor.tokenizer.unk_token_id]],\n            return_dict_in_generate=True,\n        )\n\n        self.val_outputs.extend(\n            self.processor.tokenizer.batch_decode(outputs.sequences)\n        )\n        self.val_ids.extend(batch[\"id\"])\n\n    def on_validation_start(self) -> None:\n        self.val_outputs, self.val_ids = [], []\n\n    def validation_epoch_end(self, outputs: List[Any]) -> None:\n\n        metrics = validation_metrics(self.val_outputs, self.val_ids, self.gt_df)\n        print(\"\\n\", metrics)\n\n        self.log_dict(metrics)\n\n        self.val_outputs, self.val_ids = [], []\n\n    def configure_optimizers(self) -> torch.optim.Optimizer:\n        \n        optimizer = torch.optim.Adam(self.parameters(), lr=CFG.lr)\n\n        return optimizer","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:28.066262Z","iopub.execute_input":"2024-04-18T08:09:28.066731Z","iopub.status.idle":"2024-04-18T08:09:29.033921Z","shell.execute_reply.started":"2024-04-18T08:09:28.066685Z","shell.execute_reply":"2024-04-18T08:09:29.032912Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!pip install GPUtil\n\nimport torch\nfrom GPUtil import showUtilization as gpu_usage\nfrom numba import cuda\n\ndef free_gpu_cache():\n    print(\"Initial GPU Usage\")\n    gpu_usage()                             \n\n    torch.cuda.empty_cache()\n\n    cuda.select_device(0)\n    cuda.close()\n    cuda.select_device(0)\n\n    print(\"GPU Usage after emptying the cache\")\n    gpu_usage()\n\nfree_gpu_cache()   ","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Create model and ground truth dataframe","metadata":{}},{"cell_type":"code","source":"gt_chart_type = val_gt_ds[\"chart-type\"]\ngt_x = [json.loads(_) for _ in val_gt_ds[\"x\"]]\ngt_y = [json.loads(_) for _ in val_gt_ds[\"y\"]]\ngt_ids = val_gt_ds[\"id\"]\n\nindex = [f\"{id_}_x\" for id_ in gt_ids] + [f\"{id_}_y\" for id_ in gt_ids]\ngt_df = pd.DataFrame(\n    index=index,\n    data={\n        \"data_series\": gt_x + gt_y,\n        \"chart_type\": gt_chart_type * 2,\n    },\n)\n","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:29.881290Z","iopub.execute_input":"2024-04-18T08:09:29.882184Z","iopub.status.idle":"2024-04-18T08:09:29.901863Z","shell.execute_reply.started":"2024-04-18T08:09:29.882141Z","shell.execute_reply":"2024-04-18T08:09:29.900771Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model = VisionEncoderDecoderModel.from_pretrained(\n    \"naver-clova-ix/donut-base\", config=config\n)","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:30.741032Z","iopub.execute_input":"2024-04-18T08:09:30.742141Z","iopub.status.idle":"2024-04-18T08:09:36.733572Z","shell.execute_reply.started":"2024-04-18T08:09:30.742094Z","shell.execute_reply":"2024-04-18T08:09:36.732585Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model.decoder.resize_token_embeddings(len(processor.tokenizer))","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:36.735911Z","iopub.execute_input":"2024-04-18T08:09:36.736470Z","iopub.status.idle":"2024-04-18T08:09:37.825599Z","shell.execute_reply.started":"2024-04-18T08:09:36.736421Z","shell.execute_reply":"2024-04-18T08:09:37.824446Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Load the model weights from the .pt file\nmodel.load_state_dict(torch.load('/kaggle/input/nb_general_v0416/pytorch/nb_general_epoch1/1/NB_general_finetuned_epoch1_v0416_weights.pt'))\n#model.load_state_dict(torch.load('/kaggle/working/NB_general_finetuned_epoch3_v0417_weights.pt'))","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:48.463120Z","iopub.execute_input":"2024-04-18T08:09:48.463654Z","iopub.status.idle":"2024-04-18T08:09:55.949244Z","shell.execute_reply.started":"2024-04-18T08:09:48.463613Z","shell.execute_reply":"2024-04-18T08:09:55.948147Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"model_module = DonutModelPLModule(processor, model, gt_df, num_training_steps)","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:55.950785Z","iopub.execute_input":"2024-04-18T08:09:55.951063Z","iopub.status.idle":"2024-04-18T08:09:55.956527Z","shell.execute_reply.started":"2024-04-18T08:09:55.951036Z","shell.execute_reply":"2024-04-18T08:09:55.955306Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"CFG.output_path","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:58.900438Z","iopub.execute_input":"2024-04-18T08:09:58.901285Z","iopub.status.idle":"2024-04-18T08:09:58.908429Z","shell.execute_reply.started":"2024-04-18T08:09:58.901225Z","shell.execute_reply":"2024-04-18T08:09:58.907255Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# gt_df.to_csv(CFG.output_path + 'ex_out.csv')","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:09:59.945881Z","iopub.execute_input":"2024-04-18T08:09:59.946981Z","iopub.status.idle":"2024-04-18T08:09:59.951782Z","shell.execute_reply.started":"2024-04-18T08:09:59.946920Z","shell.execute_reply":"2024-04-18T08:09:59.950650Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Train!","metadata":{}},{"cell_type":"code","source":"from pytorch_lightning.callbacks import ModelCheckpoint\nfrom pytorch_lightning.loggers import WandbLogger, TensorBoardLogger\n\ncheckpoint_cb = ModelCheckpoint(CFG.output_path)\n\ntb_logger = TensorBoardLogger('logs/')\n\n# loggers = []\n# if CFG.use_wandb:\n#     import wandb\n#     from kaggle_secrets import UserSecretsClient\n#     user_secrets = UserSecretsClient()\n#     key = user_secrets.get_secret(\"wandb\")\n#     wandb.login(key=key)\n    \n#     loggers.append(WandbLogger(project=\"benetech\"))","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:10:00.799151Z","iopub.execute_input":"2024-04-18T08:10:00.799976Z","iopub.status.idle":"2024-04-18T08:10:00.807049Z","shell.execute_reply.started":"2024-04-18T08:10:00.799933Z","shell.execute_reply":"2024-04-18T08:10:00.805872Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainer = pl.Trainer(\n        accelerator=\"gpu\",\n        devices=1,\n        max_epochs=1,\n        val_check_interval=CFG.val_check_interval,\n        check_val_every_n_epoch=CFG.check_val_every_n_epoch,\n        gradient_clip_val=CFG.gradient_clip_val,\n        precision=16, # if you have tensor cores (t4, v100, a100, etc.) training will be 2x faster\n        num_sanity_val_steps=5,\n        callbacks=[checkpoint_cb], \n        logger=tb_logger\n)","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:10:02.380977Z","iopub.execute_input":"2024-04-18T08:10:02.381922Z","iopub.status.idle":"2024-04-18T08:10:02.462061Z","shell.execute_reply.started":"2024-04-18T08:10:02.381878Z","shell.execute_reply":"2024-04-18T08:10:02.461087Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainer.fit(model_module, train_dataloaders=train_dataloader)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Could also use the checkpoint callback\nmodel_module.model.save_pretrained(CFG.output_path)\nmodel_module.processor.save_pretrained(CFG.output_path)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainer.validate(model_module, dataloaders=DataLoader(\n    train_ds.select(range(24)),\n    batch_size=CFG.batch_size,\n    shuffle=False,\n    collate_fn=collate_fn,\n    num_workers=CFG.num_workers,\n))","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:14:58.076114Z","iopub.execute_input":"2024-04-18T08:14:58.076865Z","iopub.status.idle":"2024-04-18T08:15:39.828962Z","shell.execute_reply.started":"2024-04-18T08:14:58.076818Z","shell.execute_reply":"2024-04-18T08:15:39.827780Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainer.validate(model_module, dataloaders=val_dataloader)","metadata":{"execution":{"iopub.status.busy":"2024-04-18T08:10:12.280807Z","iopub.execute_input":"2024-04-18T08:10:12.281766Z","iopub.status.idle":"2024-04-18T08:13:15.423678Z","shell.execute_reply.started":"2024-04-18T08:10:12.281717Z","shell.execute_reply":"2024-04-18T08:13:15.422443Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Could also use the checkpoint callback\nmodel_module.model.save_pretrained(CFG.output_path)\nmodel_module.processor.save_pretrained(CFG.output_path)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"trainer.save_checkpoint(CFG.output_path + \"NB_general_finetuned_epoch3_v0417.ckpt\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\n\n# Define the file path you want to delete\nfile_to_delete = CFG.output_path + \"NB_general_finetuned_epoch1_v0416.pt\"\n\n# Check if the file exists before attempting to delete it\nif os.path.exists(file_to_delete):\n    # Attempt to delete the file\n    try:\n        os.remove(file_to_delete)\n        print(f\"File '{file_to_delete}' successfully deleted.\")\n    except OSError as e:\n        print(f\"Error: {file_to_delete} : {e.strerror}\")\nelse:\n    print(f\"File '{file_to_delete}' does not exist.\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\n\n# Define the file path you want to delete\nfile_to_delete = CFG.output_path + \"ex_out.csv\"\n\n# Check if the file exists before attempting to delete it\nif os.path.exists(file_to_delete):\n    # Attempt to delete the file\n    try:\n        os.remove(file_to_delete)\n        print(f\"File '{file_to_delete}' successfully deleted.\")\n    except OSError as e:\n        print(f\"Error: {file_to_delete} : {e.strerror}\")\nelse:\n    print(f\"File '{file_to_delete}' does not exist.\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Define the path where you want to save the model weights as a .pt file\nmodel_weights_path = CFG.output_path + \"NB_general_finetuned_epoch5_v0417_weights.pt\"\n\n# Save the model's state_dict as a .pt file\ntorch.save(model_module.model.state_dict(), model_weights_path)","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}