demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from modeling_layoutlm import LayoutLMForTokenClassification\n",
    "from transformers import (\n",
    "    BertConfig,\n",
    "    BertTokenizer,\n",
    ")\n",
    "from utils_docvqa import (\n",
    "    read_docvqa_examples,\n",
    "    convert_examples_to_features)\n",
    "from torch.utils.data import DataLoader, SequentialSampler, TensorDataset\n",
    "from transformers.data.processors.squad import SquadResult\n",
    "from tqdm import tqdm\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODEL_FOLDER = \"./models/\"\n",
    "SAMPLE_DATA = \"./models/sample_data.json\"\n",
    "LABELS = [\"start\",\"end\"]\n",
    "pad_token_label_id=-100\n",
    "labels = [\"start\",\"end\"]\n",
    "max_seq_length = 512\n",
    "max_query_length = 64\n",
    "doc_stride = 128"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\")\n",
    "torch.cuda.set_device(device)\n",
    "model_class = LayoutLMForTokenClassification\n",
    "config_class = BertConfig\n",
    "tokenizer_class = BertTokenizer\n",
    "config = config_class.from_pretrained(MODEL_FOLDER,num_labels=2,cache_dir=None)\n",
    "model = model_class.from_pretrained(MODEL_FOLDER)\n",
    "tokenizer = tokenizer_class.from_pretrained(MODEL_FOLDER,do_lower_case=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /mnt/anisha/code/docvqa/utils_docvqa.py:95: The name tf.gfile.Open is deprecated. Please use tf.io.gfile.GFile instead.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "examples = read_docvqa_examples(SAMPLE_DATA, is_training=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "features = convert_examples_to_features(\n",
    "            examples=examples,\n",
    "            label_list=labels,\n",
    "            tokenizer=tokenizer,\n",
    "            max_seq_length=max_seq_length,\n",
    "            doc_stride=doc_stride,\n",
    "            max_query_length=max_query_length,\n",
    "            is_training=False,\n",
    "            pad_token_label_id=pad_token_label_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long)\n",
    "all_input_mask = torch.tensor([f.input_mask for f in features], dtype=torch.long)\n",
    "all_segment_ids = torch.tensor([f.segment_ids for f in features], dtype=torch.long)\n",
    "all_bboxes = torch.tensor([f.boxes for f in features], dtype=torch.long)\n",
    "all_example_index = torch.arange(all_input_ids.size(0), dtype=torch.long)\n",
    "\n",
    "eval_dataset = TensorDataset(\n",
    "        all_input_ids, all_input_mask, all_segment_ids,all_bboxes,all_example_index)\n",
    "eval_batch_size = 1\n",
    "eval_sampler = (\n",
    "        SequentialSampler(eval_dataset))\n",
    "\n",
    "eval_dataloader = DataLoader(\n",
    "        eval_dataset, sampler=eval_sampler, batch_size=eval_batch_size\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Evaluating: 100%|██████████| 1/1 [00:00<00:00, 32.54it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "what is the contact person name mentioned in letter ?\n",
      "maura payne\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "model.to(device)\n",
    "all_results = []\n",
    "\n",
    "def to_list(tensor):\n",
    "    return tensor.detach().cpu().tolist()\n",
    "\n",
    "for batch in tqdm(eval_dataloader, desc=\"Evaluating\"):\n",
    "    model.eval()\n",
    "    batch = tuple(t.to(device) for t in batch)\n",
    "    with torch.no_grad():\n",
    "        inputs = {\n",
    "            \"input_ids\": batch[0],\n",
    "            \"attention_mask\": batch[1],\n",
    "        }\n",
    "        inputs[\"bbox\"] = batch[3]\n",
    "        inputs[\"token_type_ids\"] = (batch[2])\n",
    "        outputs = model(**inputs)\n",
    "    example_indices = batch[4]\n",
    "\n",
    "    for i, example_index in enumerate(example_indices):\n",
    "        eval_feature = features[example_index.item()]\n",
    "        unique_id = int(eval_feature.unique_id)\n",
    "\n",
    "        output = [to_list(output[i]) for output in outputs]\n",
    "\n",
    "        start_logits, end_logits = output\n",
    "        result = SquadResult(unique_id, start_logits, end_logits)\n",
    "        all_results.append(result)\n",
    "predictions_json = {}\n",
    "assert len(all_results)==len(features)\n",
    "for i in range(len(all_results)):\n",
    "    start_index = np.argmax(all_results[i].start_logits)\n",
    "    end_index = np.argmax(all_results[i].end_logits)\n",
    "    pred_answer = features[i].tokens[start_index:end_index+1]\n",
    "    pred_answer = ' '.join([x for x in pred_answer])\n",
    "    pred_text = pred_answer.replace(' ##', '')\n",
    "    question = features[i].tokens[1:features[i].tokens.index('[SEP]')]\n",
    "    question_text = ' '.join([x for x in question])\n",
    "    question_text = question_text.replace(' ##', '')\n",
    "    print(question_text)\n",
    "    print(pred_text)\n",
    "   \n",
    "        "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "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.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}