BERT_NER.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "BERT_NER.ipynb",
      "provenance": [],
      "collapsed_sections": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU",
    "widgets": {
      "application/vnd.jupyter.widget-state+json": {
        "9256085ae60e4f1488ff4dc0fe116be4": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_25e22404d4c64ad7ae39424400e367e5",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_826a36ccc89e4cdfae1eaee1e10178ce",
              "IPY_MODEL_d262560ac6854e61a5c3e3304c6ae5ef"
            ]
          }
        },
        "25e22404d4c64ad7ae39424400e367e5": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "826a36ccc89e4cdfae1eaee1e10178ce": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_1f438624a3354b86ab56ff773de404be",
            "_dom_classes": [],
            "description": "100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 5396,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 5396,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_fead003eb0434288b82317850504be1b"
          }
        },
        "d262560ac6854e61a5c3e3304c6ae5ef": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_85cada4d329f477093755bfe805a954c",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 5396/5396 [22:51&lt;00:00,  3.94it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_13cbfbdaee4e41fd8d856706267f5bfc"
          }
        },
        "1f438624a3354b86ab56ff773de404be": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "fead003eb0434288b82317850504be1b": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "85cada4d329f477093755bfe805a954c": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "13cbfbdaee4e41fd8d856706267f5bfc": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "95289137035748f1bea68d2fc6bc0571": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HBoxModel",
          "state": {
            "_view_name": "HBoxView",
            "_dom_classes": [],
            "_model_name": "HBoxModel",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "box_style": "",
            "layout": "IPY_MODEL_440abc96491f41998b033eb95b8ab813",
            "_model_module": "@jupyter-widgets/controls",
            "children": [
              "IPY_MODEL_4f6da0e6dc4743eb978c938e6a41129d",
              "IPY_MODEL_3bae128bd65549b39a812a930ec3e4ae"
            ]
          }
        },
        "440abc96491f41998b033eb95b8ab813": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "4f6da0e6dc4743eb978c938e6a41129d": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "FloatProgressModel",
          "state": {
            "_view_name": "ProgressView",
            "style": "IPY_MODEL_a0f79e011cae4b1bb87a8cba509b813b",
            "_dom_classes": [],
            "description": "100%",
            "_model_name": "FloatProgressModel",
            "bar_style": "success",
            "max": 600,
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": 600,
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "orientation": "horizontal",
            "min": 0,
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_e3a925e3db0d4378b6bedf2225e2282e"
          }
        },
        "3bae128bd65549b39a812a930ec3e4ae": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "HTMLModel",
          "state": {
            "_view_name": "HTMLView",
            "style": "IPY_MODEL_6aad75268a7d4606ae6a5612e7b4bb1a",
            "_dom_classes": [],
            "description": "",
            "_model_name": "HTMLModel",
            "placeholder": "​",
            "_view_module": "@jupyter-widgets/controls",
            "_model_module_version": "1.5.0",
            "value": " 600/600 [00:43&lt;00:00, 13.86it/s]",
            "_view_count": null,
            "_view_module_version": "1.5.0",
            "description_tooltip": null,
            "_model_module": "@jupyter-widgets/controls",
            "layout": "IPY_MODEL_4f4a058573e443e1b026929699f775ab"
          }
        },
        "a0f79e011cae4b1bb87a8cba509b813b": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "ProgressStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "ProgressStyleModel",
            "description_width": "initial",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "bar_color": null,
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "e3a925e3db0d4378b6bedf2225e2282e": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        },
        "6aad75268a7d4606ae6a5612e7b4bb1a": {
          "model_module": "@jupyter-widgets/controls",
          "model_name": "DescriptionStyleModel",
          "state": {
            "_view_name": "StyleView",
            "_model_name": "DescriptionStyleModel",
            "description_width": "",
            "_view_module": "@jupyter-widgets/base",
            "_model_module_version": "1.5.0",
            "_view_count": null,
            "_view_module_version": "1.2.0",
            "_model_module": "@jupyter-widgets/controls"
          }
        },
        "4f4a058573e443e1b026929699f775ab": {
          "model_module": "@jupyter-widgets/base",
          "model_name": "LayoutModel",
          "state": {
            "_view_name": "LayoutView",
            "grid_template_rows": null,
            "right": null,
            "justify_content": null,
            "_view_module": "@jupyter-widgets/base",
            "overflow": null,
            "_model_module_version": "1.2.0",
            "_view_count": null,
            "flex_flow": null,
            "width": null,
            "min_width": null,
            "border": null,
            "align_items": null,
            "bottom": null,
            "_model_module": "@jupyter-widgets/base",
            "top": null,
            "grid_column": null,
            "overflow_y": null,
            "overflow_x": null,
            "grid_auto_flow": null,
            "grid_area": null,
            "grid_template_columns": null,
            "flex": null,
            "_model_name": "LayoutModel",
            "justify_items": null,
            "grid_row": null,
            "max_height": null,
            "align_content": null,
            "visibility": null,
            "align_self": null,
            "height": null,
            "min_height": null,
            "padding": null,
            "grid_auto_rows": null,
            "grid_gap": null,
            "max_width": null,
            "order": null,
            "_view_module_version": "1.2.0",
            "grid_template_areas": null,
            "object_position": null,
            "object_fit": null,
            "grid_auto_columns": null,
            "margin": null,
            "display": null,
            "left": null
          }
        }
      }
    }
  },
  "cells": [
    {
      "cell_type": "code",
      "metadata": {
        "id": "S1MyhheU9zrk",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 355
        },
        "outputId": "970e5a48-4224-449a-d444-3a7356ab18dd"
      },
      "source": [
        "!pip install transformers"
      ],
      "execution_count": 48,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Requirement already satisfied: transformers in /usr/local/lib/python3.6/dist-packages (3.0.2)\n",
            "Requirement already satisfied: regex!=2019.12.17 in /usr/local/lib/python3.6/dist-packages (from transformers) (2019.12.20)\n",
            "Requirement already satisfied: sacremoses in /usr/local/lib/python3.6/dist-packages (from transformers) (0.0.43)\n",
            "Requirement already satisfied: tokenizers==0.8.1.rc1 in /usr/local/lib/python3.6/dist-packages (from transformers) (0.8.1rc1)\n",
            "Requirement already satisfied: numpy in /usr/local/lib/python3.6/dist-packages (from transformers) (1.18.5)\n",
            "Requirement already satisfied: packaging in /usr/local/lib/python3.6/dist-packages (from transformers) (20.4)\n",
            "Requirement already satisfied: requests in /usr/local/lib/python3.6/dist-packages (from transformers) (2.23.0)\n",
            "Requirement already satisfied: tqdm>=4.27 in /usr/local/lib/python3.6/dist-packages (from transformers) (4.41.1)\n",
            "Requirement already satisfied: filelock in /usr/local/lib/python3.6/dist-packages (from transformers) (3.0.12)\n",
            "Requirement already satisfied: dataclasses; python_version < \"3.7\" in /usr/local/lib/python3.6/dist-packages (from transformers) (0.7)\n",
            "Requirement already satisfied: sentencepiece!=0.1.92 in /usr/local/lib/python3.6/dist-packages (from transformers) (0.1.91)\n",
            "Requirement already satisfied: six in /usr/local/lib/python3.6/dist-packages (from sacremoses->transformers) (1.12.0)\n",
            "Requirement already satisfied: joblib in /usr/local/lib/python3.6/dist-packages (from sacremoses->transformers) (0.16.0)\n",
            "Requirement already satisfied: click in /usr/local/lib/python3.6/dist-packages (from sacremoses->transformers) (7.1.2)\n",
            "Requirement already satisfied: pyparsing>=2.0.2 in /usr/local/lib/python3.6/dist-packages (from packaging->transformers) (2.4.7)\n",
            "Requirement already satisfied: idna<3,>=2.5 in /usr/local/lib/python3.6/dist-packages (from requests->transformers) (2.10)\n",
            "Requirement already satisfied: urllib3!=1.25.0,!=1.25.1,<1.26,>=1.21.1 in /usr/local/lib/python3.6/dist-packages (from requests->transformers) (1.24.3)\n",
            "Requirement already satisfied: chardet<4,>=3.0.2 in /usr/local/lib/python3.6/dist-packages (from requests->transformers) (3.0.4)\n",
            "Requirement already satisfied: certifi>=2017.4.17 in /usr/local/lib/python3.6/dist-packages (from requests->transformers) (2020.6.20)\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sxkgAOu498PX",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "outputId": "b9c19ccc-4962-4993-d14f-71e8716e5d51"
      },
      "source": [
        "from google.colab import drive\n",
        "drive.mount('/content/gdrive')"
      ],
      "execution_count": 49,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Drive already mounted at /content/gdrive; to attempt to forcibly remount, call drive.mount(\"/content/gdrive\", force_remount=True).\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "X6vyP9GC9ncH",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "import os\n",
        "import torch\n",
        "import torch.nn as nn\n",
        "import transformers\n",
        "from transformers import BertForTokenClassification,AdamW,BertTokenizer\n",
        "from transformers import get_linear_schedule_with_warmup\n",
        "from torchsummary import summary\n",
        "\n",
        "\n",
        "from tqdm.auto import tqdm \n",
        "import time\n",
        "import pandas as pd\n",
        "import numpy as np\n",
        "import joblib\n",
        "\n",
        "from sklearn import metrics\n",
        "from sklearn import model_selection, preprocessing"
      ],
      "execution_count": 50,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jyTC6rXG-Enh",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "MAX_LEN=128\n",
        "TRAIN_BATCH_SIZE = 8\n",
        "VALID_BATCH_SIZE = 8\n",
        "EPOCHS = 1\n",
        "MODEL_PATH=\"/content/gdrive/My Drive/bert_ner_model.bin\"\n",
        "OUTPUT_LOG=\"/content/gdrive/My Drive/bert_ner_train.log\"\n",
        "\n",
        "#From kaggle dataset https://www.kaggle.com/abhinavwalia95/entity-annotated-corpus\n",
        "TRAINING_FILE=\"/content/gdrive/My Drive/ner_dataset.csv\""
      ],
      "execution_count": 51,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "pV0u2JO8-bs-",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "outputId": "e9fae91a-bb67-4e19-98b5-723258ce0f5f"
      },
      "source": [
        "if torch.cuda.is_available() :\n",
        "  device = torch.device(\"cuda\")\n",
        "  print('We will use the GPU:',torch.cuda.get_device_name(0))\n",
        "else:\n",
        "  print('No GPU available, using the CPU instead')\n",
        "  device = torch.device(\"cpu\")"
      ],
      "execution_count": 52,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "We will use the GPU: Tesla P100-PCIE-16GB\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "O9gNc9JP-6J1",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 35
        },
        "outputId": "1f543b19-ed88-49e7-feda-2f81a27ad67c"
      },
      "source": [
        "# Load the BERT tokenizer.\n",
        "print('Loading BERT tokenizer...')\n",
        "tokenizer = BertTokenizer.from_pretrained('bert-base-uncased', do_lower_case=True)"
      ],
      "execution_count": 53,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "Loading BERT tokenizer...\n"
          ],
          "name": "stdout"
        }
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "IaD_ZYOk_DBH",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def log(value):\n",
        "  if os.path.exists(OUTPUT_LOG) :\n",
        "    f= open(OUTPUT_LOG,\"a\")\n",
        "    f.write(value+\"\\n\")\n",
        "    f.close()\n",
        "  else:\n",
        "    f= open(OUTPUT_LOG,\"w\")\n",
        "    f.write(value+ \"\\n\")\n",
        "    f.close()\n"
      ],
      "execution_count": 54,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HPfwZU1K_IBx",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "#Read Data\n",
        "\n",
        "def read_data(input_file):\n",
        "    df = pd.read_csv(input_file, encoding=\"latin-1\")\n",
        "    df.loc[:, \"Sentence #\"] = df[\"Sentence #\"].fillna(method=\"ffill\")\n",
        "\n",
        "    enc_pos = preprocessing.LabelEncoder()\n",
        "    enc_tag = preprocessing.LabelEncoder()\n",
        "\n",
        "    df.loc[:, \"POS\"] = enc_pos.fit_transform(df[\"POS\"])\n",
        "    df.loc[:, \"Tag\"] = enc_tag.fit_transform(df[\"Tag\"])\n",
        "\n",
        "    sentences = df.groupby(\"Sentence #\")[\"Word\"].apply(list).values\n",
        "    pos = df.groupby(\"Sentence #\")[\"POS\"].apply(list).values\n",
        "    tag = df.groupby(\"Sentence #\")[\"Tag\"].apply(list).values\n",
        "\n",
        "    print(len(sentences))\n",
        "    print(len(pos))\n",
        "    print(len(tag))\n",
        "    return sentences, pos, tag, enc_pos, enc_tag"
      ],
      "execution_count": 55,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FSUBF3-IBpJC",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "class BERTDataset:\n",
        "    def __init__(self, input_text, pos, tag):\n",
        "        self.input_text= input_text\n",
        "        self.pos = pos\n",
        "        self.tag = tag\n",
        "        self.tokenizer = tokenizer\n",
        "        #self.max_len = MAX_LEN\n",
        "    \n",
        "    def __len__(self):\n",
        "        return len(self.input_text)\n",
        "    \n",
        "    def __getitem__(self, item):\n",
        "        max_len=MAX_LEN\n",
        "        text = self.input_text[item]\n",
        "        pos = self.pos[item]\n",
        "        tag = self.tag[item]\n",
        "\n",
        "        ids = []\n",
        "        target_pos = []\n",
        "        target_tag =[]\n",
        "\n",
        "        for i, s in enumerate(text):\n",
        "            inputs = tokenizer.encode(\n",
        "                s,\n",
        "                add_special_tokens=False\n",
        "            )\n",
        "          \n",
        "            input_len = len(inputs)\n",
        "            ids.extend(inputs)\n",
        "            target_pos.extend([pos[i]] * input_len)\n",
        "            target_tag.extend([tag[i]] * input_len)\n",
        "\n",
        "        ids = ids[:max_len - 2]\n",
        "        target_pos = target_pos[:max_len - 2]\n",
        "        target_tag = target_tag[:max_len - 2]\n",
        "\n",
        "        ids = [101] + ids + [102]\n",
        "        target_pos = [0] + target_pos + [0]\n",
        "        target_tag = [0] + target_tag + [0]\n",
        "\n",
        "        mask = [1] * len(ids)\n",
        "        token_type_ids = [0] * len(ids)\n",
        "\n",
        "        padding_len = max_len - len(ids)\n",
        "\n",
        "        ids = ids + ([0] * padding_len)\n",
        "        mask = mask + ([0] * padding_len)\n",
        "        token_type_ids = token_type_ids + ([0] * padding_len)\n",
        "        target_pos = target_pos + ([0] * padding_len)\n",
        "        target_tag = target_tag + ([0] * padding_len)\n",
        "\n",
        "        return {\n",
        "            \"ids\": torch.tensor(ids, dtype=torch.long),\n",
        "            \"mask\": torch.tensor(mask, dtype=torch.long),\n",
        "            \"token_type_ids\": torch.tensor(token_type_ids, dtype=torch.long),\n",
        "            \"target_pos\": torch.tensor(target_pos, dtype=torch.long),\n",
        "            \"target_tag\": torch.tensor(target_tag, dtype=torch.long),\n",
        "        }"
      ],
      "execution_count": 56,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "85epjD4jCjS_",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def train_fn(data_loader, model, optimizer, device, scheduler):\n",
        "    model.train()\n",
        "    final_loss = 0\n",
        "    for data in tqdm(data_loader, total=len(data_loader)):\n",
        "        for k, v in data.items():\n",
        "            data[k] = v.to(device)\n",
        "        optimizer.zero_grad()\n",
        "        _, _, loss = model(**data)\n",
        "        loss.backward()\n",
        "        optimizer.step()\n",
        "        scheduler.step()\n",
        "        final_loss += loss.item()\n",
        "    return final_loss / len(data_loader)\n",
        "\n",
        "\n",
        "def eval_fn(data_loader, model, device):\n",
        "    model.eval()\n",
        "    final_loss = 0\n",
        "    for data in tqdm(data_loader, total=len(data_loader)):\n",
        "        for k, v in data.items():\n",
        "            data[k] = v.to(device)\n",
        "        _, _, loss = model(**data)\n",
        "        final_loss += loss.item()\n",
        "    return final_loss / len(data_loader)\n",
        "\n",
        "def loss_fn(output, target, mask, num_labels):\n",
        "    lfn = nn.CrossEntropyLoss()\n",
        "    active_loss = mask.view(-1) == 1\n",
        "    active_logits = output[0].view(-1, num_labels)\n",
        "    active_labels = torch.where(\n",
        "        active_loss,\n",
        "        target.view(-1),\n",
        "        torch.tensor(lfn.ignore_index).type_as(target)\n",
        "    )\n",
        "    loss = lfn(active_logits, active_labels)\n",
        "    return loss"
      ],
      "execution_count": 57,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "colab_type": "code",
        "id": "ZbIo7mP-m5jf",
        "colab": {}
      },
      "source": [
        "class BERTNerModel(nn.Module):\n",
        "    def __init__(self, num_tag, num_pos):\n",
        "        super(BERTNerModel, self).__init__()\n",
        "        self.num_tag = num_tag\n",
        "        self.num_pos = num_pos\n",
        "        self.bert_tag = transformers.BertForTokenClassification.from_pretrained(\"bert-base-uncased\",num_labels=num_tag)\n",
        "        self.bert_pos = transformers.BertForTokenClassification.from_pretrained(\"bert-base-uncased\",num_labels=num_pos)\n",
        "        \n",
        "\n",
        "    def forward(\n",
        "        self, \n",
        "        ids, \n",
        "        mask, \n",
        "        token_type_ids, \n",
        "        target_pos, \n",
        "        target_tag\n",
        "    ):\n",
        "        tag= self.bert_tag(\n",
        "            ids, \n",
        "            attention_mask=mask, \n",
        "            token_type_ids=token_type_ids\n",
        "        )\n",
        "\n",
        "        pos= self.bert_pos(\n",
        "            ids, \n",
        "            attention_mask=mask, \n",
        "            token_type_ids=token_type_ids\n",
        "        )\n",
        "\n",
        "        loss_tag = loss_fn(tag, target_tag, mask, self.num_tag)\n",
        "        loss_pos = loss_fn(pos, target_pos, mask, self.num_pos)\n",
        "\n",
        "        loss = (loss_tag + loss_pos) / 2\n",
        "\n",
        "        return tag, pos, loss"
      ],
      "execution_count": 58,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "qavtanFARmJc",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def train():\n",
        "  sentences, pos, tag, enc_pos, enc_tag = read_data(TRAINING_FILE)\n",
        "\n",
        "  meta_data = {\n",
        "    \"enc_pos\": enc_pos,\n",
        "    \"enc_tag\": enc_tag\n",
        "  }\n",
        "\n",
        "  joblib.dump(meta_data, \"meta.bin\")\n",
        "\n",
        "  num_pos = len(list(enc_pos.classes_))\n",
        "  num_tag = len(list(enc_tag.classes_))\n",
        "\n",
        "  (\n",
        "    train_sentences,\n",
        "    test_sentences,\n",
        "    train_pos,\n",
        "    test_pos,\n",
        "    train_tag,\n",
        "    test_tag\n",
        "  ) = model_selection.train_test_split(\n",
        "    sentences, \n",
        "    pos, \n",
        "    tag, \n",
        "    random_state=42, \n",
        "    test_size=0.1\n",
        "  )\n",
        "\n",
        "\n",
        "  train_dataset = BERTDataset(\n",
        "    input_text=train_sentences, pos=train_pos, tag=train_tag\n",
        "    )\n",
        "\n",
        "  train_data_loader = torch.utils.data.DataLoader(\n",
        "    train_dataset, batch_size=TRAIN_BATCH_SIZE, num_workers=4\n",
        "  )\n",
        "\n",
        "  valid_dataset = BERTDataset(\n",
        "    input_text=test_sentences, pos=test_pos, tag=test_tag\n",
        "  )\n",
        "\n",
        "  valid_data_loader = torch.utils.data.DataLoader(\n",
        "    valid_dataset, batch_size=VALID_BATCH_SIZE, num_workers=1\n",
        "  )\n",
        "\n",
        "  device = torch.device(\"cuda\")\n",
        "  model = BERTNerModel(num_tag=num_tag, num_pos=num_pos)\n",
        "  model.to(device)\n",
        "\n",
        "  param_optimizer = list(model.named_parameters())\n",
        "  no_decay = [\"bias\", \"LayerNorm.bias\", \"LayerNorm.weight\"]\n",
        "  optimizer_parameters = [\n",
        "    {\n",
        "        \"params\": [\n",
        "            p for n, p in param_optimizer if not any(\n",
        "                nd in n for nd in no_decay\n",
        "            )\n",
        "        ],\n",
        "        \"weight_decay\": 0.001,\n",
        "    },\n",
        "        {\n",
        "        \"params\": [\n",
        "            p for n, p in param_optimizer if any(\n",
        "                nd in n for nd in no_decay\n",
        "            )\n",
        "        ],\n",
        "        \"weight_decay\": 0.0,\n",
        "    },\n",
        "  ]\n",
        "\n",
        "  num_train_steps = int(\n",
        "    len(train_sentences) / TRAIN_BATCH_SIZE * EPOCHS\n",
        "  )\n",
        "  optimizer = AdamW(optimizer_parameters, lr=3e-5)\n",
        "  scheduler = get_linear_schedule_with_warmup(\n",
        "    optimizer, \n",
        "    num_warmup_steps=0, \n",
        "    num_training_steps=num_train_steps\n",
        "  )\n",
        "\n",
        "  best_loss = np.inf\n",
        "  for epoch in range(EPOCHS):\n",
        "    train_loss = train_fn(\n",
        "        train_data_loader, \n",
        "        model, \n",
        "        optimizer, \n",
        "        device, \n",
        "        scheduler\n",
        "    )\n",
        "    test_loss = eval_fn(\n",
        "        valid_data_loader,\n",
        "        model,\n",
        "        device\n",
        "    )\n",
        "    print(f\"Train Loss = {train_loss} Valid Loss = {test_loss}\")\n",
        "    if test_loss < best_loss:\n",
        "        torch.save(model.state_dict(), MODEL_PATH)\n",
        "        best_loss = test_loss"
      ],
      "execution_count": 59,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VcXCeEd_ULsu",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        ""
      ],
      "execution_count": 59,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Nv6-tC3XmLwq",
        "colab_type": "code",
        "colab": {}
      },
      "source": [
        "def predict():\n",
        "  meta_data = joblib.load(\"meta.bin\")\n",
        "  enc_pos = meta_data[\"enc_pos\"]\n",
        "  enc_tag = meta_data[\"enc_tag\"]\n",
        "\n",
        "  num_pos = len(list(enc_pos.classes_))\n",
        "  num_tag = len(list(enc_tag.classes_))\n",
        "\n",
        "  sentence = \"Ganga is a river in North India that flows from Himalayas to Bay of Bengal\"\n",
        "  tokenized_sentence = tokenizer.encode(sentence)\n",
        "\n",
        "  sentence = sentence.split()\n",
        "  print(sentence)\n",
        "\n",
        "  print(tokenized_sentence)\n",
        "  print(tokenizer.convert_ids_to_tokens(tokenized_sentence))\n",
        "\n",
        "  test_dataset = BERTDataset(\n",
        "    input_text=[sentence], \n",
        "    pos=[[0] * len(sentence)], \n",
        "    tag=[[0] * len(sentence)]\n",
        "  )\n",
        "\n",
        "  device = torch.device(\"cuda\")\n",
        "  model = BERTNerModel(num_tag=num_tag, num_pos=num_pos)\n",
        "  model.load_state_dict(torch.load(MODEL_PATH))\n",
        "  model.to(device)\n",
        "\n",
        "  with torch.no_grad():\n",
        "    data = test_dataset[0]\n",
        "    for k, v in data.items():\n",
        "      data[k] = v.to(device).unsqueeze(0)\n",
        "    tag, pos, _ = model(**data)\n",
        "\n",
        "    print(\n",
        "        enc_tag.inverse_transform(\n",
        "            tag[0].argmax(2).cpu().numpy().reshape(-1)\n",
        "        )[:len(tokenized_sentence)]\n",
        "    )\n",
        "    print(\n",
        "        enc_pos.inverse_transform(\n",
        "            pos[0].argmax(2).cpu().numpy().reshape(-1)\n",
        "        )[:len(tokenized_sentence)]\n",
        "    )\n"
      ],
      "execution_count": 60,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4-zt4HHynV0w",
        "colab_type": "code",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 687,
          "referenced_widgets": [
            "9256085ae60e4f1488ff4dc0fe116be4",
            "25e22404d4c64ad7ae39424400e367e5",
            "826a36ccc89e4cdfae1eaee1e10178ce",
            "d262560ac6854e61a5c3e3304c6ae5ef",
            "1f438624a3354b86ab56ff773de404be",
            "fead003eb0434288b82317850504be1b",
            "85cada4d329f477093755bfe805a954c",
            "13cbfbdaee4e41fd8d856706267f5bfc",
            "95289137035748f1bea68d2fc6bc0571",
            "440abc96491f41998b033eb95b8ab813",
            "4f6da0e6dc4743eb978c938e6a41129d",
            "3bae128bd65549b39a812a930ec3e4ae",
            "a0f79e011cae4b1bb87a8cba509b813b",
            "e3a925e3db0d4378b6bedf2225e2282e",
            "6aad75268a7d4606ae6a5612e7b4bb1a",
            "4f4a058573e443e1b026929699f775ab"
          ]
        },
        "outputId": "1d5ad4d4-4e85-4feb-8297-bc176c6708ee"
      },
      "source": [
        "train()\n",
        "predict()"
      ],
      "execution_count": 61,
      "outputs": [
        {
          "output_type": "stream",
          "text": [
            "47959\n",
            "47959\n",
            "47959\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertForTokenClassification: ['cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.decoder.weight', 'cls.seq_relationship.weight', 'cls.seq_relationship.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias']\n",
            "- This IS expected if you are initializing BertForTokenClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPretraining model).\n",
            "- This IS NOT expected if you are initializing BertForTokenClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
            "Some weights of BertForTokenClassification were not initialized from the model checkpoint at bert-base-uncased and are newly initialized: ['classifier.weight', 'classifier.bias']\n",
            "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
            "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertForTokenClassification: ['cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.decoder.weight', 'cls.seq_relationship.weight', 'cls.seq_relationship.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias']\n",
            "- This IS expected if you are initializing BertForTokenClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPretraining model).\n",
            "- This IS NOT expected if you are initializing BertForTokenClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
            "Some weights of BertForTokenClassification were not initialized from the model checkpoint at bert-base-uncased and are newly initialized: ['classifier.weight', 'classifier.bias']\n",
            "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "9256085ae60e4f1488ff4dc0fe116be4",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, max=5396.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "display_data",
          "data": {
            "application/vnd.jupyter.widget-view+json": {
              "model_id": "95289137035748f1bea68d2fc6bc0571",
              "version_minor": 0,
              "version_major": 2
            },
            "text/plain": [
              "HBox(children=(FloatProgress(value=0.0, max=600.0), HTML(value='')))"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "stream",
          "text": [
            "\n",
            "Train Loss = 0.13304480957540457 Valid Loss = 0.08998058945406229\n",
            "['Ganga', 'is', 'a', 'river', 'in', 'North', 'India', 'that', 'flows', 'from', 'Himalayas', 'to', 'Bay', 'of', 'Bengal']\n",
            "[101, 28646, 2003, 1037, 2314, 1999, 2167, 2634, 2008, 6223, 2013, 26779, 2000, 3016, 1997, 8191, 102]\n",
            "['[CLS]', 'ganga', 'is', 'a', 'river', 'in', 'north', 'india', 'that', 'flows', 'from', 'himalayas', 'to', 'bay', 'of', 'bengal', '[SEP]']\n"
          ],
          "name": "stdout"
        },
        {
          "output_type": "stream",
          "text": [
            "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertForTokenClassification: ['cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.decoder.weight', 'cls.seq_relationship.weight', 'cls.seq_relationship.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias']\n",
            "- This IS expected if you are initializing BertForTokenClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPretraining model).\n",
            "- This IS NOT expected if you are initializing BertForTokenClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
            "Some weights of BertForTokenClassification were not initialized from the model checkpoint at bert-base-uncased and are newly initialized: ['classifier.weight', 'classifier.bias']\n",
            "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n",
            "Some weights of the model checkpoint at bert-base-uncased were not used when initializing BertForTokenClassification: ['cls.predictions.bias', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.decoder.weight', 'cls.seq_relationship.weight', 'cls.seq_relationship.bias', 'cls.predictions.transform.LayerNorm.weight', 'cls.predictions.transform.LayerNorm.bias']\n",
            "- This IS expected if you are initializing BertForTokenClassification from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPretraining model).\n",
            "- This IS NOT expected if you are initializing BertForTokenClassification from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
            "Some weights of BertForTokenClassification were not initialized from the model checkpoint at bert-base-uncased and are newly initialized: ['classifier.weight', 'classifier.bias']\n",
            "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n"
          ],
          "name": "stderr"
        },
        {
          "output_type": "stream",
          "text": [
            "['B-art' 'B-geo' 'O' 'O' 'O' 'O' 'B-geo' 'I-geo' 'O' 'O' 'O' 'B-geo' 'O'\n",
            " 'B-geo' 'I-geo' 'I-geo' 'B-art']\n",
            "['$' 'NNP' 'VBZ' 'DT' 'NN' 'IN' 'NNP' 'NNP' 'WDT' 'VBZ' 'IN' 'NNP' 'TO'\n",
            " 'NNP' 'IN' 'NNP' '$']\n"
          ],
          "name": "stdout"
        }
      ]
    }
  ]
}