inference/donut_inference.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/envs/pytorch/lib/python3.9/site-packages/tqdm/auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n",
      "/opt/conda/envs/pytorch/lib/python3.9/site-packages/torch/functional.py:568: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at  /opt/conda/conda-bld/pytorch_1646755888534/work/aten/src/ATen/native/TensorShape.cpp:2228.)\n",
      "  return _VF.meshgrid(tensors, **kwargs)  # type: ignore[attr-defined]\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from transformers import DonutProcessor,VisionEncoderDecoderModel\n",
    "\n",
    "\n",
    "model = VisionEncoderDecoderModel.from_pretrained(\"naver-clova-ix/donut-base\")\n",
    "\n",
    "device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n",
    "model.to(device)\n",
    "\n",
    "# Load processor\n",
    "processor = DonutProcessor.from_pretrained(\"naver-clova-ix/donut-base\")\n",
    "\n",
    "# set decoder start token id\n",
    "model.config.decoder_start_token_id = processor.tokenizer.convert_tokens_to_ids(['<s>'])[0]\n",
    "model.config.pad_token_id = processor.tokenizer.pad_token_id\n",
    "\n",
    "max_length = 512\n",
    "image_size = [960, 720]\n",
    "\n",
    "# update image_size of the encoder\n",
    "\n",
    "# model.config.encoder.image_size = processor.feature_extractor.size[::-1] # (height, width)\n",
    "model.config.encoder.image_size = image_size # (height, width)\n",
    "model.config.decoder.max_length = max_length\n",
    "\n",
    "# resize processor and model to match\n",
    "processor.feature_extractor.size = image_size[::-1] # should be (width, height)\n",
    "processor.feature_extractor.do_align_long_axis = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Random sample"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No config specified, defaulting to: cats-image/image\n",
      "Reusing dataset cats-image (/home/ubuntu/.cache/huggingface/datasets/huggingface___cats-image/image/1.9.0/68fbc793fb10cd165e490867f5d61fa366086ea40c73e549a020103dcb4f597e)\n",
      "100%|██████████| 1/1 [00:00<00:00, 513.00it/s]\n"
     ]
    }
   ],
   "source": [
    "from datasets import load_dataset\n",
    "dataset = load_dataset(\"huggingface/cats-image\")\n",
    "image = dataset[\"test\"][\"image\"][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "pixel_values = processor.feature_extractor(image, return_tensors=\"pt\").pixel_values.to(device)\n",
    "\n",
    "\n",
    "labels = processor.tokenizer(\n",
    "    \"an image of two cats chilling on a couch\",\n",
    "    return_tensors=\"pt\",\n",
    ").input_ids.to(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# the forward function automatically creates the correct decoder_input_ids\n",
    "loss = model(pixel_values=pixel_values, labels=labels).loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(7.5516, device='cuda:0', grad_fn=<NllLossBackward0>)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SROIE example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Cloning into 'ICDAR-2019-SROIE'...\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "huggingface/tokenizers: The current process just got forked, after parallelism has already been used. Disabling parallelism to avoid deadlocks...\n",
      "To disable this warning, you can either:\n",
      "\t- Avoid using `tokenizers` before the fork if possible\n",
      "\t- Explicitly set the environment variable TOKENIZERS_PARALLELISM=(true | false)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Updating files: 100% (1980/1980), done.\n"
     ]
    }
   ],
   "source": [
    "%%bash \n",
    "# clone repository\n",
    "git clone https://github.com/zzzDavid/ICDAR-2019-SROIE.git\n",
    "# copy data\n",
    "cp -r ICDAR-2019-SROIE/data ./\n",
    "# clean up\n",
    "rm -rf ICDAR-2019-SROIE\n",
    "rm -rf data/box"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "from pathlib import Path\n",
    "import shutil\n",
    "\n",
    "# define paths\n",
    "base_path = Path(\"data\")\n",
    "metadata_path = base_path.joinpath(\"key\")\n",
    "image_path = base_path.joinpath(\"img\")\n",
    "# define metadata list\n",
    "metadata_list = []\n",
    "\n",
    "# parse metadata\n",
    "for file_name in metadata_path.glob(\"*.json\"):\n",
    "  with open(file_name, \"r\") as json_file:\n",
    "    # load json file\n",
    "    data = json.load(json_file)\n",
    "    # create \"text\" column with json string\n",
    "    text = json.dumps(data)\n",
    "    # add to metadata list if image exists\n",
    "    if image_path.joinpath(f\"{file_name.stem}.jpg\").is_file():    \n",
    "      metadata_list.append({\"text\":text,\"file_name\":f\"{file_name.stem}.jpg\"})\n",
    "      # delete json file\n",
    "      \n",
    "# write jsonline file\n",
    "with open(image_path.joinpath('metadata.jsonl'), 'w') as outfile:\n",
    "    for entry in metadata_list:\n",
    "        json.dump(entry, outfile)\n",
    "        outfile.write('\\n')\n",
    "\n",
    "# remove old meta data\n",
    "shutil.rmtree(metadata_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Resolving data files: 100%|██████████| 627/627 [00:00<00:00, 371970.10it/s]\n",
      "Using custom data configuration default-b9550808e775c2fe\n",
      "Reusing dataset imagefolder (/home/ubuntu/.cache/huggingface/datasets/imagefolder/default-b9550808e775c2fe/0.0.0/0fc50c79b681877cc46b23245a6ef5333d036f48db40d53765a68034bc48faff)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Dataset has 626 images\n",
      "Dataset features are: dict_keys(['image', 'text'])\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "from datasets import load_dataset\n",
    "\n",
    "# define paths\n",
    "base_path = Path(\"data\")\n",
    "metadata_path = base_path.joinpath(\"key\")\n",
    "image_path = base_path.joinpath(\"img\")\n",
    "# Load dataset\n",
    "dataset = load_dataset(\"imagefolder\", data_dir=image_path, split=\"train\")\n",
    "\n",
    "print(f\"Dataset has {len(dataset)} images\")\n",
    "print(f\"Dataset features are: {dataset.features.keys()}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = dataset[\"image\"][0]\n",
    "text = dataset[\"text\"][0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Embedding(57533, 1024)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "processor.tokenizer.add_special_tokens({\"additional_special_tokens\": ['<s_total>', '</s_total>', '<s_date>', '</s_date>', '<s_company>', '</s_company>', '<s_address>', '</s_address>', '<s>', '</s>']})\n",
    "model.decoder.resize_token_embeddings(len(processor.tokenizer))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "image = dataset[\"image\"][45]\n",
    "text = \"<s><s_total>$6.90</s_total><s_date>27 MAR 2018</s_date><s_company>UNIHAKKA INTERNATIONAL SDN BHD</s_company><s_address>12, JALAN TAMPOI 7/4,KAWASAN PARINDUSTRIAN TAMPOI,81200 JOHOR BAHRU,JOHOR</s_address></s>\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "pixel_values = processor.feature_extractor(image, return_tensors=\"pt\").pixel_values.to(device)\n",
    "\n",
    "\n",
    "input_ids = processor.tokenizer(\n",
    "    text,\n",
    "    return_tensors=\"pt\",\n",
    "    add_special_tokens=False,\n",
    "    max_length=512,\n",
    "    padding=\"max_length\",\n",
    ").input_ids.to(device)\n",
    "\n",
    "labels = input_ids.clone()\n",
    "labels[labels == processor.tokenizer.pad_token_id] = -100  # model doesn't need to predict pad token\n",
    "labels = labels.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Pixel values: torch.Size([1, 3, 960, 720])\n",
      "Labels: torch.Size([1, 512])\n"
     ]
    }
   ],
   "source": [
    "print(f\"Pixel values: {pixel_values.shape}\")\n",
    "print(f\"Labels: {labels.shape}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[    0, 57525, 38282, 55144, 39539, 43112, 57526, 57527, 10558, 35521,\n",
       "         34453, 57528, 57529, 48503,  3782, 46318, 56023, 28428, 13663, 43529,\n",
       "         37127, 14706, 57530, 57531, 40474, 35815, 57103, 35794, 38946, 52532,\n",
       "         44977, 38100, 16908, 35815, 46318, 42252, 37877, 35181, 49023, 48540,\n",
       "         36757, 53233, 38946, 52532, 44977, 35815,  8739, 46051, 24271, 18081,\n",
       "          1338, 37127, 25340, 17712, 35815, 34452, 18081,  1338, 57532,     2,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,\n",
       "          -100,  -100]])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'<s><s_total> $6.90</s_total><s_date> 27 MAR 2018</s_date><s_company> UNIHAKKA INTERNATIONAL SDN BHD</s_company><s_address> 12, JALAN TAMPOI 7/4,KAWASAN PARINDUSTRIAN TAMPOI,81200 JOHOR BAHRU,JOHOR</s_address></s>'"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "processor.tokenizer.decode(labels[0][labels[0]!=-100])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(11.7751, grad_fn=<NllLossBackward0>)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# the forward function automatically creates the correct decoder_input_ids\n",
    "loss = model(pixel_values=pixel_values, labels=labels).loss\n",
    "\n",
    "loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.13 ('pytorch')",
   "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.9.13"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "2d58e898dde0263bc564c6968b04150abacfd33eed9b19aaa8e45c040360e146"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}