Pos-tagging with Bert Fine-tuning.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[BERT](https://arxiv.org/abs/1810.04805) is known to be good at Sequence tagging tasks like Named Entity Recognition. Let's see if it's true for POS-tagging."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "__author__ = \"kyubyong\"\n",
    "__address__ = \"https://github.com/kyubyong/nlp_made_easy\"\n",
    "__email__ = \"kbpark.linguist@gmail.com\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from tqdm import tqdm_notebook as tqdm\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.utils import data\n",
    "import torch.optim as optim\n",
    "from pytorch_pretrained_bert import BertTokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'1.0.0'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torch.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data preparation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thanks to the great NLTK, we don't have to worry about datasets. Some of Penn Tree Banks are included in it. I believe they serves for the purpose."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3914"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import nltk\n",
    "tagged_sents = nltk.corpus.treebank.tagged_sents()\n",
    "len(tagged_sents)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('Pierre', 'NNP'),\n",
       " ('Vinken', 'NNP'),\n",
       " (',', ','),\n",
       " ('61', 'CD'),\n",
       " ('years', 'NNS'),\n",
       " ('old', 'JJ'),\n",
       " (',', ','),\n",
       " ('will', 'MD'),\n",
       " ('join', 'VB'),\n",
       " ('the', 'DT'),\n",
       " ('board', 'NN'),\n",
       " ('as', 'IN'),\n",
       " ('a', 'DT'),\n",
       " ('nonexecutive', 'JJ'),\n",
       " ('director', 'NN'),\n",
       " ('Nov.', 'NNP'),\n",
       " ('29', 'CD'),\n",
       " ('.', '.')]"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "tagged_sents[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "tags = list(set(word_pos[1] for sent in tagged_sents for word_pos in sent))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\"JJS,VBN,SYM,NNP,POS,'',-NONE-,WP,CD,.,VBZ,RBS,RB,-RRB-,NNPS,FW,WDT,DT,WRB,PRP$,:,MD,JJ,$,EX,RBR,VBD,VBP,NN,PRP,-LRB-,LS,NNS,RP,#,TO,,,``,IN,VBG,CC,JJR,PDT,UH,VB,WP$\""
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\",\".join(tags)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# By convention, the 0'th slot is reserved for padding.\n",
    "tags = [\"<pad>\"] + tags"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "tag2idx = {tag:idx for idx, tag in enumerate(tags)}\n",
    "idx2tag = {idx:tag for idx, tag in enumerate(tags)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(3522, 392)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Let's split the data into train and test (or eval)\n",
    "from sklearn.model_selection import train_test_split\n",
    "train_data, test_data = train_test_split(tagged_sents, test_size=.1)\n",
    "len(train_data), len(test_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cuda' if torch.cuda.is_available() else 'cpu'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Data loader\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "tokenizer = BertTokenizer.from_pretrained('bert-base-cased', do_lower_case=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PosDataset(data.Dataset):\n",
    "    def __init__(self, tagged_sents):\n",
    "        sents, tags_li = [], [] # list of lists\n",
    "        for sent in tagged_sents:\n",
    "            words = [word_pos[0] for word_pos in sent]\n",
    "            tags = [word_pos[1] for word_pos in sent]\n",
    "            sents.append([\"[CLS]\"] + words + [\"[SEP]\"])\n",
    "            tags_li.append([\"<pad>\"] + tags + [\"<pad>\"])\n",
    "        self.sents, self.tags_li = sents, tags_li\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.sents)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        words, tags = self.sents[idx], self.tags_li[idx] # words, tags: string list\n",
    "\n",
    "        # We give credits only to the first piece.\n",
    "        x, y = [], [] # list of ids\n",
    "        is_heads = [] # list. 1: the token is the first piece of a word\n",
    "        for w, t in zip(words, tags):\n",
    "            tokens = tokenizer.tokenize(w) if w not in (\"[CLS]\", \"[SEP]\") else [w]\n",
    "            xx = tokenizer.convert_tokens_to_ids(tokens)\n",
    "\n",
    "            is_head = [1] + [0]*(len(tokens) - 1)\n",
    "\n",
    "            t = [t] + [\"<pad>\"] * (len(tokens) - 1)  # <PAD>: no decision\n",
    "            yy = [tag2idx[each] for each in t]  # (T,)\n",
    "\n",
    "            x.extend(xx)\n",
    "            is_heads.extend(is_head)\n",
    "            y.extend(yy)\n",
    "\n",
    "        assert len(x)==len(y)==len(is_heads), \"len(x)={}, len(y)={}, len(is_heads)={}\".format(len(x), len(y), len(is_heads))\n",
    "\n",
    "        # seqlen\n",
    "        seqlen = len(y)\n",
    "\n",
    "        # to string\n",
    "        words = \" \".join(words)\n",
    "        tags = \" \".join(tags)\n",
    "        return words, x, is_heads, tags, y, seqlen\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pad(batch):\n",
    "    '''Pads to the longest sample'''\n",
    "    f = lambda x: [sample[x] for sample in batch]\n",
    "    words = f(0)\n",
    "    is_heads = f(2)\n",
    "    tags = f(3)\n",
    "    seqlens = f(-1)\n",
    "    maxlen = np.array(seqlens).max()\n",
    "\n",
    "    f = lambda x, seqlen: [sample[x] + [0] * (seqlen - len(sample[x])) for sample in batch] # 0: <pad>\n",
    "    x = f(1, maxlen)\n",
    "    y = f(-2, maxlen)\n",
    "\n",
    "\n",
    "    f = torch.LongTensor\n",
    "\n",
    "    return words, f(x), is_heads, tags, f(y), seqlens"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pytorch_pretrained_bert import BertModel"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self, vocab_size=None):\n",
    "        super().__init__()\n",
    "        self.bert = BertModel.from_pretrained('bert-base-cased')\n",
    "\n",
    "        self.fc = nn.Linear(768, vocab_size)\n",
    "        self.device = device\n",
    "\n",
    "    def forward(self, x, y):\n",
    "        '''\n",
    "        x: (N, T). int64\n",
    "        y: (N, T). int64\n",
    "        '''\n",
    "        x = x.to(device)\n",
    "        y = y.to(device)\n",
    "        \n",
    "        if self.training:\n",
    "            self.bert.train()\n",
    "            encoded_layers, _ = self.bert(x)\n",
    "            enc = encoded_layers[-1]\n",
    "        else:\n",
    "            self.bert.eval()\n",
    "            with torch.no_grad():\n",
    "                encoded_layers, _ = self.bert(x)\n",
    "                enc = encoded_layers[-1]\n",
    "        \n",
    "        logits = self.fc(enc)\n",
    "        y_hat = logits.argmax(-1)\n",
    "        return logits, y, y_hat"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train an evaluate"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, iterator, optimizer, criterion):\n",
    "    model.train()\n",
    "    for i, batch in enumerate(iterator):\n",
    "        words, x, is_heads, tags, y, seqlens = batch\n",
    "        _y = y # for monitoring\n",
    "        optimizer.zero_grad()\n",
    "        logits, y, _ = model(x, y) # logits: (N, T, VOCAB), y: (N, T)\n",
    "\n",
    "        logits = logits.view(-1, logits.shape[-1]) # (N*T, VOCAB)\n",
    "        y = y.view(-1)  # (N*T,)\n",
    "\n",
    "        loss = criterion(logits, y)\n",
    "        loss.backward()\n",
    "\n",
    "        optimizer.step()\n",
    "\n",
    "        if i%10==0: # monitoring\n",
    "            print(\"step: {}, loss: {}\".format(i, loss.item()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval(model, iterator):\n",
    "    model.eval()\n",
    "\n",
    "    Words, Is_heads, Tags, Y, Y_hat = [], [], [], [], []\n",
    "    with torch.no_grad():\n",
    "        for i, batch in enumerate(iterator):\n",
    "            words, x, is_heads, tags, y, seqlens = batch\n",
    "\n",
    "            _, _, y_hat = model(x, y)  # y_hat: (N, T)\n",
    "\n",
    "            Words.extend(words)\n",
    "            Is_heads.extend(is_heads)\n",
    "            Tags.extend(tags)\n",
    "            Y.extend(y.numpy().tolist())\n",
    "            Y_hat.extend(y_hat.cpu().numpy().tolist())\n",
    "\n",
    "    ## gets results and save\n",
    "    with open(\"result\", 'w') as fout:\n",
    "        for words, is_heads, tags, y_hat in zip(Words, Is_heads, Tags, Y_hat):\n",
    "            y_hat = [hat for head, hat in zip(is_heads, y_hat) if head == 1]\n",
    "            preds = [idx2tag[hat] for hat in y_hat]\n",
    "            assert len(preds)==len(words.split())==len(tags.split())\n",
    "            for w, t, p in zip(words.split()[1:-1], tags.split()[1:-1], preds[1:-1]):\n",
    "                fout.write(\"{} {} {}\\n\".format(w, t, p))\n",
    "            fout.write(\"\\n\")\n",
    "            \n",
    "    ## calc metric\n",
    "    y_true =  np.array([tag2idx[line.split()[1]] for line in open('result', 'r').read().splitlines() if len(line) > 0])\n",
    "    y_pred =  np.array([tag2idx[line.split()[2]] for line in open('result', 'r').read().splitlines() if len(line) > 0])\n",
    "\n",
    "    acc = (y_true==y_pred).astype(np.int32).sum() / len(y_true)\n",
    "\n",
    "    print(\"acc=%.2f\"%acc)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load model and train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "model = Net(vocab_size=len(tag2idx))\n",
    "model.to(device)\n",
    "model = nn.DataParallel(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = PosDataset(train_data)\n",
    "eval_dataset = PosDataset(test_data)\n",
    "\n",
    "train_iter = data.DataLoader(dataset=train_dataset,\n",
    "                             batch_size=8,\n",
    "                             shuffle=True,\n",
    "                             num_workers=1,\n",
    "                             collate_fn=pad)\n",
    "test_iter = data.DataLoader(dataset=eval_dataset,\n",
    "                             batch_size=8,\n",
    "                             shuffle=False,\n",
    "                             num_workers=1,\n",
    "                             collate_fn=pad)\n",
    "\n",
    "optimizer = optim.Adam(model.parameters(), lr = 0.0001)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss(ignore_index=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "step: 0, loss: 0.027864959090948105\n",
      "step: 10, loss: 0.03902581334114075\n",
      "step: 20, loss: 0.029155433177947998\n",
      "step: 30, loss: 0.036159448325634\n",
      "step: 40, loss: 0.04948236793279648\n",
      "step: 50, loss: 0.034221794456243515\n",
      "step: 60, loss: 0.017331236973404884\n",
      "step: 70, loss: 0.06194368004798889\n",
      "step: 80, loss: 0.01584777608513832\n",
      "step: 90, loss: 0.05200301483273506\n",
      "step: 100, loss: 0.042910996824502945\n",
      "step: 110, loss: 0.01104726456105709\n",
      "step: 120, loss: 0.09724321961402893\n",
      "step: 130, loss: 0.03911526873707771\n",
      "step: 140, loss: 0.01710551604628563\n",
      "step: 150, loss: 0.06321573257446289\n",
      "step: 160, loss: 0.03924640640616417\n",
      "step: 170, loss: 0.018429234623908997\n",
      "step: 180, loss: 0.08669907599687576\n",
      "step: 190, loss: 0.03778192400932312\n",
      "step: 200, loss: 0.06231529265642166\n",
      "step: 210, loss: 0.03337831050157547\n",
      "step: 220, loss: 0.02998737245798111\n",
      "step: 230, loss: 0.042920369654893875\n",
      "step: 240, loss: 0.03866969794034958\n",
      "step: 250, loss: 0.0313352607190609\n",
      "step: 260, loss: 0.07101577520370483\n",
      "step: 270, loss: 0.10232127457857132\n",
      "step: 280, loss: 0.02433393895626068\n",
      "step: 290, loss: 0.04439501464366913\n",
      "step: 300, loss: 0.040027499198913574\n",
      "step: 310, loss: 0.027676744386553764\n",
      "step: 320, loss: 0.07696736603975296\n",
      "step: 330, loss: 0.0451495386660099\n",
      "step: 340, loss: 0.047669027000665665\n",
      "step: 350, loss: 0.03548085317015648\n",
      "step: 360, loss: 0.05764956399798393\n",
      "step: 370, loss: 0.018015533685684204\n",
      "step: 380, loss: 0.01151026040315628\n",
      "step: 390, loss: 0.05745118111371994\n",
      "step: 400, loss: 0.08683514595031738\n",
      "step: 410, loss: 0.020299362018704414\n",
      "step: 420, loss: 0.03219173103570938\n",
      "step: 430, loss: 0.0664878711104393\n",
      "step: 440, loss: 0.059365227818489075\n",
      "acc=0.98\n"
     ]
    }
   ],
   "source": [
    "train(model, train_iter, optimizer, criterion)\n",
    "eval(model, test_iter)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['Bonds NNS NNS',\n",
       " 'due JJ JJ',\n",
       " 'in IN IN',\n",
       " '2005 CD CD',\n",
       " 'have VBP VBP',\n",
       " 'a DT DT',\n",
       " '7 CD CD',\n",
       " '1\\\\/2 CD CD',\n",
       " '% NN NN',\n",
       " 'coupon NN NN',\n",
       " 'and CC CC',\n",
       " 'are VBP VBP',\n",
       " 'priced VBN VBN',\n",
       " '*-1 -NONE- -NONE-',\n",
       " 'at IN IN',\n",
       " 'par NN NN',\n",
       " '. . .',\n",
       " '',\n",
       " 'Mr. NNP NNP',\n",
       " 'Sidak NNP NNP',\n",
       " 'served VBD VBD',\n",
       " 'as IN IN',\n",
       " 'an DT DT',\n",
       " 'attorney NN NN',\n",
       " 'in IN IN',\n",
       " 'the DT DT',\n",
       " 'Reagan NNP NNP',\n",
       " 'administration NN NN',\n",
       " '. . .',\n",
       " '',\n",
       " 'Municipal NNP NNP',\n",
       " 'Issues NNPS NNPS',\n",
       " '',\n",
       " 'Viacom NNP NNP',\n",
       " 'denies VBZ VBZ',\n",
       " '0 -NONE- -NONE-',\n",
       " 'it PRP PRP',\n",
       " \"'s VBZ VBZ\",\n",
       " 'using VBG VBG',\n",
       " 'pressure NN NN',\n",
       " 'tactics NNS NNS',\n",
       " '. . .',\n",
       " '',\n",
       " 'Tokyo NNP NNP',\n",
       " \"'s POS POS\",\n",
       " 'leading VBG VBG',\n",
       " 'program NN NN',\n",
       " 'traders NNS NNS',\n",
       " 'are VBP VBP',\n",
       " 'the DT DT',\n",
       " 'big JJ JJ',\n",
       " 'U.S. NNP NNP',\n",
       " 'securities NNS NNS',\n",
       " 'houses NNS NNS',\n",
       " ', , ,',\n",
       " 'though IN IN',\n",
       " 'the DT DT',\n",
       " 'Japanese NNP NNS',\n",
       " 'are VBP VBP',\n",
       " 'playing VBG VBG',\n",
       " 'catch-up NN JJ',\n",
       " '. . .',\n",
       " '',\n",
       " 'That DT DT',\n",
       " \"'s VBZ VBZ\",\n",
       " 'why WRB WRB',\n",
       " 'Columbia NNP NNP',\n",
       " 'just RB RB',\n",
       " 'wrote VBD VBD',\n",
       " 'off RP RP',\n",
       " '$ $ $',\n",
       " '130 CD CD',\n",
       " 'million CD CD',\n",
       " '*U* -NONE- -NONE-',\n",
       " 'of IN IN',\n",
       " 'its PRP$ PRP$',\n",
       " 'junk NN NN',\n",
       " 'and CC CC',\n",
       " 'reserved VBD VBD',\n",
       " '$ $ $',\n",
       " '227 CD CD',\n",
       " 'million CD CD',\n",
       " '*U* -NONE- -NONE-',\n",
       " 'for IN IN',\n",
       " 'future JJ JJ',\n",
       " 'junk NN NN',\n",
       " 'losses NNS NNS',\n",
       " '*T*-1 -NONE- -NONE-',\n",
       " '. . .',\n",
       " '',\n",
       " 'Allergan NNP NNP',\n",
       " 'Inc. NNP NNP',\n",
       " 'said VBD VBD',\n",
       " '0 -NONE- -NONE-',\n",
       " 'it PRP PRP',\n",
       " 'received VBD VBD',\n",
       " 'Food NNP NNP',\n",
       " 'and CC CC',\n",
       " 'Drug NNP NNP',\n",
       " 'Administration NNP NNP']"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "open('result', 'r').read().splitlines()[:100]"
   ]
  },
  {
   "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.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}