NER-TF.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Named Entity Recognition (NER)\n",
    "\n",
    "This notebook is from [AI for Beginners Curriculum](http://aka.ms/ai-beginners).\n",
    "\n",
    "In this example, we will learn how to train NER model on [Annotated Corpus for Named Entity Recognition](https://www.kaggle.com/datasets/abhinavwalia95/entity-annotated-corpus) Dataset from Kaggle. Before procedding, please donwload [ner_dataset.csv](https://www.kaggle.com/datasets/abhinavwalia95/entity-annotated-corpus?resource=download&select=ner_dataset.csv) file into current directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from tensorflow import keras\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Preparing the Dataset \n",
    "\n",
    "We will start by reading the dataset into a dataframe. If you want to learn more about using Pandas, visit a [lesson on data processing](https://github.com/microsoft/Data-Science-For-Beginners/tree/main/2-Working-With-Data/07-python) in our [Data Science for Beginners](http://aka.ms/datascience-beginners)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>Sentence #</th>\n",
       "      <th>Word</th>\n",
       "      <th>POS</th>\n",
       "      <th>Tag</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Sentence: 1</td>\n",
       "      <td>Thousands</td>\n",
       "      <td>NNS</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>NaN</td>\n",
       "      <td>of</td>\n",
       "      <td>IN</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>NaN</td>\n",
       "      <td>demonstrators</td>\n",
       "      <td>NNS</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>NaN</td>\n",
       "      <td>have</td>\n",
       "      <td>VBP</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>NaN</td>\n",
       "      <td>marched</td>\n",
       "      <td>VBN</td>\n",
       "      <td>O</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "    Sentence #           Word  POS Tag\n",
       "0  Sentence: 1      Thousands  NNS   O\n",
       "1          NaN             of   IN   O\n",
       "2          NaN  demonstrators  NNS   O\n",
       "3          NaN           have  VBP   O\n",
       "4          NaN        marched  VBN   O"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = pd.read_csv('ner_dataset.csv',encoding='unicode-escape')\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's get unique tags and create lookup dictionaries that we can use to convert tags into class numbers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array(['O', 'B-geo', 'B-gpe', 'B-per', 'I-geo', 'B-org', 'I-org', 'B-tim',\n",
       "       'B-art', 'I-art', 'I-per', 'I-gpe', 'I-tim', 'B-nat', 'B-eve',\n",
       "       'I-eve', 'I-nat'], dtype=object)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tags = df.Tag.unique()\n",
    "tags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'O'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "id2tag = dict(enumerate(tags))\n",
    "tag2id = { v : k for k,v in id2tag.items() }\n",
    "\n",
    "id2tag[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we need to do the same with vocabulary. For simplicity, we will create vocabulary without taking word frequency into account; in real life you might want to use Keras vectorizer, and limit the number of words."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "vocab = set(df['Word'].apply(lambda x: x.lower()))\n",
    "id2word = { i+1 : v for i,v in enumerate(vocab) }\n",
    "id2word[0] = '<UNK>'\n",
    "vocab.add('<UNK>')\n",
    "word2id = { v : k for k,v in id2word.items() }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need to create a dataset of sentences for training. Let's loop through the original dataset and separate all individual sentences into `X` (lists of words) and `Y` (list of tokens):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [],
   "source": [
    "X,Y = [],[]\n",
    "s,t = [],[]\n",
    "for i,row in df[['Sentence #','Word','Tag']].iterrows():\n",
    "    if pd.isna(row['Sentence #']):\n",
    "        s.append(row['Word'])\n",
    "        t.append(row['Tag'])\n",
    "    else:\n",
    "        if len(s)>0:\n",
    "            X.append(s)\n",
    "            Y.append(t)\n",
    "        s,t = [row['Word']],[row['Tag']]\n",
    "X.append(s)\n",
    "Y.append(t)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now vectorize all words and tokens:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([10386,\n",
       "  23515,\n",
       "  4134,\n",
       "  29620,\n",
       "  7954,\n",
       "  13583,\n",
       "  21193,\n",
       "  12222,\n",
       "  27322,\n",
       "  18258,\n",
       "  5815,\n",
       "  15880,\n",
       "  5355,\n",
       "  25242,\n",
       "  31327,\n",
       "  18258,\n",
       "  27067,\n",
       "  23515,\n",
       "  26444,\n",
       "  14412,\n",
       "  358,\n",
       "  26551,\n",
       "  5011,\n",
       "  30558],\n",
       " [0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0])"
      ]
     },
     "execution_count": 93,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def vectorize(seq):\n",
    "    return [word2id[x.lower()] for x in seq]\n",
    "\n",
    "def tagify(seq):\n",
    "    return [tag2id[x] for x in seq]\n",
    "\n",
    "Xv = list(map(vectorize,X))\n",
    "Yv = list(map(tagify,Y))\n",
    "\n",
    "Xv[0], Yv[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For simplicity, we will pad all sentences with 0 tokens to the maximum length. In real life, we might want to use more clever strategy, and pad sequences only within one minibatch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_data = keras.preprocessing.sequence.pad_sequences(Xv,padding='post')\n",
    "Y_data = keras.preprocessing.sequence.pad_sequences(Yv,padding='post')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Defining Token Classification Network\n",
    "\n",
    "We will use two-layer bidirectional LSTM network for token classification. In order to apply dense classifier to each of the output of the last LSTM layer, we will use `TimeDistributed` construction, which replicates the same dense layer to each of the outputs of LSTM at each step: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 94,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Model: \"sequential_3\"\n",
      "_________________________________________________________________\n",
      " Layer (type)                Output Shape              Param #   \n",
      "=================================================================\n",
      " embedding_4 (Embedding)     (None, 104, 300)          9545400   \n",
      "                                                                 \n",
      " bidirectional_6 (Bidirectio  (None, 104, 200)         320800    \n",
      " nal)                                                            \n",
      "                                                                 \n",
      " bidirectional_7 (Bidirectio  (None, 104, 200)         240800    \n",
      " nal)                                                            \n",
      "                                                                 \n",
      " time_distributed_3 (TimeDis  (None, 104, 17)          3417      \n",
      " tributed)                                                       \n",
      "                                                                 \n",
      "=================================================================\n",
      "Total params: 10,110,417\n",
      "Trainable params: 10,110,417\n",
      "Non-trainable params: 0\n",
      "_________________________________________________________________\n"
     ]
    }
   ],
   "source": [
    "maxlen = X_data.shape[1]\n",
    "vocab_size = len(vocab)\n",
    "num_tags = len(tags)\n",
    "model = keras.models.Sequential([\n",
    "    keras.layers.Embedding(vocab_size, 300, input_length=maxlen),\n",
    "    keras.layers.Bidirectional(keras.layers.LSTM(units=100, activation='tanh', return_sequences=True)),\n",
    "    keras.layers.Bidirectional(keras.layers.LSTM(units=100, activation='tanh', return_sequences=True)),\n",
    "    keras.layers.TimeDistributed(keras.layers.Dense(num_tags, activation='softmax'))\n",
    "])\n",
    "model.compile(loss='sparse_categorical_crossentropy',optimizer='adam',metrics=['acc'])\n",
    "model.summary()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note here that we are explicity specifying `maxlen` for our dataset - in case we want the network to be able to handle variable length sequences, we need to be a bit more clever when defining the network.\n",
    "\n",
    "Let's now train the model. For speed, we will only train for one epoch, but you may try training for longer time. Also, you may want to separate some part of the dataset as training dataset, to observe validation accuracy."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1499/1499 [==============================] - 740s 488ms/step - loss: 0.0667 - acc: 0.9841\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<keras.callbacks.History at 0x16f0bb2a310>"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.fit(X_data,Y_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Testing the Result\n",
    "\n",
    "Let's now see how our entity recognition model works on a sample sentence: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 91,
   "metadata": {},
   "outputs": [],
   "source": [
    "sent = 'John Smith went to Paris to attend a conference in cancer development institute'\n",
    "words = sent.lower().split()\n",
    "v = keras.preprocessing.sequence.pad_sequences([[word2id[x] for x in words]],padding='post',maxlen=maxlen)\n",
    "res = model(v)[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 92,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "john -> B-per\n",
      "smith -> I-per\n",
      "went -> O\n",
      "to -> O\n",
      "paris -> B-geo\n",
      "to -> O\n",
      "attend -> O\n",
      "a -> O\n",
      "conference -> O\n",
      "in -> O\n",
      "cancer -> B-org\n",
      "development -> I-org\n",
      "institute -> I-org\n"
     ]
    }
   ],
   "source": [
    "r = np.argmax(res.numpy(),axis=1)\n",
    "for i,w in zip(r,words):\n",
    "    print(f\"{w} -> {id2tag[i]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Takeaway\n",
    "\n",
    "Even simple LSTM model shows reasonable results at NER. However, to get much better results, you may want to use large pre-trained language models such as BERT. Training BERT for NER using Huggingface Transformers library is described [here](https://huggingface.co/course/chapter7/2?fw=pt)."
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "16af2a8bbb083ea23e5e41c7f5787656b2ce26968575d8763f2c4b17f9cd711f"
  },
  "kernelspec": {
   "display_name": "Python 3.8.12 ('py38')",
   "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.8.12"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}