Layoutlm_v3_on_custom_token_classification.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# LayoutLMv3 for financial report NER\n",
    "\n",
    "We now cover the latest model in the LayoutLM family. \n",
    "\n",
    "An essential difference to other models is that bounding box coordinates do not have to be passed per word not on word level but on segment level. Using this grouping procedure (because segments are coarser than words), one expects that for entities consisting of multiple tokens, predictions will be pushed towards giving equal labels to words from equal segments. As our labels `fund_name` or `umbrella` consist of many tokens, it is interesting to explore whether this leads to further improvement.\n",
    "\n",
    "Where do we get the segment information from? One possibility is to use a textline detector and use the results for segments. \n",
    "\n",
    "FRFPE was labeled so that we used a layout detector fine-tuned on fund documents. The segment results are available as `ImageAnnotation` in ground truth. With that, relations to the segments and words were created using the `MatchingService`. \n",
    "\n",
    "During training (as well as in the evaluation or pipelines) it is possible to use the segments that one wants to use as replacement for the `Word` bounding boxes. \n",
    "\n",
    "We will now use these procedures to fine-tune LayoutLMv3 correctly. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import deepdoctection as dd\n",
    "from collections import defaultdict\n",
    "import wandb\n",
    "from transformers import RobertaTokenizerFast"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m[0608 19:15.35 @file_utils.py:33]\u001b[0m  \u001b[32mINF\u001b[0m  \u001b[97mPyTorch version 1.9.0+cu111 available.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "@dd.object_types_registry.register(\"ner_first_page\")\n",
    "class FundsFirstPage(dd.ObjectTypes):\n",
    "\n",
    "    REPORT_DATE = \"report_date\"\n",
    "    UMBRELLA = \"umbrella\"\n",
    "    REPORT_TYPE = \"report_type\"\n",
    "    FUND_NAME = \"fund_name\"\n",
    "\n",
    "dd.update_all_types_dict()\n",
    "\n",
    "@dd.curry\n",
    "def overwrite_location_and_load(dp, image_dir, load_image):\n",
    "    image_file = image_dir / dp.file_name.replace(\"pdf\",\"png\")\n",
    "    dp.location = image_file.as_posix()\n",
    "    if load_image:\n",
    "        dp.image = dd.load_image_from_file(image_file)\n",
    "    return dp\n",
    "\n",
    "class NerBuilder(dd.DataFlowBaseBuilder):\n",
    "\n",
    "    def build(self, **kwargs) -> dd.DataFlow:\n",
    "        load_image = kwargs.get(\"load_image\", False)\n",
    "        filter_languages = kwargs.get(\"filter_languages\")\n",
    "\n",
    "        ann_files_dir = self.get_workdir()\n",
    "        image_dir = self.get_workdir() / \"image\"\n",
    "\n",
    "        df = dd.SerializerFiles.load(ann_files_dir,\".json\")   # get a stream of .json files\n",
    "        df = dd.MapData(df, dd.Image.from_file)   # load .json file\n",
    "\n",
    "        df = dd.MapData(df, overwrite_location_and_load(image_dir, load_image))\n",
    "\n",
    "        if self.categories.is_filtered():\n",
    "            df = dd.MapData(\n",
    "                df,\n",
    "                dd.filter_cat(\n",
    "                    self.categories.get_categories(as_dict=False, filtered=True),\n",
    "                    self.categories.get_categories(as_dict=False, filtered=False),\n",
    "                ),\n",
    "            )\n",
    "        df = dd.MapData(df,dd.re_assign_cat_ids(cat_to_sub_cat_mapping=self.categories.get_sub_categories(\n",
    "                                                 categories=dd.LayoutType.WORD,\n",
    "                                                 sub_categories={dd.LayoutType.WORD: dd.WordType.TOKEN_CLASS},\n",
    "                                                 keys = False,\n",
    "                                                 values_as_dict=True,\n",
    "                                                 name_as_key=True)))\n",
    "        \n",
    "        if filter_languages:\n",
    "            df = dd.MapData(df, dd.filter_summary({\"language\": [dd.get_type(lang) for lang in filter_languages]},\n",
    "                                                 mode=\"value\"))\n",
    "\n",
    "        return df\n",
    "    \n",
    "ner = dd.CustomDataset(name = \"FRFPE\",\n",
    "                 dataset_type=dd.DatasetType.TOKEN_CLASSIFICATION,\n",
    "                 location=\"FRFPE\",\n",
    "                 init_categories=[dd.LayoutType.TEXT, dd.LayoutType.TITLE, dd.LayoutType.LIST, dd.LayoutType.TABLE,\n",
    "                                  dd.LayoutType.FIGURE, dd.LayoutType.LINE, dd.LayoutType.WORD],\n",
    "                 init_sub_categories={dd.LayoutType.WORD: {dd.WordType.TOKEN_CLASS: [FundsFirstPage.REPORT_DATE,\n",
    "                                                                                     FundsFirstPage.REPORT_TYPE,\n",
    "                                                                                     FundsFirstPage.UMBRELLA,\n",
    "                                                                                     FundsFirstPage.FUND_NAME,\n",
    "                                                                                     dd.TokenClasses.OTHER],\n",
    "                                                           dd.WordType.TAG: []}},\n",
    "                 dataflow_builder=NerBuilder)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "|                                                                                                                                                                                              |357/?[00:00<00:00,53473.56it/s]\n",
      "\u001b[32m[0608 19:15.37 @base.py:250]\u001b[0m  \u001b[32mINF\u001b[0m  \u001b[97mWill used dataflow from previously explicitly passed configuration\u001b[0m\n",
      "|                                                                                                                                                                                                 |357/?[00:29<00:00,12.28it/s]\n"
     ]
    }
   ],
   "source": [
    "df = ner.dataflow.build(load_image=True)\n",
    "\n",
    "merge = dd.MergeDataset(ner)\n",
    "merge.explicit_dataflows(df)\n",
    "merge.buffer_datasets()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.init(project=\"FRFPE_layoutlmv1\", resume=True)\n",
    "artifact = wandb.use_artifact('jm76/FRFPE_layoutlmv1/merge_FRFPE:v0', type='dataset')\n",
    "table = artifact.get(\"split\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m[0608 19:16.09 @base.py:250]\u001b[0m  \u001b[32mINF\u001b[0m  \u001b[97mWill used dataflow from previously explicitly passed configuration\u001b[0m\n",
      "|                                                                                                                                                                                                 |357/?[00:28<00:00,12.69it/s]\n"
     ]
    }
   ],
   "source": [
    "split_dict = defaultdict(list)\n",
    "for row in table.data:\n",
    "    split_dict[row[0]].append(row[1])\n",
    "\n",
    "merge.create_split_by_id(split_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.finish()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So not forget to download the model if it is not in you .cache yet."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "path_config_json = dd.ModelCatalog.get_full_path_configs(\"microsoft/layoutlmv3-base/pytorch_model.bin\")\n",
    "path_weights = dd.ModelCatalog.get_full_path_weights(\"microsoft/layoutlmv3-base/pytorch_model.bin\")\n",
    "\n",
    "metric = dd.get_metric(\"f1\")\n",
    "metric.set_categories(sub_category_names={\"word\": [\"token_class\"]})\n",
    "\n",
    "dd.train_hf_layoutlm(path_config_json,\n",
    "                     merge,\n",
    "                     path_weights,\n",
    "                     config_overwrite=[\"max_steps=2000\",\n",
    "                                       \"per_device_train_batch_size=8\",\n",
    "                                       \"eval_steps=100\",\n",
    "                                       \"save_steps=400\",\n",
    "                                       \"use_wandb=True\",\n",
    "                                       \"wandb_project=FRFPE_layoutlmv3\"],\n",
    "                     log_dir=\"/path/to/dir/Experiments/FRFPE/layoutlmv3\",\n",
    "                     dataset_val=merge,\n",
    "                     metric=metric,\n",
    "                     use_token_tag=False,\n",
    "                     pipeline_component_name=\"LMTokenClassifierService\",\n",
    "                     segment_positions=[dd.LayoutType.TITLE, \n",
    "                                        dd.LayoutType.TEXT, \n",
    "                                        dd.LayoutType.TABLE, \n",
    "                                        dd.LayoutType.LIST])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "wandb.finish()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation\n",
    "\n",
    "Evaluation on the test split drops significantly. This is quite surprising as we haven't seen a F1-score drop of this size before. \n",
    "Especially `fund_name` and `other` have a significant drop. As there are much more `fund_name` labels in at least one sample\n",
    "it looks like the model gets confused due to the segment bounding boxes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m[0608 19:56.22 @eval.py:113]\u001b[0m  \u001b[32mINF\u001b[0m  \u001b[97mBuilding multi threading pipeline component to increase prediction throughput. Using 2 threads\u001b[0m\n",
      "\u001b[32m[0608 19:56.23 @eval.py:225]\u001b[0m  \u001b[32mINF\u001b[0m  \u001b[97mPredicting objects...\u001b[0m\n",
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 26/26 [00:01<00:00, 20.03it/s]\n",
      "\u001b[32m[0608 19:56.24 @eval.py:207]\u001b[0m  \u001b[32mINF\u001b[0m  \u001b[97mStarting evaluation...\u001b[0m\n",
      "\u001b[32m[0608 19:56.24 @accmetric.py:373]\u001b[0m  \u001b[32mINF\u001b[0m  \u001b[97mF1 results:\n",
      " \u001b[36m|     key     | category_id   | val      | num_samples   |\n",
      "|:-----------:|:--------------|:---------|:--------------|\n",
      "|    word     | 1             | 1        | 1505          |\n",
      "| token_class | 1             | 0.962162 | 89            |\n",
      "| token_class | 2             | 0.931298 | 69            |\n",
      "| token_class | 3             | 0.728571 | 86            |\n",
      "| token_class | 4             | 0.565341 | 490           |\n",
      "| token_class | 5             | 0.822703 | 771           |\u001b[0m\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "categories = ner.dataflow.categories.get_sub_categories(categories=\"word\",\n",
    "                                                        sub_categories={\"word\": [\"token_class\"]},\n",
    "                                                        keys=False)[\"word\"][\"token_class\"]\n",
    "\n",
    "path_config_json = \"/path/to/dir/Experiments/FRFPE/layoutlmv3/checkpoint-2000/config.json\"\n",
    "path_weights = \"/path/to/dir/Experiments/FRFPE/layoutlmv3/checkpoint-1600/model.safetensors\"\n",
    "\n",
    "layoutlm_classifier = dd.HFLayoutLmv3TokenClassifier(path_config_json,\n",
    "                                                     path_weights,\n",
    "                                                     categories=categories)\n",
    "\n",
    "tokenizer_fast = RobertaTokenizerFast.from_pretrained(\"roberta-base\", add_prefix_space=True)\n",
    "\n",
    "pipe_component = dd.LMTokenClassifierService(tokenizer_fast,\n",
    "                                             layoutlm_classifier,\n",
    "                                             use_other_as_default_category=True)\n",
    "\n",
    "evaluator = dd.Evaluator(merge, pipe_component, metric)\n",
    "_ = evaluator.run(split=\"test\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Many `fund_name` token have been mis-classified as `other`. And this happens particularly with segments that are rather large. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m[0608 20:06.23 @eval.py:113]\u001b[0m  \u001b[32mINF\u001b[0m  \u001b[97mBuilding multi threading pipeline component to increase prediction throughput. Using 2 threads\u001b[0m\n",
      "\u001b[32m[0608 20:06.24 @eval.py:225]\u001b[0m  \u001b[32mINF\u001b[0m  \u001b[97mPredicting objects...\u001b[0m\n",
      "100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 26/26 [00:01<00:00, 20.13it/s]\n",
      "\u001b[32m[0608 20:06.26 @eval.py:207]\u001b[0m  \u001b[32mINF\u001b[0m  \u001b[97mStarting evaluation...\u001b[0m\n",
      "\u001b[32m[0608 20:06.26 @accmetric.py:431]\u001b[0m  \u001b[32mINF\u001b[0m  \u001b[97mConfusion matrix: \n",
      " \u001b[36m|    predictions ->  |   1 |   2 |   3 |   4 |   5 |\n",
      "|     ground truth | |     |     |     |     |     |\n",
      "|                  v |     |     |     |     |     |\n",
      "|-------------------:|----:|----:|----:|----:|----:|\n",
      "|                  1 |  89 |   0 |   0 |   0 |   0 |\n",
      "|                  2 |   0 |  61 |   0 |   0 |   8 |\n",
      "|                  3 |   0 |   0 |  51 |  13 |  22 |\n",
      "|                  4 |   0 |   0 |   3 | 199 | 288 |\n",
      "|                  5 |   7 |   1 |   0 |   2 | 761 |\u001b[0m\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "metric = dd.get_metric(\"confusion\")\n",
    "metric.set_categories(sub_category_names={\"word\": [\"token_class\"]})\n",
    "\n",
    "evaluator = dd.Evaluator(merge, pipe_component, metric)\n",
    "_ = evaluator.run(split=\"test\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "evaluator.compare(interactive=True, split=\"test\", show_words=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Conclusion\n",
    "\n",
    "The results show that LayoutLMv3 is not the best choice for this dataset and it is being outperformed by LayoutXLM.\n",
    "\n",
    "It is likely to get better results with text line segment bounding boxes. This assumption is backed by the fact that the model has difficulties to deliver consistent results especially when the segment bounding box is too large. To confirm this assumption, however, one would have to adjust the dataset."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "dd_pt_033",
   "language": "python",
   "name": "dd_pt_033"
  },
  "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.9.19"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}