downstream_pipeline_diff_seqLen_pooling_DCNN.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\monkeydc\\.conda\\envs\\561\\lib\\site-packages\\tqdm\\auto.py:22: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import json\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from pandas import read_parquet\n",
    "from transformers import BertModel, BertTokenizerFast\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torch.optim import Adam\n",
    "from tqdm import tqdm\n",
    "from sklearn.metrics import precision_recall_fscore_support\n",
    "from torch.nn.utils.rnn import pad_sequence\n",
    "import evaluate\n",
    "import pandas as pd\n",
    "from TorchCRF import CRF"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_path = \"data/mBERT/fine\"\n",
    "\n",
    "tokenizer = BertTokenizerFast.from_pretrained(model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = read_parquet(\"data/merge/train.parquet\")\n",
    "dev_data = read_parquet(\"data/merge/dev.parquet\")\n",
    "test_data = read_parquet(\"data/merge/test.parquet\")\n",
    "\n",
    "with open(\"data/merge/tags_2_idx.json\", \"r\") as f:\n",
    "    tags2idx = json.load(f)\n",
    "\n",
    "idx2tags = {idx: tag for tag, idx in tags2idx.items()}\n",
    "\n",
    "with open(\"data/merge/chars2idx.json\", \"r\") as f:\n",
    "    chars2idx = json.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "sentences_train = train_data[\"tokens\"].values.tolist()\n",
    "tags_train = train_data[\"ner_tags\"].values.tolist()\n",
    "\n",
    "sentences_dev = dev_data[\"tokens\"].values.tolist()\n",
    "tags_dev = dev_data[\"ner_tags\"].values.tolist()\n",
    "\n",
    "sentences_test = test_data[\"tokens\"].values.tolist()\n",
    "tags_test = test_data[\"ner_tags\"].values.tolist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Some weights of the model checkpoint at data/mBERT/fine were not used when initializing BertModel: ['cls.seq_relationship.weight', 'cls.predictions.transform.dense.weight', 'cls.predictions.transform.dense.bias', 'cls.predictions.decoder.weight', 'cls.predictions.transform.LayerNorm.bias', 'cls.seq_relationship.bias', 'cls.predictions.bias', 'cls.predictions.transform.LayerNorm.weight']\n",
      "- This IS expected if you are initializing BertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
      "- This IS NOT expected if you are initializing BertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "BertModel(\n",
       "  (embeddings): BertEmbeddings(\n",
       "    (word_embeddings): Embedding(119547, 768, padding_idx=0)\n",
       "    (position_embeddings): Embedding(512, 768)\n",
       "    (token_type_embeddings): Embedding(2, 768)\n",
       "    (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "    (dropout): Dropout(p=0.1, inplace=False)\n",
       "  )\n",
       "  (encoder): BertEncoder(\n",
       "    (layer): ModuleList(\n",
       "      (0): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          (intermediate_act_fn): GELUActivation()\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (1): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          (intermediate_act_fn): GELUActivation()\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (2): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          (intermediate_act_fn): GELUActivation()\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (3): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          (intermediate_act_fn): GELUActivation()\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (4): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          (intermediate_act_fn): GELUActivation()\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (5): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          (intermediate_act_fn): GELUActivation()\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (6): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          (intermediate_act_fn): GELUActivation()\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (7): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          (intermediate_act_fn): GELUActivation()\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (8): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          (intermediate_act_fn): GELUActivation()\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (9): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          (intermediate_act_fn): GELUActivation()\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (10): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          (intermediate_act_fn): GELUActivation()\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "      (11): BertLayer(\n",
       "        (attention): BertAttention(\n",
       "          (self): BertSelfAttention(\n",
       "            (query): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (key): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (value): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "          (output): BertSelfOutput(\n",
       "            (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "            (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "            (dropout): Dropout(p=0.1, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (intermediate): BertIntermediate(\n",
       "          (dense): Linear(in_features=768, out_features=3072, bias=True)\n",
       "          (intermediate_act_fn): GELUActivation()\n",
       "        )\n",
       "        (output): BertOutput(\n",
       "          (dense): Linear(in_features=3072, out_features=768, bias=True)\n",
       "          (LayerNorm): LayerNorm((768,), eps=1e-12, elementwise_affine=True)\n",
       "          (dropout): Dropout(p=0.1, inplace=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (pooler): BertPooler(\n",
       "    (dense): Linear(in_features=768, out_features=768, bias=True)\n",
       "    (activation): Tanh()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "bert = BertModel.from_pretrained(\"data/mBERT/fine\").to(device)\n",
    "bert.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pooling_embedding(batch_word_ids, batch_embeddings):\n",
    "    processed_embeddings = []\n",
    "    \n",
    "    for word_ids, embeddings in zip(batch_word_ids, batch_embeddings):\n",
    "        shrunk_embeddings = []\n",
    "        current_embeddings = []\n",
    "        previous_word_idx = None\n",
    "        for i, word_idx in enumerate(word_ids):\n",
    "            if word_idx is None:\n",
    "                continue\n",
    "            if word_idx == previous_word_idx:\n",
    "                current_embeddings.append(embeddings[i])\n",
    "            else:\n",
    "                if current_embeddings:\n",
    "                    shrunk_embeddings.append(\n",
    "                        torch.mean(torch.stack(current_embeddings), dim=0)\n",
    "                    )\n",
    "                    current_embeddings.clear()\n",
    "                current_embeddings.append(embeddings[i])\n",
    "                previous_word_idx = word_idx\n",
    "        if current_embeddings:\n",
    "            shrunk_embeddings.append(\n",
    "                torch.mean(torch.stack(current_embeddings), dim=0)\n",
    "            )\n",
    "    \n",
    "        processed_embeddings.append(torch.stack(shrunk_embeddings))\n",
    "\n",
    "    return processed_embeddings\n",
    "                    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MultilingualDataset(Dataset):\n",
    "    def __init__(self, sentences, labels):\n",
    "        self.sentences = sentences\n",
    "        self.labels = labels\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.labels)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        return self.sentences[idx].tolist(), torch.tensor(self.labels[idx].tolist(), dtype=int), idx\n",
    "    \n",
    "def collate_fn(batch):\n",
    "    sentences, labels, ids = zip(*batch)\n",
    "    # print([len(x) for x in sentences])\n",
    "    B = len(labels)\n",
    "    tokenized_inputs_list = [\n",
    "        tokenizer(sentence, is_split_into_words=True, truncation=True, return_tensors=\"pt\")\n",
    "        for sentence in sentences\n",
    "    ]\n",
    "    input_ids = [x[\"input_ids\"][0] for x in tokenized_inputs_list]\n",
    "    input_ids = pad_sequence(input_ids, batch_first=True, padding_value=0).to(device)\n",
    "    attention_mask = [x[\"attention_mask\"][0] for x in tokenized_inputs_list]\n",
    "    attention_mask = pad_sequence(attention_mask, batch_first=True, padding_value=0).to(device)\n",
    "    with torch.no_grad():\n",
    "        bert_output = bert(input_ids, attention_mask)\n",
    "    bert_embeddings = bert_output[\"last_hidden_state\"]\n",
    "    \n",
    "    word_ids = [x.word_ids() for x in tokenized_inputs_list]\n",
    "    pooled_embeddings = pooling_embedding(word_ids, bert_embeddings)\n",
    "    batch_embeddings = pad_sequence(pooled_embeddings, batch_first=True, padding_value=0)\n",
    "    # print(batch_embeddings.shape)\n",
    "    T = batch_embeddings.shape[1]   # The original length might be larger than T (at most 512)\n",
    "    labels = [x[:T] for x in labels]\n",
    "    sentences = [x[:T] for x in sentences]\n",
    "    batch_labels = pad_sequence(labels, batch_first=True, padding_value=-100)\n",
    "\n",
    "    L = max(len(word) for sentence in sentences for word in sentence)\n",
    "    batch_char_ids = torch.zeros(B, T, L, dtype=int)\n",
    "    for i in range(B):\n",
    "        for j in range(len(sentences[i])):\n",
    "            cur_word = sentences[i][j]\n",
    "            for k in range(len(cur_word)):\n",
    "                batch_char_ids[i][j][k] = chars2idx.get(cur_word[k], chars2idx[\"<unk>\"])\n",
    "    try:\n",
    "        assert batch_embeddings.shape[1] == batch_labels.shape[1], f\"batch_embeddings: {batch_embeddings.shape} batch_labels: {batch_labels.shape}\"\n",
    "        assert batch_embeddings.shape[1] == batch_char_ids.shape[1]\n",
    "    except:\n",
    "        print(T)\n",
    "        print(ids)\n",
    "        print(batch_embeddings.shape)\n",
    "        print(batch_char_ids.shape)\n",
    "    return batch_embeddings, batch_labels, batch_char_ids"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "the next 2 blocks are used for testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 17, 768]) torch.Size([2, 17]) torch.Size([2, 17, 1])\n"
     ]
    }
   ],
   "source": [
    "train_dataset = MultilingualDataset(sentences_train, tags_train)\n",
    "train_loader = DataLoader(\n",
    "    dataset=train_dataset,\n",
    "    batch_size=2,\n",
    "    collate_fn=collate_fn\n",
    ")\n",
    "for data in train_loader:\n",
    "    bert_embeddings, labels, char_ids = data\n",
    "    print(bert_embeddings.shape, labels.shape, char_ids.shape)\n",
    "    break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "original labels size: [9, 3]\n",
      "[11, 5]\n"
     ]
    }
   ],
   "source": [
    "sentences = [train_dataset[2418][0], train_dataset[2419][0]]\n",
    "labels = [train_dataset[2418][1], train_dataset[2419][1]]\n",
    "print(f\"original labels size: {[len(x) for x in labels]}\")\n",
    "tokenized_inputs_list = [\n",
    "    tokenizer(sentence, is_split_into_words=True, truncation=True, return_tensors=\"pt\")\n",
    "    for sentence in sentences\n",
    "]\n",
    "input_ids = [x[\"input_ids\"][0] for x in tokenized_inputs_list]\n",
    "print([len(x) for x in input_ids])\n",
    "input_ids = pad_sequence(input_ids, batch_first=True, padding_value=0).to(device)\n",
    "attention_mask = [x[\"attention_mask\"][0] for x in tokenized_inputs_list]\n",
    "attention_mask = pad_sequence(attention_mask, batch_first=True, padding_value=0).to(device)\n",
    "with torch.no_grad():\n",
    "    bert_output = bert(input_ids, attention_mask)\n",
    "bert_embeddings = bert_output[\"last_hidden_state\"]\n",
    "word_ids = [x.word_ids() for x in tokenized_inputs_list]\n",
    "pooled_embeddings = pooling_embedding(word_ids, bert_embeddings)\n",
    "batch_embeddings = pad_sequence(pooled_embeddings, batch_first=True, padding_value=0)\n",
    "# print(batch_embeddings.shape)\n",
    "B = 2\n",
    "T = batch_embeddings.shape[1]\n",
    "labels = [x[:T] for x in labels]\n",
    "sentences = [x[:T] for x in sentences]\n",
    "batch_labels = pad_sequence(labels, batch_first=True, padding_value=-100)\n",
    "L = max(len(word) for sentence in sentences for word in sentence)\n",
    "batch_char_ids = torch.zeros(B, T, L, dtype=int)\n",
    "for i in range(B):\n",
    "    for j in range(len(sentences[i])):\n",
    "        cur_word = sentences[i][j]\n",
    "        for k in range(len(cur_word)):\n",
    "            try: \n",
    "                batch_char_ids[i][j][k] = chars2idx.get(cur_word[k], chars2idx[\"<unk>\"])\n",
    "            except:\n",
    "                print(i, j, k)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BiLSTM model training and evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# class IDCNN(nn.Module):\n",
    "#     def __init__(self, input_dim=768, output_dim=8, hidden_dim=128, num_layers=3, kernel_size=3, dilation=[1, 2, 3]):\n",
    "#         super(IDCNN, self).__init__()\n",
    "#         self.conv_layers = nn.ModuleList()\n",
    "#         in_channels = input_dim\n",
    "#         for i in range(num_layers):\n",
    "#             dilation_rate = dilation[i % len(dilation)]\n",
    "#             padding = dilation_rate * (kernel_size - 1) // 2\n",
    "#             conv_layer = nn.Conv1d(in_channels, hidden_dim, kernel_size, padding=padding, dilation=dilation_rate)\n",
    "#             in_channels = hidden_dim\n",
    "#             self.conv_layers.append(conv_layer)\n",
    "#         self.fc = nn.Linear(hidden_dim, output_dim)\n",
    "\n",
    "#     def forward(self, x):\n",
    "#         x = x.permute(0, 2, 1)  # 调整维度顺序 [batch_size, embedding_dim, sequence_length]\n",
    "#         for conv_layer in self.conv_layers:\n",
    "#             x = torch.relu(conv_layer(x))\n",
    "#         x = torch.max(x, dim=2)[0]  # 最大池化\n",
    "#         x = self.fc(x)\n",
    "#         return x\n",
    "\n",
    "\n",
    "class IDCNN(nn.Module):\n",
    "    def __init__(self, input_dim=768, output_dim=8, num_layers=3, kernel_size=3, dilation=[1, 2, 4]):\n",
    "        super(IDCNN, self).__init__()\n",
    "        self.conv_layers = nn.ModuleList()\n",
    "        in_channels = input_dim\n",
    "        for i in range(num_layers):\n",
    "            dilation_rate = dilation[i % len(dilation)]\n",
    "            padding = dilation_rate * (kernel_size - 1) // 2\n",
    "            conv_layer = nn.Conv1d(in_channels, in_channels, kernel_size, padding=padding, dilation=dilation_rate)\n",
    "            in_channels = in_channels\n",
    "            self.conv_layers.append(conv_layer)\n",
    "        self.fc = nn.Linear(in_channels, output_dim)\n",
    "        self.crf = CRF(output_dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x = x.permute(0, 2, 1)  # 调整维度顺序 [batch_size, embedding_dim, sequence_length]\n",
    "#         print(x.size())\n",
    "        for conv_layer in self.conv_layers:\n",
    "            x = torch.relu(conv_layer(x))\n",
    "#         x = torch.max(x, dim=2)[0]  # 最大池化\n",
    "        x = x.permute(0, 2, 1)\n",
    "        x = self.fc(x)\n",
    "#         x = self.crf(x)\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_dataset = MultilingualDataset(sentences_train, tags_train)\n",
    "dev_dataset = MultilingualDataset(sentences_dev, tags_dev)\n",
    "test_dataset = MultilingualDataset(sentences_test, tags_test)\n",
    "train_loader = DataLoader(\n",
    "    dataset=train_dataset,\n",
    "    batch_size=8,\n",
    "    shuffle=True,\n",
    "    collate_fn=collate_fn\n",
    ")\n",
    "dev_loader = DataLoader(\n",
    "    dataset=dev_dataset,\n",
    "    batch_size=8,\n",
    "    shuffle=False,\n",
    "    collate_fn=collate_fn\n",
    ")\n",
    "# lstm_model = BiLSTM(\n",
    "#     input_size=768,\n",
    "#     lstm_hidden_dim=256, lstm_num_layers=2, lstm_dropout=0.33,\n",
    "#     linear_output_dim=128, label_size=len(idx2tags)\n",
    "# )\n",
    "lstm_model = IDCNN()\n",
    "\n",
    "lstm_model = lstm_model.to(device)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss(ignore_index=-100)\n",
    "optimizer = Adam(lstm_model.parameters(), lr=1e-3)\n",
    "epochs = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch: 1/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch: 1/10: 100%|████████████████████████████████████████| 10013/10013 [06:51<00:00, 24.33it/s, acc=0.864, loss=0.414]\n",
      "Epoch: 1/10: 100%|█████████████████████████████████████████████████████████████████| 5013/5013 [02:48<00:00, 29.78it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_loss: 0.4136, train_acc: 0.8642, val_loss: 0.2696, val_acc: 0.8895\n",
      "val_precision: 0.8465, val_recall: 0.8187, val_f1: 0.8310\n",
      "Epoch: 2/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch: 2/10: 100%|█████████████████████████████████████████| 10013/10013 [06:46<00:00, 24.60it/s, acc=0.89, loss=0.323]\n",
      "Epoch: 2/10: 100%|█████████████████████████████████████████████████████████████████| 5013/5013 [02:48<00:00, 29.76it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_loss: 0.3228, train_acc: 0.8903, val_loss: 0.2583, val_acc: 0.8943\n",
      "val_precision: 0.8572, val_recall: 0.8207, val_f1: 0.8369\n",
      "Epoch: 3/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch: 3/10: 100%|████████████████████████████████████████| 10013/10013 [07:16<00:00, 22.92it/s, acc=0.899, loss=0.293]\n",
      "Epoch: 3/10: 100%|█████████████████████████████████████████████████████████████████| 5013/5013 [02:46<00:00, 30.18it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_loss: 0.2930, train_acc: 0.8994, val_loss: 0.2502, val_acc: 0.9002\n",
      "val_precision: 0.8602, val_recall: 0.8314, val_f1: 0.8440\n",
      "Epoch: 4/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch: 4/10: 100%|████████████████████████████████████████| 10013/10013 [06:52<00:00, 24.28it/s, acc=0.907, loss=0.268]\n",
      "Epoch: 4/10: 100%|█████████████████████████████████████████████████████████████████| 5013/5013 [02:49<00:00, 29.65it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_loss: 0.2683, train_acc: 0.9068, val_loss: 0.2601, val_acc: 0.9000\n",
      "val_precision: 0.8596, val_recall: 0.8319, val_f1: 0.8447\n",
      "Epoch: 5/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch: 5/10: 100%|████████████████████████████████████████| 10013/10013 [06:58<00:00, 23.93it/s, acc=0.913, loss=0.249]\n",
      "Epoch: 5/10: 100%|█████████████████████████████████████████████████████████████████| 5013/5013 [02:48<00:00, 29.79it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_loss: 0.2488, train_acc: 0.9127, val_loss: 0.2590, val_acc: 0.8967\n",
      "val_precision: 0.8669, val_recall: 0.8158, val_f1: 0.8389\n",
      "Epoch: 6/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch: 6/10: 100%|████████████████████████████████████████| 10013/10013 [06:47<00:00, 24.55it/s, acc=0.918, loss=0.232]\n",
      "Epoch: 6/10: 100%|█████████████████████████████████████████████████████████████████| 5013/5013 [02:47<00:00, 29.93it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_loss: 0.2316, train_acc: 0.9181, val_loss: 0.2513, val_acc: 0.9048\n",
      "val_precision: 0.8646, val_recall: 0.8410, val_f1: 0.8525\n",
      "Epoch: 7/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch: 7/10: 100%|████████████████████████████████████████| 10013/10013 [06:47<00:00, 24.55it/s, acc=0.922, loss=0.216]\n",
      "Epoch: 7/10: 100%|█████████████████████████████████████████████████████████████████| 5013/5013 [02:47<00:00, 30.02it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_loss: 0.2161, train_acc: 0.9225, val_loss: 0.2685, val_acc: 0.9043\n",
      "val_precision: 0.8567, val_recall: 0.8466, val_f1: 0.8514\n",
      "Epoch: 8/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch: 8/10: 100%|██████████████████████████████████████████| 10013/10013 [06:57<00:00, 24.00it/s, acc=0.928, loss=0.2]\n",
      "Epoch: 8/10: 100%|█████████████████████████████████████████████████████████████████| 5013/5013 [02:57<00:00, 28.24it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_loss: 0.2003, train_acc: 0.9275, val_loss: 0.2693, val_acc: 0.9012\n",
      "val_precision: 0.8444, val_recall: 0.8492, val_f1: 0.8462\n",
      "Epoch: 9/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch: 9/10: 100%|████████████████████████████████████████| 10013/10013 [07:04<00:00, 23.61it/s, acc=0.932, loss=0.188]\n",
      "Epoch: 9/10: 100%|█████████████████████████████████████████████████████████████████| 5013/5013 [02:59<00:00, 27.88it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_loss: 0.1884, train_acc: 0.9315, val_loss: 0.2984, val_acc: 0.9041\n",
      "val_precision: 0.8698, val_recall: 0.8314, val_f1: 0.8491\n",
      "Epoch: 10/10\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epoch: 10/10: 100%|███████████████████████████████████████| 10013/10013 [07:16<00:00, 22.96it/s, acc=0.936, loss=0.176]\n",
      "Epoch: 10/10: 100%|████████████████████████████████████████████████████████████████| 5013/5013 [03:02<00:00, 27.44it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "train_loss: 0.1759, train_acc: 0.9359, val_loss: 0.3008, val_acc: 0.9039\n",
      "val_precision: 0.8694, val_recall: 0.8347, val_f1: 0.8507\n"
     ]
    }
   ],
   "source": [
    "best_loss = torch.inf\n",
    "best_acc = 0\n",
    "for epoch in range(epochs):\n",
    "    print(f'Epoch: {epoch + 1}/{epochs}')\n",
    "    train_loss = 0\n",
    "    cur_total = 0\n",
    "    correct = 0\n",
    "    amount = 0\n",
    "    lstm_model.train()\n",
    "    train_loop = tqdm(train_loader, desc=f'Epoch: {epoch + 1}/{epochs}')\n",
    "    for bert_embeddings, labels, char_ids in train_loop:\n",
    "\n",
    "        B = bert_embeddings.shape[0]\n",
    "        bert_embeddings = bert_embeddings.to(device)\n",
    "        labels = labels.to(device)\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "        logits = lstm_model(bert_embeddings)\n",
    "        logits_reshape = logits.view(-1, logits.shape[-1])\n",
    "        labels_reshape = labels.view(-1)\n",
    "        loss = criterion(logits_reshape, labels_reshape)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "        train_loss += loss.cpu().item() * B\n",
    "\n",
    "        mask = labels_reshape != -100\n",
    "        \n",
    "        logits_non_pad = logits_reshape[mask]\n",
    "        labels_non_pad = labels_reshape[mask]\n",
    "\n",
    "        _, predictions = torch.max(logits_non_pad, 1)\n",
    "        correct += sum(predictions == labels_non_pad)\n",
    "        amount += len(labels_non_pad)\n",
    "\n",
    "        cur_total += B\n",
    "        running_loss = train_loss / cur_total\n",
    "        running_acc = correct / amount\n",
    "        train_loop.set_postfix(loss=running_loss, acc=running_acc.item())\n",
    "\n",
    "    train_loss /= len(train_dataset)\n",
    "    train_acc = correct / amount\n",
    "\n",
    "    dev_loss = 0\n",
    "    y_true, y_pred = [], []\n",
    "    correct = 0\n",
    "    amount = 0\n",
    "    lstm_model.eval()\n",
    "    dev_loop = tqdm(dev_loader, desc=f'Epoch: {epoch + 1}/{epochs}')\n",
    "    with torch.no_grad():\n",
    "        for bert_embeddings, labels, char_ids in dev_loop:\n",
    "            B = bert_embeddings.shape[0]\n",
    "            bert_embeddings = bert_embeddings.to(device)\n",
    "            labels = labels.to(device)\n",
    "\n",
    "            logits = lstm_model(bert_embeddings)\n",
    "            logits_reshape = logits.view(-1, logits.shape[-1])\n",
    "            labels_reshape = labels.view(-1)\n",
    "            loss = criterion(logits_reshape, labels_reshape)\n",
    "            dev_loss += loss.cpu().item() * B\n",
    "            \n",
    "            mask = labels_reshape != -100\n",
    "            logits_non_pad = logits_reshape[mask]\n",
    "            labels_non_pad = labels_reshape[mask]\n",
    "\n",
    "            _, predictions = torch.max(logits_non_pad, 1)\n",
    "            correct += sum(predictions == labels_non_pad)\n",
    "            amount += len(labels_non_pad)\n",
    "            y_pred.extend(predictions.cpu())\n",
    "            y_true.extend(labels_non_pad.cpu())\n",
    "\n",
    "    dev_loss /= len(dev_dataset)\n",
    "    val_acc = correct / amount\n",
    "    val_precision, val_recall, val_f1, _ = precision_recall_fscore_support(y_true, y_pred, average='macro', zero_division=0)\n",
    "    print('train_loss: {:.4f}, train_acc: {:.4f}, val_loss: {:.4f}, val_acc: {:.4f}'.format(train_loss, train_acc, dev_loss, val_acc))\n",
    "    print('val_precision: {:.4f}, val_recall: {:.4f}, val_f1: {:.4f}'.format(val_precision, val_recall, val_f1))\n",
    "\n",
    "#     if dev_loss < best_loss:\n",
    "#         print(f\"update dev_loss from {best_loss} -> {dev_loss}\")\n",
    "#         best_loss = dev_loss\n",
    "#         torch.save(lstm_model.state_dict(), \"best_DCNN.pt\")\n",
    "        \n",
    "    if val_acc > best_acc:\n",
    "#         print(f\"update dev_loss from {best_loss} -> {dev_loss}\")\n",
    "        best_acc = val_acc\n",
    "        torch.save(lstm_model.state_dict(), \"best_DCNN.pt\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# state_dict = torch.load(\"best_DCNN.pt\")\n",
    "# lstm_model.load_state_dict(state_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "seqeval = evaluate.load(\"seqeval\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████| 80100/80100 [24:22<00:00, 54.77it/s]\n"
     ]
    }
   ],
   "source": [
    "train_predictions = []\n",
    "train_references = []\n",
    "lstm_model.eval()\n",
    "train_loader = DataLoader(\n",
    "    dataset=train_dataset,\n",
    "    batch_size=1,   # use batch_size=1 => no padding needed\n",
    "    shuffle=False,\n",
    "    collate_fn=collate_fn\n",
    ")\n",
    "for bert_embeddings, labels, char_ids in tqdm(train_loader):\n",
    "    bert_embeddings = bert_embeddings.to(device)\n",
    "    with torch.no_grad():\n",
    "        logits = lstm_model(bert_embeddings)\n",
    "    logit = logits[0].cpu()\n",
    "    label = labels[0]\n",
    "    _, prediction = torch.max(logit, dim=1)\n",
    "    assert len(prediction) == len(label)\n",
    "    prediction_tag, reference_tag = [], []\n",
    "    for p_idx, r_idx in zip(prediction, label):\n",
    "        prediction_tag.append(idx2tags[p_idx.item()])\n",
    "        reference_tag.append(idx2tags[r_idx.item()])\n",
    "    train_predictions.append(prediction_tag)\n",
    "    train_references.append(reference_tag)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'LOC': {'precision': 0.7509615606071075,\n",
       "  'recall': 0.8488935173131997,\n",
       "  'f1': 0.796930183675714,\n",
       "  'number': 38410},\n",
       " 'ORG': {'precision': 0.6441947565543071,\n",
       "  'recall': 0.7104731214296177,\n",
       "  'f1': 0.6757125741829428,\n",
       "  'number': 34135},\n",
       " 'PER': {'precision': 0.7513860307468829,\n",
       "  'recall': 0.8889270779629947,\n",
       "  'f1': 0.8143901127015574,\n",
       "  'number': 34914},\n",
       " 'overall_precision': 0.7182584108980069,\n",
       " 'overall_recall': 0.8179305595622517,\n",
       " 'overall_f1': 0.7648609842057172,\n",
       " 'overall_accuracy': 0.94069062728435}"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "train_result = seqeval.compute(predictions=train_predictions, references=train_references)\n",
    "train_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████| 40100/40100 [12:13<00:00, 54.69it/s]\n"
     ]
    }
   ],
   "source": [
    "dev_predictions = []\n",
    "dev_references = []\n",
    "lstm_model.eval()\n",
    "dev_loader = DataLoader(\n",
    "    dataset=dev_dataset,\n",
    "    batch_size=1,   # use batch_size=1 => no padding needed\n",
    "    shuffle=False,\n",
    "    collate_fn=collate_fn\n",
    ")\n",
    "for bert_embeddings, labels, char_ids in tqdm(dev_loader):\n",
    "    bert_embeddings = bert_embeddings.to(device)\n",
    "    with torch.no_grad():\n",
    "        logits = lstm_model(bert_embeddings)\n",
    "    logit = logits[0].cpu()\n",
    "    label = labels[0]\n",
    "    _, prediction = torch.max(logit, dim=1)\n",
    "    assert len(prediction) == len(label)\n",
    "    prediction_tag, reference_tag = [], []\n",
    "    for p_idx, r_idx in zip(prediction, label):\n",
    "        prediction_tag.append(idx2tags[p_idx.item()])\n",
    "        reference_tag.append(idx2tags[r_idx.item()])\n",
    "    dev_predictions.append(prediction_tag)\n",
    "    dev_references.append(reference_tag)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'LOC': {'precision': 0.6770175121863152,\n",
       "  'recall': 0.7782504928919788,\n",
       "  'f1': 0.724112961622013,\n",
       "  'number': 19274},\n",
       " 'ORG': {'precision': 0.5287422037422037,\n",
       "  'recall': 0.6326492537313433,\n",
       "  'f1': 0.5760475651189128,\n",
       "  'number': 16080},\n",
       " 'PER': {'precision': 0.6784286630804915,\n",
       "  'recall': 0.8303835699184536,\n",
       "  'f1': 0.7467543049595307,\n",
       "  'number': 16555},\n",
       " 'overall_precision': 0.6312136103407451,\n",
       " 'overall_recall': 0.7497736423356258,\n",
       " 'overall_f1': 0.685404339250493,\n",
       " 'overall_accuracy': 0.9023121711792634}"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dev_result = seqeval.compute(predictions=dev_predictions, references=dev_references)\n",
    "dev_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|████████████████████████████████████████████████████████████████████████████| 40100/40100 [12:25<00:00, 53.81it/s]\n"
     ]
    }
   ],
   "source": [
    "test_predictions = []\n",
    "test_references = []\n",
    "lstm_model.eval()\n",
    "test_loader = DataLoader(\n",
    "    dataset=test_dataset,\n",
    "    batch_size=1,   # use batch_size=1 => no padding needed\n",
    "    shuffle=False,\n",
    "    collate_fn=collate_fn\n",
    ")\n",
    "for bert_embeddings, labels, char_ids in tqdm(test_loader):\n",
    "    bert_embeddings = bert_embeddings.to(device)\n",
    "    with torch.no_grad():\n",
    "        logits = lstm_model(bert_embeddings)\n",
    "    logit = logits[0].cpu()\n",
    "    label = labels[0]\n",
    "    _, prediction = torch.max(logit, dim=1)\n",
    "    assert len(prediction) == len(label)\n",
    "    prediction_tag, reference_tag = [], []\n",
    "    for p_idx, r_idx in zip(prediction, label):\n",
    "        prediction_tag.append(idx2tags[p_idx.item()])\n",
    "        reference_tag.append(idx2tags[r_idx.item()])\n",
    "    test_predictions.append(prediction_tag)\n",
    "    test_references.append(reference_tag)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'LOC': {'precision': 0.6582594964470567,\n",
       "  'recall': 0.7668761817159794,\n",
       "  'f1': 0.7084287299077112,\n",
       "  'number': 19569},\n",
       " 'ORG': {'precision': 0.5075318066157761,\n",
       "  'recall': 0.5900485149686427,\n",
       "  'f1': 0.5456883344276647,\n",
       "  'number': 16902},\n",
       " 'PER': {'precision': 0.6760729816600491,\n",
       "  'recall': 0.8315697674418605,\n",
       "  'f1': 0.7458024820106373,\n",
       "  'number': 17200},\n",
       " 'overall_precision': 0.6176183887805798,\n",
       " 'overall_recall': 0.7319222671461311,\n",
       " 'overall_f1': 0.6699296525261139,\n",
       " 'overall_accuracy': 0.9019221895230236}"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "test_result = seqeval.compute(predictions=test_predictions, references=test_references)\n",
    "test_result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_tokens = []\n",
    "assert len(train_dataset) == len(train_predictions)\n",
    "for i in range(len(train_dataset)):\n",
    "    token = train_dataset[i][0]\n",
    "    prediction = train_predictions[i]\n",
    "    assert len(token) >= len(prediction)\n",
    "    if len(token) != len(prediction):\n",
    "        token = token[:len(prediction)]\n",
    "    train_tokens.append(token)\n",
    "\n",
    "dev_tokens = []\n",
    "assert len(dev_dataset) == len(dev_predictions)\n",
    "for i in range(len(dev_dataset)):\n",
    "    token = dev_dataset[i][0]\n",
    "    prediction = dev_predictions[i]\n",
    "    assert len(token) >= len(prediction)\n",
    "    if len(token) != len(prediction):\n",
    "        token = token[:len(prediction)]\n",
    "    dev_tokens.append(token)\n",
    "\n",
    "test_tokens = []\n",
    "assert len(test_dataset) == len(test_predictions)\n",
    "for i in range(len(test_dataset)):\n",
    "    token = test_dataset[i][0]\n",
    "    prediction = test_predictions[i]\n",
    "    assert len(token) >= len(prediction)\n",
    "    if len(token) != len(prediction):\n",
    "        token = token[:len(prediction)]\n",
    "    test_tokens.append(token)\n",
    "\n",
    "train_df = pd.DataFrame({\n",
    "    \"tokens\": train_tokens,\n",
    "    \"predictions\": train_predictions\n",
    "})\n",
    "dev_df = pd.DataFrame({\n",
    "    \"tokens\": dev_tokens,\n",
    "    \"predictions\": dev_predictions\n",
    "})\n",
    "test_df = pd.DataFrame({\n",
    "    \"tokens\": test_tokens,\n",
    "    \"predictions\": test_predictions\n",
    "})\n",
    "\n",
    "train_df.to_csv(\"DCNN_train.csv\")\n",
    "dev_df.to_csv(\"DCNN_dev.csv\")\n",
    "test_df.to_csv(\"DCNN_test.csv\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py375",
   "language": "python",
   "name": "561"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}