pos_xlmroberta.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8_rzsBs_Egpn"
   },
   "outputs": [],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "V-9Pw5AwwHDH"
   },
   "outputs": [],
   "source": [
    "from google.colab import drive\n",
    "drive.mount('/content/drive')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "tmI1bkMsDJHZ"
   },
   "outputs": [],
   "source": [
    "!pip install transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "fCxwmP3eGlpu"
   },
   "outputs": [],
   "source": [
    "import math\n",
    "import pickle\n",
    "import random\n",
    "\n",
    "import numpy as np\n",
    "from sklearn.metrics import classification_report \n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset\n",
    "from torch.utils.data.dataloader import DataLoader, Sampler\n",
    "from torch.nn.utils.rnn import pad_sequence\n",
    "import transformers\n",
    "from transformers import XLMRobertaModel, XLMRobertaTokenizer\n",
    "\n",
    "\n",
    "train_data_dir = '/content/drive/My Drive/Colab Notebooks/ezafe/data/bijankhan_corpus.tsv'\n",
    "pickled_data_dir = '/content/drive/My Drive/Colab Notebooks/ezafe/data/roberta_data.pickle'\n",
    "model_dir = '/content/drive/My Drive/Colab Notebooks/ezafe/bert_model_pos/model.pickle'\n",
    "last_i_dir = '/content/drive/My Drive/Colab Notebooks/ezafe/bert_model_pos/last_i'\n",
    "\n",
    "HIDDEN_SIZE = 768\n",
    "NUM_EPOCHS = 100\n",
    "BATCH_SIZE = 8\n",
    "LEARNING_RATE = 2e-5\n",
    "DROPOUT_RATE = .1\n",
    "BEGIN_TOKEN = 0\n",
    "END_TOKEN = 2\n",
    "PAD_TOKEN = 1\n",
    "TAG_PAD_OR_MASK_TOKEN = 14\n",
    "\n",
    "random.seed(17)\n",
    "\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "print('device:', device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "Czh-4NV6gzJI"
   },
   "outputs": [],
   "source": [
    "def data_reader(directory):\n",
    "    sents, sent = [], [BEGIN_TOKEN]\n",
    "    all_pos_tags, pos_tags = [], [TAG_PAD_OR_MASK_TOKEN]\n",
    "    pos_to_id = {}\n",
    "    with open(directory) as bijankhan_corpus:\n",
    "        for line in bijankhan_corpus:\n",
    "            if line != '\\n':\n",
    "                word, pos_tag, ezafe_tag = line.strip().split('\\t')\n",
    "\n",
    "                word = word.replace('ي', 'ی').replace('ك', 'ک').replace('ة', 'ه')\n",
    "                word = tokenizer.encode(word, add_special_tokens=False)\n",
    "                sent.extend(word)\n",
    "\n",
    "                if pos_tag not in pos_to_id:\n",
    "                    pos_to_id[pos_tag] = len(pos_to_id)\n",
    "\n",
    "                pos_tags.extend([pos_to_id[pos_tag]] + [TAG_PAD_OR_MASK_TOKEN for _ in range(len(word) - 1)])\n",
    "            else:\n",
    "                assert len(sent) == len(pos_tags)\n",
    "                sents.append(sent + [END_TOKEN])\n",
    "                all_pos_tags.append(pos_tags + [TAG_PAD_OR_MASK_TOKEN])\n",
    "                \n",
    "                sent = [BEGIN_TOKEN]\n",
    "                pos_tags = [TAG_PAD_OR_MASK_TOKEN]\n",
    "\n",
    "    return sents, all_pos_tags\n",
    "\n",
    "\n",
    "class MySampler(Sampler):\n",
    "    def __init__(self, data, i=0):\n",
    "        random.shuffle(data)\n",
    "        self.seq = list(range(len(data)))[i * BATCH_SIZE:]\n",
    "\n",
    "    def __iter__(self):\n",
    "        return iter(self.seq)\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.seq)\n",
    "\n",
    "\n",
    "class POSDataset(Dataset):\n",
    "    def __init__(self, data):\n",
    "       self.samples = data\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.samples)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        sample = self.samples[idx]\n",
    "        x = torch.tensor(sample[0])\n",
    "        mask = torch.sign(x)\n",
    "        y = torch.tensor(sample[1])\n",
    "\n",
    "        return x, mask, y\n",
    "\n",
    "def collate_fn(batch):\n",
    "    xs, masks, ys = zip(*batch)\n",
    "    padded_xs = pad_sequence(xs, batch_first=True, padding_value=PAD_TOKEN)\n",
    "    padded_masks = pad_sequence(masks, batch_first=True)\n",
    "    padded_ys = pad_sequence(ys, batch_first=True, padding_value=TAG_PAD_OR_MASK_TOKEN)\n",
    "\n",
    "    return {'input_ids': padded_xs, 'attention_masks': padded_masks}, padded_ys"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "ZokGlh4adnhp"
   },
   "outputs": [],
   "source": [
    "tokenizer = XLMRobertaTokenizer.from_pretrained('xlm-roberta-base')\n",
    "\n",
    "try:\n",
    "    with open(pickled_data_dir, 'rb') as out_file:\n",
    "        sents, ezafe_tags = pickle.load(out_file)\n",
    "except FileNotFoundError:\n",
    "    sents, ezafe_tags = data_reader(train_data_dir) \n",
    "    with open(pickled_data_dir, 'wb') as out_file:\n",
    "        pickle.dump([sents, ezafe_tags], out_file)\n",
    "\n",
    "random.seed(17)\n",
    "ids = list(range(len(sents)))\n",
    "random.shuffle(ids)\n",
    "\n",
    "sents, ezafe_tags = zip(*[(sents[i], ezafe_tags[i]) for i in ids])\n",
    "\n",
    "sents, ezafe_tags = zip(*[(sent, ezafe) for sent, ezafe in zip(sents, ezafe_tags) if len(sent) <= 512])\n",
    "\n",
    "data_split_1 = int(len(sents) * .1)\n",
    "data_split_2 = int(len(sents) * .2)\n",
    "\n",
    "test_data = list(zip(sents[:data_split_1], ezafe_tags[:data_split_1]))\n",
    "valid_data = list(zip(sents[data_split_1:data_split_2], ezafe_tags[data_split_1:data_split_2]))\n",
    "train_data = list(zip(sents[data_split_2:], ezafe_tags[data_split_2:]))\n",
    "\n",
    "try:\n",
    "    with open(model_dir + '/last_i') as in_file:\n",
    "        last_i = int(in_file.read().strip())\n",
    "except FileNotFoundError:\n",
    "    last_i = 0\n",
    "\n",
    "print('last i:', last_i)\n",
    "\n",
    "train_dataset = POSDataset(train_data)\n",
    "train_sampler = MySampler(train_dataset, last_i)\n",
    "\n",
    "valid_dataset = POSDataset(valid_data)\n",
    "\n",
    "print(len(train_dataset))\n",
    "print(len(valid_dataset))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "EwFWvRPND67-"
   },
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(Net, self).__init__()\n",
    "\n",
    "        self.xlmroberta_model = xlmroberta_model\n",
    "        self.dropout = nn.Dropout(DROPOUT_RATE)\n",
    "        self.pos_classifier = nn.Linear(HIDDEN_SIZE, 15)\n",
    "\n",
    "    def forward(self, x, mask):\n",
    "        sequence_output, _ = self.xlmroberta_model(x, mask)\n",
    "        sequence_output = self.dropout(sequence_output)\n",
    "        pos_pred = self.pos_classifier(sequence_output)\n",
    "        return pos_pred\n",
    "\n",
    "try:\n",
    "    model = torch.load(model_dir + str(last_i))    \n",
    "except:\n",
    "    xlmroberta_model = XLMRobertaModel.from_pretrained('xlm-roberta-base')\n",
    "    model = Net().to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "SrptV6bQC7_e"
   },
   "outputs": [],
   "source": [
    "train_data_loader = DataLoader(dataset=train_dataset,                                                         \n",
    "                               batch_size=BATCH_SIZE,                                            \n",
    "                               collate_fn=collate_fn,\n",
    "                               sampler=train_sampler,\n",
    "                               shuffle=False,\n",
    "                               num_workers=4)\n",
    "\n",
    "valid_data_loader = DataLoader(dataset=valid_dataset,                                                         \n",
    "                              batch_size=BATCH_SIZE,                                            \n",
    "                              collate_fn=collate_fn,\n",
    "                              shuffle=False,\n",
    "                              num_workers=4)\n",
    "\n",
    "criterion = nn.CrossEntropyLoss(ignore_index=TAG_PAD_OR_MASK_TOKEN).to(device)\n",
    "optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)\n",
    "\n",
    "for epoch in range(NUM_EPOCHS):\n",
    "    for i, (x, pos_true) in enumerate(train_data_loader):\n",
    "        model.train()\n",
    "        i += last_i\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        pos_pred = model(x['input_ids'].to(device), x['attention_masks'].to(device))\n",
    "\n",
    "        loss = criterion(pos_pred.view(-1, pos_pred.shape[-1]).cpu(), pos_true.view(-1))\n",
    " \n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    " \n",
    "        if i % 2500 == 0 and (epoch != 0 or i >= 10_000):\n",
    "            print('testing...')\n",
    "            model.evaluate()\n",
    "            # ezafe_true_labels, ezafe_pred_labels = [], []\n",
    "            pos_true_labels, pos_pred_labels = [], []  \n",
    "            with torch.no_grad():\n",
    "                for x, pos_true in valid_data_loader:\n",
    "                    pos_pred = model(x['input_ids'].to(device), x['attention_masks'].to(device))     \n",
    "                    pos_true_labels.extend(pos_true.reshape(-1).numpy().tolist())\n",
    "                    pos_pred_labels.extend(torch.argmax(pos_pred, -1).reshape(-1).cpu().numpy().tolist())\n",
    " \n",
    "            pos_true_labels_, pos_pred_labels_ = np.array(pos_true_labels), np.array(pos_pred_labels)\n",
    " \n",
    "            ezafe_pos_mask_or_pad_ids = np.where(pos_true_labels_ == TAG_PAD_OR_MASK_TOKEN)[0]\n",
    "\n",
    "            pos_true_labels_ = np.delete(pos_true_labels_, ezafe_pos_mask_or_pad_ids)\n",
    "            pos_pred_labels_ = np.delete(pos_pred_labels_, ezafe_pos_mask_or_pad_ids)\n",
    " \n",
    "            pos_f1 = classification_report(pos_true_labels_, pos_pred_labels_, digits=4)\n",
    "          \n",
    "            summary = f'epoch: {epoch + 1} | step: {i} | POS_F1:\\n{pos_f1}\\n\\n'\n",
    "            print(summary)\n",
    "                            \n",
    "            # save everything\n",
    "            with open(model_dir + 'summary.txt', 'a+') as out_file:\n",
    "                out_file.write(summary + '\\n')           \n",
    "                 \n",
    "            with open(model_dir + 'last_i', 'w+') as out_file:\n",
    "                out_file.write(str(i))          \n",
    "\n",
    "            os.mkdir(model_dir + str(i))\n",
    "            torch.save(model, model_dir + str(i))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "9pxghNmAVXwr"
   },
   "outputs": [],
   "source": [
    "# testing\n",
    "test_dataset = POSDataset(test_data)\n",
    "\n",
    "test_data_loader = DataLoader(dataset=test_dataset,                                                         \n",
    "                              batch_size=BATCH_SIZE,                                            \n",
    "                              collate_fn=collate_fn,\n",
    "                              shuffle=False,\n",
    "                              num_workers=4)\n",
    "\n",
    "i = 17500\n",
    "model = XLMRobertaForTokenClassification.from_pretrained(model_dir + str(i))\n",
    "\n",
    "model.to(device)\n",
    "model.evaluate()\n",
    "\n",
    "with torch.no_grad():\n",
    "    for x, y_true in test_data_loader:\n",
    "        y_true = y_true.to(device)\n",
    "\n",
    "        y_pred = model(x['input_ids'].to(device), x['attention_masks'].to(device))[0]\n",
    "\n",
    "        test_true_labels.extend(y_true.reshape(-1).cpu().numpy().tolist())\n",
    "        test_pred_labels.extend(torch.argmax(y_pred, -1).reshape(-1).cpu().numpy().tolist())\n",
    "\n",
    "test_true_labels_, test_pred_labels_ = np.array(test_true_labels), np.array(test_pred_labels)\n",
    "test_pad_mask_ids = np.where(test_true_labels_ == TAG_PAD_OR_MASK_TOKEN)[0]\n",
    "test_true_labels_ = np.delete(test_true_labels_, test_pad_mask_ids)\n",
    "test_pred_labels_ = np.delete(test_pred_labels_, test_pad_mask_ids)\n",
    "\n",
    "test_f1 = classification_report(test_true_labels_, test_pred_labels_, digits=4)\n",
    "\n",
    "summary = f'epoch: {epoch + 1} | step: {i}:\\n\\n{test_f1}\\n'\n",
    "print(summary)"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "pos_xlmroberta.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}