transformer_mtl2.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ae55af32",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "device: mps\n",
      "random seed: 1234\n"
     ]
    }
   ],
   "source": [
    "import platform\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import random\n",
    "\n",
    "import torch \n",
    "from torch import optim \n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data import DataLoader\n",
    "from torch import nn\n",
    "\n",
    "from tqdm.notebook import tqdm\n",
    "from transformers import AutoTokenizer\n",
    "\n",
    "# enable tqdm in pandas\n",
    "tqdm.pandas()\n",
    "\n",
    "# select device\n",
    "if torch.cuda.is_available():\n",
    "    device = torch.device('cuda')\n",
    "elif 'arm64' in platform.platform():\n",
    "    device = torch.device('mps') # 'mps'\n",
    "else:\n",
    "    device = torch.device('cpu')\n",
    "print(f'device: {device.type}') \n",
    "\n",
    "# random seed\n",
    "seed = 1234\n",
    "\n",
    "# pytorch ignores this label in the loss\n",
    "ignore_index = -100\n",
    "\n",
    "# set random seed\n",
    "if seed is not None:\n",
    "    print(f'random seed: {seed}')\n",
    "    random.seed(seed)\n",
    "    np.random.seed(seed)\n",
    "    torch.manual_seed(seed)\n",
    "\n",
    "# which transformer to use\n",
    "transformer_name = \"bert-base-cased\" # 'xlm-roberta-base' # 'distilbert-base-cased'\n",
    "tokenizer = AutoTokenizer.from_pretrained(transformer_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4beaebbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# map labels to the first token in each word\n",
    "def align_labels(word_ids, labels, label_to_index):\n",
    "    label_ids = []\n",
    "    previous_word_id = None\n",
    "    for word_id in word_ids:\n",
    "        if word_id is None or word_id == previous_word_id:\n",
    "            # ignore if not a word or word id has already been seen\n",
    "            label_ids.append(ignore_index)\n",
    "        else:\n",
    "            # get label id for corresponding word\n",
    "            label_id = label_to_index[labels[word_id]]\n",
    "            label_ids.append(label_id)\n",
    "        # remember this word id\n",
    "        previous_word_id = word_id\n",
    "    \n",
    "    return label_ids\n",
    "            \n",
    "# build a set of labels in the dataset            \n",
    "def read_label_set(fn):\n",
    "    labels = set()\n",
    "    with open(fn) as f:\n",
    "        for index, line in enumerate(f):\n",
    "            line = line.strip()\n",
    "            tokens = line.split()\n",
    "            if tokens != []:\n",
    "                label = tokens[-1]\n",
    "                labels.add(label)\n",
    "    return labels\n",
    "\n",
    "# converts a two-column file in the basic MTL format (\"word \\t label\") into a dataframe\n",
    "def read_dataframe(fn, label_to_index, task_id):\n",
    "    # now build the actual dataframe for this dataset\n",
    "    data = {'words': [], 'str_labels': [], 'input_ids': [], 'word_ids': [], 'labels': [], 'task_ids': []}\n",
    "    with open(fn) as f:\n",
    "        sent_words = []\n",
    "        sent_labels = [] \n",
    "        for index, line in tqdm(enumerate(f)):\n",
    "            line = line.strip()\n",
    "            tokens = line.split()\n",
    "            if tokens == []:\n",
    "                data['words'].append(sent_words)\n",
    "                data['str_labels'].append(sent_labels)\n",
    "                \n",
    "                # tokenize each sentence\n",
    "                token_input = tokenizer(sent_words, is_split_into_words = True)  \n",
    "                token_ids = token_input['input_ids']\n",
    "                word_ids = token_input.word_ids(batch_index = 0)\n",
    "                \n",
    "                # map labels to the first token in each word\n",
    "                token_labels = align_labels(word_ids, sent_labels, label_to_index)\n",
    "                \n",
    "                data['input_ids'].append(token_ids)\n",
    "                data['word_ids'].append(word_ids)\n",
    "                data['labels'].append(token_labels)\n",
    "                data['task_ids'].append(task_id)\n",
    "                sent_words = []\n",
    "                sent_labels = [] \n",
    "            else:\n",
    "                sent_words.append(tokens[0])\n",
    "                sent_labels.append(tokens[1])\n",
    "    return pd.DataFrame(data)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "b5b3e017",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Task():\n",
    "    def __init__(self, task_id, task_name, train_file_name, dev_file_name, test_file_name):\n",
    "        self.task_id = task_id\n",
    "        self.task_name = task_name\n",
    "        # we need an index of labels first\n",
    "        self.labels = read_label_set(train_file_name)\n",
    "        self.index_to_label = {i:t for i,t in enumerate(self.labels)}\n",
    "        self.label_to_index = {t:i for i,t in enumerate(self.labels)}\n",
    "        self.num_labels = len(self.index_to_label)\n",
    "        # create data frames for the datasets\n",
    "        self.train_df = read_dataframe(train_file_name, self.label_to_index, self.task_id)\n",
    "        self.dev_df = read_dataframe(dev_file_name, self.label_to_index, self.task_id)\n",
    "        self.test_df = read_dataframe(test_file_name, self.label_to_index, self.task_id)\n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2d43b086",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "fc27ad30babe4ec4bd18daca8d642148",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "f224703f99204ee3adb91dc77f4d8543",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "d468e7b083474e0ca610d6b6d8234834",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "9a15031458364843a182c2de719489a2",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "64f88d438d434ce6882154390603f193",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "ac3a983cf30149c9982c6d8fa5a3dbfd",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "ner_task = Task(0, \"NER\", \"data/conll-ner/train.txt\", \"data/conll-ner/dev.txt\", \"data/conll-ner/test.txt\")\n",
    "pos_task = Task(1, \"POS\", \"data/pos/train.txt\", \"data/pos/dev.txt\", \"data/pos/test.txt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "90478984",
   "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>words</th>\n",
       "      <th>str_labels</th>\n",
       "      <th>input_ids</th>\n",
       "      <th>word_ids</th>\n",
       "      <th>labels</th>\n",
       "      <th>task_ids</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>[EU, rejects, German, call, to, boycott, Briti...</td>\n",
       "      <td>[B-ORG, O, B-MISC, O, O, O, B-MISC, O, O]</td>\n",
       "      <td>[101, 7270, 22961, 1528, 1840, 1106, 21423, 14...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 5, 6, 7, 7, 8, None]</td>\n",
       "      <td>[-100, 2, 5, 6, 5, 5, 5, 6, 5, -100, 5, -100]</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>[Peter, Blackburn]</td>\n",
       "      <td>[B-PER, I-PER]</td>\n",
       "      <td>[101, 1943, 14428, 102]</td>\n",
       "      <td>[None, 0, 1, None]</td>\n",
       "      <td>[-100, 4, 8, -100]</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>[BRUSSELS, 1996-08-22]</td>\n",
       "      <td>[B-LOC, O]</td>\n",
       "      <td>[101, 26660, 13329, 12649, 15928, 1820, 118, 4...</td>\n",
       "      <td>[None, 0, 0, 0, 0, 1, 1, 1, 1, 1, None]</td>\n",
       "      <td>[-100, 3, -100, -100, -100, 5, -100, -100, -10...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>[The, European, Commission, said, on, Thursday...</td>\n",
       "      <td>[O, B-ORG, I-ORG, O, O, O, O, O, O, B-MISC, O,...</td>\n",
       "      <td>[101, 1109, 1735, 2827, 1163, 1113, 9170, 1122...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...</td>\n",
       "      <td>[-100, 5, 2, 1, 5, 5, 5, 5, 5, 5, 6, 5, 5, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>[Germany, 's, representative, to, the, Europea...</td>\n",
       "      <td>[B-LOC, O, O, O, O, B-ORG, I-ORG, O, O, O, B-P...</td>\n",
       "      <td>[101, 1860, 112, 188, 4702, 1106, 1103, 1735, ...</td>\n",
       "      <td>[None, 0, 1, 1, 2, 3, 4, 5, 6, 7, 7, 8, 9, 10,...</td>\n",
       "      <td>[-100, 3, 5, -100, 5, 5, 5, 2, 1, 5, -100, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>5</th>\n",
       "      <td>[\", We, do, n't, support, any, such, recommend...</td>\n",
       "      <td>[O, O, O, O, O, O, O, O, O, O, O, O, O, O, O, ...</td>\n",
       "      <td>[101, 107, 1284, 1202, 183, 112, 189, 1619, 12...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 3, 3, 4, 5, 6, 7, 8, 9, 10,...</td>\n",
       "      <td>[-100, 5, 5, 5, 5, -100, -100, 5, 5, 5, 5, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>6</th>\n",
       "      <td>[He, said, further, scientific, study, was, re...</td>\n",
       "      <td>[O, O, O, O, O, O, O, O, O, O, O, O, O, O, O, ...</td>\n",
       "      <td>[101, 1124, 1163, 1748, 3812, 2025, 1108, 2320...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...</td>\n",
       "      <td>[-100, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>7</th>\n",
       "      <td>[He, said, a, proposal, last, month, by, EU, F...</td>\n",
       "      <td>[O, O, O, O, O, O, O, B-ORG, O, O, B-PER, I-PE...</td>\n",
       "      <td>[101, 1124, 1163, 170, 5835, 1314, 2370, 1118,...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...</td>\n",
       "      <td>[-100, 5, 5, 5, 5, 5, 5, 5, 2, 5, 5, 4, 8, -10...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>8</th>\n",
       "      <td>[Fischler, proposed, EU-wide, measures, after,...</td>\n",
       "      <td>[B-PER, O, B-MISC, O, O, O, O, B-LOC, O, B-LOC...</td>\n",
       "      <td>[101, 17355, 9022, 2879, 3000, 7270, 118, 2043...</td>\n",
       "      <td>[None, 0, 0, 0, 1, 2, 2, 2, 3, 4, 5, 6, 7, 8, ...</td>\n",
       "      <td>[-100, 4, -100, -100, 5, 6, -100, -100, 5, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>9</th>\n",
       "      <td>[But, Fischler, agreed, to, review, his, propo...</td>\n",
       "      <td>[O, B-PER, O, O, O, O, O, O, O, B-ORG, O, O, O...</td>\n",
       "      <td>[101, 1252, 17355, 9022, 2879, 2675, 1106, 318...</td>\n",
       "      <td>[None, 0, 1, 1, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10,...</td>\n",
       "      <td>[-100, 5, 4, -100, -100, 5, 5, 5, 5, 5, 5, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>10</th>\n",
       "      <td>[Spanish, Farm, Minister, Loyola, de, Palacio,...</td>\n",
       "      <td>[B-MISC, O, O, B-PER, I-PER, I-PER, O, O, O, B...</td>\n",
       "      <td>[101, 2124, 6865, 2110, 23828, 1260, 19585, 17...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 5, 5, 5, 6, 7, 8, 9, 9, ...</td>\n",
       "      <td>[-100, 6, 5, 5, 4, 8, 8, -100, -100, 5, 5, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>11</th>\n",
       "      <td>[.]</td>\n",
       "      <td>[O]</td>\n",
       "      <td>[101, 119, 102]</td>\n",
       "      <td>[None, 0, None]</td>\n",
       "      <td>[-100, 5, -100]</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>12</th>\n",
       "      <td>[Only, France, and, Britain, backed, Fischler,...</td>\n",
       "      <td>[O, B-LOC, O, B-LOC, O, B-PER, O, O, O]</td>\n",
       "      <td>[101, 2809, 1699, 1105, 2855, 5534, 17355, 902...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 5, 5, 5, 6, 6, 7, 8, None]</td>\n",
       "      <td>[-100, 5, 3, 5, 3, 5, 4, -100, -100, 5, -100, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>13</th>\n",
       "      <td>[The, EU, 's, scientific, veterinary, and, mul...</td>\n",
       "      <td>[O, B-ORG, O, O, O, O, O, O, O, O, O, O, O, O,...</td>\n",
       "      <td>[101, 1109, 7270, 112, 188, 3812, 27431, 1105,...</td>\n",
       "      <td>[None, 0, 1, 2, 2, 3, 4, 5, 6, 6, 7, 8, 9, 10,...</td>\n",
       "      <td>[-100, 5, 2, 5, -100, 5, 5, 5, 5, -100, 5, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>14</th>\n",
       "      <td>[Sheep, have, long, been, known, to, contract,...</td>\n",
       "      <td>[O, O, O, O, O, O, O, O, O, O, O, O, O, O, B-M...</td>\n",
       "      <td>[101, 1153, 8043, 1138, 1263, 1151, 1227, 1106...</td>\n",
       "      <td>[None, 0, 0, 1, 2, 3, 4, 5, 6, 7, 7, 8, 9, 10,...</td>\n",
       "      <td>[-100, 5, -100, 5, 5, 5, 5, 5, 5, 5, -100, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>15</th>\n",
       "      <td>[British, farmers, denied, on, Thursday, there...</td>\n",
       "      <td>[B-MISC, O, O, O, O, O, O, O, O, O, O, O, O, O...</td>\n",
       "      <td>[101, 1418, 6915, 5762, 1113, 9170, 1175, 1108...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...</td>\n",
       "      <td>[-100, 6, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>16</th>\n",
       "      <td>[\", What, we, have, to, be, extremely, careful...</td>\n",
       "      <td>[O, O, O, O, O, O, O, O, O, O, O, O, O, O, O, ...</td>\n",
       "      <td>[101, 107, 1327, 1195, 1138, 1106, 1129, 4450,...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...</td>\n",
       "      <td>[-100, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>17</th>\n",
       "      <td>[Bonn, has, led, efforts, to, protect, public,...</td>\n",
       "      <td>[B-LOC, O, O, O, O, O, O, O, O, O, O, O, O, O,...</td>\n",
       "      <td>[101, 21066, 1144, 1521, 3268, 1106, 3244, 147...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...</td>\n",
       "      <td>[-100, 3, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>18</th>\n",
       "      <td>[Germany, imported, 47,600, sheep, from, Brita...</td>\n",
       "      <td>[B-LOC, O, O, O, O, B-LOC, O, O, O, O, O, O, O...</td>\n",
       "      <td>[101, 1860, 11095, 3862, 117, 4372, 8892, 1121...</td>\n",
       "      <td>[None, 0, 1, 2, 2, 2, 3, 4, 5, 6, 7, 8, 9, 10,...</td>\n",
       "      <td>[-100, 3, 5, 5, -100, -100, 5, 5, 3, 5, 5, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>19</th>\n",
       "      <td>[It, brought, in, 4,275, tonnes, of, British, ...</td>\n",
       "      <td>[O, O, O, O, O, O, B-MISC, O, O, O, O, O, O, O...</td>\n",
       "      <td>[101, 1135, 1814, 1107, 125, 117, 18783, 10992...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 3, 3, 4, 5, 6, 7, 7, 7, 8, ...</td>\n",
       "      <td>[-100, 5, 5, 5, 5, -100, -100, 5, 5, 6, 5, -10...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>20</th>\n",
       "      <td>[-DOCSTART-]</td>\n",
       "      <td>[O]</td>\n",
       "      <td>[101, 118, 141, 9244, 9272, 12426, 1942, 118, ...</td>\n",
       "      <td>[None, 0, 0, 0, 0, 0, 0, 0, None]</td>\n",
       "      <td>[-100, 5, -100, -100, -100, -100, -100, -100, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>21</th>\n",
       "      <td>[Rare, Hendrix, song, draft, sells, for, almos...</td>\n",
       "      <td>[O, B-PER, O, O, O, O, O, O, O, O]</td>\n",
       "      <td>[101, 25574, 22609, 1461, 5039, 16695, 1111, 1...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 8, 8, 9, None]</td>\n",
       "      <td>[-100, 5, 4, 5, 5, 5, 5, 5, 5, 5, -100, -100, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>22</th>\n",
       "      <td>[LONDON, 1996-08-22]</td>\n",
       "      <td>[B-LOC, O]</td>\n",
       "      <td>[101, 149, 11414, 2137, 11414, 1820, 118, 4775...</td>\n",
       "      <td>[None, 0, 0, 0, 0, 1, 1, 1, 1, 1, None]</td>\n",
       "      <td>[-100, 3, -100, -100, -100, 5, -100, -100, -10...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>23</th>\n",
       "      <td>[A, rare, early, handwritten, draft, of, a, so...</td>\n",
       "      <td>[O, O, O, O, O, O, O, O, O, B-LOC, O, O, B-PER...</td>\n",
       "      <td>[101, 138, 4054, 1346, 1289, 16674, 5039, 1104...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 3, 4, 5, 6, 7, 8, 9, 9, 9, ...</td>\n",
       "      <td>[-100, 5, 5, 5, 5, -100, 5, 5, 5, 5, 5, 3, -10...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>24</th>\n",
       "      <td>[A, Florida, restaurant, paid, 10,925, pounds,...</td>\n",
       "      <td>[O, B-LOC, O, O, O, O, O, O, O, O, O, O, O, O,...</td>\n",
       "      <td>[101, 138, 2631, 4382, 3004, 1275, 117, 5556, ...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 4, 4, 4, 5, 6, 7, 8, 8, ...</td>\n",
       "      <td>[-100, 5, 3, 5, 5, 5, -100, -100, -100, 5, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>25</th>\n",
       "      <td>[At, the, end, of, a, January, 1967, concert, ...</td>\n",
       "      <td>[O, O, O, O, O, O, O, O, O, O, B-MISC, O, O, B...</td>\n",
       "      <td>[101, 1335, 1103, 1322, 1104, 170, 1356, 2573,...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...</td>\n",
       "      <td>[-100, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 6, 5, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>26</th>\n",
       "      <td>[Buyers, also, snapped, up, 16, other, items, ...</td>\n",
       "      <td>[O, O, O, O, O, O, O, O, O, O, O, O, O, O, B-P...</td>\n",
       "      <td>[101, 26123, 1468, 1145, 5267, 1146, 1479, 116...</td>\n",
       "      <td>[None, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11...</td>\n",
       "      <td>[-100, 5, -100, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>27</th>\n",
       "      <td>[They, included, a, black, lacquer, and, mothe...</td>\n",
       "      <td>[O, O, O, O, O, O, O, O, O, O, O, O, O, B-PER,...</td>\n",
       "      <td>[101, 1220, 1529, 170, 1602, 2495, 1665, 19061...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 4, 4, 5, 6, 7, 8, 9, 9, ...</td>\n",
       "      <td>[-100, 5, 5, 5, 5, 5, -100, -100, 5, 5, 5, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>28</th>\n",
       "      <td>[The, guitarist, died, of, a, drugs, overdose,...</td>\n",
       "      <td>[O, O, O, O, O, O, O, O, O, O, O, O]</td>\n",
       "      <td>[101, 1109, 5506, 1452, 1104, 170, 5557, 1166,...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 5, 6, 6, 6, 7, 8, 9, 10,...</td>\n",
       "      <td>[-100, 5, 5, 5, 5, 5, 5, 5, -100, -100, 5, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>29</th>\n",
       "      <td>[-DOCSTART-]</td>\n",
       "      <td>[O]</td>\n",
       "      <td>[101, 118, 141, 9244, 9272, 12426, 1942, 118, ...</td>\n",
       "      <td>[None, 0, 0, 0, 0, 0, 0, 0, None]</td>\n",
       "      <td>[-100, 5, -100, -100, -100, -100, -100, -100, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>30</th>\n",
       "      <td>[China, says, Taiwan, spoils, atmosphere, for,...</td>\n",
       "      <td>[B-LOC, O, B-LOC, O, O, O, O, O]</td>\n",
       "      <td>[101, 1975, 1867, 6036, 188, 5674, 8825, 6814,...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 3, 3, 4, 5, 6, 7, None]</td>\n",
       "      <td>[-100, 3, 5, 3, 5, -100, -100, 5, 5, 5, 5, -100]</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>31</th>\n",
       "      <td>[BEIJING, 1996-08-22]</td>\n",
       "      <td>[B-LOC, O]</td>\n",
       "      <td>[101, 139, 27514, 4538, 15740, 1820, 118, 4775...</td>\n",
       "      <td>[None, 0, 0, 0, 0, 1, 1, 1, 1, 1, None]</td>\n",
       "      <td>[-100, 3, -100, -100, -100, 5, -100, -100, -10...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>32</th>\n",
       "      <td>[China, on, Thursday, accused, Taipei, of, spo...</td>\n",
       "      <td>[B-LOC, O, O, O, B-LOC, O, O, O, O, O, O, O, O...</td>\n",
       "      <td>[101, 1975, 1113, 9170, 4806, 14512, 1104, 188...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 5, 6, 6, 6, 7, 8, 9, 10,...</td>\n",
       "      <td>[-100, 3, 5, 5, 5, 3, 5, 5, -100, -100, 5, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>33</th>\n",
       "      <td>[Speaking, only, hours, after, Chinese, state,...</td>\n",
       "      <td>[O, O, O, O, B-MISC, O, O, O, O, O, O, O, O, O...</td>\n",
       "      <td>[101, 14072, 1178, 2005, 1170, 1922, 1352, 239...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...</td>\n",
       "      <td>[-100, 5, 5, 5, 5, 6, 5, 5, 5, 5, 5, 5, 5, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>34</th>\n",
       "      <td>[State, media, quoted, China, 's, top, negotia...</td>\n",
       "      <td>[O, O, O, B-LOC, O, O, O, O, B-LOC, O, B-PER, ...</td>\n",
       "      <td>[101, 1426, 2394, 9129, 1975, 112, 188, 1499, ...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 4, 5, 6, 6, 6, 7, 8, 9, ...</td>\n",
       "      <td>[-100, 5, 5, 5, 3, 5, -100, 5, 5, -100, -100, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>35</th>\n",
       "      <td>[\", Now, is, the, time, for, the, two, sides, ...</td>\n",
       "      <td>[O, O, O, O, O, O, O, O, O, O, O, O, O, O, O]</td>\n",
       "      <td>[101, 107, 1986, 1110, 1103, 1159, 1111, 1103,...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...</td>\n",
       "      <td>[-100, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>36</th>\n",
       "      <td>[that, is, to, end, the, state, of, hostility,...</td>\n",
       "      <td>[O, O, O, O, O, O, O, O, O, O, O, O, O, O, O, ...</td>\n",
       "      <td>[101, 1115, 1110, 1106, 1322, 1103, 1352, 1104...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...</td>\n",
       "      <td>[-100, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, -10...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>37</th>\n",
       "      <td>[The, foreign, ministry, 's, Shen, told, Reute...</td>\n",
       "      <td>[O, O, O, O, B-ORG, O, B-ORG, I-ORG, O, O, O, ...</td>\n",
       "      <td>[101, 1109, 2880, 8382, 112, 188, 26197, 1500,...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 3, 4, 5, 6, 6, 7, 8, 9, 10,...</td>\n",
       "      <td>[-100, 5, 5, 5, 5, -100, 2, 5, 2, -100, 1, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>38</th>\n",
       "      <td>[China, ,, which, has, long, opposed, all, Tai...</td>\n",
       "      <td>[B-LOC, O, O, O, O, O, O, B-LOC, O, O, O, O, O...</td>\n",
       "      <td>[101, 1975, 117, 1134, 1144, 1263, 4151, 1155,...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...</td>\n",
       "      <td>[-100, 3, 5, 5, 5, 5, 5, 5, 3, 5, 5, 5, 5, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>39</th>\n",
       "      <td>[-DOCSTART-]</td>\n",
       "      <td>[O]</td>\n",
       "      <td>[101, 118, 141, 9244, 9272, 12426, 1942, 118, ...</td>\n",
       "      <td>[None, 0, 0, 0, 0, 0, 0, 0, None]</td>\n",
       "      <td>[-100, 5, -100, -100, -100, -100, -100, -100, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>40</th>\n",
       "      <td>[China, says, time, right, for, Taiwan, talks, .]</td>\n",
       "      <td>[B-LOC, O, O, O, O, B-LOC, O, O]</td>\n",
       "      <td>[101, 1975, 1867, 1159, 1268, 1111, 6036, 7430...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 5, 6, 7, None]</td>\n",
       "      <td>[-100, 3, 5, 5, 5, 5, 3, 5, 5, -100]</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>41</th>\n",
       "      <td>[BEIJING, 1996-08-22]</td>\n",
       "      <td>[B-LOC, O]</td>\n",
       "      <td>[101, 139, 27514, 4538, 15740, 1820, 118, 4775...</td>\n",
       "      <td>[None, 0, 0, 0, 0, 1, 1, 1, 1, 1, None]</td>\n",
       "      <td>[-100, 3, -100, -100, -100, 5, -100, -100, -10...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>42</th>\n",
       "      <td>[China, has, said, it, was, time, for, politic...</td>\n",
       "      <td>[B-LOC, O, O, O, O, O, O, O, O, O, B-LOC, O, O...</td>\n",
       "      <td>[101, 1975, 1144, 1163, 1122, 1108, 1159, 1111...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...</td>\n",
       "      <td>[-100, 3, 5, 5, 5, 5, 5, 5, 5, 5, 5, 3, 5, 5, ...</td>\n",
       "      <td>0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                words  \\\n",
       "0   [EU, rejects, German, call, to, boycott, Briti...   \n",
       "1                                  [Peter, Blackburn]   \n",
       "2                              [BRUSSELS, 1996-08-22]   \n",
       "3   [The, European, Commission, said, on, Thursday...   \n",
       "4   [Germany, 's, representative, to, the, Europea...   \n",
       "5   [\", We, do, n't, support, any, such, recommend...   \n",
       "6   [He, said, further, scientific, study, was, re...   \n",
       "7   [He, said, a, proposal, last, month, by, EU, F...   \n",
       "8   [Fischler, proposed, EU-wide, measures, after,...   \n",
       "9   [But, Fischler, agreed, to, review, his, propo...   \n",
       "10  [Spanish, Farm, Minister, Loyola, de, Palacio,...   \n",
       "11                                                [.]   \n",
       "12  [Only, France, and, Britain, backed, Fischler,...   \n",
       "13  [The, EU, 's, scientific, veterinary, and, mul...   \n",
       "14  [Sheep, have, long, been, known, to, contract,...   \n",
       "15  [British, farmers, denied, on, Thursday, there...   \n",
       "16  [\", What, we, have, to, be, extremely, careful...   \n",
       "17  [Bonn, has, led, efforts, to, protect, public,...   \n",
       "18  [Germany, imported, 47,600, sheep, from, Brita...   \n",
       "19  [It, brought, in, 4,275, tonnes, of, British, ...   \n",
       "20                                       [-DOCSTART-]   \n",
       "21  [Rare, Hendrix, song, draft, sells, for, almos...   \n",
       "22                               [LONDON, 1996-08-22]   \n",
       "23  [A, rare, early, handwritten, draft, of, a, so...   \n",
       "24  [A, Florida, restaurant, paid, 10,925, pounds,...   \n",
       "25  [At, the, end, of, a, January, 1967, concert, ...   \n",
       "26  [Buyers, also, snapped, up, 16, other, items, ...   \n",
       "27  [They, included, a, black, lacquer, and, mothe...   \n",
       "28  [The, guitarist, died, of, a, drugs, overdose,...   \n",
       "29                                       [-DOCSTART-]   \n",
       "30  [China, says, Taiwan, spoils, atmosphere, for,...   \n",
       "31                              [BEIJING, 1996-08-22]   \n",
       "32  [China, on, Thursday, accused, Taipei, of, spo...   \n",
       "33  [Speaking, only, hours, after, Chinese, state,...   \n",
       "34  [State, media, quoted, China, 's, top, negotia...   \n",
       "35  [\", Now, is, the, time, for, the, two, sides, ...   \n",
       "36  [that, is, to, end, the, state, of, hostility,...   \n",
       "37  [The, foreign, ministry, 's, Shen, told, Reute...   \n",
       "38  [China, ,, which, has, long, opposed, all, Tai...   \n",
       "39                                       [-DOCSTART-]   \n",
       "40  [China, says, time, right, for, Taiwan, talks, .]   \n",
       "41                              [BEIJING, 1996-08-22]   \n",
       "42  [China, has, said, it, was, time, for, politic...   \n",
       "\n",
       "                                           str_labels  \\\n",
       "0           [B-ORG, O, B-MISC, O, O, O, B-MISC, O, O]   \n",
       "1                                      [B-PER, I-PER]   \n",
       "2                                          [B-LOC, O]   \n",
       "3   [O, B-ORG, I-ORG, O, O, O, O, O, O, B-MISC, O,...   \n",
       "4   [B-LOC, O, O, O, O, B-ORG, I-ORG, O, O, O, B-P...   \n",
       "5   [O, O, O, O, O, O, O, O, O, O, O, O, O, O, O, ...   \n",
       "6   [O, O, O, O, O, O, O, O, O, O, O, O, O, O, O, ...   \n",
       "7   [O, O, O, O, O, O, O, B-ORG, O, O, B-PER, I-PE...   \n",
       "8   [B-PER, O, B-MISC, O, O, O, O, B-LOC, O, B-LOC...   \n",
       "9   [O, B-PER, O, O, O, O, O, O, O, B-ORG, O, O, O...   \n",
       "10  [B-MISC, O, O, B-PER, I-PER, I-PER, O, O, O, B...   \n",
       "11                                                [O]   \n",
       "12            [O, B-LOC, O, B-LOC, O, B-PER, O, O, O]   \n",
       "13  [O, B-ORG, O, O, O, O, O, O, O, O, O, O, O, O,...   \n",
       "14  [O, O, O, O, O, O, O, O, O, O, O, O, O, O, B-M...   \n",
       "15  [B-MISC, O, O, O, O, O, O, O, O, O, O, O, O, O...   \n",
       "16  [O, O, O, O, O, O, O, O, O, O, O, O, O, O, O, ...   \n",
       "17  [B-LOC, O, O, O, O, O, O, O, O, O, O, O, O, O,...   \n",
       "18  [B-LOC, O, O, O, O, B-LOC, O, O, O, O, O, O, O...   \n",
       "19  [O, O, O, O, O, O, B-MISC, O, O, O, O, O, O, O...   \n",
       "20                                                [O]   \n",
       "21                 [O, B-PER, O, O, O, O, O, O, O, O]   \n",
       "22                                         [B-LOC, O]   \n",
       "23  [O, O, O, O, O, O, O, O, O, B-LOC, O, O, B-PER...   \n",
       "24  [O, B-LOC, O, O, O, O, O, O, O, O, O, O, O, O,...   \n",
       "25  [O, O, O, O, O, O, O, O, O, O, B-MISC, O, O, B...   \n",
       "26  [O, O, O, O, O, O, O, O, O, O, O, O, O, O, B-P...   \n",
       "27  [O, O, O, O, O, O, O, O, O, O, O, O, O, B-PER,...   \n",
       "28               [O, O, O, O, O, O, O, O, O, O, O, O]   \n",
       "29                                                [O]   \n",
       "30                   [B-LOC, O, B-LOC, O, O, O, O, O]   \n",
       "31                                         [B-LOC, O]   \n",
       "32  [B-LOC, O, O, O, B-LOC, O, O, O, O, O, O, O, O...   \n",
       "33  [O, O, O, O, B-MISC, O, O, O, O, O, O, O, O, O...   \n",
       "34  [O, O, O, B-LOC, O, O, O, O, B-LOC, O, B-PER, ...   \n",
       "35      [O, O, O, O, O, O, O, O, O, O, O, O, O, O, O]   \n",
       "36  [O, O, O, O, O, O, O, O, O, O, O, O, O, O, O, ...   \n",
       "37  [O, O, O, O, B-ORG, O, B-ORG, I-ORG, O, O, O, ...   \n",
       "38  [B-LOC, O, O, O, O, O, O, B-LOC, O, O, O, O, O...   \n",
       "39                                                [O]   \n",
       "40                   [B-LOC, O, O, O, O, B-LOC, O, O]   \n",
       "41                                         [B-LOC, O]   \n",
       "42  [B-LOC, O, O, O, O, O, O, O, O, O, B-LOC, O, O...   \n",
       "\n",
       "                                            input_ids  \\\n",
       "0   [101, 7270, 22961, 1528, 1840, 1106, 21423, 14...   \n",
       "1                             [101, 1943, 14428, 102]   \n",
       "2   [101, 26660, 13329, 12649, 15928, 1820, 118, 4...   \n",
       "3   [101, 1109, 1735, 2827, 1163, 1113, 9170, 1122...   \n",
       "4   [101, 1860, 112, 188, 4702, 1106, 1103, 1735, ...   \n",
       "5   [101, 107, 1284, 1202, 183, 112, 189, 1619, 12...   \n",
       "6   [101, 1124, 1163, 1748, 3812, 2025, 1108, 2320...   \n",
       "7   [101, 1124, 1163, 170, 5835, 1314, 2370, 1118,...   \n",
       "8   [101, 17355, 9022, 2879, 3000, 7270, 118, 2043...   \n",
       "9   [101, 1252, 17355, 9022, 2879, 2675, 1106, 318...   \n",
       "10  [101, 2124, 6865, 2110, 23828, 1260, 19585, 17...   \n",
       "11                                    [101, 119, 102]   \n",
       "12  [101, 2809, 1699, 1105, 2855, 5534, 17355, 902...   \n",
       "13  [101, 1109, 7270, 112, 188, 3812, 27431, 1105,...   \n",
       "14  [101, 1153, 8043, 1138, 1263, 1151, 1227, 1106...   \n",
       "15  [101, 1418, 6915, 5762, 1113, 9170, 1175, 1108...   \n",
       "16  [101, 107, 1327, 1195, 1138, 1106, 1129, 4450,...   \n",
       "17  [101, 21066, 1144, 1521, 3268, 1106, 3244, 147...   \n",
       "18  [101, 1860, 11095, 3862, 117, 4372, 8892, 1121...   \n",
       "19  [101, 1135, 1814, 1107, 125, 117, 18783, 10992...   \n",
       "20  [101, 118, 141, 9244, 9272, 12426, 1942, 118, ...   \n",
       "21  [101, 25574, 22609, 1461, 5039, 16695, 1111, 1...   \n",
       "22  [101, 149, 11414, 2137, 11414, 1820, 118, 4775...   \n",
       "23  [101, 138, 4054, 1346, 1289, 16674, 5039, 1104...   \n",
       "24  [101, 138, 2631, 4382, 3004, 1275, 117, 5556, ...   \n",
       "25  [101, 1335, 1103, 1322, 1104, 170, 1356, 2573,...   \n",
       "26  [101, 26123, 1468, 1145, 5267, 1146, 1479, 116...   \n",
       "27  [101, 1220, 1529, 170, 1602, 2495, 1665, 19061...   \n",
       "28  [101, 1109, 5506, 1452, 1104, 170, 5557, 1166,...   \n",
       "29  [101, 118, 141, 9244, 9272, 12426, 1942, 118, ...   \n",
       "30  [101, 1975, 1867, 6036, 188, 5674, 8825, 6814,...   \n",
       "31  [101, 139, 27514, 4538, 15740, 1820, 118, 4775...   \n",
       "32  [101, 1975, 1113, 9170, 4806, 14512, 1104, 188...   \n",
       "33  [101, 14072, 1178, 2005, 1170, 1922, 1352, 239...   \n",
       "34  [101, 1426, 2394, 9129, 1975, 112, 188, 1499, ...   \n",
       "35  [101, 107, 1986, 1110, 1103, 1159, 1111, 1103,...   \n",
       "36  [101, 1115, 1110, 1106, 1322, 1103, 1352, 1104...   \n",
       "37  [101, 1109, 2880, 8382, 112, 188, 26197, 1500,...   \n",
       "38  [101, 1975, 117, 1134, 1144, 1263, 4151, 1155,...   \n",
       "39  [101, 118, 141, 9244, 9272, 12426, 1942, 118, ...   \n",
       "40  [101, 1975, 1867, 1159, 1268, 1111, 6036, 7430...   \n",
       "41  [101, 139, 27514, 4538, 15740, 1820, 118, 4775...   \n",
       "42  [101, 1975, 1144, 1163, 1122, 1108, 1159, 1111...   \n",
       "\n",
       "                                             word_ids  \\\n",
       "0          [None, 0, 1, 2, 3, 4, 5, 6, 7, 7, 8, None]   \n",
       "1                                  [None, 0, 1, None]   \n",
       "2             [None, 0, 0, 0, 0, 1, 1, 1, 1, 1, None]   \n",
       "3   [None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...   \n",
       "4   [None, 0, 1, 1, 2, 3, 4, 5, 6, 7, 7, 8, 9, 10,...   \n",
       "5   [None, 0, 1, 2, 3, 3, 3, 4, 5, 6, 7, 8, 9, 10,...   \n",
       "6   [None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...   \n",
       "7   [None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...   \n",
       "8   [None, 0, 0, 0, 1, 2, 2, 2, 3, 4, 5, 6, 7, 8, ...   \n",
       "9   [None, 0, 1, 1, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10,...   \n",
       "10  [None, 0, 1, 2, 3, 4, 5, 5, 5, 6, 7, 8, 9, 9, ...   \n",
       "11                                    [None, 0, None]   \n",
       "12   [None, 0, 1, 2, 3, 4, 5, 5, 5, 6, 6, 7, 8, None]   \n",
       "13  [None, 0, 1, 2, 2, 3, 4, 5, 6, 6, 7, 8, 9, 10,...   \n",
       "14  [None, 0, 0, 1, 2, 3, 4, 5, 6, 7, 7, 8, 9, 10,...   \n",
       "15  [None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...   \n",
       "16  [None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...   \n",
       "17  [None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...   \n",
       "18  [None, 0, 1, 2, 2, 2, 3, 4, 5, 6, 7, 8, 9, 10,...   \n",
       "19  [None, 0, 1, 2, 3, 3, 3, 4, 5, 6, 7, 7, 7, 8, ...   \n",
       "20                  [None, 0, 0, 0, 0, 0, 0, 0, None]   \n",
       "21   [None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 8, 8, 9, None]   \n",
       "22            [None, 0, 0, 0, 0, 1, 1, 1, 1, 1, None]   \n",
       "23  [None, 0, 1, 2, 3, 3, 4, 5, 6, 7, 8, 9, 9, 9, ...   \n",
       "24  [None, 0, 1, 2, 3, 4, 4, 4, 4, 5, 6, 7, 8, 8, ...   \n",
       "25  [None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...   \n",
       "26  [None, 0, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11...   \n",
       "27  [None, 0, 1, 2, 3, 4, 4, 4, 5, 6, 7, 8, 9, 9, ...   \n",
       "28  [None, 0, 1, 2, 3, 4, 5, 6, 6, 6, 7, 8, 9, 10,...   \n",
       "29                  [None, 0, 0, 0, 0, 0, 0, 0, None]   \n",
       "30         [None, 0, 1, 2, 3, 3, 3, 4, 5, 6, 7, None]   \n",
       "31            [None, 0, 0, 0, 0, 1, 1, 1, 1, 1, None]   \n",
       "32  [None, 0, 1, 2, 3, 4, 5, 6, 6, 6, 7, 8, 9, 10,...   \n",
       "33  [None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...   \n",
       "34  [None, 0, 1, 2, 3, 4, 4, 5, 6, 6, 6, 7, 8, 9, ...   \n",
       "35  [None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...   \n",
       "36  [None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...   \n",
       "37  [None, 0, 1, 2, 3, 3, 4, 5, 6, 6, 7, 8, 9, 10,...   \n",
       "38  [None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...   \n",
       "39                  [None, 0, 0, 0, 0, 0, 0, 0, None]   \n",
       "40               [None, 0, 1, 2, 3, 4, 5, 6, 7, None]   \n",
       "41            [None, 0, 0, 0, 0, 1, 1, 1, 1, 1, None]   \n",
       "42  [None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 1...   \n",
       "\n",
       "                                               labels  task_ids  \n",
       "0       [-100, 2, 5, 6, 5, 5, 5, 6, 5, -100, 5, -100]         0  \n",
       "1                                  [-100, 4, 8, -100]         0  \n",
       "2   [-100, 3, -100, -100, -100, 5, -100, -100, -10...         0  \n",
       "3   [-100, 5, 2, 1, 5, 5, 5, 5, 5, 5, 6, 5, 5, 5, ...         0  \n",
       "4   [-100, 3, 5, -100, 5, 5, 5, 2, 1, 5, -100, 5, ...         0  \n",
       "5   [-100, 5, 5, 5, 5, -100, -100, 5, 5, 5, 5, 5, ...         0  \n",
       "6   [-100, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, ...         0  \n",
       "7   [-100, 5, 5, 5, 5, 5, 5, 5, 2, 5, 5, 4, 8, -10...         0  \n",
       "8   [-100, 4, -100, -100, 5, 6, -100, -100, 5, 5, ...         0  \n",
       "9   [-100, 5, 4, -100, -100, 5, 5, 5, 5, 5, 5, 5, ...         0  \n",
       "10  [-100, 6, 5, 5, 4, 8, 8, -100, -100, 5, 5, 5, ...         0  \n",
       "11                                    [-100, 5, -100]         0  \n",
       "12  [-100, 5, 3, 5, 3, 5, 4, -100, -100, 5, -100, ...         0  \n",
       "13  [-100, 5, 2, 5, -100, 5, 5, 5, 5, -100, 5, 5, ...         0  \n",
       "14  [-100, 5, -100, 5, 5, 5, 5, 5, 5, 5, -100, 5, ...         0  \n",
       "15  [-100, 6, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, ...         0  \n",
       "16  [-100, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, ...         0  \n",
       "17  [-100, 3, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, ...         0  \n",
       "18  [-100, 3, 5, 5, -100, -100, 5, 5, 3, 5, 5, 5, ...         0  \n",
       "19  [-100, 5, 5, 5, 5, -100, -100, 5, 5, 6, 5, -10...         0  \n",
       "20  [-100, 5, -100, -100, -100, -100, -100, -100, ...         0  \n",
       "21  [-100, 5, 4, 5, 5, 5, 5, 5, 5, 5, -100, -100, ...         0  \n",
       "22  [-100, 3, -100, -100, -100, 5, -100, -100, -10...         0  \n",
       "23  [-100, 5, 5, 5, 5, -100, 5, 5, 5, 5, 5, 3, -10...         0  \n",
       "24  [-100, 5, 3, 5, 5, 5, -100, -100, -100, 5, 5, ...         0  \n",
       "25  [-100, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 6, 5, 5, ...         0  \n",
       "26  [-100, 5, -100, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, ...         0  \n",
       "27  [-100, 5, 5, 5, 5, 5, -100, -100, 5, 5, 5, 5, ...         0  \n",
       "28  [-100, 5, 5, 5, 5, 5, 5, 5, -100, -100, 5, 5, ...         0  \n",
       "29  [-100, 5, -100, -100, -100, -100, -100, -100, ...         0  \n",
       "30   [-100, 3, 5, 3, 5, -100, -100, 5, 5, 5, 5, -100]         0  \n",
       "31  [-100, 3, -100, -100, -100, 5, -100, -100, -10...         0  \n",
       "32  [-100, 3, 5, 5, 5, 3, 5, 5, -100, -100, 5, 5, ...         0  \n",
       "33  [-100, 5, 5, 5, 5, 6, 5, 5, 5, 5, 5, 5, 5, 5, ...         0  \n",
       "34  [-100, 5, 5, 5, 3, 5, -100, 5, 5, -100, -100, ...         0  \n",
       "35  [-100, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, ...         0  \n",
       "36  [-100, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, 5, -10...         0  \n",
       "37  [-100, 5, 5, 5, 5, -100, 2, 5, 2, -100, 1, 5, ...         0  \n",
       "38  [-100, 3, 5, 5, 5, 5, 5, 5, 3, 5, 5, 5, 5, 5, ...         0  \n",
       "39  [-100, 5, -100, -100, -100, -100, -100, -100, ...         0  \n",
       "40               [-100, 3, 5, 5, 5, 5, 3, 5, 5, -100]         0  \n",
       "41  [-100, 3, -100, -100, -100, 5, -100, -100, -10...         0  \n",
       "42  [-100, 3, 5, 5, 5, 5, 5, 5, 5, 5, 5, 3, 5, 5, ...         0  "
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ner_task.train_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "a4a0c793",
   "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>words</th>\n",
       "      <th>str_labels</th>\n",
       "      <th>input_ids</th>\n",
       "      <th>word_ids</th>\n",
       "      <th>labels</th>\n",
       "      <th>task_ids</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>[Pierre, Vinken, ,, 61, years, old, ,, will, j...</td>\n",
       "      <td>[NNP, NNP, ,, CD, NNS, JJ, ,, MD, VB, DT, NN, ...</td>\n",
       "      <td>[101, 4855, 25354, 6378, 117, 5391, 1201, 1385...</td>\n",
       "      <td>[None, 0, 1, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11...</td>\n",
       "      <td>[-100, 10, 10, -100, 31, 35, 23, 0, 31, 13, 36...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>[Mr., Vinken, is, chairman, of, Elsevier, N.V....</td>\n",
       "      <td>[NNP, NNP, VBZ, NN, IN, NNP, NNP, ,, DT, NNP, ...</td>\n",
       "      <td>[101, 1828, 119, 25354, 6378, 1110, 3931, 1104...</td>\n",
       "      <td>[None, 0, 0, 1, 1, 2, 3, 4, 5, 5, 5, 6, 6, 6, ...</td>\n",
       "      <td>[-100, 10, -100, 10, -100, 21, 33, 38, 10, -10...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>[Rudolph, Agnew, ,, 55, years, old, and, forme...</td>\n",
       "      <td>[NNP, NNP, ,, CD, NNS, JJ, CC, JJ, NN, IN, NNP...</td>\n",
       "      <td>[101, 19922, 138, 8376, 2246, 117, 3731, 1201,...</td>\n",
       "      <td>[None, 0, 1, 1, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10,...</td>\n",
       "      <td>[-100, 10, 10, -100, -100, 31, 35, 23, 0, 43, ...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>[A, form, of, asbestos, once, used, to, make, ...</td>\n",
       "      <td>[DT, NN, IN, NN, RB, VBN, TO, VB, NNP, NN, NNS...</td>\n",
       "      <td>[101, 138, 1532, 1104, 1112, 12866, 11990, 151...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 3, 3, 4, 5, 6, 7, 8, 9, 10,...</td>\n",
       "      <td>[-100, 2, 33, 38, 33, -100, -100, 18, 24, 27, ...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>[The, asbestos, fiber, ,, crocidolite, ,, is, ...</td>\n",
       "      <td>[DT, NN, NN, ,, NN, ,, VBZ, RB, JJ, IN, PRP, V...</td>\n",
       "      <td>[101, 1109, 1112, 12866, 11990, 12753, 117, 17...</td>\n",
       "      <td>[None, 0, 1, 1, 1, 2, 3, 4, 4, 4, 4, 4, 5, 6, ...</td>\n",
       "      <td>[-100, 2, 33, -100, -100, 33, 31, 33, -100, -1...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>196</th>\n",
       "      <td>[On, the, other, hand, ,, had, it, existed, th...</td>\n",
       "      <td>[IN, DT, JJ, NN, ,, VBD, PRP, VBN, RB, ,, NNP,...</td>\n",
       "      <td>[101, 1212, 1103, 1168, 1289, 117, 1125, 1122,...</td>\n",
       "      <td>[None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 10, 1...</td>\n",
       "      <td>[-100, 38, 2, 0, 33, 31, 34, 3, 24, 18, 31, 10...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>197</th>\n",
       "      <td>[Mr., Cray, ,, who, could, n't, be, reached, f...</td>\n",
       "      <td>[NNP, NNP, ,, WP, MD, RB, VB, VBN, IN, NN, ,, ...</td>\n",
       "      <td>[101, 1828, 119, 140, 6447, 117, 1150, 1180, 1...</td>\n",
       "      <td>[None, 0, 0, 1, 1, 2, 3, 4, 5, 5, 5, 6, 7, 8, ...</td>\n",
       "      <td>[-100, 10, -100, 10, -100, 31, 7, 13, 18, -100...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>198</th>\n",
       "      <td>[Regarded, as, the, father, of, the, supercomp...</td>\n",
       "      <td>[VBN, IN, DT, NN, IN, DT, NN, ,, NNP, NNP, VBD...</td>\n",
       "      <td>[101, 23287, 26541, 1112, 1103, 1401, 1104, 11...</td>\n",
       "      <td>[None, 0, 0, 1, 2, 3, 4, 5, 6, 6, 6, 6, 7, 8, ...</td>\n",
       "      <td>[-100, 24, -100, 38, 2, 33, 38, 2, 33, -100, -...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>199</th>\n",
       "      <td>[At, Cray, Computer, ,, he, will, be, paid, $,...</td>\n",
       "      <td>[IN, NNP, NNP, ,, PRP, MD, VB, VBN, $, CD, .]</td>\n",
       "      <td>[101, 1335, 140, 6447, 6701, 117, 1119, 1209, ...</td>\n",
       "      <td>[None, 0, 1, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 9, ...</td>\n",
       "      <td>[-100, 38, 10, -100, 10, 31, 3, 13, 36, 24, 41...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>200</th>\n",
       "      <td>[(, ), blah, merde, uh, pdt, #, ls, @, .]</td>\n",
       "      <td>[-LRB-, -RRB-, RBS, FW, UH, PDT, #, LS, SYM, .]</td>\n",
       "      <td>[101, 113, 114, 171, 10358, 1143, 14407, 14863...</td>\n",
       "      <td>[None, 0, 1, 2, 2, 3, 3, 4, 5, 5, 6, 7, 7, 8, ...</td>\n",
       "      <td>[-100, 25, 5, 19, -100, 8, -100, 28, 9, -100, ...</td>\n",
       "      <td>1</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>201 rows × 6 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                 words  \\\n",
       "0    [Pierre, Vinken, ,, 61, years, old, ,, will, j...   \n",
       "1    [Mr., Vinken, is, chairman, of, Elsevier, N.V....   \n",
       "2    [Rudolph, Agnew, ,, 55, years, old, and, forme...   \n",
       "3    [A, form, of, asbestos, once, used, to, make, ...   \n",
       "4    [The, asbestos, fiber, ,, crocidolite, ,, is, ...   \n",
       "..                                                 ...   \n",
       "196  [On, the, other, hand, ,, had, it, existed, th...   \n",
       "197  [Mr., Cray, ,, who, could, n't, be, reached, f...   \n",
       "198  [Regarded, as, the, father, of, the, supercomp...   \n",
       "199  [At, Cray, Computer, ,, he, will, be, paid, $,...   \n",
       "200          [(, ), blah, merde, uh, pdt, #, ls, @, .]   \n",
       "\n",
       "                                            str_labels  \\\n",
       "0    [NNP, NNP, ,, CD, NNS, JJ, ,, MD, VB, DT, NN, ...   \n",
       "1    [NNP, NNP, VBZ, NN, IN, NNP, NNP, ,, DT, NNP, ...   \n",
       "2    [NNP, NNP, ,, CD, NNS, JJ, CC, JJ, NN, IN, NNP...   \n",
       "3    [DT, NN, IN, NN, RB, VBN, TO, VB, NNP, NN, NNS...   \n",
       "4    [DT, NN, NN, ,, NN, ,, VBZ, RB, JJ, IN, PRP, V...   \n",
       "..                                                 ...   \n",
       "196  [IN, DT, JJ, NN, ,, VBD, PRP, VBN, RB, ,, NNP,...   \n",
       "197  [NNP, NNP, ,, WP, MD, RB, VB, VBN, IN, NN, ,, ...   \n",
       "198  [VBN, IN, DT, NN, IN, DT, NN, ,, NNP, NNP, VBD...   \n",
       "199      [IN, NNP, NNP, ,, PRP, MD, VB, VBN, $, CD, .]   \n",
       "200    [-LRB-, -RRB-, RBS, FW, UH, PDT, #, LS, SYM, .]   \n",
       "\n",
       "                                             input_ids  \\\n",
       "0    [101, 4855, 25354, 6378, 117, 5391, 1201, 1385...   \n",
       "1    [101, 1828, 119, 25354, 6378, 1110, 3931, 1104...   \n",
       "2    [101, 19922, 138, 8376, 2246, 117, 3731, 1201,...   \n",
       "3    [101, 138, 1532, 1104, 1112, 12866, 11990, 151...   \n",
       "4    [101, 1109, 1112, 12866, 11990, 12753, 117, 17...   \n",
       "..                                                 ...   \n",
       "196  [101, 1212, 1103, 1168, 1289, 117, 1125, 1122,...   \n",
       "197  [101, 1828, 119, 140, 6447, 117, 1150, 1180, 1...   \n",
       "198  [101, 23287, 26541, 1112, 1103, 1401, 1104, 11...   \n",
       "199  [101, 1335, 140, 6447, 6701, 117, 1119, 1209, ...   \n",
       "200  [101, 113, 114, 171, 10358, 1143, 14407, 14863...   \n",
       "\n",
       "                                              word_ids  \\\n",
       "0    [None, 0, 1, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11...   \n",
       "1    [None, 0, 0, 1, 1, 2, 3, 4, 5, 5, 5, 6, 6, 6, ...   \n",
       "2    [None, 0, 1, 1, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10,...   \n",
       "3    [None, 0, 1, 2, 3, 3, 3, 4, 5, 6, 7, 8, 9, 10,...   \n",
       "4    [None, 0, 1, 1, 1, 2, 3, 4, 4, 4, 4, 4, 5, 6, ...   \n",
       "..                                                 ...   \n",
       "196  [None, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 10, 1...   \n",
       "197  [None, 0, 0, 1, 1, 2, 3, 4, 5, 5, 5, 6, 7, 8, ...   \n",
       "198  [None, 0, 0, 1, 2, 3, 4, 5, 6, 6, 6, 6, 7, 8, ...   \n",
       "199  [None, 0, 1, 1, 2, 3, 4, 5, 6, 7, 8, 9, 9, 9, ...   \n",
       "200  [None, 0, 1, 2, 2, 3, 3, 4, 5, 5, 6, 7, 7, 8, ...   \n",
       "\n",
       "                                                labels  task_ids  \n",
       "0    [-100, 10, 10, -100, 31, 35, 23, 0, 31, 13, 36...         1  \n",
       "1    [-100, 10, -100, 10, -100, 21, 33, 38, 10, -10...         1  \n",
       "2    [-100, 10, 10, -100, -100, 31, 35, 23, 0, 43, ...         1  \n",
       "3    [-100, 2, 33, 38, 33, -100, -100, 18, 24, 27, ...         1  \n",
       "4    [-100, 2, 33, -100, -100, 33, 31, 33, -100, -1...         1  \n",
       "..                                                 ...       ...  \n",
       "196  [-100, 38, 2, 0, 33, 31, 34, 3, 24, 18, 31, 10...         1  \n",
       "197  [-100, 10, -100, 10, -100, 31, 7, 13, 18, -100...         1  \n",
       "198  [-100, 24, -100, 38, 2, 33, 38, 2, 33, -100, -...         1  \n",
       "199  [-100, 38, 10, -100, 10, 31, 3, 13, 36, 24, 41...         1  \n",
       "200  [-100, 25, 5, 19, -100, 8, -100, 28, 9, -100, ...         1  \n",
       "\n",
       "[201 rows x 6 columns]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "pos_task.train_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "798f2b24",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers.modeling_outputs import TokenClassifierOutput\n",
    "from transformers.models.bert.modeling_bert import BertModel, BertPreTrainedModel\n",
    "from transformers import PreTrainedModel\n",
    "from transformers import AutoConfig, AutoModel\n",
    "import os\n",
    "\n",
    "# This class is adapted from: https://towardsdatascience.com/how-to-create-and-train-a-multi-task-transformer-model-18c54a146240\n",
    "class TokenClassificationModel(BertPreTrainedModel):    \n",
    "    def __init__(self, config):\n",
    "        super().__init__(config)\n",
    "        self.bert = BertModel(config, add_pooling_layer=False)\n",
    "        self.output_heads = nn.ModuleDict() # these are initialized in add_heads\n",
    "        self.init_weights()\n",
    "        \n",
    "    def add_heads(self, tasks):\n",
    "        for task in tasks:\n",
    "            head = TokenClassificationHead(self.bert.config.hidden_size, task.num_labels, config.hidden_dropout_prob)\n",
    "            # ModuleDict requires keys to be strings\n",
    "            self.output_heads[str(task.task_id)] = head\n",
    "        return self\n",
    "    \n",
    "    def summarize_heads(self):\n",
    "        print(f'Found {len(self.output_heads)} heads')\n",
    "        for task_id in self.output_heads:\n",
    "            self.output_heads[task_id].summarize(task_id)\n",
    "        \n",
    "    def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, labels=None, task_ids=None, **kwargs):\n",
    "        outputs = self.bert(\n",
    "            input_ids,\n",
    "            attention_mask=attention_mask,\n",
    "            token_type_ids=token_type_ids,\n",
    "            **kwargs,\n",
    "        )\n",
    "        sequence_output = outputs[0]\n",
    "        \n",
    "        #print(f'batch size = {len(input_ids)}')\n",
    "        #print(f'task_ids in this batch: {task_ids}')\n",
    "        \n",
    "        # generate specific predictions and losses for each task head\n",
    "        unique_task_ids_list = torch.unique(task_ids).tolist()\n",
    "        logits = None\n",
    "        loss_list = []\n",
    "        for unique_task_id in unique_task_ids_list:\n",
    "            task_id_filter = task_ids == unique_task_id\n",
    "            filtered_sequence_output = sequence_output[task_id_filter]\n",
    "            filtered_labels = None if labels is None else labels[task_id_filter]\n",
    "            filtered_attention_mask = None if attention_mask is None else attention_mask[task_id_filter]\n",
    "            #print(f'size of batch for task {unique_task_id} is: {len(filtered_sequence_output)}')\n",
    "            logits, task_loss = self.output_heads[str(unique_task_id)].forward(\n",
    "                filtered_sequence_output, None,\n",
    "                filtered_labels,\n",
    "                filtered_attention_mask,\n",
    "            )\n",
    "            if filtered_labels is not None:\n",
    "                loss_list.append(task_loss)\n",
    "                \n",
    "        loss = None if len(loss_list) == 0 else torch.stack(loss_list)\n",
    "                    \n",
    "        # logits are only used for eval, in which case we handle a single task at a time\n",
    "        # TODO: allow all tasks in the forward pass at inference                     \n",
    "        return TokenClassifierOutput(\n",
    "            loss=loss.mean(),\n",
    "            logits=logits,\n",
    "            hidden_states=outputs.hidden_states,\n",
    "            attentions=outputs.attentions,\n",
    "        )\n",
    "    \n",
    "\n",
    "    def save_task(task_head, task, task_checkpoint):\n",
    "        numpy_weights = task_head.classifier.weight.cpu().detach().numpy()\n",
    "        numpy_bias = task_head.classifier.bias.cpu().detach().numpy()\n",
    "        labels = task.labels\n",
    "        #print(f\"Shape of weights: {numpy_weights.shape}\")\n",
    "        #print(f\"Weights are:\\n{numpy_weights}\")\n",
    "        #print(f\"Shape of bias: {numpy_bias.shape}\")\n",
    "        #print(f\"Bias is: {numpy_bias}\")\n",
    "        #print(f\"Labels are: {labels}\")\n",
    "        \n",
    "        os.makedirs(task_checkpoint, exist_ok = True)\n",
    "        lf = open(task_checkpoint + \"/labels\", \"w\")\n",
    "        for label in labels:\n",
    "            lf.write(f'{label}\\n')\n",
    "        lf.close()\n",
    "        \n",
    "        wf = open(task_checkpoint + \"/weights\", \"w\")\n",
    "        wf.write(f'# {numpy_weights.shape[0]} {numpy_weights.shape[1]}\\n')\n",
    "        for i, x in enumerate(numpy_weights):\n",
    "            for j, y in enumerate(x):\n",
    "                wf.write(f'{y} ')\n",
    "            wf.write('\\n')\n",
    "        wf.close()\n",
    "        \n",
    "        bf = open(task_checkpoint + \"/biases\", \"w\")\n",
    "        bf.write(f'# {numpy_bias.shape[0]}\\n')\n",
    "        for i, x in enumerate(numpy_bias):\n",
    "            bf.write(f'{x} ')\n",
    "        bf.write('\\n')\n",
    "        bf.close()\n",
    "\n",
    "    def onnx_save(self, checkpoint):\n",
    "        orig_words = [\"Using\", \"transformers\", \"with\", \"ONNX\", \"runtime\"]\n",
    "        token_input = tokenizer(orig_words, is_split_into_words = True, return_tensors = \"pt\")\n",
    "        print(token_input)\n",
    "        token_ids = token_input['input_ids'].to(device)\n",
    "        \n",
    "        inputs = (token_ids) \n",
    "        input_names = [\"token_ids\"] \n",
    "        output_names = [\"sequence_output\"]\n",
    "        \n",
    "        torch.onnx.export(self.bert,\n",
    "            inputs,\n",
    "            checkpoint,\n",
    "            export_params=True,\n",
    "            do_constant_folding=True,\n",
    "            input_names = input_names,\n",
    "            output_names = output_names,\n",
    "            opset_version=10, \n",
    "            dynamic_axes = {\"token_ids\": {1: 'sent length'}}\n",
    "        )\n",
    "    \n",
    "    # exports model in a format friendly for ingestion on the JVM\n",
    "    def export_model(self, tasks, checkpoint_dir):\n",
    "        # save the weights/bias in each linear layer\n",
    "        for task in tasks:\n",
    "            task_checkpoint = checkpoint_dir + \"/task_\" + task.task_name\n",
    "            save_task(model.output_heads[str(task.task_id)], task, task_checkpoint)\n",
    "    \n",
    "        # save the encoder as an ONNX model\n",
    "        onnx_checkpoint = checkpoint_dir + '/encoder.onnx'\n",
    "        onnx_save(model, onnx_checkpoint)\n",
    "        \n",
    "\n",
    "class TokenClassificationHead(nn.Module):\n",
    "    def __init__(self, hidden_size, num_labels, dropout_p=0.1):\n",
    "        super().__init__()\n",
    "        self.dropout = nn.Dropout(dropout_p)\n",
    "        self.classifier = nn.Linear(hidden_size, num_labels)\n",
    "        self.num_labels = num_labels\n",
    "\n",
    "        self._init_weights()\n",
    "\n",
    "    def _init_weights(self):\n",
    "        self.classifier.weight.data.normal_(mean=0.0, std=0.02)\n",
    "        if self.classifier.bias is not None:\n",
    "            self.classifier.bias.data.zero_()\n",
    "            \n",
    "    def summarize(self, task_id):\n",
    "        print(f\"Task {task_id} with {self.num_labels} labels.\")\n",
    "        print(f'Dropout is {self.dropout}')\n",
    "        print(f'Classifier layer is {self.classifier}')\n",
    "\n",
    "    def forward(self, sequence_output, pooled_output, labels=None, attention_mask=None, **kwargs):\n",
    "        sequence_output_dropout = self.dropout(sequence_output)\n",
    "        logits = self.classifier(sequence_output_dropout)\n",
    "        \n",
    "        loss = None\n",
    "        if labels is not None:\n",
    "            loss_fn = nn.CrossEntropyLoss()            \n",
    "            inputs = logits.view(-1, self.num_labels)\n",
    "            targets = labels.view(-1)\n",
    "            loss = loss_fn(inputs, targets)\n",
    "\n",
    "        return logits, loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "c9945a6c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found 2 heads\n",
      "Task 0 with 9 labels.\n",
      "Dropout is Dropout(p=0.1, inplace=False)\n",
      "Classifier layer is Linear(in_features=768, out_features=9, bias=True)\n",
      "Task 1 with 45 labels.\n",
      "Dropout is Dropout(p=0.1, inplace=False)\n",
      "Classifier layer is Linear(in_features=768, out_features=45, bias=True)\n"
     ]
    }
   ],
   "source": [
    "tasks = [ner_task, pos_task]\n",
    "config = AutoConfig.from_pretrained(transformer_name)\n",
    "model= TokenClassificationModel.from_pretrained(transformer_name, config=config).add_heads(tasks)\n",
    "model.summarize_heads()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "9c45587e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import accuracy_score\n",
    "\n",
    "def compute_metrics(eval_pred):\n",
    "    # gold labels\n",
    "    label_ids = eval_pred.label_ids\n",
    "    # predictions\n",
    "    pred_ids = np.argmax(eval_pred.predictions, axis=-1)\n",
    "    # collect gold and predicted labels, ignoring ignore_index label\n",
    "    y_true, y_pred = [], []\n",
    "    batch_size, seq_len = pred_ids.shape\n",
    "    for i in range(batch_size):\n",
    "        for j in range(seq_len):\n",
    "            if label_ids[i, j] != ignore_index:\n",
    "                y_true.append(label_ids[i][j]) #index_to_label[label_ids[i][j]])\n",
    "                y_pred.append(pred_ids[i][j]) #index_to_label[pred_ids[i][j]])\n",
    "    # return computed metrics\n",
    "    return {'accuracy': accuracy_score(y_true, y_pred)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "49167b36",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DatasetDict({\n",
       "    train: Dataset({\n",
       "        features: ['words', 'str_labels', 'input_ids', 'word_ids', 'labels', 'task_ids', '__index_level_0__'],\n",
       "        num_rows: 244\n",
       "    })\n",
       "    validation: Dataset({\n",
       "        features: ['words', 'str_labels', 'input_ids', 'word_ids', 'labels', 'task_ids', '__index_level_0__'],\n",
       "        num_rows: 8504\n",
       "    })\n",
       "    test: Dataset({\n",
       "        features: ['words', 'str_labels', 'input_ids', 'word_ids', 'labels', 'task_ids', '__index_level_0__'],\n",
       "        num_rows: 6101\n",
       "    })\n",
       "})"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from datasets import Dataset, DatasetDict\n",
    "\n",
    "ds = DatasetDict()\n",
    "ds['train'] = Dataset.from_pandas(pd.concat([ner_task.train_df, pos_task.train_df]))\n",
    "ds['validation'] = Dataset.from_pandas(pd.concat([ner_task.dev_df, pos_task.dev_df]))\n",
    "ds['test'] = Dataset.from_pandas(pd.concat([ner_task.test_df, pos_task.test_df]))\n",
    "\n",
    "# these are no longer needed; discard them to save memory\n",
    "ner_task.train_df = None\n",
    "pos_task.train_df = None\n",
    "\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "48554618",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.metrics import classification_report\n",
    "\n",
    "# compute accuracy\n",
    "def evaluation_classification_report(trainer, task, name, useTest=False):\n",
    "    print(f\"Test classification report for task {name}:\")\n",
    "    num_labels = task.num_labels\n",
    "    df = task.test_df if useTest == False else task.dev_df\n",
    "    ds = Dataset.from_pandas(df)\n",
    "    output = trainer.predict(ds)\n",
    "    label_ids = output.label_ids.reshape(-1)\n",
    "    predictions = output.predictions.reshape(-1, num_labels)\n",
    "    predictions = np.argmax(predictions, axis=-1)\n",
    "    mask = label_ids != ignore_index\n",
    "    \n",
    "    y_true = label_ids[mask]\n",
    "    y_pred = predictions[mask]\n",
    "    target_names = [task.index_to_label.get(ele, \"\") for ele in range(num_labels)]\n",
    "    print(target_names)\n",
    "    \n",
    "    total = 0\n",
    "    correct = 0\n",
    "    for(t, p) in zip(y_true, y_pred):\n",
    "        total = total + 1\n",
    "        if t == p:\n",
    "            correct = correct + 1\n",
    "    accuracy = correct / total\n",
    "    \n",
    "    report = classification_report(\n",
    "        y_true, y_pred,\n",
    "        target_names=target_names\n",
    "    )\n",
    "    print(report)\n",
    "    print(f'locally computed accuracy: {accuracy}')\n",
    "    return accuracy\n",
    "\n",
    "# compute loss and accuracy\n",
    "def evaluate(trainer, task, name):\n",
    "    print(f\"Evaluating on the validation dataset for task {name}:\")\n",
    "    ds = Dataset.from_pandas(task.dev_df)\n",
    "    scores = trainer.evaluate(ds)\n",
    "    acc = evaluation_classification_report(trainer, task, name, useTest = False)\n",
    "    return scores, acc"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "4fd912e2",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "STARTING EPOCH 1\n",
      "Found 2 heads\n",
      "Task 0 with 9 labels.\n",
      "Dropout is Dropout(p=0.1, inplace=False)\n",
      "Classifier layer is Linear(in_features=768, out_features=9, bias=True)\n",
      "Task 1 with 45 labels.\n",
      "Dropout is Dropout(p=0.1, inplace=False)\n",
      "Classifier layer is Linear(in_features=768, out_features=45, bias=True)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/msurdeanu/miniconda3/envs/transformers-mps/lib/python3.9/site-packages/transformers/optimization.py:306: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'train_runtime': 2.8605, 'train_samples_per_second': 85.3, 'train_steps_per_second': 0.699, 'train_loss': 0.9945342540740967, 'epoch': 1.0}\n",
      "Elapsed time for epoch 1: 0:00:02.867345\n",
      "Evaluating on the validation dataset for task NER:\n",
      "{'eval_loss': 0.3739100992679596, 'eval_runtime': 11.8579, 'eval_samples_per_second': 292.296, 'eval_steps_per_second': 2.361, 'epoch': 1.0}\n",
      "Test classification report for task NER:\n",
      "['I-MISC', 'I-ORG', 'B-ORG', 'B-LOC', 'B-PER', 'O', 'B-MISC', 'I-LOC', 'I-PER']\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "      I-MISC       0.00      0.01      0.01       216\n",
      "       I-ORG       0.07      0.11      0.09       835\n",
      "       B-ORG       0.02      0.02      0.02      1661\n",
      "       B-LOC       0.00      0.00      0.00      1668\n",
      "       B-PER       0.06      0.08      0.07      1617\n",
      "           O       0.83      0.86      0.84     38554\n",
      "      B-MISC       0.01      0.01      0.01       702\n",
      "       I-LOC       0.00      0.00      0.00       257\n",
      "       I-PER       0.03      0.00      0.00      1156\n",
      "\n",
      "    accuracy                           0.72     46666\n",
      "   macro avg       0.11      0.12      0.12     46666\n",
      "weighted avg       0.69      0.72      0.70     46666\n",
      "\n",
      "locally computed accuracy: 0.7169459563708053\n",
      "Evaluating on the validation dataset for task POS:\n",
      "{'eval_loss': 1.1830933094024658, 'eval_runtime': 20.0523, 'eval_samples_per_second': 251.242, 'eval_steps_per_second': 1.995, 'epoch': 1.0}\n",
      "Test classification report for task POS:\n",
      "['JJ', 'VBP', 'DT', 'PRP', 'RP', '-RRB-', ':', 'WP', 'FW', 'PDT', 'NNP', 'EX', 'WP$', 'MD', 'WDT', 'PRP$', \"''\", 'VBG', 'RB', 'RBS', '#', 'VBZ', 'NNPS', 'NNS', 'VBN', '-LRB-', 'RBR', 'TO', 'UH', 'SYM', 'POS', ',', '.', 'NN', 'VBD', 'CD', 'VB', 'JJR', 'IN', 'JJS', 'LS', '$', 'WRB', 'CC', '``']\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "          JJ       0.02      0.00      0.00      3707\n",
      "         VBP       0.04      0.05      0.04       802\n",
      "          DT       0.00      0.00      0.00      4816\n",
      "         PRP       0.03      0.03      0.03      1055\n",
      "          RP       0.00      0.00      0.00       182\n",
      "       -RRB-       0.00      0.01      0.00        72\n",
      "           :       0.00      0.00      0.00       241\n",
      "          WP       0.01      0.05      0.01       111\n",
      "          FW       0.00      0.00      0.00        20\n",
      "         PDT       0.00      0.00      0.00        18\n",
      "         NNP       0.33      0.18      0.23      5918\n",
      "          EX       0.00      0.02      0.00        57\n",
      "         WP$       0.00      0.00      0.00        21\n",
      "          MD       0.00      0.00      0.00       584\n",
      "         WDT       0.03      0.07      0.05       280\n",
      "        PRP$       0.00      0.00      0.00       509\n",
      "          ''       0.00      0.00      0.00       512\n",
      "         VBG       0.02      0.00      0.00       817\n",
      "          RB       0.00      0.00      0.00      2053\n",
      "         RBS       0.00      0.00      0.00        30\n",
      "           #       0.00      0.00      0.00         5\n",
      "         VBZ       0.00      0.00      0.00      1233\n",
      "        NNPS       0.00      0.00      0.00        42\n",
      "         NNS       0.03      0.00      0.01      3507\n",
      "         VBN       0.04      0.09      0.05      1099\n",
      "       -LRB-       0.00      0.00      0.00        72\n",
      "         RBR       0.00      0.00      0.00       112\n",
      "          TO       0.00      0.00      0.00      1240\n",
      "          UH       0.00      0.00      0.00         8\n",
      "         SYM       0.00      0.00      0.00         1\n",
      "         POS       0.00      0.00      0.00       549\n",
      "           ,       0.02      0.01      0.01      3061\n",
      "           .       0.00      0.00      0.00      2363\n",
      "          NN       0.22      0.40      0.28      7568\n",
      "         VBD       0.01      0.01      0.01      1836\n",
      "          CD       0.23      0.24      0.23      1924\n",
      "          VB       0.02      0.01      0.01      1580\n",
      "         JJR       0.00      0.00      0.00       198\n",
      "          IN       0.23      0.16      0.19      5854\n",
      "         JJS       0.00      0.05      0.01       127\n",
      "          LS       0.00      0.00      0.00         4\n",
      "           $       0.00      0.00      0.00       375\n",
      "         WRB       0.00      0.00      0.00       133\n",
      "          CC       0.23      0.13      0.17      1373\n",
      "          ``       0.03      0.01      0.01       531\n",
      "\n",
      "    accuracy                           0.11     56600\n",
      "   macro avg       0.03      0.03      0.03     56600\n",
      "weighted avg       0.11      0.11      0.10     56600\n",
      "\n",
      "locally computed accuracy: 0.10514134275618374\n",
      "DEV MACRO LOSS FOR EPOCH 1: 0.7785017043352127\n",
      "\n",
      "\n",
      "DEV MACRO ACC FOR EPOCH 1: 0.4110436495634945\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/msurdeanu/miniconda3/envs/transformers-mps/lib/python3.9/site-packages/sklearn/metrics/_classification.py:1327: UndefinedMetricWarning: Precision and F-score are ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, msg_start, len(result))\n",
      "/Users/msurdeanu/miniconda3/envs/transformers-mps/lib/python3.9/site-packages/sklearn/metrics/_classification.py:1327: UndefinedMetricWarning: Precision and F-score are ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, msg_start, len(result))\n",
      "/Users/msurdeanu/miniconda3/envs/transformers-mps/lib/python3.9/site-packages/sklearn/metrics/_classification.py:1327: UndefinedMetricWarning: Precision and F-score are ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, msg_start, len(result))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'input_ids': tensor([[  101,  7993, 11303,  1468,  1114, 21748,  2249,  3190,  1576,  4974,\n",
      "           102]]), 'token_type_ids': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}\n",
      "STARTING EPOCH 2\n",
      "Resuming from checkpoint bert-base-cased-mtl/mtl_model_epoch1\n",
      "Found 2 heads\n",
      "Task 0 with 9 labels.\n",
      "Dropout is Dropout(p=0.1, inplace=False)\n",
      "Classifier layer is Linear(in_features=768, out_features=9, bias=True)\n",
      "Task 1 with 45 labels.\n",
      "Dropout is Dropout(p=0.1, inplace=False)\n",
      "Classifier layer is Linear(in_features=768, out_features=45, bias=True)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/msurdeanu/miniconda3/envs/transformers-mps/lib/python3.9/site-packages/transformers/optimization.py:306: FutureWarning: This implementation of AdamW is deprecated and will be removed in a future version. Use the PyTorch implementation torch.optim.AdamW instead, or set `no_deprecation_warning=True` to disable this warning\n",
      "  warnings.warn(\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'train_runtime': 2.8854, 'train_samples_per_second': 84.563, 'train_steps_per_second': 0.693, 'train_loss': 0.8514912128448486, 'epoch': 1.0}\n",
      "Elapsed time for epoch 2: 0:00:02.890583\n",
      "Evaluating on the validation dataset for task NER:\n",
      "{'eval_loss': 0.25297117233276367, 'eval_runtime': 11.8226, 'eval_samples_per_second': 293.168, 'eval_steps_per_second': 2.368, 'epoch': 1.0}\n",
      "Test classification report for task NER:\n",
      "['I-MISC', 'I-ORG', 'B-ORG', 'B-LOC', 'B-PER', 'O', 'B-MISC', 'I-LOC', 'I-PER']\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "      I-MISC       0.00      0.00      0.00       216\n",
      "       I-ORG       0.13      0.00      0.00       835\n",
      "       B-ORG       0.00      0.00      0.00      1661\n",
      "       B-LOC       0.00      0.00      0.00      1668\n",
      "       B-PER       0.00      0.00      0.00      1617\n",
      "           O       0.83      1.00      0.90     38554\n",
      "      B-MISC       0.00      0.00      0.00       702\n",
      "       I-LOC       0.00      0.00      0.00       257\n",
      "       I-PER       0.00      0.00      0.00      1156\n",
      "\n",
      "    accuracy                           0.82     46666\n",
      "   macro avg       0.11      0.11      0.10     46666\n",
      "weighted avg       0.68      0.82      0.75     46666\n",
      "\n",
      "locally computed accuracy: 0.8240689152702182\n",
      "Evaluating on the validation dataset for task POS:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/msurdeanu/miniconda3/envs/transformers-mps/lib/python3.9/site-packages/sklearn/metrics/_classification.py:1327: UndefinedMetricWarning: Precision and F-score are ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, msg_start, len(result))\n",
      "/Users/msurdeanu/miniconda3/envs/transformers-mps/lib/python3.9/site-packages/sklearn/metrics/_classification.py:1327: UndefinedMetricWarning: Precision and F-score are ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, msg_start, len(result))\n",
      "/Users/msurdeanu/miniconda3/envs/transformers-mps/lib/python3.9/site-packages/sklearn/metrics/_classification.py:1327: UndefinedMetricWarning: Precision and F-score are ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, msg_start, len(result))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'eval_loss': 1.126561164855957, 'eval_runtime': 20.1129, 'eval_samples_per_second': 250.486, 'eval_steps_per_second': 1.989, 'epoch': 1.0}\n",
      "Test classification report for task POS:\n",
      "['JJ', 'VBP', 'DT', 'PRP', 'RP', '-RRB-', ':', 'WP', 'FW', 'PDT', 'NNP', 'EX', 'WP$', 'MD', 'WDT', 'PRP$', \"''\", 'VBG', 'RB', 'RBS', '#', 'VBZ', 'NNPS', 'NNS', 'VBN', '-LRB-', 'RBR', 'TO', 'UH', 'SYM', 'POS', ',', '.', 'NN', 'VBD', 'CD', 'VB', 'JJR', 'IN', 'JJS', 'LS', '$', 'WRB', 'CC', '``']\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "          JJ       0.03      0.00      0.00      3707\n",
      "         VBP       0.04      0.10      0.05       802\n",
      "          DT       0.10      0.00      0.00      4816\n",
      "         PRP       0.04      0.05      0.04      1055\n",
      "          RP       0.00      0.00      0.00       182\n",
      "       -RRB-       0.00      0.00      0.00        72\n",
      "           :       0.00      0.00      0.00       241\n",
      "          WP       0.01      0.03      0.01       111\n",
      "          FW       0.00      0.00      0.00        20\n",
      "         PDT       0.00      0.00      0.00        18\n",
      "         NNP       0.24      0.55      0.33      5918\n",
      "          EX       0.00      0.00      0.00        57\n",
      "         WP$       0.00      0.00      0.00        21\n",
      "          MD       0.00      0.00      0.00       584\n",
      "         WDT       0.08      0.05      0.06       280\n",
      "        PRP$       0.00      0.00      0.00       509\n",
      "          ''       0.00      0.00      0.00       512\n",
      "         VBG       0.00      0.00      0.00       817\n",
      "          RB       0.25      0.00      0.00      2053\n",
      "         RBS       0.00      0.00      0.00        30\n",
      "           #       0.00      0.00      0.00         5\n",
      "         VBZ       0.00      0.00      0.00      1233\n",
      "        NNPS       0.00      0.00      0.00        42\n",
      "         NNS       0.05      0.04      0.04      3507\n",
      "         VBN       0.06      0.03      0.04      1099\n",
      "       -LRB-       0.00      0.00      0.00        72\n",
      "         RBR       0.00      0.00      0.00       112\n",
      "          TO       0.00      0.00      0.00      1240\n",
      "          UH       0.00      0.00      0.00         8\n",
      "         SYM       0.00      0.00      0.00         1\n",
      "         POS       0.00      0.00      0.00       549\n",
      "           ,       0.15      0.04      0.07      3061\n",
      "           .       0.16      0.02      0.04      2363\n",
      "          NN       0.23      0.47      0.31      7568\n",
      "         VBD       0.01      0.01      0.01      1836\n",
      "          CD       0.40      0.38      0.39      1924\n",
      "          VB       0.07      0.03      0.05      1580\n",
      "         JJR       0.00      0.00      0.00       198\n",
      "          IN       0.25      0.43      0.31      5854\n",
      "         JJS       0.01      0.04      0.01       127\n",
      "          LS       0.00      0.00      0.00         4\n",
      "           $       0.00      0.00      0.00       375\n",
      "         WRB       0.00      0.00      0.00       133\n",
      "          CC       0.45      0.10      0.17      1373\n",
      "          ``       0.13      0.02      0.03       531\n",
      "\n",
      "    accuracy                           0.19     56600\n",
      "   macro avg       0.06      0.05      0.04     56600\n",
      "weighted avg       0.15      0.19      0.14     56600\n",
      "\n",
      "locally computed accuracy: 0.18996466431095407\n",
      "DEV MACRO LOSS FOR EPOCH 2: 0.6897661685943604\n",
      "\n",
      "\n",
      "DEV MACRO ACC FOR EPOCH 2: 0.5070167897905862\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/msurdeanu/miniconda3/envs/transformers-mps/lib/python3.9/site-packages/sklearn/metrics/_classification.py:1327: UndefinedMetricWarning: Precision and F-score are ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, msg_start, len(result))\n",
      "/Users/msurdeanu/miniconda3/envs/transformers-mps/lib/python3.9/site-packages/sklearn/metrics/_classification.py:1327: UndefinedMetricWarning: Precision and F-score are ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, msg_start, len(result))\n",
      "/Users/msurdeanu/miniconda3/envs/transformers-mps/lib/python3.9/site-packages/sklearn/metrics/_classification.py:1327: UndefinedMetricWarning: Precision and F-score are ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, msg_start, len(result))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'input_ids': tensor([[  101,  7993, 11303,  1468,  1114, 21748,  2249,  3190,  1576,  4974,\n",
      "           102]]), 'token_type_ids': tensor([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1]])}\n"
     ]
    }
   ],
   "source": [
    "from transformers import TrainingArguments\n",
    "from transformers import Trainer\n",
    "from transformers import DataCollatorForTokenClassification\n",
    "import time\n",
    "from datetime import timedelta\n",
    "\n",
    "epochs = 4\n",
    "batch_size = 32\n",
    "weight_decay = 0.01\n",
    "use_mps_device = True if str(device) == 'mps' else False\n",
    "model_name = f'{transformer_name}-mtl'\n",
    "\n",
    "data_collator = DataCollatorForTokenClassification(tokenizer)\n",
    "last_checkpoint = None\n",
    "\n",
    "for epoch in range(1, epochs + 1):\n",
    "    print(f'STARTING EPOCH {epoch}')\n",
    "    if last_checkpoint != None:\n",
    "        print(f'Resuming from checkpoint {last_checkpoint}')\n",
    "            \n",
    "    training_args = TrainingArguments(\n",
    "        output_dir=model_name,\n",
    "        log_level='error',\n",
    "        num_train_epochs=1, # one epoch at a time\n",
    "        per_device_train_batch_size=batch_size,\n",
    "        per_device_eval_batch_size=batch_size,\n",
    "        # evaluation_strategy='epoch',\n",
    "        do_eval=False, # we will evaluate each task explicitly\n",
    "        weight_decay=weight_decay,\n",
    "        resume_from_checkpoint = last_checkpoint,\n",
    "        use_mps_device = use_mps_device\n",
    "    )\n",
    "    \n",
    "    trainer = Trainer(\n",
    "        model=model,\n",
    "        args=training_args,\n",
    "        data_collator=data_collator,\n",
    "        # compute_metrics=compute_metrics,\n",
    "        train_dataset=ds['train'],\n",
    "        # eval_dataset=ds['validation'],\n",
    "        tokenizer=tokenizer\n",
    "    )\n",
    "    \n",
    "    model.summarize_heads()\n",
    "\n",
    "    start_time = time.monotonic()\n",
    "    trainer.train()\n",
    "    end_time = time.monotonic()\n",
    "    print(f\"Elapsed time for epoch {epoch}: {timedelta(seconds=end_time - start_time)}\")\n",
    "\n",
    "    ner_scores, ner_acc = evaluate(trainer, ner_task, \"NER\")\n",
    "    pos_scores, pos_acc = evaluate(trainer, pos_task, \"POS\")\n",
    "    macro_loss = (ner_scores['eval_loss'] + pos_scores['eval_loss'])/2\n",
    "    print(f'DEV MACRO LOSS FOR EPOCH {epoch}: {macro_loss}\\n\\n')\n",
    "    macro_acc = (ner_acc + pos_acc)/2\n",
    "    print(f'DEV MACRO ACC FOR EPOCH {epoch}: {macro_acc}')\n",
    "\n",
    "    # save the transformer encoder + the head for each task\n",
    "    last_checkpoint = training_args.output_dir + '/mtl_model_epoch' + str(epoch)\n",
    "    trainer.save_model(last_checkpoint)\n",
    "    \n",
    "    # export for JVM\n",
    "    model.export_model(tasks, last_checkpoint + \"_export\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "ba0ddaf3",
   "metadata": {},
   "outputs": [],
   "source": [
    "#model = TokenClassificationModel.from_pretrained('bert-base-cased-mtl/mtl_model_epoch2', local_files_only=True)\n",
    "#model.summarize_heads()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "864a1bab",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Test classification report for task NER:\n",
      "['I-MISC', 'I-ORG', 'B-ORG', 'B-LOC', 'B-PER', 'O', 'B-MISC', 'I-LOC', 'I-PER']\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "      I-MISC       0.00      0.00      0.00       346\n",
      "       I-ORG       0.04      0.00      0.00       751\n",
      "       B-ORG       0.00      0.00      0.00      1341\n",
      "       B-LOC       0.00      0.00      0.00      1837\n",
      "       B-PER       0.00      0.00      0.00      1842\n",
      "           O       0.83      1.00      0.91     42975\n",
      "      B-MISC       0.00      0.00      0.00       922\n",
      "       I-LOC       0.00      0.00      0.00       257\n",
      "       I-PER       0.00      0.00      0.00      1307\n",
      "\n",
      "    accuracy                           0.83     51578\n",
      "   macro avg       0.10      0.11      0.10     51578\n",
      "weighted avg       0.69      0.83      0.76     51578\n",
      "\n",
      "locally computed accuracy: 0.8295203381286596\n",
      "Test classification report for task POS:\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/msurdeanu/miniconda3/envs/transformers-mps/lib/python3.9/site-packages/sklearn/metrics/_classification.py:1327: UndefinedMetricWarning: Precision and F-score are ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, msg_start, len(result))\n",
      "/Users/msurdeanu/miniconda3/envs/transformers-mps/lib/python3.9/site-packages/sklearn/metrics/_classification.py:1327: UndefinedMetricWarning: Precision and F-score are ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, msg_start, len(result))\n",
      "/Users/msurdeanu/miniconda3/envs/transformers-mps/lib/python3.9/site-packages/sklearn/metrics/_classification.py:1327: UndefinedMetricWarning: Precision and F-score are ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, msg_start, len(result))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['JJ', 'VBP', 'DT', 'PRP', 'RP', '-RRB-', ':', 'WP', 'FW', 'PDT', 'NNP', 'EX', 'WP$', 'MD', 'WDT', 'PRP$', \"''\", 'VBG', 'RB', 'RBS', '#', 'VBZ', 'NNPS', 'NNS', 'VBN', '-LRB-', 'RBR', 'TO', 'UH', 'SYM', 'POS', ',', '.', 'NN', 'VBD', 'CD', 'VB', 'JJR', 'IN', 'JJS', 'LS', '$', 'WRB', 'CC', '``']\n",
      "              precision    recall  f1-score   support\n",
      "\n",
      "          JJ       0.01      0.00      0.00      7394\n",
      "         VBP       0.03      0.11      0.04      1365\n",
      "          DT       0.04      0.00      0.00     10542\n",
      "         PRP       0.04      0.06      0.05      1932\n",
      "          RP       0.00      0.00      0.00       297\n",
      "       -RRB-       0.00      0.00      0.00       186\n",
      "           :       0.00      0.00      0.00       706\n",
      "          WP       0.01      0.03      0.01       278\n",
      "          FW       0.00      0.00      0.00        12\n",
      "         PDT       0.00      0.00      0.00        43\n",
      "         NNP       0.21      0.51      0.30     11563\n",
      "          EX       0.00      0.00      0.00       109\n",
      "         WP$       0.00      0.00      0.00        24\n",
      "          MD       0.00      0.00      0.00      1204\n",
      "         WDT       0.08      0.04      0.06       550\n",
      "        PRP$       0.33      0.00      0.00       920\n",
      "          ''       0.00      0.00      0.00       834\n",
      "         VBG       0.00      0.00      0.00      1836\n",
      "          RB       0.12      0.00      0.00      3708\n",
      "         RBS       0.00      0.00      0.00        55\n",
      "           #       0.00      0.00      0.00        25\n",
      "         VBZ       0.00      0.00      0.00      2434\n",
      "        NNPS       0.00      0.00      0.00       278\n",
      "         NNS       0.06      0.05      0.05      7634\n",
      "         VBN       0.06      0.02      0.03      2729\n",
      "       -LRB-       0.00      0.00      0.00       182\n",
      "         RBR       0.00      0.00      0.00       210\n",
      "          TO       0.00      0.00      0.00      2836\n",
      "          UH       0.00      0.00      0.00        11\n",
      "         SYM       0.00      0.00      0.00        10\n",
      "         POS       0.00      0.00      0.00      1153\n",
      "           ,       0.13      0.04      0.06      6173\n",
      "           .       0.23      0.03      0.05      4974\n",
      "          NN       0.26      0.49      0.34     17300\n",
      "         VBD       0.01      0.01      0.01      4250\n",
      "          CD       0.45      0.41      0.43      5031\n",
      "          VB       0.06      0.03      0.04      3367\n",
      "         JJR       0.00      0.00      0.00       451\n",
      "          IN       0.25      0.42      0.32     12639\n",
      "         JJS       0.00      0.01      0.00       232\n",
      "          LS       0.00      0.00      0.00        14\n",
      "           $       0.20      0.00      0.00       956\n",
      "         WRB       0.00      0.00      0.00       258\n",
      "          CC       0.43      0.12      0.19      3016\n",
      "          ``       0.05      0.01      0.01       852\n",
      "\n",
      "    accuracy                           0.19    120573\n",
      "   macro avg       0.07      0.05      0.04    120573\n",
      "weighted avg       0.15      0.19      0.14    120573\n",
      "\n",
      "locally computed accuracy: 0.19330198303102686\n",
      "MTL macro accuracy: 0.5114111605798433\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/msurdeanu/miniconda3/envs/transformers-mps/lib/python3.9/site-packages/sklearn/metrics/_classification.py:1327: UndefinedMetricWarning: Precision and F-score are ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, msg_start, len(result))\n",
      "/Users/msurdeanu/miniconda3/envs/transformers-mps/lib/python3.9/site-packages/sklearn/metrics/_classification.py:1327: UndefinedMetricWarning: Precision and F-score are ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, msg_start, len(result))\n",
      "/Users/msurdeanu/miniconda3/envs/transformers-mps/lib/python3.9/site-packages/sklearn/metrics/_classification.py:1327: UndefinedMetricWarning: Precision and F-score are ill-defined and being set to 0.0 in labels with no predicted samples. Use `zero_division` parameter to control this behavior.\n",
      "  _warn_prf(average, modifier, msg_start, len(result))\n"
     ]
    }
   ],
   "source": [
    "ner_acc = evaluation_classification_report(trainer, ner_task, \"NER\", useTest = True)\n",
    "pos_acc = evaluation_classification_report(trainer, pos_task, \"POS\", useTest = True)\n",
    "macro_acc = (ner_acc + pos_acc)/2\n",
    "print(f\"MTL macro accuracy: {macro_acc}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24770d64",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}