{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "view-in-github"
   },
   "source": [
    "<a href=\"https://colab.research.google.com/github/NielsRogge/Transformers-Tutorials/blob/master/Fine_tuning_LayoutLMForTokenClassification_on_FUNSD.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ngqdEv0rP01q"
   },
   "source": [
    "## Introduction\n",
    "\n",
    "In this notebook, we are going to fine-tune the LayoutLM model by Microsoft Research on the [FUNSD](https://guillaumejaume.github.io/FUNSD/) dataset, which is a collection of annotated form documents. The goal of our model is to learn the annotations of a number of labels (\"question\", \"answer\", \"header\" and \"other\") on those forms, such that it can be used to annotate unseen forms in the future.\n",
    "\n",
    "* Original LayoutLM paper: https://huggingface.co/papers/1912.13318\n",
    "\n",
    "* Original FUNSD paper: https://huggingface.co/papers/1905.13538\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "6K4S2s33ebY0"
   },
   "source": [
    "## Install libraries\n",
    "\n",
    "Currently you have to first install the `unilm` package, and then the `transformers` package (which updates the outdated `transformers` package that is included in the `unilm` package). The reason we also install the `unilm` package is because we need its preprocessing files. I've forked it, and removed some statements which introduced some issues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "5cngOTr6SqEf",
    "outputId": "6c7a2f76-682b-4f93-a3db-59ab010e5ffe"
   },
   "outputs": [],
   "source": [
    "! rm -r unilm\n",
    "! git clone https://github.com/microsoft/unilm.git"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RGMkEG5aRB0D"
   },
   "source": [
    "## Getting the data\n",
    "\n",
    "Here we download the data of the [FUNSD dataset](https://guillaumejaume.github.io/FUNSD/) from the web. This results in a directory called \"data\" being created, which has 2 subdirectories, one for training and one for testing. Each of those has 2 subdirectories in turn, one containing the images as png files and one containing the annotations in json format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "DTFnEZagQm4v",
    "outputId": "97ce03ba-a6bb-4444-8eba-77eceece44e0"
   },
   "outputs": [],
   "source": [
    "! wget https://guillaumejaume.github.io/FUNSD/dataset.zip\n",
    "! unzip dataset.zip && mv dataset data && rm -rf dataset.zip __MACOSX"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "UrNMR64LsJXm"
   },
   "source": [
    "Let's take a look at a training example. For this, we are going to use PIL (Python Image Library)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "eG-eGcj3sNPs",
    "outputId": "69ead0ea-15d6-4d5e-af61-a99a7533d31b"
   },
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 93488 characters]",
      "image/png": "[encoded data, 142508 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=762x1000>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from PIL import Image, ImageDraw, ImageFont\n",
    "import os\n",
    "\n",
    "base_path = \"./data\"\n",
    "\n",
    "image = Image.open(os.path.join(base_path, \"training_data/images/0000971160.png\"))\n",
    "image = image.convert(\"RGB\")\n",
    "image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uAVffmnZyUvw"
   },
   "source": [
    "Now let's plot its corresponding annotations. Basically, if you type `data['form']`, you get a list of all general annotations. Each general annotation has a label, a bounding box, and one or more words, which in also have their own bounding box. The bounding boxes are in [xleft, ytop, xright, ybottom] format.\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "JPKkuJQ4sdZc",
    "outputId": "c95bf306-98bb-4480-cc6b-ebb3aea548b3"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'box': [292, 91, 376, 175], 'text': 'R&D', 'label': 'other', 'words': [{'box': [292, 91, 376, 175], 'text': 'R&D'}], 'linking': [], 'id': 0}\n",
      "{'box': [219, 316, 225, 327], 'text': ':', 'label': 'question', 'words': [{'box': [219, 316, 225, 327], 'text': ':'}], 'linking': [], 'id': 1}\n",
      "{'box': [95, 355, 169, 370], 'text': 'Suggestion:', 'label': 'question', 'words': [{'box': [95, 355, 169, 370], 'text': 'Suggestion:'}], 'linking': [[2, 16]], 'id': 2}\n",
      "{'box': [482, 268, 518, 282], 'text': 'Date:', 'label': 'question', 'words': [{'box': [482, 268, 518, 282], 'text': 'Date:'}], 'linking': [[3, 12]], 'id': 3}\n",
      "{'box': [511, 309, 570, 323], 'text': 'Licensee', 'label': 'answer', 'words': [{'box': [511, 309, 570, 323], 'text': 'Licensee'}], 'linking': [[13, 4]], 'id': 4}\n",
      "{'box': [211, 651, 217, 662], 'text': '', 'label': 'question', 'words': [{'box': [211, 651, 217, 662], 'text': ''}], 'linking': [], 'id': 5}\n",
      "{'box': [461, 605, 483, 619], 'text': 'Yes', 'label': 'question', 'words': [{'box': [461, 605, 483, 619], 'text': 'Yes'}], 'linking': [[19, 6]], 'id': 6}\n",
      "{'box': [545, 603, 563, 617], 'text': 'No', 'label': 'question', 'words': [{'box': [545, 603, 563, 617], 'text': 'No'}], 'linking': [[19, 7]], 'id': 7}\n",
      "{'box': [525, 904, 641, 926], 'text': '597005708', 'label': 'other', 'words': [{'box': [525, 904, 641, 926], 'text': '597005708'}], 'linking': [], 'id': 8}\n",
      "{'text': 'R&D QUALITY IMPROVEMENT SUGGESTION/ SOLUTION FORM', 'box': [256, 201, 423, 230], 'linking': [], 'label': 'header', 'words': [{'text': 'R&D', 'box': [257, 203, 279, 214]}, {'text': 'QUALITY', 'box': [285, 203, 334, 216]}, {'text': 'IMPROVEMENT', 'box': [341, 201, 418, 211]}, {'text': 'SUGGESTION/', 'box': [256, 215, 324, 229]}, {'text': '', 'box': [324, 216, 332, 230]}, {'text': 'SOLUTION', 'box': [331, 214, 387, 228]}, {'text': 'FORM', 'box': [395, 215, 423, 228]}], 'id': 9}\n",
      "{'text': 'Name / Phone Ext. :', 'box': [89, 272, 204, 289], 'linking': [[10, 11]], 'label': 'question', 'words': [{'text': 'Name', 'box': [89, 274, 118, 289]}, {'text': '/', 'box': [117, 274, 127, 288]}, {'text': 'Phone', 'box': [128, 274, 163, 289]}, {'text': 'Ext.', 'box': [169, 272, 196, 287]}, {'text': ':', 'box': [196, 274, 204, 288]}], 'id': 10}\n",
      "{'text': 'M. Hamann P. Harper, P. Martinez', 'box': [215, 271, 451, 287], 'linking': [[10, 11]], 'label': 'answer', 'words': [{'text': 'M.', 'box': [215, 272, 230, 287]}, {'text': 'Hamann', 'box': [237, 272, 287, 286]}, {'text': 'P.', 'box': [293, 272, 307, 286]}, {'text': 'Harper,', 'box': [314, 274, 363, 285]}, {'text': 'P.', 'box': [370, 272, 384, 285]}, {'text': 'Martinez', 'box': [390, 271, 451, 282]}], 'id': 11}\n",
      "{'text': '9/ 3/ 92', 'box': [543, 264, 590, 279], 'linking': [[3, 12]], 'label': 'answer', 'words': [{'text': '9/', 'box': [543, 265, 560, 279]}, {'text': '3/', 'box': [560, 264, 575, 279]}, {'text': '92', 'box': [575, 264, 590, 279]}], 'id': 12}\n",
      "{'text': 'R&D Group:', 'box': [420, 310, 491, 323], 'linking': [[13, 4]], 'label': 'question', 'words': [{'text': 'R&D', 'box': [420, 310, 442, 323]}, {'text': 'Group:', 'box': [448, 310, 491, 323]}], 'id': 13}\n",
      "{'text': 'J. S. Wigand', 'box': [236, 313, 327, 327], 'linking': [[15, 14]], 'label': 'answer', 'words': [{'text': 'J.', 'box': [236, 313, 251, 327]}, {'text': 'S.', 'box': [256, 313, 273, 326]}, {'text': 'Wigand', 'box': [278, 313, 327, 327]}], 'id': 14}\n",
      "{'text': 'Supervisor / Manager', 'box': [91, 316, 218, 331], 'linking': [[15, 14]], 'label': 'question', 'words': [{'text': 'Supervisor', 'box': [91, 316, 161, 330]}, {'text': '/', 'box': [163, 318, 169, 331]}, {'text': 'Manager', 'box': [169, 317, 218, 327]}], 'id': 15}\n",
      "{'text': 'Discontinue coal retention analyses on licensee submitted product samples (Note : Coal Retention testing is not performed by most licensees. Other B&W physical measurements as ends stability and inspection for soft spots in ciparettes are thought to be sufficient measures to assure cigarette physical integrity. The proposed action will increase laboratory productivity . )', 'box': [190, 346, 594, 447], 'linking': [[2, 16]], 'label': 'answer', 'words': [{'text': 'Discontinue', 'box': [190, 355, 268, 366]}, {'text': 'coal', 'box': [274, 353, 303, 366]}, {'text': 'retention', 'box': [309, 352, 375, 365]}, {'text': 'analyses', 'box': [381, 351, 435, 365]}, {'text': 'on', 'box': [443, 352, 458, 363]}, {'text': 'licensee', 'box': [464, 348, 520, 362]}, {'text': 'submitted', 'box': [527, 346, 594, 361]}, {'text': 'product', 'box': [190, 369, 240, 383]}, {'text': 'samples', 'box': [247, 367, 301, 380]}, {'text': '(Note', 'box': [318, 365, 352, 379]}, {'text': ':', 'box': [352, 367, 359, 380]}, {'text': 'Coal', 'box': [373, 366, 402, 376]}, {'text': 'Retention', 'box': [408, 366, 472, 376]}, {'text': 'testing', 'box': [479, 365, 529, 376]}, {'text': 'is', 'box': [536, 363, 549, 374]}, {'text': 'not', 'box': [554, 363, 578, 374]}, {'text': 'performed', 'box': [190, 383, 256, 394]}, {'text': 'by', 'box': [261, 381, 275, 394]}, {'text': 'most', 'box': [282, 383, 311, 393]}, {'text': 'licensees.', 'box': [318, 380, 386, 391]}, {'text': 'Other', 'box': [401, 378, 437, 389]}, {'text': 'B&W', 'box': [443, 378, 465, 389]}, {'text': 'physical', 'box': [471, 377, 528, 391]}, {'text': 'measurements', 'box': [191, 398, 275, 406]}, {'text': 'as', 'box': [282, 397, 297, 405]}, {'text': 'ends', 'box': [304, 394, 332, 405]}, {'text': 'stability', 'box': [339, 394, 402, 405]}, {'text': 'and', 'box': [409, 392, 430, 402]}, {'text': 'inspection', 'box': [437, 392, 508, 403]}, {'text': 'for', 'box': [515, 391, 535, 402]}, {'text': 'soft', 'box': [542, 391, 571, 401]}, {'text': 'spots', 'box': [193, 411, 228, 422]}, {'text': 'in', 'box': [235, 409, 250, 420]}, {'text': 'ciparettes', 'box': [256, 409, 327, 419]}, {'text': 'are', 'box': [332, 408, 352, 418]}, {'text': 'thought', 'box': [360, 406, 410, 419]}, {'text': 'to', 'box': [415, 406, 430, 416]}, {'text': 'be', 'box': [436, 404, 453, 417]}, {'text': 'sufficient', 'box': [458, 405, 529, 415]}, {'text': 'measures', 'box': [535, 405, 592, 415]}, {'text': 'to', 'box': [193, 425, 208, 433]}, {'text': 'assure', 'box': [214, 423, 255, 431]}, {'text': 'cigarette', 'box': [261, 420, 325, 434]}, {'text': 'physical', 'box': [331, 419, 390, 432]}, {'text': 'integrity.', 'box': [395, 418, 463, 431]}, {'text': 'The', 'box': [478, 416, 500, 429]}, {'text': 'proposed', 'box': [506, 418, 566, 431]}, {'text': 'action', 'box': [193, 436, 236, 447]}, {'text': 'will', 'box': [240, 436, 269, 447]}, {'text': 'increase', 'box': [277, 434, 333, 445]}, {'text': 'laboratory', 'box': [339, 433, 410, 446]}, {'text': 'productivity', 'box': [418, 430, 502, 445]}, {'text': '.', 'box': [503, 433, 507, 444]}, {'text': ')', 'box': [508, 430, 514, 444]}], 'id': 16}\n",
      "{'text': 'Suggested Solutions (s) :', 'box': [95, 486, 250, 504], 'linking': [[17, 18]], 'label': 'question', 'words': [{'text': 'Suggested', 'box': [95, 489, 159, 504]}, {'text': 'Solutions', 'box': [165, 487, 222, 501]}, {'text': '(s)', 'box': [223, 486, 241, 503]}, {'text': ':', 'box': [243, 489, 250, 503]}], 'id': 17}\n",
      "{'text': 'Delete coal retention from the list of standard analyses performed on licensee submitted product samples. Special requests for coal retention testing could still be submitted on an exception basis.', 'box': [263, 483, 593, 553], 'linking': [[17, 18]], 'label': 'answer', 'words': [{'text': 'Delete', 'box': [263, 486, 306, 500]}, {'text': 'coal', 'box': [313, 486, 341, 499]}, {'text': 'retention', 'box': [348, 486, 412, 497]}, {'text': 'from', 'box': [416, 485, 447, 498]}, {'text': 'the', 'box': [453, 485, 475, 498]}, {'text': 'list', 'box': [480, 483, 508, 496]}, {'text': 'of', 'box': [515, 483, 532, 494]}, {'text': 'standard', 'box': [536, 483, 593, 494]}, {'text': 'analyses', 'box': [264, 501, 320, 514]}, {'text': 'performed', 'box': [324, 501, 392, 512]}, {'text': 'on', 'box': [397, 501, 412, 511]}, {'text': 'licensee', 'box': [419, 499, 475, 512]}, {'text': 'submitted', 'box': [482, 499, 546, 510]}, {'text': 'product', 'box': [264, 517, 314, 528]}, {'text': 'samples.', 'box': [320, 514, 374, 528]}, {'text': 'Special', 'box': [390, 513, 439, 526]}, {'text': 'requests', 'box': [446, 513, 502, 524]}, {'text': 'for', 'box': [508, 511, 530, 522]}, {'text': 'coal', 'box': [538, 510, 566, 523]}, {'text': 'retention', 'box': [263, 529, 330, 540]}, {'text': 'testing', 'box': [335, 527, 387, 540]}, {'text': 'could', 'box': [390, 527, 428, 538]}, {'text': 'still', 'box': [433, 525, 468, 536]}, {'text': 'be', 'box': [473, 525, 488, 535]}, {'text': 'submitted', 'box': [496, 524, 560, 537]}, {'text': 'on', 'box': [566, 524, 584, 537]}, {'text': 'an', 'box': [264, 543, 281, 553]}, {'text': 'exception', 'box': [286, 539, 350, 553]}, {'text': 'basis.', 'box': [355, 541, 397, 551]}], 'id': 18}\n",
      "{'text': 'Have you contacted your Manager/ Supervisor?', 'box': [96, 608, 398, 624], 'linking': [[19, 6], [19, 7]], 'label': 'header', 'words': [{'text': 'Have', 'box': [96, 612, 127, 623]}, {'text': 'you', 'box': [131, 613, 156, 624]}, {'text': 'contacted', 'box': [161, 612, 225, 623]}, {'text': 'your', 'box': [229, 610, 260, 623]}, {'text': 'Manager/', 'box': [264, 609, 314, 622]}, {'text': '', 'box': [314, 608, 322, 622]}, {'text': 'Supervisor?', 'box': [323, 608, 398, 621]}], 'id': 19}\n",
      "{'text': 'Manager Comments:', 'box': [98, 651, 211, 665], 'linking': [[20, 21], [20, 22]], 'label': 'question', 'words': [{'text': 'Manager', 'box': [98, 654, 150, 665]}, {'text': 'Comments:', 'box': [154, 651, 211, 664]}], 'id': 20}\n",
      "{'text': 'Manager, please contact suggester and forward', 'box': [232, 644, 547, 662], 'linking': [[20, 21]], 'label': 'answer', 'words': [{'text': 'Manager,', 'box': [232, 648, 288, 662]}, {'text': 'please', 'box': [296, 649, 338, 662]}, {'text': 'contact', 'box': [344, 648, 394, 662]}, {'text': 'suggester', 'box': [401, 648, 464, 661]}, {'text': 'and', 'box': [469, 647, 491, 658]}, {'text': 'forward', 'box': [497, 644, 547, 657]}], 'id': 21}\n",
      "{'text': 'comments to the Quality Council.', 'box': [99, 662, 323, 677], 'linking': [[20, 22]], 'label': 'answer', 'words': [{'text': 'comments', 'box': [99, 666, 155, 677]}, {'text': 'to', 'box': [162, 665, 177, 676]}, {'text': 'the', 'box': [183, 665, 205, 675]}, {'text': 'Quality', 'box': [211, 663, 261, 676]}, {'text': 'Council.', 'box': [267, 662, 323, 676]}], 'id': 22}\n",
      "{'text': 'qip . wp', 'box': [102, 823, 145, 838], 'linking': [], 'label': 'other', 'words': [{'text': 'qip', 'box': [102, 824, 123, 837]}, {'text': '.', 'box': [124, 824, 130, 838]}, {'text': 'wp', 'box': [130, 823, 145, 837]}], 'id': 23}\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "\n",
    "with open(os.path.join(base_path, \"training_data/annotations/0000971160.json\")) as f:\n",
    "    data = json.load(f)\n",
    "\n",
    "for annotation in data[\"form\"]:\n",
    "    print(annotation)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Hs4L3S5a2Gfb"
   },
   "source": [
    "The PIL library has a handy ImageDraw module, which -you guessed it- allows to draw things (such as rectangles) on an image:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "gWaHFM_LtKPP",
    "outputId": "c498e560-035f-4170-b0b9-85ba3956711c"
   },
   "outputs": [
    {
     "data": {
      "image/jpeg": "[encoded data, 126280 characters]",
      "image/png": "[encoded data, 164500 characters]",
      "text/plain": [
       "<PIL.Image.Image image mode=RGB size=762x1000>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "draw = ImageDraw.Draw(image, \"RGBA\")\n",
    "\n",
    "font = ImageFont.load_default()\n",
    "\n",
    "label2color = {\"question\": \"blue\", \"answer\": \"green\", \"header\": \"orange\", \"other\": \"violet\"}\n",
    "\n",
    "for annotation in data[\"form\"]:\n",
    "    label = annotation[\"label\"]\n",
    "    general_box = annotation[\"box\"]\n",
    "    draw.rectangle(general_box, outline=label2color[label], width=2)\n",
    "    draw.text((general_box[0] + 10, general_box[1] - 10), label, fill=label2color[label], font=font)\n",
    "    words = annotation[\"words\"]\n",
    "    for word in words:\n",
    "        box = word[\"box\"]\n",
    "        draw.rectangle(box, outline=label2color[label], width=1)\n",
    "\n",
    "image"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "uyWQNLSCRJN7"
   },
   "source": [
    "## Preprocessing the data\n",
    "\n",
    "Next, we need to turn the document images into individual tokens and corresponding labels (BIOES format, see further). We do this both for the training and test datasets. Make sure to run this from the `/content` directory:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "4DWRyOR9RuY6",
    "outputId": "4215a24b-8049-4b1a-a23f-5aaa48e14083"
   },
   "outputs": [],
   "source": [
    "! python unilm/layoutlm/deprecated/examples/seq_labeling/preprocess.py --data_dir data/training_data/annotations \\\n",
    "                                                      --data_split train \\\n",
    "                                                      --output_dir data \\\n",
    "                                                      --model_name_or_path microsoft/layoutlm-base-uncased \\\n",
    "                                                      --max_len 510\n",
    "\n",
    "! python unilm/layoutlm/deprecated/examples/seq_labeling/preprocess.py --data_dir data/testing_data/annotations \\\n",
    "                                                      --data_split test \\\n",
    "                                                      --output_dir data \\\n",
    "                                                      --model_name_or_path microsoft/layoutlm-base-uncased \\\n",
    "                                                      --max_len 510"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "gc4Cu0ZyO5M_"
   },
   "source": [
    "Next, we create a labels.txt file that contains the unique labels of the FUNSD dataset:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "8iGOU0s3UR2u"
   },
   "outputs": [],
   "source": [
    "! cat data/train.txt | cut -d$'\\t' -f 2 | grep -v \"^$\"| sort | uniq > data/labels.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "mC9FhkG9U8yg"
   },
   "source": [
    "## Define a PyTorch dataset\n",
    "\n",
    "First, we create a list containing the unique labels based on `data/labels.txt` (run this from the content directory):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "id": "675rRa0QXnMp"
   },
   "outputs": [],
   "source": [
    "from torch.nn import CrossEntropyLoss\n",
    "\n",
    "\n",
    "def get_labels(path):\n",
    "    with open(path, \"r\") as f:\n",
    "        labels = f.read().splitlines()\n",
    "    if \"O\" not in labels:\n",
    "        labels = [\"O\"] + labels\n",
    "    return labels\n",
    "\n",
    "\n",
    "labels = get_labels(\"data/labels.txt\")\n",
    "num_labels = len(labels)\n",
    "label_map = {i: label for i, label in enumerate(labels)}\n",
    "# Use cross entropy ignore index as padding label id so that only real label ids contribute to the loss later\n",
    "pad_token_label_id = CrossEntropyLoss().ignore_index"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kZ2LGEsez2u2"
   },
   "source": [
    "We can see that the dataset uses the so-called BIOES annotation scheme to annotate the tokens. This means that a given token can be either at the beginning (B), inside (I), outside (O), at the end (E) or start (S) of a given entity. Entities include ANSWER, QUESTION, HEADER and OTHER: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "_-qXLkP9Yq_L",
    "outputId": "32ab46a4-4cf0-400c-816b-570f950035ec"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['B-ANSWER', 'B-HEADER', 'B-QUESTION', 'E-ANSWER', 'E-HEADER', 'E-QUESTION', 'I-ANSWER', 'I-HEADER', 'I-QUESTION', 'O', 'S-ANSWER', 'S-HEADER', 'S-QUESTION']\n"
     ]
    }
   ],
   "source": [
    "print(labels)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "9_ck0ZFfZInR"
   },
   "source": [
    "Next, we can create a PyTorch dataset and corresponding dataloader (both for training and evaluation):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import os\n",
    "\n",
    "import torch\n",
    "from torch.utils.data import Dataset\n",
    "\n",
    "logger = logging.getLogger(__name__)\n",
    "\n",
    "\n",
    "class FunsdDataset(Dataset):\n",
    "    def __init__(self, args, tokenizer, labels, pad_token_label_id, mode):\n",
    "        if args.local_rank not in [-1, 0] and mode == \"train\":\n",
    "            torch.distributed.barrier()  # Make sure only the first process in distributed training process the dataset, and the others will use the cache\n",
    "\n",
    "        # Load data features from cache or dataset file\n",
    "        cached_features_file = os.path.join(\n",
    "            args.data_dir,\n",
    "            \"cached_{}_{}_{}\".format(\n",
    "                mode,\n",
    "                list(filter(None, args.model_name_or_path.split(\"/\"))).pop(),\n",
    "                str(args.max_seq_length),\n",
    "            ),\n",
    "        )\n",
    "        if os.path.exists(cached_features_file) and not args.overwrite_cache:\n",
    "            logger.info(\"Loading features from cached file %s\", cached_features_file)\n",
    "            features = torch.load(cached_features_file)\n",
    "        else:\n",
    "            logger.info(\"Creating features from dataset file at %s\", args.data_dir)\n",
    "            examples = read_examples_from_file(args.data_dir, mode)\n",
    "            features = convert_examples_to_features(\n",
    "                examples,\n",
    "                labels,\n",
    "                args.max_seq_length,\n",
    "                tokenizer,\n",
    "                cls_token_at_end=bool(args.model_type in [\"xlnet\"]),\n",
    "                # xlnet has a cls token at the end\n",
    "                cls_token=tokenizer.cls_token,\n",
    "                cls_token_segment_id=2 if args.model_type in [\"xlnet\"] else 0,\n",
    "                sep_token=tokenizer.sep_token,\n",
    "                sep_token_extra=bool(args.model_type in [\"roberta\"]),\n",
    "                # roberta uses an extra separator b/w pairs of sentences, cf. github.com/pytorch/fairseq/commit/1684e166e3da03f5b600dbb7855cb98ddfcd0805\n",
    "                pad_on_left=bool(args.model_type in [\"xlnet\"]),\n",
    "                # pad on the left for xlnet\n",
    "                pad_token=tokenizer.convert_tokens_to_ids([tokenizer.pad_token])[0],\n",
    "                pad_token_segment_id=4 if args.model_type in [\"xlnet\"] else 0,\n",
    "                pad_token_label_id=pad_token_label_id,\n",
    "            )\n",
    "            # if args.local_rank in [-1, 0]:\n",
    "            # logger.info(\"Saving features into cached file %s\", cached_features_file)\n",
    "            # torch.save(features, cached_features_file)\n",
    "\n",
    "        if args.local_rank == 0 and mode == \"train\":\n",
    "            torch.distributed.barrier()  # Make sure only the first process in distributed training process the dataset, and the others will use the cache\n",
    "\n",
    "        self.features = features\n",
    "        # Convert to Tensors and build dataset\n",
    "        self.all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)\n",
    "        self.all_input_mask = torch.tensor([f.input_mask for f in features], dtype=torch.long)\n",
    "        self.all_segment_ids = torch.tensor([f.segment_ids for f in features], dtype=torch.long)\n",
    "        self.all_label_ids = torch.tensor([f.label_ids for f in features], dtype=torch.long)\n",
    "        self.all_bboxes = torch.tensor([f.boxes for f in features], dtype=torch.long)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.features)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        return (\n",
    "            self.all_input_ids[index],\n",
    "            self.all_input_mask[index],\n",
    "            self.all_segment_ids[index],\n",
    "            self.all_label_ids[index],\n",
    "            self.all_bboxes[index],\n",
    "        )\n",
    "\n",
    "\n",
    "class InputExample(object):\n",
    "    \"\"\"A single training/test example for token classification.\"\"\"\n",
    "\n",
    "    def __init__(self, guid, words, labels, boxes, actual_bboxes, file_name, page_size):\n",
    "        \"\"\"Constructs a InputExample.\n",
    "\n",
    "        Args:\n",
    "            guid: Unique id for the example.\n",
    "            words: list. The words of the sequence.\n",
    "            labels: (Optional) list. The labels for each word of the sequence. This should be\n",
    "            specified for train and dev examples, but not for test examples.\n",
    "        \"\"\"\n",
    "        self.guid = guid\n",
    "        self.words = words\n",
    "        self.labels = labels\n",
    "        self.boxes = boxes\n",
    "        self.actual_bboxes = actual_bboxes\n",
    "        self.file_name = file_name\n",
    "        self.page_size = page_size\n",
    "\n",
    "\n",
    "class InputFeatures(object):\n",
    "    \"\"\"A single set of features of data.\"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        input_ids,\n",
    "        input_mask,\n",
    "        segment_ids,\n",
    "        label_ids,\n",
    "        boxes,\n",
    "        actual_bboxes,\n",
    "        file_name,\n",
    "        page_size,\n",
    "    ):\n",
    "        assert (\n",
    "            0 <= all(boxes) <= 1000\n",
    "        ), \"Error with input bbox ({}): the coordinate value is not between 0 and 1000\".format(boxes)\n",
    "        self.input_ids = input_ids\n",
    "        self.input_mask = input_mask\n",
    "        self.segment_ids = segment_ids\n",
    "        self.label_ids = label_ids\n",
    "        self.boxes = boxes\n",
    "        self.actual_bboxes = actual_bboxes\n",
    "        self.file_name = file_name\n",
    "        self.page_size = page_size\n",
    "\n",
    "\n",
    "def read_examples_from_file(data_dir, mode):\n",
    "    file_path = os.path.join(data_dir, \"{}.txt\".format(mode))\n",
    "    box_file_path = os.path.join(data_dir, \"{}_box.txt\".format(mode))\n",
    "    image_file_path = os.path.join(data_dir, \"{}_image.txt\".format(mode))\n",
    "    guid_index = 1\n",
    "    examples = []\n",
    "    with open(file_path, encoding=\"utf-8\") as f, open(box_file_path, encoding=\"utf-8\") as fb, open(\n",
    "        image_file_path, encoding=\"utf-8\"\n",
    "    ) as fi:\n",
    "        words = []\n",
    "        boxes = []\n",
    "        actual_bboxes = []\n",
    "        file_name = None\n",
    "        page_size = None\n",
    "        labels = []\n",
    "        for line, bline, iline in zip(f, fb, fi):\n",
    "            if line.startswith(\"-DOCSTART-\") or line == \"\" or line == \"\\n\":\n",
    "                if words:\n",
    "                    examples.append(\n",
    "                        InputExample(\n",
    "                            guid=\"{}-{}\".format(mode, guid_index),\n",
    "                            words=words,\n",
    "                            labels=labels,\n",
    "                            boxes=boxes,\n",
    "                            actual_bboxes=actual_bboxes,\n",
    "                            file_name=file_name,\n",
    "                            page_size=page_size,\n",
    "                        )\n",
    "                    )\n",
    "                    guid_index += 1\n",
    "                    words = []\n",
    "                    boxes = []\n",
    "                    actual_bboxes = []\n",
    "                    file_name = None\n",
    "                    page_size = None\n",
    "                    labels = []\n",
    "            else:\n",
    "                splits = line.split(\"\\t\")\n",
    "                bsplits = bline.split(\"\\t\")\n",
    "                isplits = iline.split(\"\\t\")\n",
    "                assert len(splits) == 2\n",
    "                assert len(bsplits) == 2\n",
    "                assert len(isplits) == 4\n",
    "                assert splits[0] == bsplits[0]\n",
    "                words.append(splits[0])\n",
    "                if len(splits) > 1:\n",
    "                    labels.append(splits[-1].replace(\"\\n\", \"\"))\n",
    "                    box = bsplits[-1].replace(\"\\n\", \"\")\n",
    "                    box = [int(b) for b in box.split()]\n",
    "                    boxes.append(box)\n",
    "                    actual_bbox = [int(b) for b in isplits[1].split()]\n",
    "                    actual_bboxes.append(actual_bbox)\n",
    "                    page_size = [int(i) for i in isplits[2].split()]\n",
    "                    file_name = isplits[3].strip()\n",
    "                else:\n",
    "                    # Examples could have no label for mode = \"test\"\n",
    "                    labels.append(\"O\")\n",
    "        if words:\n",
    "            examples.append(\n",
    "                InputExample(\n",
    "                    guid=\"%s-%d\".format(mode, guid_index),\n",
    "                    words=words,\n",
    "                    labels=labels,\n",
    "                    boxes=boxes,\n",
    "                    actual_bboxes=actual_bboxes,\n",
    "                    file_name=file_name,\n",
    "                    page_size=page_size,\n",
    "                )\n",
    "            )\n",
    "    return examples\n",
    "\n",
    "\n",
    "def convert_examples_to_features(\n",
    "    examples,\n",
    "    label_list,\n",
    "    max_seq_length,\n",
    "    tokenizer,\n",
    "    cls_token_at_end=False,\n",
    "    cls_token=\"[CLS]\",\n",
    "    cls_token_segment_id=1,\n",
    "    sep_token=\"[SEP]\",\n",
    "    sep_token_extra=False,\n",
    "    pad_on_left=False,\n",
    "    pad_token=0,\n",
    "    cls_token_box=[0, 0, 0, 0],\n",
    "    sep_token_box=[1000, 1000, 1000, 1000],\n",
    "    pad_token_box=[0, 0, 0, 0],\n",
    "    pad_token_segment_id=0,\n",
    "    pad_token_label_id=-1,\n",
    "    sequence_a_segment_id=0,\n",
    "    mask_padding_with_zero=True,\n",
    "):\n",
    "    \"\"\"Loads a data file into a list of `InputBatch`s\n",
    "    `cls_token_at_end` define the location of the CLS token:\n",
    "        - False (Default, BERT/XLM pattern): [CLS] + A + [SEP] + B + [SEP]\n",
    "        - True (XLNet/GPT pattern): A + [SEP] + B + [SEP] + [CLS]\n",
    "    `cls_token_segment_id` define the segment id associated to the CLS token (0 for BERT, 2 for XLNet)\n",
    "    \"\"\"\n",
    "\n",
    "    label_map = {label: i for i, label in enumerate(label_list)}\n",
    "\n",
    "    features = []\n",
    "    for ex_index, example in enumerate(examples):\n",
    "        file_name = example.file_name\n",
    "        page_size = example.page_size\n",
    "        width, height = page_size\n",
    "        if ex_index % 10000 == 0:\n",
    "            logger.info(\"Writing example %d of %d\", ex_index, len(examples))\n",
    "\n",
    "        tokens = []\n",
    "        token_boxes = []\n",
    "        actual_bboxes = []\n",
    "        label_ids = []\n",
    "        for word, label, box, actual_bbox in zip(example.words, example.labels, example.boxes, example.actual_bboxes):\n",
    "            word_tokens = tokenizer.tokenize(word)\n",
    "            tokens.extend(word_tokens)\n",
    "            token_boxes.extend([box] * len(word_tokens))\n",
    "            actual_bboxes.extend([actual_bbox] * len(word_tokens))\n",
    "            # Use the real label id for the first token of the word, and padding ids for the remaining tokens\n",
    "            label_ids.extend([label_map[label]] + [pad_token_label_id] * (len(word_tokens) - 1))\n",
    "\n",
    "        # Account for [CLS] and [SEP] with \"- 2\" and with \"- 3\" for RoBERTa.\n",
    "        special_tokens_count = 3 if sep_token_extra else 2\n",
    "        if len(tokens) > max_seq_length - special_tokens_count:\n",
    "            tokens = tokens[: (max_seq_length - special_tokens_count)]\n",
    "            token_boxes = token_boxes[: (max_seq_length - special_tokens_count)]\n",
    "            actual_bboxes = actual_bboxes[: (max_seq_length - special_tokens_count)]\n",
    "            label_ids = label_ids[: (max_seq_length - special_tokens_count)]\n",
    "\n",
    "        # The convention in BERT is:\n",
    "        # (a) For sequence pairs:\n",
    "        #  tokens:   [CLS] is this jack ##son ##ville ? [SEP] no it is not . [SEP]\n",
    "        #  type_ids:   0   0  0    0    0     0       0   0   1  1  1  1   1   1\n",
    "        # (b) For single sequences:\n",
    "        #  tokens:   [CLS] the dog is hairy . [SEP]\n",
    "        #  type_ids:   0   0   0   0  0     0   0\n",
    "        #\n",
    "        # Where \"type_ids\" are used to indicate whether this is the first\n",
    "        # sequence or the second sequence. The embedding vectors for `type=0` and\n",
    "        # `type=1` were learned during pre-training and are added to the wordpiece\n",
    "        # embedding vector (and position vector). This is not *strictly* necessary\n",
    "        # since the [SEP] token unambiguously separates the sequences, but it makes\n",
    "        # it easier for the model to learn the concept of sequences.\n",
    "        #\n",
    "        # For classification tasks, the first vector (corresponding to [CLS]) is\n",
    "        # used as as the \"sentence vector\". Note that this only makes sense because\n",
    "        # the entire model is fine-tuned.\n",
    "        tokens += [sep_token]\n",
    "        token_boxes += [sep_token_box]\n",
    "        actual_bboxes += [[0, 0, width, height]]\n",
    "        label_ids += [pad_token_label_id]\n",
    "        if sep_token_extra:\n",
    "            # roberta uses an extra separator b/w pairs of sentences\n",
    "            tokens += [sep_token]\n",
    "            token_boxes += [sep_token_box]\n",
    "            actual_bboxes += [[0, 0, width, height]]\n",
    "            label_ids += [pad_token_label_id]\n",
    "        segment_ids = [sequence_a_segment_id] * len(tokens)\n",
    "\n",
    "        if cls_token_at_end:\n",
    "            tokens += [cls_token]\n",
    "            token_boxes += [cls_token_box]\n",
    "            actual_bboxes += [[0, 0, width, height]]\n",
    "            label_ids += [pad_token_label_id]\n",
    "            segment_ids += [cls_token_segment_id]\n",
    "        else:\n",
    "            tokens = [cls_token] + tokens\n",
    "            token_boxes = [cls_token_box] + token_boxes\n",
    "            actual_bboxes = [[0, 0, width, height]] + actual_bboxes\n",
    "            label_ids = [pad_token_label_id] + label_ids\n",
    "            segment_ids = [cls_token_segment_id] + segment_ids\n",
    "\n",
    "        input_ids = tokenizer.convert_tokens_to_ids(tokens)\n",
    "\n",
    "        # The mask has 1 for real tokens and 0 for padding tokens. Only real\n",
    "        # tokens are attended to.\n",
    "        input_mask = [1 if mask_padding_with_zero else 0] * len(input_ids)\n",
    "\n",
    "        # Zero-pad up to the sequence length.\n",
    "        padding_length = max_seq_length - len(input_ids)\n",
    "        if pad_on_left:\n",
    "            input_ids = ([pad_token] * padding_length) + input_ids\n",
    "            input_mask = ([0 if mask_padding_with_zero else 1] * padding_length) + input_mask\n",
    "            segment_ids = ([pad_token_segment_id] * padding_length) + segment_ids\n",
    "            label_ids = ([pad_token_label_id] * padding_length) + label_ids\n",
    "            token_boxes = ([pad_token_box] * padding_length) + token_boxes\n",
    "        else:\n",
    "            input_ids += [pad_token] * padding_length\n",
    "            input_mask += [0 if mask_padding_with_zero else 1] * padding_length\n",
    "            segment_ids += [pad_token_segment_id] * padding_length\n",
    "            label_ids += [pad_token_label_id] * padding_length\n",
    "            token_boxes += [pad_token_box] * padding_length\n",
    "\n",
    "        assert len(input_ids) == max_seq_length\n",
    "        assert len(input_mask) == max_seq_length\n",
    "        assert len(segment_ids) == max_seq_length\n",
    "        assert len(label_ids) == max_seq_length\n",
    "        assert len(token_boxes) == max_seq_length\n",
    "\n",
    "        if ex_index < 5:\n",
    "            logger.info(\"*** Example ***\")\n",
    "            logger.info(\"guid: %s\", example.guid)\n",
    "            logger.info(\"tokens: %s\", \" \".join([str(x) for x in tokens]))\n",
    "            logger.info(\"input_ids: %s\", \" \".join([str(x) for x in input_ids]))\n",
    "            logger.info(\"input_mask: %s\", \" \".join([str(x) for x in input_mask]))\n",
    "            logger.info(\"segment_ids: %s\", \" \".join([str(x) for x in segment_ids]))\n",
    "            logger.info(\"label_ids: %s\", \" \".join([str(x) for x in label_ids]))\n",
    "            logger.info(\"boxes: %s\", \" \".join([str(x) for x in token_boxes]))\n",
    "            logger.info(\"actual_bboxes: %s\", \" \".join([str(x) for x in actual_bboxes]))\n",
    "\n",
    "        features.append(\n",
    "            InputFeatures(\n",
    "                input_ids=input_ids,\n",
    "                input_mask=input_mask,\n",
    "                segment_ids=segment_ids,\n",
    "                label_ids=label_ids,\n",
    "                boxes=token_boxes,\n",
    "                actual_bboxes=actual_bboxes,\n",
    "                file_name=file_name,\n",
    "                page_size=page_size,\n",
    "            )\n",
    "        )\n",
    "    return features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "id": "HUJftzeBWh2S"
   },
   "outputs": [],
   "source": [
    "from transformers import LayoutLMTokenizer\n",
    "\n",
    "# from .unilm.layoutlm.data.funsd import FunsdDataset, InputFeatures\n",
    "from torch.utils.data import DataLoader, RandomSampler, SequentialSampler\n",
    "\n",
    "batch_size = 16\n",
    "args = {\n",
    "    \"local_rank\": -1,\n",
    "    \"overwrite_cache\": True,\n",
    "    \"data_dir\": \"data/\",\n",
    "    \"model_name_or_path\": \"microsoft/layoutlm-base-uncased\",\n",
    "    \"max_seq_length\": 512,\n",
    "    \"model_type\": \"layoutlm\",\n",
    "}\n",
    "\n",
    "\n",
    "# class to turn the keys of a dict into attributes (thanks Stackoverflow)\n",
    "class AttrDict(dict):\n",
    "    def __init__(self, *args, **kwargs):\n",
    "        super(AttrDict, self).__init__(*args, **kwargs)\n",
    "        self.__dict__ = self\n",
    "\n",
    "\n",
    "args = AttrDict(args)\n",
    "\n",
    "tokenizer = LayoutLMTokenizer.from_pretrained(\"microsoft/layoutlm-base-uncased\")\n",
    "\n",
    "# the LayoutLM authors already defined a specific FunsdDataset, so we are going to use this here\n",
    "train_dataset = FunsdDataset(args, tokenizer, labels, pad_token_label_id, mode=\"train\")\n",
    "train_sampler = RandomSampler(train_dataset)\n",
    "train_dataloader = DataLoader(train_dataset, sampler=train_sampler, batch_size=batch_size)\n",
    "\n",
    "eval_dataset = FunsdDataset(args, tokenizer, labels, pad_token_label_id, mode=\"test\")\n",
    "eval_sampler = SequentialSampler(eval_dataset)\n",
    "eval_dataloader = DataLoader(eval_dataset, sampler=eval_sampler, batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "18NMUBzgOdqu",
    "outputId": "eef47b70-3a9a-4b19-be6b-95900c58337b"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(train_dataloader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "toFjxtn71B1U",
    "outputId": "f4651896-cafc-449a-98b4-c81f41177e6d"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(eval_dataloader)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 137
    },
    "id": "RhINSBw9I24G",
    "outputId": "28738ce2-617c-47d3-b8c9-f949d3066d60"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'[CLS] account agency ss : date 1 / 31 / 88 insert signed 91581919 4th new albany tribune advertisements. bookkeeper mar 08 reco affidavit of performance from newspaper new albany tribune harley davidson cigarettes lorillard media service state of county of indiana ) floyd ) before me a notary public, personally appeared holly inzer who being duly sworn, says that ( he ) ( she ) is of the abovementioned newspaper and that display ads for the above account were made through the aforesaid newspaper during the month of january, 1988 as follows : column inches exclusive advertising for harley davidson cigarettes we hereby certify charges shown above on dates per attached bill are true and correct as billed to the account in upper right hand corner of the affidavit and are exclusive sworn to and subscribed before me this day of march, 1988 in testimony whereof i have set my hand and seal the day and year aforesaid. my commission expires : 2 - 9 - 90 betty j. murphy ( notary public ) [SEP] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD] [PAD]'"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "batch = next(iter(train_dataloader))\n",
    "input_ids = batch[0][0]\n",
    "tokenizer.decode(input_ids)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "66cEmLDoUFcm"
   },
   "source": [
    "## Define and fine-tune the model\n",
    "\n",
    "As this is a sequence labeling task, we are going to load `LayoutLMForTokenClassification` (the base sized model) from the hub. We are going to fine-tune it on a downstream task, namely FUNSD."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "LoraConfig(task_type=<TaskType.TOKEN_CLS: 'TOKEN_CLS'>, peft_type=<PeftType.LORA: 'LORA'>, auto_mapping=None, base_model_name_or_path=None, revision=None, inference_mode=False, r=16, target_modules=None, exclude_modules=None, lora_alpha=16, lora_dropout=0.1, fan_in_fan_out=False, bias='all', use_rslora=False, modules_to_save=None, init_lora_weights=True, layers_to_transform=None, layers_pattern=None, rank_pattern={}, alpha_pattern={}, megatron_config=None, megatron_core='megatron.core', trainable_token_indices=None, loftq_config={}, eva_config=None, corda_config=None, use_dora=False, use_qalora=False, qalora_group_size=16, layer_replication=None, runtime_config=LoraRuntimeConfig(ephemeral_gpu_offload=False), lora_bias=False, target_parameters=None)"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from peft import get_peft_config, PeftModel, get_peft_model, LoraConfig, TaskType\n",
    "\n",
    "peft_config = LoraConfig(\n",
    "    task_type=TaskType.TOKEN_CLS, inference_mode=False, r=16, lora_alpha=16, lora_dropout=0.1, bias=\"all\"\n",
    ")\n",
    "peft_config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000,
     "referenced_widgets": [
      "da1094982d044ab28eb0effebbfcbb78",
      "513e00b619924f5693259cd919a927ab",
      "63e819a04f6e4829838c0e30e65516ed",
      "d1c3e1a66db04227a74ef8d6481d6daf",
      "c75f0da13a1e4dbe94800711d55390a6",
      "31642aacae2a44879960da09f938ecc4",
      "138a6b922e454ebbaeb315ecd5f476b8",
      "a8126ba98376402888e9ba344cf1c538"
     ]
    },
    "id": "xIdOsFBiTsuw",
    "outputId": "95e8811c-025a-41a0-9d03-4285a17f2a9b"
   },
   "outputs": [],
   "source": [
    "from transformers import LayoutLMForTokenClassification\n",
    "import torch\n",
    "from transformers import set_seed\n",
    "\n",
    "seed = 100\n",
    "set_seed(seed)\n",
    "device = torch.accelerator.current_accelerator().type if hasattr(torch, \"accelerator\") else \"cuda\"\n",
    "\n",
    "model = LayoutLMForTokenClassification.from_pretrained(\"microsoft/layoutlm-base-uncased\", num_labels=num_labels)\n",
    "model = get_peft_model(model, peft_config)\n",
    "model.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(model.model.layoutlm.encoder.layer[0].attention.self.query.weight)\n",
    "print(model.model.layoutlm.encoder.layer[0].attention.self.query.lora_A.default.weight)\n",
    "print(model.model.classifier.weight)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "3weFr_pz1mla"
   },
   "source": [
    "Now we can start training:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "Yu0qePs2cRKo",
    "outputId": "cdbb9a03-eb9b-4740-bbe3-da06b9192bae"
   },
   "outputs": [],
   "source": [
    "from transformers import get_linear_schedule_with_warmup\n",
    "from tqdm import tqdm\n",
    "\n",
    "num_train_epochs = 100\n",
    "\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=3e-3)\n",
    "lr_scheduler = get_linear_schedule_with_warmup(\n",
    "    optimizer=optimizer,\n",
    "    num_warmup_steps=0.06 * (len(train_dataloader) * num_train_epochs),\n",
    "    num_training_steps=(len(train_dataloader) * num_train_epochs),\n",
    ")\n",
    "\n",
    "\n",
    "global_step = 0\n",
    "\n",
    "t_total = len(train_dataloader) * num_train_epochs  # total number of training steps\n",
    "\n",
    "# put the model in training mode\n",
    "model.train()\n",
    "for epoch in range(num_train_epochs):\n",
    "    for batch in tqdm(train_dataloader, desc=\"Training\"):\n",
    "        input_ids = batch[0].to(device)\n",
    "        bbox = batch[4].to(device)\n",
    "        attention_mask = batch[1].to(device)\n",
    "        token_type_ids = batch[2].to(device)\n",
    "        labels = batch[3].to(device)\n",
    "\n",
    "        # forward pass\n",
    "        outputs = model(\n",
    "            input_ids=input_ids, bbox=bbox, attention_mask=attention_mask, token_type_ids=token_type_ids, labels=labels\n",
    "        )\n",
    "        loss = outputs.loss\n",
    "\n",
    "        # print loss every 100 steps\n",
    "        if global_step % 10 == 0:\n",
    "            print(f\"Loss after {global_step} steps: {loss.item()}\")\n",
    "\n",
    "        # backward pass to get the gradients\n",
    "        loss.backward()\n",
    "\n",
    "        # print(\"Gradients on classification head:\")\n",
    "        # print(model.classifier.weight.grad[6,:].sum())\n",
    "\n",
    "        # update\n",
    "        optimizer.step()\n",
    "        lr_scheduler.step()\n",
    "        optimizer.zero_grad()\n",
    "        global_step += 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "u1rNslap5Y3N",
    "outputId": "877183d4-1d29-4d09-bd3a-0e5f88611dc8"
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating: 100%|██████████| 4/4 [00:00<00:00, 10.90it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'loss': 1.467050313949585, 'precision': 0.7295341474445952, 'recall': 0.806903451725863, 'f1': 0.7662707838479811}\n"
     ]
    }
   ],
   "source": [
    "import numpy as np\n",
    "from seqeval.metrics import (\n",
    "    classification_report,\n",
    "    f1_score,\n",
    "    precision_score,\n",
    "    recall_score,\n",
    ")\n",
    "\n",
    "eval_loss = 0.0\n",
    "nb_eval_steps = 0\n",
    "preds = None\n",
    "out_label_ids = None\n",
    "\n",
    "# put model in evaluation mode\n",
    "model.eval()\n",
    "for batch in tqdm(eval_dataloader, desc=\"Evaluating\"):\n",
    "    with torch.no_grad():\n",
    "        input_ids = batch[0].to(device)\n",
    "        bbox = batch[4].to(device)\n",
    "        attention_mask = batch[1].to(device)\n",
    "        token_type_ids = batch[2].to(device)\n",
    "        labels = batch[3].to(device)\n",
    "\n",
    "        # forward pass\n",
    "        outputs = model(\n",
    "            input_ids=input_ids, bbox=bbox, attention_mask=attention_mask, token_type_ids=token_type_ids, labels=labels\n",
    "        )\n",
    "        # get the loss and logits\n",
    "        tmp_eval_loss = outputs.loss\n",
    "        logits = outputs.logits\n",
    "\n",
    "        eval_loss += tmp_eval_loss.item()\n",
    "        nb_eval_steps += 1\n",
    "\n",
    "        # compute the predictions\n",
    "        if preds is None:\n",
    "            preds = logits.detach().cpu().numpy()\n",
    "            out_label_ids = labels.detach().cpu().numpy()\n",
    "        else:\n",
    "            preds = np.append(preds, logits.detach().cpu().numpy(), axis=0)\n",
    "            out_label_ids = np.append(out_label_ids, labels.detach().cpu().numpy(), axis=0)\n",
    "\n",
    "# compute average evaluation loss\n",
    "eval_loss = eval_loss / nb_eval_steps\n",
    "preds = np.argmax(preds, axis=2)\n",
    "\n",
    "out_label_list = [[] for _ in range(out_label_ids.shape[0])]\n",
    "preds_list = [[] for _ in range(out_label_ids.shape[0])]\n",
    "\n",
    "for i in range(out_label_ids.shape[0]):\n",
    "    for j in range(out_label_ids.shape[1]):\n",
    "        if out_label_ids[i, j] != pad_token_label_id:\n",
    "            out_label_list[i].append(label_map[out_label_ids[i][j]])\n",
    "            preds_list[i].append(label_map[preds[i][j]])\n",
    "\n",
    "results = {\n",
    "    \"loss\": eval_loss,\n",
    "    \"precision\": precision_score(out_label_list, preds_list),\n",
    "    \"recall\": recall_score(out_label_list, preds_list),\n",
    "    \"f1\": f1_score(out_label_list, preds_list),\n",
    "}\n",
    "print(results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "trainable params: 702,733 || all params: 113,237,786 || trainable%: 0.6206\n"
     ]
    }
   ],
   "source": [
    "model.print_trainable_parameters()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.save_pretrained(\"peft_layoutlm\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 40,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.8M\t./peft_layoutlm/adapter_model.safetensors\n"
     ]
    }
   ],
   "source": [
    "!du -h ./peft_layoutlm/adapter_model.safetensors"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "include_colab_link": true,
   "name": "Fine-tuning LayoutLMForTokenClassification on FUNSD.ipynb",
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "aee8b7b246df8f9039afb4144a1f6fd8d2ca17a180786b69acc140d282b71a49"
   }
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "138a6b922e454ebbaeb315ecd5f476b8": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "DescriptionStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "DescriptionStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "description_width": ""
     }
    },
    "31642aacae2a44879960da09f938ecc4": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "513e00b619924f5693259cd919a927ab": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "63e819a04f6e4829838c0e30e65516ed": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "FloatProgressModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "FloatProgressModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "ProgressView",
      "bar_style": "success",
      "description": "Downloading: 100%",
      "description_tooltip": null,
      "layout": "IPY_MODEL_31642aacae2a44879960da09f938ecc4",
      "max": 453093832,
      "min": 0,
      "orientation": "horizontal",
      "style": "IPY_MODEL_c75f0da13a1e4dbe94800711d55390a6",
      "value": 453093832
     }
    },
    "a8126ba98376402888e9ba344cf1c538": {
     "model_module": "@jupyter-widgets/base",
     "model_name": "LayoutModel",
     "state": {
      "_model_module": "@jupyter-widgets/base",
      "_model_module_version": "1.2.0",
      "_model_name": "LayoutModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "LayoutView",
      "align_content": null,
      "align_items": null,
      "align_self": null,
      "border": null,
      "bottom": null,
      "display": null,
      "flex": null,
      "flex_flow": null,
      "grid_area": null,
      "grid_auto_columns": null,
      "grid_auto_flow": null,
      "grid_auto_rows": null,
      "grid_column": null,
      "grid_gap": null,
      "grid_row": null,
      "grid_template_areas": null,
      "grid_template_columns": null,
      "grid_template_rows": null,
      "height": null,
      "justify_content": null,
      "justify_items": null,
      "left": null,
      "margin": null,
      "max_height": null,
      "max_width": null,
      "min_height": null,
      "min_width": null,
      "object_fit": null,
      "object_position": null,
      "order": null,
      "overflow": null,
      "overflow_x": null,
      "overflow_y": null,
      "padding": null,
      "right": null,
      "top": null,
      "visibility": null,
      "width": null
     }
    },
    "c75f0da13a1e4dbe94800711d55390a6": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "ProgressStyleModel",
     "state": {
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "ProgressStyleModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/base",
      "_view_module_version": "1.2.0",
      "_view_name": "StyleView",
      "bar_color": null,
      "description_width": "initial"
     }
    },
    "d1c3e1a66db04227a74ef8d6481d6daf": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HTMLModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HTMLModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HTMLView",
      "description": "",
      "description_tooltip": null,
      "layout": "IPY_MODEL_a8126ba98376402888e9ba344cf1c538",
      "placeholder": "​",
      "style": "IPY_MODEL_138a6b922e454ebbaeb315ecd5f476b8",
      "value": " 453M/453M [00:15&lt;00:00, 30.0MB/s]"
     }
    },
    "da1094982d044ab28eb0effebbfcbb78": {
     "model_module": "@jupyter-widgets/controls",
     "model_name": "HBoxModel",
     "state": {
      "_dom_classes": [],
      "_model_module": "@jupyter-widgets/controls",
      "_model_module_version": "1.5.0",
      "_model_name": "HBoxModel",
      "_view_count": null,
      "_view_module": "@jupyter-widgets/controls",
      "_view_module_version": "1.5.0",
      "_view_name": "HBoxView",
      "box_style": "",
      "children": [
       "IPY_MODEL_63e819a04f6e4829838c0e30e65516ed",
       "IPY_MODEL_d1c3e1a66db04227a74ef8d6481d6daf"
      ],
      "layout": "IPY_MODEL_513e00b619924f5693259cd919a927ab"
     }
    }
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
